[Pallas] Ordered carry store for jagged row tiles

stack-info: PR: pytorch#2719, branch: thcmbs/stack/3

Author: thcmbs

helion/_compiler/pallas/ordered_carry.py

@@ -12,6 +12,8 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    import torch
+
     from helion._compiler.inductor_lowering import CodegenState
 
 
@@ -122,27 +124,151 @@ def begin_end_vars(state: CodegenState, block_id: int) -> tuple[str, str] | None
 
 def emit_carry_store(
     state: CodegenState,
-    tensor: object,
+    tensor: torch.Tensor,
     subscript: list[object] | tuple[object, ...],
     name: str,
     idx_str: str,
     value: object,
 ) -> bool:
-    """Store to a jagged row tile with the boundary carry.
+    """Store a row tile to a jagged output, stitching the shared boundary row.
+
+    Rounding each group's rows out to a sublane-aligned window makes neighbouring
+    groups overlap on one row, and each group zeros the rows it does not own.  A
+    one-row VMEM ``carry`` keeps that shared row across groups:
+
+      * fold (first tile of a group, ``begin`` not S-aligned):
+            ``out[:S] += carry``    add the previous group's piece back in
+      * save (last tile of a group, ``end`` not S-aligned):
+            ``carry = out[-S:]``    hand this group's piece forward
 
-    Not implemented yet.  Returns False for a normal store; for a jagged-row
-    store it raises, because writing it directly would overwrite the previous
-    group's data in the row they share.
+    Saving the folded value keeps it cumulative, so more than two tiny groups in
+    one row still work.
     """
+    from helion._compiler.ast_extension import statement_from_string
+    from helion._compiler.host_function import HostFunction
+
     fn = state.device_function
     patterns = state.fx_node.meta.get("indexing_patterns") if state.fx_node else None
     if not patterns:
         return False
-    for pat in patterns:
+
+    # Find the carried (jagged) row dim, if this store has one.
+    row_dim = None
+    row_block_id = None
+    for d, pat in enumerate(patterns):
         bid = getattr(pat, "block_id", None)
         if bid is not None and bid in fn.carry_tiles:
-            raise NotImplementedError(
-                "Pallas ordered carry store for a jagged row tile is not "
-                "implemented yet."
-            )
-    return False
+            row_dim, row_block_id = d, bid
+            break
+    if row_block_id is None:
+        return False
+
+    # The fold guard below counts groups with pl.program_id(0), so the group must
+    # be the only grid dim.
+    # TODO(tcombes): thread the group's real program id through CarryRowTile so
+    # multi-grid kernels work, instead of refusing them here.
+    grid_block_ids = HostFunction.current().device_ir.grid_block_ids
+    n_grid = sum(len(bids) for bids in grid_block_ids)
+    if n_grid != 1:
+        raise NotImplementedError(
+            "Pallas ordered carry assumes the group loop is the only grid "
+            f"dimension (program_id(0)); got {n_grid} grid dims."
+        )
+
+    # The carry only handles a 2D store with the jagged row leading.
+    if row_dim != 0 or tensor.ndim != 2:
+        raise NotImplementedError(
+            "Pallas ordered carry supports only a leading-row 2D store "
+            f"(got ndim={tensor.ndim}, jagged row at dim {row_dim})."
+        )
+    col_block_id = getattr(patterns[1], "block_id", None)
+    if col_block_id is None:
+        raise NotImplementedError(
+            "Pallas ordered carry: the column dim of a jagged store must be tiled."
+        )
+    rec = fn.carry_tiles[row_block_id]
+    S = rec.sublane
+    block_row = fn.resolved_block_size(row_block_id)
+    block_col = fn.resolved_block_size(col_block_id)
+    n_cols = tensor.size(1)
+    if not (
+        isinstance(block_row, int)
+        and isinstance(block_col, int)
+        and isinstance(n_cols, int)
+    ):
+        raise NotImplementedError(
+            "Pallas ordered carry needs static block/column sizes "
+            f"(block_row={block_row!r}, block_col={block_col!r}, cols={n_cols!r})."
+        )
+    if block_row % S != 0:
+        raise NotImplementedError(
+            f"Pallas ordered carry: block_row ({block_row}) must be a multiple of the "
+            f"sublane tiling S ({S})."
+        )
+    # TODO(tcombes): block_row larger than the whole tensor (L < block_row) is not
+    # handled and crashes downstream with a shape mismatch; only L >= block_row
+    # is exercised.
+
+    # Only the sublane (first) dim of a VMEM ref can be sliced at runtime on TPU,
+    # not the column (last) dim.  So give each column tile its own boundary row,
+    # stacked on dim 0: carry[num_col_tiles*S, block_col].
+    num_col_tiles = (n_cols + block_col - 1) // block_col
+    carry = fn.carry_scratch.get(row_block_id)
+    if carry is None:
+        carry = fn.register_scratch(
+            (num_col_tiles * S, block_col),
+            tensor.dtype,
+            name_hint="carry",
+            scratch_type="vmem",
+        )
+        fn.carry_scratch[row_block_id] = carry
+
+    row_off = state.codegen.offset_var(row_block_id)
+    col_off = state.codegen.offset_var(col_block_id)
+    begin, end = rec.begin_var, rec.end_var
+    a_start = f"({begin} - {begin} % {S})"
+    a_end = f"(({end} + {S} - 1) // {S} * {S})"
+    # The first group (program_id 0) has nothing before it, so it never folds:
+    # the carry scratch is not written yet.  Later groups fold only after a
+    # predecessor saved, so the scratch is always written first.
+    # TODO(tcombes): assumes contiguous group offsets (no gaps); a gap would fold
+    # a stale carry from a non-adjacent predecessor.
+    fold_guard = (
+        f"(({row_off} == {a_start}) & ({begin} % {S} != 0) & (pl.program_id(0) != 0))"
+    )
+    save_guard = f"(({row_off} + {block_row} >= {a_end}) & ({end} % {S} != 0))"
+    col_sub = f"pl.multiple_of(({col_off}) // {block_col} * {S}, {S})"
+    # carry is a normal (non-buffered) scratch, so a partial-row slice is fine.
+    carry_slice = f"{carry}[pl.ds({col_sub}, {S}), :]"
+    # The group's last row block sits at offset (a_end - S - row_off) inside the
+    # last tile: S-aligned and within [0, block_row - S].
+    last_off = f"pl.multiple_of({a_end} - {S} - ({row_off}), {S})"
+    out_last = f"{name}[pl.ds({last_off}, {S}), :]"
+
+    # Add the carry into the store value so the output gets one full-tile write;
+    # Pallas rejects a partial write to the double-buffered output.  So pad the
+    # [S, block_col] carry up to the full [block_row, block_col] tile (carry on
+    # top, zeros below) and add the whole thing.
+    n = block_row // S
+    if n > 1:
+        promoted = (
+            f"jnp.concatenate([{carry_slice}] + "
+            f"[jnp.zeros_like({carry_slice})] * {n - 1}, 0)"
+        )
+    else:
+        promoted = carry_slice
+    state.codegen.add_statement(
+        statement_from_string(
+            f"{name}[{idx_str}] = ({{value}}) + jnp.where({fold_guard}, {promoted}, 0)",
+            value=value,  # pyrefly: ignore[bad-argument-type]
+        )
+    )
+    # Save this group's last row block for the next group's fold.  When that
+    # block is also the group's first (folded) one, saving it keeps the carry
+    # cumulative across several tiny groups in one row.
+    state.codegen.add_statement(
+        statement_from_string(
+            f"{carry_slice} = jnp.where({save_guard}, {out_last}, {carry_slice})"
+        )
+    )
+    return True

test/test_pallas_load_store.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
 import torch
+from torch.testing._internal.common_utils import instantiate_parametrized_tests
+from torch.testing._internal.common_utils import parametrize
 
 import helion
 from helion import exc
@@ -34,29 +36,104 @@ def jagged_dense_bmm(
     return out
 
 
+def _ref(off: torch.Tensor, j: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+    out = torch.empty((j.shape[0], d.shape[2]), dtype=j.dtype, device=j.device)
+    for g in range(d.shape[0]):
+        s, e = int(off[g]), int(off[g + 1])
+        if e > s:
+            out[s:e] = j[s:e] @ d[g]
+    return out
+
+
+def _inputs(offsets: list[int], D: int, K: int, dtype: torch.dtype):
+    off = torch.tensor(offsets, dtype=torch.int32, device=DEVICE)
+    L, B = offsets[-1], len(offsets) - 1
+    torch.manual_seed(0)
+    j = torch.randn((L, D), dtype=dtype, device=DEVICE)
+    d = torch.randn((B, D, K), dtype=dtype, device=DEVICE)
+    return off, j, d
+
+
+def _run(off, j, d, block_sizes):
+    return code_and_output(
+        jagged_dense_bmm,
+        (off, j, d),
+        block_sizes=block_sizes,
+        pallas_loop_type="emit_pipeline",
+    )
+
+
+# bf16 with the kernel's f32-accumulating MXU matches a single torch matmul of
+# the same group bit-for-bit (single D-tile), so the default tolerances hold.
 @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".
+    # Simple cases that show the logic.
 
-    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",
-            )
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_single_group(self, dtype: torch.dtype) -> None:
+        # One group: exercises the aligned-enclosing read and the tail mask,
+        # with no carry between groups.
+        off, j, d = _inputs([0, 20], 128, 128, dtype)
+        _code, out = _run(off, j, d, [16, 128, 128])
+        torch.testing.assert_close(out, _ref(off, j, d))
+
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_aligned_groups_no_carry(self, dtype: torch.dtype) -> None:
+        # Sublane-aligned boundaries: no over-read and no shared row, so the
+        # carry fold/save guards never fire.
+        off, j, d = _inputs([0, 16, 32], 128, 128, dtype)
+        _code, out = _run(off, j, d, [16, 128, 128])
+        torch.testing.assert_close(out, _ref(off, j, d))
+
+    def test_carry_keeps_both_groups(self) -> None:
+        # Two groups [0, 3) and [3, 16) share the [0, 16) boundary row.  With
+        # identity weights out == jagged, so a clobbered carry would be obvious.
+        off = torch.tensor([0, 3, 16], dtype=torch.int32, device=DEVICE)
+        j = (
+            torch.arange(1, 17, dtype=torch.bfloat16, device=DEVICE)[:, None]
+            .expand(16, 128)
+            .contiguous()
+        )
+        eye = torch.eye(128, dtype=torch.bfloat16, device=DEVICE)
+        _code, out = _run(off, j, torch.stack([eye, eye]), [16, 128, 128])
+        torch.testing.assert_close(out, j)
+
+    # The carry across the full configuration matrix.
 
-    def test_elementwise_store_reaches_carry(self) -> None:
-        # A non-matmul map-axis store is accepted too, so it also reaches the
-        # unimplemented carry.
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    @parametrize(
+        "offsets",
+        [
+            [0, 13, 25],  # shared boundary row
+            [0, 20, 50],  # multi-row groups
+            [0, 3, 7, 16],  # several tiny groups in one row (cumulative carry)
+        ],
+    )
+    def test_bmm_block_eq_sublane(self, dtype: torch.dtype, offsets: list[int]) -> None:
+        off, j, d = _inputs(offsets, 128, 128, dtype)
+        code, out = _run(off, j, d, [16, 128, 128])
+        self.assertIn("pltpu.emit_pipeline", code)
+        torch.testing.assert_close(out, _ref(off, j, d))
+
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_bmm_block_gt_group(self, dtype: torch.dtype) -> None:
+        # Two groups (13 rows) are smaller than block_row=32, total L=200 >> block_row.
+        off, j, d = _inputs([0, 13, 100, 113, 200], 128, 128, dtype)
+        _code, out = _run(off, j, d, [32, 128, 128])
+        torch.testing.assert_close(out, _ref(off, j, d))
+
+    @parametrize("dtype", [torch.float32, torch.bfloat16])
+    def test_bmm_multi_k_tile(self, dtype: torch.dtype) -> None:
+        # K=256 with block_col=128 gives two output-column tiles; the carry stacks
+        # the per-column-tile boundary rows along its scratch row dim.
+        off, j, d = _inputs([0, 17, 40, 71], 128, 256, dtype)
+        _code, out = _run(off, j, d, [16, 128, 128])
+        torch.testing.assert_close(out, _ref(off, j, d))
+
+    def test_elementwise_map_axis(self) -> None:
+        # The non-matmul map-axis store from commit 2 now runs: out[st] = 2 *
+        # jagged[st], carried across the shared boundary like the matmul case.
         @helion.kernel(backend="pallas")
         def jagged_scale(
             seq_offsets: torch.Tensor, jagged: torch.Tensor
@@ -74,13 +151,15 @@ def jagged_scale(
 
         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",
-            )
+        _code, out = code_and_output(
+            jagged_scale,
+            (off, j),
+            block_sizes=[16, 128],
+            pallas_loop_type="emit_pipeline",
+        )
+        torch.testing.assert_close(out, j * 2)
+
+    # Rejection paths.
 
     def test_jagged_reduction_over_row_f32(self) -> None:
         # Summing the jagged row away to a dense output reduces it (not a map
@@ -115,6 +194,45 @@ def jagged_row_sum(
         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_block_not_multiple_of_sublane_raises(self) -> None:
+        # bf16 sublane S=16; block_row=8 is not a multiple, so the carry rejects it
+        # loudly instead of clobbering the boundary with a plain store.
+        off, j, d = _inputs([0, 13, 25], 128, 128, torch.bfloat16)
+        with self.assertRaisesRegex(
+            exc.InductorLoweringError, "block_row .* must be a multiple"
+        ):
+            _run(off, j, d, [8, 128, 128])
+
+    def test_multi_grid_group_rejected(self) -> None:
+        # The fold guard keys off program_id(0), so a second grid dimension is
+        # refused rather than folding against the wrong program id.
+        @helion.kernel(backend="pallas")
+        def two_grid_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 _b, g in hl.grid([2, 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
+
+        off, j, d = _inputs([0, 13, 25], 128, 128, torch.bfloat16)
+        with self.assertRaisesRegex(exc.InductorLoweringError, "only grid dimension"):
+            code_and_output(
+                two_grid_bmm,
+                (off, j, d),
+                block_sizes=[16, 128, 128],
+                pallas_loop_type="emit_pipeline",
+            )
+
     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.
@@ -140,3 +258,6 @@ def static_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
             pallas_loop_type="emit_pipeline",
         )
         torch.testing.assert_close(out, a @ b)
+
+
+instantiate_parametrized_tests(TestPallasJaggedCarry)