From a0156a4858da537cdaaae6df47159e55c8180062 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Tue, 25 Aug 2026 15:15:55 +0300 Subject: [PATCH] stream: keep the ONNX registry off the operator import path iron/common/stream/__init__.py already states the rule -- its submodules need onnx/pyyaml, "so importing an operator must not pull them in" -- and requirements_stream.txt promises the stream operator's test skips itself when stream-dse is absent. Neither holds. iron/operators/__init__.py re-exports SwiGLUPrefillStream, swiglu_prefill_stream/op.py imports ELTWISE_MUL, GEMM and SILU from iron.common.stream.ops, and that module imports onnx and onnxscript at module scope for the exporter registry. So a core install cannot import iron.operators at all. The failure is ModuleNotFoundError at collection, which takes out every operator's test, not just the one operator the optional dependency belongs to. The three names op.py wants are StreamKernel instances -- a stream-dse key, a kernel source and its operand layouts. Nothing about them is ONNX. Move that half to iron/common/stream/kernels.py and leave ops.py the registry it is named for: the private-domain schemas, the translation table, and TORCH_OPS binding a torch ATen op to one of those kernels. ops.py imports the kernels it references, so the one-op-per-entry structure is unchanged; kernel_layouts.py takes its layouts from the new module. A module that does not import onnx cannot regain the dependency by accident, which is the point of splitting rather than deferring the import into the functions that use it. iron/tests/core_install.py pins the contract: a child interpreter with onnx and onnxscript hidden from sys.meta_path must still import iron.operators and iron.common.stream.kernels. It fails before this change and passes after. The ONNX path is unaffected -- TORCH_OPS still resolves three ops and translation_table three entries. --- iron/common/stream/__init__.py | 4 +- iron/common/stream/kernels.py | 122 +++++++++++++++++++++ iron/common/stream/ops.py | 112 +------------------ iron/operators/swiglu_prefill_stream/op.py | 2 +- iron/tests/core_install.py | 48 ++++++++ iron/tests/stream/kernel_layouts.py | 4 +- 6 files changed, 180 insertions(+), 112 deletions(-) create mode 100644 iron/common/stream/kernels.py create mode 100644 iron/tests/core_install.py diff --git a/iron/common/stream/__init__.py b/iron/common/stream/__init__.py index 1e35c3faee..4149a336e2 100644 --- a/iron/common/stream/__init__.py +++ b/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. diff --git a/iron/common/stream/kernels.py b/iron/common/stream/kernels.py new file mode 100644 index 0000000000..6fc6853a4d --- /dev/null +++ b/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", +) diff --git a/iron/common/stream/ops.py b/iron/common/stream/ops.py index 518e7133da..4204211790 100644 --- a/iron/common/stream/ops.py +++ b/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//.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//.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") diff --git a/iron/operators/swiglu_prefill_stream/op.py b/iron/operators/swiglu_prefill_stream/op.py index 96a60f2db5..b616a476b1 100644 --- a/iron/operators/swiglu_prefill_stream/op.py +++ b/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 diff --git a/iron/tests/core_install.py b/iron/tests/core_install.py new file mode 100644 index 0000000000..2cd194eb90 --- /dev/null +++ b/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 diff --git a/iron/tests/stream/kernel_layouts.py b/iron/tests/stream/kernel_layouts.py index d6efdad5be..3bec906baf 100644 --- a/iron/tests/stream/kernel_layouts.py +++ b/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,