From 70f46f9166cd3fcd274f78ae49859db7f0c1ea83 Mon Sep 17 00:00:00 2001 From: Dunfan Lu Date: Tue, 28 Jul 2026 03:41:19 -0700 Subject: [PATCH] Embed in-kernel Pallas runtime helpers into generated code The generated Pallas kernel calls two helion-defined, pure-`jax` helpers -- `divide_filter_topk` (aten.topk lowering, via `_helion_divide_filter_topk`) and `flatten_worklist` (compact-worklist builder) -- which were pulled in with `from helion... import ...`, leaving a helion dependency that blocks precompiling to a helion-free standalone. Add a `Backend.embedded_helper_source(body)` hook (default `""`) that `BoundKernel.to_code` injects between the imports and the kernel body. The Pallas backend overrides it to inline the source of whichever helper `body` references -- module docstring and `from __future__` stripped by `_embed_source` (located via `ast`, so a docstring containing a triple-quote can't corrupt the output) -- plus a `_helion_divide_filter_topk = divide_filter_topk` alias. The topk `library_imports` entry and the inline `flatten_worklist` import in `render_build_worklist` are dropped. Both helper modules import only `jax`/stdlib (and now carry a note to stay free of any `helion` import, which would trip the precompiler's substring guard), so the embedded module stays dependency-free. This runs for regular kernels too, keeping a single code path. Tests: the two builder tests that `exec` the rendered `_build_worklist` now supply `flatten_worklist` (no longer imported inline); `test_matching_kernel_generates_compact` and `test_topk_divide_and_filter_lowering` assert the embedded def + alias are present, the old helion import is gone, and the module still `ast.parse`s (guarding against a corrupt embed). Also repairs the window-guard test's stale call to the 2-arg `_get_vmem_limit_bytes`. Verified on TPU: topk (3 passed), compact guard/render/builder (10 passed); full `test_pallas_compact_worklist.py` shows only the 6 pre-existing `scratch_types` JAX-skew failures (identical to base), 66 passed. `ruff`/`pyrefly` clean (83 baseline); H100 Triton unaffected (`embedded_helper_source` is a no-op there). stack-info: PR: https://github.com/pytorch/helion/pull/3187, branch: AmesingFlank/stack/101 --- helion/_compiler/backend.py | 11 ++++ helion/_compiler/pallas/backend.py | 56 ++++++++++++++++++++- helion/_compiler/pallas/compact_worklist.py | 1 - helion/_compiler/pallas/topk_impl.py | 5 ++ helion/runtime/compact_worklist.py | 6 +++ helion/runtime/kernel.py | 5 ++ test/test_pallas.py | 7 +++ test/test_pallas_worklist.py | 17 ++++++- 8 files changed, 104 insertions(+), 4 deletions(-) diff --git a/helion/_compiler/backend.py b/helion/_compiler/backend.py index 431c7768c..423ed941c 100644 --- a/helion/_compiler/backend.py +++ b/helion/_compiler/backend.py @@ -703,6 +703,17 @@ def library_imports(self) -> dict[str, str]: """ ... + def embedded_helper_source(self, body: str) -> str: + """Source of in-kernel runtime helpers to inline into the generated module. + + Backends that call a helion-defined helper from inside the generated + kernel can return its source here (instead of importing it) so the output + is self-contained -- which lets the precompiler produce a helion-free + standalone. Only helpers actually referenced in ``body`` should be + emitted. Injected between the imports and the kernel body. Default: none. + """ + return "" + def launcher_keyword_args(self, config: Config, *, has_barrier: bool) -> list[str]: return [] diff --git a/helion/_compiler/pallas/backend.py b/helion/_compiler/pallas/backend.py index 118e7c774..e094bc090 100644 --- a/helion/_compiler/pallas/backend.py +++ b/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 @@ -35,6 +36,35 @@ InductorOpOverrides = OpsHandler[Any] +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", @@ -169,9 +199,33 @@ 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", - "_helion_divide_filter_topk": "from helion._compiler.pallas.topk_impl import divide_filter_topk as _helion_divide_filter_topk", } + 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. Embedding their source -- instead of importing them from helion -- + makes the generated module self-contained (so it precompiles to a + helion-free standalone) and applies to regular kernels too. + """ + 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) + # 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. diff --git a/helion/_compiler/pallas/compact_worklist.py b/helion/_compiler/pallas/compact_worklist.py index 030e02224..a2b382daf 100644 --- a/helion/_compiler/pallas/compact_worklist.py +++ b/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}", diff --git a/helion/_compiler/pallas/topk_impl.py b/helion/_compiler/pallas/topk_impl.py index 119814603..50b8a246c 100644 --- a/helion/_compiler/pallas/topk_impl.py +++ b/helion/_compiler/pallas/topk_impl.py @@ -37,6 +37,11 @@ if TYPE_CHECKING: import jax +# NOTE: this module is embedded verbatim into helion.precompile standalones (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 diff --git a/helion/runtime/compact_worklist.py b/helion/runtime/compact_worklist.py index 6b811f1c2..308e97c6d 100644 --- a/helion/runtime/compact_worklist.py +++ b/helion/runtime/compact_worklist.py @@ -27,6 +27,12 @@ if TYPE_CHECKING: from jax import Array +# NOTE: this module is embedded verbatim into helion.precompile standalones (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. diff --git a/helion/runtime/kernel.py b/helion/runtime/kernel.py index 27451f40f..f75a279e0 100644 --- a/helion/runtime/kernel.py +++ b/helion/runtime/kernel.py @@ -843,6 +843,11 @@ def to_code( ast.fix_missing_locations(body_root) imports = "\n".join(import_lines) body = unparse(body_root, output_origin_lines=output_origin_lines) + # Inline any in-kernel runtime helpers (e.g. Pallas topk / + # compact-worklist) so the generated module is self-contained. + embedded = self.env.backend.embedded_helper_source(body) + if embedded: + body = f"{embedded}\n\n\n{body}" if imports: return f"from __future__ import annotations\n\n{imports}\n\n{body}" return f"from __future__ import annotations\n\n{body}" diff --git a/test/test_pallas.py b/test/test_pallas.py index aee4fe366..3b0adb57f 100644 --- a/test/test_pallas.py +++ b/test/test_pallas.py @@ -900,6 +900,13 @@ 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) the helper source is embedded (no helion import) so the kernel is + # self-contained / precompilable, and the embed doesn't corrupt the module. + self.assertIn("def divide_filter_topk(", code) + self.assertIn("_helion_divide_filter_topk = divide_filter_topk", code) + self.assertNotIn("from helion._compiler.pallas.topk_impl import", code) + self.assertNotIn("import helion", code) + ast.parse(code) # (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() diff --git a/test/test_pallas_worklist.py b/test/test_pallas_worklist.py index f4c5882c2..276b35c8d 100644 --- a/test/test_pallas_worklist.py +++ b/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, "", "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, "", "exec"), namespace) meta = namespace["_build_worklist"]( jnp.asarray(lo.numpy()), jnp.asarray(hi.numpy()) @@ -1064,6 +1070,13 @@ 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) + # ``flatten_worklist`` is embedded at module scope (no Helion import), + # so the generated module is self-contained for precompilation rather + # than importing the runtime helper inside ``_build_worklist``. + self.assertIn("def flatten_worklist(", code) + self.assertNotIn("from helion.runtime.compact_worklist import", code) + # The embedded helper source must not corrupt the generated module. + ast.parse(code) # 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)