[Pallas] Implement is_row_map_axis legality gate for jagged carry

stack-info: PR: pytorch#2718, branch: thcmbs/stack/2

Author: thcmbs

helion/_compiler/pallas/ordered_carry.py

@@ -48,8 +48,29 @@ def is_row_map_axis(state: CodegenState, block_id: int) -> bool:
     never summed, scattered, or shifted; the over-read and carry are only
     correct in that case.
     """
-    # TODO(implement): the map-axis check above; for now reject every jagged tile.
-    return False
+    # TODO(tcombes): conservative.  Scans how the row is indexed and only accepts
+    # a straight store; exercised only on the bmm and elementwise forms, so
+    # revisit for completeness (aliasing, multiple stores, broadcast/expand).
+    from helion._compiler.device_ir import ReductionLoopGraphInfo
+    from helion._compiler.pallas.plan_tiling import TilePattern
+    from helion.language.memory_ops import store
+
+    has_straight_store = False
+    for ginfo in state.codegen.codegen_graphs:
+        if isinstance(ginfo, ReductionLoopGraphInfo) and block_id in ginfo.block_ids:
+            return False  # the row is summed away
+        for node in ginfo.graph.nodes:
+            patterns = node.meta.get("indexing_patterns")
+            if not patterns:
+                continue
+            for dim, pat in enumerate(patterns):
+                if getattr(pat, "block_id", None) != block_id:
+                    continue
+                if not isinstance(pat, TilePattern):
+                    return False  # the row is offset or scattered, not straight
+                if node.target is store and dim == 0:
+                    has_straight_store = True
+    return has_straight_store
 
 
 def needs_ordered_carry(state: CodegenState, block_id: int) -> bool:

test/test_pallas_load_store.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import torch
+
+import helion
+from helion import exc
+from helion._testing import DEVICE
+from helion._testing import TestCase
+from helion._testing import code_and_output
+from helion._testing import onlyBackends
+from helion._testing import skipUnlessPallas
+import helion.language as hl
+
+
+# out[s:e] = jagged[s:e] @ dense[g] for each group g delimited by seq_offsets.
+# The row tile hl.tile(s, e) has runtime bounds; unaligned group boundaries make
+# adjacent groups share an output row, which the ordered carry stitches.
+@helion.kernel(backend="pallas")
+def jagged_dense_bmm(
+    seq_offsets: torch.Tensor, jagged: torch.Tensor, dense: torch.Tensor
+) -> torch.Tensor:
+    L, D = jagged.shape
+    B, D, K = dense.shape
+    out = torch.empty((L, K), dtype=jagged.dtype, device=jagged.device)
+    for g in hl.grid(B):
+        s = seq_offsets[g]
+        e = seq_offsets[g + 1]
+        for st in hl.tile(s, e):
+            for kt in hl.tile(0, K):
+                acc = hl.zeros([st, kt], dtype=torch.float32)
+                for dt in hl.tile(0, D):
+                    acc = acc + torch.matmul(jagged[st, dt], dense[g, dt, kt])
+                out[st, kt] = acc
+    return out
+
+
+@onlyBackends(["pallas"])
+@skipUnlessPallas("JAX/Pallas TPU not available")
+class TestPallasJaggedCarry(TestCase):
+    # The legality gate is in place but the carry store is not, so a legal
+    # map-axis store reaches the unimplemented carry and raises "carry".
+
+    def test_bmm_store_reaches_carry(self) -> None:
+        # Identity store on the jagged row is a map axis: the gate accepts it, so
+        # it reaches the unimplemented carry.
+        off = torch.tensor([0, 13, 25], dtype=torch.int32, device=DEVICE)
+        j = torch.randn((25, 128), dtype=torch.bfloat16, device=DEVICE)
+        d = torch.randn((2, 128, 128), dtype=torch.bfloat16, device=DEVICE)
+        with self.assertRaisesRegex(exc.InductorLoweringError, "carry"):
+            code_and_output(
+                jagged_dense_bmm,
+                (off, j, d),
+                block_sizes=[16, 128, 128],
+                pallas_loop_type="emit_pipeline",
+            )
+
+    def test_elementwise_store_reaches_carry(self) -> None:
+        # A non-matmul map-axis store is accepted too, so it also reaches the
+        # unimplemented carry.
+        @helion.kernel(backend="pallas")
+        def jagged_scale(
+            seq_offsets: torch.Tensor, jagged: torch.Tensor
+        ) -> torch.Tensor:
+            L, D = jagged.shape
+            B = seq_offsets.shape[0] - 1
+            out = torch.empty((L, D), dtype=jagged.dtype, device=jagged.device)
+            for g in hl.grid(B):
+                s = seq_offsets[g]
+                e = seq_offsets[g + 1]
+                for st in hl.tile(s, e):
+                    for dt in hl.tile(0, D):
+                        out[st, dt] = jagged[st, dt] * 2
+            return out
+
+        off = torch.tensor([0, 13, 25], dtype=torch.int32, device=DEVICE)
+        j = torch.randn((25, 128), dtype=torch.bfloat16, device=DEVICE)
+        with self.assertRaisesRegex(exc.InductorLoweringError, "carry"):
+            code_and_output(
+                jagged_scale,
+                (off, j),
+                block_sizes=[16, 128],
+                pallas_loop_type="emit_pipeline",
+            )
+
+    def test_jagged_reduction_over_row_f32(self) -> None:
+        # Summing the jagged row away to a dense output reduces it (not a map
+        # axis), so it takes no carry.  For f32 the row reads directly (single
+        # lane tile) and reduces.  bf16's aligned-enclosing reduction read is a
+        # separate, not-yet-supported case (Mosaic E2003), so it is out of scope.
+        @helion.kernel(backend="pallas")
+        def jagged_row_sum(
+            seq_offsets: torch.Tensor, jagged: torch.Tensor
+        ) -> torch.Tensor:
+            L, D = jagged.shape
+            B = seq_offsets.shape[0] - 1
+            out = torch.zeros((B, D), dtype=torch.float32, device=jagged.device)
+            for g in hl.grid(B):
+                s = seq_offsets[g]
+                e = seq_offsets[g + 1]
+                for dt in hl.tile(0, D):
+                    acc = hl.zeros([dt], dtype=torch.float32)
+                    for st in hl.tile(s, e):
+                        acc = acc + jagged[st, dt].to(torch.float32).sum(0)
+                    out[g, dt] = acc
+            return out
+
+        off = torch.tensor([0, 13, 25], dtype=torch.int32, device=DEVICE)
+        j = torch.randn((25, 128), dtype=torch.float32, device=DEVICE)
+        _code, out = code_and_output(
+            jagged_row_sum,
+            (off, j),
+            block_sizes=[128, 16],
+            pallas_loop_type="emit_pipeline",
+        )
+        expected = torch.stack([j[0:13].sum(0), j[13:25].sum(0)])
+        torch.testing.assert_close(out, expected, rtol=2e-3, atol=2e-3)
+
+    def test_non_jagged_emit_pipeline_unaffected(self) -> None:
+        # Static (compile-time) tile bounds never trip the gate: a plain
+        # emit_pipeline matmul still lowers and runs correctly.
+        @helion.kernel(backend="pallas", static_shapes=True)
+        def static_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+            M, K = a.shape
+            K, N = b.shape
+            out = torch.empty((M, N), dtype=a.dtype, device=a.device)
+            for mt in hl.tile(M):
+                for nt in hl.tile(N):
+                    acc = hl.zeros([mt, nt], dtype=torch.float32)
+                    for kt in hl.tile(K):
+                        acc = acc + torch.matmul(a[mt, kt], b[kt, nt])
+                    out[mt, nt] = acc.to(a.dtype)
+            return out
+
+        a = torch.randn((64, 128), dtype=torch.bfloat16, device=DEVICE)
+        b = torch.randn((128, 256), dtype=torch.bfloat16, device=DEVICE)
+        _code, out = code_and_output(
+            static_matmul,
+            (a, b),
+            block_sizes=[32, 128, 128],
+            pallas_loop_type="emit_pipeline",
+        )
+        torch.testing.assert_close(out, a @ b)