forked from ccf-ai-infra/TileOPs-Metax
019c128a99
## Summary - add generic TileLang average pooling kernels for 1D, 2D, and 3D workloads under the TileOPs L1/L2 architecture - export the new avg pool ops and add correctness coverage against PyTorch - add avg pool benchmarks and move PyTorch layout permutes outside `ref_program` so baseline timing excludes data-layout conversion ## Validation - `python -m pytest tests/ops/test_avg_pool1d.py tests/ops/test_avg_pool2d.py tests/ops/test_avg_pool3d.py -v` - `python -m pytest benchmarks/ops/bench_avg_pool1d.py -v` - `python -m pytest benchmarks/ops/bench_avg_pool2d.py -v` - `python -m pytest benchmarks/ops/bench_avg_pool3d.py -v` - reran benchmark comparisons after removing baseline permute cost ## Benchmark Snapshot Source: `profile_run.log` generated on 2026-04-01 17:13:52. Values below follow the rounded display precision in that report. ### avg_pool1d | Case | TileOps latency (ms) | Torch latency (ms) | TileOps TFLOPS | Torch TFLOPS | TileOps BW (TB/s) | Torch BW (TB/s) | | --- | ---: | ---: | ---: | ---: | ---: | ---: | | audio-downsample-fp16 | 0.01 | 0.02 | 0.42 | 0.14 | 0.84 | 0.28 | | long-temporal-fp16 | 0.03 | 0.10 | 0.75 | 0.20 | 1.50 | 0.39 | | ceil-bf16 | 0.00 | 0.01 | 0.24 | 0.15 | 0.36 | 0.23 | ### avg_pool2d | Case | TileOps latency (ms) | Torch latency (ms) | TileOps TFLOPS | Torch TFLOPS | TileOps BW (TB/s) | Torch BW (TB/s) | | --- | ---: | ---: | ---: | ---: | ---: | ---: | | vision-3x3-s2 | 0.01 | 0.01 | 0.31 | 0.27 | 0.34 | 0.30 | | vision-5x5-s2 | 0.01 | 0.01 | 0.56 | 0.36 | 0.22 | 0.14 | | ceil-divisor-bf16 | 0.01 | 0.01 | 0.34 | 0.32 | 0.22 | 0.21 | ### avg_pool3d | Case | TileOps latency (ms) | Torch latency (ms) | TileOps TFLOPS | Torch TFLOPS | TileOps BW (TB/s) | Torch BW (TB/s) | | --- | ---: | ---: | ---: | ---: | ---: | ---: | | video-2x2x2 | 0.00 | 0.01 | 0.40 | 0.12 | 0.90 | 0.26 | | ceil-video | 0.01 | 0.02 | 0.28 | 0.16 | 0.21 | 0.12 | | divisor-bf16 | 0.00 | 0.01 | 0.08 | 0.05 | 0.10 | 0.07 | Closes #678. |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| avg_pool1d.py | ||
| avg_pool2d.py | ||
| avg_pool3d.py | ||
| common.py | ||