diff --git a/aie_kernels/generic/activations.h b/aie_kernels/generic/activations.h new file mode 100644 index 0000000000..b9da2f1c00 --- /dev/null +++ b/aie_kernels/generic/activations.h @@ -0,0 +1,112 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Branch-free activations over an aie::vector, for mm_fused's fused +// epilogue. +// +// These are templated on the vector width and return a vector, so they compose +// inside an existing vector loop -- unlike aie_kernels//{gelu,silu}.cc, +// which are whole-buffer entry points over a fixed 32-wide vector. mm_fused's +// epilogue needs the former: it applies the activation to the f32 accumulator +// it already has in registers, without a second pass over L1. +// +// All three are built on tanh, so the sigmoid below is exact-by-identity rather +// than a polynomial fit: +// sigmoid(x) == (tanh(x/2) + 1) / 2 +// GELU then uses the sigmoid approximation gelu(x) ~= x * sigmoid(1.702x), +// which is NOT the same curve as gelu.cc's tanh approximation +// 0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3))) +// The two agree to well within bf16 precision over the range that matters, but +// they are different functions -- do not expect bit-identical results if you +// compare against the gelu operator. +// +// Where that tanh comes from is the one part of mm_fused that is genuinely +// architecture-specific: +// +// AIE2P has a native vector tanh (aie::tanh) evaluated on an f32 accumulator. +// AIE2 (Phoenix) does not, so it falls back to the piecewise-linear LUT in +// mlir-aie's aie_runtime_lib/AIE2/lut_based_ops.h -- the same getTanhBf16 that +// aie_kernels/aie2/{tanh,sigmoid}.cc already use. That LUT is fixed at 16 +// lanes, which is exactly the epilogue's vector width, so nothing needs +// splitting; the static_assert below pins that assumption. +// +// The two paths therefore do NOT produce bit-identical results, and the AIE2 +// path carries the LUT's approximation error on top of bf16 rounding. test.py +// sets the accuracy budget per architecture accordingly. +#ifndef __ACTIVATIONS_H__ +#define __ACTIVATIONS_H__ +#include + +#if __AIE_ARCH__ >= 21 +#define ACTIVATIONS_NATIVE_TANH 1 +#else +#define ACTIVATIONS_NATIVE_TANH 0 +// Supplies getTanhBf16. Resolved from the runtime-lib include directory the +// build adds for the target arch (aie_runtime_lib/AIE2), not from this file's +// own directory. +#include "lut_based_ops.h" +#endif + +// All three take and return f32, and mm_fused converts to bf16 once at the end. +// Rounding the accumulator to bf16 BEFORE the activation, as an earlier version +// did, rounds twice and lets the activation's slope amplify the first rounding; +// measured against an exact f64 evaluation that costs 1.35x the error on silu, +// 1.17x on gelu and 1.08x on sigmoid. +// +// tanh is the one step that cannot always stay f32: AIE2P has a native f32 +// tanh, while AIE2's LUT is bf16-only, so on AIE2 the tanh result is rounded +// and widened back. Everything around it -- the x/2 and 1.702x scalings, the +// (t+1)/2, and silu/gelu's outer multiply by x -- stays f32 on both, and that +// outer multiply is where most of the silu gain comes from. + +// tanh of an f32 vector, on whichever path this architecture has. +template +__attribute__((always_inline)) aie::vector tanh_vec(aie::vector x) +{ + // bf16 out on both paths, widened back to f32. Asking AIE2P for the f32 + // tanh instead defeats aiecc's stack measurement -- it reports a spurious + // "__start -> _main_init -> core -> _main_init" recursion -- and tanh's + // output is in [-1, 1], where bf16 costs at most 2^-9 absolute anyway. The + // arithmetic AROUND it is where the f32 actually pays. + aie::accum widened; +#if ACTIVATIONS_NATIVE_TANH + widened.from_vector(aie::tanh(x)); +#else + static_assert(vec_size == 16, + "AIE2's LUT tanh is fixed at 16 lanes, which is mm_fused's " + "epilogue width; widening V needs an explicit split here"); + aie::accum narrowed; + narrowed.from_vector(x); + aie::vector tanh_bf16 = getTanhBf16(narrowed.template to_vector()); + widened.from_vector(tanh_bf16); +#endif + return widened.template to_vector(); +} + +// sigmoid(x) = (tanh(x/2) + 1) / 2 +template +__attribute__((always_inline)) aie::vector sigmoid_vec(aie::vector x) +{ + const aie::vector v_half = aie::broadcast(0.5f); + const aie::vector v_one = aie::broadcast(1.0f); + aie::vector t = tanh_vec(aie::mul(x, v_half).template to_vector()); + return aie::mul(aie::add(t, v_one), v_half).template to_vector(); +} + +// silu(x) = x * sigmoid(x) +template +__attribute__((always_inline)) aie::vector silu_vec(aie::vector x) +{ + return aie::mul(x, sigmoid_vec(x)).template to_vector(); +} + +// gelu(x) ~= x * sigmoid(1.702x) +template +__attribute__((always_inline)) aie::vector gelu_vec(aie::vector x) +{ + const aie::vector v_scale = aie::broadcast(1.702f); + aie::vector scaled = aie::mul(x, v_scale).template to_vector(); + return aie::mul(x, sigmoid_vec(scaled)).template to_vector(); +} + +#endif // __ACTIVATIONS_H__ diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc new file mode 100644 index 0000000000..719c09361b --- /dev/null +++ b/aie_kernels/generic/mm_fused.cc @@ -0,0 +1,190 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// bf16 GEMM compute kernel. Each compute tile owns an m x n slice of C and +// accumulates over K into an f32 accumulator that stays in L1 for the whole +// reduction. +// +// Three entry points, each called once per iteration of a loop nest that lives +// in the design (iron/operators/flm/gemm/design.py) rather than here: +// +// mm_fused_acc_init zero the accumulator, once per output tile +// mm_fused_k_step multiply one A band by one B chunk into it +// mm_fused_epilogue_chunk drain one chunk of it to a bf16 C object +// +// The nest lives in the design so that every level of it has an ObjectFifo +// acquire point, a fifo consumer having to acquire once per object. +// +// Tile geometry arrives as -D flags from design.py, which is the single source +// of truth for it: the same constants size the design's buffers and set its +// unroll factors. +#include "../aie_kernel_utils.h" +#include "activations.h" +#include "mm_fused_mmul.h" +#include "zero.cc" + +#include +#include + +#if !defined(MM_FUSED_TILE_M) || !defined(MM_FUSED_TILE_K) || !defined(MM_FUSED_TILE_N) || !defined(MM_FUSED_CT_K) +#error "design.py must pass -DMM_FUSED_TILE_M / _TILE_K / _TILE_N / _CT_K" +#endif +#if !defined(MM_FUSED_OUT_CHUNK) || !defined(MM_FUSED_C_DEPTH) +#error "design.py must pass -DMM_FUSED_OUT_CHUNK / -DMM_FUSED_C_DEPTH" +#endif + +// Epilogue selection. 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid, matching +// Epilogue.mode in design.py. +#ifndef MM_FUSED_EPILOGUE_MODE +#define MM_FUSED_EPILOGUE_MODE 0 +#endif +#ifndef MM_FUSED_CLAMP +#define MM_FUSED_CLAMP 0 +#endif +#ifndef MM_FUSED_CLAMP_MIN +#define MM_FUSED_CLAMP_MIN 0.0f +#endif +#ifndef MM_FUSED_CLAMP_MAX +#define MM_FUSED_CLAMP_MAX 0.0f +#endif + +namespace +{ +constexpr int M = MM_FUSED_TILE_M; +// Asymmetric tile buffering: the A tile spans MA rows while the accumulator +// spans M, so the core folds RHO = M / MA A bands into one C tile before +// releasing it. A dies as soon as it is consumed while C must live across the +// whole K reduction, so sizing both to M would pay the peak L1 cost twice. +// MA == M is the symmetric case. +// +// Technique from "Can Asymmetric Tile Buffering Be Beneficial?", C. Wang, +// W. Pang, X. Wu, G. Jun, L. Romero, E. Taka, D. Marculescu, T. Nowatzki, +// P. Vasireddy, J. Melber, D. Chen, J. Cong, arXiv:2511.16041 (2025), +// https://arxiv.org/abs/2511.16041. Reference AIE implementation is +// Xilinx/mlir-aie PR #3076 by @ChengyueWang, in +// programming_examples/ml/block_datatypes/gemm_asymmetric_tile_buffering. +// Those configs accumulate in bf16/bfp16, which is what affords their larger C +// tiles; this kernel keeps an f32 accumulator, so here the win comes from +// spending the freed L1 on a deeper k slice rather than on a wider C tile. +constexpr int MA = MM_FUSED_TILE_MA; +constexpr int K = MM_FUSED_TILE_K; +constexpr int N = MM_FUSED_TILE_N; +// Register tiling, and how much of K one compute tile holds at a time. Both are +// design.py's to choose -- CT_K in particular trades against the n width for a +// fixed L1 budget. +constexpr int R = MM_FUSED_R; +constexpr int S = MM_FUSED_S; +constexpr int T = MM_FUSED_T; +constexpr int CT_K = MM_FUSED_CT_K; + +// Output stage geometry. +constexpr int CHUNK = MM_FUSED_OUT_CHUNK; +constexpr int C_DEPTH = MM_FUSED_C_DEPTH; +constexpr int V = 16; // one 512-bit bf16 vector +static_assert(CHUNK % V == 0, "output chunk must be a whole number of vectors"); + +// Same divisibility conditions mm.cc asserts for its own 2x2 mmul, plus the +// two the k blocking adds. +static_assert(M % MA == 0, "tile_m must be a whole number of A bands"); +static_assert(MA % (2 * R) == 0, "tile_ma must be a multiple of 2*r (2x2 mmul)"); +static_assert(N % (2 * T) == 0, "tile_n must be a multiple of 2*t (2x2 mmul)"); +static_assert(K % CT_K == 0, "tile_k must be a multiple of the k slice"); +static_assert(CT_K % S == 0, "k slice must be a multiple of s"); + +// The core powers up in rounding_mode::floor, so a kernel that converts must +// choose explicitly. Truncation biases every conversion the same direction, so +// the error accumulates over the K reduction instead of cancelling -- ~1% of +// the result, against ~0.02% for round-to-nearest-even, which is far more than +// the bfp16 emulation itself costs. Every entry point that converts sets it: +// the mmul and the epilogue's f32->bf16 store, both below. +// +// Flag name and polarity follow mm.cc, so the two kernels are configured the +// same way; the operator passes -DROUND_CONV_EVEN by default. +#ifdef ROUND_CONV_EVEN +constexpr aie::rounding_mode round_mode = aie::rounding_mode::conv_even; +#else +constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; +#endif +} // namespace + +extern "C" { + +// Zero the f32 accumulator, before the k loop starts accumulating into it. +// +// A bias is deliberately not supported: initialising the accumulator from one +// would mean consuming an extra object through the handshake the B ObjectFifo +// owns, which desynchronises that fifo and hangs rather than mis-computing. +void mm_fused_acc_init(float *y_acc) +{ + // zero_vectorized brackets itself in event0/event1 for tracing. + zero_vectorized(y_acc); +} + +// One step of the k loop: one B chunk multiplied against one A band, +// accumulated into y_acc. +// +// Takes no locks. A is a single object spanning every z slice of the mmul, and +// the A and B fifos own the handshake, so the core body acquires around this +// call rather than the kernel acquiring inside it. +// mm_fused_b_elem_t is bfp16ebs8 or bfloat16 depending on how B is stored, +// which mm_fused_mmul.h selects from the architecture. One signature either +// way, so the design's Kernel declaration does not have to care. +void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, int32_t band) +{ + ::aie::set_rounding(round_mode); + // The accumulator is [row-block][col-block][r*t], so band b starts at + // b * MA * N -- b*(MA/R) row-blocks in, each colB*(r*t) wide. + mm_fused_mmul_2x2<(MA / R), (CT_K / S), (N / T), R, S, T>(a_buf, b_buf, y_acc + band * (MA * N)); +} + +// The output stage: convert chunk (outer * C_DEPTH + half) of the f32 +// accumulator into a bf16 C object the core body has already acquired from the +// C ObjectFifo, optionally applying an activation and a clamp on the way out. +// +// Fusing the activation here is the point: the values are already in registers +// after the f32 -> bf16 conversion, so gelu/silu/sigmoid costs one more vector +// op per 16 elements instead of a separate pass over L1 (which is what chaining +// a standalone activation operator after a GEMM would cost). The mode and clamp +// are compile-time, so the inner loop below is branch-free. +// +// The chunk index is split in two because the core body unrolls the drain by +// the C fifo depth to keep the acquired buffer index a compile-time constant; +// passing both parts avoids doing that arithmetic up there. +void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half) +{ + // The store below is a conversion, so it obeys the same rounding mode the + // mmul does and must agree with it. + ::aie::set_rounding(round_mode); + const float *__restrict src = y_acc + (outer * C_DEPTH + half) * CHUNK; + +#if MM_FUSED_CLAMP + const aie::vector lo = aie::broadcast(MM_FUSED_CLAMP_MIN); + const aie::vector hi = aie::broadcast(MM_FUSED_CLAMP_MAX); +#endif + + AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) + for (int j = 0; j < CHUNK / V; j++) { + // The accumulator stays f32 through the activation and the clamp, and + // is converted to bf16 exactly once, on the store. Converting first + // would round twice and let the activation's slope amplify the first + // rounding -- see activations.h. + aie::vector f = aie::load_v(src + j * V); +#if MM_FUSED_EPILOGUE_MODE == 1 + f = gelu_vec(f); +#elif MM_FUSED_EPILOGUE_MODE == 2 + f = silu_vec(f); +#elif MM_FUSED_EPILOGUE_MODE == 3 + f = sigmoid_vec(f); +#endif +#if MM_FUSED_CLAMP + f = aie::max(aie::min(f, hi), lo); +#endif + aie::accum out; + out.from_vector(f); + // The assignment is the conversion: to_v16bfloat16 yields a raw + // v16bfloat16, not an aie::vector. + aie::vector v = to_v16bfloat16(out); + aie::store_v(y_out + j * V, v); + } +} +} diff --git a/aie_kernels/generic/mm_fused_mmul.h b/aie_kernels/generic/mm_fused_mmul.h new file mode 100644 index 0000000000..28645db652 --- /dev/null +++ b/aie_kernels/generic/mm_fused_mmul.h @@ -0,0 +1,204 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// The 2x2 blocked mmul at the heart of mm_fused.cc. +// +// rowA/colA/colB count r x s (A), s x t (B) and r x t (C) sub-tiles, not +// elements. A is one contiguous buffer spanning every z slice, so it is indexed +// directly rather than re-based per slice, and the mmul takes no locks. B is +// laid out column-major over blocks -- block (i, j) at (j * colA + i) -- which +// is the order the operator's pack_B emits. +// +// Each iteration of the j loop keeps four accumulators live (C00/C01/C10/C11) +// over a 2x2 block of output tiles, so each pair of A loads and each pair of B +// loads feeds four macs. That 2x2 is load amortization, not loop unrolling: +// dropping to 1x1 would halve the arithmetic per load, and this kernel is +// load-port bound. It is also why the j and z loops step by two and cannot be +// rolled tighter without giving that up. +// +// There is one definition per B storage format, because the two are not both +// compilable on both architectures -- see the type selection below. They share +// mm_fused_store_2x2 and this documentation; the surrounding loop nest is +// written out in each rather than shared through a macro, which would hide the +// operand handling that is the only real difference between them. +#ifndef __MM_FUSED_MMUL_H__ +#define __MM_FUSED_MMUL_H__ +#include "../aie_kernel_utils.h" + +#include + +// B's element type. bfp16ebs8 is usable only on AIE2P: +// __AIE_API_SCALAR_BFP_TYPES__ is defined in aie_api/detail/aie2p/config.hpp, +// mmul_bfp16_bfp16.hpp exists only under aie2p/, and on AIE2 aie_api/types.hpp +// declares `struct bfp16ebs8 {}` -- an empty placeholder that turns any use of +// it into a compile error rather than a slow fallback. So AIE2 keeps B in bf16, +// and both forms below are live: one per architecture, not one plus a spare. +#ifdef MM_FUSED_BFP16_B +using mm_fused_b_elem_t = bfp16ebs8; +#else +using mm_fused_b_elem_t = bfloat16; +#endif + +// Write one 2x2 block of accumulators back to C. Templated on the accumulator +// type because the two forms hold it differently -- aie::mmul in one, +// aie::accum in the other -- while presenting the same to_vector. +template +__aie_inline void mm_fused_store_2x2(float *__restrict pC1, + float *__restrict pC2, + const Acc &C00, + const Acc &C01, + const Acc &C10, + const Acc &C11) +{ + aie::store_v(pC1, C00.template to_vector()); + aie::store_v(pC1 + sizeC, C01.template to_vector()); + aie::store_v(pC2, C10.template to_vector()); + aie::store_v(pC2 + sizeC, C11.template to_vector()); +} + +#ifdef MM_FUSED_BFP16_B + +// B arrives already quantized to bfp16ebs8. +// +// The hardware only multiplies bfp16, so the bf16 form below has to convert B +// inside every mac -- transpose, widen, to_v64bfp16ebs8. B is static weights, +// so pack_B does that conversion once on the host instead. The values are +// unchanged: this hoists a rounding that already happened rather than adding +// one. It also makes B 9 bytes per 8 elements instead of 16, which is the point +// -- the operator is data-movement bound. +// +// B is streamed rather than pointer-indexed because a block_vector cannot be +// aie::load_v'd, and because bfp16ebs8 pointer arithmetic counts bytes rather +// than blocks. TODO: index it directly once llvm-aie#1232 ("sizeof(bfp16ebs8) +// == 1, not 9") is fixed; open as of 2026-09-10. +template +__aie_inline void mm_fused_mmul_2x2(const bfloat16 *__restrict pA, const bfp16ebs8 *__restrict pB, float *__restrict pC) +{ + constexpr unsigned sizeA = r * s; + constexpr unsigned sizeB = s * t; + constexpr unsigned sizeC = r * t; + event0(); + AIE_LOOP_MAX_ITERATION_COUNT(rowA / 2) + for (unsigned z = 0; z < rowA; z += 2) { + float *__restrict pC1 = pC + (z * colB) * sizeC; + float *__restrict pC2 = pC + ((z + 1) * colB) * sizeC; + const bfloat16 *__restrict pA_cur = pA + (z >> 1) * (2 * r * colA * s); + + AIE_LOOP_MAX_ITERATION_COUNT(colB / 2) + for (unsigned j = 0; j < colB; j += 2) { + const bfloat16 *__restrict pA1 = pA_cur; + const bfloat16 *__restrict pA2 = pA_cur + colA * sizeA; + + aie::block_vector_input_buffer_stream pB1(pB); + aie::block_vector_input_buffer_stream pB2(pB); + pB1.seek(j * colA); + pB2.seek((j + 1) * colA); + + aie::accum C00(aie::load_v(pC1)); + aie::accum C01(aie::load_v(pC1 + sizeC)); + aie::accum C10(aie::load_v(pC2)); + aie::accum C11(aie::load_v(pC2 + sizeC)); + + aie::vector A0; + aie::vector A1; + aie::accum accA0; + aie::accum accA1; + + // Keep this loop rolled and its body minimal. Extra live state across it + // costs more than it saves, because Peano's register allocator runs + // before the post-RA pipeliner and manufactures false loop-carried + // anti-dependences. TODO: re-measure a hand-unroll once llvm-aie#1066 is + // fixed; open as of 2026-09-10. + AIE_LOOP_MAX_ITERATION_COUNT(colA) + for (unsigned i = 0; i < colA; i++) { + // One conversion per A operand per i, reused by both j accumulators, + // rather than one inside each of the four macs. Same values. + // + // The two operands are widened by different routes, as mlir-aie's + // mm_bfp_mixed.cc also does: widening both by assignment makes Peano's + // AIE2P backend abort with "Use not jointly dominated by defs". + // mul_elem_64 by one is the same arithmetic and does codegen. + A0 = aie::load_v(pA1); + pA1 += sizeA; + A1 = aie::load_v(pA2); + pA2 += sizeA; + accA0 = A0; + accA1 = mul_elem_64(A1, concat(broadcast_one_to_v32bfloat16(), broadcast_one_to_v32bfloat16())); + + aie::block_vector B0 = pB1.pop(); + aie::block_vector B1 = pB2.pop(); + + C00 = mac_8x8_8x8T(accA0.template to_vector(), B0, C00); + C01 = mac_8x8_8x8T(accA0.template to_vector(), B1, C01); + C10 = mac_8x8_8x8T(accA1.template to_vector(), B0, C10); + C11 = mac_8x8_8x8T(accA1.template to_vector(), B1, C11); + } + mm_fused_store_2x2(pC1, pC2, C00, C01, C10, C11); + pC1 += 2 * sizeC; + pC2 += 2 * sizeC; + } + } + event1(); +} + +#else + +// B arrives as plain bf16, in row-major s x t blocks. Used on AIE2, which has +// no bfp16 hardware and composes this shape from four native 4x8x4 macs. +template +__aie_inline void mm_fused_mmul_2x2(const bfloat16 *__restrict pA, const bfloat16 *__restrict pB, float *__restrict pC) +{ + using MMUL = aie::mmul; + static_assert(r * s == MMUL::size_A); + event0(); + AIE_LOOP_MAX_ITERATION_COUNT(rowA / 2) + for (unsigned z = 0; z < rowA; z += 2) { + float *__restrict pC1 = pC + (z * colB) * MMUL::size_C; + float *__restrict pC2 = pC + ((z + 1) * colB) * MMUL::size_C; + const bfloat16 *__restrict pA_cur = pA + (z >> 1) * (2 * r * colA * s); + + aie::vector A0; + aie::vector A1; + aie::vector B0; + aie::vector B1; + + AIE_LOOP_MAX_ITERATION_COUNT(colB / 2) + for (unsigned j = 0; j < colB; j += 2) { + const bfloat16 *__restrict pA1 = pA_cur; + const bfloat16 *__restrict pA2 = pA_cur + colA * MMUL::size_A; + const bfloat16 *__restrict pB1 = pB + (j * colA) * MMUL::size_B; + const bfloat16 *__restrict pB2 = pB + ((j + 1) * colA) * MMUL::size_B; + + MMUL C00(aie::load_v(pC1)); + MMUL C01(aie::load_v(pC1 + MMUL::size_C)); + MMUL C10(aie::load_v(pC2)); + MMUL C11(aie::load_v(pC2 + MMUL::size_C)); + + // Rolled, for the same reason as the bfp16 form above (llvm-aie#1066). + AIE_LOOP_MAX_ITERATION_COUNT(colA) + for (unsigned i = 0; i < colA; i++) { + A0 = aie::load_v(pA1); + pA1 += MMUL::size_A; + A1 = aie::load_v(pA2); + pA2 += MMUL::size_A; + B0 = aie::load_v(pB1); + pB1 += MMUL::size_B; + B1 = aie::load_v(pB2); + pB2 += MMUL::size_B; + + C00.mac(A0, B0); + C01.mac(A0, B1); + C10.mac(A1, B0); + C11.mac(A1, B1); + } + mm_fused_store_2x2(pC1, pC2, C00, C01, C10, C11); + pC1 += 2 * MMUL::size_C; + pC2 += 2 * MMUL::size_C; + } + } + event1(); +} + +#endif // MM_FUSED_BFP16_B + +#endif // __MM_FUSED_MMUL_H__ diff --git a/aie_kernels/generic/passThrough.cc b/aie_kernels/generic/passThrough.cc index f4a784de54..f2d3cbd65c 100644 --- a/aie_kernels/generic/passThrough.cc +++ b/aie_kernels/generic/passThrough.cc @@ -10,6 +10,11 @@ #include #include +// Element width in bits, chosen by the caller with -DBIT_WIDTH; 32 if unset. +#ifndef BIT_WIDTH +#define BIT_WIDTH 32 +#endif + template __attribute__((noinline)) void passThrough_aie(T *restrict in, T *restrict out, const int32_t height, const int32_t width) diff --git a/iron/common/__init__.py b/iron/common/__init__.py index cb2ff31bec..f448a6d654 100644 --- a/iron/common/__init__.py +++ b/iron/common/__init__.py @@ -16,6 +16,8 @@ KernelArchiveArtifact, SourceArtifact, PythonGeneratedMLIRArtifact, + RemoteFileArtifact, + InstsBinArtifact, DesignGenerator, ) from .layout import Stride, TiledStride, TiledStridedLayout, tiled_2d diff --git a/iron/common/compilation/__init__.py b/iron/common/compilation/__init__.py index d4e06c2e66..c1fb11855d 100644 --- a/iron/common/compilation/__init__.py +++ b/iron/common/compilation/__init__.py @@ -17,10 +17,12 @@ KernelObjectArtifact, KernelArchiveArtifact, PythonGeneratedMLIRArtifact, + RemoteFileArtifact, CompilationCommand, ShellCompilationCommand, PythonCallbackCompilationCommand, CompilationRule, + DownloadCompilationRule, GenerateMLIRFromPythonCompilationRule, AieccCompilationRule, AieccFullElfCompilationRule, diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index af06dd1289..f20e7c2439 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -37,9 +37,10 @@ from collections import deque from collections.abc import Iterator, Sequence from pathlib import Path +import hashlib import os.path import shutil -import zlib +import urllib.request import logging import subprocess import importlib.util @@ -408,6 +409,30 @@ def __init__( super().__init__(filename, dependencies=[SourceArtifact(generator.source_path)]) +def _sha256_of(path: Path) -> str: + with open(path, "rb") as f: + return hashlib.file_digest(f, "sha256").hexdigest() + + +class RemoteFileArtifact(CompilationArtifact): + """A file downloaded from a URL and pinned by its SHA-256 digest. + + The digest pins the content, so ``url`` must name an immutable revision of + the file -- a commit SHA rather than a branch. + """ + + def __init__(self, filename: str, url: str, sha256: str) -> None: + super().__init__(filename) + self.url = url + self.sha256 = sha256 + + def is_available_in_filesystem(self) -> bool: + # A stale file with a matching mtime is still the wrong file, so + # compare content rather than timestamps. + path = Path(self.filename) + return path.exists() and _sha256_of(path) == self.sha256 + + # Compilation Command # ########################################################################## @@ -485,6 +510,40 @@ def compile(self, artifacts: CompilationArtifactGraph) -> list[CompilationComman pass +class DownloadCompilationRule(CompilationRule): + """Fetch RemoteFileArtifacts over HTTPS and check their digest.""" + + def matches(self, graph): + return any(graph.get_worklist(RemoteFileArtifact)) + + def compile(self, graph): + commands = [] + for artifact in graph.get_worklist(RemoteFileArtifact): + commands.append( + PythonCallbackCompilationCommand(partial(self.download, artifact)) + ) + artifact.available = True + return commands + + @staticmethod + def download(artifact): + if not artifact.url.startswith("https://"): + raise ValueError(f"refusing to download over {artifact.url!r}") + # Download beside the target and rename, so an interrupted fetch cannot + # leave a truncated file that a later run reports as a digest mismatch. + target = Path(artifact.filename) + partial_path = target.with_suffix(target.suffix + ".part") + with urllib.request.urlopen(artifact.url, timeout=60) as response: + partial_path.write_bytes(response.read()) + digest = _sha256_of(partial_path) + if digest != artifact.sha256: + partial_path.unlink() + raise RuntimeError( + f"{artifact.url} has SHA-256 {digest}, expected {artifact.sha256}" + ) + partial_path.replace(target) + + class GenerateMLIRFromPythonCompilationRule(CompilationRule): def matches(self, graph): return any(graph.get_worklist(PythonGeneratedMLIRArtifact)) diff --git a/iron/common/context.py b/iron/common/context.py index a7a5136ca6..c2e4b2e84b 100644 --- a/iron/common/context.py +++ b/iron/common/context.py @@ -53,6 +53,7 @@ def compilation_rules(self): return [ comp.FusePythonGeneratedMLIRCompilationRule(), comp.GenerateMLIRFromPythonCompilationRule(), + comp.DownloadCompilationRule(), comp.KernelCompilationRule(peano_dir, mlir_aie_dir, use_chess=use_chess), comp.ArchiveCompilationRule(peano_dir, mlir_aie_dir), comp.AieccXclbinInstsCompilationRule(use_chess=use_chess), diff --git a/iron/common/utils.py b/iron/common/utils.py index 99ad7327b2..c3b3bbe73c 100644 --- a/iron/common/utils.py +++ b/iron/common/utils.py @@ -35,3 +35,36 @@ def float_to_name(v: float) -> str: 1e-10 -> '1en10' """ return repr(v).replace(".", "p").replace("-", "n").replace("+", "") + + +# Widest wrap a shim or mem tile DMA buffer descriptor's size field can encode. +# Not exposed by the Python bindings (AIETargetModel::getDmaBdWrapBits is +# unbound), so it is written down here rather than in each design; gemv, +# repeat and mha all hardcoded the same 1023 independently. +# +# This is the same 10 bits on every target model this repo builds for -- +# BaseNPU1TargetModel and BaseNPU2TargetModel both inherit it unmodified from +# AIE2TargetModel::getDmaBdWrapBits, which does not override it per device -- +# so callers do not need to look it up per-device. It is NOT the same for +# every tile type, though: core tiles get an 8-bit wrap (max 255), not 10-bit. +# This constant is only valid for shim/mem tile descriptors, which is what +# every current caller (gemv, repeat, mha, flm.GEMM) uses it for. +DMA_BD_MAX_WRAP = (1 << 10) - 1 + + +def split_run(run: int, max_wrap: int = DMA_BD_MAX_WRAP) -> list[tuple[int, int]]: + """Encode a contiguous run of ``run`` elements as BD (size, stride) dims. + + One dimension suffices while the run fits the BD's size field; a longer run + splits into two at the cost of one of the four available dimensions. + + >>> split_run(512) + [(512, 1)] + >>> split_run(2048) + [(2, 1024), (1024, 1)] + """ + if run <= max_wrap: + return [(run, 1)] + if run % 2: + raise ValueError(f"cannot split an odd run ({run}) exceeding {max_wrap}") + return [(2, run // 2), (run // 2, 1)] diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index e7a05a7d61..744278f696 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -27,11 +27,16 @@ "Repeat": "repeat", } -__all__ = sorted(_OPERATOR_MODULES) +# Sub-packages whose operator names would collide with the table above. +_SUBPACKAGES = ("flm",) + +__all__ = sorted(set(_OPERATOR_MODULES) | set(_SUBPACKAGES)) def __getattr__(name): """Import the operator that defines `name`, on first access.""" + if name in _SUBPACKAGES: + return importlib.import_module(f".{name}", __name__) module = _OPERATOR_MODULES.get(name) if module is None: raise AttributeError(f"module {__name__!r} has no attribute {name!r}") @@ -39,4 +44,4 @@ def __getattr__(name): def __dir__(): - return sorted(set(globals()) | set(_OPERATOR_MODULES)) + return sorted(set(globals()) | set(_OPERATOR_MODULES) | set(_SUBPACKAGES)) diff --git a/iron/operators/flm/__init__.py b/iron/operators/flm/__init__.py new file mode 100644 index 0000000000..3bc3767ae7 --- /dev/null +++ b/iron/operators/flm/__init__.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Ports of FastFlowLM overlays to IRON. + +Operators are re-exported lazily (PEP 562): + + from iron.operators.flm import GEMM # imports iron.operators.flm.gemm.op +""" + +import importlib + +_OPERATOR_MODULES = { + # The port, built from source for the current device. + "GEMM": "gemm", + # The shipped overlay itself, downloaded as a pinned binary. NPU2 only; + # exists so the port can be measured against what it was ported from. + "MMPrebuilt": "mm_prebuilt", +} + +__all__ = sorted(_OPERATOR_MODULES) + + +def __getattr__(name): + """Import the operator that defines `name`, on first access.""" + module = _OPERATOR_MODULES.get(name) + if module is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + return getattr(importlib.import_module(f".{module}.op", __name__), name) + + +def __dir__(): + return sorted(set(globals()) | set(_OPERATOR_MODULES)) diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md new file mode 100644 index 0000000000..e185873f72 --- /dev/null +++ b/iron/operators/flm/gemm/README.md @@ -0,0 +1,323 @@ + + +# `iron.operators.flm.GEMM` — bf16 GEMM with a fused epilogue + +```python +from iron.operators.flm import GEMM +from iron.operators.flm.gemm.design import Epilogue + +op = GEMM(M=1024, K=1536, N=6144, epilogue=Epilogue.SILU, context=ctx) +op.compile() +op.get_callable()(A, op.pack_B(B), C_out) +``` + +`epilogue` and `rounding` are `StrEnum`s, so the bare strings `"silu"` / +`"conv_even"` are accepted too. + +A second GEMM implementation alongside [`iron.operators.GEMM`](../../gemm), +specialised for transformer projection shapes and ported from FastFlowLM's `mm` +overlay. + +The overall dataflow is the same whole-array shape as `iron.operators.GEMM`'s — +A broadcast along each compute row, B down each column, C joined through the +memtile — so those are *not* what distinguishes it. What does: + +| | `iron.operators.GEMM` | `flm.GEMM` | +|---|---|---| +| tiling | parameterized tiles, 1–8 columns | fixed m=64 k=512, r/s/t 8/8/8; `n` selectable | +| epilogue | none (separate `convert_copy`) | fused f32→bf16 + activation + clamp | +| B layout | plain `(K, N)` | **pre-packed by the caller**, see below | +| A tile height | tied to the accumulator | decoupled (asymmetric tile buffering) | + +Pick this one for a projection-shaped GEMM that wants an activation folded in +and can pack its weights once. Pick `iron.operators.GEMM` when you need tiling +control or cannot pre-pack B. + +The shipped overlay itself is available as +[`iron.operators.flm.MMPrebuilt`](../mm_prebuilt) for comparison; `benchmark.py` +measures the two against each other and against `iron.operators.GEMM`. + +**Not yet a drop-in replacement for the shipped overlay in FastFlowLM itself.** +FLM's runtime selects matrix shape and activation per call via runtime +parameters (RTPs) on one compiled xclbin. `M`/`K`/`N`/`epilogue`/`rounding` here +are `GEMM(...)` constructor arguments instead -- baked into the MLIR and the +kernel's `-D` flags at compile time (see +[Matching the shipped overlay](#matching-the-shipped-fastflowlm-overlay)) -- so +each shape+epilogue combination is its own compiled kernel object, not one +kernel switchable at runtime. Using this operator inside FLM today means +precompiling and swapping between kernels per combination; making shape and +epilogue RTP-selectable is follow-up work. + +## Architectures + +Runs on both NPU2 (aie2p — Strix/Krackan) and NPU1 (aie2 — Phoenix/Hawk Point). +The tiling and the whole blocked L1 layout are shared; only the grid width and +two lowering details differ. + +| | NPU2 | NPU1 | +|---|---|---| +| grid | 4 x 8 | 4 x 4 | +| A broadcast sources | shim columns 0/2/4/6 | shim columns 0/1/2/3 | +| 8x8x8 mmul lowers to | 2 bfp16-emulated macs | 4 native 4x8x4 bf16 macs | +| `tile_n` default | 128 at K=512, else 64 | always 64 | +| epilogue `tanh` | native `aie::tanh` | `getTanhBf16` LUT | + +The mmul shape is **not** specific to the bfp16 path, despite needing +`AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16` to get the fast lowering on NPU2: +`aie::mmul<8,8,8>` decomposes onto AIE2's native 4x8x4 bf16 mac as exactly four +macs with no wasted lanes, so `pack_B`, the stream-dimension lists and +`gather_dims` are shared verbatim. On AIE2 that flag is silently ignored, so it +is only passed where it changes codegen. + +Two consequences of the native-vs-emulated split are worth knowing: + +* **NPU1 is materially more accurate**, because bfp16 emulation drops mantissa + bits and native bf16 macs accumulating in f32 do not. See + [Accuracy](#accuracy). +* **`Rounding.FLOOR` reproduces the shipped FastFlowLM overlay bit-for-bit on + NPU2 only.** NPU1 sums the K reduction in a different order, so it matches the + rounding *mode* but not the exact results. + +## Shape constraints + +`M % 256 == 0`, `K % 512 == 0`, `N % tile_n == 0` (so 64 by default). + +N only has to tile to `tile_n`, not to the grid's `tile_n * cols` stride: a +trailing group of fewer column-blocks than the grid is wide is handled by giving +the columns different trip counts. That matters in practice — a transformer's +`o` and `down` projections have N = model dim, which is essentially never a +multiple of the full stride. + +## B must be pre-packed + +```python +op = GEMM(M=M, K=K, N=N, context=ctx) +op.compile() +op.get_callable()(A, op.pack_B(B), C_out) +``` + +`pack_B` reorders a row-major `(K, N)` matrix into the order the compute tiles +read it, so each fill is one contiguous run. On NPU2 it also quantizes to +bfp16ebs8 and returns a flat `uint8` tensor rather than bf16; on NPU1 it stays +bf16. The layout itself lives in +[`iron/operators/flm/packing.py`](../packing.py). + +Call it on the operator — `op.pack_B(B)` — not on the class: the layout depends +on the resolved `tile_n` and on the device. + +This is deliberately the caller's job rather than something the fill descriptor +does. The same reorder *is* expressible as a strided descriptor over an unpacked +B, but its innermost run is then `t` bf16 values = 16 bytes, so each 128 KB +transfer becomes thousands of scattered bursts. B is ~70% of the bytes a +dispatch moves, so the operator ran at ~10 GB/s instead of ~47 — a 5.4x +end-to-end penalty. Weights are packed once and reused across dispatches, so the +cost belongs at the caller. + +## Matching the shipped FastFlowLM overlay + +`Rounding.FLOOR` reproduces the shipped `mm.xclbin` **bit for bit**. The AIE +core powers up in `rounding_mode::floor` and the original kernel never calls +`set_rounding`, so that is the arithmetic it ships with. + +```python +GEMM(M=M, K=K, N=N, rounding=Rounding.FLOOR, context=ctx) # matches shipped +GEMM(M=M, K=K, N=N, context=ctx) # conv_even, default +``` + +Verified against the shipped overlay on identical inputs, driven through +[`flm.MMPrebuilt`](../mm_prebuilt), which runs that xclbin unmodified: with +`floor` and no activation, output is **bit-identical across all 6291456 +elements**. With the `conv_even` default it differs everywhere, and is far more +accurate — see [Accuracy](#accuracy). + +**The activations deliberately do not match bit for bit**, even under `floor`. +The overlay rounds its accumulator to bf16 and then applies the activation to +that; this operator applies the activation to the f32 accumulator and rounds +once, on the store. Rounding before a nonlinearity rounds twice and lets the +activation's slope amplify the first rounding, so the overlay's order is the +less accurate one and is not worth reproducing. The cost of diverging is +visible — at M=256 K=512 N=1024, 117582/262144 silu elements differ from the +overlay — and so is the benefit: against an exact f64 evaluation, mean |err| +improves and gelu's worst case drops 5.5%. Measured perf-neutral (0.993-1.006x, +inside the run-to-run spread). + +The shipped kernel selects its activation -- and its shape -- from runtime +parameters, one overlay serving every projection; this operator bakes both in +at compile time instead (activation keeps the shipped 0/1/2/3 mapping), which +is what lets its inner loop be branch-free. See the FLM-compatibility note near +the top of this file for what that means for using this operator inside FLM. + +`clamp` has no counterpart in the shipped overlay to compare against — its +`generate_seq` never writes the clamp RTP words, so clamping is always off +there. + +## Accuracy + +Everything about accuracy is in this section; other sections link here. + +**How to measure it.** A pure elementwise *relative* tolerance is not meaningful: +with signed A the K-term sum cancels by ~sqrt(K), so |C| is ~20x smaller than the +accumulated magnitude while the error tracks that magnitude, leaving near-zero +outputs relatively uncheckable. Bound the error against the accumulated mass +`K * mean|a| * mean|b|` instead, as `test.py` does. + +**It is architecture-dependent**, because the same 8x8x8 mmul lowers differently: +NPU2 emulates it with two bfp16 macs, which drop mantissa bits, while NPU1 uses +four native 4x8x4 bf16 macs, which do not. `test.py` sets the budget per +architecture — inheriting NPU2's on NPU1 would leave ~70x of slack. + +Mean |err| against the accumulated mass, random signed A / non-negative B: + +| | NPU2 | NPU1 | +|---|---|---| +| `Rounding.CONV_EVEN` (default) | 0.000241 | < 1e-6 | +| `Rounding.FLOOR` | 0.009867 | 0.00015 | + +**Prefer the `conv_even` default.** Truncation biases every conversion the same +direction, so the error accumulates over the K reduction instead of cancelling: +~41x more error for no measured speed difference. Use `floor` only to reproduce +the shipped overlay (see [below](#matching-the-shipped-fastflowlm-overlay)). + +On NPU2 this operator and `iron.operators.GEMM` in its comparable mode +(`emulate=True, prio_accuracy=True`) are numerically **indistinguishable** — +identical mean error, signed bias and maximum, at both `tile_n` values. Do not +compare against `iron.operators.GEMM`'s own test tolerances, though: those assert +on its exact `r=4` path, which this operator does not offer. + +## Choosing `tile_n` + +`tile_n` defaults to `None`, which picks per shape and per device: on NPU2, +**128 when `K == 512`, otherwise 64**; on NPU1, **always 64**. Override only if +you have measured a reason to. + +`n=64` gives the mmul `colA=8` rather than 4, halving accumulator traffic per +mac. `n=128` instead halves A fetches, because the grid then covers twice as +many columns of N per pass. Which wins depends on whether compute or data +movement is the critical path, and on NPU2 that turns on how much K there is +to reduce over -- with a single k iteration there is not enough compute to +hide the extra A traffic. + +On NPU2, `tile_n=128` wins only at `k_iters=1`, and by ~3%; at `k_iters>=2` it +is 1.2-1.7x slower. Both follow from `tile_n=128` giving up resident B — its +`mt_b` is 128 KB, so `k_iters` copies do not fit the memtile — and the more k +there is to reduce over, the more that costs. + +NPU1 never reaches that crossover. It has half the columns *and* a quarter of +the per-tile bf16 mac throughput, so it stays compute-bound at every K, and +`n=128` also costs it a much larger f32 accumulator. `n=64` wins at every +`k_iters`, by 1.2-1.8x. + +`benchmark.py` measures both settings across the full shape sweep; prefer its +CSV output to any table reproduced here. + +## Performance + +### NPU2 + +M=1024 K=1536 N=6144, min of per-run medians: + +| | bytes moved | latency | err/mass | +|---|---|---|---| +| `flm.GEMM` (`tile_n=64`) | 47 MB | **1143 us** | 2.39e-04 | +| `flm.MMPrebuilt` (the shipped overlay) | 107 MB | 2175 us | 9.87e-03 | +| `iron.operators.GEMM` (same emulated mode) | 126 MB | 3353 us | 2.41e-04 | + +**1.90x the shipped overlay, and 41x more accurate than it** — the accuracy +comes from `conv_even` rounding, which the overlay does not set (see above). +`benchmark.py` reproduces this table, and covers 30 shapes rather than one. + +Three choices account for most of the gap, and none of them helps alone: + +* **`pack_B` emits the final consumption order**, so both B hops are linear + descriptors. Worth nothing by itself — it is what frees the descriptor + dimensions the other two need. +* **Asymmetric tile buffering**, which pays for a k slice deep enough to halve + the accumulator traffic per mac. +* **A rolled mmul inner loop**, which is faster than hand-unrolling it here + (see `mm_fused_mmul.h`). + +Storing B in bfp16 is numerically free: the NPU2 mmul only multiplies bfp16, so +quantizing on the host hoists a rounding that already happened on every mac +call. It does have to reproduce the core's rounding *mode* to be free — see +`iron/operators/flm/packing.py`. + +> **Measuring this.** Dispatch latency on this part is *bimodal*, with modes +> about 6% apart, and both show up for every configuration. A batch that lands +> wholly in one mode turns min-of-medians into a mode selector rather than a +> measurement. Compare configurations **interleaved** round-robin rather than +> one after the other, use at least 8 rounds each, and believe a difference only +> when the min and the median agree on it. `benchmark.py` does this. + +### NPU1 + +Against `iron.operators.GEMM` at its own defaults (64/64/64 over all 4 +columns), min of 5 interleaved rounds of 20 dispatches each: + +| M / K / N | `flm.GEMM` | `iron.operators.GEMM` | speedup | GFLOP/s | +|---|---|---|---|---| +| 256 / 512 / 512 | **218 us** | 227 us | 1.04x | 615 | +| 512 / 512 / 1024 | **390 us** | 466 us | 1.19x | 1377 | +| 512 / 1024 / 1024 | **639 us** | 798 us | 1.25x | 1680 | +| 1024 / 1024 / 1024 | **1071 us** | 1418 us | 1.32x | 2005 | +| 1024 / 2048 / 1024 | **2134 us** | 2738 us | 1.28x | 2013 | +| 512 / 1536 / 1536 | **1498 us** | 1699 us | 1.13x | 1613 | +| 1024 / 2560 / 2560 | **6364 us** | 8232 us | 1.29x | 2109 | + +`K=1536` is the weak shape, at 1.13x against 1.25-1.32x for its neighbours. +It is the one place the L1 configuration below does not suit NPU1: `k_iters=3` +leaves the k slice partly unused. Worth a look if NPU1 throughput matters. + +The margin is smaller than NPU2's ~1.9x, and that is expected rather than a +port problem: much of the NPU2 win comes from the bfp16 fast path (two macs per +8x8x8 shape against four) and from spreading A across eight columns. On NPU1 +both operators lower to the same native 4x8x4 mac, so what remains is the +cheaper transfers below — pre-packed B and one transfer per column-block — +which is why the gap grows with the problem size rather than being flat. + +Unlike NPU2, compute here is **not** hidden behind the transfers, so on NPU1 +both a faster mmul and less traffic pay off, where on NPU2 only the latter does. +Resident B is also a no-op on NPU1 — see [Resident B](#resident-b). + +### Why the transfers are cheap + +Two things, both in the runtime sequence rather than the kernel: + +* **B arrives pre-packed**, so the contiguous run per transfer is 128 KB for B + and 1 KB for A, against 128 bytes on every leg for `iron.operators.GEMM`, + which reorders in the descriptor instead. +* **Each of A, B and C goes out as one transfer per column-block**, not one per + fifo object. A single fill or drain may span many objects; issuing per object + means a host await per row-block, and a C await waits on the cores. + Collapsing them is also what makes overlapping column-blocks affordable — a + block then costs 3 shim buffer descriptors instead of `1 + 2*k_iters`, so two + can be in flight without exhausting the 16 available. + +### Resident B + +Where a whole column-block's B fits in the memtile double-buffered +(`k_iters <= 2`, i.e. K <= 1024 at `tile_n=64`) it is held there and replayed +per row-block, so DDR reads it once instead of `m_row_blocks` times -- about +43% less traffic. Larger K falls back to re-reading it, unchanged. + +On NPU2 this is a latency win as well as a power one, because there the +operator is close to DDR-bandwidth bound, and it grows with the height of the +problem since B's re-reads scale with `m_row_blocks`. At K=1024 N=4096, min of +per-round medians over 10 interleaved rounds of 20 dispatches, `npu_time`, +power mode `turbo`: + +| M | row-blocks | non-resident | resident | | +|---|---|---|---|---| +| 512 | 2 | 527.0 us | **461.2 us** | 12.5% | +| 1024 | 4 | 1025.8 us | **857.3 us** | 16.4% | +| 2048 | 8 | 1958.0 us | **1579.5 us** | 19.3% | + +Most of the available win is still on the table: `repeat_count` restarts the +memtile BD chain at every replay boundary, which costs part of the traffic +saving back. Closing that is the largest known remaining lever here. + +On NPU1 residency is neither a latency nor a power win — measured off-versus-on +at M=2048 it is at best a no-op and marginally negative at K=1024, well inside +the 1.4-4.4% round spread. diff --git a/iron/operators/flm/gemm/benchmark.py b/iron/operators/flm/gemm/benchmark.py new file mode 100644 index 0000000000..00279cf817 --- /dev/null +++ b/iron/operators/flm/gemm/benchmark.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare ``flm.GEMM`` against the overlay it was ported from and IRON's GEMM. + +Up to three implementations run per shape, on identical inputs: + + flm :class:`iron.operators.flm.GEMM`, the port + gemm :class:`iron.operators.GEMM`, left at its defaults, which are the + same emulated-bfp16 mmul and conv_even rounding -- a like-for-like + comparison rather than one against a more accurate, slower build + prebuilt :class:`iron.operators.flm.MMPrebuilt`, FastFlowLM's shipped + ``mm.xclbin``, downloaded and pinned by digest. NPU2 only, because + that binary is a fixed 8-column NPU2 overlay -- on any other device + it is dropped and the flm-vs-gemm comparison still runs. + +Nothing here needs an external install or a host-specific path: the overlay is +a ``RemoteFileArtifact``, so it is fetched into the (gitignored) build dir like +any other artifact, and its instruction stream is generated by IRON. The binary +is never checked in -- it is pinned by SHA-256 against an immutable FastFlowLM +commit and downloaded on demand. + +Marked ``extensive`` for documentation, but pytest never collects it either way: +``pytest.ini`` sets ``python_files = test.py``, and this file is named +``benchmark.py`` so no CI job -- extensive or otherwise -- runs it. That is +deliberate: it is a timing comparison, meant to be invoked directly, not a +correctness gate. The correctness half now lives in +``iron/operators/flm/mm_prebuilt/test.py``, which IS a collected ``test.py`` +and so IS reached by the extensive job. + +Timing is the runtime's own ``npu_time`` (device-side), the same source +``iron.common.test_utils.run_test`` reports, rather than a host wall clock: it +excludes host dispatch and so compares the designs rather than the driver. + +The shapes below are the ones this operator exists to serve, so they overlap +with ``test.py``'s by construction. They are not redundant with it: ``test.py`` +asserts correctness on one implementation, this compares latency across three +frozen binaries, and neither can stand in for the other. + +Usage:: + + pytest iron/operators/flm/gemm/benchmark.py --no-short + pytest iron/operators/flm/gemm/benchmark.py --no-short -k E2B --csv-output flm.csv +""" + +import statistics +from pathlib import Path + +import numpy as np +import pytest +import torch + +import aie.utils as aie_utils +from aie.utils.hostruntime.xrtruntime.tensor import XRTTensor + +from iron.operators import GEMM as IronGEMM +from iron.operators.flm import GEMM as FLMGEMM +from iron.operators.flm import MMPrebuilt + +# Opt-in only: this module downloads the overlay, so keep it out of the default +# run. See the note in the module docstring. +pytestmark = pytest.mark.extensive + +_dev = aie_utils.get_current_device() +# The shipped overlay is a fixed 8-column NPU2 binary. Where that does not +# match the device, drop that ONE candidate rather than skipping the module: +# flm vs iron.operators.GEMM is measurable on every supported device. +HAVE_PREBUILT = _dev is not None and _dev.resolve().name == "npu2" and _dev.cols >= 8 + +# Every projection of both Gemma4 variants FastFlowLM ships, at three prefill +# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. +# +# These two are the right coverage for the shipped mm.xclbin. Checked against +# FastFlowLM f81eba71: Gemma4-E2B-IT-NPU2, Gemma4-E4B-IT-NPU2 and Gemma3-4B-NPU2 +# all ship the SAME mm.xclbin blob (git 4727df98, 512220 bytes) -- one overlay +# serving several models -- so E2B and E4B between them already exercise it. +# +# Gemma4-12B-IT-NPU2 does not ship an mm.xclbin at all. Its overlays are a +# different set (attn_global, attn_sliding, audio_image_mm, dequant_mm, layer, +# lm_head), so its projections go through a quantized matmul rather than this +# bf16 one and cannot be compared against MMPrebuilt. +# proj, K, N +E2B_PROJ = [ + ("q", 1536, 4096), + ("kv", 1536, 512), + ("o", 4096, 1536), + ("gateup", 1536, 6144), + ("down", 6144, 1536), +] +E4B_PROJ = [ + ("q", 2560, 4096), + ("kv", 2560, 1024), + ("o", 4096, 2560), + ("gateup", 2560, 10240), + ("down", 10240, 2560), +] +PREFILL_LENGTHS = [256, 1024, 2048] + +# Interleaved rounds per test, and timed dispatches per implementation per +# round. 6 rounds is the floor at which min and median stopped disagreeing. +ROUNDS = 6 +ITERS = 30 +WARMUP = 20 + +# err/mass budgets. flm.GEMM and IRON's GEMM both round conv_even; the shipped +# overlay never calls set_rounding, so it runs in the core's power-up floor mode +# and carries a ~1% truncation bias that is not a bug to fix here. +BUDGET_CONV_EVEN = 4e-3 +BUDGET_FLOOR = 2e-2 + + +def get_params(): + # No shape is skipped: the four E4B projections with a 10240-wide + # dimension at M > 256 once overflowed the shim BD's 20-bit mega_row + # iteration step, but flm.GEMM and IRON's GEMM both now split that leg + # into per-mega_row transfers (see design.py's a_split/c_split and + # test_gemm_split_leg_windowing in test.py). + params = [] + for model, projections in (("E2B", E2B_PROJ), ("E4B", E4B_PROJ)): + for M in PREFILL_LENGTHS: + for proj, K, N in projections: + params.append( + pytest.param(model, proj, M, K, N, id=f"{model}-{proj}-M{M}") + ) + return params + + +def make_inputs(M, K, N): + """Identical data for all three, and the reference to check them against.""" + torch.manual_seed(1234) + A = (torch.randn(M, K) * 4).to(torch.bfloat16) + B = (torch.rand(K, N) * 4).to(torch.bfloat16) + Af, Bf = A.float(), B.float() + # Error is bounded against accumulated mass rather than relatively: with + # signed A the K-sum cancels by ~sqrt(K), so |C| ends up far smaller than + # the magnitude the bfp16 error actually tracks, leaving near-zero outputs + # relatively uncheckable. Same rationale as test.py's bound. + return A, B, Af @ Bf, float((Af.abs() @ Bf.abs()).mean()) + + +class Candidate: + """One implementation under test, with its buffers already bound so the + timed section contains nothing but the dispatch.""" + + def __init__(self, name, op, A, B, M, N, budget, ctx): + self.name = name + self.budget = budget + self.round_medians = [] + + op.compile() + self.xclbin = Path(op.xclbin_artifact.filename) + self.c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) + run = op.get_callable() + # Only the flm operators take B pre-packed. iron.operators.GEMM reorders + # in the descriptor instead, so it wants plain row-major (K, N) -- + # b_col_maj defaults False. + packed_b = op.pack_B(B) if hasattr(op, "pack_B") else B + args = [ + XRTTensor.from_torch(A.flatten()), + XRTTensor.from_torch(packed_b.flatten()), + self.c_bo, + ] + self.run = lambda: run(*args) + + def verify(self, M, N, expected, mass): + self.run() + C = self.c_bo.to_torch().reshape(M, N).float() + self.err = float((C - expected).abs().mean()) / mass + return self.err < self.budget + + def time_round(self): + # npu_time is the device-side execution time the runtime reports, in ns. + ts = [self.run().npu_time / 1e3 for _ in range(ITERS)] + self.round_medians.append(statistics.median(ts)) + + @property + def us(self): + # Minimum of the per-round medians: the median rejects the tail within + # a round, the min rejects rounds that landed in the slow mode. + return min(self.round_medians) + + @property + def jitter_pct(self): + """Spread of the per-round medians -- how bimodal this run actually was.""" + return (max(self.round_medians) - self.us) / self.us * 100.0 + + +@pytest.mark.metrics( + FLMLatency=r"flm latency \(us\): (?P[\d\.]+)", + PrebuiltLatency=r"prebuilt latency \(us\): (?P[\d\.]+)", + GEMMLatency=r"gemm latency \(us\): (?P[\d\.]+)", + SpeedupVsPrebuilt=r"speedup vs prebuilt: (?P[\d\.]+)", + SpeedupVsGEMM=r"speedup vs gemm: (?P[\d\.]+)", + # Accuracy is asserted against a budget below, but that budget is loose + # enough that a toolchain or kernel change could move the error a long way + # inside it unnoticed. Record the numbers too, so a dependency bump can be + # diffed on accuracy and not only on speed. + FLMErr=r"flm err/mass: (?P[\d\.e\+-]+)", + PrebuiltErr=r"prebuilt err/mass: (?P[\d\.e\+-]+)", + GEMMErr=r"gemm err/mass: (?P[\d\.e\+-]+)", + FLMThroughput=r"flm throughput: (?P[\d\.e\+-]+) GFLOP/s", + FLMJitterPct=r"flm jitter \(%\): (?P[\d\.]+)", + FLMXclbinKB=r"flm xclbin \(KB\): (?P[\d\.]+)", + GEMMXclbinKB=r"gemm xclbin \(KB\): (?P[\d\.]+)", +) +@pytest.mark.parametrize("model,proj,M,K,N", get_params()) +def test_gemm_vs_prebuilt(model, proj, M, K, N, aie_context): + A, B, expected, mass = make_inputs(M, K, N) + + # Build everything before timing anything. Comparing frozen binaries is the + # only way an A/B here means what it says. + candidates = [ + Candidate( + "flm", + FLMGEMM(M=M, K=K, N=N, context=aie_context), + A, + B, + M, + N, + BUDGET_CONV_EVEN, + aie_context, + ), + Candidate( + "gemm", + IronGEMM(M=M, K=K, N=N, context=aie_context), + A, + B, + M, + N, + BUDGET_CONV_EVEN, + aie_context, + ), + ] + if HAVE_PREBUILT: + candidates.append( + Candidate( + "prebuilt", + MMPrebuilt(M=M, K=K, N=N, context=aie_context), + A, + B, + M, + N, + BUDGET_FLOOR, + aie_context, + ) + ) + + bad = [c for c in candidates if not c.verify(M, N, expected, mass)] + assert not bad, "; ".join( + f"{c.name} err/mass {c.err:.3g} exceeds {c.budget:g}" for c in bad + ) + + for c in candidates: + for _ in range(WARMUP): + c.run() + # Round-robin, never all of one then all of another. Dispatch latency on + # this part is bimodal with modes about 6% apart, so a batch that lands + # wholly in one mode turns min-of-medians into a mode selector rather than a + # measurement -- that is how a change later shown to do nothing at all once + # produced a convincing 5% "win". + for _ in range(ROUNDS): + for c in candidates: + c.time_round() + + by_name = {c.name: c for c in candidates} + flm = by_name["flm"] + + print() + for c in candidates: + print(f"{c.name} latency (us): {c.us:.1f}") + print(f"{c.name} err/mass: {c.err:.3e}") + print(f"{c.name} xclbin (KB): {c.xclbin.stat().st_size / 1024:.1f}") + if "prebuilt" in by_name: + print(f"speedup vs prebuilt: {by_name['prebuilt'].us / flm.us:.3f}") + print(f"speedup vs gemm: {by_name['gemm'].us / flm.us:.3f}") + print(f"flm throughput: {2.0 * M * K * N / (flm.us * 1e-6) / 1e9:.6e} GFLOP/s") + print(f"flm jitter (%): {flm.jitter_pct:.2f}") + print() diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py new file mode 100644 index 0000000000..864e55f625 --- /dev/null +++ b/iron/operators/flm/gemm/design.py @@ -0,0 +1,938 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""bf16 GEMM over a 4-row compute-tile grid, as wide as the device. + +A second GEMM design alongside ``iron.operators.gemm``, specialised for +transformer projection shapes. The overall dataflow is the same whole-array +shape as that operator's -- A broadcast along each compute row, B down each +column, C joined through the memtile -- so those are NOT what distinguishes it. +What does: + + * **B is quantized to bfp16ebs8 on NPU2** by ``GEMM.pack_B``, not bf16. + ``iron.operators.GEMM`` only ever moves bf16. This is not primarily a DMA + saving: it is what makes NPU2's fast mmul lowering available at all -- + ``aie::mmul<8,8,8>`` needs bfp16 operands to decompose into two emulated + macs instead of four, which is most of the NPU2 speedup (see NPU1 in + README.md's Performance section, where B stays bf16 and the margin over + ``iron.operators.GEMM`` is correspondingly smaller). Quantizing is + numerically free -- the mmul only multiplies bfp16 regardless, so this + hoists a rounding that already happened on every mac -- provided it + reproduces the core's rounding mode; see ``packing.py``. + * **The tile shape is fixed, not parameterised** -- but fixedness alone is + not the advantage: ``iron.operators.GEMM`` is equally fixed once compiled + with a choice of tile args. What differs is *which* shape is fixed. r/s/t + stays 8/8/8 on both architectures for the reason below. m/k = 64/512 (n + defaults to 64) is chosen for the L1-budget tradeoff documented next to + ``CT_MAX_K_FOR_N`` below: n=64 gives the mmul a colA of 8 rather than 4, + which wins whenever compute is the critical path, at the cost of A being + re-read more often. Only the grid WIDTH varies with the device: 8 columns + on NPU2, 4 on NPU1. + * **A fused epilogue.** The f32->bf16 conversion, an optional activation and + an optional clamp all happen while the values are still in registers, on the + way into the C object, instead of a separate pass over L1. + * **B arrives pre-packed** by ``GEMM.pack_B``, in the order the cores consume + it, so both B hops are plain linear descriptors instead of the 128-byte + scattered bursts ``iron.operators.GEMM`` reorders in the descriptor. + * **Asymmetric tile buffering (ATB)**, so the A tile and the accumulator need + not share a height -- this is what buys the deep k slice (K_TILE=512) + within the L1 budget; see README.md's ATB reference. + +None of these four helps alone -- see README.md's Performance section for the +measured, per-choice breakdown of the gap against both the shipped FastFlowLM +overlay and ``iron.operators.GEMM``. + +The constants below are the single source of truth: ``op.py`` passes them to the +kernels as -D flags, so the C++ and the dataflow cannot drift apart. + +r/s/t stays 8/8/8 on both architectures. AIE2's native bf16 mac is 4x8x4, but +``aie::mmul<8,8,8>`` decomposes onto it as exactly four native macs with no +wasted lanes, so the whole blocked L1 layout -- ``pack_B``, the four stream +dimension lists below, and ``gather_dims`` -- is shared verbatim. Only AIE2P has +the bfp16-emulated path that does the same shape in two macs, which is why +``op.py`` passes ``AIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16`` there and not here. +""" + +import argparse +from enum import StrEnum +from functools import partial + +import numpy as np +from ml_dtypes import bfloat16 + +from aie.helpers.util import v8bfp16ebs8 + +from aie.helpers.taplib import TensorAccessPattern +from aie.iron import ( + Buffer, + Kernel, + ObjectFifo, + Program, + Runtime, + TaskGroup, + Worker, +) +from aie.iron.controlflow import range_ +from aie.dialects.aie import get_target_model +from aie.dialects._aie_enum_gen import AIEArch +from aie.iron.device import NPU1, NPU2, Tile +from iron.common.utils import split_run +from iron.operators._trace import maybe_enable_trace + +# --- Fixed geometry ------------------------------------------------------- +# GEMM tiling per compute tile, and the register tiling inside it. +M_TILE, K_TILE = 64, 512 +# Default n tile. 64 gives the mmul a colA of 8 rather than 4, halving the +# accumulator traffic per mac, at the cost of doubling A fetches (the grid +# then covers 512 columns of N per pass instead of 1024). That trade wins +# whenever compute is the critical path, which is the usual case; see +# README.md for the measured sweep, including the small-K shape where it +# loses. +N_TILE_DEFAULT = 64 +# How much of K one compute tile holds at a time, per n width. This is a fixed +# L1 budget split two ways, so a wider n tile leaves less room for B's k slice +# and the product stays roughly constant. op.py passes the chosen value to the +# kernel as -DMM_FUSED_CT_K, making this table the only place it is decided. +# n=256 is deliberately absent: at that width the f32 accumulator alone +# (M_TILE * 256 * 4 = 65536 bytes) already fills the whole of L1, before A, B +# or C are even counted, so no ct_max_k could ever make it fit. +CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32} +# Register tiling, shared by both architectures. These set the blocked L1 +# layout, so ``pack_B``, the four stream-dimension lists below and +# ``gather_dims`` all key off them; changing one without the others is silently +# wrong rather than a build error. +# +# Matching AIE2's native 4x8x4 mac shape instead is a measured dead end, at +# 22-30% slower: the kernel is load-port bound rather than shuffle bound, and +# 4/8/4 needs 1.62 loads per mac against 8/8/8's 1.06, because the 2x2 register +# block amortizes each load over four macs either way but over far less work. +# Retrying it needs a wider register block on the native shape, i.e. a 4x4 mmul +# kernel, not just a different r/s/t here. +R, S, T = 8, 8, 8 + + +def compute_rows(dev): + """Compute-tile rows: the array less the shim row and the memtile rows.""" + tm = get_target_model(dev.resolve()) + return tm.rows() - 1 - tm.get_num_mem_tile_rows() + + +CT_OUT_LEN = 512 # the core's C slice, streamed out in chunks this size +C_DEPTH = 2 # C fifo depth; also the core-body unroll +B_DEPTH = 2 # B fifo depth; also the core-body unroll +A_DEPTH = 2 +# How many column-blocks the runtime sequence keeps in flight. A block costs 3 +# shim buffer descriptors on a column (A + B + C) against 16 available, so the +# ceiling is 5; 2 is enough to keep the fills ahead of the cores. +OVERLAP_DEFAULT = 2 + + +class Epilogue(StrEnum): + """Activation folded into the C drain. + + Declaration order is the wire format -- it is both the kernel's + ``-DMM_FUSED_EPILOGUE_MODE`` and the shipped overlay's ``output_mode``. + """ + + NONE = "none" + GELU = "gelu" + SILU = "silu" + SIGMOID = "sigmoid" + + @property + def mode(self) -> int: + """The integer the kernel and the shipped overlay both select on.""" + return list(Epilogue).index(self) + + +class Rounding(StrEnum): + """Rounding for every f32->bf16 conversion. + + The core powers up in floor; conv_even is the default because truncation + biases every conversion the same way and the error then accumulates over + the K reduction. floor reproduces the shipped overlay. + """ + + CONV_EVEN = "conv_even" + FLOOR = "floor" + + +# The epilogue entry point, shared by the design and op.py (which needs it +# to mark the symbol alwaysinline when building the inline .ll variant). +EPILOGUE_SYMBOL = "mm_fused_epilogue_chunk" + +# Minimum problem size in K. The minimum in M is M_TILE * compute_rows(dev) and +# in N is the chosen n tile, both of which depend on the device or the config. +MIN_K = K_TILE # 512 + + +# B values per element of the MLIR type, and the bytes they occupy: v8bfp16ebs8 +# packs 8 values into 8 mantissa bytes plus one shared exponent. mlir-aie +# exposes no width query on the type, hence the literals. +BFP16_GROUP, BFP16_GROUP_BYTES = 8, 9 + + +def _b_bytes(elems, bfp16_b): + """Bytes B occupies in L1/L2. bfp16ebs8 packs 8 values as 8 mantissa bytes + plus one shared exponent; bf16 is a plain 2 bytes each.""" + if not bfp16_b: + return elems * 2 + assert elems % BFP16_GROUP == 0 + return elems // BFP16_GROUP * BFP16_GROUP_BYTES + + +# --- Shim DMA limits ------------------------------------------------------ +# +# Hardware facts the Python bindings do not expose: gemm() reads AIETargetModel +# directly for the L1 ceiling, BD count and grid, but neither getDmaBdStepBits +# nor getDmaBdWrapSizeBits is bound, and nothing models the channel task queue. +# gemv/design.py and repeat/design.py hardcode the same fields. +# +# Step field width. An IR-level bf16-element stride S is re-expressed as +# (S - 1) * 2 bytes / 4-byte granularity before AIEXDialect.cpp checks it. +_SHIM_STEP_BITS = 20 +_BF16_BYTES = 2 +_ADDR_GRANULARITY_BYTES = 4 +# Entries in a shim DMA channel's task queue. AIEDmaToNpu's NpuPushQueueOp +# pushes unconditionally, so overrunning this is a silent device hang rather +# than a diagnostic. Measured at K=10240 M=1024: 4 outstanding tasks on one +# channel run, 8 hang. +SHIM_TASK_QUEUE = 4 + + +def _hw_stride_ok(stride_elems): + hw_stride = (stride_elems - 1) * _BF16_BYTES // _ADDR_GRANULARITY_BYTES + return hw_stride <= (1 << _SHIM_STEP_BITS) - 1 + + +def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): + """Pick (A-tile height, L1 B depth) -- the largest working set that fits. + + ``b_elem_bytes`` is 9/8 where B is bfp16ebs8 and 2 where it is bf16, and + ``budget`` is the core's data memory, so the search below reflects what B + actually costs on this device. + + No stack is reserved out of ``budget``: the cores leave ``stack_size`` + unset and aiecc measures each core's requirement and fails the build if it + does not fit, so the stack is the toolchain's to enforce. This kernel + measures 192 bytes against the >=6 KB the search leaves unused anyway. + + A dies as soon as it is consumed while the accumulator lives across the + whole K reduction, so they need not share a height; shrinking A is what + pays for a k slice deep enough to halve the accumulator traffic per mac. + B's depth is searched too because at n=128 the k=128 slice makes the B + object 18 KB, and a double-buffered pair simply does not fit -- giving that + up is what buys colA=16 there, and colA is worth far more than B's L1 + prefetch (3.67 -> 2.28 cycles per mac, measured). + + Deeper B first, then the tallest A that still fits, so the n=64 default is + unchanged at (32, 2). + """ + acc = M_TILE * n_tile * 4 + cout = CT_OUT_LEN * 2 * C_DEPTH + for b_depth in (B_DEPTH, 1): + b = int(ct_max_k * n_tile * b_elem_bytes) * b_depth + for t_ma in (M_TILE, M_TILE // 2, M_TILE // 4): + if t_ma < 2 * R: + continue + a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH + if acc + a + b + cout <= budget: + return t_ma, b_depth + raise ValueError(f"nothing fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}") + + +def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget): + """Deepest B fifo depth that fits L1 alongside an explicit A-tile height. + + ``_default_l1`` picks L1_B_DEPTH together with the t_ma IT chooses; that + pairing need not fit a caller-overridden t_ma; a taller A tile leaves less + L1 for B, and can push a working set that fit at the default t_ma over + budget. Raise rather than silently reusing a depth that doesn't fit. + + How much the depth is worth, measured on npu2 (turbo, 12 interleaved rounds + of 20 dispatches, min of per-round medians) by forcing depth 1 against the + default: 1.2% at M=1024 K=1536 N=6144, and within noise at K=1024 N=4096 and + K=512 N=1024. So the prefetch earns its L1 at the largest shapes and is + close to free elsewhere -- worth keeping, but not worth contorting the + search for. + """ + acc = M_TILE * n_tile * 4 + cout = CT_OUT_LEN * 2 * C_DEPTH + a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH + for b_depth in (B_DEPTH, 1): + b = int(ct_max_k * n_tile * b_elem_bytes) * b_depth + if acc + a + b + cout <= budget: + return b_depth + raise ValueError( + f"tile_ma={t_ma} does not fit L1 for tile_n={n_tile} " + f"(ct_max_k={ct_max_k}); even single-buffered B overflows the budget" + ) + + +def gemm( + dev, + M, + K, + N, + epilogue=Epilogue.NONE, + tile_n=N_TILE_DEFAULT, + tile_ma=None, + overlap=None, + kernel_object="mm_fused.o", + trace_size=0, +): + """Emit the MLIR module for an M x K @ K x N bf16 GEMM. + + A is (M, K) row-major, B is (K, N) row-major and C is (M, N) row-major, all + bf16 and all plain dense tensors, except that B must arrive pre-packed by + ``GEMM.pack_B`` -- it emits B in the order the cores consume it, so both + B hops are plain linear descriptors. + """ + if tile_n not in CT_MAX_K_FOR_N: + raise ValueError( + f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {tile_n}" + ) + # Everything shape-related below comes from the device rather than a + # constant, so the same dataflow covers NPU2's 4x8 and NPU1's 4x4. + tm = get_target_model(dev.resolve()) + COLS, ROWS = dev.cols, compute_rows(dev) + MIN_M = M_TILE * ROWS + SHIM_BDS = tm.get_num_bds(0, 0) + N_TILE = tile_n + CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + # B is bfp16ebs8 on AIE2P and bf16 on AIE2: the scalar BFP types are gated + # on __AIE_API_SCALAR_BFP_TYPES__, which only aie_api/detail/aie2p/config.hpp + # defines, so on AIE2 B stays bf16 and the mmul lowers onto four native + # 4x8x4 macs. That choice drives every B type and extent below. B_GROUP is + # the number of B values per element of the MLIR type, so a length in values + # becomes a length in elements by dividing. + BFP16_B = dev.arch == AIEArch.AIE2p + B_GROUP = BFP16_GROUP if BFP16_B else 1 + b_elem_bytes = BFP16_GROUP_BYTES / BFP16_GROUP if BFP16_B else 2 + # Asymmetric tile buffering: the A tile spans T_MA rows while the + # accumulator spans M_TILE, so the core folds RHO bands into one C tile. + # A is dead the moment it is consumed while C lives across the whole K + # reduction, so sizing both to M_TILE pays the peak L1 cost twice. + if tile_ma is None: + T_MA, L1_B_DEPTH = _default_l1( + N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + ) + else: + T_MA = tile_ma + if M_TILE % T_MA or T_MA % (2 * R): + raise ValueError( + f"tile_ma ({T_MA}) must divide {M_TILE} and be a multiple of {2 * R}" + ) + L1_B_DEPTH = _b_depth_for( + T_MA, N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + ) + RHO = M_TILE // T_MA + OVERLAP = OVERLAP_DEFAULT if overlap is None else overlap + K_DIV_CT_K_MAX = K_TILE // CT_MAX_K + CT_A_LEN = 2 * R * CT_MAX_K # one z slice + CT_A_OBJ = CT_A_LEN * (T_MA // R // 2) # every z slice of one mmul + C_SLICE_LEN = M_TILE * N_TILE # one compute tile's C contribution + O_CHUNKS = C_SLICE_LEN // CT_OUT_LEN # C objects an accumulator drains as + B_ITERS = K_TILE // CT_MAX_K # B chunks consumed per k step + MIN_N = N_TILE * COLS + + epilogue = Epilogue(epilogue) + # A compute tile does a whole m x n block or nothing, so M and K must tile + # exactly. N need only be a multiple of N_TILE: a trailing group of fewer + # than COLS blocks is handled by giving the columns different trip counts + # (see col_work / col_drain below). That matters in practice -- for a + # transformer the o and down projections have N = model dim, which is + # essentially never a multiple of N_TILE*COLS. + for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, N_TILE)): + if value % unit != 0: + raise ValueError(f"{name} ({value}) must be a multiple of {unit}") + + bf16_ty = np.dtype[bfloat16] + f32 = np.dtype[np.float32] + + # How many times the whole grid sweeps, in each dimension. + m_row_blocks = M // MIN_M + k_iters = K // K_TILE + # A mega_row dimension with stride ROWS*M_TILE*{K,N} lands in the shim + # BD's ITERATION field, whose step is 20 bits, so it overflows once K or N + # crosses ~8191 elements -- E4B's FFN width (10240) does, E2B's max (6144) + # does not. Only relevant once M walks more than one mega_row; at M=256 the + # dimension is degenerate (size 1) and is stripped before the stride is + # ever encoded, so it never fails there regardless of K/N. + # + # Such a leg is instead issued as m_row_blocks separate transfers, each + # carrying the mega_row jump in its OFFSET (unbounded) rather than a shared + # STRIDE. The two legs are independent -- E4B's down-proj overflows on K + # (A only) and its gate/up on N (C only) -- so neither shape pays for both. + # + # Those transfers must stay live in their TaskGroup until awaited. + # TaskGroup.finish() emits dma_free_task, which returns the buffer + # descriptor id to a COMPILE-TIME allocator that does not check the + # transfer finished (mlir-aie AIEAssignRuntimeSequenceBDIDs::recycle, + # isAwait=false); ids are per shim TILE, shared across channels and + # directions, so retiring one early lets the next task reprogram a live + # descriptor. Verify with aie-opt --aie-substitute-shim-dma-allocations + # --aie-assign-runtime-sequence-bd-ids: the ids on a shim tile must be + # distinct. + a_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * K) + c_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * N) + # A split leg issues one transfer per mega_row back to back on ONE channel, + # so they must also fit that channel's task queue -- a limit nothing in the + # toolchain models, and overrunning it hangs rather than diagnoses. + # Measured at K=10240 M=1024: 4 outstanding run, 8 hang. + # + # So a split block is emitted in WINDOWS of at most SHIM_TASK_QUEUE + # mega_rows, each window awaited before the next is issued (see sequence()). + # Awaiting is what makes the window's descriptors safe to reuse, and it + # bounds both resources at once. M<=1024 is a single window, so the shapes + # that already worked are unaffected. + MB_WINDOW = min(m_row_blocks, SHIM_TASK_QUEUE) if (a_split or c_split) else 1 + bds_per_block = 1 + (MB_WINDOW if a_split else 1) + (MB_WINDOW if c_split else 1) + # Unreachable while SHIM_TASK_QUEUE is 4 (the worst case is 1 + 4 + 4 = 9 + # of 16), so this guards a future retune of the window rather than any + # shape reachable today. test_flm_gemm_split_leg_windowing asserts the same + # arithmetic from the outside. + if bds_per_block > SHIM_BDS: + raise ValueError( + f"M={M} K={K} N={N} needs {bds_per_block} shim buffer descriptors " + f"per window (1 B + {MB_WINDOW if a_split else 1} A + " + f"{MB_WINDOW if c_split else 1} C) but a shim tile has only " + f"{SHIM_BDS}." + ) + # Cross-block overlap only applies to the unsplit path; a split block + # already awaits inside itself, so keeping a second one in flight would + # refill the very queue the windowing just drained. + if a_split or c_split: + OVERLAP = 1 + else: + OVERLAP = max(1, min(OVERLAP, SHIM_BDS // bds_per_block)) + # Sweeps where all COLS columns have work, plus a trailing group of + # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. + n_full = N // MIN_N + rem_blocks = (N % MIN_N) // N_TILE + # Columns that participate at all. With no full sweep (N below the grid's + # COLS*N_TILE stride) only the first rem_blocks columns do, and the rest + # are not instantiated -- giving them fifos that nothing ever drains builds + # dead dataflow, which newer mlir-aie rejects outright with + # "objectfifo.pool op segment 0 has no drainer". + n_active_cols = COLS if n_full else rem_blocks + # Per column: how many column-blocks it computes, and whether it sits out a + # trailing one while still draining the A broadcast for its row. That drain + # only arises for a column that exists and skips the trailing block, which + # requires at least one full sweep. + col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(n_active_cols)] + col_drain = [ + 1 if (rem_blocks and n_full and c >= rem_blocks) else 0 + for c in range(n_active_cols) + ] + + # B's element type: one v8bfp16ebs8 per 8 values on AIE2P, one bf16 per + # value on AIE2. Every B extent below is therefore in values // B_GROUP. + b_elem_ty = np.dtype[v8bfp16ebs8] if BFP16_B else bf16_ty + # L1 (per compute tile) + ct_a_obj_ty = np.ndarray[(CT_A_OBJ,), bf16_ty] + ct_b_ty = np.ndarray[(CT_MAX_K * N_TILE // B_GROUP,), b_elem_ty] + ct_out_ty = np.ndarray[(CT_OUT_LEN,), bf16_ty] + ct_acc_ty = np.ndarray[(M_TILE * N_TILE,), f32] + # L2 (per memtile) + mt_a_ty = np.ndarray[(M_TILE * K_TILE,), bf16_ty] + mt_a_bytes = M_TILE * K_TILE * 2 + mt_b_ty = np.ndarray[(K_TILE * N_TILE // B_GROUP,), b_elem_ty] + mt_b_bytes = _b_bytes(K_TILE * N_TILE, BFP16_B) + mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] + mt_out_bytes = C_SLICE_LEN * ROWS * 2 + # L3 (DDR), flat -- the taps below index them linearly. + a_l3_ty = np.ndarray[(M * K,), bf16_ty] + b_l3_ty = np.ndarray[(K * N // B_GROUP,), b_elem_ty] + c_l3_ty = np.ndarray[(M * N,), bf16_ty] + + acc_init = Kernel("mm_fused_acc_init", kernel_object, [ct_acc_ty]) + # The trailing int32 is the A band index: under asymmetric tile buffering + # the core folds RHO A bands into one accumulator, so the kernel needs to + # know which band it is writing. + k_step = Kernel( + "mm_fused_k_step", + kernel_object, + [ct_a_obj_ty, ct_b_ty, ct_acc_ty, np.int32], + ) + # Same object as the mmul: the epilogue is compiled into mm_fused.cc, so + # one -D flag set and one artifact cover both. + epilogue_chunk = Kernel( + EPILOGUE_SYMBOL, + kernel_object, + [ct_out_ty, ct_acc_ty, np.int32, np.int32], + ) + + # --- Data movement ---------------------------------------------------- + # + # These stream-dimension lists are the load-bearing part of the design: + # they are what turns a row-major DDR tile into the r x s / s x t blocked + # layout the mmul indexes, and they are tightly coupled to it. A mismatch + # here produces silently wrong results, not a build error. + + # C: de-block each core's r x t tiled output back into row-major within its + # 64x128 slice, on the way into the memtile. + gather_dims = [(M_TILE // R, R * N_TILE), (N_TILE // T, T), (R, N_TILE), (T, 1)] + # B: DDR row-major (k x n) -> s x t blocks (recv), then split into the + # CT_MAX_K-deep chunks a single mmul call consumes (send). + # B needs no reblocking on either hop: pack_B already emits it in the + # order the cores consume, so the memtile just streams it through. That + # frees every descriptor dimension B used to spend -- which is what lets + # CT_MAX_K reach 128 (the innermost run would otherwise overflow the BD's + # 10-bit size field and need a split dimension) at the same time as + # residency (which spends one on its outer k walk). + b_recv_dims = None + b_send_dims = None + # A: same idea, r x s blocks. + a_recv_dims = [(M_TILE // R, R * K_TILE), (R, S), (K_TILE // S, R * S), (S, 1)] + a_send_dims = [ + (K_DIV_CT_K_MAX, R * CT_MAX_K), + (M_TILE // R, R * K_TILE), + ] + split_run(R * CT_MAX_K) + + # C: one join per column. Each of the ROWS cores in the column drops its + # slice at its own offset in a single memtile buffer, which then drains to + # DDR as one contiguous (ROWS*M_TILE) x N_TILE block. + c_l2l3_fifos = [] + c_prod = {} + for c in range(n_active_cols): + of_c = ObjectFifo(mt_out_ty, name=f"C_L2L3_{c}", depth=C_DEPTH) + c_l2l3_fifos.append(of_c) + sub = of_c.prod().join( + [C_SLICE_LEN * r for r in range(ROWS)], + obj_types=[ct_out_ty] * ROWS, + names=[f"C_L1L2_{c}_{r}" for r in range(ROWS)], + dims_from_stream=[gather_dims] * ROWS, + ) + for r in range(ROWS): + c_prod[(r, c)] = sub[r] + + # A: shim -> memtile -> broadcast along the compute row. The reblocking + # rides the forward(): inbound on cons(dims_from_stream=), outbound on + # forward(dims_to_stream=), sharing one memtile buffer. + a_l3l2_fifos = [] + a_cons = {} + for r in range(ROWS): + of_a_in = ObjectFifo(mt_a_ty, name=f"A_L3L2_{r}", depth=A_DEPTH) + a_l3l2_fifos.append(of_a_in) + of_a = of_a_in.cons(dims_from_stream=a_recv_dims).forward( + obj_type=ct_a_obj_ty, + depth=A_DEPTH, + name=f"A_L2L1_{r}", + dims_to_stream=a_send_dims, + ) + # One cons() handle per active column; every tile in the row sees + # this object, so inactive columns must not be consumers at all. + for c in range(n_active_cols): + a_cons[(r, c)] = of_a.cons() + + # B: shim -> memtile -> broadcast down the compute column. + # + # Where it fits, a whole column-block's B is held in the memtile as ONE + # object and re-walked per row-block, so DDR reads it once instead of + # m_row_blocks times. B is the dominant DDR leg, so that is roughly 43% less + # total traffic; the latency it buys grows with M, because B's re-reads + # scale with m_row_blocks. Larger K does not fit and falls back to + # re-reading. See README.md for the measured effect. + # + # Three things here are load-bearing rather than tuning: + # + # * ONE object spanning every k-block, not a pool of k_iters objects. + # Iterating a pool replays each object in turn (k0,k0,k1,k1,...) rather + # than the k0..kn sequence the cores accumulate in. + # * repeat_count on the forward() below is what re-sends an object. + # iter_count only bounds how many times an end cycles through all its + # buffers, so it is in units of depth-cycles; getting it wrong hangs + # rather than mis-computing. + # * The depth search takes the deepest that fits, not depth 1. Single + # buffering stops the next column-block prefetching behind this one's + # replay, which measures worse than not being resident at all. + # + # The budget must count what A and C actually occupy: at tile_n=128 C + # doubles and a resident B is 432 KB, and a budget that assumed C's size + # admitted a configuration that then failed address assignment. Placement + # has zero slack -- the eight memtiles pack to exactly 512 KB, relying on + # aie-objectfifo-allocate spilling one buffer to an adjacent tile -- so + # re-verify it after any change to the A, B or C buffer sizes. + mt_free = tm.get_mem_tile_size() - mt_a_bytes * A_DEPTH - mt_out_bytes * C_DEPTH + MT_B_DEPTH = next( + (d for d in (B_DEPTH, 1) if k_iters * mt_b_bytes * d <= mt_free), 0 + ) + b_resident = MT_B_DEPTH > 0 + if b_resident: + # Just a bigger buffer. With B packed in consumption order the walk is + # linear, so spanning every k-block needs no extra descriptor + # dimension -- the objects simply come out in k order. (The previous + # blocked layout had to widen one dim inbound and add an outermost k + # dim outbound, which is what collided with CT_MAX_K=128.) + mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE // B_GROUP,), b_elem_ty] + + b_l3l2_fifos = [] + b_cons = {} + for c in range(n_active_cols): + of_b_in = ObjectFifo( + mt_b_ty, name=f"B_L3L2_{c}", depth=MT_B_DEPTH if b_resident else B_DEPTH + ) + b_l3l2_fifos.append(of_b_in) + of_b = of_b_in.cons(dims_from_stream=b_recv_dims).forward( + # The one placement pin this design keeps. Everything else -- the + # workers, the accumulator buffers, the C join, the A forward and + # the shim ends -- is left to the placer, and measures the same. + # + # Without it, aie-place-tiles merges the 20 logical memtiles (4 A + # relays + 8 B relays + 8 C joins) onto the 8 physical ones in a way + # that aie-objectFifo-stateful-transform then rejects with "number + # of input DMA channel exceeded". Spreading B one-per-column is + # enough to steer it to a legal assignment; see the mlir-aie issue + # referenced in README.md. Reproduces at M=1024 K=2048 N=2048, which + # test.py covers. + tile=Tile(c, 1), + obj_type=ct_b_ty, + depth=L1_B_DEPTH, + name=f"B_L2L1_{c}", + dims_to_stream=b_send_dims, + # Replay the resident memtile object once per row-block. This is + # the mechanism that actually re-sends an object; iter_count only + # bounds how many chain iterations happen in total. Correct + # ordering depends on the memtile holding ONE object spanning every + # k-block: replicating a pool of k_iters smaller objects would + # emit k0,k0,k1,k1,... rather than the k0..kn sequence the cores + # accumulate in. + repeat_count=m_row_blocks if b_resident else None, + ) + for r in range(ROWS): + b_cons[(r, c)] = of_b.cons() + + # --- Compute ---------------------------------------------------------- + def core_fn(n_work, n_drain, acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): + """Core body for a column that computes ``n_work`` column-blocks and + then drains A for ``n_drain`` more (0 or 1). + + n_work/n_drain are bound per column via functools.partial below; they + are compile-time constants, so the trip counts below fold away. + """ + # The loop nest lives here rather than inside the kernel so that + # every level has an ObjectFifo acquire point. With M/K/N known at + # compile time all the trip counts are constants. + if n_work: + for _ in range_(n_work): + for _ in range_(m_row_blocks): + init_k(acc) + for _ in range_(k_iters): + # The l loop is unrolled by the B fifo depth so the + # acquired buffer index stays a compile-time + # constant. + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + # One B chunk feeds every A band, so B is + # acquired once around the band loop. + b = b_h.acquire(1) + for band in range(RHO): + a = a_h.acquire(1) + kstep_k(a, b, acc, band) + a_h.release(1) + b_h.release(1) + # Drain the accumulator. Unrolled by C_DEPTH for the + # same reason; a full O_CHUNKS unroll overflows program + # memory. + for chunk in range_(O_CHUNKS // C_DEPTH): + for half in range(C_DEPTH): + o = o_h.acquire(1) + epi_k(o, acc, chunk, half) + o_h.release(1) + if n_drain: + # The trailing partial column-block, for a column that sits it + # out. A is broadcast along the whole compute row, so this + # column must still consume its share or the columns that DO + # have work stall waiting for the fifo to advance. No B and no + # C here -- the runtime sequence issues neither for it. + for _ in range_(n_drain): + for _ in range_(m_row_blocks): + for _ in range_(k_iters): + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + for _ in range(RHO): + a_h.acquire(1) + a_h.release(1) + + workers = [] + for r in range(ROWS): + for c in range(n_active_cols): + acc = Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}") + workers.append( + Worker( + partial(core_fn, col_work[c], col_drain[c]), + [ + acc, + c_prod[(r, c)].prod(), + b_cons[(r, c)], + a_cons[(r, c)], + acc_init, + k_step, + epilogue_chunk, + ], + ) + ) + + # --- Runtime ---------------------------------------------------------- + # + # Every wrap below stays under the shim's wrap/size field (DMA_BD_MAX_WRAP): the + # largest are K_TILE=512 and ROWS*M_TILE=256. + # One transfer per (column-block, leg) instead of one per object. + # + # A single fill/drain may span MANY fifo objects -- the descriptor just + # walks them in the order the cores consume -- so a whole column-block's + # worth of A, B and C each go out as one task. Issuing per object instead + # meant a host-side await for every sweep, and those awaits were the + # serialisation: the next sweep could not start until the previous sweep's + # C had come all the way back. + # + # Dimension order must match the core loop nest exactly: for each + # column-block it walks mega_row, then k. Every wrap stays under the shim's + # wrap/size field (largest are K_TILE=512 and ROWS*M_TILE=256). + def a_taps(mega_col, r, mbs): + # Every (mega_row, k) block this compute row consumes for one + # column-block. A does not depend on mega_col; it is re-fetched per + # column-block because the cores re-consume it. + # + # Returns a LIST: one 4D descriptor normally, or one 3D descriptor per + # mega_row when the mega_row stride would overflow the shim BD's 20-bit + # iteration step (see a_split above). The split form carries the + # mega_row jump in the offset, which has no such limit. + if not a_split: + return [ + TensorAccessPattern( + tensor_dims=(M * K,), + offset=r * M_TILE * K, + sizes=[m_row_blocks, k_iters, M_TILE, K_TILE], + strides=[ROWS * M_TILE * K, K_TILE, K, 1], + ) + ] + return [ + TensorAccessPattern( + tensor_dims=(M * K,), + offset=mb * ROWS * M_TILE * K + r * M_TILE * K, + sizes=[1, k_iters, M_TILE, K_TILE], + strides=[0, K_TILE, K, 1], + ) + for mb in mbs + ] + + def b_tap(mega_col, c): + # Every (mega_row, k) chunk this column consumes. B does not depend on + # mega_row, hence the 0 stride: the same k-blocks are replayed for each + # row-block, which is what the cores expect. + # + # B must arrive PRE-PACKED (see GEMM.pack_B) so each k-block is one + # contiguous run. Expressing that reorder in the descriptor instead + # gives an innermost run of T=8 bf16, turning each 128 KB transfer into + # 8192 scattered bursts -- measured 5.4x slower end to end. + return TensorAccessPattern( + tensor_dims=(K * N // B_GROUP,), + offset=(mega_col * COLS + c) * N_TILE * K // B_GROUP, + sizes=( + [1, 1, 1, k_iters * K_TILE * N_TILE // B_GROUP] + if b_resident + else [m_row_blocks, k_iters, 1, K_TILE * N_TILE // B_GROUP] + ), + strides=( + [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // B_GROUP, 0, 1] + ), + ) + + def c_taps(mega_col, c, mbs): + # Every joined block this column produces for one column-block: one + # ROWS*M_TILE x N_TILE block per row-block. Returns a LIST, for the + # same reason a_taps does -- one descriptor per mega_row when N makes + # the mega_row stride overflow the shim BD's iteration step. + if c_split: + return [ + TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * N_TILE + mb * ROWS * M_TILE * N, + sizes=[1, 1, ROWS * M_TILE, N_TILE], + strides=[0, 0, N, 1], + ) + for mb in mbs + ] + return [_c_tap_unsplit(mega_col, c)] + + def _c_tap_unsplit(mega_col, c): + return TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * N_TILE, + sizes=[1, m_row_blocks, ROWS * M_TILE, N_TILE], + strides=[0, ROWS * M_TILE * N, N, 1], + ) + + def sequence(A, B, C, a_prods, b_prods, c_conses): + # Column-blocks 0..n_full-1 use every column; the trailing one (when N + # is not a multiple of N_TILE*COLS) uses only the first rem_blocks. A + # is always issued for every row, because the columns sitting the + # trailing block out still drain their share of the broadcast. + blocks = [(mc, COLS) for mc in range(n_full)] + if rem_blocks: + blocks.append((n_full, rem_blocks)) + + # One task per (column-block, leg): three per column instead of one per + # object, so a whole column-block retires on a single await rather than + # one per row-block. The C drain is issued first and retired last -- it + # is an S2MM that simply waits for the cores, so keeping it outstanding + # is what overlaps compute with write-back, and it must not share a + # group with the fills it depends on. + # Depth-2: issue column-block i+1 before retiring i, so its transfers + # are already moving while i computes. Retiring a block before issuing + # the next serialises on the C await, which waits for the cores. + # + # This is affordable only because each leg is a single task per + # column-block. Per-object tasks need 1 + 2*k_iters and cannot be + # overlapped at all. + # Keep OVERLAP column-blocks in flight, against SHIM_BDS buffer + # descriptors per column; the operator is DDR-rate bound rather than + # byte bound, so how deeply the fills are pipelined is what decides the + # rate. + # + # Every task of a block stays LIVE in its TaskGroup until the block is + # retired here. That is load-bearing, not tidiness -- see the + # dma_free_task note on a_split above. + all_mb = list(range(m_row_blocks)) + + # One emitter per leg, so the two paths below differ only in HOW they + # group and retire, not in how a leg is issued. + def issue_a(mega_col, mbs, group, wait=False): + for r in range(ROWS): + for tap in a_taps(mega_col, r, mbs): + a_prods[r].fill(A, tap, group=group, wait=wait) + + def issue_b(mega_col, active_cols, group): + for c in range(active_cols): + b_prods[c].fill(B, b_tap(mega_col, c), group=group) + + def issue_c(mega_col, active_cols, mbs, group): + for c in range(active_cols): + for tap in c_taps(mega_col, c, mbs): + c_conses[c].drain(C, tap, group=group, wait=True) + + def emit_unsplit(): + pending = [] + for mega_col, active_cols in blocks: + # C in its own group, issued first and retired last: it is an + # S2MM that waits on the cores, so keeping it outstanding is + # what overlaps compute with write-back, and it must not share + # a group with the fills it depends on. + tg_c = TaskGroup() + issue_c(mega_col, active_cols, all_mb, tg_c) + tg_f = TaskGroup() + issue_a(mega_col, all_mb, tg_f) + issue_b(mega_col, active_cols, tg_f) + + pending.append([tg_f, tg_c]) + while len(pending) >= OVERLAP: + for tg in pending.pop(0): + tg.finish() + + for group in pending: + for tg in group: + tg.finish() + + def emit_split(): + # A split leg is one transfer per mega_row, so a whole block at + # once would overrun the shim channel's task queue. Emit MB_WINDOW + # mega_rows at a time and retire each window before the next, which + # both drains the queue and -- because the window's transfers are + # awaited, not merely freed -- makes its descriptors safe to reuse. + # + # B stays live across the whole block: its descriptor replays over + # every mega_row, so freeing it per window would hand its + # descriptor away mid-flight. It is retired last, after every + # window's C has been awaited, which is what guarantees it drained. + for mega_col, active_cols in blocks: + tg_b = TaskGroup() + issue_b(mega_col, active_cols, tg_b) + + # The leg that did NOT split is still one task for the whole + # block -- its single descriptor already spans every mega_row, + # so re-issuing it per window would transfer the block twice. + # It stays live alongside the windows and retires with them. + tg_whole = TaskGroup() + if not c_split: + issue_c(mega_col, active_cols, all_mb, tg_whole) + if not a_split: + issue_a(mega_col, all_mb, tg_whole) + + for w in range(0, m_row_blocks, MB_WINDOW): + mbs = all_mb[w : w + MB_WINDOW] + tg_w = TaskGroup() + if c_split: + issue_c(mega_col, active_cols, mbs, tg_w) + if a_split: + # wait=True: the await is what makes this window's + # descriptors reusable by the next. + issue_a(mega_col, mbs, tg_w, wait=True) + tg_w.finish() + + tg_whole.finish() + tg_b.finish() + + emit_split() if (a_split or c_split) else emit_unsplit() + + rt = Runtime( + sequence, + [ + a_l3_ty, + b_l3_ty, + c_l3_ty, + [f.prod() for f in a_l3l2_fifos], + [f.prod() for f in b_l3l2_fifos], + [f.cons() for f in c_l2l3_fifos], + ], + ) + + my_program = Program(dev, rt, workers=workers) + maybe_enable_trace(my_program, trace_size, workers) + return my_program.resolve_program() + + +def main(): + argparser = argparse.ArgumentParser( + prog="FLM GEMM MLIR Design", + description="Emits MLIR code for a row-broadcast bf16 GEMM of the given input size", + ) + argparser.add_argument("--dev", type=str, choices=["npu1", "npu2"], default="npu2") + argparser.add_argument("-M", type=int, default=MIN_M) + argparser.add_argument("-K", type=int, default=MIN_K) + argparser.add_argument("-N", type=int, default=1024) + argparser.add_argument( + "--tile-n", + type=int, + choices=sorted(CT_MAX_K_FOR_N), + default=N_TILE_DEFAULT, + ) + argparser.add_argument( + "--tile-ma", + type=int, + default=None, + help="Rows of A held in L1 at a time; defaults to the largest that fits", + ) + argparser.add_argument( + "--epilogue", type=Epilogue, choices=list(Epilogue), default=Epilogue.NONE + ) + argparser.add_argument("--trace_size", type=int, default=0) + + args = argparser.parse_args() + print( + gemm( + NPU1() if args.dev == "npu1" else NPU2(), + args.M, + args.K, + args.N, + epilogue=args.epilogue, + tile_n=args.tile_n, + tile_ma=args.tile_ma, + trace_size=args.trace_size, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py new file mode 100644 index 0000000000..f7c3be3cab --- /dev/null +++ b/iron/operators/flm/gemm/op.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field + +import numpy as np +from typing import ClassVar, Dict + +from iron.common import ( + MLIROperator, + AIERuntimeArgSpec, + KernelArchiveArtifact, + KernelObjectArtifact, + SourceArtifact, + PythonGeneratedMLIRArtifact, + DesignGenerator, +) +from aie.dialects.aie import get_target_model +from aie.dialects._aie_enum_gen import AIEArch +from iron.common.device_utils import get_kernel_dir +from iron.common.operator_bases import lut_based_ops_artifacts +from iron.common.utils import float_to_name +import aie.utils as aie_utils + +from iron.operators.flm.packing import pack_b, packed_b_size +from iron.operators.flm.gemm.design import ( + BFP16_GROUP, + BFP16_GROUP_BYTES, + CT_MAX_K_FOR_N, + C_DEPTH, + compute_rows, + CT_OUT_LEN, + Epilogue, + K_TILE, + MIN_K, + M_TILE, + R, + Rounding, + S, + T, + _default_l1, +) + + +@dataclass +class GEMM(MLIROperator): + """AIE-accelerated bf16 GEMM on a 4-row grid, with a fused epilogue. + + A row-broadcast / C memtile-join design with fixed 64/512/128 tiling. See + ``design.py`` for how it differs from the more general ``GEMM`` operator. + Unlike ``GEMM`` this exposes no tiling knobs, but folds an activation and an + optional clamp into the output stage. + + The grid is as wide as the device: 8 columns on NPU2, 4 on NPU1 (Phoenix). + Only the width varies -- the tiling and the blocked L1 layout are shared. + """ + + # Every field below is repr=True, so MLIROperator.name derives the artifact + # stem from all of them. That is not cosmetic: each one changes the emitted + # MLIR or the kernel object, and this repo's build cache keys on filename, + # so a variant that shared a stem would be silently satisfied by another + # variant's cached build. + M: int + K: int + N: int + # Activation fused into the C drain. + epilogue: Epilogue = Epilogue.NONE + # Optional (min, max) applied after the activation. + clamp: tuple[float, float] | None = None + # n tile width. 64 halves the mmul's accumulator traffic per mac; 128 + # halves A fetches instead. __post_init__ resolves None per device and + # shape; see the comment there and README.md. + tile_n: int | None = None + # A-tile rows, decoupled from the accumulator's M_TILE (asymmetric tile + # buffering). __post_init__ resolves None to whatever L1 affords. + tile_ma: int | None = None + # Rounding for every f32->bf16 conversion; see Rounding in design.py. + rounding: Rounding = Rounding.CONV_EVEN + context: object = field(default=None, repr=False) + + _name_aliases: ClassVar[Dict[str, str]] = { + **MLIROperator._name_aliases, + "epilogue": "epi", + "tile_n": "tn", + "tile_ma": "ma", + "rounding": "rnd", + } + + def __post_init__(self): + # Resolve both tile knobs to concrete values here, so the dataclass + # fields hold what the build actually uses. The resolved tile_ma in + # particular must reach the artifact name: the design sizes the A object + # from it while the kernel derives the mmul's rowA from it, so an + # artifact built for one value must never satisfy a request for another. + dev = aie_utils.get_current_device() + if self.tile_n is None: + # n=64 gives the mmul colA=8 instead of 4, halving accumulator + # traffic per mac; n=128 halves A fetches instead. Which wins + # depends on whether compute or data movement is the critical path. + # + # On NPU2 that flips with K: with a single k iteration there is too + # little compute to hide the extra A traffic, so n=128 wins there + # (~9%), while n=64 wins by ~20% at K >= 1024. + # + # NPU1 never reaches that crossover. It has half the columns AND a + # quarter of the per-tile bf16 mac throughput (four native 4x8x4 + # macs per 8x8x8 shape, against NPU2's two bfp16-emulated ones), so + # it stays compute-bound at every K, and n=128's 32 KB f32 + # accumulator also overflows bank-aware L1 allocation. n=64 wins + # everywhere there by 1.21-1.38x, including at K=512 where NPU2's + # rule would pick 128. + single_k_iter = self.K // K_TILE <= 1 + self.tile_n = 128 if (dev.arch == AIEArch.AIE2p and single_k_iter) else 64 + elif self.tile_n not in CT_MAX_K_FOR_N: + raise ValueError( + f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {self.tile_n}" + ) + if self.tile_ma is None: + self.tile_ma = _default_l1( + self.tile_n, + CT_MAX_K_FOR_N[self.tile_n], + self._b_elem_bytes, + get_target_model(dev.resolve()).get_local_memory_size(), + )[0] + # N only needs to tile to N_TILE: a trailing group of fewer than + # COLS column-blocks is handled by giving the columns different trip + # counts. See design.py. + for name, value, unit in ( + ("M", self.M, M_TILE * compute_rows(dev)), + ("K", self.K, MIN_K), + ("N", self.N, self.tile_n), + ): + if value % unit != 0: + raise ValueError(f"{name} ({value}) must be a multiple of {unit}") + # Coerce so callers may pass the bare string; the enums are StrEnum, so + # the resolved fields still serialize into artifact names unchanged. + self.epilogue = Epilogue(self.epilogue) + self.rounding = Rounding(self.rounding) + if self.clamp is not None: + lo, hi = self.clamp + if lo > hi: + raise ValueError(f"clamp min ({lo}) must be <= max ({hi})") + + MLIROperator.__init__(self, context=self.context) + + @property + def name(self) -> str: + """Artifact stem, prefixed to disambiguate from ``iron.operators.GEMM``. + + ``MLIROperator.name`` derives the stem from ``type(self).__name__``, + which is ``GEMM`` for both operators. This repo's build cache keys on + filename and mtime rather than on source or flags, so two operators + sharing a stem in one build dir would silently satisfy each other. + """ + return f"FLM_{super().name}" + + @property + def _bfp16_b(self) -> bool: + """Whether B is stored as bfp16ebs8 rather than bf16. + + AIE2P only, and the reason both mmul templates in the kernel header are + live rather than one being dead code: on AIE2 the scalar BFP types do + not exist, so B stays bf16 and the mmul lowers onto four native 4x8x4 + macs. + """ + return aie_utils.get_current_device().arch == AIEArch.AIE2p + + @property + def _b_elem_bytes(self) -> float: + """Bytes per B element in L1/L2: bfp16ebs8 packs 8 values into 9 bytes.""" + return BFP16_GROUP_BYTES / BFP16_GROUP if self._bfp16_b else 2 + + @property + def _kernel_object(self) -> str: + """Object name over every flag that changes the emitted code. + + Everything the -D flags in ``get_kernel_artifacts`` carry has to appear + here: this repo's build cache keys on filename and mtime rather than on + source or flags, so an object built for one configuration would + otherwise silently satisfy a request for another. That includes r/t, + which set the blocked layout, and the epilogue flags, which since the + epilogue was folded into this translation unit shape the same object. + """ + clamp = "" + if self.clamp is not None: + clamp = "_clamp" + "_".join(float_to_name(float(v)) for v in self.clamp) + return ( + f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" + f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" + f"_epi{self.epilogue}{clamp}.o" + ) + + @property + def _link_file(self) -> str: + """What the design names as its kernel: the bare object, or the archive + bundling it with the tanh LUT tables. + + Only AIE2 evaluates the activations through a LUT (AIE2P has a native + vector tanh), and only an activation references tanh at all -- the plain + epilogue just converts and stores. When this is wrong the failure is a + LINK error for tanh_lut_ab/tanh_lut_cd rather than a compile error, so + it surfaces late. + """ + if self.epilogue is not Epilogue.NONE and get_kernel_dir() == "aie2": + return f"{self.name}_kernels.a" + return self._kernel_object + + def get_mlir_artifact(self): + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator( + self.operator_dir / "design.py", + "gemm", + (), + { + "dev": aie_utils.get_current_device(), + "M": self.M, + "K": self.K, + "N": self.N, + "tile_n": self.tile_n, + "tile_ma": self.tile_ma, + "epilogue": self.epilogue, + "kernel_object": self._link_file, + "trace_size": 0, + }, + ), + ) + + def get_kernel_artifacts(self): + kernel_dir = get_kernel_dir() + base_dir = self.context.base_dir + generic = base_dir / "aie_kernels" / "generic" + + # mm_fused.cc includes zero.cc, which is genuinely per-architecture + # (AIE2 stores 256 bits at a time, AIE2P 512). A quoted include searches + # the including file's own directory first -- now generic/ -- so the + # arch directory has to be on the include path for it to resolve there. + arch_include = [f"-I{base_dir / 'aie_kernels' / kernel_dir}"] + + # The 8x8x8 mmul shape this design uses exists on both architectures, + # but by different routes: AIE2P lowers it onto two bfp16-emulated macs, + # which is what this flag selects, while AIE2 lowers it onto four native + # 4x8x4 bf16 macs and ignores the flag entirely (it has no bfp16 + # hardware). Passing it on AIE2 would be harmless but misleading, so it + # is scoped to the architecture where it actually changes codegen. + # + # MM_FUSED_BFP16_B rides along with it: storing B as bfp16ebs8 needs the + # scalar BFP types, which only AIE2P has. See _bfp16_b. + flags = [ + # Tile geometry and register tiling, for the mmul. + f"-DMM_FUSED_TILE_M={M_TILE}", + f"-DMM_FUSED_TILE_K={K_TILE}", + f"-DMM_FUSED_TILE_N={self.tile_n}", + f"-DMM_FUSED_TILE_MA={self.tile_ma}", + f"-DMM_FUSED_R={R}", + f"-DMM_FUSED_S={S}", + f"-DMM_FUSED_T={T}", + # The k slice. Passed rather than looked up in the kernel so that + # CT_MAX_K_FOR_N is the only place it is chosen. + f"-DMM_FUSED_CT_K={CT_MAX_K_FOR_N[self.tile_n]}", + # Output stage. + f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", + f"-DMM_FUSED_C_DEPTH={C_DEPTH}", + f"-DMM_FUSED_EPILOGUE_MODE={self.epilogue.mode}", + ] + arch_include + if self._bfp16_b: + flags += [ + "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", + "-DMM_FUSED_BFP16_B", + ] + if self.clamp is not None: + lo, hi = self.clamp + # repr() rather than :g -- the latter renders -4.0 as "-4", and + # "-4f" is not a valid C float literal. + flags += [ + "-DMM_FUSED_CLAMP=1", + f"-DMM_FUSED_CLAMP_MIN={float(lo)!r}f", + f"-DMM_FUSED_CLAMP_MAX={float(hi)!r}f", + ] + if self.rounding is Rounding.CONV_EVEN: + # ROUND_CONV_EVEN is mm.cc's flag, reused rather than inventing a + # second spelling, and its polarity is mm.cc's too: absent means the + # core's power-up floor mode, even though this operator defaults the + # other way. It covers both conversions in the kernel -- the mmul + # and the epilogue's f32->bf16 store -- which must agree. + flags.append("-DROUND_CONV_EVEN") + + kernel_obj = KernelObjectArtifact( + self._kernel_object, + dependencies=[ + SourceArtifact(generic / "mm_fused.cc"), + SourceArtifact(generic / "mm_fused_mmul.h"), + SourceArtifact(generic / "activations.h"), + SourceArtifact(base_dir / "aie_kernels" / "aie_kernel_utils.h"), + SourceArtifact(base_dir / "aie_kernels" / kernel_dir / "zero.cc"), + ], + extra_flags=flags, + ) + if self._link_file == self._kernel_object: + return [kernel_obj] + # The tanh LUT's coefficient tables live in their own translation unit + # in mlir-aie's runtime lib, so on AIE2 the kernel object alone leaves + # tanh_lut_ab/tanh_lut_cd undefined at link time. See _link_file. + return [ + KernelArchiveArtifact( + self._link_file, + dependencies=[kernel_obj] + lut_based_ops_artifacts(kernel_dir), + ) + ] + + def pack_B(self, B): + """Reorder a row-major ``(K, N)`` weight matrix into consumption order. + + Returns a flat uint8 tensor of bfp16ebs8 blocks on NPU2, where B is also + quantized, and a flat bf16 tensor on NPU1. Bound to the operator rather + than a static method because the layout depends on the resolved + ``tile_n`` and on the device; call ``op.pack_B(B)``. + + Packing all the way to consumption order is what makes both B hops + linear descriptors (design.py's b_recv_dims and b_send_dims are both + None), which in turn leaves the descriptor dimensions for a k slice deep + enough to halve the accumulator traffic while B is also memtile-resident. + See :mod:`iron.operators.flm.packing` for the layout itself. + """ + return pack_b( + B, + k_tile=K_TILE, + n_tile=self.tile_n, + s=S, + t=T, + ct_k=CT_MAX_K_FOR_N[self.tile_n], + bfp16=self._bfp16_b, + round_conv_even=self.rounding is Rounding.CONV_EVEN, + ) + + def packed_B_size(self, K, N): + """Elements (bf16) or bytes (bfp16ebs8) that ``pack_B`` returns.""" + return packed_b_size(K, N, self._bfp16_b) + + def get_arg_spec(self): + return [ + AIERuntimeArgSpec("in", (self.M, self.K)), # A + # B arrives pre-packed by pack_B. On AIE2P it is also quantized to + # bfp16ebs8 -- 9 bytes per 8 values rather than bf16's 16 -- so it + # is declared in BYTES there, sizing the buffer from what pack_B + # actually returns; a (K, N) bf16 spec would over-allocate the + # largest buffer by 1.78x. On AIE2 B stays bf16 and the spec is the + # plain element count. + ( + AIERuntimeArgSpec( + "in", (self.packed_B_size(self.K, self.N),), dtype=np.uint8 + ) + if self._bfp16_b + else AIERuntimeArgSpec("in", (self.K, self.N)) + ), # B (weights) + AIERuntimeArgSpec("out", (self.M, self.N)), # C + ] + + def reference(self, A, B): + """CPU reference: ``C = epilogue(A @ B)``.""" + from iron.operators.flm.gemm.reference import reference + + return reference(A, B, self.epilogue, self.clamp) diff --git a/iron/operators/flm/gemm/reference.py b/iron/operators/flm/gemm/reference.py new file mode 100644 index 0000000000..24eda2335f --- /dev/null +++ b/iron/operators/flm/gemm/reference.py @@ -0,0 +1,84 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from iron.common.test_utils import torch_dtype_map +from iron.operators.flm.gemm.design import Epilogue + + +def apply_epilogue(C, epilogue=Epilogue.NONE, clamp=None): + """The fused output stage alone, applied to an already-accumulated C. + + Separate from ``reference`` because a test that wants to check the epilogue + without the accumulation needs exactly this -- see + ``mm_prebuilt/test.py``'s accumulator comparison, where the device's own + output is the input. + + ``gelu`` is the sigmoid approximation ``x * sigmoid(1.702x)``, matching the + kernel -- NOT torch's erf-exact gelu, and not the tanh approximation the + standalone gelu operator uses. + """ + match Epilogue(epilogue): + case Epilogue.NONE: + pass + case Epilogue.GELU: + C = C * torch.sigmoid(1.702 * C) + case Epilogue.SILU: + C = C * torch.sigmoid(C) + case Epilogue.SIGMOID: + C = torch.sigmoid(C) + if clamp is not None: + C = torch.clamp(C, clamp[0], clamp[1]) + return C + + +def reference(input_a, input_b, epilogue=Epilogue.NONE, clamp=None): + """CPU reference ``C = clamp(activation(A @ B))``. + + Follows the kernel's order of operations rather than an idealized one: the + matmul accumulates in fp32, mirroring the f32 accumulator, and the result is + converted to the output dtype BEFORE the activation and clamp, because that + is what the kernel does -- ``mm_fused_epilogue_chunk`` does + ``to_v16bfloat16(acc)`` and then applies the activation to that bf16 vector. + + Measured end to end this ordering is second order, because the accumulator's + own error dominates: at M=256 K=512 N=1024 it moves mean |err| by under + 2e-4 either way. It is worth doing because it models what the kernel does, + and it is clearly visible once the accumulator is taken out of the + comparison -- against the device's OWN accumulator on the shipped overlay it + moves silu's worst-case disagreement from 0.043 to 0.031. + + Still not bit-exact, and cannot be. The remaining gap is the hardware's own + activation approximation -- a LUT on AIE2, a native instruction on AIE2P -- + worth up to ~0.02 absolute there, which no CPU reference built on exact + ``torch.sigmoid`` can reproduce. Tolerances have to absorb that part. + """ + out_dtype = input_a.dtype + C = torch.matmul(input_a.float(), input_b.float()).to(out_dtype) + return apply_epilogue(C, epilogue, clamp) + + +def generate_golden_reference( + M: int, + K: int, + N: int, + dtype="bf16", + seed=42, + epilogue=Epilogue.NONE, + clamp=None, + scale=4.0, +): + """Random A (signed) and B (non-negative), scaled by ``scale``. + + ``scale`` matters for the epilogue tests: the result grows like + ``sqrt(K) * scale**2``, and at the default scale a K=512 product lands + around +-200, where gelu/silu are indistinguishable from the identity (or + from zero). Activation tests pass a smaller scale so the result sits in the + range where the curve is actually interesting. + """ + torch.manual_seed(seed) + dtype_torch = torch_dtype_map[dtype] + input_a = torch.randn(M, K, dtype=dtype_torch) * scale + input_b = torch.rand(K, N, dtype=dtype_torch) * scale + output = reference(input_a, input_b, epilogue, clamp) + return {"input": input_a, "input_b": input_b, "output": output} diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py new file mode 100644 index 0000000000..0dd4312e1e --- /dev/null +++ b/iron/operators/flm/gemm/test.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import aie.utils as aie_utils + +from aie.dialects.aie import get_target_model +from aie.dialects._aie_enum_gen import AIEArch + +from iron.operators import GEMM as GenericGEMM +from iron.operators.flm.gemm.design import ( + BFP16_GROUP, + BFP16_GROUP_BYTES, + CT_MAX_K_FOR_N, + Epilogue, + M_TILE, + R, + Rounding, + _b_depth_for, + _default_l1, +) +from iron.operators.flm.gemm.op import GEMM +from iron.operators.flm.gemm.reference import generate_golden_reference +from iron.common.test_utils import run_test + +# Unpacked so the parameter tables below stay column-aligned. +NONE, GELU, SILU, SIGMOID = Epilogue +CONV_EVEN, FLOOR = Rounding + +# Activation tests run at a smaller scale so the result lands where the curve +# is not flat. generate_golden_reference grows the result like sqrt(K)*scale**2, +# so at the default 4.0 a K=512 product sits around +-200, where gelu and silu +# are indistinguishable from the identity. +INPUT_SCALE = 4.0 +ACTIVATION_INPUT_SCALE = 0.5 + + +def get_params(): + dev = aie_utils.get_current_device() + if dev is None: + return [] + dev_name = dev.resolve().name + if dev_name not in ("npu1", "npu2"): + return [] + + # The grid is 4 rows by as many columns as the device has, so the width of + # one full sweep -- N_TILE * COLS -- differs per device, and so do the N + # values that leave a trailing PARTIAL column-block. That trailing case is + # the interesting one: some columns compute the block while the rest only + # drain the A broadcast, and real transformer o/down projections always + # land there, since N = model dim is essentially never a multiple of the + # sweep width. + # + # At K = 512 there is a single k iteration, so tile_n defaults to 128 and a + # sweep is 1024 wide on NPU2 and 512 on NPU1; at K >= 1024 tile_n drops to + # 64, halving both. + # fmt: off + if dev_name == "npu2": + # M, K, N, epilogue, clamp, rounding + regular_params = [ + ( 256, 512, 1024, NONE, None, CONV_EVEN), # smallest full sweep + ( 512, 1024, 2048, NONE, None, CONV_EVEN), + ( 256, 512, 1536, NONE, None, CONV_EVEN), # remainder: 4 of 8 cols + ( 256, 512, 128, NONE, None, CONV_EVEN), # remainder only: 1 col + ( 256, 512, 1024, SILU, None, CONV_EVEN), + ( 256, 512, 1024, GELU, None, CONV_EVEN), + ( 256, 512, 1024, NONE, (-2.0, 2.0), CONV_EVEN), + # floor reproduces the shipped FastFlowLM overlay's rounding mode + # (bit for bit on NPU2; NPU1 sums the K reduction in a different + # order). It is much less accurate, so it gets its own bound below. + ( 256, 512, 1024, NONE, None, FLOOR), + ] + extensive_params = [ + ( 1024, 2048, 2048, NONE, None, CONV_EVEN), + ( 2048, 2048, 2048, NONE, None, CONV_EVEN), + ( 1024, 2560, 2560, NONE, None, CONV_EVEN), # E4B o-proj + ( 512, 1536, 1536, SILU, None, CONV_EVEN), # E2B down-proj + ( 256, 512, 1024, SIGMOID, None, CONV_EVEN), + ( 512, 1024, 2048, SILU, (-4.0, 4.0), CONV_EVEN), + ( 256, 512, 1024, SILU, None, FLOOR), + # K or N = 10240 at M > 256 overflows the shim BD's 20-bit + # mega_row iteration step, so that leg is issued as one transfer + # per mega_row, retired in windows. These are the real E4B FFN + # projections and were unsupported until that landed; they are + # the regression cover for it. M=2048 needs two windows, which is + # what exercises the windowing. + ( 1024, 10240, 2560, NONE, None, CONV_EVEN), # E4B down + ( 1024, 2560, 10240, NONE, None, CONV_EVEN), # E4B gateup + ( 2048, 10240, 2560, NONE, None, CONV_EVEN), # A, 2 windows + ( 2048, 2560, 10240, NONE, None, CONV_EVEN), # C, 2 windows + ] + else: # npu1: _default_tile_n always returns 64 here, so with 4 columns + # every sweep is N_TILE*COLS = 256 wide, not the 128*4=512 an + # NPU2-shaped sweep would give. + # M, K, N, epilogue, clamp, rounding + regular_params = [ + ( 256, 512, 256, NONE, None, CONV_EVEN), # smallest full sweep + ( 512, 1024, 512, NONE, None, CONV_EVEN), + ( 256, 512, 128, NONE, None, CONV_EVEN), # remainder: 2 of 4 cols + ( 256, 512, 64, NONE, None, CONV_EVEN), # remainder only: 1 col + ( 256, 512, 320, NONE, None, CONV_EVEN), # full sweep + 1 col + ( 256, 512, 512, SILU, None, CONV_EVEN), + ( 256, 512, 512, GELU, None, CONV_EVEN), + ( 256, 512, 512, NONE, (-2.0, 2.0), CONV_EVEN), + ( 256, 512, 512, NONE, None, FLOOR), + ] + extensive_params = [ + ( 1024, 2048, 1024, NONE, None, CONV_EVEN), + ( 2048, 2048, 1024, NONE, None, CONV_EVEN), + ( 1024, 2560, 2560, NONE, None, CONV_EVEN), # E4B o-proj + ( 512, 1536, 1536, SILU, None, CONV_EVEN), # E2B down-proj + ( 256, 512, 512, SIGMOID, None, CONV_EVEN), + ( 512, 1024, 1024, SILU, (-4.0, 4.0), CONV_EVEN), + ( 256, 512, 512, SILU, None, FLOOR), + ] + # fmt: on + + params = [] + for p in regular_params: + params.append(pytest.param(*p)) + for p in extensive_params: + params.append(pytest.param(*p, marks=[pytest.mark.extensive])) + return params + + +def check_on_device(operator, golden_ref, K, rounding=CONV_EVEN): + """Run ``operator`` against its golden reference and return run_test's result. + + Bounds the error in ABSOLUTE terms as a fraction of the accumulated mass, + i.e. the expected size of the K reduction before cancellation, + K * mean|a| * mean|b|. A plain relative tolerance cannot work: with signed A + the K-sum cancels by ~sqrt(K), so |C| ends up far smaller than the mass + while the error tracks the mass, leaving near-zero outputs uncheckable. + + The fraction is per-architecture, because the two lower the same 8x8x8 mmul + onto very different arithmetic: NPU2 emulates it with bfp16, which drops + mantissa bits, while NPU1 has no bfp16 and lowers onto four native bf16 macs + accumulating in f32 -- exact up to the f32->bf16 store, so ~20x tighter. + floor truncates rather than rounding to nearest, so its bias accumulates + over the K reduction instead of cancelling and gets a looser bound on both. + """ + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + if aie_utils.get_current_device().resolve().name == "npu1": + budget = 0.002 if rounding is FLOOR else 0.0002 + else: + budget = 0.05 if rounding is FLOOR else 0.004 + return run_test( + operator, + { + "A": golden_ref["input"].flatten(), + # B is consumed pre-packed; see GEMM.pack_B. + "B": operator.pack_B(golden_ref["input_b"]), + }, + {"C": golden_ref["output"].flatten()}, + rel_tol=0.04, + abs_tol=float(budget * mass), + ) + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", + Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", +) +@pytest.mark.parametrize("M,K,N,epilogue,clamp,rounding", get_params()) +def test_gemm(M, K, N, epilogue, clamp, rounding, aie_context): + scale = INPUT_SCALE if epilogue is NONE else ACTIVATION_INPUT_SCALE + golden_ref = generate_golden_reference( + M=M, K=K, N=N, epilogue=epilogue, clamp=clamp, scale=scale + ) + + operator = GEMM( + M=M, + K=K, + N=N, + epilogue=epilogue, + clamp=clamp, + rounding=rounding, + context=aie_context, + ) + + errors, latency_us, bandwidth_gbps = check_on_device( + operator, golden_ref, K, rounding + ) + + gflops = (2.0 * M * K * N) / (latency_us * 1e-6) / 1e9 + print(f"\nLatency (us): {latency_us:.1f}") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s") + print(f"Throughput: {gflops:.6e} GFLOP/s\n") + + assert not errors, "Test failed" + + +def test_gemm_split_leg_windowing(aie_context): + """K or N = 10240 at M > 256 overflows the shim BD's 20-bit mega_row + iteration step, so that leg is issued as one transfer per mega_row, + retired in windows of at most SHIM_TASK_QUEUE. Two shim resources bound it + and NEITHER is modelled by the toolchain -- the BD ids (16/tile, freed + without a completion check) and the channel task queue (4 deep, pushed + unconditionally) -- so overrunning either is a silent device hang rather + than a diagnostic. + + Windowing keeps both inside their limits for every shape: at most + 1 B + 4 A + 4 C = 9 of 16 descriptors, and at most 4 outstanding per + channel. Assert that arithmetic here, since the numbers come from the + hardware and a future retune of SHIM_TASK_QUEUE could break it silently. + """ + from aie.dialects.aie import get_target_model + from iron.operators.flm.gemm.design import SHIM_TASK_QUEUE + + dev = aie_utils.get_current_device() + available = get_target_model(dev.resolve()).get_num_bds(0, 0) + worst = 1 + 2 * SHIM_TASK_QUEUE + assert worst <= available, ( + f"a fully split block needs {worst} shim BDs of {available}; " + "windowing no longer fits and the split shapes will hang" + ) + + # The square case splits BOTH legs, which the real Gemma shapes never do + # (E4B's down-proj overflows on K and its gate/up on N, never both), so it + # is the only cover for the two-sided path. + GEMM(M=512, K=10240, N=10240, context=aie_context).compile() + + +def test_gemm_split_leg_windowing_runs(aie_context): + """Execute the two-sided split path, not just compile it. + + test_gemm_split_leg_windowing above only compiles this shape: the failure + mode it guards against -- BD-id aliasing and shim task-queue overrun (see + that test's docstring) -- is a runtime device hang or silent corruption, + which compiling the MLIR cannot exercise. This dispatches the same shape on + hardware and checks the result. + + Regular rather than extensive despite being the largest shape here. What it + catches is a hang or silently wrong output, not a wrong number, and its + compile-only sibling is already regular, so leaving the executing half out + of the default run is the wrong side to err on. Costs ~8s against the + regular suite's ~13s. + """ + M, K, N = 512, 10240, 10240 + golden_ref = generate_golden_reference(M=M, K=K, N=N) + + operator = GEMM(M=M, K=K, N=N, context=aie_context) + + errors, _latency_us, _bandwidth_gbps = check_on_device(operator, golden_ref, K) + assert not errors, "Test failed" + + +def tile_option_params(): + """Every (tile_n, tile_ma) the design accepts on this device. + + The shape parameters above exercise only the DEFAULT tile geometry, because + __post_init__ resolves both knobs from the shape and the device. These cover + the knobs themselves, which change the blocked L1 layout: tile_n selects + CT_MAX_K and the B object width, tile_ma sets the mmul's rowA and the A + object height, and pack_B, the four stream-dimension lists and gather_dims + all key off them. A mismatch is silently wrong output rather than a build + error, so each combination has to actually run on hardware. + + The default tile_ma per tile_n stays in the regular suite; the overrides are + extensive, since each is its own kernel object and xclbin. + """ + dev = aie_utils.get_current_device() + if dev is None or dev.resolve().name not in ("npu1", "npu2"): + return [] + l1 = get_target_model(dev.resolve()).get_local_memory_size() + b_elem = BFP16_GROUP_BYTES / BFP16_GROUP if dev.arch == AIEArch.AIE2p else 2 + + params = [] + for tile_n, ct_k in sorted(CT_MAX_K_FOR_N.items()): + default_ma = _default_l1(tile_n, ct_k, b_elem, l1)[0] + # One full sweep of the grid at this tile_n, so every column has work. + M, K, N = 256, 512, tile_n * dev.cols + for tile_ma in (16, 32, 64): + if M_TILE % tile_ma or tile_ma % (2 * R): + continue + try: + _b_depth_for(tile_ma, tile_n, ct_k, b_elem, l1) + except ValueError: + continue # this A height leaves no room for B at this width + marks = [] if tile_ma == default_ma else [pytest.mark.extensive] + params.append( + pytest.param( + M, + K, + N, + tile_n, + tile_ma, + marks=marks, + id=f"tn{tile_n}-ma{tile_ma}" + ("-default" if not marks else ""), + ) + ) + return params + + +@pytest.mark.parametrize("M,K,N,tile_n,tile_ma", tile_option_params()) +def test_gemm_tile_options(M, K, N, tile_n, tile_ma, aie_context): + """Each accepted (tile_n, tile_ma) computes the right answer on hardware.""" + golden_ref = generate_golden_reference(M=M, K=K, N=N, scale=INPUT_SCALE) + operator = GEMM(M=M, K=K, N=N, tile_n=tile_n, tile_ma=tile_ma, context=aie_context) + assert operator.tile_n == tile_n and operator.tile_ma == tile_ma + errors, _latency_us, _bandwidth_gbps = check_on_device(operator, golden_ref, K) + assert not errors, "Test failed" + + +@pytest.mark.parametrize("M,K,N", [(256, 512, 1024), (512, 1024, 2048)]) +def test_artifact_stem_differs_from_generic_gemm(M, K, N, aie_context): + """``flm.GEMM`` must never share an artifact stem with ``GEMM``. + + Both classes are named ``GEMM``, and MLIROperator.name derives the stem + from the class name, while this repo's build cache keys on filename and + mtime rather than on source or flags -- so a shared stem would let the two + operators silently satisfy each other's builds in one build dir. + """ + assert ( + GEMM(M=M, K=K, N=N, context=aie_context).name + != GenericGEMM(M=M, K=K, N=N, context=aie_context).name + ) diff --git a/iron/operators/flm/mm_prebuilt/README.md b/iron/operators/flm/mm_prebuilt/README.md new file mode 100644 index 0000000000..730fcb0b9a --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/README.md @@ -0,0 +1,64 @@ + + +# `iron.operators.flm.MMPrebuilt` — FastFlowLM's shipped `mm` overlay + +```python +from iron.operators.flm import MMPrebuilt + +op = MMPrebuilt(M=1024, K=1536, N=6144, epilogue="silu", context=ctx) +op.compile() +op.get_callable()(A, op.pack_B(B), C_out) +``` + +Runs FastFlowLM's `mm.xclbin` **unmodified**. It exists so that +[`flm.GEMM`](../gemm), the IRON port of that overlay, can be measured against +what it was ported from, on identical inputs and through the same host path. +[`../gemm/benchmark.py`](../gemm/benchmark.py) does exactly that. + +**NPU2 only** — the overlay is an 8-column NPU2 binary. Constructing it +elsewhere raises `NotImplementedError`. + +## How it is obtained + +The xclbin is not checked in. It is a `RemoteFileArtifact`: downloaded on demand +into the (gitignored) build directory and pinned by SHA-256 against an immutable +FastFlowLM commit, so the fetch is reproducible and a substituted file is +rejected. + +Because this is the only thing in the tree that touches the network, the +benchmark that uses it is marked `extensive` and is not reached by the default +`-m "not extensive"` run. + +## What this operator supplies + +The overlay ships as a binary, so every core program, memtile buffer and +stream-switch route comes from the xclbin. This operator emits only the +host-side half of a dispatch: + +* **The runtime parameters.** One overlay serves every projection in a model, so + the shape, the activation and the clamp arrive as words in each core's data + memory. A core blocks on a lock until the sequence releases it, so a dispatch + that writes no parameters hangs. +* **The shim DMA transfers**, reproducing the overlay's fixed channel map. + +## Differences from `flm.GEMM` + +| | `flm.MMPrebuilt` | `flm.GEMM` | +|---|---|---| +| provenance | shipped binary, downloaded | built from source in this repo | +| devices | NPU2 only | NPU2 and NPU1 | +| `tile_n` | fixed at 128 | 64 or 128, chosen per shape and device | +| epilogue selected | at runtime, by parameter | at compile time | +| rounding | core power-up `floor` | `conv_even` by default | +| B | pre-packed bf16 | pre-packed, bfp16 on NPU2 | + +The epilogue difference is the interesting one. Selecting at runtime means one +build serves every activation; baking it in, as `flm.GEMM` does, costs a build +per activation but leaves the inner loop branch-free. The rounding difference is +why `flm.GEMM` is ~41x more accurate by default — see +[the port's README](../gemm/README.md#matching-the-shipped-fastflowlm-overlay), +which also records that `flm.GEMM(rounding="floor")` reproduces this overlay bit +for bit. diff --git a/iron/operators/flm/mm_prebuilt/design.py b/iron/operators/flm/mm_prebuilt/design.py new file mode 100644 index 0000000000..713864985d --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/design.py @@ -0,0 +1,192 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Runtime sequence for the prebuilt FastFlowLM ``mm`` overlay. + +The overlay ships as a binary xclbin, so this module emits only the host-side +half of a dispatch: the shim DMA transfers and the runtime parameters. Every +core program, memtile buffer and stream-switch route comes from the xclbin. + +Three properties of the overlay set what this sequence must do, and none of +them are visible in the xclbin: + + * **The cores read their shape from runtime parameters.** One overlay serves + every GEMM in a model, so ``K/K_TILE``, ``M`` and ``N``, the activation + and the clamp all arrive as words in each core's data memory. A core + blocks on :data:`RTP_LOCK_ID` until the sequence releases it, so a + dispatch that writes no parameters hangs. + * **The shim channel map is fixed.** A arrives on MM2S channel 0 of columns + 0, 2, 4 and 6; B on MM2S channel 1 of every column; C leaves on S2MM + channel 0 of every column. The allocations below reproduce that map. + * **B arrives pre-packed**, in the order :func:`iron.operators.flm.packing` + produces. + +``iron.operators.flm.gemm`` is a port of this overlay, so the two agree on +tiling, on the byte order of each transfer and on the packed B layout. Its own +instruction stream still cannot drive this xclbin: it writes no runtime +parameters, and its lowering puts B on MM2S channel 0 in the odd columns. +""" + +import numpy as np + +from aie.dialects import aie, aiex +from aie.dialects.aie import DMAChannelDir +from aie.extras.context import mlir_mod_ctx +from aie.ir import BF16Type, MemRefType + +from iron.operators.flm.gemm.design import ( + Epilogue, + K_TILE, + M_TILE, +) + +# The shipped overlay is a fixed 4x8 NPU2 binary built with n=128, so unlike +# flm.gemm these do NOT follow the device -- they describe the artifact. Every +# other tiling knob matches flm.gemm, whose constants are imported above. +N_TILE = 128 +COLS = 8 +ROWS = 4 +# Which shim column sources the A broadcast for each compute row. Unlike +# flm.gemm -- which lets the placer choose -- this must match the placement +# baked into the downloaded xclbin: the four A streams go to alternate columns +# so each gets its own shim MM2S path and never contends with a B fill. +A_SOURCE_COL = [2 * r for r in range(ROWS)] + +# Core data memory holding the runtime parameters, and the lock a core waits +# on before it reads them. Both are baked into the overlay's core programs. +RTP_ADDRESS = 4096 +LOCK_ADDRESS_BASE = 0x1F000 +RTP_LOCK_ID = 10 + +# Outstanding transfers per shim channel. The overlay's memtiles hold two +# objects per stream, so a third transfer would overwrite one still in use. +QUEUE_DEPTH = 2 + +MIN_M = M_TILE * ROWS +MIN_K = K_TILE + + +def mm_prebuilt(dev, M, K, N, epilogue=Epilogue.NONE, clamp=None): + """Emit the MLIR module whose runtime sequence drives the overlay. + + A is ``(M, K)`` row-major and C is ``(M, N)`` row-major, both bf16. B is + ``(K, N)`` reordered by :meth:`MMPrebuilt.pack_B`. + """ + epilogue = Epilogue(epilogue) + for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, N_TILE)): + if value % unit != 0: + raise ValueError(f"{name} ({value}) must be a multiple of {unit}") + + k_iters = K // K_TILE + m_row_blocks = M // MIN_M + # Sweeps of the whole grid, plus a trailing group of rem_blocks columns. + # The columns outside that group still receive A, because A is broadcast + # along a whole compute row and the row stalls if one column stops + # draining it. + n_full = N // (N_TILE * COLS) + rem_blocks = (N % (N_TILE * COLS)) // N_TILE + + clamp_min, clamp_max = clamp if clamp is not None else (0.0, 0.0) + parameters = [ + (RTP_ADDRESS + 0, k_iters), + (RTP_ADDRESS + 4, M), + (RTP_ADDRESS + 8, N), + (RTP_ADDRESS + 12, 0), # bias, which this operator does not expose + (RTP_ADDRESS + 16, epilogue.mode), + (RTP_ADDRESS + 20, 1 if clamp is not None else 0), + (RTP_ADDRESS + 24, int(np.float32(clamp_min).view(np.int32))), + (RTP_ADDRESS + 28, int(np.float32(clamp_max).view(np.int32))), + ] + + with mlir_mod_ctx() as ctx: + bf16 = BF16Type.get() + a_ty = MemRefType.get((M * K,), bf16) + b_ty = MemRefType.get((K * N,), bf16) + c_ty = MemRefType.get((M * N,), bf16) + + @aie.device(dev.resolve()) + def device_body(): + shim = [aie.tile(c, 0) for c in range(COLS)] + for r in range(ROWS): + aie.shim_dma_allocation( + f"A_{r}", shim[A_SOURCE_COL[r]], DMAChannelDir.MM2S, 0 + ) + for c in range(COLS): + aie.shim_dma_allocation(f"B_{c}", shim[c], DMAChannelDir.MM2S, 1) + aie.shim_dma_allocation(f"C_{c}", shim[c], DMAChannelDir.S2MM, 0) + + @aiex.runtime_sequence(a_ty, b_ty, c_ty) + def sequence(A, B, C): + # Every core gets the same parameters; the overlay derives the + # per-tile work from its own coordinates. + for row in range(2, 2 + ROWS): + for col in range(COLS): + for address, value in parameters: + aiex.npu_write32(address, value, column=col, row=row) + aiex.npu_write32( + LOCK_ADDRESS_BASE + 16 * RTP_LOCK_ID, + 1, + column=col, + row=row, + ) + + outstanding = {} + + def transfer(allocation, buffer, offset, sizes, strides): + queue = outstanding.setdefault(allocation, []) + if len(queue) == QUEUE_DEPTH: + aiex.dma_await_task(queue.pop(0)) + task = aiex.shim_dma_single_bd_task( + allocation, + buffer, + offset=offset, + sizes=sizes, + strides=strides, + issue_token=True, + ) + aiex.dma_start_task(task) + queue.append(task) + + # One transfer per (column-block, row-block, leg), matching the + # order the overlay's memtiles consume: column-block outermost, + # then row-block, then column. + for mega_col in range(n_full + (1 if rem_blocks else 0)): + active = rem_blocks if (rem_blocks and mega_col == n_full) else COLS + for mega_row in range(m_row_blocks): + for c in range(COLS): + if c in A_SOURCE_COL: + r = A_SOURCE_COL.index(c) + transfer( + f"A_{r}", + A, + mega_row * ROWS * M_TILE * K + r * M_TILE * K, + [1, k_iters, M_TILE, K_TILE], + [0, K_TILE, K, 1], + ) + if c >= active: + continue + # One contiguous run: pack_B has already put this + # column's k-blocks in the order the memtile + # writes them. + transfer( + f"B_{c}", + B, + (mega_col * COLS + c) * N_TILE * K, + [1, 1, 1, k_iters * K_TILE * N_TILE], + [0, 0, 0, 1], + ) + transfer( + f"C_{c}", + C, + mega_col * COLS * N_TILE + + mega_row * ROWS * M_TILE * N + + c * N_TILE, + [1, 1, ROWS * M_TILE, N_TILE], + [0, 0, N, 1], + ) + + for queue in outstanding.values(): + for task in queue: + aiex.dma_await_task(task) + + return str(ctx.module) diff --git a/iron/operators/flm/mm_prebuilt/op.py b/iron/operators/flm/mm_prebuilt/op.py new file mode 100644 index 0000000000..17d34bfced --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/op.py @@ -0,0 +1,176 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass, field +from typing import Any, Callable, ClassVar, Dict + +import aie.utils as aie_utils +from aie.utils.npukernel import NPUKernel + +from iron.common import ( + AIERuntimeArgSpec, + DesignGenerator, + InstsBinArtifact, + MLIROperator, + PythonGeneratedMLIRArtifact, + RemoteFileArtifact, +) + +from iron.operators.flm.packing import pack_b +from iron.operators.flm.gemm.design import Epilogue, K_TILE, S, T +from iron.operators.flm.mm_prebuilt.design import MIN_K, MIN_M, N_TILE + +# The FastFlowLM revision the overlay is taken from. A commit SHA rather than +# a branch, so the digest below stays valid. +FASTFLOWLM_COMMIT = "f81eba7140decef5e4eda670d02a91b9d6402ee9" +XCLBIN_PATH = "src/xclbins/Gemma4-E4B-IT-NPU2/mm.xclbin" +XCLBIN_URL = ( + f"https://raw.githubusercontent.com/ROCm/FastFlowLM/{FASTFLOWLM_COMMIT}/" + f"{XCLBIN_PATH}" +) +XCLBIN_SHA256 = "6f1e5507b84d4545536c9b8281002d0e0e10ed241f8593cb4db50eee63876e5f" +XCLBIN_KERNEL_NAME = "MLIR_AIE" + +# The overlay's k slice. Its B layout is fixed by the shipped binary, so unlike +# flm.gemm this is not a tuning knob. +CT_K = K_TILE + + +@dataclass +class MMPrebuilt(MLIROperator): + """bf16 GEMM running FastFlowLM's shipped ``mm`` overlay unmodified. + + The overlay is downloaded rather than built: it exists only as a binary + xclbin. This operator supplies the other half of a dispatch -- the runtime + parameters and the shim DMA transfers -- so that the shipped kernel can be + measured against :class:`iron.operators.flm.GEMM`, the IRON port of it, at + the same shapes and on the same inputs. + + NPU2 only: the overlay is built for the 8-column grid. + + B must be pre-packed; use :meth:`pack_B`. + + Note the epilogue here is selected through a RUNTIME parameter, because one + overlay serves every projection in a model. ``flm.GEMM`` bakes it in at + compile time instead, which is what lets its inner loop be branch-free; the + cost is one build per activation rather than one build for all of them. + """ + + M: int + K: int + N: int + # Activation, selected through a runtime parameter rather than at compile + # time as in flm.GEMM. + epilogue: Epilogue = Epilogue.NONE + # Optional (min, max) applied after the activation. + clamp: tuple[float, float] | None = None + context: object = field(default=None, repr=False) + + _name_aliases: ClassVar[Dict[str, str]] = { + **MLIROperator._name_aliases, + "epilogue": "epi", + } + + def __post_init__(self): + for name, value, unit in ( + ("M", self.M, MIN_M), + ("K", self.K, MIN_K), + ("N", self.N, N_TILE), + ): + if value % unit != 0: + raise ValueError(f"{name} ({value}) must be a multiple of {unit}") + self.epilogue = Epilogue(self.epilogue) + if self.clamp is not None and self.clamp[0] > self.clamp[1]: + raise ValueError( + f"clamp min ({self.clamp[0]}) must be <= max ({self.clamp[1]})" + ) + device = aie_utils.get_current_device() + if device.resolve().name != "npu2" or device.cols < 8: + raise NotImplementedError( + "flm.MMPrebuilt runs a prebuilt NPU2 overlay and needs the 8 " + f"columns of NPU2 (aie2p); got {device.resolve().name!r} with " + f"{device.cols} columns" + ) + + MLIROperator.__init__(self, context=self.context) + + @property + def name(self) -> str: + """Artifact stem. Prefixed for the same reason as flm.GEMM's.""" + return f"FLM_{super().name}" + + def get_mlir_artifact(self): + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator( + self.operator_dir / "design.py", + "mm_prebuilt", + (), + { + "dev": aie_utils.get_current_device(), + "M": self.M, + "K": self.K, + "N": self.N, + "epilogue": self.epilogue, + "clamp": self.clamp, + }, + ), + ) + + def get_kernel_artifacts(self): + # None to build: every core program is inside the downloaded xclbin. + return [] + + def set_up_artifacts(self) -> None: + mlir_artifact = self.get_mlir_artifact() + self.insts_artifact = InstsBinArtifact( + f"{self.name}.bin", + mlir_input=mlir_artifact, + dependencies=[mlir_artifact], + ) + self.xclbin_artifact = RemoteFileArtifact( + f"flm_mm_{FASTFLOWLM_COMMIT[:8]}.xclbin", + url=XCLBIN_URL, + sha256=XCLBIN_SHA256, + ) + self.add_artifacts([self.insts_artifact, self.xclbin_artifact]) + + def get_callable(self) -> Callable[..., Any]: + npu_kernel = NPUKernel( + xclbin_path=self.xclbin_artifact.filename, + kernel_name=XCLBIN_KERNEL_NAME, + insts_path=self.insts_artifact.filename, + ) + handle = aie_utils.DefaultNPURuntime.load(npu_kernel) + + def call(*args): + return aie_utils.DefaultNPURuntime.run(handle, list(args)) + + return call + + def pack_B(self, B): + """Reorder a row-major ``(K, N)`` weight matrix into the order the B + transfers read. Returns a flat bf16 tensor. + + NOT the same layout ``flm.GEMM.pack_B`` produces: the overlay's own + loop nest sweeps the two within-block k axes in the opposite order + from ``mm_fused_mmul_2x2``'s, so this needs ``overlay_order`` -- see + ``pack_b``'s docstring. + """ + return pack_b( + B, k_tile=K_TILE, n_tile=N_TILE, s=S, t=T, ct_k=CT_K, overlay_order=True + ) + + def get_arg_spec(self): + return [ + AIERuntimeArgSpec("in", (self.M, self.K)), # A + # B, pre-packed by pack_B -- same element count, different order. + AIERuntimeArgSpec("in", (self.K, self.N)), # B (weights) + AIERuntimeArgSpec("out", (self.M, self.N)), # C + ] + + def reference(self, A, B): + """CPU reference: ``C = epilogue(A @ B)``.""" + from iron.operators.flm.gemm.reference import reference + + return reference(A, B, self.epilogue, self.clamp) diff --git a/iron/operators/flm/mm_prebuilt/test.py b/iron/operators/flm/mm_prebuilt/test.py new file mode 100644 index 0000000000..805a66c783 --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/test.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Correctness for :class:`iron.operators.flm.MMPrebuilt`. + +Extensive only: constructing the operator downloads FastFlowLM's shipped +``mm.xclbin`` over the network (see ``op.py``), and a developer running the +default operator suite should not trip that download. Unlike +``iron/operators/flm/gemm/benchmark.py`` -- which times this operator against +``flm.GEMM`` and IRON's ``GEMM`` at production shapes but is never collected, +named as it is -- this module IS named ``test.py``, so the extensive CI job +actually runs it. +""" + +import numpy as np +import pytest +import torch + +import aie.utils as aie_utils + +from iron.common.test_utils import run_test +from iron.operators.flm.gemm.reference import ( + apply_epilogue, + generate_golden_reference, +) +from iron.operators.flm.gemm.design import Epilogue +from iron.operators.flm.mm_prebuilt.op import MMPrebuilt + +NONE, GELU, SILU, SIGMOID = Epilogue + +# Largest |d/dx| of each epilogue, used to carry the accumulator's error bound +# through to the output. sigmoid's is exactly 1/4; silu and gelu both peak at +# 1.0998 (gelu here being the x*sigmoid(1.702x) approximation the overlay +# implements, whose derivative happens to share silu's maximum), rounded up. +MAX_SLOPE = {NONE: 1.0, SIGMOID: 0.25, SILU: 1.1, GELU: 1.1} + +pytestmark = pytest.mark.extensive + +_dev = aie_utils.get_current_device() +if _dev.resolve().name != "npu2" or _dev.cols < 8: + pytest.skip( + "the prebuilt FastFlowLM overlay is an 8-column NPU2 binary; " + f"this device is {_dev.resolve().name!r} with {_dev.cols} columns", + allow_module_level=True, + ) + +# The overlay never calls set_rounding, so it runs in the core's power-up floor +# mode and carries a ~1% truncation bias -- not a bug. See gemm/benchmark.py. +BUDGET_FLOOR = 2e-2 + + +@pytest.mark.parametrize( + "M,K,N,epilogue,clamp", + [ + (256, 512, 1024, NONE, None), # exactly one full 8-column sweep + (512, 1024, 2048, NONE, None), # two full sweeps + (256, 512, 640, NONE, None), # remainder only: 5 of 8 cols + (256, 512, 1280, NONE, None), # full sweep + remainder: 1 of 8 cols + (256, 512, 1024, SILU, None), + (256, 512, 1024, GELU, None), + ], +) +def test_mm_prebuilt(M, K, N, epilogue, clamp, aie_context): + golden_ref = generate_golden_reference( + M=M, K=K, N=N, epilogue=epilogue, clamp=clamp + ) + + operator = MMPrebuilt( + M=M, K=K, N=N, epilogue=epilogue, clamp=clamp, context=aie_context + ) + + input_buffers = { + "A": golden_ref["input"].flatten(), + # B is consumed pre-packed; see MMPrebuilt.pack_B. + "B": operator.pack_B(golden_ref["input_b"]), + } + output_buffers = {"C": golden_ref["output"].flatten()} + + # The overlay's error is made in the ACCUMULATOR -- it runs in the core's + # power-up floor rounding, worth about BUDGET_FLOOR of the accumulated mass + # -- and the epilogue then maps that accumulator through an activation. So + # the output bound is the accumulator bound carried through the activation, + # |f(x+e) - f(x)| <= max|f'| * |e|, rather than a tolerance invented in the + # output domain. + # + # Only the unbounded epilogues are checked this way. For sigmoid and clamp + # no bound over this reference can be both correct and useful -- the + # accumulator error alone exceeds their whole output range -- so they are + # covered functionally by test_mm_prebuilt_epilogue_matches_accumulator. + mass = float( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + abs_tol = MAX_SLOPE[epilogue] * BUDGET_FLOOR * mass + errors, latency_us, bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=0.04, + abs_tol=abs_tol, + ) + assert not errors, "Test failed" + + +@pytest.mark.parametrize( + "epilogue,clamp", + [ + (SIGMOID, None), + (NONE, (-2.0, 2.0)), + (SILU, None), + (GELU, None), + ], +) +def test_mm_prebuilt_epilogue_matches_accumulator(epilogue, clamp, aie_context): + """The epilogue is the right function of the accumulator the device produced. + + Checking a bounded epilogue against the idealized CPU reference cannot work. + The overlay accumulates in the core's power-up floor rounding, worth ~2% of + the accumulated mass, which here is ~65 -- larger than sigmoid's entire (0,1) + range and than this clamp's (-2, 2). Any bound wide enough to admit that + accumulator error also admits an all-zero result, and any bound tight enough + to reject all-zeros also rejects correct hardware. That is why the earlier + flat tolerance failed on working arithmetic. + + So compare the epilogue against the device's OWN accumulator instead: run + the same inputs with no epilogue, apply the activation and clamp to that on + the host, and require the epilogue build to agree. The accumulator error is + then common to both sides and cancels, leaving only the epilogue under test. + An all-zero result still fails, because the reference side is not zero. + """ + M, K, N = 256, 512, 1024 + # A small input scale keeps the accumulator in the range where these curves + # are actually curved; at the default scale the product lands around +-900, + # where gelu and silu are indistinguishable from the identity. + golden_ref = generate_golden_reference(M=M, K=K, N=N, scale=0.5) + A = golden_ref["input"] + B = golden_ref["input_b"] + + def run(epi, clm): + op = MMPrebuilt(M=M, K=K, N=N, epilogue=epi, clamp=clm, context=aie_context) + op.compile() + tensor = aie_utils.DEFAULT_TENSOR_CLASS + out = tensor((M, N), dtype=np.dtype("bfloat16")) + op.get_callable()( + tensor.from_torch(A.flatten()), tensor.from_torch(op.pack_B(B)), out + ) + return out.to_torch().reshape(M, N).float() + + acc = run(NONE, None) + got = run(epilogue, clamp) + expected = apply_epilogue(acc, epilogue, clamp) + + # Both sides see the same accumulator, so what is left is the epilogue. + # Two terms, and they are different in kind. + # + # The bf16 term is per element rather than one global number: the + # accumulator read back is bf16, good to ~2^-8 RELATIVELY, and clamp is only + # sensitive near its boundary, so a tolerance taken from the accumulator's + # largest magnitude would be wider there than the clamp range itself -- i.e. + # vacuous. + # + # The activation term covers what bf16 rounding does NOT explain. Checked by + # bounding the true accumulator to its bf16 rounding interval and evaluating + # the epilogue across it: clamp lands inside for all 262144 elements, but + # sigmoid, silu and gelu land outside for about half, by up to 0.018. That + # residual is the overlay's own activation approximation -- a LUT or native + # instruction, not exact math -- which no reference built on torch.sigmoid + # can reproduce. 0.05 is ~3x the measured worst case and still ~20x below + # where the bound would go vacuous; the assertion at the end pins that down. + approx = 0.0 if epilogue is NONE else 0.05 + tol = MAX_SLOPE[epilogue] * acc.abs() * 2.0**-8 + 2.0**-8 + approx + err = (got - expected).abs() + over = err > tol + assert not over.any(), ( + f"{epilogue} clamp={clamp}: {int(over.sum())} of {over.numel()} elements " + f"differ from epilogue(device accumulator) by more than the bf16 bound; " + f"worst {float((err - tol).max()):.4f} over" + ) + # The bound must not be wide enough to admit a dead device. + assert ( + expected.abs() > tol + ).any(), f"{epilogue}: tolerance is vacuous -- an all-zero result would pass" diff --git a/iron/operators/flm/packing.py b/iron/operators/flm/packing.py new file mode 100644 index 0000000000..538ecfeb3f --- /dev/null +++ b/iron/operators/flm/packing.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Weight packing shared by the FastFlowLM-derived operators. + +Both ``flm.GEMM`` and ``flm.MMPrebuilt`` consume B pre-packed into the order +the compute tiles read it, so their B transfers are plain linear descriptors. +The reorder is deliberately the caller's job: expressing it as a strided +descriptor over an unpacked B leaves an innermost run of ``t`` bf16 values, so +each transfer becomes thousands of scattered bursts -- measured 5.4x slower end +to end. Weights are packed once and reused across dispatches, so the cost +belongs here. +""" + +import numpy as np +import torch + + +def f32_to_bfp16ebs8(a, round_conv_even=True): + """float32 -> bfp16ebs8, matching the hardware's to_v64bfp16ebs8. + + Blocks of 8 share the max f32 exponent in the block; each mantissa is the + 24-bit magnitude with the implicit bit made explicit, shifted right by + 17 + (maxExp - exp) to land on the shared exponent. + + That shift OBEYS THE CORE'S ROUNDING MODE. mlir-aie's reference + ``floatToBfp16`` (``programming_examples/ml/block_datatypes/helper.h``) + hardcodes truncation and says AIE2P always truncates -- true only of the + power-up ``floor`` mode. The operators here call ``set_rounding(conv_even)``, + so the kernel's own conversion rounds to nearest with ties to even, and + matching it here is what makes packing B on the host numerically free. + Measured on hardware: 14.9375 -> 15 (rounds up) while 106.5 -> 106 and + 94.5 -> 94 (ties to even), which truncation cannot produce. + + Layout per block: one shared-exponent byte then the 8 mantissa bytes. + """ + flat = np.ascontiguousarray(a, dtype=np.float32).reshape(-1, 8) + u = flat.view(np.uint32) + sign = (u & 0x80000000) != 0 + exp = ((u >> 23) & 0xFF).astype(np.int32) + man = (u & 0x007FFFFF).astype(np.uint32) + man = np.where(exp != 0, man | 0x00800000, man).astype(np.uint32) + max_exp = exp.max(axis=1, keepdims=True) + # signed magnitude; rounding below must see the sign to tie correctly + mag = np.where(sign, -man.astype(np.int64), man.astype(np.int64)) + # The two shifts compose: 17 to keep 7 mantissa bits plus the sign, then + # (maxExp - exp) to bring the value onto the block's shared exponent. + shift = (max_exp - exp).astype(np.int64) + total = np.clip(17 + shift, 0, 62) + if round_conv_even: + # np.rint is round-half-to-even. man < 2**24 and the divisor is a power + # of two, so the quotient is exact in float64 and the only rounding is + # the intended one. + v8 = np.rint(mag.astype(np.float64) / np.exp2(total.astype(np.float64))) + else: + v8 = mag >> total + v8 = np.where(shift >= 32, np.where(sign, -1, 0), v8) + # Rounding can carry the block's largest magnitude from 127 to 128, which + # does not fit the signed 8-bit mantissa; saturate rather than wrap. + v8 = np.clip(v8, -128, 127) + out = np.empty((flat.shape[0], 9), dtype=np.uint8) + out[:, 0] = max_exp[:, 0].astype(np.uint8) + out[:, 1:] = v8.astype(np.int8).view(np.uint8) + return torch.from_numpy(out.reshape(-1)) + + +def pack_b( + B, + k_tile, + n_tile, + s, + t, + ct_k, + bfp16=False, + round_conv_even=True, + overlay_order=False, +): + """Reorder a row-major ``(K, N)`` weight matrix into consumption order. + + ``s``/``t`` are the mmul's register tiling and ``ct_k`` the k slice one + compute tile holds at a time; all three set the blocked layout, and packing + with the wrong value produces a wrongly ordered buffer of the RIGHT SIZE, so + it mis-computes silently rather than raising. + + With ``bfp16`` the result is a flat uint8 tensor of bfp16ebs8 blocks (9 + bytes per 8 values); otherwise a flat bf16 tensor. The quantization is not a + loss this adds: the AIE2P mmul only multiplies bfp16, so the bf16 path would + convert B inside every mac anyway. Doing it here hoists a rounding that + already happened, and makes B 9 bytes per 8 values instead of 16. + + ``overlay_order`` swaps the two within-block k axes (``i`` and ``s_in`` + below). It exists solely for :class:`iron.operators.flm.MMPrebuilt`, whose + B stream is read by FastFlowLM's shipped ``mm.xclbin``, not by a kernel + built here: that overlay's own loop nest sweeps ``s_in`` outer and ``i`` + inner, the reverse of ``mm_fused_mmul_2x2``'s ``i``-outer loop. The + now-deleted ``flm_pack_B`` (see ``bench_vs_flm.py`` history) verified this + ordering against the overlay via ``tile.reshape(...).transpose(2, 1, 0, + 3)``; incompatible with ``bfp16``, which only the IRON-built kernel uses. + """ + if overlay_order and bfp16: + raise ValueError("overlay_order is bf16-only; the overlay never takes bfp16 B") + K, N = B.shape + if K % k_tile or N % n_tile: + raise ValueError(f"B ({K}, {N}) must tile to ({k_tile}, {n_tile}) to be packed") + col_a = ct_k // s + blocked = B.reshape( + K // k_tile, k_tile // ct_k, col_a, s, N // n_tile, n_tile // t, t + ) + # (kb, kslice, i, s_in, cb, tb, t_in) + if not bfp16: + if overlay_order: + # -> (cb, kb, kslice, tb, s_in, i, t_in) + out = blocked.permute(4, 0, 1, 5, 3, 2, 6).reshape(-1).contiguous() + else: + # -> (cb, kb, kslice, tb, i, s_in, t_in) + # Row-major s x t within the block, which is what the plain mmul + # loads. + out = blocked.permute(4, 0, 1, 5, 2, 3, 6).reshape(-1).contiguous() + # Callers may pass B in whatever dtype they have it in (e.g. a model's + # native f32 weight); the kernels and get_arg_spec() assume the result + # is bf16, so guarantee that here rather than silently returning + # whatever B.dtype was. + return out.to(torch.bfloat16) + # -> (cb, kb, kslice, tb, i, t_in, s_in) + # t-major within the block: the mixed mmul hands B straight to + # mac_8x8_8x8T without the transpose the bf16 form applies, so the transpose + # happens here instead. It also puts the 8 values that share a bfp16 + # exponent (8 consecutive k for one n) adjacent, which is what makes the + # block grouping match the kernel's. Grouping over n instead measures + # 1.95e-02 against this layout's 2.69e-04. + blocked = blocked.permute(4, 0, 1, 5, 2, 6, 3).reshape(-1, 8).contiguous() + return f32_to_bfp16ebs8(blocked.float().numpy(), round_conv_even=round_conv_even) + + +def packed_b_size(K, N, bfp16): + """Elements (bf16) or bytes (bfp16ebs8) that ``pack_b`` returns.""" + return K * N // 8 * 9 if bfp16 else K * N diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index c8a4f6e7c1..ab390d1cac 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -20,7 +20,7 @@ str_to_dtype, ) from aie.iron.device import NPU1Col1, NPU1Col2, NPU1, NPU2, Tile -from aie.helpers.taplib import TensorAccessSequence, TensorTiler2D, TensorAccessPattern +from aie.helpers.taplib import TensorTiler2D, TensorAccessPattern from aie.iron.controlflow import range_ from iron.operators._trace import maybe_enable_trace @@ -77,12 +77,6 @@ def main(): default="bf16", ) argparser.add_argument("--trace_size", type=int, default=0) - argparser.add_argument( - "--generate-taps", - action="store_true", - help="Generate TensorAccessPatterns, a Python object to represent each data transfer" - "of the input/output matrices. These objects can be used for visualization.", - ) argparser.add_argument( "--output-file-path", "-o", @@ -91,7 +85,7 @@ def main(): ) args = argparser.parse_args() - maybe_module = my_matmul( + module = my_matmul( args.dev, args.M, args.K, @@ -111,16 +105,11 @@ def main(): args.trace_size, args.archive, "", - args.generate_taps, ) - if args.generate_taps: - return maybe_module - else: - output_file_path = Path(args.output_file_path) - - with open(output_file_path, "w") as f: - f.write(str(maybe_module)) + output_file_path = Path(args.output_file_path) + with open(output_file_path, "w") as f: + f.write(str(module)) def ceildiv(a, b): @@ -147,7 +136,6 @@ def my_matmul( trace_size, kernel_object=None, func_prefix="", - generate_taps=False, ): n_aie_rows = 4 @@ -172,6 +160,17 @@ def my_matmul( mem_tile_m_C = m * n_aie_rows mem_tile_n = n * n_aie_cols + # A shim BD's outermost descriptor dimension lands in the ITERATION field, + # whose step is 20 bits wide (AIETargetModel::getDmaBdStepBits for + # ShimNOCTile). An element stride S is re-expressed as (S - 1) * itemsize + # / 4-byte address granularity before the check, so a wide N pushes C's row + # stride past it: M=1024 K=2560 N=10240 needs mem_tile_m_C * N = 2621440 + # and aiecc rejects the build with "Stride 3 exceeds the [1:1048576] + # range". See the C drain below for how that is split, and flm_gemm's + # design.py for the same fix worked through in more detail. + def _hw_stride_ok(stride_elems, itemsize): + return (stride_elems - 1) * itemsize // 4 <= (1 << 20) - 1 + if prio_accuracy: assert ( dtype_out_str == "bf16" @@ -262,12 +261,6 @@ def my_matmul( else: dev_ty = NPU2() - # These will hold TensorAccessPattern objects that represent the runtime - # npu_dma_memcpy_nd operations of this design. They are only used if generate_taps is true - A_taps = [] - B_taps = [] - C_taps = [] - # Define tensor types A_ty = np.ndarray[(M * K,), np.dtype[dtype_in]] B_ty = np.ndarray[(K * N,), np.dtype[dtype_in]] @@ -597,39 +590,69 @@ def sequence(A, B, C, A_prods, B_prods, C_conses): # | | # | | # ---------------- + # Normally one descriptor walks all current_tb_n_rows + # row-blocks. When that outermost stride overflows the + # shim's 20-bit iteration step (see _hw_stride_ok + # above), issue one descriptor per row-block instead, + # carrying the row jump in the OFFSET -- which has no + # such limit -- and leaving the outer dimension + # degenerate. Same bytes, same order, same number of + # objects; only the descriptor is reshaped. + # + # These extra tasks are safe against the two shim + # limits neither the toolchain nor the verifier models. + # BD ids: all of a (tb, pingpong) iteration's tasks stay + # live until tg.finish() below, so they stay distinct -- + # 2 iterations x (2 C + 2 A + 2 B) = 12 of 16. Channel + # task queue: the C channel goes from 2 outstanding to + # current_tb_n_rows x 2 = 4, which is where A and B + # already sit. + C_rows = [(row_base, current_tb_n_rows)] if not c_col_maj: - C_row_offset = row_base * mem_tile_m_C * N - C_col_offset = col * n - C_offset = C_col_offset + C_row_offset - C_sizes = [ - current_tb_n_rows, - N // mem_tile_n, - mem_tile_m_C, - n, - ] - C_strides = [mem_tile_m_C * N, mem_tile_n, N, 1] - else: - C_row_offset = row_base * mem_tile_m_C - C_col_offset = col * n * M - C_offset = C_col_offset + C_row_offset - C_sizes = [N // mem_tile_n, n_aie_rows, n, m] - C_strides = [M * mem_tile_n, m, M, 1] - C_tile = TensorAccessPattern( - (N, M) if c_col_maj else (M, N), - offset=C_offset, - sizes=C_sizes, - strides=C_strides, - ) + row_stride = mem_tile_m_C * N + if current_tb_n_rows > 1 and not _hw_stride_ok( + row_stride, np.dtype(dtype_out).itemsize + ): + C_rows = [ + (row_base + r, 1) for r in range(current_tb_n_rows) + ] - # This line does not change MLIR output at all - it's just for recording data movement - C_taps.append(C_tile) + for c_row_base, c_n_rows in C_rows: + if not c_col_maj: + C_row_offset = c_row_base * mem_tile_m_C * N + C_col_offset = col * n + C_offset = C_col_offset + C_row_offset + C_sizes = [ + c_n_rows, + N // mem_tile_n, + mem_tile_m_C, + n, + ] + C_strides = [ + mem_tile_m_C * N if c_n_rows > 1 else 0, + mem_tile_n, + N, + 1, + ] + else: + C_row_offset = c_row_base * mem_tile_m_C + C_col_offset = col * n * M + C_offset = C_col_offset + C_row_offset + C_sizes = [N // mem_tile_n, n_aie_rows, n, m] + C_strides = [M * mem_tile_n, m, M, 1] + C_tile = TensorAccessPattern( + (N, M) if c_col_maj else (M, N), + offset=C_offset, + sizes=C_sizes, + strides=C_strides, + ) - C_conses[col].drain( - C, - tap=C_tile, - wait=True, - group=tg, - ) + C_conses[col].drain( + C, + tap=C_tile, + wait=True, + group=tg, + ) for tile_row in range(current_tb_n_rows): if separate_c_tiles: @@ -683,9 +706,6 @@ def sequence(A, B, C, A_prods, B_prods, C_conses): wait=True, group=tg, ) - # This line does not change MLIR output at all - it's just for recording data movement - C_taps.append(C_tile) - # A input transfer: # # The smallest transfer unit is a (m*n_A_tiles_per_shim)-sized sub-tile of the input matrix. @@ -742,10 +762,6 @@ def sequence(A, B, C, A_prods, B_prods, C_conses): tap=B_tiles[col], group=tg, ) - - # These lines do not change MLIR output at all - they are just for recording data movement - A_taps.append(A_tiles[tile_offset]) - B_taps.append(B_tiles[col]) if tb > 0 or (tb == 0 and pingpong > 0): tg.finish() tg = TaskGroup() @@ -771,20 +787,7 @@ def sequence(A, B, C, A_prods, B_prods, C_conses): maybe_enable_trace(my_program, trace_size, workers) # Place components (assign them resources on the device) and generate an MLIR module. - # This is what runs the sequence body, so it must happen before the taps it - # records are read. - module = my_program.resolve_program() - - if generate_taps: - # If generate taps is true, return a representation of tensor access patterns - # representing all the npu_dma_memcpy_nd runtime sequence operations per input/ouput tensor. - return ( - TensorAccessSequence.from_taps(A_taps), - TensorAccessSequence.from_taps(B_taps), - TensorAccessSequence.from_taps(C_taps), - ) - - return module + return my_program.resolve_program() if __name__ == "__main__": diff --git a/iron/operators/gemm/op.py b/iron/operators/gemm/op.py index 58b5f14b21..cfeddbe949 100644 --- a/iron/operators/gemm/op.py +++ b/iron/operators/gemm/op.py @@ -105,7 +105,6 @@ def get_mlir_artifact(self): "prio_accuracy": self.prio_accuracy, "separate_c_tiles": int(self.separate_c_tiles), "trace_size": 0, - "generate_taps": False, "kernel_object": f"gemm_{self.tile_m}x{self.tile_k}x{self.tile_n}_{int(self.b_col_maj)}_{int(self.c_col_maj)}{self._kernel_flags_suffix}.o", }, ), diff --git a/iron/operators/gemm/test.py b/iron/operators/gemm/test.py index bbd41b00ae..100b9c2ca9 100755 --- a/iron/operators/gemm/test.py +++ b/iron/operators/gemm/test.py @@ -49,6 +49,11 @@ def get_params(): (2048, 8192, 2048, 2, False, True, 64, 64, 64, 0, 1), (2048, 64, 2048, 2, False, True, 64, 64, 64, 0, 1), (2048, 64, 8192, 2, False, True, 64, 64, 64, 0, 1), + # N wide enough that C's row stride (mem_tile_m_C * N) overflows the + # shim BD's 20-bit iteration step, so the drain is issued as one + # descriptor per row-block. Cover for that split. + (1024, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), + (2048, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), ] # fmt: on diff --git a/requirements.txt b/requirements.txt index 1558581d35..7f7c43b370 100755 --- a/requirements.txt +++ b/requirements.txt @@ -13,8 +13,8 @@ --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.3.dev60+gc80b88c -llvm-aie==22.0.0.2026082001+84660bc3 +mlir_aie==1.4.3.dev85+gdf48abc +llvm-aie==22.0.0.2026090701+3e93bf7b black reuse