[Pallas] Aligned-enclosing addressing for jagged row tiles (gate closed)

stack-info: PR: pytorch#2717, branch: thcmbs/stack/1

Author: thcmbs

helion/_compiler/backend.py

@@ -4,6 +4,7 @@
 import ast
 import base64
 import contextlib
+import enum
 import functools
 import hashlib
 from itertools import starmap
@@ -1439,6 +1440,36 @@ def reserved_launch_param_names() -> frozenset[str]:
 _PALLAS_UNSUPPORTED_DTYPES = frozenset({torch.int64, torch.uint64, torch.float64})
 
 
+class SliceAddressing(enum.Enum):
+    """How a dynamic-offset slice on a tensor dim must be emitted on TPU."""
+
+    DIRECT = enum.auto()  # offset used as-is              -> plain pl.ds
+    ALIGNED = enum.auto()  # offset rounded to sublane tile -> aligned-enclosing (+carry on a store)
+
+
+def _slice_addressing(
+    tensor: torch.Tensor, dim: int, lane_block: int | None = None
+) -> SliceAddressing:
+    """Whether a dynamic slice on ``tensor``'s ``dim`` can take any offset.
+
+    TPU tiles only the last two dims into ``(8, 128)`` blocks; any earlier dim is
+    plain row-major, so a runtime offset there reads directly (DIRECT).  A slice
+    on the 2nd-minor (sublane) dim normally must land on a tile boundary
+    (ALIGNED) -- except f32 whose access spans a single lane tile (``lane_block``
+    <= 128) is row-major-contiguous there and reads any offset too (DIRECT).
+    ``lane_block`` is the access extent on the last dim (its block size, or the
+    full width if untiled); unknown (None) stays conservative (ALIGNED).
+    """
+    if dim < tensor.ndim - 2:
+        return SliceAddressing.DIRECT  # major dim: row-major, any offset
+    if dim == tensor.ndim - 2:  # 2nd-minor (sublane) dim
+        f32 = tensor.dtype.itemsize >= 4
+        if f32 and isinstance(lane_block, int) and lane_block <= 128:
+            return SliceAddressing.DIRECT  # f32, single lane tile -> contiguous
+        return SliceAddressing.ALIGNED  # bf16 pack, or f32 spanning >1 lane tile
+    return SliceAddressing.ALIGNED  # lane dim (rare): conservative
+
+
 class PallasBackend(Backend):
     """Pallas (JAX) code generation backend for TPU."""
 
@@ -1765,6 +1796,18 @@ def _get_pallas_required_alignment(
             return 8
         return 1  # No requirements for other dimensions
 
+    def sublane_tiling(self, dtype: torch.dtype) -> int:
+        """Native sublane (2nd-minor) tile for ``dtype``: f32->8, bf16->16, i8->32.
+
+        The jagged carry slices its emit_pipeline VMEM refs at this
+        granularity, and such a ref must be accessed as a *whole* native tile:
+        a smaller slice (e.g. 8 rows of a bf16 ref, whose tile is 16) is
+        rejected by Mosaic ("E2003: unproven memory access alignment"),
+        independent of offset.
+        """
+        bitwidth = min(dtype.itemsize * 8, 32)
+        return 8 * (32 // bitwidth)
+
     fake_tensor_loads: list[tuple[torch.Tensor, list[object]]]
 
     def process_fake_tensor_load(

helion/_compiler/device_function.py

@@ -50,6 +50,7 @@
     from .generate_ast import GenerateAST
     from .indexing_strategy import IndexingStrategy
     from .program_id import ProgramIDs
+    from helion._compiler.pallas.ordered_carry import CarryRowTile
     from helion._compiler.pallas.plan_tiling import DimensionTiling
 
     _P = TypeVar("_P", bound="TensorPropertyArg")
@@ -337,6 +338,12 @@ def __init__(
         # Pallas: id(fake_tensor) → {dim: (block_id, extra_pad)} for dims
         # using pl.ds() that may need host-side padding.
         self.pallas_pad_info: dict[int, dict[int, tuple[int, int]]] = {}
+        # Pallas ordered carry: jagged row block_id -> CarryRowTile.  Filled by
+        # the emit_pipeline codegen when the tile is a legal map axis; read by
+        # the store codegen to stitch the boundary across neighbouring groups.
+        self.carry_tiles: dict[int, CarryRowTile] = {}
+        # row block_id -> carry scratch var name (allocated lazily at the store).
+        self.carry_scratch: dict[int, str] = {}
 
     def allocate_store_index(self) -> int:
         """Bump store counters and return the indexing strategy slot."""

helion/_compiler/pallas/ordered_carry.py

@@ -0,0 +1,127 @@
+"""Ordered carry for jagged row tiles on the Pallas emit_pipeline path.
+
+A ``hl.tile(s, e)`` with runtime bounds rounds its rows out to a sublane-aligned
+window and masks the extra.  Neighbouring groups then share one boundary row,
+which a small VMEM ``carry`` stitches back together.  Only valid when each
+output row comes from its own input row.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from helion._compiler.inductor_lowering import CodegenState
+
+
+@dataclasses.dataclass(frozen=True)
+class CarryRowTile:
+    """What the block-spec, mask, and store code share about one carried tile."""
+
+    block_id: int
+    begin_var: str  # host var for the tile's runtime begin (e.g. ``s``)
+    end_var: str  # host var for the tile's runtime end (e.g. ``e``)
+    sublane: int  # S = backend.sublane_tiling(dtype)
+    carry_scratch_name: str | None = None
+
+
+def is_dynamic_bound_tile(state: CodegenState, block_id: int) -> bool:
+    """True for a jagged tile: its begin and end are runtime tensor values.
+
+    A static ``hl.tile(0, K)`` keeps a constant ``end_expr``; a ``hl.tile(s, e)``
+    with runtime bounds has both ``begin_expr`` and ``end_expr`` set to None.
+    """
+    loops = state.codegen.active_device_loops.get(block_id)
+    if not loops:
+        return False
+    info = loops[-1].block_id_to_info.get(block_id)
+    if info is None:
+        return False
+    return info.begin_expr is None and info.end_expr is None
+
+
+def is_row_map_axis(state: CodegenState, block_id: int) -> bool:
+    """Whether each output row comes only from its own input row (a map axis).
+
+    True only when the row is written straight back (row i to output row i),
+    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
+
+
+def needs_ordered_carry(state: CodegenState, block_id: int) -> bool:
+    """Whether this dynamic row tile needs sublane carry stitching.
+
+    Carry stitches the shared boundary tile of a per-row map store, so it
+    applies only when this jagged tile is a straight map axis (is_row_map_axis)
+    *and* it feeds a 2-D store whose leading dim is the tile and whose trailing
+    (column) dim is tiled -- that tiled column is the block the carry folds and
+    saves into.  An untiled ``:`` column (a BlockSpec store) cannot carry and
+    does not need to; a higher-rank ``[L, H, D]`` store has no shared row; a
+    reduction or scatter axis is not a map and never carries.
+    """
+    from helion._compiler.pallas.plan_tiling import TilePattern
+    from helion.language.memory_ops import store
+
+    if not is_row_map_axis(state, block_id):
+        return False
+    for ginfo in state.codegen.codegen_graphs:
+        for node in ginfo.graph.nodes:
+            if node.target is not store:
+                continue
+            patterns = node.meta.get("indexing_patterns")
+            if not patterns or len(patterns) != 2:
+                continue
+            row_pat, col_pat = patterns
+            if getattr(row_pat, "block_id", None) != block_id:
+                continue
+            if isinstance(row_pat, TilePattern) and (
+                getattr(col_pat, "block_id", None) is not None
+            ):
+                return True
+    return False
+
+
+def begin_end_vars(state: CodegenState, block_id: int) -> tuple[str, str] | None:
+    """Return the (begin_var, end_var) host names for a dynamic row tile."""
+    loops = state.codegen.active_device_loops.get(block_id)
+    if not loops:
+        return None
+    info = loops[-1].block_id_to_info.get(block_id)
+    if info is None or info.begin_var_name is None:
+        return None
+    end_var = info.end_var_name
+    if end_var is None:
+        return None
+    return info.begin_var_name, end_var
+
+
+def emit_carry_store(
+    state: CodegenState,
+    tensor: object,
+    subscript: list[object] | tuple[object, ...],
+    name: str,
+    idx_str: str,
+    value: object,
+) -> bool:
+    """Store to a jagged row tile with the boundary carry.
+
+    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.
+    """
+    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:
+        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