From 46bf7005057cef06c99e873d60186230ba8e21f2 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 21 May 2026 23:57:21 -0600 Subject: [PATCH] Expose target architecture via Device.arch (#2485) Wraps AIETargetModel::getTargetArch in the C API (aieTargetModelGetTargetArch), binds it on PyAieTargetModel as get_target_arch(), and adds an arch property on iron.Device that returns the AIEArch enum. Replaces the isinstance/AIEDevice cascade in utils/jit.py with a single device.arch lookup. --- include/aie-c/TargetModel.h | 5 +++++ lib/CAPI/TargetModel.cpp | 4 ++++ python/AIEMLIRModule.cpp | 4 ++++ python/iron/device/device.py | 7 ++++++- python/utils/jit.py | 21 +++++++++++---------- 5 files changed, 30 insertions(+), 11 deletions(-) diff --git a/include/aie-c/TargetModel.h b/include/aie-c/TargetModel.h index 837dd1df0f5..6cb29ffc966 100644 --- a/include/aie-c/TargetModel.h +++ b/include/aie-c/TargetModel.h @@ -55,6 +55,11 @@ MLIR_CAPI_EXPORTED int aieTargetModelRows(AieTargetModel targetModel); /// Returns true if this is an NPU target model. MLIR_CAPI_EXPORTED bool aieTargetModelIsNPU(AieTargetModel targetModel); +/// Returns the AIE architecture (as the underlying value of xilinx::AIE::AIEArch: +/// AIE1=1, AIE2=2, AIE2p=3). +MLIR_CAPI_EXPORTED uint32_t +aieTargetModelGetTargetArch(AieTargetModel targetModel); + /// Returns the tile type for the given coordinates. MLIR_CAPI_EXPORTED uint32_t aieTargetModelGetTileType(AieTargetModel targetModel, int col, int row); diff --git a/lib/CAPI/TargetModel.cpp b/lib/CAPI/TargetModel.cpp index 499fd8c1b02..87b6b8b4302 100644 --- a/lib/CAPI/TargetModel.cpp +++ b/lib/CAPI/TargetModel.cpp @@ -201,6 +201,10 @@ bool aieTargetModelIsNPU(AieTargetModel targetModel) { return unwrap(targetModel).hasProperty(xilinx::AIE::AIETargetModel::IsNPU); } +uint32_t aieTargetModelGetTargetArch(AieTargetModel targetModel) { + return static_cast(unwrap(targetModel).getTargetArch()); +} + uint32_t aieTargetModelGetColumnShift(AieTargetModel targetModel) { return unwrap(targetModel).getColumnShift(); } diff --git a/python/AIEMLIRModule.cpp b/python/AIEMLIRModule.cpp index 2f1050eef11..cf9b46e018c 100644 --- a/python/AIEMLIRModule.cpp +++ b/python/AIEMLIRModule.cpp @@ -358,6 +358,10 @@ NB_MODULE(_aie, m) { [](PyAieTargetModel &self) { return aieTargetModelIsNPU(self.get()); }) + .def("get_target_arch", + [](PyAieTargetModel &self) { + return aieTargetModelGetTargetArch(self.get()); + }) .def("get_column_shift", [](PyAieTargetModel &self) { return aieTargetModelGetColumnShift(self.get()); diff --git a/python/iron/device/device.py b/python/iron/device/device.py index 6b1035949c5..be9bd6e3a38 100644 --- a/python/iron/device/device.py +++ b/python/iron/device/device.py @@ -9,7 +9,7 @@ from typing import Generator from ... import ir # type: ignore -from ...dialects._aie_enum_gen import AIETileType, WireBundle # type: ignore +from ...dialects._aie_enum_gen import AIEArch, AIETileType, WireBundle # type: ignore from ...dialects.aie import ( AIEDevice, logical_tile, @@ -45,6 +45,11 @@ def rows(self) -> int: """Number of rows in the device tile array.""" return self._tm.rows() + @property + def arch(self) -> AIEArch: + """AIE architecture of the device (AIE1, AIE2, or AIE2p).""" + return AIEArch(self._tm.get_target_arch()) + def _validate_coordinates(self, col, row): """Raise ValueError if coordinates are outside the device grid.""" if col < 0 or col >= self._tm.columns() or row < 0 or row >= self._tm.rows(): diff --git a/python/utils/jit.py b/python/utils/jit.py index 4ee11094b63..73abdcd7e41 100644 --- a/python/utils/jit.py +++ b/python/utils/jit.py @@ -15,7 +15,6 @@ from aie.extras.context import mlir_mod_ctx from .compile import compile_mlir_module, compile_external_kernel from .npukernel import NPUKernel -from aie.dialects.aie import AIEDevice from .compile.cache.circular_cache import CircularCache from .compile.cache.utils import _create_function_cache_key, file_lock from .compile import NPU_CACHE_HOME @@ -48,8 +47,10 @@ def jit(function=None, use_cache=True): @functools.wraps(function) def decorator(*args, **kwargs): - from aie.iron.device import NPU1, NPU2, NPU1Col1, NPU2Col1 + from aie.iron.device import Device from aie.iron.kernel import ExternalFunction + from aie.dialects._aie_enum_gen import AIEArch + from aie.dialects.aie import get_target_model from . import DefaultNPURuntime if DefaultNPURuntime is None: @@ -118,17 +119,17 @@ def decorator(*args, **kwargs): current_device = DefaultNPURuntime.device() - # Determine target architecture based on device type - if isinstance(current_device, (NPU2, NPU2Col1)): - target_arch = "aie2p" - elif isinstance(current_device, (NPU1, NPU1Col1)): - target_arch = "aie2" - elif current_device in (AIEDevice.npu2, AIEDevice.npu2_1col): + # Determine target architecture from the device's target model. + if isinstance(current_device, Device): + arch = current_device.arch + else: + arch = AIEArch(get_target_model(current_device).get_target_arch()) + if arch == AIEArch.AIE2p: target_arch = "aie2p" - elif current_device in (AIEDevice.npu1, AIEDevice.npu1_1col): + elif arch == AIEArch.AIE2: target_arch = "aie2" else: - raise RuntimeError(f"Unsupported device type: {type(current_device)}") + raise RuntimeError(f"Unsupported device arch: {arch}") # Hash of the IR string, ExternalFunction compiler options, and target architecture module_hash = hash_module(mlir_module, external_kernels, target_arch)