IRON Python Configurations¶
There are several options that exist to configure the IRON Python programming environment.
Default IRON Tensor Class¶
This is a variable that controls the types of aie.utils.NpuTensors that are produced by the utility functions tensor, ones, etc. The available tensor implementations are CPUOnlyTensor, XRTTensor, HRXTensor and HSATensor; the last two are selected via NPU_RUNTIME.
By default, if pyxrt is available, the DEFAULT_TENSOR_CLASS is set to XRTTensor. HRXTensor is never selected automatically: it is chosen only by setting NPU_RUNTIME=hrx explicitly (see HRX runtime). You can also set the class directly through set_tensor_class(), e.g.:
>>> import numpy as np
>>> print(aie.utils.tensor.DEFAULT_TENSOR_CLASS.__name__)
XRTTensor
>>> type(iron.tensor((2, 2), np.int32))
<class 'aie.utils.xrtruntime.tensor.XRTTensor'>
>>> aie.utils.set_tensor_class(aie.utils.tensor.CPUOnlyTensor)
>>> print(aie.utils.tensor.DEFAULT_TENSOR_CLASS.__name__)
CPUOnlyTensor
>>> type(aie.utils.tensor((2, 2), np.int32))
<class 'aie.utils.tensor.CPUOnlyTensor'>
Default IRON Device¶
If the IRON device is not set, many designs fetch it on demand from the DefaultNPURuntime (a CachedXRTRuntime instance), which queries XRT for the attached NPU. Select an explicit target with iron.set_current_device():
>>> import aie.iron as iron
>>> iron.set_current_device(iron.device.NPU1())
>>> iron.get_current_device()
<abc.NPU1 object at 0x722a65903a10>
Cross-compiling for a different NPU¶
The target architecture (aie2 / aie2p) is part of the per-design
cache key (see compilation_stages.md
§Lowering — _compute_artifact_hash mixes target_arch in along with
peano + aiecc mtimes), so set_current_device(...) before .compile()
is enough to drive a cross-arch build. Each arch lands in its own
cache subdirectory, no collision with the binary for whatever NPU is
physically attached:
import aie.iron as iron
from aie.iron.device import NPU1Col1, NPU2Col1
# Same generator, two arches → two distinct cache dirs.
for dev_cls in (NPU1Col1, NPU2Col1):
iron.set_current_device(dev_cls())
my_design.specialize(N=4096).compile()
Useful for building Strix binaries on Phoenix-only CI hosts (or vice versa), and for shipping pre-built xclbins for multiple NPU generations without needing each one attached at build time.
Arch-aware kernel introspection (mac_dims)¶
Some aie.iron.kernels factories pick a different MMUL geometry per
arch — kernels.mm(int16, int16) is (r, s, t) = (4, 4, 4) on AIE2
(Phoenix) but (4, 4, 8) on AIE2P (Strix). The chosen geometry is
exposed on the returned ExternalFunction as .mac_dims, so designs
can drive their DMA-layout transforms from the kernel itself instead
of hardcoding for one arch:
import aie.iron as iron
from aie.iron.device import NPU2Col1
iron.set_current_device(NPU2Col1())
mm = kernels.mm(dim_m=64, dim_k=64, dim_n=64,
input_dtype=np.int16, output_dtype=np.int16)
r, s, t = mm.mac_dims
Set the target before constructing an arch-sensitive factory. The factory selects its source and MMUL geometry when it is created.
The full per-arch table lives in
python/iron/kernels/linalg.py
(_MM_MAC_DIMS). Combined with the cross-compile pattern above, the
same generator file produces a correct binary for each arch without
ever editing the source.
IRON Cache Location¶
The IRON jit feature caches compiled objects in a directory defined by NPU_CACHE_HOME. By default this value is the user's home directory.
aiecc Executable Override (AIECC_PATH)¶
aie.utils.config.aiecc_path() resolves the aiecc executable used by JIT
compilation: the AIECC_PATH environment variable first (an explicit full
path to the aiecc executable), then the MLIR-AIE bin directory, then
PATH. Set AIECC_PATH when a consumer needs to pin a specific aiecc
without relying on PATH search order.
IRON XRT Runtime Cache Size¶
The CachedXRTRuntime caches XRT contexts to improve performance. The size of this cache can be configured using the XRT_CONTEXT_CACHE_SIZE environment variable. This is particularly useful in CI environments where multiple tests run in parallel and might exhaust the available NPU contexts.
Host-runtime backend selection (NPU_RUNTIME)¶
IRON dispatches designs through a host runtime that consumes the aiecc
artifacts. XRT and HRX load final.xclbin + insts.bin; HSA loads the PDI
(main.pdi) + insts.bin from the same build. NPU_RUNTIME selects which
backend is used:
| Value | Behavior |
|---|---|
auto (default) |
Use XRT when pyxrt imports, else fall back to CPU-only tensors. Never selects HRX or HSA. |
xrt |
Force the XRT backend (falls back to CPU tensors if pyxrt is missing). |
hrx |
Force the HRX (amdxdna / libhrx) backend. Errors at import if libhrx.so cannot be located. |
hsa |
Force the HSA/ROCR backend. Errors at import if libhsa-runtime64.so cannot be located. |
An unset value defaults to auto; an explicitly invalid value is a hard error
(a typo must not silently resolve to another backend). NPU_RUNTIME is read
before any capability probe, so a forced backend only probes itself (hrx
and hsa never import pyxrt, and xrt/auto never run HRX or HSA
discovery).
NPU_RUNTIME is a single selector shared with the C++ example make flow: the
same NPU_RUNTIME=hrx also builds the HRX host stack (-DUSE_HRX=ON) for
make-driven examples (see programming_examples/makefile-common). The auto
and cpu values are meaningful only to the Python flow; the build system treats
an unset selector as xrt. hsa is a Python-flow backend only; there is no
C++ host stack for it, so make-driven examples ignore it.
HRX Runtime (amdxdna) Configuration¶
These variables apply when the HRX backend is active (NPU_RUNTIME=hrx). See
the HRX Runtime (amdxdna) overview for the architecture and
enabling instructions, and the
HRX runtime README for the
full step-by-step flow.
Library discovery¶
Python locates libhrx in this order (filesystem only — no dlopen). Explicit
HRX_* hints match FindHRX.cmake; CMake does not search site-packages.
| Source | Meaning |
|---|---|
HRX_LIBHRX |
Explicit full path to libhrx.so / hrx.dll. |
LIBHRX_DIR |
Directory containing the library (e.g. set by activate_env.sh). |
| pip site-packages | A package that ships lib/libhrx.so* or bin/hrx.dll. Filesystem only — the package is not imported. |
HRX_DIR |
HRX install prefix ($HRX_DIR/lib/libhrx.so, $HRX_DIR/include/...). |
sibling / FindHRX roots |
A sibling hrx checkout, $HOME/hrx, /opt/hrx, /usr/local/hrx, then the loader's search path. |
Runtime behavior¶
| Variable | Default | Meaning |
|---|---|---|
IRON_HRX_DEVICE |
auto-detect | Force the amdxdna device generation (npu1 / npu2) instead of detecting it from sysfs PCI IDs. |
HRX_EXE_CACHE_SIZE |
32 |
Max number of amdxdna executables the CachedHRXRuntime keeps (LRU). |
IRON_HRX_TIMEOUT |
0 (disabled) |
Watchdog timeout, in seconds, bounding the wait in hrx_stream_synchronize. 0, unset, or an invalid value disables the watchdog. On expiry a diagnosable error is raised (the underlying sync cannot be cancelled). |
# Force npu2, cap the executable cache, and fail a wedged sync after 30s.
NPU_RUNTIME=hrx IRON_HRX_DEVICE=npu2 HRX_EXE_CACHE_SIZE=8 \
IRON_HRX_TIMEOUT=30 python my_script.py
HSA/ROCR Runtime Configuration¶
These variables apply when the HSA backend is active (NPU_RUNTIME=hsa). See
HSA Runtime (ROCR) for the architecture and the full
step-by-step flow.
Getting ROCm¶
The recommended source is the pip wheel from TheRock, installed into the environment you run designs from:
The base rocm package is enough: it pulls rocm-sdk-core, which carries
libhsa-runtime64.so with the AIE support this backend needs. The
[libraries] and [device-gfx…] extras are for GPU workloads and are not
required for the NPU. Discovery finds it in site-packages with nothing to
configure.
ROCm discovery¶
libhsa-runtime64.so is located inside a ROCm installation root, looked for
in three places, in this order:
| Source | Meaning |
|---|---|
ROCM_PATH |
An explicit installation root, overriding the wheel. It names a ROCm tree, not a library file, so the same variable serves every ROCm component. |
| pip-installed ROCm | The wheel above, whose runtime tree ships inside site-packages. |
| System install | /opt/rocm. |
The first root that actually contains the library wins. Both the unversioned
libhsa-runtime64.so and the bare SONAME libhsa-runtime64.so.1 are accepted:
TheRock's runtime wheels contain no symlinks, so a pip-installed ROCm provides
only the versioned name.
With no wheel and no
ROCM_PATH, a system/opt/rocmis what you get. If it predates AIE support, the failure surfaces later as an opaque HSA error rather than a version complaint. Installing the wheel avoids this; setROCM_PATHonly when you need a specific ROCm instead.
Runtime behavior¶
| Variable | Default | Meaning |
|---|---|---|
IRON_HSA_DEVICE |
auto-detect | Force the device generation (npu1 / npu2) instead of detecting it from the HSA agent name. ROCR names the agent after its ISA (aie2 = Phoenix/npu1, aie2p = Strix/npu2); an unrecognized name is an error rather than a guess, because a wrong generation compiles for the wrong architecture and wedges the NPU. |
HSA_EXE_CACHE_SIZE |
32 |
Max number of loaded designs the CachedHSAHostRuntime keeps (LRU). A non-integer value is ignored, with a warning, in favour of the default. |
IRON_HSA_TIMEOUT |
0 (disabled) |
Timeout, in seconds, bounding the two waits IRON itself performs: the completion-signal wait and the full-queue wait in enqueue. 0, unset, or an invalid value disables it. Implemented with hsa_signal_wait's own timeout, so arming it costs nothing per dispatch. On expiry a diagnosable error is raised (the underlying dispatch cannot be cancelled). This is not a watchdog for a wedged dispatch: ROCR's AIE doorbell submits and blocks until completion, so a hung dispatch stalls inside the doorbell store — before either bounded wait is reached — on an internal wait this setting cannot reach. |
# Point at a specific ROCm, force npu2, cap the cache, bound IRON's waits at 30s.
NPU_RUNTIME=hsa ROCM_PATH=/opt/rocm-7.0 IRON_HSA_DEVICE=npu2 \
HSA_EXE_CACHE_SIZE=8 IRON_HSA_TIMEOUT=30 python my_script.py
Limitations¶
- Trace capture is not supported. A design with a
trace_configis rejected up front; useNPU_RUNTIME=xrtfor trace-enabled designs. - Dispatches must be serialized. The single in-order AIE queue and its doorbell are not safe for concurrent dispatch from multiple threads (the same constraint HRX documents).
Diagnostic Output and Log Level¶
The aie library uses Python's standard logging module for all diagnostic output. Set
AIE_LOG_LEVEL to control verbosity. Valid values: DEBUG, INFO, WARNING
(default), ERROR, CRITICAL.
AIE_LOG_LEVEL=DEBUG python my_script.py # show debug messages
AIE_LOG_LEVEL=INFO python my_script.py # show info and above
AIE_LOG_LEVEL=ERROR python my_script.py # errors only
For per-module control or routing to a file, use the logging API directly:
import logging
logging.getLogger("aie").setLevel(logging.ERROR)
# Route aie logs to a file instead of the console
handler = logging.FileHandler("aie.log")
handler.setFormatter(logging.Formatter("%(asctime)s %(name)s %(levelname)s %(message)s"))
logging.getLogger("aie").addHandler(handler)
logging.getLogger("aie").propagate = False # don't also send to root logger
Helper utilities cheat-sheet¶
The helpers most designs reach for, grouped by where they live. Each
has a docstring; help(obj) or print(obj.__doc__) shows it.
Host-side tensor + design construction (aie.iron)¶
| Helper | What it does |
|---|---|
iron.{tensor, zeros, ones, full, rand, randint, arange, zeros_like} |
Host-tensor factories (default device="npu", also accept device="cpu"); see also set_tensor_class. |
iron.{In, Out, InOut, CompileTime} |
Type annotation markers for @iron.jit generator parameters (see compilation_stages.md §Appendix A). |
iron.{jit, CompilableDesign, CallableDesign} |
The JIT decorator + the two design wrapper classes (CompilableDesign is the recipe; CallableDesign is the ready-to-run wrapper). |
iron.ceildiv(a, b) |
Pure-integer ceiling division. Same value on every code path; here so designs don't redefine it locally. |
iron.{set_current_device, get_current_device} |
Read/write the active Device (see §Default IRON Device). |
iron.kernels.* |
Pre-packaged kernel factories — mm, conv2dk1, conv2dk3, passthrough, eltwise, etc. Each returns an ExternalFunction ready to bind in a Worker. |
iron.{Buffer, Lock, Flow, TileDma, DmaChannel, Bd, Acquire, Release} |
IRON-Python peers of ObjectFifo for designs that want to hand-wire DMA + sync (canonical example: programming_examples/basic/chaining_channels/). |
iron.algorithms.{transform, transform_binary, transform_parallel, for_each} |
Element-wise dataflow templates — handle Worker / ObjectFifo / Runtime plumbing for one-arg / two-arg / multi-column / fill-and-drain patterns. |
iron.{compile_context, get_compile_arg} |
Dynamic compile-time arg injection. See §compile_context below. |
Argparse + runtime glue (aie.iron.device, aie.utils)¶
| Helper | What it does |
|---|---|
aie.utils.hostruntime.argparse.device_from_args(args) |
Resolve an explicit parsed args.dev to a Device — collapses from_name(args.dev, n_cols=...) boilerplate. Returns None when target selection is automatic, so run_design_cli() can bind the attached runtime family. n_cols="auto" reads args.n_cols if present, otherwise defaults to 1. |
aie.utils.DefaultNPURuntime |
Module-level CachedXRTRuntime instance; auto-detects NPU1 / NPU2 via XRT. Used by iron.tensor(..., device="npu") and @iron.jit runtime binding. |
aie.utils.hostruntime.argparse.{add_compile_args, add_runtime_args} |
Add the standard --xclbin-path/--insts-path and --xclbin/--instr/-k/--trace_size flags to a parser. |
aie.utils.test.create_npu_kernel(opts) |
Build an NPUKernel (plus optional TraceConfig) from a parsed argparse.Namespace. See programming_examples/basic/vector_scalar_mul/test.py for the canonical use pattern. |
aie.utils.benchmark.{run_iters, print_benchmark} |
Warmup + timed iterations + summary stats. See section-4/section-4a/README.md. |
aie.utils.verify.{nearly_equal, count_mismatches} |
Tolerance-aware compare for LUT / saturating / bf16 outputs that don't match exactly. |
Compile-pipeline introspection (aie.utils.compile)¶
| Helper | What it does |
|---|---|
aie.utils.compile.NPU_CACHE_HOME |
Cache root Path; defaults to ~/.npu/cache, override with the NPU_CACHE_HOME env var. |
aie.utils.compile.jit._dma_size_parser.parse_dma_sizes(kernel_dir) |
Per-host-arg element counts read from the entry-point aie.runtime_sequence in input_with_addresses.mlir. Backs the tensor-arg validation in stage 5. |
logging.getLogger("aie.utils.compile").setLevel(logging.DEBUG) |
Print every clang++ / aiecc subprocess invocation. See compilation_stages.md §Watching the compile pipeline. |
compile_context for nested generator helpers¶
Most designs supply compile-time values via the explicit
CompileTime[T]-annotated generator signature. Some patterns — composite
generators, helper functions reused across designs — want to inject
values through a helper that doesn't take them as explicit kwargs.
compile_context opens a per-thread context that get_compile_arg
can read out:
from aie.iron import compile_context, get_compile_arg
def make_fifo_pair(line_ty):
# Helper has no name_prefix kwarg, but can read one from context.
prefix = get_compile_arg("prefix", default="")
return (ObjectFifo(line_ty, name=f"{prefix}in"),
ObjectFifo(line_ty, name=f"{prefix}out"))
with compile_context(prefix="layer1_"):
of_in, of_out = make_fifo_pair(line_ty) # layer1_in, layer1_out
of_in, of_out = make_fifo_pair(line_ty) # in, out (default)
Contexts nest — inner values shadow outer ones for the duration of the
inner with block. Implemented on top of contextvars, so it is
thread- and async-safe. CompilableDesign.compile() uses the same
mechanism internally to surface the bound CompileTime[T] kwargs to the
generator body.
Prefer explicit CompileTime[T] parameters when you can; reserve
compile_context for the cases where threading the value through every
helper signature would obscure the design.