From 8e05b14eb5a99b761b5f8f8d5079d4d0157b8ee9 Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Fri, 24 Jul 2026 08:41:14 -0700 Subject: [PATCH 1/7] [Pallas] Elide redundant fast-math matmul masks --- docs/api/settings.md | 5 +- helion/_compiler/aten_lowering.py | 13 ++++ helion/_compiler/generate_ast.py | 13 +--- helion/_compiler/pallas/compact_worklist.py | 41 ++++-------- helion/_compiler/pallas/tracing_ops.py | 22 ++----- helion/runtime/settings.py | 3 +- test/test_masking.py | 23 +++++++ test/test_pallas_worklist.py | 72 ++++++++++++++++++++- 8 files changed, 133 insertions(+), 59 deletions(-) diff --git a/docs/api/settings.md b/docs/api/settings.md index e7034d9b0..414fba106 100644 --- a/docs/api/settings.md +++ b/docs/api/settings.md @@ -109,7 +109,8 @@ def my_kernel(x: torch.Tensor) -> torch.Tensor: If ``True``, enable fast math approximations. This activates both Helion-level optimizations (e.g. fast sigmoid) and Inductor-level fast math (flush-to-zero exp, fast online softmax, - etc.). May reduce numerical precision. Default is ``False``. Controlled by ``HELION_FAST_MATH=1``. + etc.). May reduce numerical precision and change NaN/Inf behavior. Default is ``False``. + Controlled by ``HELION_FAST_MATH=1``. .. autoattribute:: Settings.persistent_reserved_sms @@ -319,7 +320,7 @@ Built-in values for ``HELION_AUTOTUNER`` include ``"LFBOTreeSearch"`` (default), | ``JAX_DEFAULT_MATMUL_PRECISION`` | ``dot_precision`` | Accepts JAX matmul precision values for Pallas dot products (``"default"``, ``"high"``, ``"highest"``, etc., mapped to Helion ``DotPrecision``). On TPU these values emit Pallas default precision. This variable does not apply to the Triton backend. | | ``HELION_INDEX_DTYPE`` | ``index_dtype`` | Choose the index dtype (accepts any ``torch.`` name, e.g. ``int64``), or set to ``auto``/unset to allow Helion to pick ``int32`` vs ``int64`` based on input sizes. | | ``HELION_STATIC_SHAPES`` | ``static_shapes`` | Set to ``0``/``false`` to disable global static shape specialization. | -| ``HELION_FAST_MATH`` | ``fast_math`` | Set to ``1`` to enable fast math approximations (Helion-level and Inductor-level). May reduce numerical precision. | +| ``HELION_FAST_MATH`` | ``fast_math`` | Set to ``1`` to enable fast math approximations (Helion-level and Inductor-level). May reduce numerical precision and change NaN/Inf behavior. | | ``HELION_PERSISTENT_RESERVED_SMS`` | ``persistent_reserved_sms`` | Reserve this many streaming multiprocessors when launching persistent kernels (``0`` uses all available SMs). | | ``HELION_FORCE_AUTOTUNE`` | ``force_autotune`` | Force the autotuner to run even when explicit configs are provided. The result is saved to the cache. | | ``HELION_AUTOTUNE_FORCE_PERSISTENT`` | ``autotune_force_persistent`` | Restrict ``pid_type`` to persistent kernel strategies during config search. | diff --git a/helion/_compiler/aten_lowering.py b/helion/_compiler/aten_lowering.py index 69a0a0eda..dddd5a5f4 100644 --- a/helion/_compiler/aten_lowering.py +++ b/helion/_compiler/aten_lowering.py @@ -494,15 +494,28 @@ def apply_dot_requirements(lowering: AtenLowering, node: Node) -> Lowering: return lowering +def matmul_masked_value(node: Node) -> float | bool | None: + """Propagate zero padding through a matmul under fast-math semantics.""" + if not CompileEnvironment.current().settings.fast_math: + # In strict mode, zero multiplied by NaN or infinity is not zero. + return None + lhs, rhs = cast("tuple[Node, Node]", node.args[:2]) + return ( + 0 if cached_masked_value(lhs) == 0 and cached_masked_value(rhs) == 0 else None + ) + + bmm_lowering = register_lowering( torch.ops.aten.bmm.default, apply_dot_requirements, + masked_value_fn=matmul_masked_value, ) mm_lowering = register_lowering( torch.ops.aten.mm.default, apply_dot_requirements, + masked_value_fn=matmul_masked_value, ) diff --git a/helion/_compiler/generate_ast.py b/helion/_compiler/generate_ast.py index 9a01e992e..320fe6d28 100644 --- a/helion/_compiler/generate_ast.py +++ b/helion/_compiler/generate_ast.py @@ -59,17 +59,8 @@ class ResidentPrepLowering: hoist: ResidentPrepHoist resident_window_name: str cache_name: str - # The value the refill writes into the cache's padded (out-of-range) tail; it also - # serves as the elision key, since a downstream per-tile ``_mask_to`` with this same - # fill is then redundant and may be dropped (letting Mosaic fold the transpose into - # the matmul push). Every prep that installs a refill -- all of them today -- must - # declare a finite value: ``_emit_resident_prep_refill`` emits it as a bare literal - # and asserts it is non-None and finite. The ``None`` default is a construction guard - # only -- a new prep kind that forgets to set a fill is not silently opted into elision - # (the elision dict skips ``None``) and trips that refill assert -- NOT a usable "write - # a fill but keep the load mask" mode. Expressing that needs the field split into a - # required tail-write value plus an optional (separate) elision fill. - tail_fill_value: float | None = None + # Fill written to the padded tail and used to identify redundant masks. + tail_fill_value: float class GenerateAST(NodeVisitor, CodegenInterface): diff --git a/helion/_compiler/pallas/compact_worklist.py b/helion/_compiler/pallas/compact_worklist.py index 8b2532376..0c8312363 100644 --- a/helion/_compiler/pallas/compact_worklist.py +++ b/helion/_compiler/pallas/compact_worklist.py @@ -476,45 +476,28 @@ def elide_installed_prep_load_masks( graph: torch.fx.Graph, load_tail_fills: Mapping[str, float], ) -> None: - """Drop the redundant per-tile OOB masks on loads whose prep cache was installed. - - ``load_tail_fills`` maps the load node name of each ACTUALLY-installed prep lowering - to the value that lowering's refill writes into the cache's padded tail (its - ``tail_fill_value``). Because the refill already wrote that value there, a - downstream ``_mask_to`` with the same fill is redundant; deleting it also lets - Mosaic fold the transpose into the matmul push. - - Called from prep-lowering installation with only the lowerings that survived - validation, so a prep that fell back to resident-only leaves its load's mask in - place (correctness is never coupled to admission-time optimism). Elision is keyed - on the declared ``tail_fill_value``: a flash-style ``_mask_to(scores, -inf)`` (fill - != the cache's tail fill, and downstream of the dot) is preserved automatically. - Non-resident (streamed) deferred loads keep their unknown masked value untouched. + """Drop redundant masks after resident prep lowerings are installed. + + ``load_tail_fills`` maps each installed load to the value written into its prep + cache's padded tail. Deferred loads without an installed prep remain unknown and + keep their masks. """ from ..host_function import HostFunction from ..node_masking import remove_unnecessary_masking if not load_tail_fills: return - # ``remove_unnecessary_masking`` recomputes masked values, which for loop-carried - # (phi) nodes walks the enclosing graphs via ``DeviceIR.current()``; make the device - # IR current since prep-lowering install runs during codegen, outside the pass. + # Loop placeholders resolve masked values through the enclosing device IR. with HostFunction.current().device_ir: for node in graph.nodes: - fill = load_tail_fills.get(node.name) - if fill is not None: - # The refill wrote ``fill`` into the padded tail: declare it so the - # matching-fill ``_mask_to`` downstream is judged redundant. - node.meta["masked_value"] = fill - elif ( - node.meta.get("masked_value") is None - and "pallas_deferred_mask_block_ids" in node.meta - ): - # A non-resident deferred load: keep its unknown masked value so its - # (still-needed) deferred mask is preserved. + if node.op == "placeholder": + # Prep installation changes inner loads and their descendants only. continue + if node.name in load_tail_fills: + node.meta["masked_value"] = load_tail_fills[node.name] + elif "pallas_deferred_mask_block_ids" in node.meta: + node.meta["masked_value"] = None else: - # Drop the stale cache so masked values recompute from the loads above. node.meta.pop("masked_value", None) remove_unnecessary_masking(graph) diff --git a/helion/_compiler/pallas/tracing_ops.py b/helion/_compiler/pallas/tracing_ops.py index e06407c08..8cdefd2c2 100644 --- a/helion/_compiler/pallas/tracing_ops.py +++ b/helion/_compiler/pallas/tracing_ops.py @@ -304,16 +304,11 @@ def _prepare_resident_prep_lowerings( tail_fill_value=0.0, ) ) - # These lowerings are now installed (all validation above passed; any fallback - # returned early), so it is safe to drop the redundant per-tile masks on exactly - # these loads -- keyed on each lowering's declared tail fill, which also preserves - # a flash-style _mask_to(scores, -inf) whose fill differs. Coupling elision to the - # installed set (not admission) keeps correctness off the "prep always emits" path. - load_tail_fills: dict[str, float] = {} - for lw in lowerings: - fill = lw.tail_fill_value - if fill is not None: - load_tail_fills[lw.hoist.load_node_name] = fill + # Fallback paths above return before declaring any tail-fill guarantees. + load_tail_fills = { + lowering.hoist.load_node_name: lowering.tail_fill_value + for lowering in lowerings + } elide_installed_prep_load_masks(graph_info.graph, load_tail_fills) if common_statements is not None: state.codegen.grouped_resident_prep_lowering_cache[cache_key] = lowerings @@ -357,12 +352,9 @@ def _stmt(src: str) -> ast.stmt: refill_tail_stmts: list[ast.stmt] = [] for lowering in lowerings: assert isinstance(lowering, ResidentPrepLowering) - # The tail fill is emitted as a bare literal below, so it must be a finite - # number. A non-finite (inf/-inf/nan) or undeclared (None) fill would need - # deliberate literal formatting; today only the transpose prep (0.0) reaches - # here, so assert rather than silently mis-emit. + # Generated fill literals currently support finite values only. tail_fill = lowering.tail_fill_value - assert tail_fill is not None and -float("inf") < tail_fill < float("inf"), ( + assert -float("inf") < tail_fill < float("inf"), ( "resident prep refill supports only finite numeric tail_fill_value" ) perm = lowering.hoist.perm diff --git a/helion/runtime/settings.py b/helion/runtime/settings.py index e2cb6c140..0cff37218 100644 --- a/helion/runtime/settings.py +++ b/helion/runtime/settings.py @@ -623,7 +623,8 @@ class Settings(_Settings): "dot_precision": "Precision for dot products. For Triton backend, see `triton.language.dot` (can be 'tf32', 'tf32x3', 'ieee'). For JAX/Pallas backend, accepted values emit Pallas default precision on TPU. Unified mappings exist so that any value can be used on any backend.", "fast_math": ( "If True, enable fast math approximations (Helion-level and Inductor-level). " - "May reduce numerical precision. Set HELION_FAST_MATH=1 to enable." + "May reduce numerical precision and change NaN/Inf behavior. " + "Set HELION_FAST_MATH=1 to enable." ), "static_shapes": ( "If True, use static shapes for all tensors. This is a performance optimization. " diff --git a/test/test_masking.py b/test/test_masking.py index 834248d5a..2b9e89ced 100644 --- a/test/test_masking.py +++ b/test/test_masking.py @@ -1,5 +1,6 @@ from __future__ import annotations +import types import unittest from unittest.mock import patch @@ -7,6 +8,8 @@ import helion from helion import _compat +from helion._compiler.aten_lowering import matmul_masked_value +from helion._compiler.compile_environment import CompileEnvironment from helion._testing import DEVICE from helion._testing import RefEagerTestBase from helion._testing import TestCase @@ -18,6 +21,26 @@ from helion.runtime.settings import _get_backend +class TestMatmulMaskedValue(unittest.TestCase): + def test_requires_fast_math_and_two_zero_padded_operands(self) -> None: + graph = torch.fx.Graph() + lhs = graph.placeholder("lhs") + rhs = graph.placeholder("rhs") + lhs.meta["masked_value"] = 0 + rhs.meta["masked_value"] = 0 + mm = graph.call_function(torch.ops.aten.mm.default, (lhs, rhs)) + + env = types.SimpleNamespace(settings=types.SimpleNamespace(fast_math=False)) + with patch.object(CompileEnvironment, "current", return_value=env): + self.assertIsNone(matmul_masked_value(mm)) + + env.settings.fast_math = True + self.assertEqual(matmul_masked_value(mm), 0) + + rhs.meta["masked_value"] = None + self.assertIsNone(matmul_masked_value(mm)) + + @onlyBackends(["triton", "cute"]) class TestMasking(RefEagerTestBase, TestCase): def test_mask_dot(self): diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 5d43f7ff6..6026ed8b9 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -1545,6 +1545,35 @@ def test_fully_jagged_with_empty_kv_matches_eager(self): ) torch.testing.assert_close(out.cpu(), ref, rtol=2e-2, atol=2e-2) + @skipIfPallasInterpret( + "the resident-cache ordered KV path is validated on real TPU, not " + "Pallas interpret mode" + ) + def test_grouped_fast_math_mask_elision_matches_eager(self): + H, D, block = 2, 16, 16 + qo = _offsets([10, 23, 7, 40]) + kvo = _offsets([16, 5, 0, 33]) + lq, lkv = int(qo[-1]), int(kvo[-1]) + torch.manual_seed(0) + q = torch.randn(lq, H, D, device=DEVICE, dtype=torch.bfloat16) + k = torch.randn(lkv, H, D, device=DEVICE, dtype=torch.bfloat16) + v = torch.randn(lkv, H, D, device=DEVICE, dtype=torch.bfloat16) + ref = _eager_fully_jagged(q.cpu(), k.cpu(), v.cpu(), qo, kvo) + kernel = helion.kernel( + _fully_jagged_kernel.fn, + backend="pallas", + static_shapes=True, + fast_math=True, + ) + + code, out = code_and_output( + kernel, + (q, k, v, qo.to(DEVICE), kvo.to(DEVICE)), + **_worklist_config([block, block], loop_type="unroll", grouping=2), + ) + self.assertEqual(code.count("lax.dot_general(scores"), 2) + torch.testing.assert_close(out.cpu(), ref, rtol=2e-2, atol=2e-2) + @skipIfPallasInterpret( "the resident-cache ordered KV path is validated on real TPU, not " "Pallas interpret mode" @@ -1830,6 +1859,17 @@ def _resident_args(self): kvo, ) + def _fast_math_resident_code(self, grouping: int) -> str: + kernel = helion.kernel( + _fully_jagged_kernel.fn, + backend="pallas", + static_shapes=True, + fast_math=True, + ) + return kernel.bind(self._resident_args()).to_triton_code( + _worklist_config([8, 8], grouping=grouping) + ) + def test_resident_prep_zero_fill_load_mask_elided_from_reduction(self): # The prep-hoisted resident K load reads a zero-filled cache (the refill writes # tail_fill_value=0 once), so its per-tile fill-0 mask is redundant and dropped. @@ -1852,12 +1892,42 @@ def test_resident_prep_zero_fill_load_mask_elided_from_reduction(self): self.assertRegex(body, rf"\b{src} = \w+_prep\[") self.assertNotRegex(body, rf"\b{src} = jnp\.where") + def test_fast_math_elides_zero_score_masks_for_both_groupings(self): + for grouping in (1, 2): + with self.subTest(grouping=grouping): + code = self._fast_math_resident_code(grouping) + score_masks = [ + line + for line in code.splitlines() + if "jnp.where" in line and ", scores" in line + ] + self.assertEqual(score_masks, []) + self.assertEqual(code.count("lax.dot_general(scores"), grouping) + + def test_strict_math_keeps_zero_score_mask(self): + code = _fully_jagged_kernel.bind(self._resident_args()).to_triton_code( + _worklist_config([8, 8]) + ) + score_masks = [ + line + for line in code.splitlines() + if "jnp.where" in line and ", scores" in line + ] + self.assertEqual(len(score_masks), 1) + self.assertNotIn("lax.dot_general(scores", code) + def test_flash_resident_prep_keeps_softmax_neg_inf_mask(self): # Flash's fill-0 K/V load masks elide (prep cache zeroed), but the amax # reduction's softmax mask fills -inf (!= the cache's 0 tail) and is downstream # of the dot, so it is preserved. Assert the score mask specifically: a # jnp.where whose fill is a -inf full (not merely the m_i init's -inf full). - code = _flash_prep_kernel.bind(self._resident_args()).to_triton_code( + kernel = helion.kernel( + _flash_prep_kernel.fn, + backend="pallas", + static_shapes=True, + fast_math=True, + ) + code = kernel.bind(self._resident_args()).to_triton_code( _worklist_config([8, 8]) ) self.assertIn("_rc_prep_refill", code) # prep cache installed From 3daa0856427fe713b149631c1a9f2b773331c96c Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Tue, 28 Jul 2026 19:45:48 -0700 Subject: [PATCH 2/7] Support dependent Pallas tile ends --- helion/_compiler/pallas/backend.py | 9 - helion/_compiler/pallas/compact_worklist.py | 84 ++++- helion/_compiler/pallas/tracing_ops.py | 363 +++++++++++++++++--- helion/language/builtin_ops.py | 43 ++- test/test_pallas.py | 61 ++++ test/test_pallas_worklist.py | 196 ++++++++++- 6 files changed, 677 insertions(+), 79 deletions(-) diff --git a/helion/_compiler/pallas/backend.py b/helion/_compiler/pallas/backend.py index 53ce08625..38f1b5a44 100644 --- a/helion/_compiler/pallas/backend.py +++ b/helion/_compiler/pallas/backend.py @@ -1292,15 +1292,6 @@ def pre_codegen( ) env = CompileEnvironment.current() - if ( - grouping == 0 - and config.get("pallas_loop_type", "unroll") == "unroll" - and env.config_spec.has_symbolic_or_data_dependent_bounds - ): - raise exc.InvalidConfig( - "pallas_loop_type='unroll' requires static inner-loop bounds or " - "pallas_worklist_grouping in (1, 2)." - ) plan_tiling(graphs, config, tile_strategy) diff --git a/helion/_compiler/pallas/compact_worklist.py b/helion/_compiler/pallas/compact_worklist.py index 0c8312363..0e4d8e80e 100644 --- a/helion/_compiler/pallas/compact_worklist.py +++ b/helion/_compiler/pallas/compact_worklist.py @@ -116,6 +116,9 @@ class CompactWorklistPlan: num_owners_expr: str = "" # Number of consecutive compact base tiles combined into one work item. grouping: int = 1 + # The ordered loop computes only through the current compact tile's end, + # while its metadata and resident window retain the full source range. + ordered_end_clamped_to_compact: bool = False @property def owner_axis(self) -> Axis: @@ -861,6 +864,73 @@ def _length_ast(begin: ast.AST, end: ast.AST) -> ast.AST: return _plain(f"({ast.unparse(end)}) - ({ast.unparse(begin)})") +def _same_ast(a: ast.AST, b: ast.AST) -> bool: + return ast.dump(a, include_attributes=False) == ast.dump( + b, include_attributes=False + ) + + +def _is_named_call(expr: ast.AST, name: str, qualifier: str) -> bool: + if not isinstance(expr, ast.Call): + return False + if isinstance(expr.func, ast.Name): + return expr.func.id == name + return ( + isinstance(expr.func, ast.Attribute) + and expr.func.attr == name + and isinstance(expr.func.value, ast.Name) + and expr.func.value.id == qualifier + ) + + +def _is_tile_end(expr: ast.AST, tile_var: str) -> bool: + if ( + isinstance(expr, ast.Attribute) + and expr.attr == "end" + and isinstance(expr.value, ast.Name) + and expr.value.id == tile_var + ): + return True + return ( + _is_named_call(expr, "tile_end", "hl") + and isinstance(expr, ast.Call) + and len(expr.args) == 1 + and not expr.keywords + and isinstance(expr.args[0], ast.Name) + and expr.args[0].id == tile_var + ) + + +def _ordered_source_end( + begin: ast.AST, + end: ast.AST, + compact_var: str, + compact_begin: ast.AST, + compact_end: ast.AST, +) -> tuple[ast.AST, bool]: + """Separate a compact-clamped ordered end from its full source end.""" + if _is_tile_end(end, compact_var): + if not _same_ast(begin, compact_begin): + raise exc.InvalidConfig( + "compact_worklist: an ordered bound of tile.end requires the " + "ordered and compact tile begins to match." + ) + return compact_end, True + + if ( + _is_named_call(end, "min", "builtins") + and isinstance(end, ast.Call) + and len(end.args) == 2 + and not end.keywords + ): + lhs_is_end = _is_tile_end(end.args[0], compact_var) + rhs_is_end = _is_tile_end(end.args[1], compact_var) + if lhs_is_end != rhs_is_end: + return (end.args[1] if lhs_is_end else end.args[0]), True + + return end, False + + def _packed_consecutive(begin: ast.AST, end: ast.AST, owner_var: str) -> bool: """True if ``(begin, end)`` are ``T[e]`` / ``T[e+1]`` on the SAME tensor ``T``. @@ -874,11 +944,6 @@ def _packed_consecutive(begin: ast.AST, end: ast.AST, owner_var: str) -> bool: over-allocates metadata, which is safe; a too-small UPPER would not be). """ - def _same_ast(a: ast.AST, b: ast.AST) -> bool: - return ast.dump(a, include_attributes=False) == ast.dump( - b, include_attributes=False - ) - def _owner_plus_const(node: ast.AST) -> int | None: """Return k for exactly ``owner_var + k`` forms, else None.""" if isinstance(node, ast.Name) and node.id == owner_var: @@ -1148,6 +1213,7 @@ def detect_compact_worklist_plan( # An inner tile loop is only supported when it is the ordered (carried-state) # axis: a name assigned in its body, read in its body, and read after it. + ordered_end_clamped_to_compact = False if ordered_loop is not None: if not isinstance(ordered_loop.target, ast.Name): raise exc.InvalidConfig("compact_worklist: ordered tile var is not a Name.") @@ -1243,6 +1309,13 @@ def detect_compact_worklist_plan( } o_begin = _inline(_to_plain(ordered_call.args[0]), ordered_prologue) o_end = _inline(_to_plain(ordered_call.args[1]), ordered_prologue) + o_end, ordered_end_clamped_to_compact = _ordered_source_end( + o_begin, + o_end, + compact_var, + c_begin, + c_end, + ) ordered_block_id = _block_id_of_loop(ordered_loop) if ordered_block_id is None: raise exc.InvalidConfig( @@ -1298,6 +1371,7 @@ def detect_compact_worklist_plan( tensor_policies=policies, upper_bound_expr="", # finalized at codegen via program_id.num_pids_expr num_owners_expr=num_owners_expr, + ordered_end_clamped_to_compact=ordered_end_clamped_to_compact, ) diff --git a/helion/_compiler/pallas/tracing_ops.py b/helion/_compiler/pallas/tracing_ops.py index 8cdefd2c2..1376705cd 100644 --- a/helion/_compiler/pallas/tracing_ops.py +++ b/helion/_compiler/pallas/tracing_ops.py @@ -20,6 +20,7 @@ from torch._inductor.codegen.simd import constant_repr from ...exc import BackendUnsupported +from ...exc import InvalidConfig from ...language import _decorators from ...language._tracing_ops import _and from ...language._tracing_ops import _for_loop @@ -93,6 +94,86 @@ def _loop_carried_indices(state: CodegenState, n_args: int) -> set[int]: return carried +def _proxy_loop_parts(value: object) -> list[object]: + return list(value) if isinstance(value, (list, tuple)) else [value] + + +def _dependent_tile_end_expr(state: CodegenState, loop_dim_index: int) -> str | None: + """Render a supported enclosing-``Tile.end`` bound from its provenance.""" + from ..variable_origin import TileEndOrigin + + graph_info = state.get_graph(state.proxy_arg(0)) + block_ids = getattr(graph_info, "block_ids", ()) + if loop_dim_index >= len(block_ids): + return None + + ends = _proxy_loop_parts(state.proxy_arg(2)) + if loop_dim_index >= len(ends): + return None + end = ends[loop_dim_index] + if not isinstance(end, torch.SymInt): + return None + expr = _symint_sympy_expr(end) + + tile_ends: list[tuple[sympy.Symbol, TileEndOrigin]] = [] + for symbol in expr.free_symbols: + if not isinstance(symbol, sympy.Symbol): + return None + origin_info = HostFunction.current().expr_to_origin.get(symbol) + if origin_info is not None and isinstance(origin_info.origin, TileEndOrigin): + tile_ends.append((symbol, origin_info.origin)) + if len(tile_ends) != 1: + return None + + tile_end_symbol, tile_end_origin = tile_ends[0] + if ( + tile_end_origin.block_id in block_ids + or not state.codegen.active_device_loops.get(tile_end_origin.block_id) + ): + return None + tile_end_expr = tile_end_origin.host_str() + if expr == tile_end_symbol: + return tile_end_expr + if ( + expr.func is not sympy.Min + or tile_end_symbol not in expr.args + or len(expr.args) != 2 + ): + return None + source_end = next(arg for arg in expr.args if arg != tile_end_symbol) + if not isinstance(source_end, sympy.Expr): + return None + for symbol in source_end.free_symbols: + origin_info = HostFunction.current().expr_to_origin.get(symbol) + if origin_info is None or not origin_info.origin.is_host(): + return None + return CompileEnvironment.current().backend.minimum_expr( + state.sympy_expr(source_end), tile_end_expr + ) + + +def _has_supported_dependent_tile_end(state: CodegenState) -> bool: + """Whether this loop has one supported enclosing-``Tile.end`` bound.""" + graph_info = state.get_graph(state.proxy_arg(0)) + block_ids = getattr(graph_info, "block_ids", ()) + return len(block_ids) == 1 and _dependent_tile_end_expr(state, 0) is not None + + +def _has_dynamic_unroll_bound(state: CodegenState) -> bool: + bounds = [ + *_proxy_loop_parts(state.proxy_arg(1)), + *_proxy_loop_parts(state.proxy_arg(2)), + ] + return any(isinstance(bound, (torch.SymInt, torch.Tensor)) for bound in bounds) + + +def _raise_unsupported_dynamic_unroll() -> None: + raise InvalidConfig( + "pallas_loop_type='unroll' requires static inner-loop bounds, an " + "enclosing Tile.end bound, or pallas_worklist_grouping in (1, 2)." + ) + + def _extract_subscript_vals(subscript: object) -> list[object]: """Extract meta values from a subscript argument in an FX graph. @@ -137,6 +218,10 @@ def _(state: CodegenState) -> object: return _codegen_emit_pipeline(state) if pallas_loop_type == "fori_loop": return _codegen_fori_loop(state) + if _has_supported_dependent_tile_end(state): + return _codegen_dynamic_unroll(state) + if _has_dynamic_unroll_bound(state): + _raise_unsupported_dynamic_unroll() # unroll: fall through to common codegen path # pyrefly: ignore[bad-return] return state.get_graph(state.proxy_arg(0)).codegen(state) @@ -170,6 +255,11 @@ def _(state: CodegenState) -> None: if pallas_loop_type == "fori_loop": _codegen_fori_loop(state) return None + if _has_supported_dependent_tile_end(state): + _codegen_dynamic_unroll(state) + return None + if _has_dynamic_unroll_bound(state): + _raise_unsupported_dynamic_unroll() # pyrefly: ignore[bad-return] return state.get_graph(state.proxy_arg(0)).codegen(state) @@ -179,7 +269,7 @@ def _codegen_resident_cache(state: CodegenState) -> object: The ordered operand is held in a per-range resident ``pl.Element(C)`` window keyed on ``range_start`` (``C`` is the compile-threaded physical window). - Optional prep-cache descriptors are handled inside ``_codegen_fori_loop``. + Optional prep-cache descriptors are installed by the dynamic resident loop. Ranges longer than ``C`` are NOT handled in-kernel: the torch launcher raises (``runtime._compact_raise_if_range_exceeds_window``), while JAX export keeps @@ -187,7 +277,7 @@ def _codegen_resident_cache(state: CodegenState) -> object: """ decision = CompileEnvironment.current().compact_worklist_resident_cache_decision assert decision is not None and decision.active - return _codegen_fori_loop(state) + return _codegen_dynamic_unroll(state) def _resident_prep_fallback(reason: str) -> None: @@ -318,7 +408,6 @@ def _prepare_resident_prep_lowerings( def _emit_resident_prep_refill( state: CodegenState, block_ids: list[int], - grid_parts: list[str], lowerings: list[ResidentPrepLowering], ) -> None: """Emit once-per-prep-key cache refill for active descriptors.""" @@ -343,7 +432,7 @@ def _emit_resident_prep_refill( f"({ref}[_wid] != {ref}[jnp.maximum(_wid - 1, 0)])" for ref in prep_key_refs ) range_len_ref = metadata_ref_for_field(plan, "range_len") - num_ordered_tiles = grid_parts[0] + num_ordered_tiles = f"(({range_len_ref}[_wid] + {blk} - 1) // {blk})" def _stmt(src: str) -> ast.stmt: return cast("ast.stmt", statement_from_string(src)) @@ -425,6 +514,35 @@ def _stmt(src: str) -> ast.stmt: state.add_statement(refill_fn) +def _emit_resident_prep_refill_once( + state: CodegenState, + block_ids: list[int], + lowerings: list[ResidentPrepLowering], +) -> None: + if not lowerings: + return + refill_key = tuple( + ( + lowering.hoist.graph_id, + lowering.hoist.prep_node_name, + lowering.cache_name, + ) + for lowering in lowerings + ) + common_statements = state.codegen.grouped_compact_common_statements + if ( + common_statements is not None + and refill_key in state.codegen.grouped_resident_prep_refill_cache + ): + return + if common_statements is None: + _emit_resident_prep_refill(state, block_ids, lowerings) + return + with state.codegen.set_statements(common_statements): + _emit_resident_prep_refill(state, block_ids, lowerings) + state.codegen.grouped_resident_prep_refill_cache[refill_key] = "emitted" + + def _classify_loop_tensors( graph_info: object, state: object, @@ -750,7 +868,15 @@ def _compact_worklist_bounds( assert plan is not None ref_names = compact_ref_names if kind == "compact" else ordered_ref_names begin_ref, extent_ref = (f"{n}_ref" for n in ref_names(plan)) - return f"{begin_ref}[_wid]", f"{begin_ref}[_wid] + {extent_ref}[_wid]" + begin = f"{begin_ref}[_wid]" + end = f"{begin} + {extent_ref}[_wid]" + if kind == "ordered" and plan.ordered_end_clamped_to_compact: + compact_begin, compact_extent = ( + f"{name}_ref" for name in compact_ref_names(plan) + ) + compact_end = f"{compact_begin}[_wid] + {compact_extent}[_wid]" + end = f"jnp.minimum({end}, {compact_end})" + return begin, end def _get_loop_begin_and_end( @@ -771,6 +897,7 @@ def _get_loop_begin_and_end( remap = _compact_worklist_bounds(state, loop_dim_index) if remap is not None: return remap + dependent_end = _dependent_tile_end_expr(state, loop_dim_index) ast_begins = state.ast_args[1] ast_ends = state.ast_args[2] begins = list(ast_begins) if isinstance(ast_begins, (list, tuple)) else [ast_begins] @@ -781,7 +908,9 @@ def _to_str(value: object) -> str: return ast.unparse(value) return str(value) - return _to_str(begins[loop_dim_index]), _to_str(ends[loop_dim_index]) + return _to_str(begins[loop_dim_index]), ( + dependent_end if dependent_end is not None else _to_str(ends[loop_dim_index]) + ) def _get_loop_numel(state: CodegenState, loop_dim_index: int) -> str: @@ -927,6 +1056,18 @@ def _compute_pipeline_or_dma_extra_pad( return bs_val - 1 +def _active_loop_begin_expr(state: CodegenState, block_id: int) -> str: + loops = state.codegen.active_device_loops.get(block_id) + if not loops: + return "0" + info = loops[-1].block_id_to_info.get(block_id) + if info is None: + return "0" + if info.begin_expr is not None: + return str(info.begin_expr) + return info.begin_var_name or "0" + + def _scratch_read(state: CodegenState, sname: str) -> str: """Read expression for a scratch buffer, slicing if padded for TPU.""" sl = state.device_function.scratch_read_slice(sname) @@ -2197,6 +2338,12 @@ def _make_block_spec( bs_var = state.device_function.block_size_var(bid) if bs_var: block_shape_parts.append(bs_var) + from ...language.memory_ops import _record_pad_info + + extra_pad = _compute_pipeline_or_dma_extra_pad( + _active_loop_begin_expr(state, bid), bid, env, state + ) + _record_pad_info(state, fake, dim_idx, bid, extra_pad) else: block_shape_parts.append(str(int(shape[dim_idx]))) lambda_parts.append(pid_var) @@ -2377,12 +2524,14 @@ def _make_hbm_slice( # Build block_id_to_info for the pipeline state block_id_to_info: dict[int, LoopDimInfo] = {} - for block_id in block_ids: + for i, block_id in enumerate(block_ids): block_size = env.block_sizes[block_id] # when the block_size.size is None, we cannot form a SymPy expr for the numel sympy_end_expr = block_size.numel if block_size.size is not None else None + begin_expr, end_expr = _get_loop_begin_and_end(state, i) block_id_to_info[block_id] = LoopDimInfo( - end_var_name=None, + begin_var_name=begin_expr, + end_var_name=end_expr, end_expr=sympy_end_expr, ) @@ -2697,12 +2846,7 @@ def _classify_pipelined_tensors( outer_access_targets = ATOMIC_OPS | {_load_op, _store_op} - all_tensor_info: list[tuple[torch.Tensor, list[object], str]] = [] - for key, (fake, _tensor_node, sub_meta) in loaded_tensors.items(): - if key not in stored_tensors: - all_tensor_info.append((fake, sub_meta, "load")) - for fake, _tensor_node, sub_meta in stored_tensors.values(): - all_tensor_info.append((fake, sub_meta, "store")) + all_tensor_info = _resident_loop_tensor_info(loaded_tensors, stored_tensors) vmem_shapes = _compute_vmem_shapes( all_tensor_info, block_ids, slice_size_exprs, env, state ) @@ -2755,6 +2899,152 @@ def _classify_pipelined_tensors( return all_tensor_info, vmem_shapes, pipelined_ids +def _resident_loop_tensor_info( + loaded_tensors: dict[int, tuple[torch.Tensor, torch.fx.Node, list[object]]], + stored_tensors: dict[int, tuple[torch.Tensor, torch.fx.Node, list[object]]], +) -> list[tuple[torch.Tensor, list[object], str]]: + """Tensor access records needed by optional resident prep lowering.""" + result = [ + (fake, sub_meta, "load") + for key, (fake, _tensor_node, sub_meta) in loaded_tensors.items() + if key not in stored_tensors + ] + result.extend( + (fake, sub_meta, "store") + for fake, _tensor_node, sub_meta in stored_tensors.values() + ) + return result + + +def _codegen_dynamic_unroll(state: CodegenState) -> object: + """Run the ordinary resident unroll body with a dynamic trip count.""" + from ..device_ir import ForLoopGraphInfo + from ..generate_ast import GenerateAST + from ..inductor_lowering import codegen_call_with_graph + from ..tile_strategy import ForiLoopState + from ..tile_strategy import LoopDimInfo + + graph_info = state.get_graph(state.proxy_arg(0)) + assert isinstance(graph_info, ForLoopGraphInfo) + assert isinstance(state.codegen, GenerateAST) + block_ids = graph_info.block_ids + if len(block_ids) != 1: + raise InvalidConfig( + "dynamic pallas unroll currently supports one inner tile dimension" + ) + + args = state.ast_args[-1] + assert isinstance(args, list) + assert all(isinstance(arg, ast.AST) for arg in args) + + env = CompileEnvironment.current() + grid_parts, block_size_vars = _compute_grid_and_block_sizes(state, block_ids, env) + begin_exprs, iter_step_exprs, _ = _pallas_loop_begin_and_step_exprs( + state, block_ids, block_size_vars + ) + strategy = _find_strategy(state, block_ids) + loop_var = state.device_function.new_var("_j") + body_stmts: list[ast.AST] = [] + _emit_inner_loop_offset_indices( + state, + strategy, + block_ids, + block_size_vars, + begin_exprs, + iter_step_exprs, + [loop_var], + env, + body_stmts, + ) + _setup_inner_loop_masks( + state, + strategy, + block_ids, + block_size_vars, + env, + body_stmts, + offset_expr_fn=lambda _i, bs: f"{loop_var} * {bs} + jnp.arange({bs})", + ) + + begin_expr, end_expr = _get_loop_begin_and_end(state, 0) + block_id_to_info = { + block_ids[0]: LoopDimInfo( + begin_var_name=begin_expr, + end_var_name=end_expr, + ) + } + body_fn_name = state.device_function.new_var("_dynamic_unroll_body") + fori_state = ForiLoopState( + strategy=strategy, # pyrefly: ignore[bad-argument-type] + block_id_to_info=block_id_to_info, + body_fn_name=body_fn_name, + loop_var_name=loop_var, + inner_statements=body_stmts, + ) + + loaded_tensors, stored_tensors = _classify_loop_tensors(graph_info, state) + resident_prep_lowerings = _prepare_resident_prep_lowerings( + state, + block_ids, + _resident_loop_tensor_info(loaded_tensors, stored_tensors), + ) + _emit_resident_prep_refill_once(state, block_ids, resident_prep_lowerings) + + carried = sorted(_loop_carried_indices(state, len(args))) + body_args = [*args] + for carry_index, arg_index in enumerate(carried): + body_args[arg_index] = expr_from_string(f"_carry[{carry_index}]") + + with state.codegen.add_fori_loop(fori_state): + with state.codegen.resident_prep_lowering_scope(resident_prep_lowerings): + graph_results = codegen_call_with_graph( + state.codegen, + graph_info.graph, + body_args, + ) + assert len(graph_results) == len(carried) + assert all(isinstance(result, ast.AST) for result in graph_results) + if graph_results: + ast_results = cast("list[ast.AST]", graph_results) + return_values = ", ".join(ast.unparse(result) for result in ast_results) + if len(graph_results) == 1: + return_values += "," + state.codegen.add_statement( + statement_from_string(f"return ({return_values})") + ) + else: + state.codegen.add_statement(statement_from_string("return _carry")) + + _emit_nonlocal_scratch_declarations(state, body_stmts) + body_fn = statement_from_string( + f"def {body_fn_name}({loop_var}, _carry):\n pass" + ) + assert isinstance(body_fn, ast.FunctionDef) + body_fn.body = cast("list[ast.stmt]", body_stmts) + state.add_statement(body_fn) + + initial_values = ", ".join(ast.unparse(args[index]) for index in carried) + if len(carried) == 1: + initial_values += "," + initial_carry = f"({initial_values})" + if not carried: + state.add_statement( + statement_from_string( + f"jax.lax.fori_loop(0, {grid_parts[0]}, {body_fn_name}, ())" + ) + ) + return None + + result_var = state.device_function.new_var("_dynamic_unroll_result") + state.add_statement( + statement_from_string( + f"{result_var} = jax.lax.fori_loop(0, {grid_parts[0]}, " + f"{body_fn_name}, {initial_carry})" + ) + ) + return [expr_from_string(f"{result_var}[{i}]") for i in range(len(carried))] + + def _codegen_fori_loop(state: CodegenState) -> object: """Emit inner device loops using jax.lax.fori_loop. @@ -2970,12 +3260,14 @@ def _codegen_fori_loop(state: CodegenState) -> object: # Build block_id_to_info block_id_to_info: dict[int, LoopDimInfo] = {} - for block_id in block_ids: + for i, block_id in enumerate(block_ids): block_size = env.block_sizes[block_id] # when the block_size.size is None, we cannot form a SymPy expr for the numel sympy_end_expr = block_size.numel if block_size.size is not None else None + begin_expr, end_expr = _get_loop_begin_and_end(state, i) block_id_to_info[block_id] = LoopDimInfo( - end_var_name=None, + begin_var_name=begin_expr, + end_var_name=end_expr, end_expr=sympy_end_expr, ) @@ -3018,38 +3310,7 @@ def _codegen_fori_loop(state: CodegenState) -> object: ) if resident_prep_lowerings: assert len(grid_parts) == 1 - refill_key = tuple( - ( - lowering.hoist.graph_id, - lowering.hoist.prep_node_name, - lowering.cache_name, - ) - for lowering in resident_prep_lowerings - ) - common_statements = state.codegen.grouped_compact_common_statements - num_ordered_tiles = ( - state.codegen.grouped_resident_prep_refill_cache.get(refill_key) - if common_statements is not None - else None - ) - if num_ordered_tiles is None: - num_ordered_tiles = state.device_function.new_var("_rc_num_ordered_tiles") - target = common_statements - with state.codegen.set_statements(target): - state.add_statement( - statement_from_string(f"{num_ordered_tiles} = {grid_parts[0]}") - ) - _emit_resident_prep_refill( - state, - block_ids, - [num_ordered_tiles], - resident_prep_lowerings, - ) - if common_statements is not None: - state.codegen.grouped_resident_prep_refill_cache[refill_key] = ( - num_ordered_tiles - ) - grid_parts = [num_ordered_tiles] + _emit_resident_prep_refill_once(state, block_ids, resident_prep_lowerings) def _build_dma_slices( fake: torch.Tensor, @@ -3126,6 +3387,12 @@ def _build_dma_slices( if bs_var: hbm_parts.append(f"pl.ds({offset}, {bs_var})") hbm_needs_slice = True + from ...language.memory_ops import _record_pad_info + + extra_pad = _compute_pipeline_or_dma_extra_pad( + _active_loop_begin_expr(state, bid), bid, env, state + ) + _record_pad_info(state, fake, dim_idx, bid, extra_pad) else: hbm_parts.append(":") else: diff --git a/helion/language/builtin_ops.py b/helion/language/builtin_ops.py index 89d547549..69ddec82f 100644 --- a/helion/language/builtin_ops.py +++ b/helion/language/builtin_ops.py @@ -1,17 +1,15 @@ from __future__ import annotations import builtins -from typing import TYPE_CHECKING +from typing import cast import sympy +import torch from .._compiler.compile_environment import CompileEnvironment from .._compiler.compile_environment import _to_sympy from . import _decorators -if TYPE_CHECKING: - import torch - def compute_symbolic_min_max( args: tuple[int | torch.SymInt, ...], op: object @@ -33,20 +31,43 @@ def compute_symbolic_min_max( return shape_env.create_symintnode(expr, hint=hint) # type: ignore[return-value] +def _compute_scalar_tensor_min( + args: tuple[int | torch.SymInt | torch.Tensor, ...], +) -> torch.Tensor: + reference = next(arg for arg in args if isinstance(arg, torch.Tensor)) + assert isinstance(reference, torch.Tensor) + if any(isinstance(arg, torch.Tensor) and arg.ndim != 0 for arg in args): + raise TypeError("device min/max only supports scalar tensor arguments") + tensor_args = [ + arg + if isinstance(arg, torch.Tensor) + else torch.full_like(reference, cast("int", arg)) + for arg in args + ] + result = tensor_args[0] + for arg in tensor_args[1:]: + result = torch.minimum(result, arg) + return result + + @_decorators.device_func_replacement(builtins.min) -def _builtin_min(*args: int | torch.SymInt) -> torch.SymInt | int: - """Device replacement for builtin min() that supports symbolic integers. +def _builtin_min( + *args: int | torch.SymInt | torch.Tensor, +) -> torch.SymInt | torch.Tensor | int: + """Device replacement for min() over symbolic ints or scalar tensors. - Returns the minimum value among the provided arguments, preserving - symbolic integer expressions when present. + A scalar tensor result is used when any input is a scalar tensor; otherwise + symbolic integer expressions are preserved. Args: - *args: Integer arguments, which may be concrete ints or symbolic SymInts + *args: Concrete ints, symbolic SymInts, or scalar tensors. Returns: - The minimum value, as a SymInt if any argument is symbolic, otherwise int + The minimum value with the corresponding scalar representation. """ - return compute_symbolic_min_max(args, op=builtins.min) + if any(isinstance(arg, torch.Tensor) for arg in args): + return _compute_scalar_tensor_min(args) + return compute_symbolic_min_max(args, op=builtins.min) # type: ignore[arg-type] @_decorators.device_func_replacement(builtins.max) diff --git a/test/test_pallas.py b/test/test_pallas.py index b0984140d..f763eb626 100644 --- a/test/test_pallas.py +++ b/test/test_pallas.py @@ -515,6 +515,20 @@ def pallas_row_scale_mul(x: torch.Tensor, r: torch.Tensor) -> torch.Tensor: return out +@helion.kernel(backend="pallas", static_shapes=True) +def pallas_causal_prefix_sum(x: torch.Tensor) -> torch.Tensor: + """Sum each row through its causal diagonal using a dependent tile end.""" + n = x.size(0) + out = torch.empty([n], dtype=x.dtype, device=x.device) + for tile_q in hl.tile(n): + acc = hl.zeros([tile_q], dtype=torch.float32) + for tile_k in hl.tile(0, min(x.size(1), tile_q.end)): + causal = tile_k.index[None, :] <= tile_q.index[:, None] + acc += torch.where(causal, x[tile_q, tile_k], 0.0).sum(-1) + out[tile_q] = acc.to(out.dtype) + return out + + @helion.kernel(backend="pallas", static_shapes=True) def pallas_reduce_non_pow2(x: torch.Tensor) -> torch.Tensor: """Softmax over a non-power-of-2 reduction dim. @@ -3953,6 +3967,53 @@ def per_block_reduction(x: torch.Tensor) -> torch.Tensor: ref = x.view(8, 8, 384).sum(1) torch.testing.assert_close(result, ref, rtol=1e-3, atol=1e-3) + def test_dependent_tile_end_unroll_uses_resident_value_carry(self) -> None: + x = torch.randn(256, 256, device=DEVICE, dtype=torch.float32) + code, result = code_and_output( + pallas_causal_prefix_sum, + (x,), + block_sizes=[128, 128], + pallas_loop_type="unroll", + ) + + self.assertIn("def _dynamic_unroll_body", code) + self.assertIn("jax.lax.fori_loop", code) + self.assertNotIn("pltpu.make_async_copy", code) + self.assertNotIn("dma_semaphore", code) + self.assertNotIn("scratch_", code) + torch.testing.assert_close(result, torch.tril(x).sum(-1)) + + def test_direct_tile_end_unroll_handles_partial_outer_tile(self) -> None: + x = torch.randn(70, 128, device=DEVICE, dtype=torch.float32) + r = torch.randn(70, 1, device=DEVICE, dtype=torch.float32) + code, result = code_and_output( + pallas_row_scale_mul, + (x, r), + block_sizes=[8], + pallas_loop_type="unroll", + ) + + self.assertIn("def _dynamic_unroll_body", code) + self.assertNotIn("pltpu.make_async_copy", code) + torch.testing.assert_close(result, x * r) + + def test_dependent_tile_end_composes_with_streaming_loop_types(self) -> None: + x = torch.randn(192, 192, device=DEVICE, dtype=torch.float32) + bound = pallas_causal_prefix_sum.bind((x,)) + for loop_type, marker in ( + ("fori_loop", "jax.lax.fori_loop"), + ("emit_pipeline", "pltpu.emit_pipeline"), + ): + with self.subTest(loop_type=loop_type): + code = bound.to_triton_code( + helion.Config( + block_sizes=[128, 128], + pallas_loop_type=loop_type, + ) + ) + self.assertIn(marker, code) + self.assertIn("(0, 0, 128, 0)", code) + @xfailIfPallasInterpret( "JAX interpret cannot trace dynamic shapes (TypeError: JitTracer ~int32[])" ) diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 6026ed8b9..3bcd72b34 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -408,6 +408,31 @@ def _fully_jagged_kernel(q, k, v, q_offsets, kv_offsets): return out +@helion.kernel(backend="pallas", static_shapes=True) +def _causal_jagged_kernel(q, k, v, offsets): + H = hl.specialize(q.size(1)) + D = hl.specialize(q.size(2)) + num_sequences = offsets.size(0) - 1 + out = torch.empty_like(q) + for seq_idx in hl.grid(num_sequences): + q_start = offsets[seq_idx] + q_end = offsets[seq_idx + 1] + k_start = offsets[seq_idx] + k_end = offsets[seq_idx + 1] + for tile_q in hl.tile(q_start, q_end): + q_blk = q[tile_q, :, :].transpose(0, 1) + acc = hl.zeros([H, tile_q, D], dtype=torch.float32) + for tile_k in hl.tile(k_start, min(k_end, tile_q.end)): + k_blk = k[tile_k, :, :].transpose(0, 1) + v_blk = v[tile_k, :, :].transpose(0, 1) + scores = torch.bmm(q_blk, k_blk.transpose(-2, -1)) + causal = tile_k.index[None, None, :] <= tile_q.index[None, :, None] + scores = torch.where(causal, scores, 0.0) + acc = torch.baddbmm(acc, scores.to(v.dtype), v_blk) + out[tile_q, :, :] = acc.transpose(0, 1).to(out.dtype) + return out + + @helion.kernel(backend="pallas", static_shapes=True) def _flash_prep_kernel(q, k, v, q_offsets, kv_offsets): """Jagged flash attention: a max reduction (amax) whose padded scores must be -inf, @@ -636,6 +661,53 @@ def _dense_kv_plan(self): with bk.env: return detect_compact_worklist_plan(bk.host_function) + def test_direct_tile_end_requires_matching_begins(self): + from helion._compiler.pallas.compact_worklist import _ordered_source_end + + source_end, clamped = _ordered_source_end( + _expr("offsets[seq_idx]"), + _expr("tile_q.end"), + "tile_q", + _expr("offsets[seq_idx]"), + _expr("offsets[seq_idx + 1]"), + ) + self.assertTrue(clamped) + self.assertEqual(ast.unparse(source_end), "offsets[seq_idx + 1]") + + source_end, clamped = _ordered_source_end( + _expr("offsets[seq_idx]"), + _expr("builtins.min(offsets[seq_idx + 1], hl.tile_end(tile_q))"), + "tile_q", + _expr("offsets[seq_idx]"), + _expr("offsets[seq_idx + 1]"), + ) + self.assertTrue(clamped) + self.assertEqual(ast.unparse(source_end), "offsets[seq_idx + 1]") + + for unsupported in ( + "foo.tile_end(tile_q)", + "foo.min(offsets[seq_idx + 1], tile_q.end)", + ): + with self.subTest(unsupported=unsupported): + source_end, clamped = _ordered_source_end( + _expr("offsets[seq_idx]"), + _expr(unsupported), + "tile_q", + _expr("offsets[seq_idx]"), + _expr("offsets[seq_idx + 1]"), + ) + self.assertFalse(clamped) + self.assertEqual(ast.unparse(source_end), unsupported) + + with self.assertRaisesRegex(exc.InvalidConfig, "begins to match"): + _ordered_source_end( + _expr("kv_offsets[seq_idx]"), + _expr("tile_q.end"), + "tile_q", + _expr("q_offsets[seq_idx]"), + _expr("q_offsets[seq_idx + 1]"), + ) + def _fully_jagged_plan(self): qo = _offsets([16, 16, 16, 16]) lq = int(qo[-1]) @@ -1148,16 +1220,66 @@ def _fully_jagged_args(): kvo, ) - def test_unroll_uses_resident_fori(self): + @staticmethod + def _causal_jagged_args(): + offsets = _offsets([12, 20, 5, 30]) + length = int(offsets[-1]) + return ( + torch.randn(length, 4, 128), + torch.randn(length, 4, 128), + torch.randn(length, 4, 128), + offsets, + ) + + def test_dependent_tile_end_preserves_source_range(self): + args = self._causal_jagged_args() + bound = _causal_jagged_kernel.bind(args) + assert bound.host_function is not None + with bound.env: + plan = detect_compact_worklist_plan(bound.host_function) + self.assertTrue(plan.ordered_end_clamped_to_compact) + ordered = plan.ordered_axis + assert ordered is not None + self.assertEqual( + ast.unparse(ordered.length), + "offsets[seq_idx + 1] - offsets[seq_idx]", + ) + + def test_dependent_tile_end_composes_with_loop_types_and_grouping(self): + args = self._causal_jagged_args() + for grouping in (1, 2): + for loop_type, marker in ( + ("unroll", "def _dynamic_unroll_body"), + ("fori_loop", "jax.lax.fori_loop"), + ("emit_pipeline", "pltpu.emit_pipeline"), + ): + with self.subTest(grouping=grouping, loop_type=loop_type): + code = _causal_jagged_kernel.bind(args).to_triton_code( + _worklist_config([8, 8], grouping=grouping, loop_type=loop_type) + ) + self.assertIn(marker, code) + self.assertIn( + "jnp.minimum(k_begin_ref[_wid] + k_len_ref[_wid], " + "q_begin_ref[_wid] + q_extent_ref[_wid])", + code, + ) + if loop_type == "unroll": + self.assertNotIn("pltpu.make_async_copy", code) + self.assertNotIn("dma_semaphore", code) + self.assertNotIn("scratch_0", code) + + def test_unroll_uses_resident_value_carry(self): code = _fully_jagged_kernel.bind(self._fully_jagged_args()).to_triton_code( _worklist_config([8, 8]) ) - # A flattened unroll reduction lowers via the resident-cache fori path, - # reuses the unified launcher, and keeps q/out in aligned windows. The - # transpose-cache structure is covered by TestResidentPrepHoistCodegen. + # A flattened unroll reduction uses a dynamic value-carried resident + # loop, reuses the unified launcher, and keeps q/out in aligned windows. + # The transpose-cache structure is covered by TestResidentPrepHoistCodegen. + self.assertIn("def _dynamic_unroll_body", code) self.assertIn("lax.fori_loop", code) self.assertNotIn("pltpu.emit_pipeline(", code) + self.assertNotIn("scratch_0", code) self.assertIn("_compact_aligned_arg_indices=", code) def test_grouping_two_emits_static_compact_variants(self): @@ -1301,6 +1423,18 @@ def _eager_fully_jagged(q, k, v, qo, kvo): return out +def _eager_causal_jagged(q, k, v, offsets): + out = torch.empty_like(q) + for s in range(len(offsets) - 1): + begin, end = int(offsets[s]), int(offsets[s + 1]) + qb = q[begin:end].transpose(0, 1) + kb = k[begin:end].transpose(0, 1) + vb = v[begin:end].transpose(0, 1) + scores = torch.bmm(qb, kb.transpose(-2, -1)) + out[begin:end] = torch.bmm(torch.tril(scores), vb).transpose(0, 1) + return out + + @onlyBackends(["pallas"]) class TestUpperBoundPacking(unittest.TestCase): """Only packed-offset bounds are accepted; non-packed is rejected. @@ -1457,6 +1591,52 @@ def test_dense_kv_unaligned_matches_eager(self): ) torch.testing.assert_close(out.cpu(), ref, rtol=2e-2, atol=2e-2) + def test_causal_dependent_tile_end_matches_eager(self): + H, D, block = 2, 128, 8 + offsets = _offsets([20]) + length = int(offsets[-1]) + torch.manual_seed(0) + q = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + k = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + v = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + ref = _eager_causal_jagged(q.cpu(), k.cpu(), v.cpu(), offsets) + for grouping in (1, 2): + with self.subTest(grouping=grouping): + _, out = code_and_output( + _causal_jagged_kernel, + (q, k, v, offsets.to(DEVICE)), + **_worklist_config( + [block, block], loop_type="unroll", grouping=grouping + ), + ) + torch.testing.assert_close(out.cpu(), ref, rtol=1e-3, atol=1e-3) + + @skipIfPallasInterpret( + "dynamic worklist streaming is validated on real TPU, not Pallas interpret" + ) + def test_causal_dependent_tile_end_streaming_matches_eager(self): + H, D, block = 2, 128, 8 + offsets = _offsets([20, 13]) + length = int(offsets[-1]) + torch.manual_seed(0) + q = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + k = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + v = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + ref = _eager_causal_jagged(q.cpu(), k.cpu(), v.cpu(), offsets) + for grouping in (1, 2): + for loop_type in ("fori_loop", "emit_pipeline"): + with self.subTest(grouping=grouping, loop_type=loop_type): + _, out = code_and_output( + _causal_jagged_kernel, + (q, k, v, offsets.to(DEVICE)), + **_worklist_config( + [block, block], + loop_type=loop_type, + grouping=grouping, + ), + ) + torch.testing.assert_close(out.cpu(), ref, rtol=1e-3, atol=1e-3) + def test_dense_kv_empty_batch_zero_grid(self): # total_q == 0 => num_work == 0 => dynamic grid=(0,). End-to-end guard # that the empty-batch launch returns an empty output (Mosaic tolerates @@ -1823,8 +2003,12 @@ def test_jagged_gdpa_emits_resident_prep_cache(self): ) self.assertIn("_rc_prep_refill", code) self.assertIn("jnp.maximum(_wid - 1, 0)", code) - self.assertIn("_rc_num_ordered_tiles", code) self.assertIn("_rc_full_ordered_tiles", code) + self.assertIn( + "_rc_full_ordered_tiles < (kv_len_ref[_wid] + _BLOCK_SIZE_2 - 1) " + "// _BLOCK_SIZE_2", + code, + ) self.assertNotIn("_rc_full_nkv", code) self.assertIn("kv_len_ref[_wid]", code) refill_guard = next( @@ -1880,7 +2064,7 @@ def test_resident_prep_zero_fill_load_mask_elided_from_reduction(self): code = _fully_jagged_kernel.bind(self._resident_args()).to_triton_code( _worklist_config([8, 8]) ) - body = code[code.index("def _fori_body_0") :] + body = code[code.index("def _dynamic_unroll_body") :] dot = re.search(r"dot_general\(\w+, (permute_\d+),", body) self.assertIsNotNone(dot, "q@káµ€ dot should read a permute operand") pvar = dot.group(1) From 9c46bc78b674ebae073da28bdef71d17d0f805c6 Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Wed, 29 Jul 2026 21:08:27 -0700 Subject: [PATCH 3/7] Use exact FP32 inputs for Pallas causal tests --- test/test_pallas_worklist.py | 31 +++++++++++++++++++------------ 1 file changed, 19 insertions(+), 12 deletions(-) diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 3bcd72b34..83c604c53 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -1435,6 +1435,21 @@ def _eager_causal_jagged(q, k, v, offsets): return out +def _exact_causal_jagged_inputs(offsets, num_heads, head_dim): + length = int(offsets[-1]) + q = torch.ones(length, num_heads, head_dim, dtype=torch.float32) + k = torch.zeros_like(q) + v = torch.zeros_like(q) + tokens = torch.arange(length) + heads = torch.arange(num_heads) + features = tokens.remainder(head_dim) + k[tokens[:, None], heads[None, :], features[:, None]] = 1.0 + v[tokens[:, None], heads[None, :], features[:, None]] = ( + (tokens[:, None] + 1) * (heads[None, :] + 1) + ).to(torch.float32) + return q.to(DEVICE), k.to(DEVICE), v.to(DEVICE) + + @onlyBackends(["pallas"]) class TestUpperBoundPacking(unittest.TestCase): """Only packed-offset bounds are accepted; non-packed is rejected. @@ -1594,11 +1609,7 @@ def test_dense_kv_unaligned_matches_eager(self): def test_causal_dependent_tile_end_matches_eager(self): H, D, block = 2, 128, 8 offsets = _offsets([20]) - length = int(offsets[-1]) - torch.manual_seed(0) - q = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) - k = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) - v = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + q, k, v = _exact_causal_jagged_inputs(offsets, H, D) ref = _eager_causal_jagged(q.cpu(), k.cpu(), v.cpu(), offsets) for grouping in (1, 2): with self.subTest(grouping=grouping): @@ -1609,7 +1620,7 @@ def test_causal_dependent_tile_end_matches_eager(self): [block, block], loop_type="unroll", grouping=grouping ), ) - torch.testing.assert_close(out.cpu(), ref, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) @skipIfPallasInterpret( "dynamic worklist streaming is validated on real TPU, not Pallas interpret" @@ -1617,11 +1628,7 @@ def test_causal_dependent_tile_end_matches_eager(self): def test_causal_dependent_tile_end_streaming_matches_eager(self): H, D, block = 2, 128, 8 offsets = _offsets([20, 13]) - length = int(offsets[-1]) - torch.manual_seed(0) - q = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) - k = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) - v = torch.randn(length, H, D, device=DEVICE, dtype=torch.float32) + q, k, v = _exact_causal_jagged_inputs(offsets, H, D) ref = _eager_causal_jagged(q.cpu(), k.cpu(), v.cpu(), offsets) for grouping in (1, 2): for loop_type in ("fori_loop", "emit_pipeline"): @@ -1635,7 +1642,7 @@ def test_causal_dependent_tile_end_streaming_matches_eager(self): grouping=grouping, ), ) - torch.testing.assert_close(out.cpu(), ref, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(out.cpu(), ref, rtol=0, atol=0) def test_dense_kv_empty_batch_zero_grid(self): # total_q == 0 => num_work == 0 => dynamic grid=(0,). End-to-end guard From a5680c7b018ef6bc75ad99bbec91cca0f2adbb06 Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Thu, 30 Jul 2026 22:09:05 -0700 Subject: [PATCH 4/7] Tighten the Pallas dependent tile-end lowering Follow-up cleanups to the dependent tile-end support: remove the ways the three inner-loop lowerings could drift, and drop a config knob that the resident lowering silently ignored. - Build LoopDimInfo through one _loop_dim_infos helper. The three lowerings each constructed it separately and the resident loop had fallen out of step, omitting end_expr. Setting it changes no generated code. - Allocate the resident loop's carry tuple through new_var. Two sibling bodies in one kernel both emitted a bare _carry; a nested one would have shadowed the enclosing carry and silently rebound reads of it. - Drop pallas_pre_broadcast when the loop type cannot apply it. The transform widens loop-carried VMEM scratch, which the resident lowering does not allocate, so both settings were autotuning as distinct configs that generate one kernel. - Spell out which bound forms _ordered_source_end accepts, and name the two other places that recognize the same shape and have to move with it: tracing_ops._dependent_tile_end_expr for kernels with no worklist plan, and tile_strategy._fold_tile_end_op for every backend. - Delete the worklist search gating tests: they assert behavior keyed on ConfigSpec.pallas_worklist_search_enabled, which does not exist. Replace the recognizer test with test_dependent_bound_grammar, covering accepted and rejected forms and both begin-matching rules. Generated code is unchanged across 42 kernel/config combinations spanning every loop type, grouping, and both dense and worklist paths, except for the intended carry rename. --- helion/_compiler/pallas/compact_worklist.py | 24 +++++- helion/_compiler/pallas/tracing_ops.py | 83 +++++++++++--------- helion/autotuner/config_spec.py | 11 +++ test/test_pallas_worklist.py | 87 ++++++++++++++++----- 4 files changed, 147 insertions(+), 58 deletions(-) diff --git a/helion/_compiler/pallas/compact_worklist.py b/helion/_compiler/pallas/compact_worklist.py index 0e4d8e80e..018be7c46 100644 --- a/helion/_compiler/pallas/compact_worklist.py +++ b/helion/_compiler/pallas/compact_worklist.py @@ -908,8 +908,27 @@ def _ordered_source_end( compact_begin: ast.AST, compact_end: ast.AST, ) -> tuple[ast.AST, bool]: - """Separate a compact-clamped ordered end from its full source end.""" + """Split an ordered end into its source end and whether it clamps. + + Accepts two forms on the ordered loop's end argument:: + + tile.end / hl.tile_end(tile) -> source end is the compact tile's own + min(, tile.end) -> source end is + + Returns the full SOURCE end -- what sizes range_len, the resident window, + and its refills -- plus whether the per-work-item compute range is clamped + to the compact tile's end. Not recognized: ``max(src, tile.begin)`` (what + a sliding window needs), offsets from the edge, and enclosing axes other + than the compact tile. + + Two other places recognize the same shape and must be extended in step: + ``tracing_ops._dependent_tile_end_expr`` matches it on the traced SymInt for + kernels with no worklist plan, and ``tile_strategy._fold_tile_end_op`` + matches the bare form for every backend. + """ if _is_tile_end(end, compact_var): + # A bare ``tile.end`` names no source end, so the compact tile's own + # source range has to be the ordered range for range_len to be right. if not _same_ast(begin, compact_begin): raise exc.InvalidConfig( "compact_worklist: an ordered bound of tile.end requires the " @@ -925,6 +944,9 @@ def _ordered_source_end( ): lhs_is_end = _is_tile_end(end.args[0], compact_var) rhs_is_end = _is_tile_end(end.args[1], compact_var) + # Exactly one side is the tile edge; the other is the source end. Both + # sides render faithfully from what the kernel wrote, so unlike the bare + # form this needs no agreement between the two begins. if lhs_is_end != rhs_is_end: return (end.args[1] if lhs_is_end else end.args[0]), True diff --git a/helion/_compiler/pallas/tracing_ops.py b/helion/_compiler/pallas/tracing_ops.py index 1376705cd..a9e5a3444 100644 --- a/helion/_compiler/pallas/tracing_ops.py +++ b/helion/_compiler/pallas/tracing_ops.py @@ -49,6 +49,7 @@ from ...runtime.config import Config from ..generate_ast import ResidentPrepLowering from ..inductor_lowering import CodegenState + from ..tile_strategy import LoopDimInfo from ..tile_strategy import TileStrategy from .compact_worklist import ResidentPrepHoist @@ -99,7 +100,14 @@ def _proxy_loop_parts(value: object) -> list[object]: def _dependent_tile_end_expr(state: CodegenState, loop_dim_index: int) -> str | None: - """Render a supported enclosing-``Tile.end`` bound from its provenance.""" + """Render a supported enclosing-``Tile.end`` bound from its provenance. + + Accepts ``tile.end`` and ``min(, tile.end)`` on the traced + SymInt, for kernels with no worklist plan. Returns ``None`` for any other + bound, leaving the caller to fall back. + ``compact_worklist._ordered_source_end`` recognizes the same two forms on + the source AST; extend the two together. + """ from ..variable_origin import TileEndOrigin graph_info = state.get_graph(state.proxy_arg(0)) @@ -871,6 +879,8 @@ def _compact_worklist_bounds( begin = f"{begin_ref}[_wid]" end = f"{begin} + {extent_ref}[_wid]" if kind == "ordered" and plan.ordered_end_clamped_to_compact: + # The source range above still spans the whole reused window; this work + # item computes only through the compact tile's current end. compact_begin, compact_extent = ( f"{name}_ref" for name in compact_ref_names(plan) ) @@ -918,6 +928,32 @@ def _get_loop_numel(state: CodegenState, loop_dim_index: int) -> str: return f"(({end}) - ({begin}))" +def _loop_dim_infos( + state: CodegenState, + block_ids: list[int], + env: CompileEnvironment, +) -> dict[int, LoopDimInfo]: + """Per-dim bounds for an inner device loop, shared by every loop lowering. + + ``tile.end``/``tile.count`` on an enclosing tile read ``end_var_name`` back + out of here, so all three lowerings must publish the same bounds they + generate code against; building them in one place keeps them from drifting. + """ + from ..tile_strategy import LoopDimInfo + + infos: dict[int, LoopDimInfo] = {} + for i, block_id in enumerate(block_ids): + block_size = env.block_sizes[block_id] + begin_expr, end_expr = _get_loop_begin_and_end(state, i) + infos[block_id] = LoopDimInfo( + begin_var_name=begin_expr, + end_var_name=end_expr, + # No SymPy numel exists when the block size has no static size. + end_expr=block_size.numel if block_size.size is not None else None, + ) + return infos + + def _is_static_int(expr: str) -> bool: """True if a begin/end expression string is a compile-time integer constant. @@ -2166,7 +2202,6 @@ def _codegen_emit_pipeline(state: CodegenState) -> object: from ..generate_ast import GenerateAST from ..inductor_lowering import codegen_call_with_graph from ..tile_strategy import EmitPipelineLoopState - from ..tile_strategy import LoopDimInfo graph_info = state.get_graph(state.proxy_arg(0)) assert isinstance(graph_info, ForLoopGraphInfo) @@ -2523,17 +2558,7 @@ def _make_hbm_slice( ] # Build block_id_to_info for the pipeline state - block_id_to_info: dict[int, LoopDimInfo] = {} - for i, block_id in enumerate(block_ids): - block_size = env.block_sizes[block_id] - # when the block_size.size is None, we cannot form a SymPy expr for the numel - sympy_end_expr = block_size.numel if block_size.size is not None else None - begin_expr, end_expr = _get_loop_begin_and_end(state, i) - block_id_to_info[block_id] = LoopDimInfo( - begin_var_name=begin_expr, - end_var_name=end_expr, - end_expr=sympy_end_expr, - ) + block_id_to_info = _loop_dim_infos(state, block_ids, env) strategy = _find_strategy(state, block_ids) # Emit offset_/indices_ at the body prologue. @@ -2922,7 +2947,6 @@ def _codegen_dynamic_unroll(state: CodegenState) -> object: from ..generate_ast import GenerateAST from ..inductor_lowering import codegen_call_with_graph from ..tile_strategy import ForiLoopState - from ..tile_strategy import LoopDimInfo graph_info = state.get_graph(state.proxy_arg(0)) assert isinstance(graph_info, ForLoopGraphInfo) @@ -2966,17 +2990,10 @@ def _codegen_dynamic_unroll(state: CodegenState) -> object: offset_expr_fn=lambda _i, bs: f"{loop_var} * {bs} + jnp.arange({bs})", ) - begin_expr, end_expr = _get_loop_begin_and_end(state, 0) - block_id_to_info = { - block_ids[0]: LoopDimInfo( - begin_var_name=begin_expr, - end_var_name=end_expr, - ) - } body_fn_name = state.device_function.new_var("_dynamic_unroll_body") fori_state = ForiLoopState( strategy=strategy, # pyrefly: ignore[bad-argument-type] - block_id_to_info=block_id_to_info, + block_id_to_info=_loop_dim_infos(state, block_ids, env), body_fn_name=body_fn_name, loop_var_name=loop_var, inner_statements=body_stmts, @@ -2991,9 +3008,12 @@ def _codegen_dynamic_unroll(state: CodegenState) -> object: _emit_resident_prep_refill_once(state, block_ids, resident_prep_lowerings) carried = sorted(_loop_carried_indices(state, len(args))) + # Uniquely named: a nested dynamic-unroll body would otherwise shadow the + # enclosing loop's carry tuple and silently rebind reads of it. + carry_var = state.device_function.new_var("_carry") body_args = [*args] for carry_index, arg_index in enumerate(carried): - body_args[arg_index] = expr_from_string(f"_carry[{carry_index}]") + body_args[arg_index] = expr_from_string(f"{carry_var}[{carry_index}]") with state.codegen.add_fori_loop(fori_state): with state.codegen.resident_prep_lowering_scope(resident_prep_lowerings): @@ -3013,11 +3033,11 @@ def _codegen_dynamic_unroll(state: CodegenState) -> object: statement_from_string(f"return ({return_values})") ) else: - state.codegen.add_statement(statement_from_string("return _carry")) + state.codegen.add_statement(statement_from_string(f"return {carry_var}")) _emit_nonlocal_scratch_declarations(state, body_stmts) body_fn = statement_from_string( - f"def {body_fn_name}({loop_var}, _carry):\n pass" + f"def {body_fn_name}({loop_var}, {carry_var}):\n pass" ) assert isinstance(body_fn, ast.FunctionDef) body_fn.body = cast("list[ast.stmt]", body_stmts) @@ -3057,7 +3077,6 @@ def _codegen_fori_loop(state: CodegenState) -> object: from ..generate_ast import GenerateAST from ..inductor_lowering import codegen_call_with_graph from ..tile_strategy import ForiLoopState - from ..tile_strategy import LoopDimInfo graph_info = state.get_graph(state.proxy_arg(0)) assert isinstance(graph_info, ForLoopGraphInfo) @@ -3259,17 +3278,7 @@ def _codegen_fori_loop(state: CodegenState) -> object: dim_idx_exprs: list[str] = loop_vars # Build block_id_to_info - block_id_to_info: dict[int, LoopDimInfo] = {} - for i, block_id in enumerate(block_ids): - block_size = env.block_sizes[block_id] - # when the block_size.size is None, we cannot form a SymPy expr for the numel - sympy_end_expr = block_size.numel if block_size.size is not None else None - begin_expr, end_expr = _get_loop_begin_and_end(state, i) - block_id_to_info[block_id] = LoopDimInfo( - begin_var_name=begin_expr, - end_var_name=end_expr, - end_expr=sympy_end_expr, - ) + block_id_to_info = _loop_dim_infos(state, block_ids, env) # Emit offset_/indices_ at the body prologue. _emit_inner_loop_offset_indices( diff --git a/helion/autotuner/config_spec.py b/helion/autotuner/config_spec.py index f46dfb000..1b5474743 100644 --- a/helion/autotuner/config_spec.py +++ b/helion/autotuner/config_spec.py @@ -1844,6 +1844,17 @@ def normalize( else: config.pop("pallas_load_buffer_count", None) + if ( + self.supports_config_key("pallas_pre_broadcast") + and self.has_pallas_inner_loops + and config.get("pallas_loop_type") not in ("fori_loop", "emit_pipeline") + ): + # The transform widens loop-carried VMEM scratch, so it only applies + # to the streaming lowerings. "unroll" carries values through the + # jax.lax.fori_loop tuple and allocates no scratch to widen; pin the + # flag off there so both settings do not autotune as distinct configs. + config.pop("pallas_pre_broadcast", None) + if self.supports_config_key("pid_type"): if "pid_type" in config: if config["pid_type"] not in VALID_PID_TYPES: diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index 83c604c53..1c769fc11 100644 --- a/test/test_pallas_worklist.py +++ b/test/test_pallas_worklist.py @@ -661,32 +661,41 @@ def _dense_kv_plan(self): with bk.env: return detect_compact_worklist_plan(bk.host_function) - def test_direct_tile_end_requires_matching_begins(self): - from helion._compiler.pallas.compact_worklist import _ordered_source_end + def test_dependent_bound_grammar(self): + """The AST recognizer accepts exactly the two clamped-end forms. - source_end, clamped = _ordered_source_end( - _expr("offsets[seq_idx]"), - _expr("tile_q.end"), - "tile_q", - _expr("offsets[seq_idx]"), - _expr("offsets[seq_idx + 1]"), - ) - self.assertTrue(clamped) - self.assertEqual(ast.unparse(source_end), "offsets[seq_idx + 1]") + The traced-SymInt recognizer for kernels without a worklist plan + (``tracing_ops._dependent_tile_end_expr``) must accept the same forms; + it is covered end-to-end by the dense tests in test_pallas.py. + """ + from helion._compiler.pallas.compact_worklist import _ordered_source_end - source_end, clamped = _ordered_source_end( - _expr("offsets[seq_idx]"), - _expr("builtins.min(offsets[seq_idx + 1], hl.tile_end(tile_q))"), - "tile_q", - _expr("offsets[seq_idx]"), - _expr("offsets[seq_idx + 1]"), - ) - self.assertTrue(clamped) - self.assertEqual(ast.unparse(source_end), "offsets[seq_idx + 1]") + for accepted in ( + "tile_q.end", + "hl.tile_end(tile_q)", + "min(offsets[seq_idx + 1], tile_q.end)", + "builtins.min(offsets[seq_idx + 1], hl.tile_end(tile_q))", + "min(hl.tile_end(tile_q), offsets[seq_idx + 1])", + ): + with self.subTest(accepted=accepted): + source_end, clamped = _ordered_source_end( + _expr("offsets[seq_idx]"), + _expr(accepted), + "tile_q", + _expr("offsets[seq_idx]"), + _expr("offsets[seq_idx + 1]"), + ) + self.assertTrue(clamped) + # Either form yields the full source end, never the clamped one. + self.assertEqual(ast.unparse(source_end), "offsets[seq_idx + 1]") for unsupported in ( "foo.tile_end(tile_q)", "foo.min(offsets[seq_idx + 1], tile_q.end)", + # A different tile's edge is not the enclosing compact tile. + "min(offsets[seq_idx + 1], tile_k.end)", + # Begin edges (sliding windows) are not recognized yet. + "max(offsets[seq_idx], tile_q.begin)", ): with self.subTest(unsupported=unsupported): source_end, clamped = _ordered_source_end( @@ -699,6 +708,11 @@ def test_direct_tile_end_requires_matching_begins(self): self.assertFalse(clamped) self.assertEqual(ast.unparse(source_end), unsupported) + def test_direct_tile_end_requires_matching_begins(self): + from helion._compiler.pallas.compact_worklist import _ordered_source_end + + # A bare tile.end names no source end, so it can only stand in for the + # compact tile's own range when the two begins agree. with self.assertRaisesRegex(exc.InvalidConfig, "begins to match"): _ordered_source_end( _expr("kv_offsets[seq_idx]"), @@ -708,6 +722,18 @@ def test_direct_tile_end_requires_matching_begins(self): _expr("q_offsets[seq_idx + 1]"), ) + # The min form renders both sides from what the kernel wrote, so + # mismatched begins are fine there. + source_end, clamped = _ordered_source_end( + _expr("kv_offsets[seq_idx]"), + _expr("min(kv_offsets[seq_idx + 1], tile_q.end)"), + "tile_q", + _expr("q_offsets[seq_idx]"), + _expr("q_offsets[seq_idx + 1]"), + ) + self.assertTrue(clamped) + self.assertEqual(ast.unparse(source_end), "kv_offsets[seq_idx + 1]") + def _fully_jagged_plan(self): qo = _offsets([16, 16, 16, 16]) lq = int(qo[-1]) @@ -1282,6 +1308,27 @@ def test_unroll_uses_resident_value_carry(self): self.assertNotIn("scratch_0", code) self.assertIn("_compact_aligned_arg_indices=", code) + def test_pre_broadcast_dropped_when_loop_type_cannot_apply_it(self): + """The transform widens loop-carried VMEM scratch, which only the + streaming lowerings allocate. Pinning the flag off for "unroll" keeps + both settings from autotuning as two configs that generate one kernel. + """ + spec = _fully_jagged_kernel.bind(self._fully_jagged_args()).env.config_spec + for loop_type, retained in ( + ("unroll", False), + ("fori_loop", True), + ("emit_pipeline", True), + ): + with self.subTest(loop_type=loop_type): + config: dict[str, object] = { + "block_sizes": [8, 8], + "pallas_loop_type": loop_type, + "pallas_worklist_grouping": 1, + "pallas_pre_broadcast": True, + } + spec.normalize(config) + self.assertEqual("pallas_pre_broadcast" in config, retained) + def test_grouping_two_emits_static_compact_variants(self): qo = _offsets([12, 20, 5, 30]) lq = int(qo[-1]) From 1ea6d8d5dd50bffc6bd98e09b27287e453fcaa9a Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Thu, 30 Jul 2026 22:09:05 -0700 Subject: [PATCH 5/7] Fix min()/max() over a tile end, count, or id _lift_sympy_arg returned the loop offset for the whole GridOrigin family, but only a tile's begin is the offset. TileEndOrigin, TileCountOrigin and TileIdOrigin each render as a compound expression, so a symbol carrying one of those origins was silently replaced by the tile's begin. This is reachable whenever such a symbol survives into a sympy expression instead of its own op. min()/max() do exactly that: they are device function replacements that fold their operands into a single sympy Min/Max at trace time, dropping the tile_end node. So for tile_k in hl.tile(0, min(n, tile_q.end)): lowered its trip count from tile_q.begin -- zero iterations on the first tile, and short by one block on every later one. tile.count collapsed the same way (begin instead of cdiv), as did tile.id (begin instead of offset // block). Plain arithmetic on a tile edge (tile.end + 8) keeps the op and was never affected, nor was min() over tile.begin, which is what the offset already meant. Dispatch on the concrete origin instead. The block id is resolved before rendering and carried into a replaced copy, because host_str() looks up offset_var and active_device_loops by block id and an aliased symbol's own block may have no active loop. The rendered edge is parenthesized: it becomes a sympy Symbol name inside the enclosing expression, so an unbracketed tile.id would turn 2 * (offset // block) into (2 * offset) // block. Covered for all three edges under both ops, on an input whose final tile is partial. The test is Triton-only: Pallas lowers a dependent tile bound through its own loop codegen and never reaches this path, so its generated code is unchanged across 42 kernel/config combinations. --- helion/_compiler/device_function.py | 22 ++++++++- test/test_loops.py | 77 +++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/helion/_compiler/device_function.py b/helion/_compiler/device_function.py index e8a35d3f4..4680dc2a6 100644 --- a/helion/_compiler/device_function.py +++ b/helion/_compiler/device_function.py @@ -43,6 +43,7 @@ from .variable_origin import GridOrigin from .variable_origin import Origin from .variable_origin import TensorSizeOrigin +from .variable_origin import TileBeginOrigin if TYPE_CHECKING: from ..runtime.config import Config @@ -595,9 +596,26 @@ def _lift_sympy_arg(self, expr: sympy.Expr) -> str: assert result is not None return result if isinstance(origin.origin, GridOrigin): - return self.codegen.offset_var( - env.resolve_codegen_block_id(origin.origin.block_id, self.codegen) + # Only the tile's begin is the loop offset. The other tile edges + # render as compound expressions (``tile.end`` clamps to the loop + # end, ``tile.count`` divides, ``tile.id`` shifts), so returning the + # offset for them would silently substitute the begin. This path is + # reached whenever such a symbol survives into a sympy expression + # rather than its own op -- e.g. ``min(n, tile.end)``, which + # ``_builtin_min`` folds into ``sympy.Min`` at trace time. The + # result becomes a sympy Symbol name, so parenthesize it to keep + # precedence intact inside the enclosing expression. + resolved = env.resolve_codegen_block_id( + origin.origin.block_id, self.codegen ) + if type(origin.origin) in (GridOrigin, TileBeginOrigin): + return self.codegen.offset_var(resolved) + # Render through the origin so each derived edge keeps its own + # formula, but on the resolved live loop: host_str() reads + # offset_var and active_device_loops by block_id, and an aliased + # symbol's own block may have no active loop. + derived = dataclasses.replace(origin.origin, block_id=resolved) + return f"({derived.host_str()})" return self.expr_arg(expr, origin.origin).name def user_sympy_expr(self, expr: sympy.Expr) -> str: diff --git a/test/test_loops.py b/test/test_loops.py index a899a9831..9b08fd3ae 100644 --- a/test/test_loops.py +++ b/test/test_loops.py @@ -1482,6 +1482,83 @@ def fn(x) -> torch.Tensor: x = torch.randn(128, 1024, dtype=torch.float32, device=DEVICE) torch.testing.assert_close(fn(x), x) + @skipIfNotTriton( + "asserts on Triton's rendered bound; Pallas lowers a dependent tile " + "bound through its own loop codegen and never reaches this path" + ) + def test_min_max_over_derived_tile_edge_keeps_its_own_formula(self): + """A tile edge folded into ``min``/``max`` must keep its own formula. + + ``min``/``max`` are device function replacements, so they fold their + operands into one sympy expression at trace time. That drops the + ``tile_end``/``tile_count``/``tile_id`` op and leaves a bare symbol, + whose origin then has to be honored when it is rendered. Emitting the + loop offset for all of them substitutes ``tile.begin``: for an end + bound that is a zero trip count on the first tile. + + The 200x200 input is deliberately not a multiple of the 128 block, so + the final tile is partial and the end has to clamp. + """ + cfg = helion.Config(block_sizes=[128, 128]) + + @helion.kernel(config=cfg) + def end_min(x) -> torch.Tensor: + out = torch.empty([x.size(0)], dtype=x.dtype, device=x.device) + for tile_q in hl.tile(x.size(0)): + acc = hl.zeros([tile_q], dtype=torch.float32) + for tile_k in hl.tile(0, min(x.size(1), tile_q.end)): + acc += x[tile_q, tile_k].sum(-1) + out[tile_q] = acc.to(out.dtype) + return out + + @helion.kernel(config=cfg) + def end_max(x) -> torch.Tensor: + out = torch.empty([x.size(0)], dtype=x.dtype, device=x.device) + for tile_q in hl.tile(x.size(0)): + acc = hl.zeros([tile_q], dtype=torch.float32) + for tile_k in hl.tile(0, max(8, tile_q.end)): + acc += x[tile_q, tile_k].sum(-1) + out[tile_q] = acc.to(out.dtype) + return out + + @helion.kernel(config=cfg) + def count_min(x) -> torch.Tensor: + out = torch.empty([x.size(0)], dtype=x.dtype, device=x.device) + for tile_q in hl.tile(x.size(0)): + acc = hl.zeros([tile_q], dtype=torch.float32) + for tile_k in hl.tile(0, min(x.size(1), tile_q.count)): + acc += x[tile_q, tile_k].sum(-1) + out[tile_q] = acc.to(out.dtype) + return out + + @helion.kernel(config=cfg) + def id_min(x) -> torch.Tensor: + out = torch.empty([x.size(0)], dtype=x.dtype, device=x.device) + for tile_q in hl.tile(x.size(0)): + acc = hl.zeros([tile_q], dtype=torch.float32) + for tile_k in hl.tile(0, min(x.size(1), 2 * tile_q.id + 1)): + acc += x[tile_q, tile_k].sum(-1) + out[tile_q] = acc.to(out.dtype) + return out + + x = torch.randn(200, 200, device=DEVICE) + # Each edge renders its own formula; the offset alone would mean begin. + # The id case also pins the parenthesization: unbracketed, ``2 * + # offset_0 // BLOCK`` would floor-divide the product instead. + for label, kernel, fragment in ( + ("end/min", end_min, "offset_0 + _BLOCK_SIZE_0"), + ("end/max", end_max, "offset_0 + _BLOCK_SIZE_0"), + ("count", count_min, "tl.cdiv("), + ("id", id_min, "2 * (offset_0 // _BLOCK_SIZE_0)"), + ): + with self.subTest(edge=label): + # Codegen the declared config, not the spec default: the + # partial final tile depends on the 128 block size. + code = kernel.bind((x,)).to_triton_code(cfg) + rendered = [ln for ln in code.splitlines() if "symnode_0 = " in ln] + self.assertTrue(rendered, f"no rendered bound in:\n{code}") + self.assertIn(fragment, rendered[0]) + @skipIfNotTriton( "tl.debug_barrier() is only emitted in Triton device codegen (not Pallas/JAX)" ) From a2c6dbe24bb7614e6df14b4aa4f194b15e290702 Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Fri, 31 Jul 2026 12:47:03 -0700 Subject: [PATCH 6/7] Skip the derived-tile-edge codegen test under ref eager test_min_max_over_derived_tile_edge_keeps_its_own_formula only inspects generated code, so under HELION_INTERPRET=1 it fails twice over: ref eager exposes no tunable block sizes, making the declared block_sizes=[128, 128] an InvalidConfig, and RefEagerTestBase asserts every test either calls run_ref or skips. Mark it skipIfRefEager like the other codegen-only tests in this file. --- test/test_loops.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/test_loops.py b/test/test_loops.py index 9b08fd3ae..bcfd388c6 100644 --- a/test/test_loops.py +++ b/test/test_loops.py @@ -1482,6 +1482,7 @@ def fn(x) -> torch.Tensor: x = torch.randn(128, 1024, dtype=torch.float32, device=DEVICE) torch.testing.assert_close(fn(x), x) + @skipIfRefEager("inspects generated code; ref eager never lowers a kernel") @skipIfNotTriton( "asserts on Triton's rendered bound; Pallas lowers a dependent tile " "bound through its own loop codegen and never reaches this path" From 30553122cf16ffe13b9e9282a56d2a0bbe2c817a Mon Sep 17 00:00:00 2001 From: Ethan Che Date: Fri, 31 Jul 2026 12:47:03 -0700 Subject: [PATCH 7/7] Drop the stale interpret xfail on jsd and kl_div Both tile a symbolic V, so under Pallas interpret they used to hit the backend-wide gate that rejected pallas_loop_type="unroll" whenever a kernel had symbolic or data-dependent bounds. "Support dependent Pallas tile ends" replaced that gate with a per-loop check, so these kernels now compile and pass in interpret mode and the xfail reports an unexpected success. Bisected across the branch: main and the mask-elision commit still xfail; the dependent-tile-end commit is where they start passing. --- test/test_examples.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/test/test_examples.py b/test/test_examples.py index 177deb3c4..98e018cfa 100644 --- a/test/test_examples.py +++ b/test/test_examples.py @@ -1771,9 +1771,6 @@ def test_swiglu(self): num_stages=3, ) - @xfailIfPallasInterpret( - "JAX interpret cannot trace dynamic shapes (TypeError: JitTracer ~int32[])" - ) def test_jsd(self): args = ( torch.randn([1024, 4096], device=DEVICE, dtype=torch.float32).log_softmax( @@ -1799,9 +1796,6 @@ def test_jsd(self): num_stages=3, ) - @xfailIfPallasInterpret( - "JAX interpret cannot trace dynamic shapes (TypeError: JitTracer ~int32[])" - ) def test_kl_div(self): args = ( torch.randn([1024, 4096], device=DEVICE, dtype=torch.float32).log_softmax(