stream: keep the ONNX registry off the operator import path

iron/common/stream/__init__.py

@@ -6,8 +6,10 @@
 An operator supplies a reference ``nn.Module`` and a placement; these modules turn
 that into everything stream-dse needs:
 
+* :mod:`~iron.common.stream.kernels` -- the AIE kernels and operand layouts a design
+  runs. Free of ``onnx``, so an operator may name its kernels at import time.
 * :mod:`~iron.common.stream.ops` -- the registry binding a torch ATen op to its ONNX
-  form, its stream-dse kernel and IRON's ``aie_kernels`` source.
+  form and to one of those kernels.
 * :mod:`~iron.common.stream.workload` -- ``torch.export`` of the module into the ONNX
   workload stream-dse optimizes.
 * :mod:`~iron.common.stream.mapping` -- the mapping YAML, named from that same graph.

iron/common/stream/kernels.py

@@ -0,0 +1,122 @@
+# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""The AIE kernels stream-dse designs run, and the operand layouts they take.
+
+Split out of :mod:`~iron.common.stream.ops` because an operator names its kernels
+at import time while the ONNX registry is only needed when a design is built. The
+registry pulls in ``onnx``/``onnxscript``, which a core install does not have, so
+nothing importable from ``iron.operators`` may reach it.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+from iron.common.layout import TiledStridedLayout, tiled_2d
+
+# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The
+# operand layouts are the contract the generated DMAs and the compiled kernel
+# objects agree on.
+# mm.cc takes an 8-row MAC tile when bf16 matmuls run on the bfp16 MACs and a
+# 4-row one when they do not.
+R, S, T = 4, 8, 8
+MAC_ROWS_BFP16 = 8
+
+# Element tile the stream-dse elementwise kernels are written against.
+ELEMENTWISE_TILE = (32, 64)
+
+
+def mac_rows(bfp16_mmul: bool) -> int:
+    """Rows of the MAC tile a kernel object compiled this way takes."""
+    return MAC_ROWS_BFP16 if bfp16_mmul else R
+
+
+def gemm_layouts(
+    m: int, k: int, n: int, bfp16_mmul: bool = False
+) -> tuple[TiledStridedLayout, ...]:
+    """Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands."""
+    rows = mac_rows(bfp16_mmul)
+    return (tiled_2d(m, k, rows, S), tiled_2d(k, n, S, T), tiled_2d(m, n, rows, T))
+
+
+def elementwise_layouts(
+    nb_operands: int, bfp16_mmul: bool = False
+) -> tuple[TiledStridedLayout, ...]:
+    """Identical tiled layout for each operand of an elementwise kernel."""
+    return (tiled_2d(*ELEMENTWISE_TILE, mac_rows(bfp16_mmul), T),) * nb_operands
+
+
+def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int):
+    """The ``mm.cc`` object specialized for one tile shape.
+
+    stream-dse emits dimension-suffixed symbols so GEMMs of different tile shapes
+    coexist in one design (``GemmKernel.function_name``/``zero_name``); rename
+    ``mm.cc``'s unsuffixed symbols to match.
+    """
+    from iron.common.compilation import KernelObjectArtifact, SourceArtifact
+
+    suffix = f"{m}_{k}_{n}"
+    return [
+        KernelObjectArtifact(
+            f"mm_{suffix}.o",
+            dependencies=[SourceArtifact(kernels_dir / kernel_dir / "mm.cc")],
+            extra_flags=[
+                f"-DDIM_M={m}",
+                f"-DDIM_K={k}",
+                f"-DDIM_N={n}",
+                "-Dbf16_bf16_ONLY",
+                # Emulating the matmul on the bfp16 MACs is what makes the 8-row
+                # MAC tile available, so it and the layouts move together.
+                "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
+                "-DROUND_CONV_EVEN",
+            ],
+            rename_symbols={
+                "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}",
+                "zero_bf16": f"zero_bf16_{suffix}",
+            },
+        )
+    ]
+
+
+@dataclass(frozen=True)
+class StreamKernel:
+    """An AIE kernel: its stream-dse identity, its source, and its operand layouts.
+
+    ``source``/``subdir`` name the file in IRON's ``aie_kernels`` library the same
+    way the hand-written operators do (``subdir=None`` means the device directory,
+    e.g. ``aie2p``). The object name must equal the kernel's ``linkwith_name`` in
+    stream-dse, since the generated MLIR links against it.
+    """
+
+    key: str  # stream-dse AIEKernels key
+    layouts: Callable[..., tuple[TiledStridedLayout, ...]]
+    source: str | None = None
+    subdir: str | None = None
+    artifacts: Callable | None = None  # overrides source/subdir when tile-specialized
+
+    def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs):
+        """Compilation artifacts building this kernel's object file."""
+        if self.artifacts is not None:
+            return self.artifacts(kernels_dir, kernel_dir, **kwargs)
+        from iron.common.compilation import KernelObjectArtifact, SourceArtifact
+
+        subdir = self.subdir or kernel_dir
+        return [
+            KernelObjectArtifact(
+                f"{self.source}.o",
+                dependencies=[
+                    SourceArtifact(kernels_dir / subdir / f"{self.source}.cc")
+                ],
+            )
+        ]
+
+
+GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts)
+SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu")
+ELTWISE_MUL = StreamKernel(
+    key="eltwise_mul",
+    layouts=lambda: elementwise_layouts(3),
+    source="mul",
+)

iron/common/stream/ops.py

@@ -12,9 +12,9 @@
 declared with :func:`custom_op`, which gives them a schema in a private domain so
 the exporter emits them as a single node.
 
-Supporting a new op is one :class:`StreamKernel` plus one :data:`TORCH_OPS` entry --
-the kernel source is IRON's existing ``aie_kernels/<dir>/<name>.cc``, exactly as the
-hand-written operators use it.
+Supporting a new op is one :class:`~iron.common.stream.kernels.StreamKernel` plus one
+:data:`TORCH_OPS` entry -- the kernel source is IRON's existing
+``aie_kernels/<dir>/<name>.cc``, exactly as the hand-written operators use it.
 """
 
 from __future__ import annotations
@@ -27,18 +27,7 @@
 from onnxscript import opset18
 from onnxscript.values import Op, Opset
 
-from iron.common.layout import TiledStridedLayout, tiled_2d
-
-# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The
-# operand layouts are the contract the generated DMAs and the compiled kernel
-# objects agree on.
-# mm.cc takes an 8-row MAC tile when bf16 matmuls run on the bfp16 MACs and a
-# 4-row one when they do not.
-R, S, T = 4, 8, 8
-MAC_ROWS_BFP16 = 8
-
-# Element tile the stream-dse elementwise kernels are written against.
-ELEMENTWISE_TILE = (32, 64)
+from iron.common.stream.kernels import ELTWISE_MUL, GEMM, SILU, StreamKernel
 
 # Private domain for ops that exist as an AIE kernel but not as an ONNX operator.
 CUSTOM_DOMAIN = Opset("com.example", 1)
@@ -59,99 +48,6 @@ def custom_op(name: str, arity: int = 1) -> Op:
     return Op(CUSTOM_DOMAIN, name, schema)
 
 
-def mac_rows(bfp16_mmul: bool) -> int:
-    """Rows of the MAC tile a kernel object compiled this way takes."""
-    return MAC_ROWS_BFP16 if bfp16_mmul else R
-
-
-def gemm_layouts(
-    m: int, k: int, n: int, bfp16_mmul: bool = False
-) -> tuple[TiledStridedLayout, ...]:
-    """Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands."""
-    rows = mac_rows(bfp16_mmul)
-    return (tiled_2d(m, k, rows, S), tiled_2d(k, n, S, T), tiled_2d(m, n, rows, T))
-
-
-def elementwise_layouts(
-    nb_operands: int, bfp16_mmul: bool = False
-) -> tuple[TiledStridedLayout, ...]:
-    """Identical tiled layout for each operand of an elementwise kernel."""
-    return (tiled_2d(*ELEMENTWISE_TILE, mac_rows(bfp16_mmul), T),) * nb_operands
-
-
-def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int):
-    """The ``mm.cc`` object specialized for one tile shape.
-
-    stream-dse emits dimension-suffixed symbols so GEMMs of different tile shapes
-    coexist in one design (``GemmKernel.function_name``/``zero_name``); rename
-    ``mm.cc``'s unsuffixed symbols to match.
-    """
-    from iron.common.compilation import KernelObjectArtifact, SourceArtifact
-
-    suffix = f"{m}_{k}_{n}"
-    return [
-        KernelObjectArtifact(
-            f"mm_{suffix}.o",
-            dependencies=[SourceArtifact(kernels_dir / kernel_dir / "mm.cc")],
-            extra_flags=[
-                f"-DDIM_M={m}",
-                f"-DDIM_K={k}",
-                f"-DDIM_N={n}",
-                "-Dbf16_bf16_ONLY",
-                # Emulating the matmul on the bfp16 MACs is what makes the 8-row
-                # MAC tile available, so it and the layouts move together.
-                "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
-                "-DROUND_CONV_EVEN",
-            ],
-            rename_symbols={
-                "matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}",
-                "zero_bf16": f"zero_bf16_{suffix}",
-            },
-        )
-    ]
-
-
-@dataclass(frozen=True)
-class StreamKernel:
-    """An AIE kernel: its stream-dse identity, its source, and its operand layouts.
-
-    ``source``/``subdir`` name the file in IRON's ``aie_kernels`` library the same
-    way the hand-written operators do (``subdir=None`` means the device directory,
-    e.g. ``aie2p``). The object name must equal the kernel's ``linkwith_name`` in
-    stream-dse, since the generated MLIR links against it.
-    """
-
-    key: str  # stream-dse AIEKernels key
-    layouts: Callable[..., tuple[TiledStridedLayout, ...]]
-    source: str | None = None
-    subdir: str | None = None
-    artifacts: Callable | None = None  # overrides source/subdir when tile-specialized
-
-    def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs):
-        """Compilation artifacts building this kernel's object file."""
-        if self.artifacts is not None:
-            return self.artifacts(kernels_dir, kernel_dir, **kwargs)
-        from iron.common.compilation import KernelObjectArtifact, SourceArtifact
-
-        subdir = self.subdir or kernel_dir
-        return [
-            KernelObjectArtifact(
-                f"{self.source}.o",
-                dependencies=[
-                    SourceArtifact(kernels_dir / subdir / f"{self.source}.cc")
-                ],
-            )
-        ]
-
-
-GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts)
-SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu")
-ELTWISE_MUL = StreamKernel(
-    key="eltwise_mul",
-    layouts=lambda: elementwise_layouts(3),
-    source="mul",
-)
-
 Silu = custom_op("Silu")
 
 

iron/operators/swiglu_prefill_stream/op.py

@@ -14,7 +14,7 @@
 )
 from iron.common.device_utils import get_kernel_dir
 from iron.common.sequence import OperatorSequence
-from iron.common.stream.ops import ELTWISE_MUL, GEMM, SILU
+from iron.common.stream.kernels import ELTWISE_MUL, GEMM, SILU
 
 
 @dataclass

iron/tests/core_install.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""What a core install must be able to import.
+
+``requirements.txt`` is the core install; ``requirements_stream.txt`` adds
+``onnx``/``onnxscript`` for the one stream-dse-backed operator, whose test skips
+itself when they are absent. That promise only holds if nothing reachable from
+``import iron.operators`` pulls them in -- a hard import there fails pytest
+collection for every operator, not just that one.
+
+The dependencies are usually installed in the environment running this, so the
+check has to happen in a child interpreter that cannot see them.
+"""
+
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+_BLOCK_AND_IMPORT = """
+import sys
+
+
+class Blocked:
+    def find_spec(self, name, path=None, target=None):
+        if name.split(".")[0] in ("onnx", "onnxscript"):
+            raise ImportError("No module named %r" % name)
+        return None
+
+
+sys.meta_path.insert(0, Blocked())
+import {module}
+"""
+
+
+@pytest.mark.parametrize("module", ["iron.operators", "iron.common.stream.kernels"])
+def test_imports_without_stream_dependencies(module):
+    result = subprocess.run(
+        [sys.executable, "-c", _BLOCK_AND_IMPORT.format(module=module)],
+        cwd=_REPO_ROOT,
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0, result.stderr

iron/tests/stream/kernel_layouts.py

@@ -6,7 +6,7 @@
 
 stream-dse generates the DMAs that feed the kernel objects IRON compiles from
 ``aie_kernels``; both sides must agree on how an operand is tiled in memory. The
-layouts declared in :mod:`iron.common.stream.ops` are that contract. They happen
+layouts declared in :mod:`iron.common.stream.kernels` are that contract. They happen
 to coincide with stream-dse's built-in kernel layouts today, so no override is
 needed -- this test fails if a future stream-dse release changes them, which would
 otherwise corrupt results silently.
@@ -20,7 +20,7 @@
 
 from stream.compiler.kernels import AIEKernels  # noqa: E402
 
-from iron.common.stream.ops import (  # noqa: E402
+from iron.common.stream.kernels import (  # noqa: E402
     ELTWISE_MUL,
     GEMM,
     SILU,