diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index a57e033aa..3bc39c220 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -119,9 +119,12 @@ def compile( build_dir: str = "build", dry_run: bool = False, ) -> None: - if not Path(build_dir).exists() and not dry_run: - Path(build_dir).mkdir(parents=True, exist_ok=True) artifacts.move_artifacts(build_dir) + if not dry_run: + # move_artifacts() may place kernel objects under a per-arch + # subdirectory of build_dir, so mkdir per artifact rather than once. + for artifact in artifacts.bfs(): + Path(artifact.filename).parent.mkdir(parents=True, exist_ok=True) artifacts.populate_availability_from_filesystem() plan_steps = plan(rules, artifacts) if not dry_run: @@ -215,10 +218,22 @@ def get_worklist(self, kind: type | tuple[type, ...]) -> list[CompilationArtifac ] def move_artifacts(self, new_root: str) -> None: - """Make all artifacts paths point into a build directory""" + """Make all artifact paths point into a build directory. + + Kernel objects/archives get an extra get_kernel_dir() segment: their + filename (e.g. "mul.o") does not encode arch, but their compiled + content does, and is_available_in_filesystem() only compares mtimes -- + so two arches sharing one path would silently reuse each other's object. + """ + kernel_dir = None for artifact in self.bfs(): if not Path(artifact.filename).is_absolute(): - artifact.filename = str(Path(new_root) / Path(artifact.filename).name) + root = new_root + if isinstance(artifact, (KernelObjectArtifact, KernelArchiveArtifact)): + if kernel_dir is None: + kernel_dir = get_kernel_dir() + root = Path(new_root) / kernel_dir + artifact.filename = str(Path(root) / Path(artifact.filename).name) def add(self, artifact: CompilationArtifact) -> None: self.artifacts.append(artifact) @@ -509,23 +524,36 @@ def _link_build_outputs_into(work_dir: Path, build_dir: Path) -> None: aiecc resolves an MLIR module's relative kernel-object references (e.g. ``link_with = "axpy.o"``, produced by KernelCompilationRule / - ArchiveCompilationRule into the flat build_dir) against work_dir, since - that's where compile_mlir_module() writes its own copy of the MLIR - source. Symlinking makes those lookups succeed without copying kernel - objects into every artifact's own work_dir. + ArchiveCompilationRule) against work_dir, since that's where + compile_mlir_module() writes its own copy of the MLIR source. Symlinking + makes those lookups succeed without copying kernel objects into every + artifact's own work_dir. + + Kernel objects live under build_dir/ (see move_artifacts), so they + are linked from there too, flattened -- the reference in the MLIR carries + no directory. Only the current arch's subdirectory is linked: walking all + of them would put both arches' "mul.o" in one work_dir and reinstate the + collision the per-arch scoping exists to prevent. """ - for entry in build_dir.iterdir(): - if entry.is_dir(): - continue - link = work_dir / entry.name - if link.exists(): - continue - target = entry.resolve() - try: - link.symlink_to(target) - except OSError: - # Windows without Developer Mode cannot create symlinks. - shutil.copy2(target, link) + + def link_files_from(directory: Path) -> None: + if not directory.is_dir(): + return + for entry in directory.iterdir(): + if entry.is_dir(): + continue + link = work_dir / entry.name + if link.exists(): + continue + target = entry.resolve() + try: + link.symlink_to(target) + except OSError: + # Windows without Developer Mode cannot create symlinks. + shutil.copy2(target, link) + + link_files_from(build_dir) + link_files_from(build_dir / get_kernel_dir()) class AieccCompilationRule(CompilationRule): diff --git a/iron/operators/dequant/design.py b/iron/operators/dequant/design.py index ad5cdb59d..e652f918e 100644 --- a/iron/operators/dequant/design.py +++ b/iron/operators/dequant/design.py @@ -8,6 +8,8 @@ from aie.helpers.taplib.tap import TensorAccessPattern from aie.iron.controlflow import range_ +from iron.common.device_utils import get_kernel_dir + def my_dequant_kernel( dev, @@ -62,7 +64,7 @@ def my_dequant_kernel( # AIE Core Function declaration dequant_kernel = Kernel( "expand_int4_to_bfloat16", - f"expand_aie2_{tile_size}.o", + f"expand_{get_kernel_dir(dev)}_{tile_size}.o", [in_tile_ty, out_tile_ty], ) diff --git a/iron/operators/dequant/op.py b/iron/operators/dequant/op.py index bb1ec1602..adc3081b3 100644 --- a/iron/operators/dequant/op.py +++ b/iron/operators/dequant/op.py @@ -14,6 +14,7 @@ PythonGeneratedMLIRArtifact, DesignGenerator, ) +from iron.common.device_utils import get_kernel_dir import aie.utils as aie_utils @@ -64,7 +65,7 @@ def get_mlir_artifact(self): def get_kernel_artifacts(self): return [ KernelObjectArtifact( - f"expand_aie2_{self.tile_size}.o", + f"expand_{get_kernel_dir()}_{self.tile_size}.o", dependencies=[ SourceArtifact( self.context.base_dir / "aie_kernels" / "generic" / "expand.cc" diff --git a/iron/tests/compilation/kernel_object_arch_isolation.py b/iron/tests/compilation/kernel_object_arch_isolation.py new file mode 100644 index 000000000..4c9f56011 --- /dev/null +++ b/iron/tests/compilation/kernel_object_arch_isolation.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A kernel object's on-disk path must be unique per target arch. + +KernelCompilationRule.compile() passes a different --target and +aie_runtime_lib -I per arch for the same output filename (e.g. "mul.o"), and +for aie_kernels/generic/ sources the very same input file compiles to +different machine code per arch. CompilationArtifact.is_available_in_filesystem() +only ever compares mtimes and never records which arch an object was built +for, so if two arches' objects resolve to the same build_dir path, whichever +was compiled last is silently handed to the other arch's link step. This is +what happened in production: an npu1 (aie2) build's mul.o was reused by a +following npu2 (aie2p) run in the same build/, producing bogus ElementwiseMul +failures. These tests build the real artifact graph for both arches and +assert their kernel object paths never collide. +""" + +from pathlib import Path + +import aie.utils as aie_utils +from aie.iron.device import NPU1, NPU2 + +from iron.common import AIEContext +from iron.common.compilation import KernelObjectArtifact +from iron.common.compilation.base import _link_build_outputs_into +from iron.operators.elementwise_mul.op import ElementwiseMul + + +def _mul_kernel_object(build_dir, device): + """Set up ElementwiseMul's artifact graph for `device` and resolve its + kernel object's build_dir path, without invoking Peano/xchesscc.""" + aie_utils.set_current_device(device) + ctx = AIEContext(build_dir=build_dir) + op = ElementwiseMul(size=4096, tile_size=4096, num_aie_columns=1, context=ctx) + op.set_up_artifacts() + op.artifacts.move_artifacts(str(ctx.build_dir)) + op.artifacts.populate_availability_from_filesystem() + for artifact in op.artifacts.bfs(): + if isinstance(artifact, KernelObjectArtifact): + return artifact + raise AssertionError("ElementwiseMul produced no KernelObjectArtifact") + + +def test_two_arches_do_not_resolve_the_same_kernel_object_path(tmp_path): + """aie_kernels/generic/mul.cc is one source shared by aie2 and aie2p + (ElementwiseMul.kernel_subdir); its object must not collide in build_dir.""" + aie2 = _mul_kernel_object(tmp_path, NPU1()) + aie2p = _mul_kernel_object(tmp_path, NPU2()) + assert aie2.filename != aie2p.filename + + +def test_a_stale_object_from_one_arch_is_not_silently_reused_by_another(tmp_path): + """Reproduces the production incident: plant a real leftover aie2 object, + then check the following aie2p build does not report it available.""" + aie2 = _mul_kernel_object(tmp_path, NPU1()) + Path(aie2.filename).parent.mkdir(parents=True, exist_ok=True) + Path(aie2.filename).write_bytes(b"aie2-machine-code") + + aie2p = _mul_kernel_object(tmp_path, NPU2()) + + assert not aie2p.is_available_in_filesystem(), ( + "aie2p build reused a leftover aie2 kernel object -- both resolved " + f"to {aie2p.filename}" + ) + + +def test_the_arch_scoped_object_is_linked_into_the_aiecc_work_dir(tmp_path): + """Scoping the object under build_dir/ must not hide it from the + link step. aiecc resolves link_with="mul.o" against its own work_dir, fed + by _link_build_outputs_into(), which skips directories -- so an object + moved into a subdirectory stops being linked and ld.lld fails with + "cannot open .../mul.o: No such file or directory".""" + obj = _mul_kernel_object(tmp_path, NPU2()) + Path(obj.filename).parent.mkdir(parents=True, exist_ok=True) + Path(obj.filename).write_bytes(b"aie2p-machine-code") + + work_dir = tmp_path / "design.mlir.d" + work_dir.mkdir() + _link_build_outputs_into(work_dir, tmp_path) + + linked = work_dir / Path(obj.filename).name + assert linked.exists(), ( + f"{Path(obj.filename).name} was not linked into the aiecc work dir; " + f"the object is at {obj.filename}" + ) + assert linked.read_bytes() == b"aie2p-machine-code"