Skip to content

Commit 4c7bcd1

Browse files
atassisatassis
authored
Merge a0156a4 into 27cf75d 
2 parents 27cf75d + a0156a4 commit 4c7bcd1

6 files changed

Lines changed: 180 additions & 112 deletions

File tree

‎iron/common/stream/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@
66
An operator supplies a reference ``nn.Module`` and a placement; these modules turn
77
that into everything stream-dse needs:
88
9+
* :mod:`~iron.common.stream.kernels` -- the AIE kernels and operand layouts a design
10+
runs. Free of ``onnx``, so an operator may name its kernels at import time.
911
* :mod:`~iron.common.stream.ops` -- the registry binding a torch ATen op to its ONNX
10-
form, its stream-dse kernel and IRON's ``aie_kernels`` source.
12+
form and to one of those kernels.
1113
* :mod:`~iron.common.stream.workload` -- ``torch.export`` of the module into the ONNX
1214
workload stream-dse optimizes.
1315
* :mod:`~iron.common.stream.mapping` -- the mapping YAML, named from that same graph.

‎iron/common/stream/kernels.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
# SPDX-FileCopyrightText: Copyright (C) 2026 KU Leuven (MICAS). All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
"""The AIE kernels stream-dse designs run, and the operand layouts they take.
5+
6+
Split out of :mod:`~iron.common.stream.ops` because an operator names its kernels
7+
at import time while the ONNX registry is only needed when a design is built. The
8+
registry pulls in ``onnx``/``onnxscript``, which a core install does not have, so
9+
nothing importable from ``iron.operators`` may reach it.
10+
"""
11+
12+
from __future__ import annotations
13+
14+
from dataclasses import dataclass
15+
from typing import Callable
16+
17+
from iron.common.layout import TiledStridedLayout, tiled_2d
18+
19+
# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The
20+
# operand layouts are the contract the generated DMAs and the compiled kernel
21+
# objects agree on.
22+
# mm.cc takes an 8-row MAC tile when bf16 matmuls run on the bfp16 MACs and a
23+
# 4-row one when they do not.
24+
R, S, T = 4, 8, 8
25+
MAC_ROWS_BFP16 = 8
26+
27+
# Element tile the stream-dse elementwise kernels are written against.
28+
ELEMENTWISE_TILE = (32, 64)
29+
30+
31+
def mac_rows(bfp16_mmul: bool) -> int:
32+
"""Rows of the MAC tile a kernel object compiled this way takes."""
33+
return MAC_ROWS_BFP16 if bfp16_mmul else R
34+
35+
36+
def gemm_layouts(
37+
m: int, k: int, n: int, bfp16_mmul: bool = False
38+
) -> tuple[TiledStridedLayout, ...]:
39+
"""Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands."""
40+
rows = mac_rows(bfp16_mmul)
41+
return (tiled_2d(m, k, rows, S), tiled_2d(k, n, S, T), tiled_2d(m, n, rows, T))
42+
43+
44+
def elementwise_layouts(
45+
nb_operands: int, bfp16_mmul: bool = False
46+
) -> tuple[TiledStridedLayout, ...]:
47+
"""Identical tiled layout for each operand of an elementwise kernel."""
48+
return (tiled_2d(*ELEMENTWISE_TILE, mac_rows(bfp16_mmul), T),) * nb_operands
49+
50+
51+
def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int):
52+
"""The ``mm.cc`` object specialized for one tile shape.
53+
54+
stream-dse emits dimension-suffixed symbols so GEMMs of different tile shapes
55+
coexist in one design (``GemmKernel.function_name``/``zero_name``); rename
56+
``mm.cc``'s unsuffixed symbols to match.
57+
"""
58+
from iron.common.compilation import KernelObjectArtifact, SourceArtifact
59+
60+
suffix = f"{m}_{k}_{n}"
61+
return [
62+
KernelObjectArtifact(
63+
f"mm_{suffix}.o",
64+
dependencies=[SourceArtifact(kernels_dir / kernel_dir / "mm.cc")],
65+
extra_flags=[
66+
f"-DDIM_M={m}",
67+
f"-DDIM_K={k}",
68+
f"-DDIM_N={n}",
69+
"-Dbf16_bf16_ONLY",
70+
# Emulating the matmul on the bfp16 MACs is what makes the 8-row
71+
# MAC tile available, so it and the layouts move together.
72+
"-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
73+
"-DROUND_CONV_EVEN",
74+
],
75+
rename_symbols={
76+
"matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}",
77+
"zero_bf16": f"zero_bf16_{suffix}",
78+
},
79+
)
80+
]
81+
82+
83+
@dataclass(frozen=True)
84+
class StreamKernel:
85+
"""An AIE kernel: its stream-dse identity, its source, and its operand layouts.
86+
87+
``source``/``subdir`` name the file in IRON's ``aie_kernels`` library the same
88+
way the hand-written operators do (``subdir=None`` means the device directory,
89+
e.g. ``aie2p``). The object name must equal the kernel's ``linkwith_name`` in
90+
stream-dse, since the generated MLIR links against it.
91+
"""
92+
93+
key: str # stream-dse AIEKernels key
94+
layouts: Callable[..., tuple[TiledStridedLayout, ...]]
95+
source: str | None = None
96+
subdir: str | None = None
97+
artifacts: Callable | None = None # overrides source/subdir when tile-specialized
98+
99+
def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs):
100+
"""Compilation artifacts building this kernel's object file."""
101+
if self.artifacts is not None:
102+
return self.artifacts(kernels_dir, kernel_dir, **kwargs)
103+
from iron.common.compilation import KernelObjectArtifact, SourceArtifact
104+
105+
subdir = self.subdir or kernel_dir
106+
return [
107+
KernelObjectArtifact(
108+
f"{self.source}.o",
109+
dependencies=[
110+
SourceArtifact(kernels_dir / subdir / f"{self.source}.cc")
111+
],
112+
)
113+
]
114+
115+
116+
GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts)
117+
SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu")
118+
ELTWISE_MUL = StreamKernel(
119+
key="eltwise_mul",
120+
layouts=lambda: elementwise_layouts(3),
121+
source="mul",
122+
)

‎iron/common/stream/ops.py‎

Lines changed: 4 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212
declared with :func:`custom_op`, which gives them a schema in a private domain so
1313
the exporter emits them as a single node.
1414
15-
Supporting a new op is one :class:`StreamKernel` plus one :data:`TORCH_OPS` entry --
16-
the kernel source is IRON's existing ``aie_kernels/<dir>/<name>.cc``, exactly as the
17-
hand-written operators use it.
15+
Supporting a new op is one :class:`~iron.common.stream.kernels.StreamKernel` plus one
16+
:data:`TORCH_OPS` entry -- the kernel source is IRON's existing
17+
``aie_kernels/<dir>/<name>.cc``, exactly as the hand-written operators use it.
1818
"""
1919

2020
from __future__ import annotations
@@ -27,18 +27,7 @@
2727
from onnxscript import opset18
2828
from onnxscript.values import Op, Opset
2929

30-
from iron.common.layout import TiledStridedLayout, tiled_2d
31-
32-
# Intrinsic MAC tile dimensions of the aie2p kernels stream-dse targets. The
33-
# operand layouts are the contract the generated DMAs and the compiled kernel
34-
# objects agree on.
35-
# mm.cc takes an 8-row MAC tile when bf16 matmuls run on the bfp16 MACs and a
36-
# 4-row one when they do not.
37-
R, S, T = 4, 8, 8
38-
MAC_ROWS_BFP16 = 8
39-
40-
# Element tile the stream-dse elementwise kernels are written against.
41-
ELEMENTWISE_TILE = (32, 64)
30+
from iron.common.stream.kernels import ELTWISE_MUL, GEMM, SILU, StreamKernel
4231

4332
# Private domain for ops that exist as an AIE kernel but not as an ONNX operator.
4433
CUSTOM_DOMAIN = Opset("com.example", 1)
@@ -59,99 +48,6 @@ def custom_op(name: str, arity: int = 1) -> Op:
5948
return Op(CUSTOM_DOMAIN, name, schema)
6049

6150

62-
def mac_rows(bfp16_mmul: bool) -> int:
63-
"""Rows of the MAC tile a kernel object compiled this way takes."""
64-
return MAC_ROWS_BFP16 if bfp16_mmul else R
65-
66-
67-
def gemm_layouts(
68-
m: int, k: int, n: int, bfp16_mmul: bool = False
69-
) -> tuple[TiledStridedLayout, ...]:
70-
"""Layouts of a GEMM's ``A[m,k]``, ``B[k,n]`` and ``C[m,n]`` operands."""
71-
rows = mac_rows(bfp16_mmul)
72-
return (tiled_2d(m, k, rows, S), tiled_2d(k, n, S, T), tiled_2d(m, n, rows, T))
73-
74-
75-
def elementwise_layouts(
76-
nb_operands: int, bfp16_mmul: bool = False
77-
) -> tuple[TiledStridedLayout, ...]:
78-
"""Identical tiled layout for each operand of an elementwise kernel."""
79-
return (tiled_2d(*ELEMENTWISE_TILE, mac_rows(bfp16_mmul), T),) * nb_operands
80-
81-
82-
def _gemm_artifacts(kernels_dir, kernel_dir, m: int, k: int, n: int):
83-
"""The ``mm.cc`` object specialized for one tile shape.
84-
85-
stream-dse emits dimension-suffixed symbols so GEMMs of different tile shapes
86-
coexist in one design (``GemmKernel.function_name``/``zero_name``); rename
87-
``mm.cc``'s unsuffixed symbols to match.
88-
"""
89-
from iron.common.compilation import KernelObjectArtifact, SourceArtifact
90-
91-
suffix = f"{m}_{k}_{n}"
92-
return [
93-
KernelObjectArtifact(
94-
f"mm_{suffix}.o",
95-
dependencies=[SourceArtifact(kernels_dir / kernel_dir / "mm.cc")],
96-
extra_flags=[
97-
f"-DDIM_M={m}",
98-
f"-DDIM_K={k}",
99-
f"-DDIM_N={n}",
100-
"-Dbf16_bf16_ONLY",
101-
# Emulating the matmul on the bfp16 MACs is what makes the 8-row
102-
# MAC tile available, so it and the layouts move together.
103-
"-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16",
104-
"-DROUND_CONV_EVEN",
105-
],
106-
rename_symbols={
107-
"matmul_bf16_bf16": f"matmul_bf16_bf16_{suffix}",
108-
"zero_bf16": f"zero_bf16_{suffix}",
109-
},
110-
)
111-
]
112-
113-
114-
@dataclass(frozen=True)
115-
class StreamKernel:
116-
"""An AIE kernel: its stream-dse identity, its source, and its operand layouts.
117-
118-
``source``/``subdir`` name the file in IRON's ``aie_kernels`` library the same
119-
way the hand-written operators do (``subdir=None`` means the device directory,
120-
e.g. ``aie2p``). The object name must equal the kernel's ``linkwith_name`` in
121-
stream-dse, since the generated MLIR links against it.
122-
"""
123-
124-
key: str # stream-dse AIEKernels key
125-
layouts: Callable[..., tuple[TiledStridedLayout, ...]]
126-
source: str | None = None
127-
subdir: str | None = None
128-
artifacts: Callable | None = None # overrides source/subdir when tile-specialized
129-
130-
def kernel_artifacts(self, kernels_dir, kernel_dir, **kwargs):
131-
"""Compilation artifacts building this kernel's object file."""
132-
if self.artifacts is not None:
133-
return self.artifacts(kernels_dir, kernel_dir, **kwargs)
134-
from iron.common.compilation import KernelObjectArtifact, SourceArtifact
135-
136-
subdir = self.subdir or kernel_dir
137-
return [
138-
KernelObjectArtifact(
139-
f"{self.source}.o",
140-
dependencies=[
141-
SourceArtifact(kernels_dir / subdir / f"{self.source}.cc")
142-
],
143-
)
144-
]
145-
146-
147-
GEMM = StreamKernel(key="gemm", layouts=gemm_layouts, artifacts=_gemm_artifacts)
148-
SILU = StreamKernel(key="silu", layouts=lambda: elementwise_layouts(2), source="silu")
149-
ELTWISE_MUL = StreamKernel(
150-
key="eltwise_mul",
151-
layouts=lambda: elementwise_layouts(3),
152-
source="mul",
153-
)
154-
15551
Silu = custom_op("Silu")
15652

15753

‎iron/operators/swiglu_prefill_stream/op.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
)
1515
from iron.common.device_utils import get_kernel_dir
1616
from iron.common.sequence import OperatorSequence
17-
from iron.common.stream.ops import ELTWISE_MUL, GEMM, SILU
17+
from iron.common.stream.kernels import ELTWISE_MUL, GEMM, SILU
1818

1919

2020
@dataclass

‎iron/tests/core_install.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
"""What a core install must be able to import.
5+
6+
``requirements.txt`` is the core install; ``requirements_stream.txt`` adds
7+
``onnx``/``onnxscript`` for the one stream-dse-backed operator, whose test skips
8+
itself when they are absent. That promise only holds if nothing reachable from
9+
``import iron.operators`` pulls them in -- a hard import there fails pytest
10+
collection for every operator, not just that one.
11+
12+
The dependencies are usually installed in the environment running this, so the
13+
check has to happen in a child interpreter that cannot see them.
14+
"""
15+
16+
import subprocess
17+
import sys
18+
from pathlib import Path
19+
20+
import pytest
21+
22+
_REPO_ROOT = Path(__file__).resolve().parents[2]
23+
24+
_BLOCK_AND_IMPORT = """
25+
import sys
26+
27+
28+
class Blocked:
29+
def find_spec(self, name, path=None, target=None):
30+
if name.split(".")[0] in ("onnx", "onnxscript"):
31+
raise ImportError("No module named %r" % name)
32+
return None
33+
34+
35+
sys.meta_path.insert(0, Blocked())
36+
import {module}
37+
"""
38+
39+
40+
@pytest.mark.parametrize("module", ["iron.operators", "iron.common.stream.kernels"])
41+
def test_imports_without_stream_dependencies(module):
42+
result = subprocess.run(
43+
[sys.executable, "-c", _BLOCK_AND_IMPORT.format(module=module)],
44+
cwd=_REPO_ROOT,
45+
capture_output=True,
46+
text=True,
47+
)
48+
assert result.returncode == 0, result.stderr

‎iron/tests/stream/kernel_layouts.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
77
stream-dse generates the DMAs that feed the kernel objects IRON compiles from
88
``aie_kernels``; both sides must agree on how an operand is tiled in memory. The
9-
layouts declared in :mod:`iron.common.stream.ops` are that contract. They happen
9+
layouts declared in :mod:`iron.common.stream.kernels` are that contract. They happen
1010
to coincide with stream-dse's built-in kernel layouts today, so no override is
1111
needed -- this test fails if a future stream-dse release changes them, which would
1212
otherwise corrupt results silently.
@@ -20,7 +20,7 @@
2020

2121
from stream.compiler.kernels import AIEKernels # noqa: E402
2222

23-
from iron.common.stream.ops import ( # noqa: E402
23+
from iron.common.stream.kernels import ( # noqa: E402
2424
ELTWISE_MUL,
2525
GEMM,
2626
SILU,

0 commit comments

Comments
 (0)