detect pair-based jagged dispatch for the Pallas backend

Adds detect_jagged_dispatch(env) -> int | None, which reads
CompileEnvironment.jagged_tile_parent_ids and returns the
items-axis block_id when every hl.jagged_tile in the kernel shares
a single hl.tile parent (i.e. the kernel has a single items axis).
Returns None otherwise so codegen can fall through.

Detector is pure (no IR mutation, no FX walk); follow-up codegen
commit reads the returned block_id to route the call through the
jagged_reduce template.

Unit tests (8) cover: no jagged tiles, single pair, multiple pairs
sharing the same parent, child with multiple parents, multiple items
axes, mixed dispatchable/non-dispatchable, and the defensive zero-
parent path.

helion/_compiler/pallas/jagged_dispatch.py

@@ -0,0 +1,64 @@
+"""Pair-based jagged dispatch detection for the Pallas backend.
+
+Helion already tracks every (``hl.jagged_tile`` child, ``hl.tile``
+parent) pair in :attr:`CompileEnvironment.jagged_tile_parent_ids`.
+This module adds a single decision on top: *is the kernel dispatchable
+to the per-item DMA-orchestrated template* at
+:func:`helion.runtime.pallas_templates.jagged_reduce_pallas`?
+
+The detector stores no new IR. Its only output is a single ``int |
+None``: the ``block_id`` of the items axis if the kernel is
+dispatchable, ``None`` otherwise. Codegen reads this and walks
+Helion's existing IR to derive everything else (jagged children of
+that axis, compute slot, flush cast, tensor args, ...).
+
+Dispatch rule
+-------------
+
+A kernel is dispatchable iff:
+
+  1. It contains at least one ``hl.jagged_tile``.
+  2. The kernel has a **single items axis** — every ``hl.jagged_tile``
+     references exactly one outer ``hl.tile``, and they all reference
+     the same one.
+
+Number of jagged children is unbounded — a kernel with N
+``hl.jagged_tile`` calls all parameterised by the same items axis
+dispatches fine (e.g. ``jagged_mean`` has two: ``tile_m`` jagged in
+features, ``tile_k`` jagged in seq length, both sharing parent
+``tile_b``).
+
+When the rule fails, the detector returns ``None`` and codegen falls
+through to the existing Pallas lowering. Kernels with multi-parent
+jagged tiles or multiple items axes are out of scope for this
+template and would need a separate extension.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..compile_environment import CompileEnvironment
+
+
+def detect_jagged_dispatch(env: CompileEnvironment) -> int | None:
+    """Return the items-axis ``block_id`` if dispatchable, else ``None``.
+
+    Reads only ``env.jagged_tile_parent_ids`` — no FX walk, no IR
+    mutation. See module docstring for the rule.
+    """
+    parents_by_child = env.jagged_tile_parent_ids
+    if not parents_by_child:
+        return None  # no hl.jagged_tile in the kernel
+
+    unique_parents: set[int] = set()
+    for parent_ids in parents_by_child.values():
+        if len(parent_ids) != 1:
+            return None  # jagged_tile has 0 or >1 parents
+        unique_parents.add(parent_ids[0])
+
+    if len(unique_parents) != 1:
+        return None  # multiple items axes (different parents)
+
+    return next(iter(unique_parents))

test/test_pallas_jagged_dispatch.py

@@ -0,0 +1,74 @@
+"""Unit tests for the pair-based jagged dispatch detector.
+
+Tests the detector in isolation by constructing a fake ``CompileEnvironment``
+with controlled ``jagged_tile_parent_ids``, so no kernel binding is needed.
+End-to-end coverage is provided by the example tests that lower
+``examples/jagged_*.py`` kernels through the Pallas backend.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from dataclasses import field
+import unittest
+
+from helion._compiler.pallas.jagged_dispatch import detect_jagged_dispatch
+
+
+@dataclass
+class _FakeEnv:
+    """Minimal stand-in — the detector reads only this one attribute."""
+
+    jagged_tile_parent_ids: dict[int, list[int]] = field(default_factory=dict)
+
+
+class TestJaggedDispatchDetect(unittest.TestCase):
+    def test_no_jagged_tile_returns_none(self) -> None:
+        # Non-jagged kernel: empty dict.
+        env = _FakeEnv()
+        self.assertIsNone(detect_jagged_dispatch(env))  # type: ignore[arg-type]
+
+    def test_single_pair_dispatches(self) -> None:
+        # jagged_sum-style: one jagged child, one parent.
+        # tile_b has block_id=0; tile_k=hl.jagged_tile(nnz) has block_id=1
+        # with parent_ids=[0].
+        env = _FakeEnv(jagged_tile_parent_ids={1: [0]})
+        self.assertEqual(detect_jagged_dispatch(env), 0)  # type: ignore[arg-type]
+
+    def test_two_pairs_same_parent_dispatches(self) -> None:
+        # jagged_mean-style: tile_m and tile_k both children of tile_b.
+        env = _FakeEnv(jagged_tile_parent_ids={1: [0], 2: [0]})
+        self.assertEqual(detect_jagged_dispatch(env), 0)  # type: ignore[arg-type]
+
+    def test_many_pairs_same_parent_dispatches(self) -> None:
+        # Unbounded number of children sharing one parent → still
+        # dispatches.
+        env = _FakeEnv(jagged_tile_parent_ids={1: [0], 2: [0], 3: [0], 4: [0]})
+        self.assertEqual(detect_jagged_dispatch(env), 0)  # type: ignore[arg-type]
+
+    def test_multi_parent_child_returns_none(self) -> None:
+        # A jagged_tile whose nnz depends on TWO outer tiles.
+        env = _FakeEnv(jagged_tile_parent_ids={2: [0, 1]})
+        self.assertIsNone(detect_jagged_dispatch(env))  # type: ignore[arg-type]
+
+    def test_zero_parent_child_returns_none(self) -> None:
+        # Defensive: should never happen (jagged_tile always has a
+        # parent), but the detector must not crash on len==0.
+        env = _FakeEnv(jagged_tile_parent_ids={1: []})
+        self.assertIsNone(detect_jagged_dispatch(env))  # type: ignore[arg-type]
+
+    def test_multiple_items_axes_returns_none(self) -> None:
+        # Two children, two distinct parents → kernel has more than one
+        # items axis. Not supported.
+        env = _FakeEnv(jagged_tile_parent_ids={2: [0], 3: [1]})
+        self.assertIsNone(detect_jagged_dispatch(env))  # type: ignore[arg-type]
+
+    def test_mixed_one_valid_one_multi_parent_returns_none(self) -> None:
+        # Defensive: even a single multi-parent child taints the
+        # dispatch — fall through.
+        env = _FakeEnv(jagged_tile_parent_ids={2: [0], 3: [0, 1]})
+        self.assertIsNone(detect_jagged_dispatch(env))  # type: ignore[arg-type]
+
+
+if __name__ == "__main__":
+    unittest.main()