MLIR-AIR Programming Examples

These programming examples demonstrate how to leverage the AIR design flow with mlir-air Python bindings and the mlir-air intermediate representation (IR) to build applications targeting AI Engines on AMD NPUs.

Operator Dashboard

Category	Operation	Datatype(s)	NPU1	NPU2	Design Example
Linear Algebra	Matrix Multiplication	bf16, i16, i8	🟢	🟢	matrix_multiplication/
Linear Algebra	Vector-Matrix Multiplication	bf16	🟢	🟢	vector_matrix_multiplication/
Linear Algebra	Matrix-Vector Multiplication	bf16	⚪	🟢	matrix_vector_multiplication/bf16/
Linear Algebra	Matrix-Vector Multiplication (Cascade)	bf16	🟢	🟢	matrix_vector_multiplication/bf16_cascade/
Linear Algebra	Matrix-Vector Multiplication (AWQ int4)	int4 weights / bf16 activations	⚪	🟢	matrix_vector_multiplication/int4_awq/
Linear Algebra	Matrix Multiplication (AWQ int4)	int4 weights / bf16 activations	⚪	🟢	matrix_multiplication/int4_awq/
Linear Algebra	Matrix Multiplication (bf16 x bfp16ebs8)	bf16 activations / bfp16ebs8 weights	⚪	🟢	matrix_multiplication/bf16_x_bfp16/
Linear Algebra	AXPY	bf16	🟢	🟢	axpy/
Element-wise	Element-wise Add	f32	🟢	🟢	eltwise_add/
Element-wise	Element-wise Add (with L2)	f32	🟢	🟢	eltwise_add_with_l2/
Element-wise	Element-wise Add (bf16)	bf16	🟢	🟢	primitives/vector_examples/vector_add/
Element-wise	Element-wise Mul	bf16	🟢	🟢	primitives/vector_examples/vector_mul/
Activation/Math	SiLU	bf16	⚪	🟢	silu/
Activation/Math	GELU	bf16	⚪	🟢	gelu/
Activation/Math	Softmax	bf16	🟢	🟢	softmax/
Activation/Math	Sine / Cosine	bf16	🟢	⚪	sine_cosine/
Activation/Math	RELU	bf16	🟢	🟢	relu/
Activation/Math	Leaky RELU	bf16	🟢	🟢	leaky_relu/
Activation/Math	Sigmoid	bf16	⚪	🟢	sigmoid/
Activation/Math	Tanh	bf16	⚪	🟢	primitives/vector_examples/vector_tanh/
Normalization	Layer Normalization	bf16	⚪	🟢	layer_norm/
Normalization	RMS Normalization	bf16	⚪	🟢	rms_norm/
Normalization	Weighted RMS Normalization	bf16	⚪	🟢	weighted_rms_norm/
Aggregation	Reduction (Add)	bf16	🟢	🟢	primitives/vector_examples/vector_reduce_add/
Pooling	MaxPool	bf16	🟢	🟢	primitives/vector_examples/vector_reduce_max/
Pooling	AveragePool	bf16	🟢	🟢	average_pool/
LLM Kernels	Multi-Head Attention (LLaMA2)	bf16	🟢	⚪	llama2_mha/
LLM Kernels	Attention (Decode)	bf16	⚪	🟢	attention_decode/
LLM Kernels	SwiGLU	bf16	⚪	🟢	swiglu/
LLM Kernels	FFN SwiGLU (Decode)	bf16	⚪	🟢	ffn_swiglu/decode/
LLM Kernels	FFN SwiGLU (Prefill)	bf16	⚪	🟢	ffn_swiglu/prefill/
LLM Kernels	RoPE (LUT-based)	bf16	⚪	🟢	rope_lut/
LLM Kernels	RoPE (On-chip Sin/Cos)	bf16	🟢	🟢	rope_sincos/
LLM Kernels	Fused Superkernel Decode	Q4NX weights / bf16 activations	⚪	🟢	fused_decode/
Attention	Flash Attention (Dataflow)	bf16	🟢	🟢	flash_attention/dataflow_based/
Attention	Flash Attention (Kernel Fusion)	bf16	🟢	🟢	flash_attention/kernel_fusion_based/
Attention	Grouped Query Attention (GQA)	bf16	🟢	🟢	flash_attention/kernel_fusion_based/
Data Movement	Passthrough (DMA)	u8, i8, i16, u16, f32, bf16	🟢	🟢	passthrough/passthrough_dma/
Data Movement	Passthrough (Channel)	u8	🟢	🟢	passthrough/passthrough_channel/
Data Movement	Passthrough (Kernel)	u8	🟢	🟢	passthrough/passthrough_kernel/
Data Movement	Shim DMA 2D	i32	🟢	🟢	shim_dma_2d/
Data Movement	Data Transfer Transpose	u32	🟢	🟢	data_transfer_transpose/
Data Movement	Transpose (bf16)	bf16	⚪	🟢	data_transfer_transpose/dma_bf16/
Data Movement	Matrix Scalar Add	i32	🟢	🟢	matrix_scalar_add/
Communication	Channel Examples	i32	🟢	🟢	channel_examples/
Communication	3D Channel with Segment Unroll	i32	⚪	🟢	channel_examples/channel_3d_segment_unroll/
Communication	Broadcast Selective Capture	i32	🟢	🟢	channel_examples/broadcast_selective_capture/
Communication	Dual-Herd Packet Switch	bf16	⚪	🟢	channel_examples/dual_herd_packet_switch/
Communication	Multi-Segment Examples	i32	🟡	🟡	multi_segment/
Communication	Cascade Reduction	i32	🟢	🟢	cascade_reduction/
Memory	Segment Alloc	i32	🟢	🟢	segment_alloc/
Spatial	Segment Unroll	i32	🟢	🟢	segment_unroll/
Dataflow	Herd Dataflow	bf16	🟢	🟢	herd_dataflow/
Control Flow	Conditional Branching	i32	🟢	🟢	conditional_branching/
CNN	2D Convolution	i32	🟢	🟢	conv2d/
CNN	Conv2d 14x14	ui8/i8	🟢	🟢	conv2d_14x14/
CNN	Bottleneck	bf16	🟢	🟢	bottleneck/
ML Pipeline	MNIST-FC (Broadcast Bias Add)	f32	⚪	🟢	mnist_fc/broadcast_bias_add/
ML Pipeline	MNIST-FC (ReLU 2D)	f32/bf16	⚪	🟢	mnist_fc/relu/
ML Pipeline	MNIST-FC (Argmax)	f32→i32	⚪	🟢	mnist_fc/argmax/
ML Pipeline	MNIST-FC (Integration)	f32	⚪	🟢	mnist_fc/integration/
Memory	Shared L1 Buffer (Multi-Herd)	bf16	🟢	⚪	shared_l1_multi_herd/
Memory	Shared L1 Buffer (Single-Herd)	bf16	🟢	⚪	shared_l1_single_herd/
Quantization	Dequant (AWQ int4→bf16)	int4/bf16	🟢	🟢	dequant_awq/
Primitives	Scalar/Vector Operations	various	🟢	🟢	primitives/

Status Legend

• 🟢 Supported and tested
• 🟡 Work in progress
• ⚪ Not yet supported

NPU1 = AMD Ryzen AI (Phoenix, AIE2)    NPU2 = AMD Ryzen AI (Strix, AIE2P)