Human Notes:
Complete axis coverage and add the multi-CTA path for large rows.
Codex:
Complete axis coverage on the shared reduction body. The general driver
represents TensorIterator geometry as runtime extent/stride pairs, allowing
one structural kernel to serve arbitrary dimensions and layouts. The column
driver vectorizes kept-axis outputs and splits tall reduced axes into
partials.
The cross-CTA driver serves few-row, large-N reductions with two fused
launches. Exact divisor splits use the optimized row body; the general
ragged split preserves coverage and global indices for prime lengths. The
dispatcher now routes row, column, cross-CTA, reduce-all, and arbitrary-axis
cases without adding a second general device kernel.
Review `tile.py` and `kernel_general.py` before `kernel_coltile.py` and
`kernel_xcta.py`.
Test Plan:
```
spin quicklint
python test/python_native/test_kernel_general.py
python test/python_native/test_kernel_xcta.py
python test/python_native/test_kernel_rowtile.py
python test/python_native/test_kernel_coltile.py
```
Authored with assistance from an AI coding agent.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/195051
Approved by: https://github.com/drisspg
ghstack dependencies: #195052, #195056
Human Notes:
Optimization for tall-and-narrow reductions - large M, small N. Potentially use TMA
for these cases where that has been measured to be faster.
Codex:
Add a one-thread-per-row configuration for tall, narrow reductions, avoiding
the idle lanes imposed by a full-warp row mapping. For wider narrow rows, TMA
stages whole rows into shared memory and rotates reads to avoid bank
conflicts.
The measured row-count and chunk gates retain direct loads where they are
faster. TMA additionally requires SM90+, fp32, a power-of-two row length, and
a lane stride at the measured direct-load cliff. The shared body validates
those requirements even when callers force the option.
Test Plan:
```
spin quicklint
python test/python_native/test_kernel_rowtile.py
```
Authored with assistance from an AI coding agent.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/195056
Approved by: https://github.com/drisspg
ghstack dependencies: #195052
Human Notes:
Add the shared reduction body and the common contiguous-row fast path.
Codex:
Add `tile.TileReduce`, the shared reduction datapath responsible for vector
loads, rolled folds, cross-thread merging, projection, and raw partial stores.
The row driver owns launch policy, geometry, and plan caching, and shares one
compiled kernel across each runtime vector class.
The body supports multiple accumulator fields and one or two projected
outputs. It is structured for later axes without introducing an
axis-specific device kernel that a later commit must replace.
Review `tile.py` before `kernel_rowtile.py` and its tests.
Test Plan:
```
spin quicklint
python test/python_native/test_kernel_rowtile.py
```
Authored with assistance from an AI coding agent.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/195052
Approved by: https://github.com/drisspg, https://github.com/ngimel
Simon Layton