Integrate Phoenix (NPU1) support with correct cross-compilation

## What changed

**Phoenix (NPU1/aie2) support:**
- `ChanneledUnaryOperator` now auto-detects the target kernel directory
  (`aie2` vs `aie2p`) at compile time via `get_kernel_dir()`, replacing the
  hardcoded `kernel_subdir` ClassVar
- Adds `needs_lut_ops: ClassVar[bool] = False` to `ChanneledUnaryOperator`;
  subclasses that require `lut_based_ops.o` on aie2 (gelu, sigmoid, tanh,
  silu, softmax) set it to `True`
- `PeanoCompilationRule` derives the compiler target triple and runtime lib
  dir from the device at compile time instead of a hardcoded `DEVICE_CONFIGS`
  dict

**Cross-compilation:**
- All compilation-path `DefaultNPURuntime.device()` calls replaced with
  `aie_utils.get_current_device()`, which respects `set_current_device()`
- Cross-compile for a specific target with
  `aie_utils.set_current_device(NPU1())` before compilation; execution paths
  (`DefaultNPURuntime.load/run`) are unchanged
- Artifact names now include the device architecture (e.g.
  `GELU_..._npu2.xclbin`, `gelu_npu2.o`) so NPU1 and NPU2 artifacts coexist
  safely in the same build directory without cache collisions
- Device detection uses `isinstance(dev, NPU2)` rather than string matching

**Cleanup:**
- Deleted `DEVICE_CONFIGS`, `get_device_name`, `get_device_type` from
  `device_utils.py`; replaced with a single `get_kernel_dir()` helper
- Removed `AIEDeviceManager` imports from all test files (class was deleted
  in this branch)
- Removed dead `get_device_name`/`get_device_type` imports from
  `gemm/design.py` and `mha/design.py`
- Fixed `kernel_archive` NameError in gelu/sigmoid/tanh design files
- Fixed bare `kernel_dir` NameError in `mha/op.py`
- Fixed hardcoded `"aie2p"` in `gemm/op.py` and `mha/op.py` (mm.cc)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: hunhoffe

conftest.py

@@ -12,7 +12,6 @@
 
 from iron.common import AIEContext
 import aie.utils as aie_utils
-from iron.common.aie_device_manager import AIEDeviceManager
 
 
 @pytest.fixture
@@ -158,7 +157,7 @@ def pytest_configure(config):
 
 
 def pytest_collection_modifyitems(config, items):
-    device = AIEDeviceManager().device_str()
+    device = aie_utils.DefaultNPURuntime.device().resolve().name
     for item in items:
         marker = item.get_closest_marker("supported_devices")
         if marker and device not in marker.args:

iron/common/base.py

@@ -128,16 +128,20 @@ def name(self) -> str:
                 for f in dataclasses.fields(self)
                 if f.repr and getattr(self, f.name) is not None
             )
-            return type(self).__name__ + "_" + "_".join(parts)
-        params = self._params
-        if params is not None:
+            base = type(self).__name__ + "_" + "_".join(parts)
+        elif self._params is not None:
             parts = (
-                f"{k}{_serialize_param(v)}" for k, v in params.items() if v is not None
+                f"{k}{_serialize_param(v)}"
+                for k, v in self._params.items()
+                if v is not None
             )
-            return type(self).__name__ + "_" + "_".join(parts)
-        raise NotImplementedError(
-            f"{type(self).__name__} must be a @dataclass or define a _params property"
-        )
+            base = type(self).__name__ + "_" + "_".join(parts)
+        else:
+            raise NotImplementedError(
+                f"{type(self).__name__} must be a @dataclass or define a _params property"
+            )
+        dev = aie_utils.get_current_device()
+        return f"{base}_{dev.resolve().name}"
 
     @abstractmethod
     def get_mlir_artifact(self) -> CompilationArtifact:
@@ -151,8 +155,13 @@ def get_artifacts(
         self, prefix: str = "", dynamic_obj_fifos: bool = False
     ) -> tuple[XclbinArtifact, InstsBinArtifact]:
         operator_name = prefix + self.name
+        arch = self.name.rsplit("_", 1)[-1]
         mlir_artifact = self.get_mlir_artifact()
         kernel_deps = self.get_kernel_artifacts()
+        for dep in kernel_deps:
+            if isinstance(dep, KernelObjectArtifact):
+                p = Path(dep.filename)
+                dep.filename = str(p.with_stem(f"{p.stem}_{arch}"))
         extra_flags = ["--dynamic-objFifos"] if dynamic_obj_fifos else []
         xclbin_artifact = XclbinArtifact(
             f"{operator_name}.xclbin",

iron/common/compilation/base.py

@@ -47,6 +47,9 @@
 from typing import Any, Callable
 import sys
 
+import aie.utils as aie_utils
+from aie.iron.device import NPU2
+
 # Global Functions
 # ##########################################################################
 
@@ -574,10 +577,9 @@ def compile(self, graph):
 
 
 class PeanoCompilationRule(CompilationRule):
-    def __init__(self, peano_dir, mlir_aie_dir, device_type=None, *args, **kwargs):
+    def __init__(self, peano_dir, mlir_aie_dir, *args, **kwargs):
         self.peano_dir = peano_dir
         self.mlir_aie_dir = mlir_aie_dir
-        self.device_type = device_type
         super().__init__(*args, **kwargs)
 
     def matches(self, artifacts):
@@ -589,15 +591,13 @@ def compile(self, artifacts):
         worklist = artifacts.get_worklist(KernelObjectArtifact)
         commands = []
 
-        if self.device_type not in DEVICE_CONFIGS:
-            raise ValueError(
-                f"Unsupported device type: {self.device_type!r} "
-                f"(supported: {', '.join(DEVICE_CONFIGS)})"
-            )
-        config = DEVICE_CONFIGS[self.device_type]
-        target = config["target"]
+        dev = aie_utils.get_current_device()
+        is_npu2 = isinstance(dev, NPU2)
+        target = "aie2p-none-unknown-elf" if is_npu2 else "aie2-none-unknown-elf"
         runtime_lib_include_path = (
-            Path(self.mlir_aie_dir) / "aie_runtime_lib" / config["runtime_lib_dir"]
+            Path(self.mlir_aie_dir)
+            / "aie_runtime_lib"
+            / ("AIE2P" if is_npu2 else "AIE2")
         )
 
         for artifact in worklist:

iron/common/device_utils.py

@@ -1,76 +1,12 @@
 # SPDX-FileCopyrightText: Copyright (C) 2025 Advanced Micro Devices, Inc. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Device utility functions for handling NPU device types and configurations."""
+import aie.utils as aie_utils
+from aie.iron.device import NPU2
 
-from aie.iron.device import NPU1, NPU1Col1, NPU1Col2, NPU2
 
-
-def get_device_name(dev):
-    """Get device name string for looking up microkernel dimensions.
-
-    Returns "npu1" for Phoenix/NPU1 devices, "npu2" for Strix/NPU2 devices.
-
-    Args:
-        dev: Either a string ("npu", "npu1", "npu2"), a device object with .resolve(),
-             or a device type object.
-
-    Returns:
-        str: "npu1" or "npu2"
-    """
-    if isinstance(dev, str):
-        if dev in ("npu", "npu1"):
-            return "npu1"
-        else:
-            return "npu2"
-    elif hasattr(dev, "resolve"):
-        # Device object from device_manager.device_type
-        return dev.resolve().name
-    else:
-        # Assume it's a device type object - check class name
-        name = type(dev).__name__
-        if "NPU1" in name:
-            return "npu1"
-        else:
-            return "npu2"
-
-
-DEVICE_CONFIGS = {
-    "npu1": {
-        "target": "aie2-none-unknown-elf",
-        "runtime_lib_dir": "AIE2",
-        "kernel_dir": "aie2",
-        "max_columns": NPU1().cols,
-    },
-    "npu2": {
-        "target": "aie2p-none-unknown-elf",
-        "runtime_lib_dir": "AIE2P",
-        "kernel_dir": "aie2p",
-        "max_columns": NPU2().cols,
-    },
-}
-
-
-def get_device_type(dev, n_aie_cols):
-    """Resolve device type to appropriate NPU device instance.
-
-    Handles both string inputs ("npu", "npu1", "npu2") and device objects.
-
-    Args:
-        dev: Either a string ("npu", "npu1", "npu2"), a device object with .resolve(),
-             or a device type object.
-        n_aie_cols: Number of AIE columns to use (1, 2, or 4).
-
-    Returns:
-        NPU device instance (NPU1, NPU1Col1, NPU1Col2, or NPU2)
-    """
-    dev_name = get_device_name(dev)
-    if dev_name == "npu1":
-        if n_aie_cols == 1:
-            return NPU1Col1()
-        elif n_aie_cols == 2:
-            return NPU1Col2()
-        else:
-            return NPU1()
-    else:
-        return NPU2()
+def get_kernel_dir(dev=None) -> str:
+    """Returns 'aie2p' for NPU2 (Strix), 'aie2' for NPU1 (Phoenix)."""
+    if dev is None:
+        dev = aie_utils.get_current_device()
+    return "aie2p" if isinstance(dev, NPU2) else "aie2"

iron/common/operator_bases.py

@@ -4,16 +4,19 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import Any, ClassVar, Dict
 
 import aie.utils as aie_utils
+
 from .base import MLIROperator, AIERuntimeArgSpec
 from .compilation import (
     KernelObjectArtifact,
     SourceArtifact,
     PythonGeneratedMLIRArtifact,
     DesignGenerator,
 )
+from .device_utils import get_kernel_dir
 from .utils import get_shim_dma_limit
 
 
@@ -24,10 +27,10 @@ class ChanneledUnaryOperator(MLIROperator):
     Assumes a single kernel source file and a standard design.py callback
     with args [device, size, num_aie_columns, num_channels, tile_size, trace_size].
 
-    Subclasses must define three ClassVar attributes:
+    Subclasses must define ClassVar attributes:
         kernel_name:   name of the kernel object file (e.g. "gelu" → gelu.o / gelu.cc)
-        kernel_subdir: subdirectory under aie_kernels/ (e.g. "aie2p", "generic")
         callback_fn:   design.py callback function name (e.g. "my_gelu")
+        needs_lut_ops: set True for operators that require lut_based_ops.o on aie2
 
     Customization points:
         - For operators with extra parameters (e.g. alpha, trace_size), add
@@ -45,8 +48,8 @@ class ChanneledUnaryOperator(MLIROperator):
     context: object = field(default=None, repr=False)
 
     kernel_name: ClassVar[str]
-    kernel_subdir: ClassVar[str]
     callback_fn: ClassVar[str]
+    needs_lut_ops: ClassVar[bool] = False
 
     def __post_init__(self) -> None:
         max_multiple = self.num_aie_columns * self.tile_size
@@ -55,7 +58,7 @@ def __post_init__(self) -> None:
                 f"size ({self.size}) must be a multiple of "
                 f"num_aie_columns * tile_size ({max_multiple})"
             )
-        dev = aie_utils.DefaultNPURuntime.device()
+        dev = aie_utils.get_current_device()
         shim_dma_limit = get_shim_dma_limit(dev)
         total_shimdma_channels = self.num_aie_columns * self.num_channels
         if total_shimdma_channels > shim_dma_limit:
@@ -78,7 +81,7 @@ def _mlir_callback_args(self) -> list[Any]:
         override this method.
         """
         return [
-            aie_utils.DefaultNPURuntime.device(),
+            aie_utils.get_current_device(),
             self.size,
             self.num_aie_columns,
             self.num_channels,
@@ -97,19 +100,37 @@ def get_mlir_artifact(self) -> PythonGeneratedMLIRArtifact:
         )
 
     def get_kernel_artifacts(self) -> list[KernelObjectArtifact]:
-        return [
+        dev = aie_utils.get_current_device()
+        kernel_dir = get_kernel_dir(dev)
+        artifacts = [
             KernelObjectArtifact(
                 f"{self.kernel_name}.o",
                 dependencies=[
                     SourceArtifact(
                         self.context.base_dir
                         / "aie_kernels"
-                        / self.kernel_subdir
+                        / kernel_dir
                         / f"{self.kernel_name}.cc"
                     )
                 ],
-            ),
+            )
         ]
+        if self.needs_lut_ops and kernel_dir == "aie2":
+            mlir_aie_dir = Path(aie_utils.config.root_path())
+            artifacts.append(
+                KernelObjectArtifact(
+                    "lut_based_ops.o",
+                    dependencies=[
+                        SourceArtifact(
+                            mlir_aie_dir
+                            / "aie_runtime_lib"
+                            / "AIE2"
+                            / "lut_based_ops.cpp"
+                        )
+                    ],
+                )
+            )
+        return artifacts
 
 
 @dataclass
@@ -123,7 +144,7 @@ class BinaryElementwiseOperator(MLIROperator):
     parameter — each core uses 2 DMA channels (one per input), so the ShimDMA
     limit is enforced as num_aie_columns * 2 <= 16.
 
-    Subclasses must define three ClassVar attributes:
+    Subclasses must define ClassVar attributes:
         kernel_name:   name of the kernel object file (e.g. "add" → add.o / add.cc)
         kernel_subdir: subdirectory under aie_kernels/ (e.g. "generic")
         callback_fn:   design.py callback function name (e.g. "my_eltwise_add")
@@ -148,7 +169,7 @@ def __post_init__(self) -> None:
                 f"size ({self.size}) must be a multiple of "
                 f"num_aie_columns * tile_size ({self.num_aie_columns * self.tile_size})"
             )
-        dev = aie_utils.DefaultNPURuntime.device()
+        dev = aie_utils.get_current_device()
         shim_dma_limit = get_shim_dma_limit(dev)
         # Binary operators use 2 ShimDMA channels per column (one per input).
         total_shimdma_channels = self.num_aie_columns * 2
@@ -168,7 +189,7 @@ def get_arg_spec(self) -> list[AIERuntimeArgSpec]:
 
     def _mlir_callback_args(self) -> list[Any]:
         return [
-            aie_utils.DefaultNPURuntime.device(),
+            aie_utils.get_current_device(),
             self.size,
             self.num_aie_columns,
             self.tile_size,

iron/operators/axpy/op.py

@@ -39,7 +39,7 @@ def get_mlir_artifact(self):
                 self.operator_dir / "design.py",
                 "my_axpy",
                 (
-                    aie_utils.DefaultNPURuntime.device(),
+                    aie_utils.get_current_device(),
                     self.size,
                     self.num_aie_columns,
                     self.tile_size,

iron/operators/axpy/test.py

@@ -3,17 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
+import aie.utils as aie_utils
 
 from iron.operators.axpy.op import AXPY
 from iron.operators.axpy.reference import generate_golden_reference
 from iron.common.test_utils import run_test
-from iron.common.aie_device_manager import AIEDeviceManager
-from iron.common.device_utils import DEVICE_CONFIGS
 
 
 def get_params():
-    device_type = AIEDeviceManager().device_str()
-    max_aie_columns = DEVICE_CONFIGS[device_type]["max_columns"]
+    max_aie_columns = aie_utils.get_current_device().cols
     input_lengths = [1024, 2048, 4096, 8192]
     scalar_factors = [3.0, 10.0]
 

iron/operators/dequant/op.py

@@ -50,7 +50,7 @@ def get_mlir_artifact(self):
                 self.operator_dir / "design.py",
                 "my_dequant_kernel",
                 (
-                    aie_utils.DefaultNPURuntime.device(),
+                    aie_utils.get_current_device(),
                     self.size,
                     self.num_aie_columns,
                     self.num_channels,

iron/operators/dequant/test.py

@@ -3,17 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
+import aie.utils as aie_utils
 
 from iron.operators.dequant.op import Dequant
 from iron.operators.dequant.reference import generate_golden_reference
 from iron.common.test_utils import run_test
-from iron.common.aie_device_manager import AIEDeviceManager
-from iron.common.device_utils import DEVICE_CONFIGS
 
 
 def get_params():
-    device_type = AIEDeviceManager().device_str()
-    max_aie_columns = DEVICE_CONFIGS[device_type]["max_columns"]
+    max_aie_columns = aie_utils.get_current_device().cols
 
     input_lengths = [1024, 2048, 4096, 8192]
     group_size = 32

iron/operators/elementwise_add/test.py

@@ -3,17 +3,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import pytest
+import aie.utils as aie_utils
 
-from iron.common.aie_device_manager import AIEDeviceManager
-from iron.common.device_utils import DEVICE_CONFIGS
 from iron.operators.elementwise_add.op import ElementwiseAdd
 from iron.operators.elementwise_add.reference import generate_golden_reference
 from iron.common.test_utils import run_test
 
 
 def get_params():
-    device_type = AIEDeviceManager().device_str()
-    max_aie_columns = DEVICE_CONFIGS[device_type]["max_columns"]
+    max_aie_columns = aie_utils.get_current_device().cols
     # Combine all lengths
     input_lengths = [1024, 2048, 4096, 8192]