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

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/<arch> (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):

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],
     )
 

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"

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/<arch> 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"