Skip to content

AIE C++ Kernel Intrinsics Cheatsheet

Before writing a kernel from this cheatsheet, check builtin_kernels.md — aie.iron.kernels already ships maintained, vectorized kernels for element-wise ops, reductions, activations, matmul, conv2d, and vision. This file is for the cases those don't cover.

All code targets aie_api/aie.hpp. Always include aie_kernel_utils.h for portable loop annotations.

Standard kernel file boilerplate

#define NOCPP
#include <cstdint>
#include <type_traits>
#include "aie_kernel_utils.h"        // copy from aie_kernels/ with its aie_arch.h, or adjust path
#include <aie_api/aie.hpp>

// Templated implementation
template <typename T_in, typename T_out, int N>
static inline void eltwise_add_impl(const T_in *__restrict a,
                                    const T_in *__restrict b,
                                    T_out      *__restrict c) {
    constexpr int VEC = 32;              // natural bf16 width; multiplies: AIE_BF16_LANES
    static_assert(N % VEC == 0, "N must be divisible by VEC");
    constexpr int F = N / VEC;

    event0();
    AIE_PREPARE_FOR_PIPELINING           // no-op under Peano; only Chess honors it
    AIE_LOOP_MIN_ITERATION_COUNT(F)      // guaranteed minimum trip count of this loop
    for (int i = 0; i < F; ++i) {
        aie::vector<T_in,  VEC> va = aie::load_v<VEC>(a); a += VEC;
        aie::vector<T_in,  VEC> vb = aie::load_v<VEC>(b); b += VEC;
        aie::vector<T_out, VEC> vc = aie::add(va, vb);
        aie::store_v(c, vc); c += VEC;
    }
    event1();
}

// extern "C" wrapper(s) — these names go into the IRON Kernel(...) declaration
// Keep __restrict (and const on inputs) on the wrapper itself, not just the
// templated impl — this is the symbol IRON links against, and __restrict here
// is what lets the compiler pipeline the loop.
extern "C" {
void eltwise_add_bf16_vector(const bfloat16 *__restrict a,
                              const bfloat16 *__restrict b,
                              bfloat16 *__restrict c) {
    eltwise_add_impl<bfloat16, bfloat16, 1024>(a, b, c);
}
} // extern "C"

Loop annotations (from aie_kernel_utils.h)

Macro What it does
AIE_PREPARE_FOR_PIPELINING Pipelining hint for Chess only; expands to nothing under Peano
AIE_LOOP_MIN_ITERATION_COUNT(n) Promise the loop runs at least n times — a trip-count hint. On AIE2 it let runtime-count loops overlap (axpy 269 → 87); under Peano it can also cost the zero-overhead loop, so check non_zol_loops in the remarks report
AIE_LOOP_MAX_ITERATION_COUNT(n) Promise the loop runs at most n times
AIE_LOOP_RANGE(min, max) Combo of the two above
AIE_LOOP_UNROLL(n) Unroll by factor n
AIE_LOOP_UNROLL_FULL Fully unroll
AIE_LOOP_NO_UNROLL Block unrolling
AIE_TRY_INITIATION_INTERVAL(n) Request an initiation interval of n (Peano only)
AIE_PREPARE_FOR_POSTPIPELINING Disable Peano's pipeliner for this loop (Peano only)
AIE_NO_PREPARE_FOR_PIPELINING Block pipelining (Chess only; rare, for setup loops)

Use the macros rather than backend-specific pragmas. Note they don't all map onto both backends: AIE_LOOP_MIN_ITERATION_COUNT / AIE_LOOP_RANGE / AIE_LOOP_UNROLL* become clang loop pragmas under Peano/AIECC and the equivalent under Chess, but AIE_PREPARE_FOR_PIPELINING and AIE_LOOP_FLATTEN map to chess:: hints only — they are no-ops under Peano. All expand to nothing on host builds. Peano pipelines inner loops without any hint, so none of these is needed to get a pipelined loop; for which hints have measurably helped, see the aie-kernel-opt skill.

Vector load / store

aie::vector<T, N> v = aie::load_v<N>(ptr);   // contiguous load of N elements
aie::store_v(ptr, v);                         // contiguous store
// Aligned loads are required for N ≥ natural width; use restrict + aligned pointers from ObjectFifo.

Element-wise arithmetic

auto sum  = aie::add(va, vb);              // vector + vector
auto prod = aie::mul(va, vb);              // returns accumulator for integer types
auto sc   = aie::add(va, scalar);          // vector + scalar (broadcasts)
auto neg  = aie::neg(va);
auto mx   = aie::max(va, vb);              // element-wise max  (use for ReLU vs zeros)
auto mn   = aie::min(va, vb);
auto absv = aie::abs(va);

Multiply-accumulate (non-MMUL)

aie::accum<acc32, VEC> acc;
acc.from_vector(aie::load_v<VEC>(c_ptr));
acc = aie::mac(acc, va, vb);                // acc += va * vb
aie::store_v(c_ptr, acc.template to_vector<T_out>(shift));  // shift = number of fractional bits to drop

Broadcast / zeros / ones

auto zeros  = aie::zeros<T, VEC>();
auto allone = aie::broadcast<T, VEC>(T(1));
auto sclrv  = aie::broadcast<T, VEC>(scalar_value);

Type conversion

aie::vector<float, 16>  vf  = ...;
aie::vector<bfloat16, 32> vbf = aie::to_v32bfloat16(vf);   // (when widths line up)
auto v_i8  = acc.template to_vector<int8_t>(shift);
acc.template to_vector<bfloat16>();                         // no shift for float-acc → bf16

MMUL — matrix multiply intrinsic (the perf path)

// 1. Define the MMUL shape: r × s @ s × t → r × t   (see architecture.md for valid shapes)
using MMUL = aie::mmul<r, s, t, T_in, T_in, accauto>;

// 2. Sizes of each operand fragment in elements:
//    MMUL::size_A = r*s, MMUL::size_B = s*t, MMUL::size_C = r*t

// 3. Construct an MMUL initialized from a partial accumulator in C
aie::vector<T_out, MMUL::size_C> c_init = aie::load_v<MMUL::size_C>(pC);
MMUL C00(c_init);

// 4. Inner reduction loop
for (unsigned k = 0; k < K_TILES; ++k) {
    auto A = aie::load_v<MMUL::size_A>(pA); pA += MMUL::size_A;
    auto B = aie::load_v<MMUL::size_B>(pB); pB += MMUL::size_B;
    C00.mac(A, B);                          // C00 += A · B
}

// 5. Write back (optionally with shift for fixed-point)
aie::store_v(pC, C00.template to_vector<T_out>());
// or: aie::store_v(pC, C00.template to_vector<T_out>(SHIFT));

Outer expansion (per-tile macro-kernel)

For real GEMMs you unroll the MMUL across an outer block (e.g., 4×4 expansion: hold 16 MMUL accumulators in registers, share each loaded A across 4 Bs and vice versa). See mlir-aie/aie_kernels/linalg/mm_aie2.h for canonical implementations.

// Sketch of 2x2 expansion
MMUL C00(load_v<MMUL::size_C>(pC1));
MMUL C01(load_v<MMUL::size_C>(pC1 + MMUL::size_C));
MMUL C10(load_v<MMUL::size_C>(pC2));
MMUL C11(load_v<MMUL::size_C>(pC2 + MMUL::size_C));

AIE_PREPARE_FOR_PIPELINING
AIE_LOOP_MIN_ITERATION_COUNT(2)
for (unsigned k = 0; k < K_TILES; ++k) {
    auto A0 = load_v<MMUL::size_A>(pA1); pA1 += MMUL::size_A;
    auto A1 = load_v<MMUL::size_A>(pA2); pA2 += MMUL::size_A;
    auto B0 = load_v<MMUL::size_B>(pB1); pB1 += MMUL::size_B * colB;
    auto B1 = load_v<MMUL::size_B>(pB2); pB2 += MMUL::size_B * colB;
    C00.mac(A0, B0); C01.mac(A0, B1);
    C10.mac(A1, B0); C11.mac(A1, B1);
}

Reduction

Horizontal vector reduction is done by repeated shift_bytes + aie::add/max:

template <typename T, int N>
T reduce_sum(const T *__restrict in, int total) {
    using V = aie::vector<T, N>;
    V acc = aie::zeros<T, N>();
    event0();
    AIE_PREPARE_FOR_PIPELINING
    AIE_LOOP_MIN_ITERATION_COUNT(8)
    for (int i = 0; i < total; i += N) {
        acc = aie::add(acc, aie::load_v<N>(in + i));
    }
    event1();
    return aie::reduce_add(acc);            // final horizontal sum
}

// reduce_max: same shape, use aie::max() and aie::reduce_max()

Saturation / rounding (for integer kernels)

::aie::set_saturation(aie::saturation_mode::saturate);    // clamp on overflow (vs wrap)
::aie::set_rounding  (aie::rounding_mode::symmetric_inf); // shift rounding mode

Call once at the top of the kernel (or in a wrapper) before any vector op that could overflow, never inside the hot loop (see pitfalls.md). aie::swap_rounding returns the previous mode for restoring it.

Accumulator type quick-pick

Input Multiply result Accumulator
int8 int16 / int32 acc32
int16 int32 acc32 (or acc64 for wide)
int32 int64 acc64
bfloat16 float accfloat (accauto in MMUL)
float float accfloat

If you're unsure, use accauto inside aie::mmul<>; the compiler picks.

Common operations table

Op Call
Vector add aie::add(va, vb)
Vector mul (returns acc) aie::mul(va, vb)
Vector MAC aie::mac(acc, va, vb)
Vector load aie::load_v<N>(ptr)
Vector store aie::store_v(ptr, v)
Broadcast scalar aie::broadcast<T, N>(s)
Zero vector aie::zeros<T, N>()
Element max aie::max(va, vb)
Element min aie::min(va, vb)
Element abs aie::abs(va)
Horizontal sum aie::reduce_add(v)
Horizontal max aie::reduce_max(v)
Acc → vector acc.template to_vector<T>(shift)
Vector → acc acc.from_vector(v)
Matrix multiply tile MMUL C(c0); C.mac(A, B); store_v(p, C.to_vector<T>());