🦾 - IRON：释放 NPU 的全部潜力 - 🦾

IRON Logo

IRON 是一个开源且贴近底层的 Python API，可在 AMD Ryzen™ AI NPU 上实现快速且高效的执行。它依赖于围绕 MLIR-AIE 方言的语言绑定。

主要特性：

通过 MLIR-AIE Python 绑定实现贴近底层的 NPU 编程
预构建算子库（GEMM、MHA、RMSNorm、RoPE、激活函数等）
算子融合以实现最佳性能
可扩展架构以支持自定义算子
端到端 LLM 推理（包含 Llama 3.2 1B 示例）

Ryzen™ AI NPU 的 IRON Python API 在以下论文中进行了描述：

E. Hunhoff, J. Melber, K. Denolf, A. Bisca, S. Bayliss, S. Neuendorffer, J. Fifield, J. Lo, P. Vasireddy, P. James-Roxby, E. Keller. "Efficiency, Expressivity, and Extensibility in a Close-to-Metal NPU Programming Interface". In 33rd IEEE International Symposium On Field-Programmable Custom Computing Machines, May 2025.

🎯 算子仪表板

章节	描述	数据类型	AIE2	AIE2P	状态	设计示例
逐元素加法	逐元素加法内核	bfloat16	✓	✓	🟢	iron/operators/elementwise_add/
逐元素乘法	逐元素乘法内核	bfloat16	✓	✓	🟢	iron/operators/elementwise_mul/
GEMM	通用矩阵乘法内核	bfloat16	✓	✓	🟢	iron/operators/gemm/
GEMV	通用矩阵-向量乘法内核	bfloat16	✓	✓	🟢	iron/operators/gemv/
GQA	分组查询注意力内核（单流水线）	bfloat16		✓	🟢	iron/operators/mha/
MHA	多头注意力内核 & 分组查询注意力	bfloat16		✓	🟢	iron/operators/mha/
RMSNorm	RMSNorm 内核	bfloat16	✓	✓	🟢	iron/operators/rms_norm/
RoPE	旋转位置编码内核	bfloat16	✓	✓	🟢	iron/operators/rope/
SiLU	Sigmoid 线性单元激活内核	bfloat16	✓	✓	🟢	iron/operators/silu/
Softmax	Softmax 内核	bfloat16	✓	✓	🟢	iron/operators/softmax/
加权 RMSNorm	加权 RMSNorm 内核	bfloat16	✓	✓	🟢	iron/operators/rms_norm/
复制	复制	bfloat16	✓	✓	🟢	iron/operators/mem_copy/
转置	转置	bfloat16	✓	✓	🟢	iron/operators/transpose/
AXPY	AXPY	bfloat16	✓	✓	🟢	iron/operators/axpy/
Reduction	Reduction	bfloat16			🟡
Dequant	Dequant Q4NX from AWQ to bfloat16	bfloat16	✓	✓	🟢	iron/operators/dequant/
RELU	RELU	bfloat16	✓	✓	🟢	iron/operators/relu/
Leaky RELU	Leaky RELU	bfloat16	✓	✓	🟢	iron/operators/leaky_relu/
GELU	GELU	bfloat16	✓	✓	🟢	iron/operators/gelu/
LayerNorm	LayerNorm	bfloat16	✓	✓	🟢	iron/operators/layer_norm/
Convolution	Convolution	bfloat16			🟡
MaxPool	MaxPool	bfloat16			⚪
AveragePool	AveragePool	bfloat16			⚪
Tanh	Tanh kernel	bfloat16	✓	✓	🟢	iron/operators/tanh/
Sigmoid	Sigmoid kernel	bfloat16	✓	✓	🟢	iron/operators/sigmoid/

使用此仪表板快速查看每个内核的状态，并定位相关的设置、构建和使用信息。

📌 图例

状态	含义
🟢	已完成
🟡	开发中
⚪	未分配

安装 (Linux)

以下说明将指导您完成在 Ryzen™ AI NPU 上构建和运行程序所需的所有步骤，从全新的基础 Ubuntu 24.04 或 Ubuntu 24.10 安装开始。

初始设置

重要：确保您的系统拥有启用 NPU 支持的最新 BIOS 版本。请查看您的笔记本电脑/迷你 PC 制造商的支持网站以获取 BIOS 更新。

如果从 Ubuntu 24.04 开始，您可能需要通过安装硬件启用 (HWE) 堆栈将 Linux 内核更新到 6.11+：

sudo apt update
sudo apt install --install-recommends linux-generic-hwe-24.04
sudo reboot

安装 XDNA™ 驱动程序和 XRT：

来自 mlir-aie 仓库的说明

安装 IRON 和 MLIR-AIE 所需的软件包：

# Python versions 3.10, 3.12 and 3.13 are currently supported by our wheels
sudo apt install \
build-essential clang clang-14 lld lld-14 python3-venv python3-pip

设置虚拟环境并激活它：

python3 -m venv ironenv
source ironenv/bin/activate
python3 -m pip install --upgrade pip

源 XRT（在第 1 步中安装）：
```
source /opt/xilinx/xrt/setup.sh
```
安装所需的 Python 包（来自 requirements.txt）：
```
pip install -r requirements.txt
```
要测试您的安装，您可以尝试构建并运行以下示例：
```
pytest ./iron/operators/axpy/
```

构建/使用与测试算子

所有可用的运算符均可在 iron/operators 中找到。这些运算符各自包含：

op.py: Python 算子接口——一个易于访问的集成点，用于将算子集成到您的项目中，它规定了如何编译算子（构建产物）以及如何在运行时调用它（缓冲区大小等）。
design.py: 算子的 NPU 代码实现。通常引用 aie_kernels 中的内核作为计算核心代码，并使用 ObjectFIFOs 描述数据移动。
reference.py: 用于验证 NPU 实现正确性的参考 CPU 实现。
test.py: 一个端到端测试，实例化并构建算子，运行它并将其输出与参考实现进行验证。

注意：请确保已加载 XRT 设置脚本并激活 Python 环境： source /opt/xilinx/xrt/setup.sh source /path/to/ironenv/bin/activate

要构建并测试所有算子：

pytest iron/operators/ -m "not extensive"

要运行完整的测试套件：

pytest iron/operators/

要运行特定操作符的测试：

pytest iron/operators/axpy/

Git Hooks（可选但推荐）

为确保代码在推送前通过 CI 代码检查，请安装 pre-push 钩子：

cp scripts/hooks/pre-push .git/hooks/pre-push
chmod +x .git/hooks/pre-push

该钩子将运行与 CI 相同的 linting 检查：

许可证检查（reuse）
Python 格式化（black）
C++ 格式化（clang-format）

如需绕过该钩子：git push --no-verify

应用

Llama 3.2 1B 推理

IRON 包含一个完整的 LLM 推理示例，演示了 NPU 加速：

位置：iron/applications/llama_3.2_1b/
模型：Meta Llama 3.2 1B
特性：多头注意力、融合算子、bfloat16 量化

请参阅 iron/applications/llama_3.2_1b/README.md 获取设置和使用说明。

架构

IRON 采用三层架构：

算子（iron/operators/）：用于 NPU 操作的高级 Python API
- 每个算子包含：op.py（接口）、design.py（MLIR-AIE 实现）、reference.py（CPU 参考实现）、test.py（验证）
AIE 内核（aie_kernels/）：底层 C++ 计算内核
- 按架构组织：generic/、aie2/、aie2p/
- 使用 AIE API 进行向量化以实现最佳性能
通用基础设施（iron/common/）：编译、设备管理和工具
- MLIR-AIE 编译流水线
- XRT 运行时集成
- 算子融合框架

性能

IRON 算子旨在实现最大的 NPU 利用率：

跨多个 AIE 列的并行执行
通过 ObjectFIFOs 优化数据移动
融合操作以最小化主机-NPU 传输
使用 AIE 内置函数进行向量化内核

运行基准测试：

# Run all operators with performance metrics stored in tests_latest.csv
pytest iron/operators/ -m "not extensive" -v

社区与支持

💬 Discord: 加入我们的 Discord 服务器进行讨论和支持
🐛 Issues: 通过 GitHub Issues 报告错误和请求功能
📖 Contributing: 参见 CONTRIBUTING.md 了解开发指南
📚 Documentation: 算子示例见 iron/operators/，内核文档见 aie_kernels/README.md

许可证

IRON 采用 Apache License 2.0 许可。详见 LICENSE。