[FlagGems Operator Development Competition] Add avg_pool3d operator

[goldenfox2025]

Summary

This PR implements the forward path of torch.nn.functional.avg_pool3d in Triton.

Supported API:

torch.nn.functional.avg_pool3d(
    input,
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode=False,
    count_include_pad=True,
    divisor_override=None,
) -> Tensor

Supported dtypes:

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

Supported input ranks:

• 4D: (C, D, H, W)
• 5D: (N, C, D, H, W)

Implementation

The implementation uses one Triton kernel. Each program handles a block of output elements, maps each flattened output offset to (n, c, od, oh, ow), gathers the valid input values inside the 3D pooling window, accumulates them in float32, computes the divisor, and stores the averaged result.

Handled features:

Feature	Support
kernel_size	int or 3-element sequence
stride	int, 3-element sequence, None, or empty sequence
padding	int or 3-element sequence
ceil_mode	supported
count_include_pad	supported
divisor_override	supported
non-contiguous input	supported through input.contiguous()
invalid inputs	explicit validation

The kernel uses float32 accumulation for numerical stability. The output is cast back to the original input dtype before storing.

Accuracy Validation

Command:

PYTHONPATH=src python -m pytest tests/test_avg_pool3d.py --ref cpu

Result:

40 passed, 39 warnings

Coverage includes:

• 4D and 5D inputs
• fp16, fp32, bf16
• scalar and tuple kernel_size / stride / padding
• ceil_mode=True/False
• count_include_pad=True/False
• divisor_override
• non-contiguous input
• NaN / Inf special values
• invalid input checks
• Aten registration through torch.nn.functional.avg_pool3d

Benchmark

Command:

PYTHONPATH=src python -m pytest benchmark/test_avg_pool3d.py -s --level comprehensive --warmup 20 --iter 50 --record log

Benchmark result summary:

DType	Speedup Range
float16	1.405x - 2.226x
float32	0.966x - 1.153x
bfloat16	1.417x - 2.249x

Most benchmark cases are faster than PyTorch. The only slightly slower case is the fp32 ceil_mode=True case, where FlagGems reaches 0.966x of PyTorch performance.

Notes

• The implementation focuses on forward acceleration for avg_pool3d.
• The code follows PyTorch behavior for output shape and divisor calculation.
• divisor_override=0 is rejected explicitly.
• Padding larger than half of the kernel size is rejected.
• The current design uses a simple generic gather-style kernel instead of multiple specialized paths.

[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.