[pallas] Support dependent Tile.end bounds in Pallas

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.<dtype>`` 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. |

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,
 )
 
 

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):

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)
 

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:
@@ -476,45 +479,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)
 
@@ -878,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``.
 
@@ -891,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:
@@ -1165,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.")
@@ -1260,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(
@@ -1315,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,
     )