Support OutputCodeOptions(allow_helion_deps=False) for the Pallas backend (TorchTPU)

helion/_compiler/pallas/backend.py

@@ -6,6 +6,7 @@
 import ast
 import dataclasses
 import enum
+import inspect
 import math
 from typing import TYPE_CHECKING
 from typing import Any
@@ -17,6 +18,7 @@
 from ... import exc
 from ..ast_extension import expr_from_string
 from ..backend import Backend
+from ..backend import LauncherInfo
 from ..backend import _loop_contains_matmul
 
 if TYPE_CHECKING:
@@ -35,6 +37,56 @@
     InductorOpOverrides = OpsHandler[Any]
 
 
+def _embedded_helper_source(body: str) -> str:
+    """Source of the in-kernel Pallas helpers referenced by ``body`` (module-level
+    so both ``PallasBackend.embedded_helper_source`` and the jax standalone builder
+    can inline them). Only helpers actually referenced are emitted."""
+    blocks: list[str] = []
+    if "_helion_divide_filter_topk" in body:
+        from . import topk_impl
+
+        blocks.extend(
+            [
+                _embed_source(inspect.getsource(topk_impl)),
+                "_helion_divide_filter_topk = divide_filter_topk",
+            ]
+        )
+    if "flatten_worklist" in body:
+        from ...runtime import compact_worklist
+
+        blocks.append(_embed_source(inspect.getsource(compact_worklist)))
+    return "\n\n\n".join(blocks)
+
+
+def _embed_source(source: str) -> str:
+    """Return a module's source ready to inline: its module docstring and
+    ``from __future__`` lines stripped (leading comments -- e.g. an SPDX header --
+    and everything else preserved), so the docstring prose can't leak into the
+    generated code and the mid-module ``from __future__`` (a SyntaxError) is gone.
+
+    The docstring span is located via ``ast`` (not a quote scan) so a docstring
+    whose prose contains a triple-quote can't corrupt the output, and a module
+    that opens with code rather than a docstring is handled correctly.
+    """
+    lines = source.split("\n")
+    doc_lines: set[int] = set()
+    tree = ast.parse(source)
+    if (
+        tree.body
+        and isinstance(first := tree.body[0], ast.Expr)
+        and isinstance(first.value, ast.Constant)
+        and isinstance(first.value.value, str)
+    ):
+        # ast line numbers are 1-based; end_lineno is the closing-quote line.
+        doc_lines = set(range(first.lineno - 1, (first.end_lineno or first.lineno)))
+    kept = [
+        line
+        for idx, line in enumerate(lines)
+        if idx not in doc_lines and not line.strip().startswith("from __future__")
+    ]
+    return "\n".join(kept).strip("\n")
+
+
 # Mapping from torch dtype to JAX dtype string (e.g., "jnp.float32")
 _TORCH_TO_JAX_DTYPE: dict[str, str] = {
     "torch.float16": "jnp.float16",
@@ -156,6 +208,18 @@ def constexpr_type(self) -> str:
     def default_launcher_name(self) -> str:
         return "_default_pallas_launcher"
 
+    @property
+    def dependency_free_launcher_info(self) -> LauncherInfo:
+        # Pallas generated code makes no ``helion.runtime.<fn>`` helper calls
+        # beyond the launcher, so the shim need only re-export the launcher itself.
+        return LauncherInfo(
+            launcher_module="helion.runtime.pallas.launcher",
+            launcher_symbol="default_pallas_launcher",
+            launcher_alias="_default_pallas_launcher",
+            deps="torch + jax",
+            runtime_helper_names=(),
+        )
+
     @property
     def library_imports(self) -> dict[str, str]:
         return {
@@ -169,9 +233,26 @@ def library_imports(self) -> dict[str, str]:
             "lax": "import jax.lax as lax",
             "pltpu": "from jax.experimental.pallas import tpu as pltpu",
             "_default_pallas_launcher": "from helion.runtime import default_pallas_launcher as _default_pallas_launcher",
+            # In-kernel helpers the generated code calls. Regular output imports
+            # them from helion (conditionally, only when referenced); the
+            # dependency-free path drops these imports and embeds the source instead
+            # (see ``embedded_helper_source`` / ``build_dependency_free_code``).
             "_helion_divide_filter_topk": "from helion._compiler.pallas.topk_impl import divide_filter_topk as _helion_divide_filter_topk",
+            "flatten_worklist": "from helion.runtime.compact_worklist import flatten_worklist",
         }
 
+    def embedded_helper_source(self, body: str) -> str:
+        """Inline the in-kernel Pallas helpers referenced by ``body``.
+
+        ``divide_filter_topk`` (aten.topk lowering) and ``flatten_worklist``
+        (compact-worklist builder) are pure-``jax`` helpers the generated kernel
+        calls. Regular output imports them from helion (see ``library_imports``);
+        this embeds their source instead, so a dependency-free / jax standalone is
+        self-contained. Called only by the standalone builders (never for regular
+        ``to_code``), which drop the corresponding helion imports.
+        """
+        return _embedded_helper_source(body)
+
     # Config keys that Pallas actually uses.  Everything else
     # (pid_type, num_warps, num_stages, maxnreg, indexing, etc.)
     # is GPU-specific and should not be tuned.

helion/_compiler/pallas/compact_worklist.py

@@ -703,7 +703,6 @@ def render_build_worklist(
     lines = [
         f"def {builder_name}({', '.join(offset_params)}):",
         "    import jax.numpy as jnp",
-        "    from helion.runtime.compact_worklist import flatten_worklist",
         f"    {owner_array} = jnp.arange({num_owners_expr}, dtype=jnp.int32)",
         f"    base = {base_src}",
         f"    length = {length_src}",

helion/_compiler/pallas/topk_impl.py

@@ -37,6 +37,11 @@
 if TYPE_CHECKING:
     import jax
 
+# NOTE: this module is embedded verbatim into helion-free `to_code(allow_helion_deps=False)` output (via
+# PallasBackend.embedded_helper_source), so it must not import anything from the
+# `helion` package (nor mention such an import in a comment) -- the precompiler's
+# helion-free guard is a substring check and would reject any topk kernel.
+
 _NUM_LANES = 128
 NUM_LANES = 128
 NUM_SUBLANES = 8

helion/runtime/compact_worklist.py

@@ -27,6 +27,12 @@
 if TYPE_CHECKING:
     from jax import Array
 
+# NOTE: this module is embedded verbatim into helion-free `to_code(allow_helion_deps=False)` output (via
+# PallasBackend.embedded_helper_source), so its non-docstring code must not import
+# anything from the `helion` package (nor mention such an import in a comment) --
+# the precompiler's helion-free guard is a substring check and would reject any
+# compact-worklist kernel.
+
 
 class CompactWorkMetadata(NamedTuple):
     """Scalar-prefetch metadata describing one compact worklist.

test/test_bound_kernel.py

@@ -10,6 +10,7 @@
 import sys
 import tempfile
 import textwrap
+from typing import Any
 import unittest
 
 import torch
@@ -19,6 +20,7 @@
 from helion._testing import TestCase
 from helion._testing import onlyBackends
 from helion._testing import skipIfRefEager
+from helion._testing import skipUnlessPallas
 import helion.language as hl
 
 _FREE = helion.OutputCodeOptions(allow_helion_deps=False)
@@ -123,6 +125,25 @@ def _run_add_no_helion(code: str, entrypoint: str, shape: tuple[int, int]) -> No
             )
 
 
+def _import_code(code: str, name: str, tmp: str) -> Any:
+    """Import generated standalone source as a module (registered in
+    ``sys.modules`` before ``exec_module`` so any inlined ``@dataclass`` can
+    resolve ``cls.__module__``)."""
+    import importlib.util
+
+    path = _write(tmp, code)
+    spec = importlib.util.spec_from_file_location(name, path)
+    assert spec is not None and spec.loader is not None
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules[name] = mod
+    try:
+        spec.loader.exec_module(mod)
+    except BaseException:
+        sys.modules.pop(name, None)
+        raise
+    return mod
+
+
 @onlyBackends(["triton"])
 @skipIfRefEager("to_code compiles real Triton code; not meaningful in ref-eager")
 class TestToCodeTriton(TestCase):
@@ -205,5 +226,80 @@ def test_kernel_named_like_runtime_helper(self) -> None:
         _run_add_no_helion(code, "get_num_sm", (128, 128))
 
 
+@helion.kernel(backend="pallas", static_shapes=True)
+def pallas_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    out = torch.empty_like(x)
+    for tile in hl.tile(out.size()):
+        out[tile] = x[tile] + y[tile]
+    return out
+
+
+@helion.kernel(
+    config=helion.Config(block_sizes=[8]), static_shapes=False, backend="pallas"
+)
+def pallas_rows_times_two(x: torch.Tensor) -> torch.Tensor:
+    """static_shapes=False Pallas (torch-tensor): the standalone must run at any
+    leading dim (grid + output shape derived from the runtime input)."""
+    n, _d = x.shape
+    out = torch.empty_like(x)
+    for tile in hl.tile(n):
+        out[tile, :] = x[tile, :] * 2.0
+    return out
+
+
+def _pallas_to_code(
+    kernel: Any, args: tuple[object, ...], options: helion.OutputCodeOptions
+) -> str:
+    """``to_code`` with an explicit config. These kernels are never run/autotuned in
+    the test, so ``to_code(config=None)`` has no implicit config to resolve; use the
+    kernel's own config when it declares one, else the backend default."""
+    bound = kernel.bind(args)
+    configs = bound.kernel.configs
+    config = configs[0] if len(configs) == 1 else bound.config_spec.default_config()
+    return bound.to_code(config, options=options)
+
+
+@skipUnlessPallas("Pallas to_code test requires the Pallas backend / TPU")
+@skipIfRefEager("to_code compiles real kernels; not meaningful in ref-eager")
+class TestToCodePallas(TestCase):
+    def test_pallas_torch_standalone_runs(self) -> None:
+        x = torch.randn([256, 256], device=DEVICE, dtype=torch.float32)
+        y = torch.randn([256, 256], device=DEVICE, dtype=torch.float32)
+        code = _pallas_to_code(pallas_add, (x, y), _FREE)
+        # The helion package is never imported; the dependency-free Pallas launcher
+        # is inlined into a local ``helion.runtime`` shim instead.
+        self.assertNotIn("import helion", code)
+        self.assertNotIn("from helion", code)
+        self.assertIn("def default_pallas_launcher(", code)
+        self.assertIn(
+            "_default_pallas_launcher = helion.runtime.default_pallas_launcher", code
+        )
+        self.assertIn("def pallas_add(", code)
+        with tempfile.TemporaryDirectory() as tmp:
+            name = "pallas_add_standalone_test"
+            mod = _import_code(code, name, tmp)
+            try:
+                torch.testing.assert_close(mod.pallas_add(x, y), x + y)
+            finally:
+                sys.modules.pop(name, None)
+
+    def test_pallas_torch_dynamic_shapes(self) -> None:
+        """A static_shapes=False Pallas (torch) standalone runs at other shapes."""
+        d = 128
+        x0 = torch.zeros(512, d, device=DEVICE, dtype=torch.float32)
+        code = _pallas_to_code(pallas_rows_times_two, (x0,), _FREE)
+        with tempfile.TemporaryDirectory() as tmp:
+            name = "pallas_rt2_test"
+            mod = _import_code(code, name, tmp)
+            try:
+                for t in (512, 128, 256):
+                    x = torch.randn(t, d, device=DEVICE, dtype=torch.float32)
+                    out = mod.pallas_rows_times_two(x)
+                    self.assertEqual(tuple(out.shape), (t, d))
+                    torch.testing.assert_close(out, x * 2.0)
+            finally:
+                sys.modules.pop(name, None)
+
+
 if __name__ == "__main__":
     unittest.main()

test/test_pallas.py

@@ -19,6 +19,7 @@
 import helion
 from helion._testing import DEVICE
 from helion._testing import TestCase
+from helion._testing import _bound_test_config
 from helion._testing import code_and_output
 from helion._testing import onlyBackends
 from helion._testing import skipIfPallasInterpret
@@ -900,6 +901,23 @@ def test_topk_divide_and_filter_lowering(self) -> None:
         # (1) the lowering emits the divide-and-filter helper, not jax.lax.top_k
         self.assertIn("_helion_divide_filter_topk", code)
         self.assertNotIn("lax.top_k", code)
+        # (1b) regular output imports the helper from helion; the module parses.
+        self.assertIn(
+            "from helion._compiler.pallas.topk_impl import divide_filter_topk", code
+        )
+        ast.parse(code)
+        # (1c) the dependency-free output embeds the helper source instead (no helion
+        # import) so the standalone is self-contained, and the embed still parses.
+        bound = _topk_pallas_kernel.bind((x, _TOPK_TEST_K))
+        free = bound.to_code(
+            _bound_test_config(bound, block_sizes=[8]),
+            options=helion.OutputCodeOptions(allow_helion_deps=False),
+        )
+        self.assertIn("def divide_filter_topk(", free)
+        self.assertIn("_helion_divide_filter_topk = divide_filter_topk", free)
+        self.assertNotIn("from helion._compiler.pallas.topk_impl import", free)
+        self.assertNotIn("import helion", free)
+        ast.parse(free)
         # (2) correctness vs the exact top-k
         ref_v, ref_i = torch.topk(x, _TOPK_TEST_K, dim=-1, largest=True)
         idx_c = idx.cpu()

test/test_pallas_worklist.py

@@ -749,7 +749,10 @@ def _render_and_exec(self, plan, offset_arrays, upper):
         src, offset_params = render_build_worklist(
             plan, block_expr=str(self.BLOCK), upper_expr=str(upper)
         )
-        namespace: dict = {}
+        # The rendered builder now calls a module-level ``flatten_worklist``
+        # (provided by the backend's embedded-helper inlining in real generated
+        # modules) rather than importing it inline, so supply it here.
+        namespace: dict = {"flatten_worklist": flatten_worklist}
         exec(compile(src, "<build_worklist>", "exec"), namespace)
         builder = namespace["_build_worklist"]
         return src, offset_params, builder(*offset_arrays)
@@ -881,7 +884,10 @@ def test_lo_hi_builder_includes_both_args(self):
         self.assertEqual(offset_params, ["lo", "hi"])
         self.assertIn("jnp.arange(lo.shape[0]", src)
 
-        namespace: dict = {}
+        # The rendered builder now calls a module-level ``flatten_worklist``
+        # (supplied by the backend's embedded-helper inlining in real modules)
+        # rather than importing it inline, so provide it here.
+        namespace: dict = {"flatten_worklist": flatten_worklist}
         exec(compile(src, "<bw>", "exec"), namespace)
         meta = namespace["_build_worklist"](
             jnp.asarray(lo.numpy()), jnp.asarray(hi.numpy())
@@ -1064,6 +1070,22 @@ def fn(q, k, v, q_offsets):
         # builder kwargs; there is no separate launcher name.
         self.assertIn("_compact_build_worklist=_build_worklist", code)
         self.assertIn("def _build_worklist(", code)
+        # Regular output imports ``flatten_worklist`` from helion (the builder calls
+        # it as a module-level name); the dependency-free path embeds it instead.
+        self.assertIn(
+            "from helion.runtime.compact_worklist import flatten_worklist", code
+        )
+        ast.parse(code)
+        # Dependency-free output embeds the helper source at module scope (no helion
+        # import), so the standalone is self-contained, and it still parses.
+        free = bound.to_code(
+            _worklist_config([8]),
+            options=helion.OutputCodeOptions(allow_helion_deps=False),
+        )
+        self.assertIn("def flatten_worklist(", free)
+        self.assertNotIn("from helion.runtime.compact_worklist import", free)
+        self.assertNotIn("import helion", free)
+        ast.parse(free)
         # Offsets arg index is non-empty (q_offsets feeds the builder).
         self.assertRegex(code, r"_compact_offset_arg_indices=\[\d")
         self.assertIn("_compact_num_scalar_prefetch=3", code)