Human Notes:
We finally get there!
Cut over the existing inner-tree implementations to the new stack.
Codex:
Route the existing ordered `aten::sum` and `aten::prod` overrides through the
shared fixed-DAG kernel. Eligibility, TensorIterator canonicalization,
keepdim, and direct/out overload handling remain in the adapter. Unsupported
dtypes and geometry continue to fall through to ATen.
Nonzero input row strides and unit-stride outputs use the shared kernel.
Noncompact rows use a dynamic leading stride; compactness and alignment key
the plan, with underaligned rows selecting the bit-neutral unstaged path.
Strided outputs retain the legacy same-order implementation.
H100 uses a targeted native-suite list; B200's existing dynamic discovery
covers the same suites without scheduling them twice. Labeling covers both
implementation and test-only changes. End-to-end tests assert both reference
bits and that the ordered route actually ran.
Test Plan:
```
spin quicklint
PYTORCH_SUM_INNER_TREE=1 python test/python_native/test_sum_cutedsl.py
python test/python_native/test_sum_inner_tree_plan.py
python test/python_native/test_inner_tree_order.py
```
Authored with assistance from an AI coding agent.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/195062
Approved by: https://github.com/ngimel
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
Simon Layton