Skip to content

Scope kernel objects to build_dir/<arch>, fixing aie2/aie2p reuse - #166

Open
atassis wants to merge 1 commit into
amd:develfrom
atassis:iron-kernel-obj-arch
Open

Scope kernel objects to build_dir/<arch>, fixing aie2/aie2p reuse#166
atassis wants to merge 1 commit into
amd:develfrom
atassis:iron-kernel-obj-arch

Conversation

@atassis

@atassis atassis commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

Problem

KernelCompilationRule.compile() (iron/common/compilation/base.py) passes a
different --target and aie_runtime_lib include path per arch, but the
output filename it writes to is f"{kernel_name}.o" with no arch component,
and build_dir is one directory for every target
(iron/common/context.py:29). For the 9 aie_kernels/generic/ kernels the
same source compiles to different machine code per arch (measured:
generic/mul.cc gives ELF flags 0x2/1568 B under --target=aie2 and
0x3/1488 B under aie2p, different sha256). CompilationArtifact.is_available_in_filesystem
only compares mtimes and never records which arch an object was built for, so
a shared path lets one arch's build silently hand its object to the other.
This produced 75 bogus ElementwiseMul failures when an npu1 (aie2) run's
mul.o was reused by the following npu2 (aie2p) run in the same build/;
rms_norm's weighted path links generic/mul.cc the same way. 11 more
kernels (gelu, rms_norm, softmax, ...) exist as different sources under
both aie_kernels/aie2/ and aie_kernels/aie2p/ and collide the same way for
the ordinary reason that a freshly-built object is newer than either source.

Separately, dequant names its object f"expand_aie2_{tile_size}.o"
(design.py, op.py) while compiling for the current target: on an aie2p
device it names an aie2p object expand_aie2_....

Fix

CompilationArtifactGraph.move_artifacts now places KernelObjectArtifact/
KernelArchiveArtifact under build_dir/<get_kernel_dir()>/; every other
artifact (MLIR, xclbin, insts.bin) keeps its current flat path, since those
already carry the device in their name via AIEOperatorBase.name. I picked
this over qualifying every object name with the arch: the object name is
independently spelled twice per operator (the Kernel(...) link_with string
in design.py and the KernelObjectArtifact filename in op.py, with no
shared source of truth), across ~12 operators, so qualifying names touches
~24 call sites; the build_dir change is two lines and fixes every current and
future kernel object at once. compile()'s single build_dir.mkdir() moves
to a per-artifact mkdir after move_artifacts, since kernel objects can now
resolve one level deeper.

Also fixed dequant's expand_aie2_ literal to use get_kernel_dir(),
independent of which shape this lands as.

Test / Evidence

iron/tests/compilation/kernel_object_arch_isolation.py builds
ElementwiseMul's real artifact graph for NPU1() and NPU2() (no Peano
invocation, no device) and asserts the two kernel object paths never collide,
and that a real leftover aie2 object is not reported available to the
following aie2p build. Both assertions fail against aab8083 (verified by
reverting compilation/base.py alone and rerunning) and pass after this
change.

Also exercised axpy, softmax (which additionally bundles lut_based_ops.o
into a KernelArchiveArtifact on aie2), rms_norm in both its plain and
weighted (generic/mul.cc-linking) forms, mem_copy, and dequant across
both devices: every kernel object/archive lands under the expected
aie2/aie2p subdirectory with no collisions, and dequant's object name no
longer says aie2 on an aie2p build.

Not run: pytest. This repo's root conftest.py
(pytest_collection_modifyitems) unconditionally calls
aie_utils.DefaultNPURuntime.device() at collection time for every test in
the suite, which probes the attached NPU over XRT regardless of which
test is selected -- out of scope for a device-free change. Confirmed by calling
the two test_* functions directly instead.

Kernel objects/archives are named without an arch component (e.g. "mul.o"),
but KernelCompilationRule.compile() passes a different --target and
runtime-lib -I per arch, and aie_kernels/generic/ sources compile to
different machine code per arch from the same input file. mtime-only
is_available_in_filesystem() can't see any of that, so a shared build_dir
lets one arch's build silently reuse the other's object; this is what
produced 75 bogus ElementwiseMul failures when an npu1 run's mul.o was
picked up by the following npu2 run.

move_artifacts() now places KernelObjectArtifact/KernelArchiveArtifact
under build_dir/<get_kernel_dir()>; every other artifact keeps its current
path. Also fixes dequant's expand_aie2_{tile_size}.o, which hardcoded the
aie2 literal while compiling for the current target.

iron/tests/compilation/kernel_object_arch_isolation.py fails on both
counts before this change.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant