[Perf][Pool] Optimize avg_pool2d kernels

[RMLYC]

Closes #1654

Summary

• Add an AvgPool2dSpatialKernel fast path for the common NCHW case with ceil_mode=False, count_include_pad=True, and no divisor_override.
• Rewrite the general AvgPool2dKernel fallback to use a flat N*C*H_out*W_out launch grid and remove the block_c tuning dimension.
• Extend AvgPool2dFixture coverage for ceil_mode, count_include_pad, and divisor_override combinations.

Test plan

• Cleared TileLang caches before the final benchmark/profiling run:

rm -rf /home/lyc/.tilelang \
  /home/lyc/Project/TileOps-workspace2/.tmp/tilelang_clear_sync_cache \
  /home/lyc/Project/TileOps-workspace2/.tmp/tilelang_ifelse_cache

• Functional tests:

CUDA_VISIBLE_DEVICES=1 conda run -n tileops-dev \
  python -m pytest tests/ops/test_pool.py -k avg_pool2d -vvs

Result:

23 passed, 28 deselected, 14 warnings in 57.08s

• Benchmark:

CUDA_VISIBLE_DEVICES=1 conda run -n tileops-dev \
  python -m pytest benchmarks/ops/bench_pool.py::test_avg_pool2d_bench -vvs

Result:

3 passed, 15 warnings in 34.03s

• Whitespace check:

git diff --check

Benchmark

Final benchmark on NVIDIA H200:

workload	kernel path	before TileOps ms	after TileOps ms	torch-ref ms	after vs before	after vs torch
vision-3x3-s2	AvgPool2dSpatialKernel	0.0505	0.0046	0.0094	10.98x faster	2.04x faster
vision-5x5-s2	AvgPool2dSpatialKernel	0.0372	0.0045	0.0102	8.27x faster	2.27x faster
ceil-divisor-bf16	AvgPool2dKernel flat general	0.0149	0.0038	0.0080	3.92x faster	2.11x faster

Autotune best configs from the final run:

workload	kernel	best config
vision-3x3-s2	AvgPool2dSpatialKernel	{"block_m": 256, "threads": 256}
vision-5x5-s2	AvgPool2dSpatialKernel	{"block_m": 512, "threads": 256}
ceil-divisor-bf16	AvgPool2dKernel	{"block_m": 256, "threads": 256}

Notes on the baseline values:

• The vision-3x3-s2 and vision-5x5-s2 "before" values are from the original TileOps kernel measured during the avg_pool2d performance investigation.
• The ceil-divisor-bf16 "before" value is the local pre-flat-general fallback re-run used for the review comparison.

Profiling summary

The benchmark improvement is consistent with fixed-config nsys and ncu data.

For vision-3x3-s2, ncu shows that TileOps executes much less work than torch:

metric	TileOps	torch
Duration	5.12 us	10.05 us
Waves/SM	1.48	1.48
Achieved Occupancy	73.61%	79.44%
SM Busy	44.53%	67.92%
Memory Throughput	628.30 GB/s	321.45 GB/s
Issued Instructions	1.09M	4.24M

Torch has higher SM busy, but it issues about 3.9x more instructions. The TileOps fast path avoids the general pooling divisor/control path and reaches about 2x higher effective memory throughput.

For vision-5x5-s2, fixed-config nsys reports:

metric	TileOps	torch
nsys median	3.808 us	9.664 us
nsys avg	3.822 us	9.672 us
Grid / Block	392 x 256	196 x 1024
Registers/thread	29	32

For ceil-divisor-bf16, the new flat general fallback is also faster than torch:

metric	TileOps	torch
nsys median	3.456 us	7.488 us
ncu Duration	4.48 us	8.96 us
Grid / Block	914 x 256	229 x 1024
Waves/SM	0.87	0.87
Achieved Occupancy	68.44%	73.10%
SM Busy	43.02%	65.50%
Eligible Warps/Scheduler	1.42	3.65
Memory Throughput	404.46 GB/s	203.60 GB/s
Issued Instructions	875K	3.28M

Torch again keeps the SMs busier, but it issues about 3.75x more instructions. The flat general kernel removes the old block_c split, launches over N*C*H_out*W_out, and completes the same workload with a shorter instruction path and higher effective memory throughput.

[gemini-code-assist]

Code Review

This pull request introduces a fast-path spatial kernel (AvgPool2dSpatialKernel) for common NCHW 2D average pooling workloads, refactoring the existing _avg_pool2d_kernel to simplify its block configuration. It also updates the forward operator to dispatch to this new spatial kernel when specific fast-path conditions are met. Feedback on the changes suggests adding a defensive check in the spatial kernel's forward method to handle cases where the output dimensions are zero, preventing potential CUDA configuration errors or crashes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[superAngGao]

Thanks for the optimization. The fast-path split and the flatter fallback look directionally good to me. I have two contract-level comments before approval:

1. Please sync the manifest with the new dispatch surface. AvgPool2dFwdOp.default_kernel_map now contains both avg_pool2d_kernel and avg_pool2d_spatial_kernel, but tileops/manifest/pool.yaml still declares only avg_pool2d_kernel. The ops design docs say default_kernel_map should match source.kernel_map verbatim, and the benchmark/reporting path is manifest-driven here, so the new spatial kernel should be listed there as well. It may also be worth exporting AvgPool2dSpatialKernel from the top-level tileops.kernels package for consistency with the other pool kernels.

2. Please double-check the custom kernel_map partial override semantics. Because dispatch_kernel() fills missing entries from default_kernel_map, a caller that passes only {"avg_pool2d_kernel": MyKernel} will still get the default AvgPool2dSpatialKernel inserted. For fast-path parameter sets, dispatch will then choose the default spatial kernel and bypass the user-provided general kernel. I think user-supplied overrides should either explicitly opt into the spatial fast path or fall back to the provided avg_pool2d_kernel; otherwise tests/experiments that override the existing kernel key can silently stop exercising their replacement. A small smoke test for this case, plus one for non-fast-path dispatch, would lock the behavior down.

The core indexing/divisor logic looks reasonable to me, and the output-size guard from the earlier bot comment appears to be addressed.

[zhen8838]

LGTM. Static review found no blocker on head 7dafac8; GitHub CI is green for gpu-smoke and preflight checks. I did not run local tests or benchmarks.

[Ibuki-wind]

Clean under a correctness-focused review. I treated manifest/export synchronization and the custom kernel_map override semantics as follow-up contract work rather than blockers for this PR.

The fast-path guard is narrow enough for the fixed-divisor spatial kernel, the general fallback keeps the PyTorch divisor/boundary semantics while flattening the launch grid, and the new output-size guard matches PyTorch error behavior for non-positive outputs. The added avg_pool2d cases cover the changed dispatch paths and the ceil/count/divisor combinations touched by the fallback rewrite.

I did not run local GPU tests or benchmarks; GitHub gpu-smoke/preflight are green and the PR includes H200 benchmark results.