IRON Programming Guide

This is the programming guide for iron, the operator and graph layer on top of mlir-aie's aie.iron. It teaches how to declare an operator, record a graph of operators, compile it, and run it on an NPU.

New to the NPU itself? This guide assumes you already know what a Worker, an ObjectFifo and @iron.jit are. If you don't, start with mlir-aie's programming guide first; iron's operators are built out of exactly those pieces, one level up.

The mental model in 60 seconds

• An operator (iron.operators.GEMM, ReLU, MHA, ...) is one declared class: its shape and tunables are param()/tunable() fields, its compute is a build() over aie.iron. Called on its own, op(*tensors) compiles once and runs on the device.
• A graph (a subclass of iron.Graph) records a model's operators, once per input shape, into a body() method. Its positional parameters are the graph's inputs; its return is its outputs. Inside body(), Op.step(...) records a step and returns a Handle in place of a tensor.
• graph.compile(**shapes) lowers one recorded shape to an image on the device: a full ELF on NPU2, an xclbin per step on NPU1. graph(*tensors) finds and calls the matching compiled version. It can compile an unseen shape automatically only while no version is loaded; compile every shape you'll need before the first device call.
• graph.reference(*tensors) runs the same body() on the host, through each operator's numpy reference(), so a design can be checked without touching the NPU.

import numpy as np
from ml_dtypes import bfloat16

import iron
from iron import Scratchpad
from iron.operators import GEMM, SiLU, ElementwiseMul


class FeedForward(iron.Graph):
    """(SiLU(x @ w_gate.T) * (x @ w_up.T)) @ w_down.T, the first n rows of x."""

    def __init__(self, w_gate, w_up, w_down):
        self.w_gate, self.w_up, self.w_down = map(iron.param, (w_gate, w_up, w_down))

    def body(self, x, *, n: Scratchpad[np.int32]):
        gate = GEMM.step(x[:n], self.w_gate, b_col_maj=True)
        up = GEMM.step(x[:n], self.w_up, b_col_maj=True)
        return GEMM.step(SiLU.step(gate) * up, self.w_down, b_col_maj=True)


ffn = FeedForward(w_gate, w_up, w_down)
ffn.compile(x=(256, hidden))     # one compile serves any n up to 256 rows
y = ffn(x, n=37)                  # the device runs the first 37 rows

Where to go next

You want...	Read
The full guide, top to bottom	Section 0 onward
An operator by itself, no graph	Section 1
Your first graph: body, handles, .step, compile, calling it	Section 2
One compile that serves any input extent	Section 3
Weights a graph holds, and state it keeps on the device	Section 4
A device loop that never comes back to the host	Section 5
Checking a design without an NPU	Section 6
Tunables, and a profile that sets them for a whole model	Section 7
From graphs to a language model (iron.lm)	Section 8
The rules a contributor or agent follows	AGENTS.md

Sections

• Section 0 — Getting started
• Section 1 — An operator on its own
• Section 2 — Your first graph
• Section 3 — Per-call values and bounds
• Section 4 — Weights and state
• Section 5 — Carried values and device loops
• Section 6 — The host reference and check mode
• Section 7 — Profiles and tunables
• Section 8 — From graphs to a language model

---

Section 0 — Getting started →