[FlagGems Operator Development Competition] Optimize conv_transpose2d forward

[goldenfox2025]

Summary

This PR implements and optimizes the forward path of torch.nn.functional.conv_transpose2d in Triton.

Supported API:

torch.nn.functional.conv_transpose2d(
    input,
    weight,
    bias=None,
    stride=1,
    padding=0,
    output_padding=0,
    groups=1,
    dilation=1,
) -> Tensor

Supported dtypes:

• torch.float16
• torch.float32
• torch.bfloat16

Autograd behavior:

• torch.no_grad() forward uses the optimized Triton implementation.
• When gradients are enabled, the operator redispatches to Aten for correctness.
• This PR focuses on forward acceleration, matching the competition schema for conv_transpose2d.

Implementation

The implementation uses three main computation strategies:

Path	Use Case	Notes
direct gather	generic fallback, stride 1, small overlapping large-stride cases	single kernel launch, each output point gathers valid input/kernel contributions
residue gather	stride 2 style cases, plus large stride3 overlap outputs	multiple residue launches; compile-time residue constants remove invalid kernel taps
scatter sparse write	large stride with no spatial overlap, where effective_kernel <= stride	initializes output, then writes only non-zero transposed-convolution contributions

There are several Triton wrapper entry points, but most of them share the same residue body. The split is intentional: it isolates autotune candidate pools by dtype and kernel-size bucket so one shape does not poison another shape's autotune result.

Residue autotune buckets:

• fp32 normal kernel
• fp32 large kernel
• fp16/bf16 normal kernel
• fp16/bf16 large kernel

Residue dispatch condition:

stride_product = stride_h * stride_w
large_stride3_output = (
    stride_h == 3
    and stride_w == 3
    and N * out_h * out_w >= 131072
)
use_residue = stride_product <= 4 or large_stride3_output

The stride3 large-output case is included because direct gather wastes substantial work on the expanded output grid. Small stride3 cases stay on direct/scatter paths to avoid the launch overhead of 9 residue kernels.

Scatter path condition:

stride_h * stride_w > 4
effective_kernel_h <= stride_h
effective_kernel_w <= stride_w
effective_kernel = dilation * (kernel - 1) + 1

This avoids atomics because neighboring input points do not overlap spatially under this condition.

Accuracy Validation

Command:

export PATH="/usr/lib/wsl/lib:$PATH"
export GEMS_VENDOR=nvidia
export PYTHONPATH=src
python -m pytest tests/test_conv_transpose2d.py -q

Result:

129 passed, 338 deselected, 37 warnings

Test Coverage Checklist

Category	Covered Cases
Input scale	1x1 small case, 24x24 regular cases, 64x64 and 128x128 larger benchmark shapes
DType	float16, float32, bfloat16 forward
API parameters	stride 1, stride 2, stride 3, stride 4, stride 8, padding, output_padding, dilation, groups, bias/no-bias
Groups	groups 1, 2, 4
Kernel sizes	1x1, 2x2, 3x3, 3x2, 5x5
Layout	contiguous and non-contiguous input/weight cases
Edge cases	1x1 tensors, asymmetric stride/padding/output_padding, grouped/depthwise-like cases, stride3 overlapping large-output residue case
Invalid inputs	bad ndim, bad groups, channel mismatch, bad bias shape, invalid stride/padding/output_padding/dilation, bad pair length
Optimized branches	direct, stride2 residue, stride3 large-output residue, large-kernel residue, scatter sparse-write path

Benchmark

Command:

export PATH="/usr/lib/wsl/lib:$PATH"
export GEMS_VENDOR=nvidia
export PYTHONPATH=src
python -m pytest benchmark/test_conv_transpose2d.py -s --level core

Core benchmark shapes:

ID	Input	Weight	Stride	Padding	Output Padding	Groups	Dilation	Main Path
1	[16, 32, 24, 24]	[32, 12, 3, 3]	1	1	0	2	1	direct
2	[16, 32, 24, 24]	[32, 12, 3, 3]	2	1	1	2	1	residue
3	[32, 64, 64, 64]	[64, 32, 3, 3]	2	1	1	1	1	residue
4	[32, 64, 128, 128]	[64, 32, 5, 5]	2	2	1	1	1	large-kernel residue
5	[4, 32, 128, 128]	[32, 32, 5, 5]	3	2	0	1	1	stride3 large-output residue
6	[16, 32, 24, 24]	[32, 12, 3, 3]	3	1	0	2	1	scatter
7	[16, 32, 24, 24]	[32, 12, 3, 3]	4	1	0	2	1	scatter
8	[8, 32, 32, 32]	[32, 32, 3, 3]	4	1	0	1	1	scatter
9	[4, 32, 32, 32]	[32, 32, 3, 3]	8	1	0	1	1	scatter

Benchmark result:

DType	Input / Weight	Params	Main Path	PyTorch (ms)	FlagGems (ms)	Speedup
float16	16x32x24x24 / 32x12x3x3	s1 p1 op0 g2 d1	direct	0.156240	0.017408	8.975x
float16	16x32x24x24 / 32x12x3x3	s2 p1 op1 g2 d1	residue	0.163136	0.050144	3.253x
float16	32x64x64x64 / 64x32x3x3	s2 p1 op1 g1 d1	residue	0.745984	0.529408	1.409x
float16	32x64x128x128 / 64x32x5x5	s2 p2 op1 g1 d1	large-kernel residue	3.836928	4.241408	0.905x
float16	4x32x128x128 / 32x32x5x5	s3 p2 op0 g1 d1	stride3 residue	0.673280	0.546816	1.231x
float16	16x32x24x24 / 32x12x3x3	s3 p1 op0 g2 d1	scatter	0.179264	0.049152	3.647x
float16	16x32x24x24 / 32x12x3x3	s4 p1 op0 g2 d1	scatter	0.364608	0.067584	5.395x
float16	8x32x32x32 / 32x32x3x3	s4 p1 op0 g1 d1	scatter	0.134144	0.083968	1.598x
float16	4x32x32x32 / 32x32x3x3	s8 p1 op0 g1 d1	scatter	0.223232	0.110080	2.028x
float32	16x32x24x24 / 32x12x3x3	s1 p1 op0 g2 d1	direct	0.137216	0.034816	3.941x
float32	16x32x24x24 / 32x12x3x3	s2 p1 op1 g2 d1	residue	0.163840	0.112640	1.455x
float32	32x64x64x64 / 64x32x3x3	s2 p1 op1 g1 d1	residue	2.267136	2.153472	1.053x
float32	32x64x128x128 / 64x32x5x5	s2 p2 op1 g1 d1	large-kernel residue	19.229696	12.160000	1.581x
float32	4x32x128x128 / 32x32x5x5	s3 p2 op0 g1 d1	stride3 residue	2.868224	1.605632	1.786x
float32	16x32x24x24 / 32x12x3x3	s3 p1 op0 g2 d1	scatter	0.215136	0.091136	2.361x
float32	16x32x24x24 / 32x12x3x3	s4 p1 op0 g2 d1	scatter	0.439280	0.121888	3.604x
float32	8x32x32x32 / 32x32x3x3	s4 p1 op0 g1 d1	scatter	0.207360	0.154624	1.341x
float32	4x32x32x32 / 32x32x3x3	s8 p1 op0 g1 d1	scatter	1.360896	0.238592	5.704x
bfloat16	16x32x24x24 / 32x12x3x3	s1 p1 op0 g2 d1	direct	0.111616	0.017408	6.412x
bfloat16	16x32x24x24 / 32x12x3x3	s2 p1 op1 g2 d1	residue	0.294080	0.044032	6.679x
bfloat16	32x64x64x64 / 64x32x3x3	s2 p1 op1 g1 d1	residue	0.734208	0.503808	1.457x
bfloat16	32x64x128x128 / 64x32x5x5	s2 p2 op1 g1 d1	large-kernel residue	3.869632	4.144128	0.934x
bfloat16	4x32x128x128 / 32x32x5x5	s3 p2 op0 g1 d1	stride3 residue	0.674816	0.565248	1.194x
bfloat16	16x32x24x24 / 32x12x3x3	s3 p1 op0 g2 d1	scatter	0.872448	0.050080	17.421x
bfloat16	16x32x24x24 / 32x12x3x3	s4 p1 op0 g2 d1	scatter	1.398784	0.066560	21.015x
bfloat16	8x32x32x32 / 32x32x3x3	s4 p1 op0 g1 d1	scatter	0.134144	0.084544	1.587x
bfloat16	4x32x32x32 / 32x32x3x3	s8 p1 op0 g1 d1	scatter	0.221696	0.110592	2.005x

All core benchmark cases are above the competition threshold of 0.9x.

Notes

• The optimized forward path is active under torch.no_grad().
• Gradient-enabled execution redispatches to Aten to preserve autograd correctness.
• The scatter path is intentionally limited to non-overlapping large-stride cases, avoiding atomic adds and reducing correctness risk.
• Overlapping large stride3 outputs use residue dispatch only above a size threshold. This keeps small cases on lower-overhead paths while fixing the large-output direct fallback weakness.
• The multiple Triton wrappers are mostly autotune buckets over shared compute bodies rather than separate algorithm implementations.

[douxetpur]

Hi, thank you for your participation in the FlagOS Open Computing Competition 🙏

After review, we have chosen to move forward with a different implementation for this operator. As a result, this PR will be closed.

We appreciate your time and contribution, and hope to see more of your submissions in the future.