mirror of
https://github.com/pytorch/pytorch.git
synced 2026-10-28 18:56:51 +08:00
Simon Layton
6f790d4343
Human Notes: Previously all datapaths use runtime trip counts, which allows a single kernel to service all N values without recompile. We will be adding in fixed-ordering (inner-tree) reductions which instead will require compile-time trip counts. This functionality is added as an optional thing when required. Codex: Add an opt-in reduction order whose addition DAG depends only on N and reproduces the existing ATen inner-tree sum/prod bits. Static fragments, streaming carries, ascending lane butterflies, and N-specific launch plans preserve that DAG while reusing the shared load, projection, and store body. This fixed order is the stack's basis for eager/Inductor bitwise parity. The default rolled folds remain runtime-shaped and continue sharing kernels by vector class. The fixed order compiles per N, covers all reference plan shapes, and rejects explicit requests it cannot honor rather than silently using another order. Golden hashes and differential tests pin all plan shapes, four dtypes, sum/prod, signed zero, batch independence, routing, and multi-field traits. Review the static datapath in `tile.py`, the plans in `kernel_rowtile.py`, then the host and device tests. Test Plan: ``` spin quicklint python test/python_native/test_tile_datapath.py python test/python_native/test_inner_tree_order.py 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/195059 Approved by: https://github.com/ngimel ghstack dependencies: #195052, #195056, #195051
88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
# Owner(s): ["module: dsl-native-ops"]
|
|
#
|
|
# Host-only tests for TileMap, the arithmetic description of a fold order's thread map.
|
|
|
|
import math
|
|
import sys
|
|
import unittest
|
|
|
|
from torch.testing._internal.common_utils import run_tests, TEST_CUTEDSL, TestCase
|
|
|
|
|
|
if not TEST_CUTEDSL:
|
|
sys.stderr.write("CuTeDSL not available\n")
|
|
if __name__ == "__main__":
|
|
sys.exit(0)
|
|
raise unittest.SkipTest("CuTeDSL not available")
|
|
|
|
from torch._native.ops.reductions import tile
|
|
|
|
|
|
class TestTileDatapath(TestCase):
|
|
def _tm(self, **kw):
|
|
return tile.TileMap(**kw)
|
|
|
|
def test_rejects_a_partial_warp(self):
|
|
# Lane shuffles require threads_per_row=1 or whole warps.
|
|
with self.assertRaises(ValueError):
|
|
self._tm(N=512, itemsize=4, threads_per_row=48, loads=1)
|
|
|
|
def test_exact_is_derived_and_overridable(self):
|
|
# Exact tiles need no load predication; batched tiles retain bounds.
|
|
covering = self._tm(
|
|
N=256, itemsize=4, threads_per_row=32, loads=2
|
|
) # 4 * 2 * 32 == 256
|
|
self.assertTrue(covering.exact)
|
|
ragged = self._tm(N=252, itemsize=4, threads_per_row=32, loads=2)
|
|
self.assertFalse(ragged.exact)
|
|
self.assertFalse(
|
|
self._tm(N=256, itemsize=4, threads_per_row=32, loads=2, exact=False).exact,
|
|
"an explicit exact=False must win over the derivation",
|
|
)
|
|
|
|
def test_vec_override_is_what_an_order_needs(self):
|
|
# Orders override N-derived vec to fix their DAG and pad the tail.
|
|
derived = self._tm(N=252, itemsize=4, threads_per_row=32, loads=2)
|
|
self.assertEqual(derived.vec, math.gcd(252, 4))
|
|
fixed = self._tm(
|
|
N=252, itemsize=4, threads_per_row=32, loads=2, vec=4, exact=False
|
|
)
|
|
self.assertEqual(fixed.vec, 4)
|
|
|
|
def test_strides_are_the_only_difference_between_the_two_orders(self):
|
|
# Orders differ only in chunk-to-warp assignment. Both coalesce innermost stride 1
|
|
# and permute the same columns into the same registers.
|
|
kw = dict(N=1024, itemsize=4, threads_per_row=128, loads=2)
|
|
row_major = self._tm(**kw)
|
|
warp_major = self._tm(**kw, warp_major=True)
|
|
self.assertNotEqual(row_major.strides(), warp_major.strides())
|
|
for tm in (row_major, warp_major):
|
|
cols = {
|
|
tm.col_base(lane, w, l)
|
|
for lane in range(32)
|
|
for w in range(tm.nw)
|
|
for l in range(tm.loads)
|
|
}
|
|
self.assertEqual(
|
|
len(cols), 32 * tm.nw * tm.loads, "two threads share a chunk"
|
|
)
|
|
self.assertEqual(min(cols), 0)
|
|
self.assertEqual(max(cols), kw["N"] - tm.vec)
|
|
|
|
def test_align_bytes_tracks_the_load_width(self):
|
|
# Alignment enables wide loads but faults if overstated. Derived vec divides N;
|
|
# an order's explicit vec may not, so both load and declaration must narrow.
|
|
wide = self._tm(N=1024, itemsize=4, threads_per_row=32, loads=1)
|
|
self.assertTrue(wide.wide_ok)
|
|
self.assertEqual(wide.align_bytes(4), wide.vec * 4)
|
|
|
|
padded = self._tm(
|
|
N=253, itemsize=4, threads_per_row=32, loads=2, vec=4, exact=False
|
|
)
|
|
self.assertFalse(padded.wide_ok, "253 is not a multiple of the order's vec")
|
|
self.assertEqual(padded.align_bytes(4), 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run_tests()
|