From 18ef81cdf078a2c0afd98dbfa90050c65929ef8d Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 14:34:51 -0600 Subject: [PATCH 01/31] Port the FLM bf16 GEMM as the flm_gemm operator Adds a second GEMM design, distinct from the existing gemm operator rather than a retuning of it: A is broadcast along each compute row from four shim columns, C is joined at the memtile, the mmul keeps A in a single ObjectFifo object, and an activation + clamp are fused into the C drain. Ported from an inert-overlay design whose host wrote shapes into an RTP buffer per dispatch. With M/K/N static per IRON operator that machinery is gone: the trip counts are compile-time constants, so the counts kernel, the RTP buffer and its sync lock, and the N-remainder drain path all disappear. The ObjectFifo dims are carried over verbatim and verified to match the original's MLIR exactly. Three things had to change to make it correct here: - B's fill descriptor reorders on the fly. The memtile expects B in t-block-major (n//T, k%S, k//S, n%T) order; the original relied on its host pre-packing the weights that way. Doing it in the descriptor keeps B an ordinary dense (K, N) tensor. - The kernel now sets rounding to conv_even. The core powers up in floor, and truncation biases every conversion the same direction, so the error accumulates coherently over K instead of cancelling: 1% of accumulated mass versus 0.042%, a 24x difference. The original never set it. - A is filled one k-block at a time and the fills are retired in batches. A single BD spanning every k-block stalls mid-transfer once k_iters exceeds the fifo depth while still holding its shim channel, and a shim tile allows only 16 active BDs -- which is why this only showed up at larger K. Accuracy now matches the existing GEMM run in the same bfp16-emulated mode (mean error 0.00042 of accumulated mass vs 0.00044). The r=8 mmul shape only exists on that emulated path, so the test bounds error against the accumulated mass rather than elementwise-relatively, which a K-term signed sum cancelling ~sqrt(K) makes meaningless. Co-Authored-By: Claude --- aie_kernels/aie2p/flm_gemm.cc | 85 ++++++ aie_kernels/aie2p/flm_gemm_epilogue.cc | 96 +++++++ aie_kernels/aie2p/flm_gemm_geometry.h | 39 +++ aie_kernels/aie2p/flm_gemm_mmul.h | 182 ++++++++++++ aie_kernels/aie2p/nonlut_based_ops.h | 66 +++++ iron/operators/__init__.py | 1 + iron/operators/flm_gemm/design.py | 365 +++++++++++++++++++++++++ iron/operators/flm_gemm/op.py | 180 ++++++++++++ iron/operators/flm_gemm/reference.py | 61 +++++ iron/operators/flm_gemm/test.py | 104 +++++++ 10 files changed, 1179 insertions(+) create mode 100644 aie_kernels/aie2p/flm_gemm.cc create mode 100644 aie_kernels/aie2p/flm_gemm_epilogue.cc create mode 100644 aie_kernels/aie2p/flm_gemm_geometry.h create mode 100644 aie_kernels/aie2p/flm_gemm_mmul.h create mode 100644 aie_kernels/aie2p/nonlut_based_ops.h create mode 100644 iron/operators/flm_gemm/design.py create mode 100644 iron/operators/flm_gemm/op.py create mode 100644 iron/operators/flm_gemm/reference.py create mode 100644 iron/operators/flm_gemm/test.py diff --git a/aie_kernels/aie2p/flm_gemm.cc b/aie_kernels/aie2p/flm_gemm.cc new file mode 100644 index 000000000..11319f0a1 --- /dev/null +++ b/aie_kernels/aie2p/flm_gemm.cc @@ -0,0 +1,85 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Compute kernel for the flm_gemm operator: a bf16 GEMM over a fixed 4x8 grid +// of compute tiles, where each tile owns an m x n slice of C and accumulates +// over K in f32. +// +// Split into per-step entry points rather than one entry point owning the whole +// loop nest. The loop nest lives in the design's core body instead +// (iron/operators/flm_gemm/design.py), which is what gives each level of it an +// ObjectFifo acquire point -- a fifo needs its consumer to acquire once per +// object, so a single entry point spanning a whole dispatch could not be fed by +// one. +// +// Tile geometry comes from the design as -D flags so design.py stays the single +// source of truth; the design's own buffer sizes and unroll factors are derived +// from the same constants. +#include "flm_gemm_mmul.h" +#include "zero.cc" +#include +#include + +#if !defined(FLM_GEMM_TILE_M) || !defined(FLM_GEMM_TILE_K) || \ + !defined(FLM_GEMM_TILE_N) +#error "design.py must pass -DFLM_GEMM_TILE_M / _TILE_K / _TILE_N" +#endif + +namespace { +constexpr int M = FLM_GEMM_TILE_M; +constexpr int K = FLM_GEMM_TILE_K; +constexpr int N = FLM_GEMM_TILE_N; +constexpr int R = 8; // register tiling r/s/t +constexpr int S = 8; +constexpr int T = 8; + +// How much of K one compute tile holds at a time, given the n width. +constexpr int CT_K = compute_CT_k_max_n(); +static_assert(CT_K > 0, "no K-blocking geometry for this tile_n"); + +static_assert(M % (2 * R) == 0, "tile_m 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 with rounding_mode::floor. Truncation biases every +// operand conversion the same direction, so the error accumulates coherently +// over the K reduction instead of cancelling -- measured as a ~1% bias in the +// result, ~20x worse than round-to-nearest-even, which is far more than the +// bfp16 emulation itself costs. Set it explicitly in every entry point that +// converts (the mmul here, and the f32->bf16 store in the epilogue). +#ifdef FLM_GEMM_ROUND_FLOOR +constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; +#else +constexpr aie::rounding_mode round_mode = aie::rounding_mode::conv_even; +#endif +} // namespace + +extern "C" { + +// Zero the f32 accumulator. Called once per mega-block-row, before the k loop +// starts accumulating into it. +// +// ADD_BIAS is deliberately not supported: initialising the accumulator from a +// bias vector would mean consuming an extra object through the same handshake +// the B ObjectFifo now owns, which would desynchronise that fifo and hang +// rather than silently mis-compute. +void flm_gemm_acc_init(float *y_acc) { + // zero_vectorized brackets itself in event0/event1 for tracing. + zero_vectorized(y_acc); +} + +// One l-step of a k iteration: one B chunk multiplied against one A object, +// accumulated into y_acc. +// +// The l loop lives in the core body so that each B chunk gets its own acquire +// point. A is a single object spanning every z slice of the mmul, so this takes +// no locks -- the A and B fifos own that handshake. +void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc) { + ::aie::set_rounding(round_mode); + constexpr int NUM_ITER = K / CT_K; + flm_gemm_mmul_2x2( + a_buf, b_buf, y_acc); +} +} diff --git a/aie_kernels/aie2p/flm_gemm_epilogue.cc b/aie_kernels/aie2p/flm_gemm_epilogue.cc new file mode 100644 index 000000000..5fee55c47 --- /dev/null +++ b/aie_kernels/aie2p/flm_gemm_epilogue.cc @@ -0,0 +1,96 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// flm_gemm's output stage: convert one chunk of the f32 accumulator into a bf16 +// output 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 every instantiation of this kernel +// has a branch-free inner loop. That differs from the design this was ported +// from, where one overlay served every activation and had to test a runtime +// mode word per chunk. +// +// This is a separate translation unit from flm_gemm.cc so it can be built with +// its own flags and, if the per-call overhead ever shows up in a trace, be +// switched to an inlined LLVM-IR kernel independently of the much larger mmul. +#include "../aie_kernel_utils.h" +#include "nonlut_based_ops.h" +#include +#include + +#if !defined(FLM_GEMM_OUT_CHUNK) || !defined(FLM_GEMM_C_DEPTH) +#error "design.py must pass -DFLM_GEMM_OUT_CHUNK / -DFLM_GEMM_C_DEPTH" +#endif + +// 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid +#ifndef FLM_GEMM_EPILOGUE_MODE +#define FLM_GEMM_EPILOGUE_MODE 0 +#endif +#ifndef FLM_GEMM_CLAMP +#define FLM_GEMM_CLAMP 0 +#endif +#ifndef FLM_GEMM_CLAMP_MIN +#define FLM_GEMM_CLAMP_MIN 0.0f +#endif +#ifndef FLM_GEMM_CLAMP_MAX +#define FLM_GEMM_CLAMP_MAX 0.0f +#endif + +namespace { +constexpr int CHUNK = FLM_GEMM_OUT_CHUNK; +constexpr int DEPTH = FLM_GEMM_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"); +} // namespace + +extern "C" { + +// Chunk (outer * DEPTH + half) of the accumulator -> one C object. +// +// 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 flm_gemm_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, + int32_t half) { + // The f32 -> bf16 store below is a conversion, so it depends on the rounding + // mode just as the mmul does; the core default is floor. See flm_gemm.cc. +#ifdef FLM_GEMM_ROUND_FLOOR + ::aie::set_rounding(aie::rounding_mode::floor); +#else + ::aie::set_rounding(aie::rounding_mode::conv_even); +#endif + const float *__restrict src = y_acc + (outer * DEPTH + half) * CHUNK; + +#if FLM_GEMM_CLAMP + const aie::vector lo = + aie::broadcast(static_cast(FLM_GEMM_CLAMP_MIN)); + const aie::vector hi = + aie::broadcast(static_cast(FLM_GEMM_CLAMP_MAX)); +#endif + + AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) + for (int j = 0; j < CHUNK / V; j++) { + aie::accum acc; + acc.from_vector(aie::load_v(src + j * V)); + // The assignment is the conversion: to_v16bfloat16 yields a raw + // v16bfloat16, not an aie::vector. + aie::vector v = to_v16bfloat16(acc); +#if FLM_GEMM_EPILOGUE_MODE == 1 + v = getGeluBf16_nonLUT(v); +#elif FLM_GEMM_EPILOGUE_MODE == 2 + v = getSiluBf16_nonLUT(v); +#elif FLM_GEMM_EPILOGUE_MODE == 3 + v = getSigmoidBf16_nonLUT(v); +#endif +#if FLM_GEMM_CLAMP + v = aie::clamp(v, lo, hi); +#endif + aie::store_v(y_out + j * V, v); + } +} +} diff --git a/aie_kernels/aie2p/flm_gemm_geometry.h b/aie_kernels/aie2p/flm_gemm_geometry.h new file mode 100644 index 000000000..fcb4e69d9 --- /dev/null +++ b/aie_kernels/aie2p/flm_gemm_geometry.h @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// K-blocking geometry for the flm_gemm design: how much of K one compute tile +// holds at a time, as a function of the n tile width. flm_gemm.cc reads this to +// size its k loop; flm_gemm_mmul.h is the mmul that consumes the result. +// +// The values are a fixed L1 budget split two ways: a wider n tile leaves less +// room for the k slice of B, so the product stays roughly constant. n=128 (the +// only width flm_gemm currently builds) caps the core's k slice at 32. +#ifndef __FLM_GEMM_GEOMETRY_H__ +#define __FLM_GEMM_GEOMETRY_H__ +#include +#include + +constexpr int CT_k_max_n_16 = 16; +constexpr int CT_k_max_n_32 = 32; +constexpr int CT_k_max_n_64 = 64; +constexpr int CT_k_max_n_128 = 32; +constexpr int CT_k_max_n_256 = 16; + +template +constexpr int compute_CT_k_max_n() { + if constexpr (N == 16) { + return CT_k_max_n_16; + } else if constexpr (N == 32) { + return CT_k_max_n_32; + } else if constexpr (N == 64) { + return CT_k_max_n_64; + } else if constexpr (N == 128) { + return CT_k_max_n_128; + } else if constexpr (N == 256) { + return CT_k_max_n_256; + } else { + return -1; + } +} + +#endif // __FLM_GEMM_GEOMETRY_H__ diff --git a/aie_kernels/aie2p/flm_gemm_mmul.h b/aie_kernels/aie2p/flm_gemm_mmul.h new file mode 100644 index 000000000..95488d680 --- /dev/null +++ b/aie_kernels/aie2p/flm_gemm_mmul.h @@ -0,0 +1,182 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// The 2x2 mmul at the heart of flm_gemm, in the "one-buffer A" form its A +// ObjectFifo leg requires. +// +// A arrives as a SINGLE ObjectFifo object spanning every z slice, rather than a +// ping/pong pair the kernel locks for itself. The fifo owns that handshake, so +// there are no acquire/release pairs in here at all, and the core body acquires +// exactly one A object per call to this function. +// +// Each iteration of the j loop keeps four MMUL accumulators live (C00/C01/C10/ +// C11) across a 2x2 block of output tiles, so each pair of A loads and each +// pair of B loads feeds four macs. That is what keeps the vector unit busy; +// dropping to a 1x1 tile would halve the arithmetic per load. +#ifndef __FLM_GEMM_MMUL_H__ +#define __FLM_GEMM_MMUL_H__ +#include "../aie_kernel_utils.h" +#include "flm_gemm_geometry.h" +#include + +// rowA/colA/colB count r x s (A), s x t (B) and r x t (C) sub-tiles, not +// elements. b_row_maj / is_b_s_t_in_row_major select B's in-L1 layout; flm_gemm +// always instantiates (B column-major with row-major s x t +// sub-blocks), which is the layout its memtile forward() produces. The other +// combinations are kept because they are `if constexpr` and cost nothing, but +// they are untested here. +template +__aie_inline void +flm_gemm_mmul_2x2(const T_in *__restrict pA, const T_in *__restrict pB, + T_out *__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) { + T_out *__restrict pC1 = pC + (z * colB) * MMUL::size_C; + T_out *__restrict pC2 = pC + ((z + 1) * colB) * MMUL::size_C; + + // A is one object spanning every z slice; index it directly. + const T_in *__restrict pA_cur_buf = 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 T_in *__restrict pA1 = pA_cur_buf; + const T_in *__restrict pA2 = pA_cur_buf + colA * MMUL::size_A; + const T_in *__restrict pB1; + const T_in *__restrict pB2; + if constexpr (b_row_maj) { + pB1 = pB + (j)*MMUL::size_B; + pB2 = pB + (j + 1) * MMUL::size_B; + } else { + pB1 = pB + (j * colA) * MMUL::size_B; + 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)); + + static_assert(colA % 2 == 0); + // Peano schedules the 2x-unrolled body better than it schedules the + // rolled one; chess does not need the hand-unroll. +#if defined(__chess__) + 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; + + if constexpr (b_row_maj) { + B0 = aie::load_v(pB1); + pB1 += MMUL::size_B * colB; + B1 = aie::load_v(pB2); + pB2 += MMUL::size_B * colB; + } else { + if constexpr (is_b_s_t_in_row_major == false) { + B0 = aie::transpose(aie::load_v(pB1), t, s); + pB1 += MMUL::size_B; + B1 = aie::transpose(aie::load_v(pB2), t, s); + pB2 += MMUL::size_B; + } else { + 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); + } +#else + AIE_LOOP_MAX_ITERATION_COUNT(colA / 2) + for (unsigned i = 0; i < colA; i += 2) { + // First iteration + A0 = aie::load_v(pA1); + pA1 += MMUL::size_A; + A1 = aie::load_v(pA2); + pA2 += MMUL::size_A; + + if constexpr (b_row_maj) { + B0 = aie::load_v(pB1); + pB1 += MMUL::size_B * colB; + B1 = aie::load_v(pB2); + pB2 += MMUL::size_B * colB; + } else { + if constexpr (is_b_s_t_in_row_major == false) { + B0 = aie::transpose(aie::load_v(pB1), t, s); + pB1 += MMUL::size_B; + B1 = aie::transpose(aie::load_v(pB2), t, s); + pB2 += MMUL::size_B; + } else { + 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); + + // Second iteration + A0 = aie::load_v(pA1); + pA1 += MMUL::size_A; + A1 = aie::load_v(pA2); + pA2 += MMUL::size_A; + + if constexpr (b_row_maj) { + B0 = aie::load_v(pB1); + pB1 += MMUL::size_B * colB; + B1 = aie::load_v(pB2); + pB2 += MMUL::size_B * colB; + } else { + if constexpr (is_b_s_t_in_row_major == false) { + B0 = aie::transpose(aie::load_v(pB1), t, s); + pB1 += MMUL::size_B; + B1 = aie::transpose(aie::load_v(pB2), t, s); + pB2 += MMUL::size_B; + } else { + 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); + } +#endif + aie::store_v(pC1, C00.template to_vector()); + pC1 += MMUL::size_C; + aie::store_v(pC1, C01.template to_vector()); + pC1 += MMUL::size_C; + aie::store_v(pC2, C10.template to_vector()); + pC2 += MMUL::size_C; + aie::store_v(pC2, C11.template to_vector()); + pC2 += MMUL::size_C; + } + } + + event1(); +} + +#endif // __FLM_GEMM_MMUL_H__ diff --git a/aie_kernels/aie2p/nonlut_based_ops.h b/aie_kernels/aie2p/nonlut_based_ops.h new file mode 100644 index 000000000..5ec591513 --- /dev/null +++ b/aie_kernels/aie2p/nonlut_based_ops.h @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Branch-free, LUT-free activations over an aie::vector. +// +// These are templated on the vector width and return a vector, so they compose +// inside an existing vector loop -- unlike aie_kernels/aie2p/{gelu,silu}.cc, +// which are whole-buffer entry points over a fixed 32-wide vector. flm_gemm's +// epilogue needs the former: it applies the activation to the same 16-wide +// vector it just converted from the f32 accumulator, without a second pass over +// L1. +// +// All three are built on tanh, which AIE2P has as a native vector op, 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. +#ifndef __NONLUT_BASED_OPS_H__ +#define __NONLUT_BASED_OPS_H__ +#include + +// sigmoid(x) = (tanh(x/2) + 1) / 2 +template +aie::vector +getSigmoidBf16_nonLUT(aie::vector x) { + const bfloat16 half = 0.5f; + const bfloat16 one = 1.0f; + aie::vector v_half = + aie::broadcast(half); + aie::vector v_one = + aie::broadcast(one); + aie::accum x_mul_half = aie::mul(x, v_half); + aie::vector tanh_x_half = + aie::tanh(x_mul_half.template to_vector()); + + aie::vector tanh_x_half_plus_one = + aie::add(tanh_x_half, v_one); + return aie::mul(tanh_x_half_plus_one, v_half); +} + +// silu(x) = x * sigmoid(x) +template +aie::vector +getSiluBf16_nonLUT(aie::vector x) { + aie::vector sigmoid_x = getSigmoidBf16_nonLUT(x); + return aie::mul(x, sigmoid_x); +} + +// gelu(x) ~= x * sigmoid(1.702x) +template +aie::vector +getGeluBf16_nonLUT(aie::vector x) { + constexpr bfloat16 x_scale = 1.702; + aie::vector v_x_scale = + aie::broadcast(x_scale); + aie::vector x_scaled = aie::mul(x, v_x_scale); + aie::vector sigmoid_x_scaled = + getSigmoidBf16_nonLUT(x_scaled); + return aie::mul(x, sigmoid_x_scaled); +} + +#endif // __NONLUT_BASED_OPS_H__ diff --git a/iron/operators/__init__.py b/iron/operators/__init__.py index e7a05a7d6..7555d1cb9 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -13,6 +13,7 @@ _OPERATOR_MODULES = { "ElementwiseAdd": "elementwise_add", "ElementwiseMul": "elementwise_mul", + "FLMGEMM": "flm_gemm", "GEMM": "gemm", "GEMV": "gemv", "MHA": "mha", diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py new file mode 100644 index 000000000..8d22a3b8e --- /dev/null +++ b/iron/operators/flm_gemm/design.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""bf16 GEMM over a fixed 4x8 compute-tile grid. + +This is a different design from ``iron.operators.gemm``, not a retuning of it. +The distinguishing choices, all of which the kernel's L1 layout depends on: + + * **A is broadcast along each compute row.** Four shim tiles (columns 0/2/4/6) + each feed one row of the grid, and every one of the 8 tiles in that row + consumes the same A object. B is broadcast down each column. So an A tile is + fetched once per row rather than once per tile. + * **C is joined at the memtile.** Each of the 4 tiles in a column writes its + own 64x128 slice into one memtile buffer, which drains to DDR as a single + 256x128 block, rather than each tile draining separately. + * **The mmul keeps A in a single ObjectFifo object** spanning every z slice + (``flm_gemm_mmul.h``), instead of a ping/pong pair the kernel locks itself. + * **The epilogue is fused**: 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. + +Geometry is fixed (m/k/n = 64/512/128, r/s/t = 8/8/8, 4x8 grid). 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. +""" + +import numpy as np +from ml_dtypes import bfloat16 + +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.iron.device import Tile +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, N_TILE = 64, 512, 128 +R, S, T = 8, 8, 8 +ROWS, COLS = 4, 8 + +# Which shim column sources the A broadcast for each compute row. Spreading +# them over alternate columns keeps four independent MM2S paths; the existing +# gemm operator pins A the same way in the 8-column case. +A_SOURCE_COL = [0, 2, 4, 6] + +# Must match compute_CT_k_max_n() in flm_gemm_geometry.h: with n=128, +# a compute tile holds 32 of K at a time. +CT_MAX_K = 32 +K_DIV_CT_K_MAX = K_TILE // CT_MAX_K + +# Buffer lengths, in elements. +CT_A_LEN = 2 * R * CT_MAX_K # one z slice +CT_A_OBJ = CT_A_LEN * (M_TILE // R // 2) # A object: every z slice of one mmul +CT_OUT_LEN = 512 # the core's C slice, streamed out in chunks this size +C_SLICE_LEN = M_TILE * N_TILE # one compute tile's C contribution +O_CHUNKS = C_SLICE_LEN // CT_OUT_LEN # C objects one accumulator drains as +C_DEPTH = 2 # C fifo depth; also the core-body unroll +B_ITERS = K_TILE // CT_MAX_K # B chunks the core consumes per k step +B_DEPTH = 2 # B fifo depth; also the core-body unroll +A_DEPTH = 2 + +STACK_SIZE = 4096 + +EPILOGUE_MODES = {"none": 0, "gelu": 1, "silu": 2, "sigmoid": 3} + +# Minimum problem size, i.e. one pass of the whole grid. +MIN_M = M_TILE * ROWS # 256 +MIN_K = K_TILE # 512 +MIN_N = N_TILE * COLS # 1024 + + +def flm_gemm( + dev, + M, + K, + N, + epilogue="none", + kernel_object="flm_gemm.o", + epilogue_object="flm_gemm_epilogue.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 -- the block-major reordering B needs on + the way in is done by the fill descriptor, not by the caller. + """ + if epilogue not in EPILOGUE_MODES: + raise ValueError( + f"epilogue must be one of {sorted(EPILOGUE_MODES)}, got {epilogue!r}" + ) + # The design has no partial-block path: a compute tile either does a full + # m x n block of work or none at all. The existing gemm operator constrains + # its shapes the same way. + for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, MIN_N)): + 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. + n_col_blocks = N // MIN_N + m_row_blocks = M // MIN_M + k_iters = K // K_TILE + + # 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,), bf16_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_b_ty = np.ndarray[(K_TILE * N_TILE,), bf16_ty] + mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] + # 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,), bf16_ty] + c_l3_ty = np.ndarray[(M * N,), bf16_ty] + + acc_init = Kernel("flm_gemm_acc_init", kernel_object, [ct_acc_ty]) + k_step = Kernel( + "flm_gemm_k_step", kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty] + ) + epilogue_chunk = Kernel( + "flm_gemm_epilogue_chunk", + epilogue_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_recv_dims = [(N_TILE // T, K_TILE * T), (T, S), (K_TILE // S, S * T), (S, 1)] + b_send_dims = [ + (K_DIV_CT_K_MAX, T * CT_MAX_K), + (N_TILE // T, K_TILE * T), + (T * CT_MAX_K, 1), + ] + # 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), + (R * CT_MAX_K, 1), + ] + + # 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(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, + tile=Tile(c, 1), + ) + 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): + src = A_SOURCE_COL[r] + 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( + tile=Tile(src, 1), + obj_type=ct_a_obj_ty, + depth=A_DEPTH, + name=f"A_L2L1_{r}", + dims_to_stream=a_send_dims, + ) + # One cons() handle per column: every tile in the row sees this object. + for c in range(COLS): + a_cons[(r, c)] = of_a.cons() + + # B: shim -> memtile -> broadcast down the compute column. + b_l3l2_fifos = [] + b_cons = {} + for c in range(COLS): + of_b_in = ObjectFifo(mt_b_ty, name=f"B_L3L2_{c}", depth=B_DEPTH) + b_l3l2_fifos.append(of_b_in) + of_b = of_b_in.cons(dims_from_stream=b_recv_dims).forward( + tile=Tile(c, 1), + obj_type=ct_b_ty, + depth=B_DEPTH, + name=f"B_L2L1_{c}", + dims_to_stream=b_send_dims, + ) + for r in range(ROWS): + b_cons[(r, c)] = of_b.cons() + + # --- Compute ---------------------------------------------------------- + def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): + # 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. + for _ in range_(n_col_blocks): + 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): + b = b_h.acquire(1) + a = a_h.acquire(1) + kstep_k(a, b, acc) + 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) + + workers = [] + for r in range(ROWS): + for c in range(COLS): + tile = Tile(c, r + 2) + acc = Buffer(tile=tile, type=ct_acc_ty, name=f"c_acc_{r}_{c}") + workers.append( + Worker( + core_fn, + [ + acc, + c_prod[(r, c)].prod(), + b_cons[(r, c)], + a_cons[(r, c)], + acc_init, + k_step, + epilogue_chunk, + ], + tile=tile, + stack_size=STACK_SIZE, + ) + ) + + # --- Runtime ---------------------------------------------------------- + # + # Every wrap below stays under the shim's 10-bit (1023) size field: the + # largest are K_TILE=512 and ROWS*M_TILE=256. + def a_tap(mega_row, r, kb): + # One M_TILE x K_TILE block of A, row-major -- which is exactly the + # order the memtile's dims_from_stream expects, so no reordering is + # needed here (unlike B). + # + # Issued one k-block at a time, like B. Packing all k_iters blocks into + # a single BD also describes the right bytes, but that one BD then has + # to stall part-way through whenever k_iters exceeds the fifo depth, + # while still holding its shim channel -- which deadlocks against the + # C drain sharing that channel. It survives k_iters <= A_DEPTH and + # hangs above it, so the failure only appears at larger K. + return TensorAccessPattern( + tensor_dims=(M * K,), + offset=(mega_row * ROWS + r) * M_TILE * K + kb * K_TILE, + sizes=[1, 1, M_TILE, K_TILE], + strides=[0, 0, K, 1], + ) + + def b_tap(mega_col, c, kb): + # One K_TILE x N_TILE chunk of B's column stripe, reordered on the fly. + # + # B is a plain row-major (K, N) tensor, but the memtile's + # dims_from_stream expects the elements in t-block-major order -- + # (n//T, k%S, k//S, n%T), outermost first -- not row-major. Rather than + # make the caller pre-pack B (which is what the design this came from + # did on the host), the gather is expressed here, so B stays an ordinary + # dense tensor. Getting this order wrong yields silently wrong results, + # not a build error. + # + # Only one k-block fits in the 4 available dimensions, so the caller + # issues one of these per k iteration; each delivers exactly one + # memtile object. + return TensorAccessPattern( + tensor_dims=(K * N,), + offset=(mega_col * COLS + c) * N_TILE + kb * K_TILE * N, + sizes=[N_TILE // T, S, K_TILE // S, T], + strides=[T, N, S * N, 1], + ) + + def c_tap(mega_col, mega_row, c): + # One joined block: ROWS*M_TILE rows of this column's N_TILE-wide slice. + return TensorAccessPattern( + tensor_dims=(M * N,), + offset=mega_row * ROWS * M_TILE * N + (mega_col * COLS + c) * N_TILE, + sizes=[1, 1, ROWS * M_TILE, N_TILE], + strides=[0, 0, N, 1], + ) + + # A shim tile supports only SHIM_BD_LIMIT simultaneously active buffer + # descriptors, and shim column 0 carries three legs at once: the A fills + # for compute row 0, the B fills for compute column 0, and the C drain for + # compute column 0. So each k iteration in flight costs 2 BDs there, plus + # one for the drain -- and exceeding the limit is a hard compile error, not + # a slowdown. Retire the fills in batches sized to stay under it. + SHIM_BD_LIMIT = 16 + K_BATCH = max(1, (SHIM_BD_LIMIT - 2) // 2) # leave room for the C drain + + def sequence(A, B, C, a_prods, b_prods, c_conses): + for mega_col in range(n_col_blocks): + for mega_row in range(m_row_blocks): + # The drain is issued first and retired last: it is an S2MM + # that simply waits for the cores to produce, so having it + # outstanding across the whole sweep is what lets compute and + # write-back overlap. It must NOT share a task group with the + # fills -- finishing a group that contains it before the fills + # it depends on have been issued would deadlock. + tg_c = TaskGroup() + for c in range(COLS): + c_conses[c].drain( + C, c_tap(mega_col, mega_row, c), group=tg_c, wait=True + ) + + # One A and one B object per k iteration, matching the core's + # k_iters x B_ITERS acquires on each leg. + for batch in range(0, k_iters, K_BATCH): + tg_f = TaskGroup() + for kb in range(batch, min(batch + K_BATCH, k_iters)): + for r in range(ROWS): + a_prods[r].fill(A, a_tap(mega_row, r, kb), group=tg_f) + for c in range(COLS): + b_prods[c].fill(B, b_tap(mega_col, c, kb), group=tg_f) + tg_f.finish() + tg_c.finish() + + rt = Runtime( + sequence, + [ + a_l3_ty, + b_l3_ty, + c_l3_ty, + [ + f.prod(tile=Tile(A_SOURCE_COL[r], 0)) + for r, f in enumerate(a_l3l2_fifos) + ], + [f.prod(tile=Tile(c, 0)) for c, f in enumerate(b_l3l2_fifos)], + [f.cons(tile=Tile(c, 0)) for c, f in enumerate(c_l2l3_fifos)], + ], + ) + + my_program = Program(dev, rt, workers=workers) + maybe_enable_trace(my_program, trace_size, workers) + return my_program.resolve_program() diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py new file mode 100644 index 000000000..ba69118c0 --- /dev/null +++ b/iron/operators/flm_gemm/op.py @@ -0,0 +1,180 @@ +# 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 ClassVar, Dict + +from iron.common import ( + MLIROperator, + AIERuntimeArgSpec, + KernelObjectArtifact, + SourceArtifact, + PythonGeneratedMLIRArtifact, + DesignGenerator, +) +from iron.common.device_utils import get_kernel_dir +import aie.utils as aie_utils + +from iron.operators.flm_gemm.design import ( + COLS, + C_DEPTH, + CT_OUT_LEN, + EPILOGUE_MODES, + K_TILE, + MIN_K, + MIN_M, + MIN_N, + M_TILE, + N_TILE, +) + + +@dataclass +class FLMGEMM(MLIROperator): + """AIE-accelerated bf16 GEMM on a fixed 4x8 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. + """ + + M: int + K: int + N: int + # "none" | "gelu" | "silu" | "sigmoid", fused into the C drain. + epilogue: str = field(default="none", repr=False) + # Optional (min, max) applied after the activation. + clamp: tuple[float, float] | None = field(default=None, repr=False) + context: object = field(default=None, repr=False) + + _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} + + def __post_init__(self): + for name, value, unit in ( + ("M", self.M, MIN_M), + ("K", self.K, MIN_K), + ("N", self.N, MIN_N), + ): + if value % unit != 0: + raise ValueError(f"{name} ({value}) must be a multiple of {unit}") + if self.epilogue not in EPILOGUE_MODES: + raise ValueError( + f"epilogue must be one of {sorted(EPILOGUE_MODES)}, " + f"got {self.epilogue!r}" + ) + 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: + # epilogue/clamp are repr=False so the plain path keeps a stable name, + # but they change the emitted kernel, so the variants must not share an + # artifact name: in a shared build dir a cached plain build would + # otherwise satisfy a fused op and silently skip the activation. + base = super().name + if self.epilogue != "none": + base = f"{base}_epi{self.epilogue}" + if self.clamp is not None: + base = f"{base}_clamp{self._clamp_tag}" + return base + + @property + def _clamp_tag(self) -> str: + lo, hi = self.clamp + return f"{lo:g}_{hi:g}".replace("-", "m").replace(".", "p") + + @property + def _epilogue_object(self) -> str: + obj = f"flm_gemm_epilogue_{self.epilogue}" + if self.clamp is not None: + obj = f"{obj}_clamp{self._clamp_tag}" + return f"{obj}.o" + + @property + def _kernel_object(self) -> str: + return f"flm_gemm_{M_TILE}x{K_TILE}x{N_TILE}.o" + + def get_mlir_artifact(self): + return PythonGeneratedMLIRArtifact( + f"{self.name}.mlir", + DesignGenerator( + self.operator_dir / "design.py", + "flm_gemm", + (), + { + "dev": aie_utils.get_current_device(), + "M": self.M, + "K": self.K, + "N": self.N, + "epilogue": self.epilogue, + "kernel_object": self._kernel_object, + "epilogue_object": self._epilogue_object, + "trace_size": 0, + }, + ), + ) + + def get_kernel_artifacts(self): + # The mmul and the epilogue are both aie2p-only: the mmul relies on the + # bf16 emulation path and the grid needs 8 columns. + kernel_dir = get_kernel_dir() + if kernel_dir != "aie2p": + raise NotImplementedError( + f"flm_gemm is only available on NPU2 (aie2p); got {kernel_dir!r}" + ) + base_dir = self.context.base_dir + aie2p = base_dir / "aie_kernels" / "aie2p" + + epilogue_flags = [ + f"-DFLM_GEMM_OUT_CHUNK={CT_OUT_LEN}", + f"-DFLM_GEMM_C_DEPTH={C_DEPTH}", + f"-DFLM_GEMM_EPILOGUE_MODE={EPILOGUE_MODES[self.epilogue]}", + ] + 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. + epilogue_flags += [ + "-DFLM_GEMM_CLAMP=1", + f"-DFLM_GEMM_CLAMP_MIN={float(lo)!r}f", + f"-DFLM_GEMM_CLAMP_MAX={float(hi)!r}f", + ] + + return [ + KernelObjectArtifact( + self._kernel_object, + dependencies=[SourceArtifact(aie2p / "flm_gemm.cc")], + extra_flags=[ + f"-DFLM_GEMM_TILE_M={M_TILE}", + f"-DFLM_GEMM_TILE_K={K_TILE}", + f"-DFLM_GEMM_TILE_N={N_TILE}", + # The r=8 mmul shape this design uses only exists on the + # bfp16-emulated path; without this the kernel will not + # compile. + "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", + ], + ), + KernelObjectArtifact( + self._epilogue_object, + dependencies=[SourceArtifact(aie2p / "flm_gemm_epilogue.cc")], + extra_flags=epilogue_flags, + ), + ] + + def get_arg_spec(self): + return [ + AIERuntimeArgSpec("in", (self.M, self.K)), # A + 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 000000000..861580089 --- /dev/null +++ b/iron/operators/flm_gemm/reference.py @@ -0,0 +1,61 @@ +# 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 + + +def _activation(x, epilogue): + if epilogue == "none": + return x + if epilogue == "gelu": + # Match the kernel, which uses the sigmoid approximation + # gelu(x) ~= x * sigmoid(1.702x) -- NOT torch's erf-exact gelu, and not + # the tanh approximation the standalone gelu operator uses. + return x * torch.sigmoid(1.702 * x) + if epilogue == "silu": + return x * torch.sigmoid(x) + if epilogue == "sigmoid": + return torch.sigmoid(x) + raise ValueError(f"unknown epilogue {epilogue!r}") + + +def reference(input_a, input_b, epilogue="none", clamp=None): + """CPU reference ``C = clamp(activation(A @ B))``. + + The matmul is accumulated in fp32 to mirror the kernel's f32 accumulator, + then cast back to the input dtype at the end, which is where the kernel + converts too. + """ + out_dtype = input_a.dtype + C = torch.matmul(input_a.float(), input_b.float()) + C = _activation(C, epilogue) + if clamp is not None: + C = torch.clamp(C, clamp[0], clamp[1]) + return C.to(out_dtype) + + +def generate_golden_reference( + M: int, + K: int, + N: int, + dtype="bf16", + seed=42, + 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 000000000..92c03ff9c --- /dev/null +++ b/iron/operators/flm_gemm/test.py @@ -0,0 +1,104 @@ +#!/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 iron.operators.flm_gemm.op import FLMGEMM +from iron.operators.flm_gemm.reference import generate_golden_reference +from iron.common.test_utils import run_test + + +def get_params(): + dev = aie_utils.get_current_device() + # The design is a fixed 4x8 grid, so it needs all 8 columns. + if dev.cols < 8 or dev.resolve().name != "npu2": + return [] + + # fmt: off + # M, K, N, epilogue, clamp + regular_params = [ + ( 256, 512, 1024, "none", None), # smallest legal shape: one grid sweep + ( 512, 1024, 2048, "none", None), + ( 256, 512, 1024, "silu", None), + ( 256, 512, 1024, "gelu", None), + ( 256, 512, 1024, "none", (-2.0, 2.0)), + ] + extensive_params = [ + ( 1024, 2048, 2048, "none", None), + ( 2048, 2048, 2048, "none", None), + ( 256, 512, 1024, "sigmoid", None), + ( 512, 1024, 2048, "silu", (-4.0, 4.0)), + ] + # 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 + + +@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", get_params()) +def test_flm_gemm(M, K, N, epilogue, clamp, aie_context): + # Keep the activation tests in the range where the curve is not flat. + scale = 4.0 if epilogue == "none" else 0.5 + golden_ref = generate_golden_reference( + M=M, K=K, N=N, epilogue=epilogue, clamp=clamp, scale=scale + ) + + operator = FLMGEMM( + M=M, + K=K, + N=N, + epilogue=epilogue, + clamp=clamp, + context=aie_context, + ) + + input_buffers = { + "A": golden_ref["input"].flatten(), + "B": golden_ref["input_b"].flatten(), + } + output_buffers = {"C": golden_ref["output"].flatten()} + + # This design's r=8 mmul shape exists only on the bfp16-emulated path, so + # its error budget is that of an emulated GEMM, not of the exact one the + # GEMM operator's test asserts on (that test opts into r=4 via + # emulate_bf16_mmul_with_bfp16=False, which is not available here). + # + # A pure relative tolerance cannot work: 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. So the error is bounded in ABSOLUTE terms + # against the accumulated mass, which is what bfp16 error actually scales + # with. Measured on this data: mean |err| is 0.00042 of the mass and the + # worst element 0.0025 -- both marginally better than the GEMM operator run + # in the same emulated mode (0.00044 / 0.0031), so the budget below is not + # papering over a regression in this port. The bound is tight enough to + # have caught a real bug: leaving the core in its default floor rounding + # mode pushes mean error to 0.0099 of mass, ~7x over. + mass = K * golden_ref["input"].abs().float().mean() * ( + golden_ref["input_b"].abs().float().mean() + ) + errors, latency_us, bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=0.04, + abs_tol=float(0.004 * mass), + ) + + 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" From bbcc1462cc0cd1c5466cd0c95d7735c6875c728d Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 14:43:17 -0600 Subject: [PATCH 02/31] flm_gemm: inline the epilogue, and support .ll kernel artifacts The epilogue is short and runs once per C object, so leaving it as a call costs real time against the per-object overhead the C ObjectFifo introduces; the design this was ported from merges it into the core for exactly that reason. Inlining the much larger mmul measures worse there, so only the epilogue is merged. aie.iron's ExternalFunction(inline=True) emits the right declaration (link_with_mode = "merge") but its own compilation is driven by CompilableDesign, which IRON's flow does not use -- so the .ll was never built and aiecc failed to find it. Make .ll/.bc kernel artifacts first-class instead: KernelCompiler infers the inline build from the suffix and passes the symbol to mark alwaysinline, which KernelObjectArtifact now carries. Both are rejected with a clear message when misused (chess has no equivalent path, and an inline build with no symbol is a silent no-op otherwise). Verified inlined rather than silently falling back: the .ll carries alwaysinline, the MLIR declares merge mode, and the core ELF exports flm_gemm_acc_init and flm_gemm_k_step but no epilogue symbol. inline_epilogue is a field so the two can be A/B'd, and it participates in the operator name -- otherwise a cached xclbin from one satisfies the other and the comparison measures a binary against itself. Co-Authored-By: Claude --- iron/common/compilation/base.py | 23 +++++++++ iron/operators/flm_gemm/design.py | 46 ++++++++++++++--- iron/operators/flm_gemm/op.py | 83 ++++++++++++++++++++++--------- 3 files changed, 122 insertions(+), 30 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 3bc39c220..6bd48a0db 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -382,11 +382,15 @@ def __init__( extra_flags: list[str] | None = None, rename_symbols: dict[str, str] | None = None, prefix_symbols: str | None = None, + inline_symbol: str | None = None, ) -> None: super().__init__(filename, dependencies) self.extra_flags = extra_flags if extra_flags is not None else [] self.rename_symbols = rename_symbols if rename_symbols is not None else {} self.prefix_symbols = prefix_symbols + # Required when filename is a .ll/.bc: the symbol to mark alwaysinline + # so aiecc inlines it after llvm-linking the module into the core. + self.inline_symbol = inline_symbol class KernelArchiveArtifact(CompilationArtifact): @@ -790,6 +794,23 @@ def compile(self, artifacts): "-Wno-missing-template-arg-list-after-template-kw" ] + compile_args + # A .ll/.bc output means the kernel is meant to be llvm-linked into + # the core and inlined (aie.iron's ExternalFunction(inline=True) + # declares it with link_with_mode = "merge") rather than left as a + # call. That is a Peano-only path; xchesscc has no equivalent. + inline = str(artifact.filename).endswith((".ll", ".bc")) + if inline and self.use_chess: + raise RuntimeError( + f"Kernel artifact '{artifact.filename}' requests an inline " + "(.ll/.bc) build, which requires the Peano compiler; this " + "compilation is configured for chess." + ) + if inline and not artifact.inline_symbol: + raise RuntimeError( + f"Kernel artifact '{artifact.filename}' is an inline " + "(.ll/.bc) build and must name the symbol to inline via " + "KernelObjectArtifact(inline_symbol=...)." + ) commands.append( PythonCallbackCompilationCommand( partial( @@ -800,6 +821,8 @@ def compile(self, artifacts): include_dirs=[str(runtime_lib_include_path)], compile_args=compile_args, use_chess=self.use_chess, + inline=inline, + symbol_name=artifact.inline_symbol if inline else None, ) ) ) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 8d22a3b8e..b619439dd 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -24,11 +24,22 @@ -D flags, so the C++ and the dataflow cannot drift apart. """ +from pathlib import Path + import numpy as np from ml_dtypes import bfloat16 from aie.helpers.taplib import TensorAccessPattern -from aie.iron import Buffer, Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker +from aie.iron import ( + Buffer, + ExternalFunction, + Kernel, + ObjectFifo, + Program, + Runtime, + TaskGroup, + Worker, +) from aie.iron.controlflow import range_ from aie.iron.device import Tile from iron.operators._trace import maybe_enable_trace @@ -63,6 +74,9 @@ STACK_SIZE = 4096 EPILOGUE_MODES = {"none": 0, "gelu": 1, "silu": 2, "sigmoid": 3} +# 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 = "flm_gemm_epilogue_chunk" # Minimum problem size, i.e. one pass of the whole grid. MIN_M = M_TILE * ROWS # 256 @@ -78,6 +92,9 @@ def flm_gemm( epilogue="none", kernel_object="flm_gemm.o", epilogue_object="flm_gemm_epilogue.o", + epilogue_source=None, + epilogue_flags=None, + inline_epilogue=False, trace_size=0, ): """Emit the MLIR module for an M x K @ K x N bf16 GEMM. @@ -123,11 +140,28 @@ def flm_gemm( k_step = Kernel( "flm_gemm_k_step", kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty] ) - epilogue_chunk = Kernel( - "flm_gemm_epilogue_chunk", - epilogue_object, - [ct_out_ty, ct_acc_ty, np.int32, np.int32], - ) + epilogue_arg_types = [ct_out_ty, ct_acc_ty, np.int32, np.int32] + if inline_epilogue: + # Compile the epilogue to alwaysinline LLVM IR so aiecc llvm-links it + # into the core instead of leaving a call. The epilogue is short and + # runs once per C object, so the call overhead the C ObjectFifo + # introduces is a real cost here -- whereas inlining the much larger + # mmul measures worse, which is why only this one is merged. + if epilogue_source is None: + raise ValueError("inline_epilogue requires epilogue_source") + epilogue_chunk = ExternalFunction( + EPILOGUE_SYMBOL, + object_file_name=str(Path(epilogue_object).with_suffix(".ll")), + source_file=str(epilogue_source), + inline=True, + arg_types=epilogue_arg_types, + include_dirs=[str(Path(epilogue_source).parent)], + compile_flags=list(epilogue_flags or []), + ) + else: + epilogue_chunk = Kernel( + EPILOGUE_SYMBOL, epilogue_object, epilogue_arg_types + ) # --- Data movement ---------------------------------------------------- # diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index ba69118c0..695b91f4c 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -20,6 +20,7 @@ C_DEPTH, CT_OUT_LEN, EPILOGUE_MODES, + EPILOGUE_SYMBOL, K_TILE, MIN_K, MIN_M, @@ -46,6 +47,9 @@ class FLMGEMM(MLIROperator): epilogue: str = field(default="none", repr=False) # Optional (min, max) applied after the activation. clamp: tuple[float, float] | None = field(default=None, repr=False) + # Compile the epilogue to alwaysinline LLVM IR merged into the core, + # rather than leaving it as a call. See design.py. + inline_epilogue: bool = field(default=True, repr=False) context: object = field(default=None, repr=False) _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} @@ -81,6 +85,13 @@ def name(self) -> str: base = f"{base}_epi{self.epilogue}" if self.clamp is not None: base = f"{base}_clamp{self._clamp_tag}" + if not self.inline_epilogue: + # Different core binary, so it must not share an artifact name with + # the inlined build -- otherwise a cached xclbin from one satisfies + # the other and an A/B of the two silently compares a binary with + # itself. (The GEMM operator has this hazard for its + # emulate_bf16_mmul_with_bfp16 / prio_accuracy flags.) + base = f"{base}_noinline" return base @property @@ -89,11 +100,36 @@ def _clamp_tag(self) -> str: return f"{lo:g}_{hi:g}".replace("-", "m").replace(".", "p") @property - def _epilogue_object(self) -> str: + def _epilogue_artifact(self) -> str: obj = f"flm_gemm_epilogue_{self.epilogue}" if self.clamp is not None: obj = f"{obj}_clamp{self._clamp_tag}" - return f"{obj}.o" + # .ll is llvm-linked into the core and inlined; .o is left as a call. + return f"{obj}.ll" if self.inline_epilogue else f"{obj}.o" + + @property + def _epilogue_source(self): + return self.context.base_dir / "aie_kernels" / "aie2p" / "flm_gemm_epilogue.cc" + + @property + def _epilogue_flags(self) -> list[str]: + """Compile flags for the epilogue, shared by the inline and + separately-compiled paths so the two cannot drift apart.""" + flags = [ + f"-DFLM_GEMM_OUT_CHUNK={CT_OUT_LEN}", + f"-DFLM_GEMM_C_DEPTH={C_DEPTH}", + f"-DFLM_GEMM_EPILOGUE_MODE={EPILOGUE_MODES[self.epilogue]}", + ] + 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 += [ + "-DFLM_GEMM_CLAMP=1", + f"-DFLM_GEMM_CLAMP_MIN={float(lo)!r}f", + f"-DFLM_GEMM_CLAMP_MAX={float(hi)!r}f", + ] + return flags @property def _kernel_object(self) -> str: @@ -113,7 +149,10 @@ def get_mlir_artifact(self): "N": self.N, "epilogue": self.epilogue, "kernel_object": self._kernel_object, - "epilogue_object": self._epilogue_object, + "epilogue_object": self._epilogue_artifact, + "epilogue_source": str(self._epilogue_source), + "epilogue_flags": self._epilogue_flags, + "inline_epilogue": self.inline_epilogue, "trace_size": 0, }, ), @@ -130,22 +169,7 @@ def get_kernel_artifacts(self): base_dir = self.context.base_dir aie2p = base_dir / "aie_kernels" / "aie2p" - epilogue_flags = [ - f"-DFLM_GEMM_OUT_CHUNK={CT_OUT_LEN}", - f"-DFLM_GEMM_C_DEPTH={C_DEPTH}", - f"-DFLM_GEMM_EPILOGUE_MODE={EPILOGUE_MODES[self.epilogue]}", - ] - 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. - epilogue_flags += [ - "-DFLM_GEMM_CLAMP=1", - f"-DFLM_GEMM_CLAMP_MIN={float(lo)!r}f", - f"-DFLM_GEMM_CLAMP_MAX={float(hi)!r}f", - ] - - return [ + artifacts = [ KernelObjectArtifact( self._kernel_object, dependencies=[SourceArtifact(aie2p / "flm_gemm.cc")], @@ -159,12 +183,23 @@ def get_kernel_artifacts(self): "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", ], ), - KernelObjectArtifact( - self._epilogue_object, - dependencies=[SourceArtifact(aie2p / "flm_gemm_epilogue.cc")], - extra_flags=epilogue_flags, - ), ] + # Inlined, this is a .ll that aiecc llvm-links into the core; otherwise + # an ordinary .o it calls. Either way the build system produces it -- + # design.py's ExternalFunction only supplies the merge-mode + # declaration in the MLIR, it does not compile anything under IRON's + # compilation flow. + artifacts.append( + KernelObjectArtifact( + self._epilogue_artifact, + dependencies=[SourceArtifact(self._epilogue_source)], + extra_flags=self._epilogue_flags, + inline_symbol=( + EPILOGUE_SYMBOL if self.inline_epilogue else None + ), + ) + ) + return artifacts def get_arg_spec(self): return [ From 80679f20b391252731ccaef7c23b49afe59b7f1e Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 15:31:00 -0600 Subject: [PATCH 03/31] flm_gemm: handle a partial trailing column-block N previously had to be a multiple of N_TILE*COLS=1024, which excluded the shapes this design exists to serve: a transformer's o and down projections have N = model dim, so 8 of Gemma4's 14 projections could not run at all. N now only has to tile to N_TILE. The trailing group of fewer than COLS column-blocks is handled by giving each column its own trip counts rather than a runtime branch: columns below the remainder width compute one block more, and the rest run an A-only drain loop for it. That drain is not optional -- A is broadcast along the whole compute row, so a column sitting the block out must still consume its share or the columns that do have work stall behind it. The runtime sequence matches, issuing A for every row but B and C only for the active columns. Also enforces a K ceiling instead of silently mis-lowering. A sweep's fills have to go in ONE task group; splitting them across groups to fit more k-iterations into a shim tile's 16 buffer descriptors returns wrong data, reproducible at M=1024 K=2560 N=2560 where one group passes and two fail. Undiagnosed, so K > 3584 now raises rather than producing bad results. Tests cover N=1536 (4 of 8 columns active), N=128 (1 column, no full sweep), and the real E4B o-projection and E2B down-projection shapes. 13/13 flm_gemm and 24/24 gemm pass. Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 184 +++++++++++++++++++++++------- iron/operators/flm_gemm/op.py | 6 +- iron/operators/flm_gemm/test.py | 11 +- 3 files changed, 157 insertions(+), 44 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index b619439dd..28b399ab9 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -107,10 +107,13 @@ def flm_gemm( raise ValueError( f"epilogue must be one of {sorted(EPILOGUE_MODES)}, got {epilogue!r}" ) - # The design has no partial-block path: a compute tile either does a full - # m x n block of work or none at all. The existing gemm operator constrains - # its shapes the same way. - for name, value, unit in (("M", M, MIN_M), ("K", K, MIN_K), ("N", N, MIN_N)): + # 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}") @@ -118,9 +121,16 @@ def flm_gemm( f32 = np.dtype[np.float32] # How many times the whole grid sweeps, in each dimension. - n_col_blocks = N // MIN_N m_row_blocks = M // MIN_M k_iters = K // K_TILE + # 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 + # Per column: how many column-blocks it computes, and whether it has to sit + # out a trailing one while still draining the A broadcast for its row. + col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(COLS)] + col_drain = [1 if (rem_blocks and c >= rem_blocks) else 0 for c in range(COLS)] # L1 (per compute tile) ct_a_obj_ty = np.ndarray[(CT_A_OBJ,), bf16_ty] @@ -244,30 +254,52 @@ def flm_gemm( b_cons[(r, c)] = of_b.cons() # --- Compute ---------------------------------------------------------- - def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): - # 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. - for _ in range_(n_col_blocks): - 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): - b = b_h.acquire(1) - a = a_h.acquire(1) - kstep_k(a, b, acc) - 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) + def make_core_fn(n_work, n_drain): + """Core body for a column that computes ``n_work`` column-blocks and + then drains A for ``n_drain`` more (0 or 1).""" + + def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): + # 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): + b = b_h.acquire(1) + a = a_h.acquire(1) + kstep_k(a, b, acc) + 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): + a_h.acquire(1) + a_h.release(1) + + return core_fn workers = [] for r in range(ROWS): @@ -276,7 +308,7 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): acc = Buffer(tile=tile, type=ct_acc_ty, name=f"c_acc_{r}_{c}") workers.append( Worker( - core_fn, + make_core_fn(col_work[c], col_drain[c]), [ acc, c_prod[(r, c)].prod(), @@ -350,34 +382,104 @@ def c_tap(mega_col, mega_row, c): # one for the drain -- and exceeding the limit is a hard compile error, not # a slowdown. Retire the fills in batches sized to stay under it. SHIM_BD_LIMIT = 16 - K_BATCH = max(1, (SHIM_BD_LIMIT - 2) // 2) # leave room for the C drain + # Cost on the worst shim tile (an A source column, which carries an A fill, + # a B fill and a C drain): 2 BDs per k-block in a fill batch, 1 per drain. + # A sweep's fills must all go in ONE task group. Splitting them across + # groups -- which is what would make larger k_iters fit -- produces + # silently WRONG results, reproducible at M=1024 K=2560 N=2560 + # (k_iters=5): one group passes, two groups fail on the same shape. Not + # yet diagnosed, so the split is not used and the limit is enforced + # instead of being silently mis-lowered. + K_BATCH = (SHIM_BD_LIMIT - 1) // 2 # 7 + if k_iters > K_BATCH: + raise ValueError( + f"K ({K}) needs {k_iters} k-iterations, but at most {K_BATCH} fit " + f"in a shim tile's {SHIM_BD_LIMIT} buffer descriptors alongside " + f"the C drain. Split the GEMM along K, or fix the multi-group " + f"fill path." + ) def sequence(A, B, C, a_prods, b_prods, c_conses): - for mega_col in range(n_col_blocks): + # Sweeps 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 have to drain their share of the broadcast. + sweeps = [(mc, COLS) for mc in range(n_full)] + if rem_blocks: + sweeps.append((n_full, rem_blocks)) + + # Retiring each sweep before issuing the next serialises the pipeline: + # the next sweep's fills cannot start until this sweep's C drain has + # come back, which waits on the cores. So issue sweep i+1 in full, and + # only THEN retire sweep i -- the depth-2 model, which is what the + # design this was ported from uses (its notes record that a depth-1 + # "retire immediately" variant hung on hardware). + # + # NOT YET DONE -- this is the operator's main performance gap. Driving + # the original overlay directly measures 8750 GFLOP/s against this + # sequence's 1687 on the same shape and kernel, and the difference is + # here, not in the kernel or the tiling. + # + # Two attempts at overlapping both hung on hardware: an opportunistic + # "retire the oldest group only when BDs run short" scheduler, and a + # strict depth-2 "issue sweep i+1, then retire sweep i". + # + # The idiom itself is not the problem. The original design's sequence + # issues the next iteration's tasks and only THEN awaits and frees the + # previous one's -- and its notes are explicit that the retirement must + # be a real await, not a bare free (freeing early returns the BD slot + # while the transfer is still in flight: the first invocation comes back + # correct and later ones progressively corrupt). TaskGroup.finish() + # awaits then frees, so it already expresses this. + # + # So the hang here has a different, undiagnosed cause -- plausibly BD + # accounting, since the cost model above assumes one BD per transfer. + # Worth knowing before spending much on it: the original hit an + # undiagnosed hang on this path too, which survived even after the + # relevant compiler fix landed, and it still ships with its own IRON + # sequence disabled. + pipelined = False + + prev = None + for mega_col, active_cols in sweeps: for mega_row in range(m_row_blocks): # The drain is issued first and retired last: it is an S2MM - # that simply waits for the cores to produce, so having it - # outstanding across the whole sweep is what lets compute and - # write-back overlap. It must NOT share a task group with the - # fills -- finishing a group that contains it before the fills - # it depends on have been issued would deadlock. + # that waits for the cores to produce, so keeping it + # outstanding across the sweep overlaps compute with + # write-back. It must not share a task group with the fills it + # depends on -- finishing those together would deadlock. tg_c = TaskGroup() - for c in range(COLS): + for c in range(active_cols): c_conses[c].drain( C, c_tap(mega_col, mega_row, c), group=tg_c, wait=True ) # One A and one B object per k iteration, matching the core's # k_iters x B_ITERS acquires on each leg. + fill_groups = [] for batch in range(0, k_iters, K_BATCH): tg_f = TaskGroup() for kb in range(batch, min(batch + K_BATCH, k_iters)): for r in range(ROWS): a_prods[r].fill(A, a_tap(mega_row, r, kb), group=tg_f) - for c in range(COLS): + for c in range(active_cols): b_prods[c].fill(B, b_tap(mega_col, c, kb), group=tg_f) - tg_f.finish() - tg_c.finish() + if pipelined: + fill_groups.append(tg_f) + else: + tg_f.finish() + + if pipelined: + if prev is not None: + for tg in prev: + tg.finish() + prev = fill_groups + [tg_c] + else: + tg_c.finish() + + if prev is not None: + for tg in prev: + tg.finish() rt = Runtime( sequence, diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 695b91f4c..b39402ca5 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -24,7 +24,6 @@ K_TILE, MIN_K, MIN_M, - MIN_N, M_TILE, N_TILE, ) @@ -55,10 +54,13 @@ class FLMGEMM(MLIROperator): _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} def __post_init__(self): + # 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, MIN_M), ("K", self.K, MIN_K), - ("N", self.N, MIN_N), + ("N", self.N, N_TILE), ): if value % unit != 0: raise ValueError(f"{name} ({value}) must be a multiple of {unit}") diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index 92c03ff9c..d9f23a303 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -16,11 +16,18 @@ def get_params(): if dev.cols < 8 or dev.resolve().name != "npu2": return [] + # N values that are NOT a multiple of N_TILE*COLS=1024 exercise the + # trailing partial column-block, where some columns compute it and the + # rest only drain the A broadcast. Real transformer o/down projections + # have N = model dim, so they always land here: 1536 leaves 4 active + # columns, 2560 leaves 4, and 128 leaves just 1. # fmt: off # M, K, N, epilogue, clamp regular_params = [ - ( 256, 512, 1024, "none", None), # smallest legal shape: one grid sweep + ( 256, 512, 1024, "none", None), # smallest full sweep ( 512, 1024, 2048, "none", None), + ( 256, 512, 1536, "none", None), # remainder: 4 of 8 columns + ( 256, 512, 128, "none", None), # remainder only: 1 column ( 256, 512, 1024, "silu", None), ( 256, 512, 1024, "gelu", None), ( 256, 512, 1024, "none", (-2.0, 2.0)), @@ -28,6 +35,8 @@ def get_params(): extensive_params = [ ( 1024, 2048, 2048, "none", None), ( 2048, 2048, 2048, "none", None), + ( 1024, 2560, 2560, "none", None), # E4B o-projection shape + ( 512, 1536, 1536, "silu", None), # E2B down-projection shape ( 256, 512, 1024, "sigmoid", None), ( 512, 1024, 2048, "silu", (-4.0, 4.0)), ] From e2c6d1319bde8d8f69f90cf3d1d7ed636175f8cb Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 16:16:19 -0600 Subject: [PATCH 04/31] flm_gemm: do not instantiate columns that have no work When N is smaller than the grid's COLS*N_TILE stride there is no full sweep, so only the first rem_blocks columns participate. The remaining columns were still getting B and C objectfifos, plus a consumer on the A broadcast, none of which anything ever drains. The pinned mlir-aie accepts that silently; a newer one rejects it outright with "objectfifo.pool op segment 0 has no drainer", which is how it surfaced -- N=128 failed to compile there while passing on the pin. It was latent dead dataflow either way. Those columns are now not built at all. Dropping them from the A broadcast also removes their A-drain obligation, so the drain loop is needed only for a column that exists and sits out the trailing block, which requires at least one full sweep. 13/13 on both the pinned wheel and mlir_aie 1.4.3.dev60 (verified with a clean build dir -- a shared one silently reuses the other toolchain's xclbins and the second run proves nothing). Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 28b399ab9..0027dc146 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -127,10 +127,21 @@ def flm_gemm( # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. n_full = N // MIN_N rem_blocks = (N % MIN_N) // N_TILE - # Per column: how many column-blocks it computes, and whether it has to sit - # out a trailing one while still draining the A broadcast for its row. - col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(COLS)] - col_drain = [1 if (rem_blocks and c >= rem_blocks) else 0 for c in range(COLS)] + # 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) + ] # L1 (per compute tile) ct_a_obj_ty = np.ndarray[(CT_A_OBJ,), bf16_ty] @@ -204,7 +215,7 @@ def flm_gemm( # DDR as one contiguous (ROWS*M_TILE) x N_TILE block. c_l2l3_fifos = [] c_prod = {} - for c in range(COLS): + 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( @@ -233,14 +244,15 @@ def flm_gemm( name=f"A_L2L1_{r}", dims_to_stream=a_send_dims, ) - # One cons() handle per column: every tile in the row sees this object. - for c in range(COLS): + # 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. b_l3l2_fifos = [] b_cons = {} - for c in range(COLS): + for c in range(n_active_cols): of_b_in = ObjectFifo(mt_b_ty, name=f"B_L3L2_{c}", depth=B_DEPTH) b_l3l2_fifos.append(of_b_in) of_b = of_b_in.cons(dims_from_stream=b_recv_dims).forward( @@ -303,7 +315,7 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): workers = [] for r in range(ROWS): - for c in range(COLS): + for c in range(n_active_cols): tile = Tile(c, r + 2) acc = Buffer(tile=tile, type=ct_acc_ty, name=f"c_acc_{r}_{c}") workers.append( From f7b4355a4a420c263f83330f59a98efb5a739fea Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 16:31:31 -0600 Subject: [PATCH 05/31] flm_gemm: consume B pre-packed, closing a 3.3x gap B was kept as a plain row-major (K, N) tensor and reordered into the memtile's expected layout by the fill descriptor. That is correct, and it was a deliberate choice to spare the caller a packing step, but its innermost run is T=8 bf16 = 16 bytes: every 128 KB B transfer became 8192 scattered bursts. B is ~70% of the bytes a dispatch moves, so the whole operator ran at ~10 GB/s against the original overlay's ~47. B is now consumed pre-packed via FLMGEMM.pack_B and each fill is one contiguous read, which is exactly why the design this came from packs its weights on the host. Weights are packed once and reused across dispatches, so the cost belongs on the caller. M=1024 K=1536 N=6144, same session: FLM peano mm.xclbin 2270 us shipped v1.0.4 2335 us flm_gemm before 11509 us -> after 3446 us IRON GEMM (r8/f32) 3465 us So 5.2x off the original becomes 1.4x, and level with the existing GEMM operator while keeping 41x better accuracy. Found by nulling the compute out: with no arithmetic at all the operator still took 11 ms, which exonerated the kernel, the mmul geometry, the epilogue and the L1 budget in one measurement, and pointed at the transfer descriptors. An earlier test of this same hypothesis had reported only 5% because it reused a cached xclbin from a build directory it had not cleared. The residual 1.4x is the missing depth-2 overlap: data movement alone is 2.07 ms and compute adds ~1.37 ms, so the two are running almost entirely serialised. Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 28 ++++++++++++++-------------- iron/operators/flm_gemm/op.py | 30 ++++++++++++++++++++++++++++++ iron/operators/flm_gemm/test.py | 3 ++- 3 files changed, 46 insertions(+), 15 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 0027dc146..e76f7f000 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -358,24 +358,24 @@ def a_tap(mega_row, r, kb): ) def b_tap(mega_col, c, kb): - # One K_TILE x N_TILE chunk of B's column stripe, reordered on the fly. + # One K_TILE x N_TILE chunk of B, read as a single contiguous run. # - # B is a plain row-major (K, N) tensor, but the memtile's - # dims_from_stream expects the elements in t-block-major order -- - # (n//T, k%S, k//S, n%T), outermost first -- not row-major. Rather than - # make the caller pre-pack B (which is what the design this came from - # did on the host), the gather is expressed here, so B stays an ordinary - # dense tensor. Getting this order wrong yields silently wrong results, - # not a build error. + # B must arrive PRE-PACKED in the memtile's expected order (see + # FLMGEMM.pack_B). Expressing that reorder in the descriptor instead -- + # a 4D gather over a plain row-major (K, N) tensor -- is correct but + # ruinous: its innermost run is T=8 bf16, so a 128 KB transfer becomes + # 8192 scattered 16-byte bursts. Measured with the compute nulled out, + # that costs 5.4x (11122 us vs 2070 us at M=1024 K=1536 N=6144) and was + # the entire gap against the original overlay, which pre-packs its + # weights on the host for exactly this reason. # - # Only one k-block fits in the 4 available dimensions, so the caller - # issues one of these per k iteration; each delivers exactly one - # memtile object. + # Weights are packed once and reused across dispatches, so this belongs + # on the caller rather than in the inner loop. return TensorAccessPattern( tensor_dims=(K * N,), - offset=(mega_col * COLS + c) * N_TILE + kb * K_TILE * N, - sizes=[N_TILE // T, S, K_TILE // S, T], - strides=[T, N, S * N, 1], + offset=(mega_col * COLS + c) * N_TILE * K + kb * K_TILE * N_TILE, + sizes=[1, 1, 1, K_TILE * N_TILE], + strides=[0, 0, 0, 1], ) def c_tap(mega_col, mega_row, c): diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index b39402ca5..8ce691ace 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -26,6 +26,8 @@ MIN_M, M_TILE, N_TILE, + S, + T, ) @@ -203,9 +205,37 @@ def get_kernel_artifacts(self): ) return artifacts + @staticmethod + def pack_B(B): + """Reorder a row-major ``(K, N)`` weight matrix into the layout the B + fill expects. Returns a flat tensor. + + Each ``K_TILE x N_TILE`` tile is emitted in t-block-major order -- the + odometer ``(n//T, k%S, k//S, n%T)``, outermost first -- with tiles + ordered by column stripe and then by k-block, so each fill is one + contiguous read. + + This is deliberately the caller's job. The same reorder is expressible + as a strided descriptor over an unpacked B, but its innermost run is + then T=8 bf16 = 16 bytes, turning each 128 KB transfer into 8192 + scattered bursts -- measured 5.4x slower end to end, and the whole of + this operator's gap against the design it was ported from, which packs + its weights on the host for the same reason. Weights are packed once + and reused across dispatches, so the cost belongs here. + """ + 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" + ) + t = B.reshape(K // K_TILE, K_TILE // S, S, N // N_TILE, N_TILE // T, T) + # (kb, kb8, s_in, cb, tb, t_in) -> (cb, kb, tb, s_in, kb8, t_in) + return t.permute(3, 0, 4, 2, 1, 5).reshape(-1).contiguous() + 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 ] diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index d9f23a303..96fee9453 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -74,7 +74,8 @@ def test_flm_gemm(M, K, N, epilogue, clamp, aie_context): input_buffers = { "A": golden_ref["input"].flatten(), - "B": golden_ref["input_b"].flatten(), + # B is consumed pre-packed; see FLMGEMM.pack_B. + "B": FLMGEMM.pack_B(golden_ref["input_b"]), } output_buffers = {"C": golden_ref["output"].flatten()} From 78eaebe2945965ef872a392812f40a6f1151ecd7 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 16:48:04 -0600 Subject: [PATCH 06/31] flm_gemm: one transfer per column-block, not per object A single fill or drain may span many fifo objects -- the descriptor walks them in the order the cores consume. The sequence was issuing one task per object instead, so every row-block ended in a host-side await on its C drain, and the next row-block's fills could not start until that C had come all the way back from DDR. Those awaits were the serialisation. Each of A, B and C now goes out as one task per column-block, with the dimension order matching the core loop nest (mega_row, then k). For M=1024 K=1536 N=6144 that is 120 tasks and 48 awaits where it was 1056 and 192. before 3446 us after 2534 us FLM mm.xclbin 2241 us, shipped 2212 us IRON GEMM (r8/f32) 3331 us So 1.13x off the original, and 1.3x faster than the existing GEMM operator at this shape, still with ~40x better accuracy. This also retires the K ceiling. The previous per-k-block tasks needed 2*k_iters + 1 buffer descriptors on a shim tile, which capped K at 3584 and forced a batching path that silently corrupted results when it split a sweep's fills across task groups. One task per leg needs three, so that whole mechanism is gone: K=4096 now builds and runs (err/mass 1.5e-4). Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 199 +++++++++--------------------- 1 file changed, 61 insertions(+), 138 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index e76f7f000..4e8a45c3a 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -339,159 +339,82 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): # # Every wrap below stays under the shim's 10-bit (1023) size field: the # largest are K_TILE=512 and ROWS*M_TILE=256. - def a_tap(mega_row, r, kb): - # One M_TILE x K_TILE block of A, row-major -- which is exactly the - # order the memtile's dims_from_stream expects, so no reordering is - # needed here (unlike B). - # - # Issued one k-block at a time, like B. Packing all k_iters blocks into - # a single BD also describes the right bytes, but that one BD then has - # to stall part-way through whenever k_iters exceeds the fifo depth, - # while still holding its shim channel -- which deadlocks against the - # C drain sharing that channel. It survives k_iters <= A_DEPTH and - # hangs above it, so the failure only appears at larger K. + # 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 + # 10-bit size field (largest are K_TILE=512 and ROWS*M_TILE=256). + def a_tap(mega_col, r): + # 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. return TensorAccessPattern( tensor_dims=(M * K,), - offset=(mega_row * ROWS + r) * M_TILE * K + kb * K_TILE, - sizes=[1, 1, M_TILE, K_TILE], - strides=[0, 0, K, 1], + offset=r * M_TILE * K, + sizes=[m_row_blocks, k_iters, M_TILE, K_TILE], + strides=[ROWS * M_TILE * K, K_TILE, K, 1], ) - def b_tap(mega_col, c, kb): - # One K_TILE x N_TILE chunk of B, read as a single contiguous run. + 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 in the memtile's expected order (see - # FLMGEMM.pack_B). Expressing that reorder in the descriptor instead -- - # a 4D gather over a plain row-major (K, N) tensor -- is correct but - # ruinous: its innermost run is T=8 bf16, so a 128 KB transfer becomes - # 8192 scattered 16-byte bursts. Measured with the compute nulled out, - # that costs 5.4x (11122 us vs 2070 us at M=1024 K=1536 N=6144) and was - # the entire gap against the original overlay, which pre-packs its - # weights on the host for exactly this reason. - # - # Weights are packed once and reused across dispatches, so this belongs - # on the caller rather than in the inner loop. + # B must arrive PRE-PACKED (see FLMGEMM.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,), - offset=(mega_col * COLS + c) * N_TILE * K + kb * K_TILE * N_TILE, - sizes=[1, 1, 1, K_TILE * N_TILE], - strides=[0, 0, 0, 1], + offset=(mega_col * COLS + c) * N_TILE * K, + sizes=[m_row_blocks, k_iters, 1, K_TILE * N_TILE], + strides=[0, K_TILE * N_TILE, 0, 1], ) - def c_tap(mega_col, mega_row, c): - # One joined block: ROWS*M_TILE rows of this column's N_TILE-wide slice. + def c_tap(mega_col, c): + # Every joined block this column produces for one column-block: one + # ROWS*M_TILE x N_TILE block per row-block. return TensorAccessPattern( tensor_dims=(M * N,), - offset=mega_row * ROWS * M_TILE * N + (mega_col * COLS + c) * N_TILE, - sizes=[1, 1, ROWS * M_TILE, N_TILE], - strides=[0, 0, N, 1], - ) - - # A shim tile supports only SHIM_BD_LIMIT simultaneously active buffer - # descriptors, and shim column 0 carries three legs at once: the A fills - # for compute row 0, the B fills for compute column 0, and the C drain for - # compute column 0. So each k iteration in flight costs 2 BDs there, plus - # one for the drain -- and exceeding the limit is a hard compile error, not - # a slowdown. Retire the fills in batches sized to stay under it. - SHIM_BD_LIMIT = 16 - # Cost on the worst shim tile (an A source column, which carries an A fill, - # a B fill and a C drain): 2 BDs per k-block in a fill batch, 1 per drain. - # A sweep's fills must all go in ONE task group. Splitting them across - # groups -- which is what would make larger k_iters fit -- produces - # silently WRONG results, reproducible at M=1024 K=2560 N=2560 - # (k_iters=5): one group passes, two groups fail on the same shape. Not - # yet diagnosed, so the split is not used and the limit is enforced - # instead of being silently mis-lowered. - K_BATCH = (SHIM_BD_LIMIT - 1) // 2 # 7 - if k_iters > K_BATCH: - raise ValueError( - f"K ({K}) needs {k_iters} k-iterations, but at most {K_BATCH} fit " - f"in a shim tile's {SHIM_BD_LIMIT} buffer descriptors alongside " - f"the C drain. Split the GEMM along K, or fix the multi-group " - f"fill path." + 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): - # Sweeps 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 have to drain their share of the broadcast. - sweeps = [(mc, COLS) for mc in range(n_full)] + # 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: - sweeps.append((n_full, rem_blocks)) - - # Retiring each sweep before issuing the next serialises the pipeline: - # the next sweep's fills cannot start until this sweep's C drain has - # come back, which waits on the cores. So issue sweep i+1 in full, and - # only THEN retire sweep i -- the depth-2 model, which is what the - # design this was ported from uses (its notes record that a depth-1 - # "retire immediately" variant hung on hardware). - # - # NOT YET DONE -- this is the operator's main performance gap. Driving - # the original overlay directly measures 8750 GFLOP/s against this - # sequence's 1687 on the same shape and kernel, and the difference is - # here, not in the kernel or the tiling. - # - # Two attempts at overlapping both hung on hardware: an opportunistic - # "retire the oldest group only when BDs run short" scheduler, and a - # strict depth-2 "issue sweep i+1, then retire sweep i". - # - # The idiom itself is not the problem. The original design's sequence - # issues the next iteration's tasks and only THEN awaits and frees the - # previous one's -- and its notes are explicit that the retirement must - # be a real await, not a bare free (freeing early returns the BD slot - # while the transfer is still in flight: the first invocation comes back - # correct and later ones progressively corrupt). TaskGroup.finish() - # awaits then frees, so it already expresses this. - # - # So the hang here has a different, undiagnosed cause -- plausibly BD - # accounting, since the cost model above assumes one BD per transfer. - # Worth knowing before spending much on it: the original hit an - # undiagnosed hang on this path too, which survived even after the - # relevant compiler fix landed, and it still ships with its own IRON - # sequence disabled. - pipelined = False - - prev = None - for mega_col, active_cols in sweeps: - for mega_row in range(m_row_blocks): - # The drain is issued first and retired last: it is an S2MM - # that waits for the cores to produce, so keeping it - # outstanding across the sweep overlaps compute with - # write-back. It must not share a task group with the fills it - # depends on -- finishing those together would deadlock. - tg_c = TaskGroup() - for c in range(active_cols): - c_conses[c].drain( - C, c_tap(mega_col, mega_row, c), group=tg_c, wait=True - ) - - # One A and one B object per k iteration, matching the core's - # k_iters x B_ITERS acquires on each leg. - fill_groups = [] - for batch in range(0, k_iters, K_BATCH): - tg_f = TaskGroup() - for kb in range(batch, min(batch + K_BATCH, k_iters)): - for r in range(ROWS): - a_prods[r].fill(A, a_tap(mega_row, r, kb), group=tg_f) - for c in range(active_cols): - b_prods[c].fill(B, b_tap(mega_col, c, kb), group=tg_f) - if pipelined: - fill_groups.append(tg_f) - else: - tg_f.finish() - - if pipelined: - if prev is not None: - for tg in prev: - tg.finish() - prev = fill_groups + [tg_c] - else: - tg_c.finish() - - if prev is not None: - for tg in prev: - tg.finish() + 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. + for mega_col, active_cols in blocks: + tg_c = TaskGroup() + for c in range(active_cols): + c_conses[c].drain(C, c_tap(mega_col, c), group=tg_c, wait=True) + + tg_f = TaskGroup() + for r in range(ROWS): + a_prods[r].fill(A, a_tap(mega_col, r), group=tg_f) + for c in range(active_cols): + b_prods[c].fill(B, b_tap(mega_col, c), group=tg_f) + tg_f.finish() + tg_c.finish() rt = Runtime( sequence, From 8615cf5a0c310e8edcc08a55689dec032a003b8b Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 17:24:11 -0600 Subject: [PATCH 07/31] flm_gemm: overlap column-blocks, reaching parity with the original Issue column-block i+1's transfers before retiring i's, so they are already moving while i computes. Retiring each block first serialises the whole pipeline on its C await, and a C await waits for the cores. This is the same depth-2 model the design was ported from, and it failed twice before for a reason that had nothing to do with the idiom: when each leg was issued per object a block cost 1 + 2*k_iters buffer descriptors on a shim column, so two in flight blew the 16-BD limit and deadlocked. Now that a leg is a single task a block costs three, and two in flight is six. Collapsing the per-object tasks was the precondition, not an alternative. M=1024 K=1536 N=6144, same session: flm_gemm 2248 us 8834 GFLOP/s FLM peano mm.xclbin 2173 us shipped v1.0.4 2314 us IRON GEMM (r8/f32) 3384 us So the operator is now level with the original overlay -- between its two builds, 1.5x faster than the existing GEMM operator -- while keeping ~40x better accuracy (err/mass 2.4e-4 vs the original's 9.9e-3, which runs in the core's default floor rounding). Two collapses that do NOT work, for the record: B cannot fold the row-block in because it does not vary with it and a 0-stride wrap of size > 1 is rejected by the BD lowering; and draining C once for the whole dispatch hangs. 37/37 pass from a clean build directory. Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 4e8a45c3a..5b1ed0e32 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -403,18 +403,32 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): # 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 now a single task: a + # block costs 3 buffer descriptors on a shim column (A + B + C), so two + # in flight is 6 of 16. Per-object tasks needed 1 + 2*k_iters and could + # not be overlapped at all. + prev = None for mega_col, active_cols in blocks: tg_c = TaskGroup() for c in range(active_cols): c_conses[c].drain(C, c_tap(mega_col, c), group=tg_c, wait=True) - tg_f = TaskGroup() for r in range(ROWS): a_prods[r].fill(A, a_tap(mega_col, r), group=tg_f) for c in range(active_cols): b_prods[c].fill(B, b_tap(mega_col, c), group=tg_f) - tg_f.finish() - tg_c.finish() + + if prev is not None: + for tg in prev: + tg.finish() + prev = [tg_f, tg_c] + + for tg in prev or []: + tg.finish() rt = Runtime( sequence, From ed666a2f5421524ee893c277cb5c76164e6a86eb Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 17:33:37 -0600 Subject: [PATCH 08/31] flm_gemm: drop the inlined epilogue, and the infra it needed The epilogue was compiled to alwaysinline LLVM IR and llvm-linked into the core, on the strength of a comment in the design this was ported from saying that is what pays back the per-object call overhead the C ObjectFifo introduces. Measured here it buys nothing: 7 runs x 100 iterations, 2189 us inlined vs 2182 us as a plain call, against a ~5% run-to-run spread. Indistinguishable. That was the only thing requiring .ll kernel artifacts, so iron/common/compilation/base.py goes back to upstream. Making .ll a first-class artifact type is a reasonable capability -- aie.iron's ExternalFunction(inline=True) emits the merge-mode declaration but does not build anything under IRON's compilation flow, so nothing produced the file -- but it should land on its own merits with a case that measurably needs it, not as shared-infrastructure drift inside an operator change. The epilogue is now an ordinary Kernel call, which also removes the inline_epilogue field, its artifact-naming special case, and the ExternalFunction branch in the design. Diff is now confined to the operator and its kernels. 37/37 pass, and performance is unchanged: min 2183 us vs the original overlay's 2175 us. Co-Authored-By: Claude --- iron/common/compilation/base.py | 23 --------------------- iron/operators/flm_gemm/design.py | 33 +++++-------------------------- iron/operators/flm_gemm/op.py | 28 ++------------------------ 3 files changed, 7 insertions(+), 77 deletions(-) diff --git a/iron/common/compilation/base.py b/iron/common/compilation/base.py index 6bd48a0db..3bc39c220 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -382,15 +382,11 @@ def __init__( extra_flags: list[str] | None = None, rename_symbols: dict[str, str] | None = None, prefix_symbols: str | None = None, - inline_symbol: str | None = None, ) -> None: super().__init__(filename, dependencies) self.extra_flags = extra_flags if extra_flags is not None else [] self.rename_symbols = rename_symbols if rename_symbols is not None else {} self.prefix_symbols = prefix_symbols - # Required when filename is a .ll/.bc: the symbol to mark alwaysinline - # so aiecc inlines it after llvm-linking the module into the core. - self.inline_symbol = inline_symbol class KernelArchiveArtifact(CompilationArtifact): @@ -794,23 +790,6 @@ def compile(self, artifacts): "-Wno-missing-template-arg-list-after-template-kw" ] + compile_args - # A .ll/.bc output means the kernel is meant to be llvm-linked into - # the core and inlined (aie.iron's ExternalFunction(inline=True) - # declares it with link_with_mode = "merge") rather than left as a - # call. That is a Peano-only path; xchesscc has no equivalent. - inline = str(artifact.filename).endswith((".ll", ".bc")) - if inline and self.use_chess: - raise RuntimeError( - f"Kernel artifact '{artifact.filename}' requests an inline " - "(.ll/.bc) build, which requires the Peano compiler; this " - "compilation is configured for chess." - ) - if inline and not artifact.inline_symbol: - raise RuntimeError( - f"Kernel artifact '{artifact.filename}' is an inline " - "(.ll/.bc) build and must name the symbol to inline via " - "KernelObjectArtifact(inline_symbol=...)." - ) commands.append( PythonCallbackCompilationCommand( partial( @@ -821,8 +800,6 @@ def compile(self, artifacts): include_dirs=[str(runtime_lib_include_path)], compile_args=compile_args, use_chess=self.use_chess, - inline=inline, - symbol_name=artifact.inline_symbol if inline else None, ) ) ) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 5b1ed0e32..47f4d3361 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -24,15 +24,12 @@ -D flags, so the C++ and the dataflow cannot drift apart. """ -from pathlib import Path - import numpy as np from ml_dtypes import bfloat16 from aie.helpers.taplib import TensorAccessPattern from aie.iron import ( Buffer, - ExternalFunction, Kernel, ObjectFifo, Program, @@ -92,9 +89,6 @@ def flm_gemm( epilogue="none", kernel_object="flm_gemm.o", epilogue_object="flm_gemm_epilogue.o", - epilogue_source=None, - epilogue_flags=None, - inline_epilogue=False, trace_size=0, ): """Emit the MLIR module for an M x K @ K x N bf16 GEMM. @@ -161,28 +155,11 @@ def flm_gemm( k_step = Kernel( "flm_gemm_k_step", kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty] ) - epilogue_arg_types = [ct_out_ty, ct_acc_ty, np.int32, np.int32] - if inline_epilogue: - # Compile the epilogue to alwaysinline LLVM IR so aiecc llvm-links it - # into the core instead of leaving a call. The epilogue is short and - # runs once per C object, so the call overhead the C ObjectFifo - # introduces is a real cost here -- whereas inlining the much larger - # mmul measures worse, which is why only this one is merged. - if epilogue_source is None: - raise ValueError("inline_epilogue requires epilogue_source") - epilogue_chunk = ExternalFunction( - EPILOGUE_SYMBOL, - object_file_name=str(Path(epilogue_object).with_suffix(".ll")), - source_file=str(epilogue_source), - inline=True, - arg_types=epilogue_arg_types, - include_dirs=[str(Path(epilogue_source).parent)], - compile_flags=list(epilogue_flags or []), - ) - else: - epilogue_chunk = Kernel( - EPILOGUE_SYMBOL, epilogue_object, epilogue_arg_types - ) + epilogue_chunk = Kernel( + EPILOGUE_SYMBOL, + epilogue_object, + [ct_out_ty, ct_acc_ty, np.int32, np.int32], + ) # --- Data movement ---------------------------------------------------- # diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 8ce691ace..8a2edb41a 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -20,7 +20,6 @@ C_DEPTH, CT_OUT_LEN, EPILOGUE_MODES, - EPILOGUE_SYMBOL, K_TILE, MIN_K, MIN_M, @@ -48,9 +47,6 @@ class FLMGEMM(MLIROperator): epilogue: str = field(default="none", repr=False) # Optional (min, max) applied after the activation. clamp: tuple[float, float] | None = field(default=None, repr=False) - # Compile the epilogue to alwaysinline LLVM IR merged into the core, - # rather than leaving it as a call. See design.py. - inline_epilogue: bool = field(default=True, repr=False) context: object = field(default=None, repr=False) _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} @@ -89,13 +85,6 @@ def name(self) -> str: base = f"{base}_epi{self.epilogue}" if self.clamp is not None: base = f"{base}_clamp{self._clamp_tag}" - if not self.inline_epilogue: - # Different core binary, so it must not share an artifact name with - # the inlined build -- otherwise a cached xclbin from one satisfies - # the other and an A/B of the two silently compares a binary with - # itself. (The GEMM operator has this hazard for its - # emulate_bf16_mmul_with_bfp16 / prio_accuracy flags.) - base = f"{base}_noinline" return base @property @@ -108,8 +97,7 @@ def _epilogue_artifact(self) -> str: obj = f"flm_gemm_epilogue_{self.epilogue}" if self.clamp is not None: obj = f"{obj}_clamp{self._clamp_tag}" - # .ll is llvm-linked into the core and inlined; .o is left as a call. - return f"{obj}.ll" if self.inline_epilogue else f"{obj}.o" + return f"{obj}.o" @property def _epilogue_source(self): @@ -117,8 +105,7 @@ def _epilogue_source(self): @property def _epilogue_flags(self) -> list[str]: - """Compile flags for the epilogue, shared by the inline and - separately-compiled paths so the two cannot drift apart.""" + """Compile flags for the epilogue.""" flags = [ f"-DFLM_GEMM_OUT_CHUNK={CT_OUT_LEN}", f"-DFLM_GEMM_C_DEPTH={C_DEPTH}", @@ -154,9 +141,6 @@ def get_mlir_artifact(self): "epilogue": self.epilogue, "kernel_object": self._kernel_object, "epilogue_object": self._epilogue_artifact, - "epilogue_source": str(self._epilogue_source), - "epilogue_flags": self._epilogue_flags, - "inline_epilogue": self.inline_epilogue, "trace_size": 0, }, ), @@ -188,19 +172,11 @@ def get_kernel_artifacts(self): ], ), ] - # Inlined, this is a .ll that aiecc llvm-links into the core; otherwise - # an ordinary .o it calls. Either way the build system produces it -- - # design.py's ExternalFunction only supplies the merge-mode - # declaration in the MLIR, it does not compile anything under IRON's - # compilation flow. artifacts.append( KernelObjectArtifact( self._epilogue_artifact, dependencies=[SourceArtifact(self._epilogue_source)], extra_flags=self._epilogue_flags, - inline_symbol=( - EPILOGUE_SYMBOL if self.inline_epilogue else None - ), ) ) return artifacts From c3bc3b1885ae78d8b455342847e723f863f4f41f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 17:40:52 -0600 Subject: [PATCH 09/31] flm_gemm: expose the rounding mode, and reproduce the shipped kernel exactly The kernels already had the #ifdef; nothing could reach it. rounding= "floor" now selects the core's power-up mode, which is what the design this was ported from runs in -- it never calls set_rounding. With it, the operator is BIT-IDENTICAL to the shipped FastFlowLM v1.0.4 mm.xclbin: all 6291456 elements match on M=1024 K=1536 N=6144, maxdiff 0. That pins the port as arithmetically faithful and isolates rounding as the only numerical difference between the two. conv_even (default) err/mass 0.000241 floor err/mass 0.009867 shipped mm.xclbin err/mass 0.009867 So the 41x accuracy gain really is the rounding mode alone, and users who need to match the shipped overlay can ask for it. The mode is part of the operator name and of both kernel object names: it changes the emitted code, so a cached build of one mode must not satisfy the other. Co-Authored-By: Claude --- iron/operators/flm_gemm/op.py | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 8a2edb41a..728f589e3 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -47,6 +47,12 @@ class FLMGEMM(MLIROperator): epilogue: str = field(default="none", repr=False) # Optional (min, max) applied after the activation. clamp: tuple[float, float] | None = field(default=None, repr=False) + # "conv_even" (round to nearest even) or "floor" (truncate). The core + # powers up in floor, and the design this was ported from never sets the + # mode, so "floor" reproduces its arithmetic exactly -- at ~40x the error, + # because truncation biases every conversion the same way and the bias + # accumulates over the K reduction instead of cancelling. + rounding: str = field(default="conv_even", repr=False) context: object = field(default=None, repr=False) _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} @@ -71,6 +77,10 @@ def __post_init__(self): lo, hi = self.clamp if lo > hi: raise ValueError(f"clamp min ({lo}) must be <= max ({hi})") + if self.rounding not in ("conv_even", "floor"): + raise ValueError( + f"rounding must be 'conv_even' or 'floor', got {self.rounding!r}" + ) MLIROperator.__init__(self, context=self.context) @@ -85,6 +95,8 @@ def name(self) -> str: base = f"{base}_epi{self.epilogue}" if self.clamp is not None: base = f"{base}_clamp{self._clamp_tag}" + if self.rounding != "conv_even": + base = f"{base}_{self.rounding}" return base @property @@ -97,8 +109,16 @@ def _epilogue_artifact(self) -> str: obj = f"flm_gemm_epilogue_{self.epilogue}" if self.clamp is not None: obj = f"{obj}_clamp{self._clamp_tag}" + if self.rounding != "conv_even": + obj = f"{obj}_{self.rounding}" return f"{obj}.o" + @property + def _rounding_flags(self) -> list[str]: + """Applies to both kernels: the mmul and the epilogue's f32->bf16 + store are both conversions and must agree.""" + return ["-DFLM_GEMM_ROUND_FLOOR"] if self.rounding == "floor" else [] + @property def _epilogue_source(self): return self.context.base_dir / "aie_kernels" / "aie2p" / "flm_gemm_epilogue.cc" @@ -120,11 +140,12 @@ def _epilogue_flags(self) -> list[str]: f"-DFLM_GEMM_CLAMP_MIN={float(lo)!r}f", f"-DFLM_GEMM_CLAMP_MAX={float(hi)!r}f", ] - return flags + return flags + self._rounding_flags @property def _kernel_object(self) -> str: - return f"flm_gemm_{M_TILE}x{K_TILE}x{N_TILE}.o" + rnd = "" if self.rounding == "conv_even" else f"_{self.rounding}" + return f"flm_gemm_{M_TILE}x{K_TILE}x{N_TILE}{rnd}.o" def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( @@ -169,7 +190,8 @@ def get_kernel_artifacts(self): # bfp16-emulated path; without this the kernel will not # compile. "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", - ], + ] + + self._rounding_flags, ), ] artifacts.append( From fea39ccb5ea2d3c08e4b65043564b3c24358a6df Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 17:46:49 -0600 Subject: [PATCH 10/31] flm_gemm: document reproducing the shipped overlay, and cover rounding in tests Adds a README covering the parts a caller cannot guess: that B must be pre-packed and why, the shape constraints, the accuracy budget of the bfp16-emulated path (and why an elementwise relative tolerance is the wrong instrument for it), and how to reproduce the shipped FastFlowLM overlay. Verified against FastFlowLM v1.0.4's Gemma4-E2B mm.xclbin, driven directly with the instruction stream from that project's own TXN generator, on identical inputs: rounding="floor" bit-identical, 6291456/6291456 elements, maxdiff 0 rounding="conv_even" differs everywhere (41x more accurate) The activation path matches too. The shipped kernel selects its activation from RTP word 4; this operator bakes it in at compile time with the same 0/1/2/3 mapping, and all four are bit-identical to the shipped kernel at output_mode 0/1/2/3 -- 1048576/1048576 elements each. Inputs were scaled down for that check so the activations sit where the curve is not flat, and each mode's output was confirmed to differ from mode 0, so a silently ignored mode could not pass. clamp has no shipped counterpart: their generate_seq never writes the clamp RTP words, so clamping is always off there. Noted in the README rather than left implicit. Tests now carry rounding as a parameter, with floor given its own error bound -- holding truncation to the conv_even budget would simply fail. Co-Authored-By: Claude --- iron/operators/flm_gemm/README.md | 126 ++++++++++++++++++++++++++++++ iron/operators/flm_gemm/test.py | 46 +++++++---- 2 files changed, 155 insertions(+), 17 deletions(-) create mode 100644 iron/operators/flm_gemm/README.md diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md new file mode 100644 index 000000000..4a210ac28 --- /dev/null +++ b/iron/operators/flm_gemm/README.md @@ -0,0 +1,126 @@ + + +# FLMGEMM — bf16 GEMM on a fixed 4x8 grid + +A second GEMM design, ported from FastFlowLM's `mm` overlay. It is a different +dataflow from [`GEMM`](../gemm), not a retuning of it: + +| | `GEMM` | `FLMGEMM` | +|---|---|---| +| geometry | parameterized tiles, 1–8 columns | fixed 64/512/128, r/s/t 8/8/8, 4x8 grid | +| A delivery | per column | broadcast along each compute row from 4 shim columns | +| C collection | per column | ObjectFifo `join` of 4 rows through the memtile | +| epilogue | separate `convert_copy` | fused f32→bf16 + activation + clamp | +| B layout | plain `(K, N)` | **pre-packed**, see below | + +On NPU2 (aie2p) only: the r=8 mmul shape exists solely on the bfp16-emulated +path, and the grid needs all 8 columns. + +## Shape constraints + +`M % 256 == 0`, `K % 512 == 0`, `N % 128 == 0`. + +N only has to tile to `N_TILE=128`, not to the grid's 1024-wide stride: a +trailing group of fewer than 8 column-blocks 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 1024. + +## B must be pre-packed + +```python +op = FLMGEMM(M=M, K=K, N=N, context=ctx) +C = op.compile().get_callable()(A, FLMGEMM.pack_B(B), C_out) +``` + +`pack_B` reorders a row-major `(K, N)` matrix into the order the memtile +expects — each `K_TILE x N_TILE` tile as the odometer `(n//T, k%S, k//S, n%T)`, +outermost first — so each fill is one contiguous read. + +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, and that was the original implementation, but its +innermost run is then `T=8` bf16 = 16 bytes: each 128 KB transfer becomes 8192 +scattered bursts. B is ~70% of the bytes a dispatch moves, so the whole +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 +FLMGEMM(M=M, K=K, N=N, rounding="floor", context=ctx) # matches shipped +FLMGEMM(M=M, K=K, N=N, context=ctx) # conv_even, default +``` + +Verified against FastFlowLM v1.0.4's +`xclbins/Gemma4-E2B-IT-NPU2/mm.xclbin`, driving it directly with the +instruction stream from that project's own TXN generator, on identical inputs: + +| | err/mass | vs shipped | +|---|---|---| +| `rounding="floor"` | 0.009867 | **bit-identical, 6291456/6291456 elements** | +| `rounding="conv_even"` (default) | 0.000241 | differs everywhere | + +All four epilogues are bit-identical to the shipped kernel too, with `floor` +(the shipped kernel selects its activation from RTP word 4; this operator +bakes it in at compile time, with the same 0/1/2/3 mapping): + +| epilogue | vs shipped `output_mode` | +|---|---| +| `none` / `gelu` / `silu` / `sigmoid` | bit-identical, 1048576/1048576 each | + +**The default is `conv_even`, not `floor`.** 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 original. + +`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 expectations + +The r=8 mmul exists only on the bfp16-emulated path, so the error budget is +that of an emulated GEMM. Do not compare against `GEMM`'s test tolerances, +which assert on the exact r=4 path (`emulate_bf16_mmul_with_bfp16=False`). + +A pure elementwise *relative* tolerance is not meaningful here: 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 instead, as `test.py` does. Reference points on random +signed A / non-negative B: + +| | mean err / mass | +|---|---| +| `FLMGEMM` (default) | 0.00024 | +| `GEMM`, same emulated mode (`emulate=True, prio_accuracy=True`) | 0.00044 | +| `GEMM`, exact r=4 path | 0.00007 | + +## Performance + +M=1024 K=1536 N=6144, 7 processes x 100 iterations, min of per-run medians +(run-to-run spread is ~5%, so differences below that are not meaningful): + +| | latency | +|---|---| +| `FLMGEMM` | 2183 us | +| shipped `mm.xclbin` | 2175 us | +| `GEMM` (`emulate=True, prio_accuracy=True`) | 3353 us | + +Two things dominate, and both are in the runtime sequence rather than the +kernel: B must be pre-packed (above), and each of A, B and C must go out as +**one transfer per column-block** rather than one per fifo object. A single +fill or drain may span many objects; issuing per object instead means a host +await per row-block, and a C await waits on the cores. Collapsing those is +also what makes overlapping column-blocks affordable — a block then costs 3 +buffer descriptors on a shim tile instead of `1 + 2*k_iters`, so two can be in +flight without exhausting the 16 available. diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index 96fee9453..5c26bae06 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -22,23 +22,27 @@ def get_params(): # have N = model dim, so they always land here: 1536 leaves 4 active # columns, 2560 leaves 4, and 128 leaves just 1. # fmt: off - # M, K, N, epilogue, clamp + # M, K, N, epilogue, clamp, rounding regular_params = [ - ( 256, 512, 1024, "none", None), # smallest full sweep - ( 512, 1024, 2048, "none", None), - ( 256, 512, 1536, "none", None), # remainder: 4 of 8 columns - ( 256, 512, 128, "none", None), # remainder only: 1 column - ( 256, 512, 1024, "silu", None), - ( 256, 512, 1024, "gelu", None), - ( 256, 512, 1024, "none", (-2.0, 2.0)), + ( 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 bit for bit; 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), - ( 2048, 2048, 2048, "none", None), - ( 1024, 2560, 2560, "none", None), # E4B o-projection shape - ( 512, 1536, 1536, "silu", None), # E2B down-projection shape - ( 256, 512, 1024, "sigmoid", None), - ( 512, 1024, 2048, "silu", (-4.0, 4.0)), + ( 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"), ] # fmt: on @@ -55,8 +59,8 @@ def get_params(): 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", get_params()) -def test_flm_gemm(M, K, N, epilogue, clamp, aie_context): +@pytest.mark.parametrize("M,K,N,epilogue,clamp,rounding", get_params()) +def test_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): # Keep the activation tests in the range where the curve is not flat. scale = 4.0 if epilogue == "none" else 0.5 golden_ref = generate_golden_reference( @@ -69,6 +73,7 @@ def test_flm_gemm(M, K, N, epilogue, clamp, aie_context): N=N, epilogue=epilogue, clamp=clamp, + rounding=rounding, context=aie_context, ) @@ -98,12 +103,19 @@ def test_flm_gemm(M, K, N, epilogue, clamp, aie_context): mass = K * golden_ref["input"].abs().float().mean() * ( golden_ref["input_b"].abs().float().mean() ) + # + # floor rounding truncates rather than rounding to nearest, so its bias + # accumulates over the K reduction instead of cancelling: ~0.0099 of mass + # rather than ~0.00042, measured, and bit-identical to the shipped overlay. + # It gets a bound to match; holding it to the conv_even budget would just + # fail. + budget = 0.05 if rounding == "floor" else 0.004 errors, latency_us, bandwidth_gbps = run_test( operator, input_buffers, output_buffers, rel_tol=0.04, - abs_tol=float(0.004 * mass), + abs_tol=float(budget * mass), ) gflops = (2.0 * M * K * N) / (latency_us * 1e-6) / 1e9 From cc03cd513ea18b74fe7b84f27b953a5fbca5f6a2 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 18:39:20 -0600 Subject: [PATCH 11/31] flm_gemm: pick the n tile from the shape, and go 20% faster Nulling the mmul out showed this operator is COMPUTE bound -- it drops from 2178 us to 1481 us with no arithmetic -- while the GEMM operator does not move at all under the same treatment (3374 vs 3353 us) and is therefore entirely data-movement bound. The two want opposite fixes, and this one wants a cheaper inner loop. n=64 gives the mmul colA=8 instead of 4, halving accumulator traffic per mac, at the cost of doubling A fetches. With compute on the critical path that trades well: M/K/N k_iters tile_n=64 tile_n=128 1024/512/4096 1 642 us 589 us 1024/1024/4096 2 850 us 1034 us 1024/1536/6144 3 1741 us 2178 us 1024/2560/4096 5 1891 us 2371 us 2048/2048/2048 4 1535 us 1924 us 256/4096/1024 8 254 us 316 us Only the single-k-iteration shape prefers 128: there is too little compute there to hide the extra A traffic. tile_n now defaults to None and picks 128 when K == 512 and 64 otherwise, which is correct on every shape measured including the K=1024 boundary. It stays overridable. At 1741 us the operator is now 1.25x faster than the shipped overlay (2175 us) and 1.9x faster than GEMM (3353 us) on the reference shape. pack_B becomes an instance method: the packing layout depends on tile_n, so a static one silently mismatches the operator it feeds. tile_n also joins the operator and kernel-object names. This vindicates the CT_MAX_K hypothesis from early on, which had been dismissed after testing it while the operator was still DMA bound at 11 ms -- where compute could not matter. Right idea, wrong regime. 39/39 pass. Co-Authored-By: Claude --- iron/operators/flm_gemm/README.md | 59 +++++++++++++++++++++++++------ iron/operators/flm_gemm/design.py | 41 +++++++++++++-------- iron/operators/flm_gemm/op.py | 34 ++++++++++++++---- iron/operators/flm_gemm/test.py | 3 +- 4 files changed, 105 insertions(+), 32 deletions(-) diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index 4a210ac28..085af8e88 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -10,8 +10,9 @@ dataflow from [`GEMM`](../gemm), not a retuning of it: | | `GEMM` | `FLMGEMM` | |---|---|---| -| geometry | parameterized tiles, 1–8 columns | fixed 64/512/128, r/s/t 8/8/8, 4x8 grid | +| geometry | parameterized tiles, 1–8 columns | fixed m=64 k=512, r/s/t 8/8/8, 4x8 grid; n selectable | | A delivery | per column | broadcast along each compute row from 4 shim columns | +| C staging | full m x n tile in L1 | streamed out in 512-element chunks | | C collection | per column | ObjectFifo `join` of 4 rows through the memtile | | epilogue | separate `convert_copy` | fused f32→bf16 + activation + clamp | | B layout | plain `(K, N)` | **pre-packed**, see below | @@ -21,23 +22,23 @@ path, and the grid needs all 8 columns. ## Shape constraints -`M % 256 == 0`, `K % 512 == 0`, `N % 128 == 0`. +`M % 256 == 0`, `K % 512 == 0`, `N % tile_n == 0` (so 64 by default). -N only has to tile to `N_TILE=128`, not to the grid's 1024-wide stride: a +N only has to tile to `tile_n`, not to the grid's `tile_n * 8` stride: a trailing group of fewer than 8 column-blocks 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 1024. +of the full stride. ## B must be pre-packed ```python op = FLMGEMM(M=M, K=K, N=N, context=ctx) -C = op.compile().get_callable()(A, FLMGEMM.pack_B(B), C_out) +op.compile().get_callable()(A, op.pack_B(B), C_out) ``` `pack_B` reorders a row-major `(K, N)` matrix into the order the memtile -expects — each `K_TILE x N_TILE` tile as the odometer `(n//T, k%S, k//S, n%T)`, +expects — each `K_TILE x tile_n` tile as the odometer `(n//T, k%S, k//S, n%T)`, outermost first — so each fill is one contiguous read. This is deliberately the caller's job rather than something the fill @@ -105,16 +106,52 @@ signed A / non-negative B: | `GEMM`, same emulated mode (`emulate=True, prio_accuracy=True`) | 0.00044 | | `GEMM`, exact r=4 path | 0.00007 | +## Choosing `tile_n` + +`tile_n` defaults to `None`, which picks per shape: **128 when `K == 512`, +otherwise 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 1024 +columns of N per pass rather than 512. Which wins depends on whether compute +or data movement is the critical path, and 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. Measured, minimum of 3 runs: + +| M / K / N | k_iters | `tile_n=64` | `tile_n=128` | +|---|---|---|---| +| 1024 / 512 / 4096 | 1 | 642 us | **589 us** | +| 1024 / 1024 / 4096 | 2 | **850 us** | 1034 us | +| 1024 / 1536 / 6144 | 3 | **1741 us** | 2178 us | +| 1024 / 2560 / 4096 | 5 | **1891 us** | 2371 us | +| 2048 / 2048 / 2048 | 4 | **1535 us** | 1924 us | +| 256 / 4096 / 1024 | 8 | **254 us** | 316 us | + +`pack_B` is bound to the operator because the packing layout depends on +`tile_n`; call `op.pack_B(B)`, not `FLMGEMM.pack_B(B)`. + ## Performance M=1024 K=1536 N=6144, 7 processes x 100 iterations, min of per-run medians (run-to-run spread is ~5%, so differences below that are not meaningful): -| | latency | -|---|---| -| `FLMGEMM` | 2183 us | -| shipped `mm.xclbin` | 2175 us | -| `GEMM` (`emulate=True, prio_accuracy=True`) | 3353 us | +| | bytes moved | latency | DMA-only (compute nulled) | +|---|---|---|---| +| `FLMGEMM` (`tile_n=64`) | 126 MB | **1741 us** | -- | +| `FLMGEMM` (`tile_n=128`) | 107 MB | 2178 us | 1481 us | +| shipped `mm.xclbin` | 107 MB | 2175 us | -- | +| `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | + +Nulling the mmul out is what makes this legible. `GEMM` does not change at all +without it (3374 vs 3353 us), so it is entirely data-movement bound; this +operator drops to 1481 us, so it is compute bound with its transfers hidden. +That is why the two respond to opposite fixes: `GEMM` would want cheaper +transfers, this design wants a cheaper inner loop -- which is what `tile_n=64` +buys. + +Its transfers are cheaper mostly because B arrives pre-packed: the contiguous +run per transfer is 128 KB for B and 1 KB for A, against 128 bytes on every +leg for `GEMM`, which reorders in the descriptor instead. Two things dominate, and both are in the runtime sequence rather than the kernel: B must be pre-packed (above), and each of A, B and C must go out as diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 47f4d3361..a49c1df45 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -43,7 +43,16 @@ # --- Fixed geometry ------------------------------------------------------- # GEMM tiling per compute tile, and the register tiling inside it. -M_TILE, K_TILE, N_TILE = 64, 512, 128 +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 +# k-slice per n width; must match compute_CT_k_max_n in flm_gemm_geometry.h +CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 64, 128: 32, 256: 16} R, S, T = 8, 8, 8 ROWS, COLS = 4, 8 @@ -52,19 +61,8 @@ # gemm operator pins A the same way in the 8-column case. A_SOURCE_COL = [0, 2, 4, 6] -# Must match compute_CT_k_max_n() in flm_gemm_geometry.h: with n=128, -# a compute tile holds 32 of K at a time. -CT_MAX_K = 32 -K_DIV_CT_K_MAX = K_TILE // CT_MAX_K - -# Buffer lengths, in elements. -CT_A_LEN = 2 * R * CT_MAX_K # one z slice -CT_A_OBJ = CT_A_LEN * (M_TILE // R // 2) # A object: every z slice of one mmul CT_OUT_LEN = 512 # the core's C slice, streamed out in chunks this size -C_SLICE_LEN = M_TILE * N_TILE # one compute tile's C contribution -O_CHUNKS = C_SLICE_LEN // CT_OUT_LEN # C objects one accumulator drains as C_DEPTH = 2 # C fifo depth; also the core-body unroll -B_ITERS = K_TILE // CT_MAX_K # B chunks the core consumes per k step B_DEPTH = 2 # B fifo depth; also the core-body unroll A_DEPTH = 2 @@ -75,10 +73,10 @@ # to mark the symbol alwaysinline when building the inline .ll variant). EPILOGUE_SYMBOL = "flm_gemm_epilogue_chunk" -# Minimum problem size, i.e. one pass of the whole grid. +# Minimum problem size, i.e. one pass of the whole grid. MIN_N depends on the +# chosen n tile, so it is computed per call. MIN_M = M_TILE * ROWS # 256 MIN_K = K_TILE # 512 -MIN_N = N_TILE * COLS # 1024 def flm_gemm( @@ -87,6 +85,7 @@ def flm_gemm( K, N, epilogue="none", + tile_n=N_TILE_DEFAULT, kernel_object="flm_gemm.o", epilogue_object="flm_gemm_epilogue.o", trace_size=0, @@ -97,6 +96,20 @@ def flm_gemm( bf16 and all plain dense tensors -- the block-major reordering B needs on the way in is done by the fill descriptor, not by the caller. """ + 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}" + ) + N_TILE = tile_n + CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + 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 * (M_TILE // 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 + if epilogue not in EPILOGUE_MODES: raise ValueError( f"epilogue must be one of {sorted(EPILOGUE_MODES)}, got {epilogue!r}" diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 728f589e3..17e5e6794 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -24,7 +24,7 @@ MIN_K, MIN_M, M_TILE, - N_TILE, + N_TILE_DEFAULT, S, T, ) @@ -47,6 +47,10 @@ class FLMGEMM(MLIROperator): epilogue: str = field(default="none", repr=False) # Optional (min, max) applied after the activation. clamp: tuple[float, float] | None = field(default=None, repr=False) + # n tile width. 64 halves the mmul's accumulator traffic per mac; 128 + # halves A fetches instead and wins only when small K makes the operator + # DMA-bound. See README.md. + tile_n: int | None = field(default=None, repr=False) # "conv_even" (round to nearest even) or "floor" (truncate). The core # powers up in floor, and the design this was ported from never sets the # mode, so "floor" reproduces its arithmetic exactly -- at ~40x the error, @@ -58,13 +62,15 @@ class FLMGEMM(MLIROperator): _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} def __post_init__(self): + if self.tile_n is None: + self.tile_n = self._default_tile_n(self.K) # 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, MIN_M), ("K", self.K, MIN_K), - ("N", self.N, N_TILE), + ("N", self.N, self.tile_n), ): if value % unit != 0: raise ValueError(f"{name} ({value}) must be a multiple of {unit}") @@ -84,6 +90,19 @@ def __post_init__(self): MLIROperator.__init__(self, context=self.context) + @staticmethod + def _default_tile_n(K: int) -> int: + """Pick the n tile from the shape. + + 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, and that is set by how + much K there is to reduce over: with a single k iteration there is too + little compute to hide the extra A traffic. Measured ~20% for n=64 at + K >= 1024 and ~9% the other way at K = 512. + """ + return 128 if K // K_TILE <= 1 else 64 + @property def name(self) -> str: # epilogue/clamp are repr=False so the plain path keeps a stable name, @@ -97,6 +116,8 @@ def name(self) -> str: base = f"{base}_clamp{self._clamp_tag}" if self.rounding != "conv_even": base = f"{base}_{self.rounding}" + if self.tile_n != self._default_tile_n(self.K): + base = f"{base}_tn{self.tile_n}" return base @property @@ -145,7 +166,7 @@ def _epilogue_flags(self) -> list[str]: @property def _kernel_object(self) -> str: rnd = "" if self.rounding == "conv_even" else f"_{self.rounding}" - return f"flm_gemm_{M_TILE}x{K_TILE}x{N_TILE}{rnd}.o" + return f"flm_gemm_{M_TILE}x{K_TILE}x{self.tile_n}{rnd}.o" def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( @@ -159,6 +180,7 @@ def get_mlir_artifact(self): "M": self.M, "K": self.K, "N": self.N, + "tile_n": self.tile_n, "epilogue": self.epilogue, "kernel_object": self._kernel_object, "epilogue_object": self._epilogue_artifact, @@ -185,7 +207,7 @@ def get_kernel_artifacts(self): extra_flags=[ f"-DFLM_GEMM_TILE_M={M_TILE}", f"-DFLM_GEMM_TILE_K={K_TILE}", - f"-DFLM_GEMM_TILE_N={N_TILE}", + f"-DFLM_GEMM_TILE_N={self.tile_n}", # The r=8 mmul shape this design uses only exists on the # bfp16-emulated path; without this the kernel will not # compile. @@ -203,8 +225,7 @@ def get_kernel_artifacts(self): ) return artifacts - @staticmethod - def pack_B(B): + def pack_B(self, B): """Reorder a row-major ``(K, N)`` weight matrix into the layout the B fill expects. Returns a flat tensor. @@ -222,6 +243,7 @@ def pack_B(B): and reused across dispatches, so the cost belongs here. """ K, N = B.shape + N_TILE = self.tile_n if K % K_TILE or N % N_TILE: raise ValueError( f"B ({K}, {N}) must tile to ({K_TILE}, {N_TILE}) to be packed" diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index 5c26bae06..b76791725 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -67,6 +67,7 @@ def test_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): M=M, K=K, N=N, epilogue=epilogue, clamp=clamp, scale=scale ) + operator = FLMGEMM( M=M, K=K, @@ -80,7 +81,7 @@ def test_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): input_buffers = { "A": golden_ref["input"].flatten(), # B is consumed pre-packed; see FLMGEMM.pack_B. - "B": FLMGEMM.pack_B(golden_ref["input_b"]), + "B": operator.pack_B(golden_ref["input_b"]), } output_buffers = {"C": golden_ref["output"].flatten()} From 053996e66e52b1d511cf9793f7defb3e24447ad5 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 20:03:14 -0600 Subject: [PATCH 12/31] flm_gemm: correct the accuracy comparison in the README The table claimed GEMM in the same emulated mode had 0.00044 mean error against this operator's 0.00024. That was wrong: 0.00044 came from a different configuration (prio_accuracy=False, a bf16 accumulator), quoted from a separate experiment and mislabelled as the f32 one. Measured on identical data, this operator and GEMM with emulate_bf16_mmul_with_bfp16=True and prio_accuracy=True are numerically indistinguishable -- err/mass 0.000241, signed bias +0.0871, max 15.44 for both, and the same at either tile_n. Same mmul shape, same emulation, same f32 accumulation, same rounding; there was never a reason for them to differ. The only real accuracy difference remains conv_even versus the shipped overlay's floor, which stands at 41x and is separately verified bit-exact. Co-Authored-By: Claude --- iron/operators/flm_gemm/README.md | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index 085af8e88..8ddcf98d0 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -102,9 +102,17 @@ signed A / non-negative B: | | mean err / mass | |---|---| -| `FLMGEMM` (default) | 0.00024 | -| `GEMM`, same emulated mode (`emulate=True, prio_accuracy=True`) | 0.00044 | -| `GEMM`, exact r=4 path | 0.00007 | +| `FLMGEMM` (default) | 0.000241 | +| `GEMM`, same mode (`emulate=True, prio_accuracy=True`) | 0.000241 | +| `GEMM`, bf16 accumulator (`prio_accuracy=False`) | 0.000445 | +| `GEMM`, exact r=4 path (`emulate=False`) | 0.00007 | + +This operator and `GEMM` in the same mode are numerically **indistinguishable** +-- identical mean error, signed bias and maximum, at both `tile_n` values. +Same mmul shape, same bfp16 emulation, same f32 accumulation, same rounding, +so there is no reason for them to differ and they do not. The only accuracy +difference worth knowing about is `conv_even` versus the shipped overlay's +`floor` (above). ## Choosing `tile_n` From a5c51c887aabd32825c7ffaaae2c72f083e02ac2 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 21:55:19 -0600 Subject: [PATCH 13/31] deps: bump mlir-aie to 1.4.3.dev60 and Peano to 2026090201 Takes the Peano pin from mlir-aie's own utils/peano-requirements.txt at c80b88c, so the two stay in step. 1.4.3 adds ObjectFifo DMA channel pinning (prod_dma_channel / cons_dma_channels, exposed as prod(channel=)/cons(channel=)), which is absent from 1.4.2.dev16. 39/39 of the flm_gemm and gemm iter0 tests pass on the new toolchain, and flm_gemm's latency is unchanged within noise at M=1024 K=1536 N=6144 (1743 us, against 1741 before the bump). Co-Authored-By: Claude --- requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index f626f2f26..d0bd8880d 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.2.dev16+g7e00b57 -llvm-aie==22.0.0.2026082001+84660bc3 +mlir_aie==1.4.3.dev60+gc80b88c +llvm-aie==22.0.0.2026090201+a36c62b9 black reuse From 7c97d5b78dbcca6dec041ea02187e4456fcddf96 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Tue, 8 Sep 2026 21:55:32 -0600 Subject: [PATCH 14/31] flm_gemm: keep B resident in the memtile, for 43% less DDR traffic Hold a whole column-block's B in the memtile as one object and replay it per row-block, so DDR reads it once rather than m_row_blocks times. B is the dominant DDR leg -- 75 MB of the 126 MB moved at M=1024 K=1536 N=6144 -- so this is ~43% less traffic overall. This operator is DDR-bandwidth bound, so that is a latency win as well as a power one, and it grows with the height of the problem because B's re-reads scale with m_row_blocks. At K=1024 N=4096, full builds, min of interleaved rounds: M=512 (2 row-blocks) 470.8 -> 468.5 us 0.5% M=1024 (4 row-blocks) 860.4 -> 846.5 us 1.6% M=2048 (8 row-blocks) 1760.1 -> 1622.8 us 7.8% With the mmul nulled out to isolate data movement, the non-resident floor at M=2048 is 1739 us to move 118 MB -- 68 GB/s, against a memcpy-measured 63-70 GB/s roof -- and residency drops it to 1135 us. Only 137 us of that 604 us is captured, because repeat_count restarts the memtile BD chain at every replay boundary. Closing that gap is the largest known remaining lever here; it is left for follow-up. The replay is repeat_count on the forward(). Earlier attempts used iter_count and hung on every shape with more than one column-block: iter_count does not replay anything, it only bounds how many times an end cycles through all of its buffers (objects = iter_count * elemNumber * repeat_count), so it counts depth-cycles, not objects, and a wrong value runs the BD chain out from under the cores. Correct k ordering needs the memtile to hold ONE object spanning every k-block; replaying a pool of k_iters smaller objects emits k0,k0,k1,k1, ... instead of the k0..kn sequence the cores accumulate in. Residency is gated on the buffer fitting double-buffered, which admits only k_iters <= 2 (K <= 1024 at tile_n=64); larger K falls back to the previous behaviour unchanged. Co-Authored-By: Claude --- iron/operators/flm_gemm/design.py | 66 ++++++++++++++++++++++++++++++- 1 file changed, 64 insertions(+), 2 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index a49c1df45..6b1c316af 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -158,6 +158,7 @@ def flm_gemm( # L2 (per memtile) mt_a_ty = np.ndarray[(M_TILE * K_TILE,), bf16_ty] mt_b_ty = np.ndarray[(K_TILE * N_TILE,), bf16_ty] + mt_b_bytes = K_TILE * N_TILE * 2 mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] # L3 (DDR), flat -- the taps below index them linearly. a_l3_ty = np.ndarray[(M * K,), bf16_ty] @@ -240,6 +241,55 @@ def flm_gemm( a_cons[(r, c)] = of_a.cons() # B: shim -> memtile -> broadcast down the compute column. + # Resident B: hold a whole column-block's B in the memtile as ONE object + # and re-walk it per row-block, so DDR sees it once instead of + # m_row_blocks times. B is the dominant DDR leg -- at M=1024 K=1536 N=6144 + # it is 75 MB of the 126 MB moved -- so this is ~43% less total traffic. + # + # It saves power and time both, and the time is worth more the taller the + # problem is, because B's DDR re-reads scale with m_row_blocks. Measured at + # K=1024 N=4096 (full builds, min of interleaved rounds): + # + # M=512 (2 row-blocks) 470.8 -> 468.5 us 0.5% + # M=1024 (4 row-blocks) 860.4 -> 846.5 us 1.6% + # M=2048 (8 row-blocks) 1760.1 -> 1622.8 us 7.8% + # + # The operator IS DDR-bandwidth bound: with the mmul nulled out, the + # non-resident floor at M=2048 is 1739 us for 118 MB, i.e. 68 GB/s, against + # a measured 63-70 GB/s roof. Residency drops that floor to 1135 us. + # + # Note the gap between that 604 us of floor and the 137 us actually + # captured: repeat_count restarts the memtile BD chain at every replay + # boundary, and ~77% of the win goes there. Closing it is the largest known + # remaining lever on this design. Do not measure it at small M -- at M=512 + # the effect is inside the noise, which is how it was first mistaken for a + # power-only optimisation. + # + # One object, not k_iters of them: iterating a pool replays each object in + # turn (k0,k0,k1,k1,...) rather than the sequence. Fitting k into the + # descriptors within the 4-dimension limit takes both hops -- inbound the + # outermost dim already steps by (N_TILE//T)*(K_TILE*T), exactly + # K_TILE*N_TILE, so widening its COUNT walks into the next k-block; + # outbound k becomes a new outermost dim, which also keeps one emitted + # object per CT_MAX_K slice. + # + # Replay is repeat_count, on the forward() below. iter_count cannot do this + # job: it only bounds how many times an end cycles through all its buffers + # (objects = iter_count * elemNumber * repeat_count), so it is in units of + # depth-cycles rather than objects, and getting it wrong hangs rather than + # mis-computes. + # + # Gated on the buffer fitting DOUBLE-buffered, so the next column-block + # still prefetches; A takes 128 KB and C 64 KB of the 512 KB memtile. + # NOTE this admits only k_iters <= 2, i.e. K <= 1024 at tile_n=64. Larger K + # silently falls back to non-resident, so check this gate before believing + # any measurement that claims to be testing residency. + b_resident = (k_iters * mt_b_bytes * 2) <= (512 - 128 - 64) * 1024 + if b_resident: + mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE,), bf16_ty] + b_recv_dims = [(k_iters * (N_TILE // T), K_TILE * T)] + b_recv_dims[1:] + b_send_dims = [(k_iters, K_TILE * N_TILE)] + b_send_dims + b_l3l2_fifos = [] b_cons = {} for c in range(n_active_cols): @@ -251,6 +301,14 @@ def flm_gemm( depth=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() @@ -364,8 +422,12 @@ def b_tap(mega_col, c): return TensorAccessPattern( tensor_dims=(K * N,), offset=(mega_col * COLS + c) * N_TILE * K, - sizes=[m_row_blocks, k_iters, 1, K_TILE * N_TILE], - strides=[0, K_TILE * N_TILE, 0, 1], + sizes=( + [1, 1, 1, k_iters * K_TILE * N_TILE] + if b_resident + else [m_row_blocks, k_iters, 1, K_TILE * N_TILE] + ), + strides=([0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE, 0, 1]), ) def c_tap(mega_col, c): From 2830a2a297992ad3b28d1adae0798c6f8fc88139 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 08:42:48 -0600 Subject: [PATCH 15/31] flm_gemm: document resident B, and how to measure this operator Three gaps in the README, all found by re-deriving its numbers. Resident B was undocumented. It is gated on a whole column-block's B fitting the memtile double-buffered (K <= 1024 at tile_n=64), so it is inactive on the shape the performance table reports and easy to miss entirely. Record what it does, where it applies, and that its benefit grows with M -- 0.5% / 1.6% / 7.8% at M=512/1024/2048 -- along with the 604 us of DMA floor it exposes against the 137 us that currently reaches the full build. The measurement recipe was unsafe. "min of per-run medians" is only sound if the runs are not all in the same place: dispatch latency here is bimodal with modes ~6% apart, so a batch landing wholly in one mode makes that statistic a mode selector. It reported a convincing 5% win for a change subsequently shown to do nothing. Say to interleave the configurations and to require the min and the median to agree. The tile_n=64 DMA-only cell was blank. Filling it in (1692 of 1741 us) shows the default configuration is itself data-movement bound, which the surrounding prose implied only of tile_n=128 -- and which is what makes "move fewer bytes" the right next lever rather than "write a faster kernel". Co-Authored-By: Claude --- iron/operators/flm_gemm/README.md | 51 ++++++++++++++++++++++++++----- 1 file changed, 43 insertions(+), 8 deletions(-) diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index 8ddcf98d0..f52acde0a 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -140,22 +140,30 @@ hide the extra A traffic. Measured, minimum of 3 runs: ## Performance -M=1024 K=1536 N=6144, 7 processes x 100 iterations, min of per-run medians -(run-to-run spread is ~5%, so differences below that are not meaningful): +M=1024 K=1536 N=6144, min of per-run medians across separate processes: | | bytes moved | latency | DMA-only (compute nulled) | |---|---|---|---| -| `FLMGEMM` (`tile_n=64`) | 126 MB | **1741 us** | -- | +| `FLMGEMM` (`tile_n=64`) | 126 MB | **1741 us** | 1692 us | | `FLMGEMM` (`tile_n=128`) | 107 MB | 2178 us | 1481 us | | shipped `mm.xclbin` | 107 MB | 2175 us | -- | | `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | +**Measure this carefully.** 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 -- that is how a change later shown to do nothing at all first +produced a convincing 5% "win". 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. + Nulling the mmul out is what makes this legible. `GEMM` does not change at all -without it (3374 vs 3353 us), so it is entirely data-movement bound; this -operator drops to 1481 us, so it is compute bound with its transfers hidden. -That is why the two respond to opposite fixes: `GEMM` would want cheaper -transfers, this design wants a cheaper inner loop -- which is what `tile_n=64` -buys. +without it (3374 vs 3353 us), so it is entirely data-movement bound. At +`tile_n=128` this operator drops to 1481 us, so *there* it is compute bound +with its transfers hidden -- which is what `tile_n=64` fixes, buying a much +cheaper inner loop at the price of more data movement. But note the default +`tile_n=64` is then data-movement bound itself (1692 of 1741 us), so further +gains there come from moving fewer bytes, not from a faster kernel. Its transfers are cheaper mostly because B arrives pre-packed: the contiguous run per transfer is 128 KB for B and 1 KB for A, against 128 bytes on every @@ -169,3 +177,30 @@ await per row-block, and a C await waits on the cores. Collapsing those is also what makes overlapping column-blocks affordable — a block then costs 3 buffer descriptors on a shim tile 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. + +This operator is DDR-bandwidth bound, so that is a latency win as well as a +power one, and it grows with the height of the problem because B's re-reads +scale with `m_row_blocks`. At K=1024 N=4096: + +| M | row-blocks | non-resident | resident | | +|---|---|---|---|---| +| 512 | 2 | 470.8 us | 468.5 us | 0.5% | +| 1024 | 4 | 860.4 us | 846.5 us | 1.6% | +| 2048 | 8 | 1760.1 us | **1622.8 us** | 7.8% | + +Do not evaluate this at small M: at M=512 the effect is inside the noise, which +is how it was first mistaken for a power-only optimisation. + +Most of the available win is still on the table. With the mmul nulled, the +non-resident floor at M=2048 is 1739 us -- 118 MB at 68 GB/s, against a +memcpy-measured 63-70 GB/s roof for mixed read/write traffic -- and residency +drops that floor to 1135 us. Only 137 us of those 604 us reaches the full +build; the rest goes to `repeat_count` restarting the memtile BD chain at every +replay boundary. Closing that is the largest known remaining lever here. From f09404ccf7d9a5056d4800ac6f499997c7d60390 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 11:20:08 -0600 Subject: [PATCH 16/31] flm_gemm: re-roll the mmul loop and keep B resident, for 1.52x Two changes that only pay together, taking M=1024 K=1536 N=6144 from 1741 us to 1434 us against the shipped overlay's 2175 us -- 1.25x to 1.52x -- at bit-identical arithmetic (err/mass 2.41e-04 either way). The mmul's colA loop was hand-unrolled 2x for the Peano path, on the premise that "Peano schedules the 2x-unrolled body better than it schedules the rolled one". That is no longer true, and the unroll had become a 23% pessimization: it halves the inner trip count to 4, too few to amortize the software pipeline's fill and drain. Measured cycles per mmul call, by hardware trace of the event0/event1 pair the kernel already emits: rolled 1875, hand-unrolled 2x 2446, compiler unroll 4 and 8 3291 and 3221. Chess always took the rolled path, so the branch is simply deleted rather than re-gated. The rolled loop is a sharp optimum -- anything adding live state across it loses more than it gains, which is the signature of llvm-aie#1066, where Peano's pipelining-unaware register allocator manufactures false loop-carried anti-dependencies. Loop hints, hoisting the bfp16 operand conversion, and prefetching the accumulator all measured neutral or worse. Widening the residency gate is what banks the win. It only counts C's 64 KB, so it admits k_iters <= 3 (K <= 1536) rather than 2, overcommitting the A-carrying memtiles by 64 KB and relying on aie-objectfifo-allocate to spill one buffer to the adjacent memtile. That packs all eight to exactly 512 KB, with zero slack -- see the comment in design.py before changing any buffer size. Neither change is worth much alone: residency was measured latency-neutral while the mmul was the critical path, and the re-rolled mmul gained little while the non-resident DMA floor was. Total latency is max(compute, DMA), so testing them separately scores both as zero. 39/39 green. --- aie_kernels/aie2p/flm_gemm_mmul.h | 77 ++++--------------------------- iron/operators/flm_gemm/README.md | 38 +++++++++++---- iron/operators/flm_gemm/design.py | 27 +++++++++-- 3 files changed, 61 insertions(+), 81 deletions(-) diff --git a/aie_kernels/aie2p/flm_gemm_mmul.h b/aie_kernels/aie2p/flm_gemm_mmul.h index 95488d680..d4dedf978 100644 --- a/aie_kernels/aie2p/flm_gemm_mmul.h +++ b/aie_kernels/aie2p/flm_gemm_mmul.h @@ -66,10 +66,15 @@ flm_gemm_mmul_2x2(const T_in *__restrict pA, const T_in *__restrict pB, MMUL C10(aie::load_v(pC2)); MMUL C11(aie::load_v(pC2 + MMUL::size_C)); - static_assert(colA % 2 == 0); - // Peano schedules the 2x-unrolled body better than it schedules the - // rolled one; chess does not need the hand-unroll. -#if defined(__chess__) + // Rolled, deliberately. An earlier 2x hand-unroll for the Peano path + // is now a pessimization: with colA/2 = 4 trip counts there are too + // few iterations to amortize the software pipeline's fill/drain. + // Measured cycles per mmul call (HW trace, event0/event1): rolled + // 1875, hand-unrolled 2x 2446, compiler unroll 4/8 3291/3221. Any + // extra live state across this loop loses more than it gains -- + // Peano's register allocator is pipelining-unaware and manufactures + // false loop-carried anti-deps (llvm-aie#1066), so keep the body + // minimal and let the pipeliner overlap the iterations. AIE_LOOP_MAX_ITERATION_COUNT(colA) for (unsigned i = 0; i < colA; i++) { A0 = aie::load_v(pA1); @@ -101,70 +106,6 @@ flm_gemm_mmul_2x2(const T_in *__restrict pA, const T_in *__restrict pB, C10.mac(A1, B0); C11.mac(A1, B1); } -#else - AIE_LOOP_MAX_ITERATION_COUNT(colA / 2) - for (unsigned i = 0; i < colA; i += 2) { - // First iteration - A0 = aie::load_v(pA1); - pA1 += MMUL::size_A; - A1 = aie::load_v(pA2); - pA2 += MMUL::size_A; - - if constexpr (b_row_maj) { - B0 = aie::load_v(pB1); - pB1 += MMUL::size_B * colB; - B1 = aie::load_v(pB2); - pB2 += MMUL::size_B * colB; - } else { - if constexpr (is_b_s_t_in_row_major == false) { - B0 = aie::transpose(aie::load_v(pB1), t, s); - pB1 += MMUL::size_B; - B1 = aie::transpose(aie::load_v(pB2), t, s); - pB2 += MMUL::size_B; - } else { - 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); - - // Second iteration - A0 = aie::load_v(pA1); - pA1 += MMUL::size_A; - A1 = aie::load_v(pA2); - pA2 += MMUL::size_A; - - if constexpr (b_row_maj) { - B0 = aie::load_v(pB1); - pB1 += MMUL::size_B * colB; - B1 = aie::load_v(pB2); - pB2 += MMUL::size_B * colB; - } else { - if constexpr (is_b_s_t_in_row_major == false) { - B0 = aie::transpose(aie::load_v(pB1), t, s); - pB1 += MMUL::size_B; - B1 = aie::transpose(aie::load_v(pB2), t, s); - pB2 += MMUL::size_B; - } else { - 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); - } -#endif aie::store_v(pC1, C00.template to_vector()); pC1 += MMUL::size_C; aie::store_v(pC1, C01.template to_vector()); diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index f52acde0a..390c2776b 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -126,6 +126,12 @@ or data movement is the critical path, and 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. Measured, minimum of 3 runs: +> **Stale:** the sweep below predates the re-rolled mmul and the widened +> residency gate, which together took 1024/1536/6144 from 1741 to 1434 us. The +> `tile_n` choice it justifies is unlikely to have changed sign (residency does +> not fit at `tile_n=128`, whose `mt_b` is 128 KB), but the absolute numbers +> are no longer right and it wants re-measuring. + | M / K / N | k_iters | `tile_n=64` | `tile_n=128` | |---|---|---|---| | 1024 / 512 / 4096 | 1 | 642 us | **589 us** | @@ -144,11 +150,14 @@ M=1024 K=1536 N=6144, min of per-run medians across separate processes: | | bytes moved | latency | DMA-only (compute nulled) | |---|---|---|---| -| `FLMGEMM` (`tile_n=64`) | 126 MB | **1741 us** | 1692 us | -| `FLMGEMM` (`tile_n=128`) | 107 MB | 2178 us | 1481 us | +| `FLMGEMM` (`tile_n=64`) | 69 MB | **1434 us** | 1231 us | | shipped `mm.xclbin` | 107 MB | 2175 us | -- | | `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | +**1.52x the shipped overlay**, at identical arithmetic (err/mass 2.41e-04 +either way). The 69 MB is with B resident in the memtile; the 126 MB figure +this table used to quote was the non-resident fallback. + **Measure this carefully.** 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 @@ -158,12 +167,25 @@ 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. Nulling the mmul out is what makes this legible. `GEMM` does not change at all -without it (3374 vs 3353 us), so it is entirely data-movement bound. At -`tile_n=128` this operator drops to 1481 us, so *there* it is compute bound -with its transfers hidden -- which is what `tile_n=64` fixes, buying a much -cheaper inner loop at the price of more data movement. But note the default -`tile_n=64` is then data-movement bound itself (1692 of 1741 us), so further -gains there come from moving fewer bytes, not from a faster kernel. +without it (3374 vs 3353 us), so it is entirely data-movement bound. + +This operator is **not**. An earlier revision of this section read the small +gap between the nulled floor and the full time (1692 of 1741 us) as proof that +`tile_n=64` was data-movement bound, and concluded that "further gains come +from moving fewer bytes, not from a faster kernel". That was backwards. The +floor sat just *above* the mmul, hiding it; the operator was compute bound the +whole time, and every attempt to move fewer bytes duly measured as worth +nothing. Re-rolling the mmul's inner loop cut it 23% (2446 -> 1875 cycles per +call, HW trace) and only then did keeping B resident pay -- together 1741 -> +1434 us. + +Two lessons worth keeping. When total latency is `max(compute, DMA)`, testing +levers **one at a time scores both as zero**: residency alone gained nothing +while the mmul was the wall, and the re-rolled mmul alone gained little while +the non-resident floor was. And a nulled-mmul floor close to the full time does +not by itself mean DMA-bound -- it equally means compute is hiding just +underneath. Compare against the *nulled floor*, not the full time, when judging +a data-movement change. Its transfers are cheaper mostly because B arrives pre-packed: the contiguous run per transfer is 128 KB for B and 1 KB for A, against 128 bytes on every diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 6b1c316af..cc19cf004 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -280,11 +280,28 @@ def flm_gemm( # mis-computes. # # Gated on the buffer fitting DOUBLE-buffered, so the next column-block - # still prefetches; A takes 128 KB and C 64 KB of the 512 KB memtile. - # NOTE this admits only k_iters <= 2, i.e. K <= 1024 at tile_n=64. Larger K - # silently falls back to non-resident, so check this gate before believing - # any measurement that claims to be testing residency. - b_resident = (k_iters * mt_b_bytes * 2) <= (512 - 128 - 64) * 1024 + # still prefetches. B_DEPTH is load-bearing, not a safety margin: + # single-buffering would let the lowering use the DMA's native repeat field + # instead of a duplicated BD chain, but it also stops the next + # column-block's B prefetching behind this one's replay, and that costs + # more than it saves -- 1720 us against 1600 at M=2048 K=1024 N=4096, + # i.e. worse than not being resident at all. + # + # The budget counts only C's 64 KB, not A's 128 KB, so it admits k_iters + # <= 3 (K <= 1536 at tile_n=64) rather than 2. That deliberately overcommits + # the A-carrying memtiles by 64 KB and relies on aie-objectfifo-allocate + # spilling one buffer to the least-loaded adjacent memtile, which packs all + # eight to exactly 512 KB (C_L2L3_6 lands on mem_tile_7_1). There is ZERO + # slack: re-verify placement after any change to the A, B or C buffer + # sizes, or the build will fail address assignment rather than silently + # mis-run. + # + # Residency only pays once compute is off the critical path -- it was + # measured latency-neutral while the mmul was the wall, and worth 260 us + # immediately after the mmul loop was re-rolled. Larger K still falls back + # to non-resident, so check this gate before believing any measurement that + # claims to be testing residency. + b_resident = (k_iters * mt_b_bytes * 2) <= (512 - 64) * 1024 if b_resident: mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE,), bf16_ty] b_recv_dims = [(k_iters * (N_TILE // T), K_TILE * T)] + b_recv_dims[1:] From acafdbae1f92ad7db298c6983b2c65517f38d0f8 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 11:47:01 -0600 Subject: [PATCH 17/31] flm_gemm: asymmetric tile buffering and a linear B layout, for 1.74x Takes M=1024 K=1536 N=6144 from 1435 us to 1252 us against the shipped overlay's 2175 us -- 1.52x to 1.74x -- at unchanged arithmetic (err/mass 2.41e-04). Three changes that only pay together. Asymmetric tile buffering decouples the A tile's height from the accumulator's. A is dead as soon as it is consumed while C must live across the whole K reduction, so sizing both to M_TILE pays the peak L1 cost twice. Giving A 16 rows against the accumulator's 64 (rho = 4) frees enough L1 for a 128-deep k slice, which halves the accumulator bytes per mac and doubles the inner loop's trip count: 3.67 -> 2.64 cycles per 8x8x8 mac by hardware trace. Idea from "Can Asymmetric Tile Buffering Be Beneficial?" (arXiv:2511.16041) and mlir-aie's gemm_asymmetric_tile_buffering examples, though those spend the freed L1 on a larger C tile, which an f32 accumulator cannot afford. That k slice was not reachable before. At CT_MAX_K=128 B's innermost contiguous run is 1024, one over the buffer descriptor's 10-bit size field, so it needs a second dimension to encode -- and residency already spends one walking k. Five dimensions against the hardware's four. pack_B now emits B in the order the cores consume it rather than an intermediate blocked order, so both B hops are linear and neither spends a dimension. On its own that is worth nothing; it is what makes the rest fit. tile_ma is chosen by fitting the L1 working set rather than fixed, so tile_n=128 (whose accumulator is already 32 KB) stays symmetric. Note the resolved tile_ma is in the operator name as well as the kernel object name. It changes both the emitted MLIR and the kernel, and the build cache is keyed on filename -- a dir holding another value's artifacts silently produced NaNs before the name encoded it. 39/39 green. --- aie_kernels/aie2p/flm_gemm.cc | 25 ++++++-- aie_kernels/aie2p/flm_gemm_geometry.h | 2 +- iron/operators/flm_gemm/README.md | 21 ++++++- iron/operators/flm_gemm/design.py | 91 +++++++++++++++++++++------ iron/operators/flm_gemm/op.py | 48 ++++++++++++-- 5 files changed, 155 insertions(+), 32 deletions(-) diff --git a/aie_kernels/aie2p/flm_gemm.cc b/aie_kernels/aie2p/flm_gemm.cc index 11319f0a1..5f3264e24 100644 --- a/aie_kernels/aie2p/flm_gemm.cc +++ b/aie_kernels/aie2p/flm_gemm.cc @@ -27,6 +27,16 @@ namespace { constexpr int M = FLM_GEMM_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 and C must live across the +// whole K reduction, so sizing both to M pays the peak cost twice. Defaults +// to M, which is the symmetric design. +#ifndef FLM_GEMM_TILE_MA +#define FLM_GEMM_TILE_MA FLM_GEMM_TILE_M +#endif +constexpr int MA = FLM_GEMM_TILE_MA; +static_assert(M % MA == 0, "tile_m must be a whole number of A bands"); constexpr int K = FLM_GEMM_TILE_K; constexpr int N = FLM_GEMM_TILE_N; constexpr int R = 8; // register tiling r/s/t @@ -34,10 +44,14 @@ constexpr int S = 8; constexpr int T = 8; // How much of K one compute tile holds at a time, given the n width. +#ifdef FLM_GEMM_CT_K +constexpr int CT_K = FLM_GEMM_CT_K; +#else constexpr int CT_K = compute_CT_k_max_n(); +#endif static_assert(CT_K > 0, "no K-blocking geometry for this tile_n"); -static_assert(M % (2 * R) == 0, "tile_m must be a multiple of 2*r (2x2 mmul)"); +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"); @@ -75,11 +89,14 @@ void flm_gemm_acc_init(float *y_acc) { // The l loop lives in the core body so that each B chunk gets its own acquire // point. A is a single object spanning every z slice of the mmul, so this takes // no locks -- the A and B fifos own that handshake. -void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc) { +void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc, + int32_t band) { ::aie::set_rounding(round_mode); constexpr int NUM_ITER = K / CT_K; - flm_gemm_mmul_2x2( - a_buf, b_buf, y_acc); + a_buf, b_buf, y_acc + band * (MA * N)); } } diff --git a/aie_kernels/aie2p/flm_gemm_geometry.h b/aie_kernels/aie2p/flm_gemm_geometry.h index fcb4e69d9..3b0b46e21 100644 --- a/aie_kernels/aie2p/flm_gemm_geometry.h +++ b/aie_kernels/aie2p/flm_gemm_geometry.h @@ -15,7 +15,7 @@ constexpr int CT_k_max_n_16 = 16; constexpr int CT_k_max_n_32 = 32; -constexpr int CT_k_max_n_64 = 64; +constexpr int CT_k_max_n_64 = 128; constexpr int CT_k_max_n_128 = 32; constexpr int CT_k_max_n_256 = 16; diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index 390c2776b..fe7373d92 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -150,13 +150,28 @@ M=1024 K=1536 N=6144, min of per-run medians across separate processes: | | bytes moved | latency | DMA-only (compute nulled) | |---|---|---|---| -| `FLMGEMM` (`tile_n=64`) | 69 MB | **1434 us** | 1231 us | +| `FLMGEMM` (`tile_n=64`) | 69 MB | **1252 us** | 1231 us | | shipped `mm.xclbin` | 107 MB | 2175 us | -- | | `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | -**1.52x the shipped overlay**, at identical arithmetic (err/mass 2.41e-04 +**1.74x the shipped overlay**, at identical arithmetic (err/mass 2.41e-04 either way). The 69 MB is with B resident in the memtile; the 126 MB figure -this table used to quote was the non-resident fallback. +this table used to quote was the non-resident fallback. At 1252 us against a +1231 us data-movement floor, this operator is now essentially DMA-bound: the +mmul is finally cheap enough to hide, so further gains have to come from +moving fewer bytes. + +Three changes compound to get there, and none of them works alone: + +* **The mmul's inner loop is rolled**, not hand-unrolled. 2446 -> 1875 cycles + per call. +* **`pack_B` emits the final consumption order**, so both B hops are linear + descriptors instead of blocked ones. Worth nothing by itself -- it is what + frees the descriptor dimensions the other two need. +* **Asymmetric tile buffering.** The A tile is 16 rows while the accumulator + is 64 (`rho = 4`), which pays for a 128-deep k slice. That halves the + accumulator traffic per mac and doubles the inner loop's trip count: + 3.67 -> 2.64 cycles per 8x8x8 mac. **Measure this carefully.** Dispatch latency on this part is *bimodal*, with modes about 6% apart, and both show up for every configuration. A batch that diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index cc19cf004..a4cf2d268 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -52,7 +52,7 @@ # loses. N_TILE_DEFAULT = 64 # k-slice per n width; must match compute_CT_k_max_n in flm_gemm_geometry.h -CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 64, 128: 32, 256: 16} +CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32, 256: 16} R, S, T = 8, 8, 8 ROWS, COLS = 4, 8 @@ -67,6 +67,8 @@ A_DEPTH = 2 STACK_SIZE = 4096 +# Usable L1 per compute tile: 64 KB less the stack and a little slack. +L1_BUDGET = 60 * 1024 EPILOGUE_MODES = {"none": 0, "gelu": 1, "silu": 2, "sigmoid": 3} # The epilogue entry point, shared by the design and op.py (which needs it @@ -79,6 +81,28 @@ MIN_K = K_TILE # 512 +def _default_tile_ma(n_tile, ct_max_k): + """Largest A-tile height whose L1 working set fits. + + A is dead 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 + (ct_max_k=128 at n=64), which measured 3.67 -> 2.64 cycles per 8x8x8 mac. + """ + acc = M_TILE * n_tile * 4 + b = ct_max_k * n_tile * 2 * B_DEPTH + cout = CT_OUT_LEN * 2 * C_DEPTH + for t_ma in (M_TILE, M_TILE // 2, M_TILE // 4, M_TILE // 8): + if t_ma < 2 * R: + break + a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH + if acc + a + b + cout <= L1_BUDGET: + return t_ma + raise ValueError( + f"no A-tile height fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}" + ) + + def flm_gemm( dev, M, @@ -86,6 +110,7 @@ def flm_gemm( N, epilogue="none", tile_n=N_TILE_DEFAULT, + tile_ma=None, kernel_object="flm_gemm.o", epilogue_object="flm_gemm_epilogue.o", trace_size=0, @@ -93,8 +118,9 @@ def flm_gemm( """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 -- the block-major reordering B needs on - the way in is done by the fill descriptor, not by the caller. + bf16 and all plain dense tensors, except that B must arrive pre-packed by + ``FLMGEMM.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( @@ -102,9 +128,19 @@ def flm_gemm( ) N_TILE = tile_n CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + # 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. + T_MA = _default_tile_ma(N_TILE, CT_MAX_K) if tile_ma is None else 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}" + ) + RHO = M_TILE // T_MA 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 * (M_TILE // R // 2) # every z slice of one mmul + 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 @@ -167,7 +203,9 @@ def flm_gemm( acc_init = Kernel("flm_gemm_acc_init", kernel_object, [ct_acc_ty]) k_step = Kernel( - "flm_gemm_k_step", kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty] + "flm_gemm_k_step", + kernel_object, + [ct_a_obj_ty, ct_b_ty, ct_acc_ty, np.int32], ) epilogue_chunk = Kernel( EPILOGUE_SYMBOL, @@ -182,24 +220,30 @@ def flm_gemm( # layout the mmul indexes, and they are tightly coupled to it. A mismatch # here produces silently wrong results, not a build error. + def _split_run(run): + # A BD wrap may not exceed 1023; the run is contiguous, so a longer one + # re-encodes as two dimensions at the cost of one of the four. + return [(run, 1)] if run <= 1023 else [(2, run // 2), (run // 2, 1)] + # 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_recv_dims = [(N_TILE // T, K_TILE * T), (T, S), (K_TILE // S, S * T), (S, 1)] - b_send_dims = [ - (K_DIV_CT_K_MAX, T * CT_MAX_K), - (N_TILE // T, K_TILE * T), - (T * CT_MAX_K, 1), - ] + # 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), - (R * CT_MAX_K, 1), - ] + ] + _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 @@ -303,9 +347,12 @@ def flm_gemm( # claims to be testing residency. b_resident = (k_iters * mt_b_bytes * 2) <= (512 - 64) * 1024 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,), bf16_ty] - b_recv_dims = [(k_iters * (N_TILE // T), K_TILE * T)] + b_recv_dims[1:] - b_send_dims = [(k_iters, K_TILE * N_TILE)] + b_send_dims b_l3l2_fifos = [] b_cons = {} @@ -349,10 +396,13 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): # 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) - a = a_h.acquire(1) - kstep_k(a, b, acc) - a_h.release(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 @@ -373,8 +423,9 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): for _ in range_(k_iters): for _ in range_(B_ITERS // B_DEPTH): for _ in range(B_DEPTH): - a_h.acquire(1) - a_h.release(1) + for _ in range(RHO): + a_h.acquire(1) + a_h.release(1) return core_fn diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 17e5e6794..14ab24212 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -51,6 +51,9 @@ class FLMGEMM(MLIROperator): # halves A fetches instead and wins only when small K makes the operator # DMA-bound. See README.md. tile_n: int | None = field(default=None, repr=False) + # A-tile rows, decoupled from the accumulator's M_TILE (asymmetric tile + # buffering). None means symmetric (T_MA == M_TILE). + tile_ma: int | None = field(default=None, repr=False) # "conv_even" (round to nearest even) or "floor" (truncate). The core # powers up in floor, and the design this was ported from never sets the # mode, so "floor" reproduces its arithmetic exactly -- at ~40x the error, @@ -118,6 +121,11 @@ def name(self) -> str: base = f"{base}_{self.rounding}" if self.tile_n != self._default_tile_n(self.K): base = f"{base}_tn{self.tile_n}" + # The RESOLVED height, not just an explicit override: it changes the + # emitted MLIR and the kernel object, so a build dir holding another + # value's artifacts must not satisfy this one. + if self._tile_ma != M_TILE: + base = f"{base}_ma{self._tile_ma}" return base @property @@ -163,10 +171,24 @@ def _epilogue_flags(self) -> list[str]: ] return flags + self._rounding_flags + @property + def _tile_ma(self) -> int: + """Resolved A-tile height. design.py picks the default, and it MUST be + the same value the kernel is compiled with -- the design sizes the A + object from it while the kernel derives the mmul's rowA from it, so a + mismatch reads past the buffer and produces garbage rather than a build + error.""" + from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N, _default_tile_ma + + if self.tile_ma is not None: + return self.tile_ma + return _default_tile_ma(self.tile_n, CT_MAX_K_FOR_N[self.tile_n]) + @property def _kernel_object(self) -> str: rnd = "" if self.rounding == "conv_even" else f"_{self.rounding}" - return f"flm_gemm_{M_TILE}x{K_TILE}x{self.tile_n}{rnd}.o" + ma = f"_ma{self._tile_ma}" + return f"flm_gemm_{M_TILE}x{K_TILE}x{self.tile_n}{rnd}{ma}.o" def get_mlir_artifact(self): return PythonGeneratedMLIRArtifact( @@ -181,6 +203,7 @@ def get_mlir_artifact(self): "K": self.K, "N": self.N, "tile_n": self.tile_n, + "tile_ma": self._tile_ma, "epilogue": self.epilogue, "kernel_object": self._kernel_object, "epilogue_object": self._epilogue_artifact, @@ -208,6 +231,7 @@ def get_kernel_artifacts(self): f"-DFLM_GEMM_TILE_M={M_TILE}", f"-DFLM_GEMM_TILE_K={K_TILE}", f"-DFLM_GEMM_TILE_N={self.tile_n}", + f"-DFLM_GEMM_TILE_MA={self._tile_ma}", # The r=8 mmul shape this design uses only exists on the # bfp16-emulated path; without this the kernel will not # compile. @@ -248,9 +272,25 @@ def pack_B(self, B): raise ValueError( f"B ({K}, {N}) must tile to ({K_TILE}, {N_TILE}) to be packed" ) - t = B.reshape(K // K_TILE, K_TILE // S, S, N // N_TILE, N_TILE // T, T) - # (kb, kb8, s_in, cb, tb, t_in) -> (cb, kb, tb, s_in, kb8, t_in) - return t.permute(3, 0, 4, 2, 1, 5).reshape(-1).contiguous() + # Emit the FINAL consumption order, not an intermediate one. The old + # layout left a 4-dimension scatter for the memtile's dims_from_stream + # to finish, which put the n-block index outside the k-slice index and + # so cost two descriptor dimensions on the way back out. Packing all + # the way here makes both B hops linear, which is what leaves room for + # a k-slice deep enough to halve the accumulator traffic (CT_MAX_K=128) + # while B is also memtile-resident. + from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N + + CT_K = CT_MAX_K_FOR_N[N_TILE] + col_a = CT_K // S + t = 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) + # -> (cb, kb, kslice, tb, i, s_in, t_in) + # The s x t block is s-major: the mmul is instantiated with + # is_b_s_t_in_row_major=true, so it loads the block straight from L1. + return t.permute(4, 0, 1, 5, 2, 3, 6).reshape(-1).contiguous() def get_arg_spec(self): return [ From 3ff319f8fde95bd009e3f989ebf0f634525aedd7 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 12:30:20 -0600 Subject: [PATCH 18/31] flm_gemm: store B in bfp16, for 1.90x -- with a known accuracy regression Takes M=1024 K=1536 N=6144 from 1436 us to 1143 us against the shipped overlay's 2175 us -- 1.74x to 1.90x. DDR drops 69 -> 47 MB because bfp16ebs8 packs 8 values in 9 bytes where bf16 needs 16, and B is the leg that residency keeps but cannot shrink. DO NOT SHIP AS IS. Accuracy moves from 2.4133e-04 to 2.6923e-04, about 12% worse, and 81% of output elements change. That contradicts the premise this was built on -- that quantizing B on the host merely hoists a rounding the mmul already performs on every mac call, and so is free. It is not free, and the residual is not yet explained. For scale: 2.69e-04 is still 37x better than the shipped overlay's 9.87e-03. But it breaks an exact tie -- flm_gemm at 2.4133e-04 matched IRON GEMM's emulated mode to five digits, and GEMM's accurate mode (emulate_bf16_mmul_with_bfp16=False) reaches 4.1321e-05, which this moves further away from. What has been ruled out as the cause, each by measurement: * Block grouping. The shared exponent covers 8 consecutive k for one n; grouping over n instead gives 1.95e-02. * Rounding. Truncation is correct; round-to-nearest gives 6.18e-03. * The A widening path. Using mul_elem_64 for one operand and direct assignment for the other, as mm_bfp_mixed.cc does, is byte for byte identical to using direct assignment for both. * The packer itself. It is byte-identical to mlir-aie's reference floatToBfp16, and single-value probes decode exactly on hardware across magnitude, sign and shared-exponent crushing. * Layout and accumulation. Single-element (k,n) probes map correctly, and partial sums over k are exact to 512 terms. Also fixes a bug found on the way in: the bfp16 mmul never advanced pC1 and pC2 across the j loop, so every j iteration accumulated into the same C slot. Every single-element probe passed anyway because they all used columns that j=0 owns -- the error only showed as a period-64 column pattern in full data. 39/39 green (the test budget is 4e-3, well above both figures). --- aie_kernels/aie2p/flm_gemm.cc | 9 ++++ aie_kernels/aie2p/flm_gemm_mmul.h | 87 +++++++++++++++++++++++++++++++ iron/operators/flm_gemm/README.md | 8 +-- iron/operators/flm_gemm/design.py | 38 +++++++++----- iron/operators/flm_gemm/op.py | 68 +++++++++++++++++++++--- 5 files changed, 189 insertions(+), 21 deletions(-) diff --git a/aie_kernels/aie2p/flm_gemm.cc b/aie_kernels/aie2p/flm_gemm.cc index 5f3264e24..7745f2bf9 100644 --- a/aie_kernels/aie2p/flm_gemm.cc +++ b/aie_kernels/aie2p/flm_gemm.cc @@ -89,14 +89,23 @@ void flm_gemm_acc_init(float *y_acc) { // The l loop lives in the core body so that each B chunk gets its own acquire // point. A is a single object spanning every z slice of the mmul, so this takes // no locks -- the A and B fifos own that handshake. +#ifdef FLM_GEMM_BFP16_B +void flm_gemm_k_step(bfloat16 *a_buf, bfp16ebs8 *b_buf, float *y_acc, +#else void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc, +#endif int32_t band) { ::aie::set_rounding(round_mode); constexpr int NUM_ITER = K / CT_K; // 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. +#ifdef FLM_GEMM_BFP16_B + flm_gemm_mmul_2x2_bfpb(a_buf, b_buf, y_acc + band * (MA * N)); +#else flm_gemm_mmul_2x2( a_buf, b_buf, y_acc + band * (MA * N)); +#endif } } diff --git a/aie_kernels/aie2p/flm_gemm_mmul.h b/aie_kernels/aie2p/flm_gemm_mmul.h index d4dedf978..44d7c4fa5 100644 --- a/aie_kernels/aie2p/flm_gemm_mmul.h +++ b/aie_kernels/aie2p/flm_gemm_mmul.h @@ -120,4 +120,91 @@ flm_gemm_mmul_2x2(const T_in *__restrict pA, const T_in *__restrict pB, event1(); } +// The same 2x2 mmul, but B arrives ALREADY in bfp16ebs8 rather than bf16. +// +// The bf16 form converts B inside every mac -- transpose, widen, then +// to_v64bfp16ebs8 -- purely to feed hardware that only multiplies bfp16. 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, it +// does not add one. It also makes B 9 bytes per 8 elements instead of 16, +// which is why it is worth doing at all -- the operator is DMA-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, not +// blocks (llvm-aie#1232). The stream sidesteps both. +template +__aie_inline void flm_gemm_mmul_2x2_bfpb(const bfloat16 *__restrict pA, + const bfp16ebs8 *__restrict pB, + T_out *__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) { + T_out *__restrict pC1 = pC + (z * colB) * sizeC; + T_out *__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; + + // Rolled for the same reason as the bf16 form: extra live state across + // this loop loses more than it gains (llvm-aie#1066). + 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, exactly as + // mlir-aie's mm_bfp_mixed.cc 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 codegens. + 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); + } + aie::store_v(pC1, C00.template to_vector()); + pC1 += sizeC; + aie::store_v(pC1, C01.template to_vector()); + pC1 += sizeC; + aie::store_v(pC2, C10.template to_vector()); + pC2 += sizeC; + aie::store_v(pC2, C11.template to_vector()); + pC2 += sizeC; + } + } + event1(); +} + #endif // __FLM_GEMM_MMUL_H__ diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index fe7373d92..a322a0fcb 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -150,12 +150,14 @@ M=1024 K=1536 N=6144, min of per-run medians across separate processes: | | bytes moved | latency | DMA-only (compute nulled) | |---|---|---|---| -| `FLMGEMM` (`tile_n=64`) | 69 MB | **1252 us** | 1231 us | +| `FLMGEMM` (`tile_n=64`) | 47 MB | **1143 us** | -- | | shipped `mm.xclbin` | 107 MB | 2175 us | -- | | `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | -**1.74x the shipped overlay**, at identical arithmetic (err/mass 2.41e-04 -either way). The 69 MB is with B resident in the memtile; the 126 MB figure +**1.90x the shipped overlay.** Accuracy is err/mass 2.69e-04 against the +shipped overlay's 9.87e-03 -- 37x better -- but see the note below: it is +2.69e-04 rather than the 2.41e-04 this operator reached before B was stored +in bfp16, and that gap is a known open bug, not a fundamental cost. The 69 MB is with B resident in the memtile; the 126 MB figure this table used to quote was the non-resident fallback. At 1252 us against a 1231 us data-movement floor, this operator is now essentially DMA-bound: the mmul is finally cheap enough to hide, so further gains have to come from diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index a4cf2d268..8b1c625e6 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -27,6 +27,8 @@ 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, @@ -81,6 +83,12 @@ MIN_K = K_TILE # 512 +def _bfp16_bytes(elems): + """bfp16ebs8 packs 8 values as 8 mantissa bytes plus one shared exponent.""" + assert elems % 8 == 0 + return elems // 8 * 9 + + def _default_tile_ma(n_tile, ct_max_k): """Largest A-tile height whose L1 working set fits. @@ -90,9 +98,13 @@ def _default_tile_ma(n_tile, ct_max_k): (ct_max_k=128 at n=64), which measured 3.67 -> 2.64 cycles per 8x8x8 mac. """ acc = M_TILE * n_tile * 4 - b = ct_max_k * n_tile * 2 * B_DEPTH + b = _bfp16_bytes(ct_max_k * n_tile) * B_DEPTH cout = CT_OUT_LEN * 2 * C_DEPTH - for t_ma in (M_TILE, M_TILE // 2, M_TILE // 4, M_TILE // 8): + # Not below 32: at t_ma = 16 the mmul's z loop has a single trip and + # Peano's AIE2P backend aborts with "Use not jointly dominated by defs" + # on the bfp16 path. 32 also halves the A object count, which is fifo + # overhead the asymmetry would otherwise add. + for t_ma in (M_TILE, M_TILE // 2): if t_ma < 2 * R: break a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH @@ -188,17 +200,17 @@ def flm_gemm( # 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,), bf16_ty] + ct_b_ty = np.ndarray[(CT_MAX_K * N_TILE // 8,), np.dtype[v8bfp16ebs8]] 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_b_ty = np.ndarray[(K_TILE * N_TILE,), bf16_ty] - mt_b_bytes = K_TILE * N_TILE * 2 + mt_b_ty = np.ndarray[(K_TILE * N_TILE // 8,), np.dtype[v8bfp16ebs8]] + mt_b_bytes = _bfp16_bytes(K_TILE * N_TILE) mt_out_ty = np.ndarray[(C_SLICE_LEN * ROWS,), bf16_ty] # 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,), bf16_ty] + b_l3_ty = np.ndarray[(K * N // 8,), np.dtype[v8bfp16ebs8]] c_l3_ty = np.ndarray[(M * N,), bf16_ty] acc_init = Kernel("flm_gemm_acc_init", kernel_object, [ct_acc_ty]) @@ -352,7 +364,7 @@ def _split_run(run): # 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,), bf16_ty] + mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE // 8,), np.dtype[v8bfp16ebs8]] b_l3l2_fifos = [] b_cons = {} @@ -488,14 +500,16 @@ def b_tap(mega_col, c): # 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,), - offset=(mega_col * COLS + c) * N_TILE * K, + tensor_dims=(K * N // 8,), + offset=(mega_col * COLS + c) * N_TILE * K // 8, sizes=( - [1, 1, 1, k_iters * K_TILE * N_TILE] + [1, 1, 1, k_iters * K_TILE * N_TILE // 8] if b_resident - else [m_row_blocks, k_iters, 1, K_TILE * N_TILE] + else [m_row_blocks, k_iters, 1, K_TILE * N_TILE // 8] + ), + strides=( + [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // 8, 0, 1] ), - strides=([0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE, 0, 1]), ) def c_tap(mega_col, c): diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 14ab24212..21f175297 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -2,6 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field + +import numpy as np +import torch from typing import ClassVar, Dict from iron.common import ( @@ -30,6 +33,39 @@ ) +def _f32_to_bfp16ebs8(a): + """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, two's-complemented + when negative, truncated to 8 bits and then arithmetically right-shifted by + (maxExp - exp). AIE2P always TRUNCATES here, so this is independent of + ``rounding``. Verified byte-identical against mlir-aie's reference + implementation (``programming_examples/ml/block_datatypes/helper.h``). + + 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) + mag = np.where(sign, ~man.astype(np.int64) + 1, 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. + # TRUNCATION, not rounding -- that is what AIE2P does, and round-to-nearest + # here measures 6.18e-03 against truncation's 2.69e-04. + shift = (max_exp - exp).astype(np.int64) + total = 17 + np.clip(shift, 0, 40) + v8 = np.where(shift >= 32, np.where(sign, -1, 0), mag >> np.clip(total, 0, 62)) + 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)) + + @dataclass class FLMGEMM(MLIROperator): """AIE-accelerated bf16 GEMM on a fixed 4x8 grid, with a fused epilogue. @@ -232,6 +268,7 @@ def get_kernel_artifacts(self): f"-DFLM_GEMM_TILE_K={K_TILE}", f"-DFLM_GEMM_TILE_N={self.tile_n}", f"-DFLM_GEMM_TILE_MA={self._tile_ma}", + "-DFLM_GEMM_BFP16_B", # The r=8 mmul shape this design uses only exists on the # bfp16-emulated path; without this the kernel will not # compile. @@ -250,8 +287,15 @@ def get_kernel_artifacts(self): return artifacts def pack_B(self, B): - """Reorder a row-major ``(K, N)`` weight matrix into the layout the B - fill expects. Returns a flat tensor. + """Reorder and quantize a row-major ``(K, N)`` weight matrix into the + layout the B fill expects. Returns a flat uint8 tensor of bfp16ebs8 + blocks, NOT a bf16 tensor. + + The quantization is not a loss this adds. The mmul only multiplies + bfp16, so the bf16 path converts B inside every mac call; doing it here + hoists a rounding that already happened and leaves the arithmetic + bit-identical. It also makes B 9 bytes per 8 values instead of 16, + which is the point -- this operator is data-movement bound. Each ``K_TILE x N_TILE`` tile is emitted in t-block-major order -- the odometer ``(n//T, k%S, k//S, n%T)``, outermost first -- with tiles @@ -287,10 +331,22 @@ def pack_B(self, B): 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) - # -> (cb, kb, kslice, tb, i, s_in, t_in) - # The s x t block is s-major: the mmul is instantiated with - # is_b_s_t_in_row_major=true, so it loads the block straight from L1. - return t.permute(4, 0, 1, 5, 2, 3, 6).reshape(-1).contiguous() + # -> (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 below match the kernel's. + # Grouping the shared exponent over 8 consecutive k (for one n) is + # verified: grouping over n instead measures 1.95e-02 against this + # layout's 2.69e-04. + t = t.permute(4, 0, 1, 5, 2, 6, 3).reshape(-1, 8).contiguous() + return _f32_to_bfp16ebs8(t.float().numpy()) + + @staticmethod + def unpack_B_size(K, N): + """Bytes ``pack_B`` returns for a ``(K, N)`` weight matrix.""" + return K * N // 8 * 9 def get_arg_spec(self): return [ From a963dc3b884ebb787b38a1d7aff222ea448a4485 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 12:38:50 -0600 Subject: [PATCH 19/31] flm_gemm: match the core's rounding mode when packing B Removes the 12% accuracy regression from the previous commit. B in bfp16 is now numerically free, as it was supposed to be: 2.3912e-04 against the bf16 build's 2.4133e-04, i.e. very slightly BETTER, and better on every shape measured (2.900 vs 2.93e-04, 2.068 vs 2.09e-04, 4.173 vs 4.20e-04, 1.846 vs 1.86e-04). 1.90x the shipped overlay, 41x more accurate than it. Two bugs, both in the host-side quantizer. The conversion obeys the core's ROUNDING MODE. mlir-aie's reference floatToBfp16 hardcodes truncation and comments that AIE2P always truncates; that is true only of the power-up floor mode. flm_gemm calls set_rounding(conv_even), so the kernel's own conversion rounds to nearest with ties to even. Reading one block back off the hardware a slot at a time shows it plainly: 14.9375 -> 15 while 106.5 -> 106 and 94.5 -> 94. Truncation cannot produce that. pack_B now follows self.rounding, so rounding="floor" still truncates. Rounding then overflows the mantissa. A block's largest magnitude sits at 127 before rounding and can carry to 128, which does not fit the signed 8-bit field and wrapped to -128. B is drawn from rand() in the tests, so values sit near their block maximum often, and the wrap alone cost 6.18e-03. Saturate instead. Found by probing a single 8-value block: A[0,ki]=1 selects slot ki so C[0,0] is exactly that slot's decoded value, which makes the kernel's conversion directly readable and comparable against the packer's. Worth keeping in mind -- the earlier single-value probes all passed because a lone nonzero owns its block's exponent and never exercises the shift. 39/39 green. --- iron/operators/flm_gemm/README.md | 9 +++---- iron/operators/flm_gemm/op.py | 39 +++++++++++++++++++++++-------- 2 files changed, 34 insertions(+), 14 deletions(-) diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index a322a0fcb..bc91568fc 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -154,10 +154,11 @@ M=1024 K=1536 N=6144, min of per-run medians across separate processes: | shipped `mm.xclbin` | 107 MB | 2175 us | -- | | `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | -**1.90x the shipped overlay.** Accuracy is err/mass 2.69e-04 against the -shipped overlay's 9.87e-03 -- 37x better -- but see the note below: it is -2.69e-04 rather than the 2.41e-04 this operator reached before B was stored -in bfp16, and that gap is a known open bug, not a fundamental cost. The 69 MB is with B resident in the memtile; the 126 MB figure +**1.90x the shipped overlay**, at err/mass 2.39e-04 against its 9.87e-03 -- +41x more accurate. Storing B in bfp16 is numerically free: the mmul only +multiplies bfp16, so quantizing on the host hoists a rounding that already +happened on every mac call. It has to reproduce the core's rounding MODE to +do so -- see ``pack_B``. The 69 MB is with B resident in the memtile; the 126 MB figure this table used to quote was the non-resident fallback. At 1252 us against a 1231 us data-movement floor, this operator is now essentially DMA-bound: the mmul is finally cheap enough to hide, so further gains have to come from diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 21f175297..7deca1e6c 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -33,15 +33,21 @@ ) -def _f32_to_bfp16ebs8(a): +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, two's-complemented - when negative, truncated to 8 bits and then arithmetically right-shifted by - (maxExp - exp). AIE2P always TRUNCATES here, so this is independent of - ``rounding``. Verified byte-identical against mlir-aie's reference - implementation (``programming_examples/ml/block_datatypes/helper.h``). + 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. flm_gemm calls ``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. """ @@ -52,14 +58,25 @@ def _f32_to_bfp16ebs8(a): 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) - mag = np.where(sign, ~man.astype(np.int64) + 1, man.astype(np.int64)) + # 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. # TRUNCATION, not rounding -- that is what AIE2P does, and round-to-nearest # here measures 6.18e-03 against truncation's 2.69e-04. shift = (max_exp - exp).astype(np.int64) - total = 17 + np.clip(shift, 0, 40) - v8 = np.where(shift >= 32, np.where(sign, -1, 0), mag >> np.clip(total, 0, 62)) + 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) @@ -341,7 +358,9 @@ def pack_B(self, B): # verified: grouping over n instead measures 1.95e-02 against this # layout's 2.69e-04. t = t.permute(4, 0, 1, 5, 2, 6, 3).reshape(-1, 8).contiguous() - return _f32_to_bfp16ebs8(t.float().numpy()) + return _f32_to_bfp16ebs8( + t.float().numpy(), round_conv_even=self.rounding == "conv_even" + ) @staticmethod def unpack_B_size(K, N): From 562abc57612d20d66a16fceeb267839b6c83793a Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 12:44:39 -0600 Subject: [PATCH 20/31] flm_gemm: declare B's runtime arg in bytes, not bf16 elements pack_B returns bfp16ebs8 -- 9 bytes per 8 values -- but the arg spec still described B as (K, N) bf16. #179 made buffer sizing follow spec.dtype (prod(shape) * dtype.itemsize in calculate_buffer_layout), so that declaration over-allocates this operator's LARGEST buffer by 1.78x: at K=1536 N=6144 it asks for 18.9 MB where pack_B produces 10.6 MB. Harmless before the merge because the spec's dtype was not consulted; worth fixing now rather than leaving a caller-visible size lie in place. 39/39 green. --- iron/operators/flm_gemm/op.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 7deca1e6c..4f295ff42 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -370,8 +370,13 @@ def unpack_B_size(K, N): 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) + # B arrives pre-packed AND quantized by pack_B: bfp16ebs8, which is + # 9 bytes per 8 values rather than bf16's 16. Declared in bytes so + # the buffer is sized from what pack_B actually returns -- a + # (K, N) bf16 spec would over-allocate the largest buffer by 1.78x. + AIERuntimeArgSpec( + "in", (self.unpack_B_size(self.K, self.N),), dtype=np.uint8 + ), # B (weights) AIERuntimeArgSpec("out", (self.M, self.N)), # C ] From a07e8505a370b0658ef58cab2c8925089bdc855f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 13:01:49 -0600 Subject: [PATCH 21/31] flm_gemm: ablation knobs that locate the remaining DMA cost All three are env-gated and inert by default; 15/15 green with nothing set. They exist because the operator is now DMA-rate bound rather than byte bound, and that is not visible from wall clock alone. FLM_NULL_MMUL=1 skip the multiply, keep every DMA and handshake FLM_C_LINEAR=1 drain C as one contiguous run per column-block FLM_OVERLAP=N column-blocks in flight (default 2, unchanged) What they establish, at M=1024 K=1536 N=6144: Compute is entirely hidden. Nulling the mmul does not make the operator faster (1125 us full against 1152 nulled), so the ~14% of compute still above its ResMII bound is worth nothing until data movement improves. The operator moves 60.9 MB in ~1140 us, i.e. 53 GB/s against a measured 63-70 GB/s roof -- 79%. The gap is not per-core ingress: the trace shows the cores' two input DMA ports only 45% and 25% busy, so they are starved rather than saturated. Nor is it pipelining: overlap depth 3 is worth under 1%, and 4 and 5 are worse as shim buffer descriptors run out. It is C's write pattern. C is drained as ROWS*M_TILE runs of N_TILE elements -- 128 bytes -- strided by N, and 128 B sits below the ~256 B cliff where DDR bursts fall apart. Draining the same bytes contiguously saves 114 us on the min and 86 us on the median over 6 interleaved rounds. That is ~8% of the whole operator, and recovering it would put it near 2.1x the shipped overlay. FLM_C_LINEAR writes C to the wrong place -- it is a measurement, not a fix. The fix needs 256 B runs with C still row-major, which means either a 128-wide n tile (but that caps the k slice at 64, making compute the wall again) or joining C across pairs of adjacent columns so one drain covers 128 columns. --- aie_kernels/aie2p/flm_gemm.cc | 6 ++++++ iron/operators/flm_gemm/design.py | 33 +++++++++++++++++++++++++------ iron/operators/flm_gemm/op.py | 12 ++++++++++- 3 files changed, 44 insertions(+), 7 deletions(-) diff --git a/aie_kernels/aie2p/flm_gemm.cc b/aie_kernels/aie2p/flm_gemm.cc index 7745f2bf9..6cd655c7d 100644 --- a/aie_kernels/aie2p/flm_gemm.cc +++ b/aie_kernels/aie2p/flm_gemm.cc @@ -95,6 +95,12 @@ void flm_gemm_k_step(bfloat16 *a_buf, bfp16ebs8 *b_buf, float *y_acc, void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc, #endif int32_t band) { +#ifdef FLM_GEMM_NULL_MMUL + // ABLATION ONLY: skip the multiply, keep every acquire, release and DMA. + // Output is garbage; never correctness-gate a build with this. + (void)a_buf; (void)b_buf; (void)y_acc; (void)band; + return; +#endif ::aie::set_rounding(round_mode); constexpr int NUM_ITER = K / CT_K; // The accumulator is [row-block][col-block][r*t], so band b starts at diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 8b1c625e6..be037e20f 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -123,6 +123,7 @@ def flm_gemm( epilogue="none", tile_n=N_TILE_DEFAULT, tile_ma=None, + overlap=None, kernel_object="flm_gemm.o", epilogue_object="flm_gemm_epilogue.o", trace_size=0, @@ -150,6 +151,9 @@ def flm_gemm( f"tile_ma ({T_MA}) must divide {M_TILE} and be a multiple of {2 * R}" ) RHO = M_TILE // T_MA + import os as _os + + OVERLAP = int(_os.environ.get("FLM_OVERLAP", "2")) 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 @@ -515,6 +519,17 @@ def b_tap(mega_col, c): def c_tap(mega_col, c): # Every joined block this column produces for one column-block: one # ROWS*M_TILE x N_TILE block per row-block. + if _os.environ.get("FLM_C_LINEAR") == "1": + # ABLATION: same byte count, one contiguous run per column-block + # instead of ROWS*M_TILE runs of N_TILE. Writes C to the WRONG + # place -- only for measuring what the scattered write costs. + blk = m_row_blocks * ROWS * M_TILE * N_TILE + return TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * blk, + sizes=[1, 1, 1, blk], + strides=[0, 0, 0, 1], + ) return TensorAccessPattern( tensor_dims=(M * N,), offset=(mega_col * COLS + c) * N_TILE, @@ -545,7 +560,12 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): # block costs 3 buffer descriptors on a shim column (A + B + C), so two # in flight is 6 of 16. Per-object tasks needed 1 + 2*k_iters and could # not be overlapped at all. - prev = None + # Keep OVERLAP column-blocks in flight. A block costs 3 shim buffer + # descriptors on a column (A + B + C) against 16 available, so the + # ceiling is 5; the operator is DDR-rate bound rather than byte bound + # (53 GB/s of a 63-70 GB/s roof), so how deeply the fills are pipelined + # is what decides the rate. + pending = [] for mega_col, active_cols in blocks: tg_c = TaskGroup() for c in range(active_cols): @@ -556,13 +576,14 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): for c in range(active_cols): b_prods[c].fill(B, b_tap(mega_col, c), group=tg_f) - if prev is not None: - for tg in prev: + pending.append([tg_f, tg_c]) + while len(pending) >= OVERLAP: + for tg in pending.pop(0): tg.finish() - prev = [tg_f, tg_c] - for tg in prev or []: - tg.finish() + for group in pending: + for tg in group: + tg.finish() rt = Runtime( sequence, diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index 4f295ff42..b6285e7df 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -224,6 +224,15 @@ def _epilogue_flags(self) -> list[str]: ] return flags + self._rounding_flags + @property + def _ablate_mmul(self) -> bool: + """ABLATION: FLM_NULL_MMUL=1 nulls the multiply, leaving all data + movement. Threaded into the object AND operator names because the + build cache is keyed on filename.""" + import os + + return os.environ.get("FLM_NULL_MMUL", "") == "1" + @property def _tile_ma(self) -> int: """Resolved A-tile height. design.py picks the default, and it MUST be @@ -240,7 +249,7 @@ def _tile_ma(self) -> int: @property def _kernel_object(self) -> str: rnd = "" if self.rounding == "conv_even" else f"_{self.rounding}" - ma = f"_ma{self._tile_ma}" + ma = f"_ma{self._tile_ma}" + ("_nomm" if self._ablate_mmul else "") return f"flm_gemm_{M_TILE}x{K_TILE}x{self.tile_n}{rnd}{ma}.o" def get_mlir_artifact(self): @@ -286,6 +295,7 @@ def get_kernel_artifacts(self): f"-DFLM_GEMM_TILE_N={self.tile_n}", f"-DFLM_GEMM_TILE_MA={self._tile_ma}", "-DFLM_GEMM_BFP16_B", + *(["-DFLM_GEMM_NULL_MMUL"] if self._ablate_mmul else []), # The r=8 mmul shape this design uses only exists on the # bfp16-emulated path; without this the kernel will not # compile. From 3a8f8a34812d5ef7c84a57a3f1e20cce697d007c Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 13:17:07 -0600 Subject: [PATCH 22/31] flm_gemm: count A and C in the resident-B gate The gate compared B's resident footprint against (512 - 64) KB, which ignores A's 128 KB entirely and hardcodes C's size. It happened to be conservative enough at tile_n=64, but at tile_n=128 -- where C doubles to 128 KB and resident B is 432 KB -- it admitted a configuration totalling 688 KB of a 512 KB memtile, which failed address assignment outright rather than falling back to non-resident. Found while measuring whether tile_n=128 could pay for itself now that ATB and bfp16 B changed the L1 and L2 budgets. It cannot: with the gate honest it correctly drops to non-resident, and B's DDR then quadruples to 42.5 MB for a total of 74 MB against tile_n=64's 60.9. Measured 2283 us against 1154. So tile_n=64 stays the default, and the sweep table's preference is confirmed for the current balance even though its absolute numbers are stale. 39/39 green. --- iron/operators/flm_gemm/design.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index be037e20f..826b7379d 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -209,9 +209,11 @@ def flm_gemm( 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 // 8,), np.dtype[v8bfp16ebs8]] mt_b_bytes = _bfp16_bytes(K_TILE * N_TILE) 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 // 8,), np.dtype[v8bfp16ebs8]] @@ -361,7 +363,13 @@ def _split_run(run): # immediately after the mmul loop was re-rolled. Larger K still falls back # to non-resident, so check this gate before believing any measurement that # claims to be testing residency. - b_resident = (k_iters * mt_b_bytes * 2) <= (512 - 64) * 1024 + # Count what A and C actually occupy rather than assuming C's 64 KB is the + # only other tenant. The old form ignored A entirely and hardcoded C's size, + # which at tile_n=128 (where C doubles and resident B is 432 KB) admitted a + # configuration that then failed address assignment outright. + b_resident = (k_iters * mt_b_bytes * B_DEPTH) <= ( + 512 * 1024 - mt_a_bytes * A_DEPTH - mt_out_bytes * C_DEPTH + ) 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 From 4c313ffe6b9afd5eae1bc5d1341bf183d3e1f72f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 13:28:34 -0600 Subject: [PATCH 23/31] flm_gemm: search L1 jointly over A-tile height and B depth _default_tile_ma picked the A-tile height against a fixed double-buffered B. That is one axis of a two-axis budget, and it cannot express the tradeoff that matters at wider n tiles: there, B's L1 object is large enough that a double-buffered pair does not fit, and giving up B's prefetch is what buys a deeper k slice -- colA 8 -> 16 is worth 3.67 -> 2.28 cycles per mac, far more than B's L1 prefetch. _default_l1 now returns both, deepest B first and then the tallest A that fits, so tile_n=64 is unchanged at (T_MA=32, B_DEPTH=2). Also lets the resident-B gate fall back to a single-buffered memtile B rather than giving up residency: at tile_n=128 B is 432 KB double-buffered and does not fit beside A and C, but 216 KB does, and there residency is worth much more than the prefetch because it is also what stops B's DDR traffic quadrupling. Neither changes the shipped configuration. 39/39 green. --- iron/operators/flm_gemm/design.py | 65 ++++++++++++++++++------------- iron/operators/flm_gemm/op.py | 4 +- 2 files changed, 41 insertions(+), 28 deletions(-) diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 826b7379d..a662386a7 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -89,30 +89,31 @@ def _bfp16_bytes(elems): return elems // 8 * 9 -def _default_tile_ma(n_tile, ct_max_k): - """Largest A-tile height whose L1 working set fits. - - A is dead 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 - (ct_max_k=128 at n=64), which measured 3.67 -> 2.64 cycles per 8x8x8 mac. +def _default_l1(n_tile, ct_max_k): + """Pick (A-tile height, L1 B depth) -- the largest working set that fits. + + 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 - b = _bfp16_bytes(ct_max_k * n_tile) * B_DEPTH cout = CT_OUT_LEN * 2 * C_DEPTH - # Not below 32: at t_ma = 16 the mmul's z loop has a single trip and - # Peano's AIE2P backend aborts with "Use not jointly dominated by defs" - # on the bfp16 path. 32 also halves the A object count, which is fifo - # overhead the asymmetry would otherwise add. - for t_ma in (M_TILE, M_TILE // 2): - if t_ma < 2 * R: - break - a = (2 * R * ct_max_k) * (t_ma // R // 2) * 2 * A_DEPTH - if acc + a + b + cout <= L1_BUDGET: - return t_ma - raise ValueError( - f"no A-tile height fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}" - ) + for b_depth in (B_DEPTH, 1): + b = _bfp16_bytes(ct_max_k * n_tile) * 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 <= L1_BUDGET: + return t_ma, b_depth + raise ValueError(f"nothing fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}") def flm_gemm( @@ -145,7 +146,8 @@ def flm_gemm( # 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. - T_MA = _default_tile_ma(N_TILE, CT_MAX_K) if tile_ma is None else tile_ma + _t_ma_fit, L1_B_DEPTH = _default_l1(N_TILE, CT_MAX_K) + T_MA = _t_ma_fit if tile_ma is None else 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}" @@ -367,9 +369,18 @@ def _split_run(run): # only other tenant. The old form ignored A entirely and hardcoded C's size, # which at tile_n=128 (where C doubles and resident B is 432 KB) admitted a # configuration that then failed address assignment outright. - b_resident = (k_iters * mt_b_bytes * B_DEPTH) <= ( - 512 * 1024 - mt_a_bytes * A_DEPTH - mt_out_bytes * C_DEPTH + # Prefer a double-buffered resident B so the next column-block prefetches + # behind this one's replay. Single-buffering costs that prefetch and was + # measured WORSE than not being resident at all -- at tile_n=64, where B is + # small enough that depth 2 fits anyway. At tile_n=128 B is twice the size + # and depth 2 does not fit, but depth 1 does; and there residency is worth + # far more, because it is also what keeps B's DDR from quadrupling. So take + # the deepest that fits rather than giving up on residency. + mt_free = 512 * 1024 - 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 @@ -381,12 +392,14 @@ def _split_run(run): 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=B_DEPTH) + 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( tile=Tile(c, 1), obj_type=ct_b_ty, - depth=B_DEPTH, + 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 diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py index b6285e7df..109e3c6a5 100644 --- a/iron/operators/flm_gemm/op.py +++ b/iron/operators/flm_gemm/op.py @@ -240,11 +240,11 @@ def _tile_ma(self) -> int: object from it while the kernel derives the mmul's rowA from it, so a mismatch reads past the buffer and produces garbage rather than a build error.""" - from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N, _default_tile_ma + from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N, _default_l1 if self.tile_ma is not None: return self.tile_ma - return _default_tile_ma(self.tile_n, CT_MAX_K_FOR_N[self.tile_n]) + return _default_l1(self.tile_n, CT_MAX_K_FOR_N[self.tile_n])[0] @property def _kernel_object(self) -> str: From 1595b5124fd0599c29c57f6d631b5d428dee417b Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Wed, 9 Sep 2026 15:13:46 -0600 Subject: [PATCH 24/31] flm_gemm: re-measure the tile_n sweep, and pin down what C's write scatter costs The sweep table had been marked stale since cf270ed -- it predated the rolled mmul, resident B, ATB and bfp16 B, which between them took 1024/1536/6144 from 1741 to 1141 us. Re-measured all six shapes against the current design, min of per-run medians over 6 rounds with the two tile_n builds interleaved round-robin (this box is bimodal ~6%, so running all of one and then all of the other measures drift rather than design). The default rule is unchanged in sign -- tile_n=128 still wins only at k_iters=1 -- but both edges moved: its margin there narrowed 8% -> 3%, and its penalty elsewhere grew from ~1.2-1.25x to 1.2-1.7x. Both follow from tile_n=128 giving up resident B, which costs more the more k there is. FLM_C_RUN2 is a new ablation next to FLM_C_LINEAR. C leaves each core as a 64-element (128 B) run, and the guess was that this sits below a ~256 B DDR write-efficiency cliff, making a column-pair join worth ~100 us. It is not: matched for byte count and footprint, halving the number of runs and doubling their length buys 14.9 us min / 13.0 us median, against 98.2 / 96.9 for removing the scatter entirely. So there is no cliff at 256 B -- the cost tracks the NUMBER of scattered runs and keeps paying well past it, and one doubling recovers only 15%. Both knobs write C to the wrong place; they are timing probes, not modes. That measurement is why the column-pair join is not being built: it would have taken 1148 -> ~1134 us (1.90x -> 1.92x vs the shipped overlay) in exchange for replacing the whole C leg with explicit Buffer/Lock/Flow/TileDma -- 8 core-side DMA programs plus 2 memtile ones -- because ObjectFifo.join() cannot place its sources interleaved and cannot reorder on the way out either. A hardware- verified reproducer for that limitation lives in /scratch/ehunhoff/objectfifo_repro. Co-Authored-By: Claude --- iron/operators/flm_gemm/README.md | 31 ++++++++++++++++++------------- iron/operators/flm_gemm/design.py | 13 +++++++++++++ 2 files changed, 31 insertions(+), 13 deletions(-) diff --git a/iron/operators/flm_gemm/README.md b/iron/operators/flm_gemm/README.md index bc91568fc..ed9598193 100644 --- a/iron/operators/flm_gemm/README.md +++ b/iron/operators/flm_gemm/README.md @@ -124,22 +124,27 @@ mac. `n=128` instead halves A fetches, because the grid then covers 1024 columns of N per pass rather than 512. Which wins depends on whether compute or data movement is the critical path, and 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. Measured, minimum of 3 runs: - -> **Stale:** the sweep below predates the re-rolled mmul and the widened -> residency gate, which together took 1024/1536/6144 from 1741 to 1434 us. The -> `tile_n` choice it justifies is unlikely to have changed sign (residency does -> not fit at `tile_n=128`, whose `mt_b` is 128 KB), but the absolute numbers -> are no longer right and it wants re-measuring. +hide the extra A traffic. Measured 2026-09-09 against the current design +(rolled mmul, resident B, ATB, bfp16 B), min of per-run medians over 6 rounds +with the two `tile_n` builds interleaved round-robin -- this box is bimodal +~6%, so running all of one and then all of the other measures drift rather +than design: | M / K / N | k_iters | `tile_n=64` | `tile_n=128` | |---|---|---|---| -| 1024 / 512 / 4096 | 1 | 642 us | **589 us** | -| 1024 / 1024 / 4096 | 2 | **850 us** | 1034 us | -| 1024 / 1536 / 6144 | 3 | **1741 us** | 2178 us | -| 1024 / 2560 / 4096 | 5 | **1891 us** | 2371 us | -| 2048 / 2048 / 2048 | 4 | **1535 us** | 1924 us | -| 256 / 4096 / 1024 | 8 | **254 us** | 316 us | +| 1024 / 512 / 4096 | 1 | 514 us | **498 us** | +| 1024 / 1024 / 4096 | 2 | **591 us** | 931 us | +| 1024 / 1536 / 6144 | 3 | **1141 us** | 1960 us | +| 1024 / 2560 / 4096 | 5 | **1239 us** | 1940 us | +| 2048 / 2048 / 2048 | 4 | **915 us** | 1581 us | +| 256 / 4096 / 1024 | 8 | **227 us** | 265 us | + +The default rule is unchanged in sign: `tile_n=128` still wins only at +`k_iters=1`. But its margin there has narrowed to 3% (was 8%) and its penalty +everywhere else has grown -- at `k_iters>=2` it is now 1.2-1.7x slower where +it used to be 1.2-1.25x. 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): the more k +there is to reduce over, the more that costs. `pack_B` is bound to the operator because the packing layout depends on `tile_n`; call `op.pack_B(B)`, not `FLMGEMM.pack_B(B)`. diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index a662386a7..3a847f5d2 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -551,6 +551,19 @@ def c_tap(mega_col, c): sizes=[1, 1, 1, blk], strides=[0, 0, 0, 1], ) + if _os.environ.get("FLM_C_RUN2") == "1": + # ABLATION: identical byte count and identical DDR footprint, but + # HALF as many runs each TWICE as long (N_TILE*2 = 256 B instead of + # 128 B), by walking every other row. Writes C to the WRONG place. + # This isolates exactly what the column-pair join would buy -- + # FLM_C_LINEAR above removes ALL scatter, so it is the ceiling for + # perfect linearisation, not for the 128 B -> 256 B step. + return TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * N_TILE, + sizes=[1, m_row_blocks, ROWS * M_TILE // 2, N_TILE * 2], + strides=[0, ROWS * M_TILE * N, N * 2, 1], + ) return TensorAccessPattern( tensor_dims=(M * N,), offset=(mega_col * COLS + c) * N_TILE, From 3bd7dbeea759473fe9d0c505b1ea5bbdbd094c5f Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 10:44:25 -0600 Subject: [PATCH 25/31] benchmarking support --- iron/operators/flm_gemm/bench_vs_flm.py | 370 ++++++++++++++++++++++++ iron/operators/flm_gemm/design.py | 69 +++++ iron/operators/flm_gemm/test.py | 17 ++ 3 files changed, 456 insertions(+) create mode 100644 iron/operators/flm_gemm/bench_vs_flm.py diff --git a/iron/operators/flm_gemm/bench_vs_flm.py b/iron/operators/flm_gemm/bench_vs_flm.py new file mode 100644 index 000000000..ef5526daf --- /dev/null +++ b/iron/operators/flm_gemm/bench_vs_flm.py @@ -0,0 +1,370 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark flm_gemm against the shipped FastFlowLM overlay and IRON's generic GEMM. + +Kept out of ``test.py`` on purpose: this file needs an external FastFlowLM +install, which the operator itself does not. ``test.py`` must stay runnable in +CI with nothing but this repo; this one skips the whole module when the install +is absent. + +Every shape is one test, covering the real Gemma4 projections (E2B and E4B x +{q, kv, o, gateup, down}) at three prefill lengths. Three competitors run per +shape: + + iron : the ``FLMGEMM`` operator (arg order A, B, C) + flm : FastFlowLM's shipped ``mm.xclbin`` + its dumped TXN insts (order C, A, B) + gemm : IRON's generic ``GEMM`` operator, same emulated-bfp16 numerics + +The box is bimodal by ~6% (see the npu-bimodal-timing note), so the three are +interleaved round-robin over several rounds and each is scored by the MINIMUM +of its per-round medians. Running all of one competitor and then all of another +fabricates differences of about that size. Everything is compiled up front; +nothing is rebuilt between rounds. + +No time is reported for a dispatch whose output was not checked. + +Usage:: + + pytest iron/operators/flm_gemm/bench_vs_flm.py --iterations 1 + pytest iron/operators/flm_gemm/bench_vs_flm.py -k E2B --csv-output flm.csv +""" + +import os +import statistics +import subprocess +import time +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.flm_gemm.op import FLMGEMM +from iron.operators.gemm.op import GEMM + +# --------------------------------------------------------------------------- +# External FastFlowLM install. Both are overridable so this is not pinned to +# one machine's layout; absent either, the module skips rather than errors. +# --------------------------------------------------------------------------- +FLM_XCLBIN = Path( + os.environ.get( + "FLM_MM_XCLBIN", + "/scratch/ehunhoff/flm-release-1.0.4/extracted/opt/fastflowlm/share/flm" + "/xclbins/Gemma4-E2B-IT-NPU2/mm.xclbin", + ) +) +# Dumps the overlay's control instructions for one (M, K, N). The generator +# baked into it is gemma4_e2b_mm_txn, so it belongs to the xclbin above. +# Previously lived in /tmp, which is tmpfs here and did not survive reboots. +TXN_DUMP = Path( + os.environ.get("FLM_TXN_DUMP", "/scratch/ehunhoff/flm_gemm_bench/txn/dump") +) +TXN_CACHE = Path(os.environ.get("FLM_TXN_CACHE", TXN_DUMP.parent)) + +# The shipped overlay is a fixed n=128 design; its B layout is not negotiable. +FLM_N_TILE = 128 +# B's memtile odometer, shared by both packers. K_TILE is the operator's. +K_TILE, S, T = 512, 8, 8 + +# Interleaved rounds per test, and timed dispatches per competitor per round. +# 6 rounds is the floor at which min and median stopped disagreeing on this box. +ROUNDS = 6 +ITERS = 30 +WARMUP = 20 + +# err/mass budgets. iron and 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 _flm_available(): + return FLM_XCLBIN.is_file() and os.access(TXN_DUMP, os.X_OK) + + +if not _flm_available(): + pytest.skip( + f"FastFlowLM install not found (looked for {FLM_XCLBIN} and {TXN_DUMP}); " + "set FLM_MM_XCLBIN / FLM_TXN_DUMP to point at one", + allow_module_level=True, + ) + +_dev = aie_utils.get_current_device() +if _dev.cols < 8 or _dev.resolve().name != "npu2": + pytest.skip( + "flm_gemm is a fixed 4x8 npu2 design; this device cannot run it", + allow_module_level=True, + ) + + +# --------------------------------------------------------------------------- +# Shapes +# --------------------------------------------------------------------------- +# Every projection of both Gemma4 variants FastFlowLM ships, at three prefill +# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. +# 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] + +# The two E4B projections with a 10240-wide dimension need a shim DMA +# descriptor stride past the AIE2p shim's 20-bit step field once M walks more +# than one mega_row. design.py rejects them at construction; see its comment +# and test.py's test_flm_gemm_stride_overflow_rejected. They are skipped here +# rather than left to raise, so the sweep reports 26 results and 4 known +# blocks instead of 4 errors. +BLOCKED = "K/N=10240 at M>256 overflows the shim's 20-bit DMA stride field" + + +def get_params(): + params = [] + for model, projections in (("E2B", E2B_PROJ), ("E4B", E4B_PROJ)): + for M in PREFILL_LENGTHS: + for proj, K, N in projections: + blocked = M > 256 and max(K, N) > 8191 + marks = [pytest.mark.skip(reason=BLOCKED)] if blocked else [] + params.append( + pytest.param( + model, proj, M, K, N, marks=marks, id=f"{model}-{proj}-M{M}" + ) + ) + return params + + +# --------------------------------------------------------------------------- +# Inputs and packing +# --------------------------------------------------------------------------- +def make_inputs(M, K, N): + """Identical data for all three competitors, and the reference to check.""" + 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, not 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, and near-zero outputs are + # relatively uncheckable. Same rationale as test.py's bound. + return A, B, Af @ Bf, float((Af.abs() @ Bf.abs()).mean()) + + +def flm_pack_B(Bt, n_tile=FLM_N_TILE): + """(K, N) row-major -> the shipped overlay's memtile order. + + Odometer (n//T, k%S, k//S, n%T), tiles ordered by column stripe then + k-block. Mirrors ``FLMGEMM.pack_B``'s tiling but stops at bf16: the + overlay consumes bf16, not the bfp16 blocks the IRON operator takes, and + is fixed at n=128 regardless of what the IRON operator chooses. + """ + b = Bt.float().numpy() + K, N = b.shape + out = [] + for cb in range(N // n_tile): + stripe = b[:, cb * n_tile : (cb + 1) * n_tile] + for kb in range(K // K_TILE): + tile = stripe[kb * K_TILE : (kb + 1) * K_TILE, :] + out.append( + tile.reshape(K_TILE // S, S, n_tile // T, T) + .transpose(2, 1, 0, 3) + .ravel() + ) + return torch.from_numpy(np.concatenate(out).astype(np.float32)).to(torch.bfloat16) + + +def flm_insts(M, K, N): + """Path to the overlay's TXN insts for one shape, dumping it if absent.""" + path = TXN_CACHE / f"txn_{M}_{K}_{N}.bin" + if not path.is_file(): + path.parent.mkdir(parents=True, exist_ok=True) + subprocess.run([str(TXN_DUMP), str(M), str(K), str(N), str(path)], check=True) + return path + + +# --------------------------------------------------------------------------- +# Competitors. Each returns (run, c_bo, label, xclbin_path) with the buffers +# already bound, so the timed section is nothing but the dispatch. +# --------------------------------------------------------------------------- +def _bind(xclbin, insts, args): + from aie.utils.npukernel import NPUKernel + + handle = aie_utils.DefaultNPURuntime.load(NPUKernel(str(xclbin), str(insts))) + return lambda: aie_utils.DefaultNPURuntime.run(handle, list(args)) + + +def setup_iron(M, K, N, A, B, ctx): + op = FLMGEMM(M=M, K=K, N=N, context=ctx) + built_before = _artifacts_exist(op, ctx) + t0 = time.perf_counter() + op.compile() + compile_s = time.perf_counter() - t0 + c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) + run = op.get_callable() + args = [ + XRTTensor.from_torch(A.flatten()), + XRTTensor.from_torch(op.pack_B(B).flatten()), + c_bo, + ] + return _Competitor( + "iron", + lambda: run(*args), + c_bo, + Path(op.xclbin_artifact.filename), + None if built_before else compile_s, + BUDGET_CONV_EVEN, + ) + + +def setup_gemm(M, K, N, A, B, ctx): + # Left at the operator's defaults, which are the same emulated-bfp16 mmul + # and conv_even rounding flm_gemm uses -- a like-for-like comparison, not + # flm_gemm against a more accurate and necessarily slower configuration. + op = GEMM(M=M, K=K, N=N, context=ctx) + built_before = _artifacts_exist(op, ctx) + t0 = time.perf_counter() + op.compile() + compile_s = time.perf_counter() - t0 + c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) + run = op.get_callable() + # b_col_maj defaults False, so B goes in as plain row-major (K, N). + args = [XRTTensor.from_torch(A.flatten()), XRTTensor.from_torch(B.flatten()), c_bo] + return _Competitor( + "gemm", + lambda: run(*args), + c_bo, + Path(op.xclbin_artifact.filename), + None if built_before else compile_s, + BUDGET_CONV_EVEN, + ) + + +def setup_flm(M, K, N, A, B, ctx): + c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) + # The shipped overlay's host contract is C, A, B -- not IRON's A, B, C. + args = [ + c_bo, + XRTTensor.from_torch(A.flatten()), + XRTTensor.from_torch(flm_pack_B(B).flatten()), + ] + run = _bind(FLM_XCLBIN, flm_insts(M, K, N), args) + # Prebuilt and shipped: there is no compile to time. + return _Competitor("flm", run, c_bo, FLM_XCLBIN, None, BUDGET_FLOOR) + + +class _Competitor: + def __init__(self, name, run, c_bo, xclbin, compile_s, budget): + self.name = name + self.run = run + self.c_bo = c_bo + self.xclbin = xclbin + self.compile_s = compile_s + self.budget = budget + self.round_medians = [] + + 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): + ts = [] + for _ in range(ITERS): + t0 = time.perf_counter() + self.run() + ts.append((time.perf_counter() - t0) * 1e6) + 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 + + +def _artifacts_exist(op, ctx): + """Whether this operator's xclbin is already built in ctx's build dir. + + Compile time is only meaningful on a genuine miss; on a hit ``compile()`` + returns in milliseconds and reporting that as a build time would be a lie. + """ + if not op.artifacts: + op.set_up_artifacts() + return (Path(ctx.build_dir) / f"{op.name}.xclbin").is_file() + + +# --------------------------------------------------------------------------- +# The benchmark +# --------------------------------------------------------------------------- +@pytest.mark.metrics( + IronLatency=r"iron latency \(us\): (?P[\d\.]+)", + FLMLatency=r"flm latency \(us\): (?P[\d\.]+)", + GEMMLatency=r"gemm latency \(us\): (?P[\d\.]+)", + SpeedupVsFLM=r"speedup vs flm: (?P[\d\.]+)", + SpeedupVsGEMM=r"speedup vs gemm: (?P[\d\.]+)", + IronThroughput=r"iron throughput: (?P[\d\.e\+-]+) GFLOP/s", + IronJitterPct=r"iron jitter \(%\): (?P[\d\.]+)", + IronXclbinKB=r"iron xclbin \(KB\): (?P[\d\.]+)", + IronCompileTime=r"iron compile \(s\): (?P[\d\.]+)", +) +@pytest.mark.parametrize("model,proj,M,K,N", get_params()) +def test_flm_gemm_vs_flm(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. + competitors = [ + setup_iron(M, K, N, A, B, aie_context), + setup_flm(M, K, N, A, B, aie_context), + setup_gemm(M, K, N, A, B, aie_context), + ] + + bad = [c for c in competitors 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 competitors: + for _ in range(WARMUP): + c.run() + # Round-robin, never all of one then all of another. + for _ in range(ROUNDS): + for c in competitors: + c.time_round() + + by_name = {c.name: c for c in competitors} + iron = by_name["iron"] + + print() + for c in competitors: + print(f"{c.name} latency (us): {c.us:.1f}") + print(f"{c.name} err/mass: {c.err:.3e}") + print(f"speedup vs flm: {by_name['flm'].us / iron.us:.3f}") + print(f"speedup vs gemm: {by_name['gemm'].us / iron.us:.3f}") + print(f"iron throughput: {2.0 * M * K * N / (iron.us * 1e-6) / 1e9:.6e} GFLOP/s") + print(f"iron jitter (%): {iron.jitter_pct:.2f}") + print(f"iron xclbin (KB): {iron.xclbin.stat().st_size / 1024:.1f}") + if iron.compile_s is not None: + print(f"iron compile (s): {iron.compile_s:.1f}") + print() diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 3a847f5d2..0f6341927 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -89,6 +89,21 @@ def _bfp16_bytes(elems): return elems // 8 * 9 +# Shim-tile DMA BD step field is 20 bits wide (AIE2p; see mlir-aie's +# AIETargetModel::getDmaBdStepBits and AIEXDialect.cpp's "Stride N exceeds" +# verifier). An IR-level bf16-element stride S is re-expressed in hardware +# units as (S - 1) * 2 bytes / 4-byte address granularity before that check, +# so S must satisfy hw_stride_bf16(S) <= (1 << 20) - 1. +_SHIM_STEP_BITS = 20 +_BF16_BYTES = 2 +_ADDR_GRANULARITY_BYTES = 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): """Pick (A-tile height, L1 B depth) -- the largest working set that fits. @@ -184,6 +199,60 @@ def flm_gemm( # How many times the whole grid sweeps, in each dimension. m_row_blocks = M // MIN_M k_iters = K // K_TILE + # The mega_row dimension's stride (ROWS*M_TILE*{K,N}) overflows the shim's + # 20-bit step field once K or N crosses ~8191 elements -- E4B's FFN width + # (10240) does, E2B's max (6144) does not. Only relevant once M actually + # walks more than one mega_row (m_row_blocks>1); at M=256 the dimension is + # degenerate (size 1) and the taplib/MLIR toolchain strips it before this + # stride is ever encoded, so it never fails regardless of K/N. + # + # Attempted fix: split the single 4D descriptor into m_row_blocks separate + # 3D fill()/drain() calls, each carrying the mega_row jump in its OFFSET + # (unbounded) instead of a shared STRIDE (20-bit limited). Compiles cleanly + # either direction (A's L3->L2 broadcast feed, or C's L2->L3 drain), but + # the dispatch hangs on real hardware (XRT: ERT_CMD_STATE_TIMEOUT) for + # BOTH once tested at the actual target shape (K or N = 10240) rather + # than a small stand-in -- a small forced-split repro (e.g. N=4096) + # appeared to work for the C direction, which was a false signal: the + # same shape at the real N=10240 hangs reproducibly on fresh builds. Not + # a buffer-depth issue either (tried objectfifo depths 2/4/8/16 for A, + # all hang). A dedicated IR-forensics pass diffed every aiecc compiler + # stage between working and broken small-shape builds and found the + # generated MLIR/BDs/locks correct and consistent throughout -- not a + # compiler bug. The kernel driver's health-report interface (amdxdna's + # `aie2_dump_ctx`, see + # /usr/src/xrt-amdxdna-*/driver/amdxdna/aie2_ctx.c and + # aie2_msg_priv.h's `struct app_health_report`) shows, on an actual hang: + # `Fatal error type: 0x0` and every exception field zero (NOT a + # crash/fault), `dpu_pc`/`txn_op_id` at their documented "not captured + # for this op type" sentinel (0xffffffff), while `ctx_pc` holds a real + # captured address -- consistent with the on-chip firmware sitting in a + # wait loop for a completion signal the array's DMA engine never raises. + # The driver's TDR watchdog (aie2_tdr.c, 2s default) is what eventually + # force-resets it. This is a genuine, silent hardware/firmware + # synchronization hang, not something splitting the host-issued task + # count works around. Root cause not fully pinned down -- the firmware + # itself is closed, unsymbolized microcode + # (/lib/firmware/amdnpu/*/npu*.sbin); further work needs AMD-internal + # firmware source or a hardware debugger, neither available here. + # + # Until a real fix lands, K/N > ~8191 at M > 256 raises a clear + # compile-time error below rather than emitting a build that hangs on + # real hardware. + if m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * K): + raise ValueError( + f"K={K} at M={M} would need a shim DMA descriptor stride " + f"(ROWS*M_TILE*K={ROWS * M_TILE * K}) that exceeds the AIE2p " + "shim's 20-bit step field (see the comment above this check) " + "-- not yet fixed." + ) + if m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * N): + raise ValueError( + f"N={N} at M={M} would need a shim DMA descriptor stride " + f"(ROWS*M_TILE*N={ROWS * M_TILE * N}) that exceeds the AIE2p " + "shim's 20-bit step field (see the comment above this check) " + "-- not yet fixed." + ) # 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 diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index b76791725..0bc79bcad 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -125,3 +125,20 @@ def test_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): print(f"Throughput: {gflops:.6e} GFLOP/s\n") assert not errors, "Test failed" + + +@pytest.mark.parametrize( + "M,K,N", + [ + (1024, 10240, 2560), # E4B down-proj: K overflows the shim's 20-bit stride + (1024, 2560, 10240), # E4B gateup-proj: N overflows it instead + ], +) +def test_flm_gemm_stride_overflow_rejected(M, K, N, aie_context): + # K or N > ~8191 at M > 256 needs a shim DMA descriptor stride that + # exceeds the AIE2p shim's 20-bit step field. Splitting the transfer + # into multiple descriptors compiles but hangs real hardware (see + # design.py's comment above this check) -- so this must keep failing + # fast at construction, not silently emit a build that hangs. + with pytest.raises(ValueError, match="20-bit step field"): + FLMGEMM(M=M, K=K, N=N, context=aie_context).compile() From 1c1cd384df19f10765cbd2c63a9580d50036b6f0 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 13:35:15 -0600 Subject: [PATCH 26/31] flm_gemm: support K/N=10240 by splitting the mega_row descriptor The four E4B FFN projections (down at K=10240, gate/up at N=10240, both at M>=1024) were rejected at construction: their mega_row dimension has stride ROWS*M_TILE*{K,N}, which lands in the shim BD's iteration field and overflows its 20-bit step. Issue that leg as one transfer per mega_row instead, carrying the jump in the offset, which has no such limit. Only the overflowing leg splits, and K and N never both overflow on a real shape. That split was tried before and recorded as an unfixable firmware hang. It was not. TaskGroup.finish() emits dma_free_task, which returns the buffer descriptor id to a compile-time allocator that never checks the transfer completed (AIEAssignRuntimeSequenceBDIDs::recycle, isAwait=false), and ids are per shim TILE, shared across channels. Retiring each mega_row immediately therefore collapsed every task on a column onto bd_id 0 and reprogrammed it mid-flight -- including B, whose descriptor streams across all mega_rows. Visible statically: dump bd_ids after --aie-assign-runtime-sequence-bd-ids and the unsplit build has {0,1,2} where the split one has {0}. A second unmodelled shim resource bounds it too: the channel task queue is 4 deep and NpuPushQueueOp pushes unconditionally, so m_row_blocks transfers back to back on one channel hang at M=2048. Both are handled by emitting a split block in windows of at most 4 mega_rows and awaiting each window before the next -- the await is what makes the descriptors reusable and drains the queue. M<=1024 is a single window, so previously working shapes are unaffected. Measured (err/mass against a 4e-3 budget): M=1024 K=10240 N=2560 6245 us 1.149e-04 M=1024 K=2560 N=10240 6169 us 2.298e-04 M=2048 K=10240 N=2560 9927 us 1.150e-04 M=2048 K=2560 N=10240 8525 us 2.299e-04 All four join test.py's extensive params and bench_vs_flm.py no longer skips them, so the prefill sweep covers all 30 shapes. Co-Authored-By: Claude --- iron/operators/flm_gemm/bench_vs_flm.py | 17 +- iron/operators/flm_gemm/design.py | 274 ++++++++++++++++-------- iron/operators/flm_gemm/test.py | 49 +++-- 3 files changed, 228 insertions(+), 112 deletions(-) diff --git a/iron/operators/flm_gemm/bench_vs_flm.py b/iron/operators/flm_gemm/bench_vs_flm.py index ef5526daf..e23b35673 100644 --- a/iron/operators/flm_gemm/bench_vs_flm.py +++ b/iron/operators/flm_gemm/bench_vs_flm.py @@ -125,13 +125,10 @@ def _flm_available(): ] PREFILL_LENGTHS = [256, 1024, 2048] -# The two E4B projections with a 10240-wide dimension need a shim DMA -# descriptor stride past the AIE2p shim's 20-bit step field once M walks more -# than one mega_row. design.py rejects them at construction; see its comment -# and test.py's test_flm_gemm_stride_overflow_rejected. They are skipped here -# rather than left to raise, so the sweep reports 26 results and 4 known -# blocks instead of 4 errors. -BLOCKED = "K/N=10240 at M>256 overflows the shim's 20-bit DMA stride field" +# All 30 shapes run. The four E4B projections with a 10240-wide dimension used +# to be skipped here: at M>256 their mega_row stride overflows the shim BD's +# 20-bit iteration step. design.py now issues that leg as one transfer per +# mega_row, retired in windows -- see its a_split comment. def get_params(): @@ -139,12 +136,8 @@ def get_params(): for model, projections in (("E2B", E2B_PROJ), ("E4B", E4B_PROJ)): for M in PREFILL_LENGTHS: for proj, K, N in projections: - blocked = M > 256 and max(K, N) > 8191 - marks = [pytest.mark.skip(reason=BLOCKED)] if blocked else [] params.append( - pytest.param( - model, proj, M, K, N, marks=marks, id=f"{model}-{proj}-M{M}" - ) + pytest.param(model, proj, M, K, N, id=f"{model}-{proj}-M{M}") ) return params diff --git a/iron/operators/flm_gemm/design.py b/iron/operators/flm_gemm/design.py index 0f6341927..eacb0b47c 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm_gemm/design.py @@ -97,6 +97,15 @@ def _bfp16_bytes(elems): _SHIM_STEP_BITS = 20 _BF16_BYTES = 2 _ADDR_GRANULARITY_BYTES = 4 +# Buffer descriptors per shim tile (AIETargetModel::getNumBDs, ShimNOCTile). +# A per-TILE resource shared by every channel and both directions, so the A, B +# and C legs on one column all draw from the same 16. +SHIM_BDS = 16 +# Entries in a shim DMA channel's task queue. Nothing in mlir-aie models this +# -- AIEDmaToNpu's NpuPushQueueOp pushes unconditionally -- so overrunning it +# is a silent device hang, not a diagnostic. Measured here at K=10240 M=1024: +# 4 outstanding tasks on one channel run, 8 hang. +SHIM_TASK_QUEUE = 4 def _hw_stride_ok(stride_elems): @@ -199,60 +208,61 @@ def flm_gemm( # How many times the whole grid sweeps, in each dimension. m_row_blocks = M // MIN_M k_iters = K // K_TILE - # The mega_row dimension's stride (ROWS*M_TILE*{K,N}) overflows the shim's - # 20-bit step field once K or N crosses ~8191 elements -- E4B's FFN width - # (10240) does, E2B's max (6144) does not. Only relevant once M actually - # walks more than one mega_row (m_row_blocks>1); at M=256 the dimension is - # degenerate (size 1) and the taplib/MLIR toolchain strips it before this - # stride is ever encoded, so it never fails regardless of K/N. + # 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. # - # Attempted fix: split the single 4D descriptor into m_row_blocks separate - # 3D fill()/drain() calls, each carrying the mega_row jump in its OFFSET - # (unbounded) instead of a shared STRIDE (20-bit limited). Compiles cleanly - # either direction (A's L3->L2 broadcast feed, or C's L2->L3 drain), but - # the dispatch hangs on real hardware (XRT: ERT_CMD_STATE_TIMEOUT) for - # BOTH once tested at the actual target shape (K or N = 10240) rather - # than a small stand-in -- a small forced-split repro (e.g. N=4096) - # appeared to work for the C direction, which was a false signal: the - # same shape at the real N=10240 hangs reproducibly on fresh builds. Not - # a buffer-depth issue either (tried objectfifo depths 2/4/8/16 for A, - # all hang). A dedicated IR-forensics pass diffed every aiecc compiler - # stage between working and broken small-shape builds and found the - # generated MLIR/BDs/locks correct and consistent throughout -- not a - # compiler bug. The kernel driver's health-report interface (amdxdna's - # `aie2_dump_ctx`, see - # /usr/src/xrt-amdxdna-*/driver/amdxdna/aie2_ctx.c and - # aie2_msg_priv.h's `struct app_health_report`) shows, on an actual hang: - # `Fatal error type: 0x0` and every exception field zero (NOT a - # crash/fault), `dpu_pc`/`txn_op_id` at their documented "not captured - # for this op type" sentinel (0xffffffff), while `ctx_pc` holds a real - # captured address -- consistent with the on-chip firmware sitting in a - # wait loop for a completion signal the array's DMA engine never raises. - # The driver's TDR watchdog (aie2_tdr.c, 2s default) is what eventually - # force-resets it. This is a genuine, silent hardware/firmware - # synchronization hang, not something splitting the host-issued task - # count works around. Root cause not fully pinned down -- the firmware - # itself is closed, unsymbolized microcode - # (/lib/firmware/amdnpu/*/npu*.sbin); further work needs AMD-internal - # firmware source or a hardware debugger, neither available here. + # Fix: issue that leg as m_row_blocks separate transfers, each carrying the + # mega_row jump in its OFFSET (unbounded) instead of 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. # - # Until a real fix lands, K/N > ~8191 at M > 256 raises a clear - # compile-time error below rather than emitting a build that hangs on - # real hardware. - if m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * K): - raise ValueError( - f"K={K} at M={M} would need a shim DMA descriptor stride " - f"(ROWS*M_TILE*K={ROWS * M_TILE * K}) that exceeds the AIE2p " - "shim's 20-bit step field (see the comment above this check) " - "-- not yet fixed." - ) - if m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * N): + # This was long recorded as an unfixable firmware hang. It was not. The + # earlier attempt retired each mega_row's TaskGroup immediately, and + # TaskGroup.finish() emits dma_free_task, which hands the buffer descriptor + # id back to a COMPILE-TIME allocator that never checks the transfer + # finished (mlir-aie AIEAssignRuntimeSequenceBDIDs::recycle, isAwait=false). + # Ids are per shim TILE, shared across channels and directions, so every + # task on a column collapsed onto bd_id 0 and reprogrammed it mid-flight -- + # including B, whose descriptor streams across every mega_row. Verify with + # aie-opt --aie-substitute-shim-dma-allocations + # --aie-assign-runtime-sequence-bd-ids: the ids on a shim tile must be + # distinct, and were all 0. + 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"N={N} at M={M} would need a shim DMA descriptor stride " - f"(ROWS*M_TILE*N={ROWS * M_TILE * N}) that exceeds the AIE2p " - "shim's 20-bit step field (see the comment above this check) " - "-- not yet fixed." + 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 @@ -573,16 +583,33 @@ def core_fn(acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): # 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 # 10-bit size field (largest are K_TILE=512 and ROWS*M_TILE=256). - def a_tap(mega_col, r): + 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. - 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], - ) + # + # 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 @@ -606,9 +633,26 @@ def b_tap(mega_col, c): ), ) - def c_tap(mega_col, c): + 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. + # 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. The + # ablation knobs below are unsplit-only; they write C to the wrong + # place by construction and exist purely for timing. + 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): if _os.environ.get("FLM_C_LINEAR") == "1": # ABLATION: same byte count, one contiguous run per column-block # instead of ROWS*M_TILE runs of N_TILE. Writes C to the WRONG @@ -659,34 +703,94 @@ def sequence(A, B, C, a_prods, b_prods, c_conses): # 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 now a single task: a - # block costs 3 buffer descriptors on a shim column (A + B + C), so two - # in flight is 6 of 16. Per-object tasks needed 1 + 2*k_iters and could - # not be overlapped at all. - # Keep OVERLAP column-blocks in flight. A block costs 3 shim buffer - # descriptors on a column (A + B + C) against 16 available, so the - # ceiling is 5; the operator is DDR-rate bound rather than byte bound - # (53 GB/s of a 63-70 GB/s roof), so how deeply the fills are pipelined - # is what decides the rate. - pending = [] - for mega_col, active_cols in blocks: - tg_c = TaskGroup() - for c in range(active_cols): - c_conses[c].drain(C, c_tap(mega_col, c), group=tg_c, wait=True) - tg_f = TaskGroup() - for r in range(ROWS): - a_prods[r].fill(A, a_tap(mega_col, r), group=tg_f) - for c in range(active_cols): - b_prods[c].fill(B, b_tap(mega_col, c), group=tg_f) - - pending.append([tg_f, tg_c]) - while len(pending) >= OVERLAP: - for tg in pending.pop(0): + # This is affordable only because each leg is a single task per + # column-block. Per-object tasks needed 1 + 2*k_iters and could not be + # overlapped at all. + # Keep OVERLAP column-blocks in flight, against 16 shim buffer + # descriptors per column; the operator is DDR-rate bound rather than + # byte bound (53 GB/s of a 63-70 GB/s roof), 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: finish() emits + # dma_free_task, which returns the buffer descriptor id to a + # compile-time allocator that does not check the transfer completed, so + # retiring a leg early lets the next task reprogram a live descriptor. + # See the a_split comment above for what that cost. + all_mb = list(range(m_row_blocks)) + + def emit_unsplit(): + pending = [] + for mega_col, active_cols in blocks: + tg_c = TaskGroup() + for c in range(active_cols): + for tap in c_taps(mega_col, c, all_mb): + c_conses[c].drain(C, tap, group=tg_c, wait=True) + tg_f = TaskGroup() + for r in range(ROWS): + for tap in a_taps(mega_col, r, all_mb): + a_prods[r].fill(A, tap, group=tg_f) + for c in range(active_cols): + b_prods[c].fill(B, b_tap(mega_col, c), group=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() - 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() + for c in range(active_cols): + b_prods[c].fill(B, b_tap(mega_col, c), group=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: + for c in range(active_cols): + for tap in c_taps(mega_col, c, all_mb): + c_conses[c].drain(C, tap, group=tg_whole, wait=True) + if not a_split: + for r in range(ROWS): + for tap in a_taps(mega_col, r, all_mb): + a_prods[r].fill(A, tap, group=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: + for c in range(active_cols): + for tap in c_taps(mega_col, c, mbs): + c_conses[c].drain(C, tap, group=tg_w, wait=True) + if a_split: + for r in range(ROWS): + for tap in a_taps(mega_col, r, mbs): + # wait=True: the await is what makes this + # window's descriptors reusable by the next. + a_prods[r].fill(A, tap, group=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, diff --git a/iron/operators/flm_gemm/test.py b/iron/operators/flm_gemm/test.py index 0bc79bcad..f9a10ce87 100644 --- a/iron/operators/flm_gemm/test.py +++ b/iron/operators/flm_gemm/test.py @@ -43,6 +43,15 @@ def get_params(): ( 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 ] # fmt: on @@ -127,18 +136,28 @@ def test_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): assert not errors, "Test failed" -@pytest.mark.parametrize( - "M,K,N", - [ - (1024, 10240, 2560), # E4B down-proj: K overflows the shim's 20-bit stride - (1024, 2560, 10240), # E4B gateup-proj: N overflows it instead - ], -) -def test_flm_gemm_stride_overflow_rejected(M, K, N, aie_context): - # K or N > ~8191 at M > 256 needs a shim DMA descriptor stride that - # exceeds the AIE2p shim's 20-bit step field. Splitting the transfer - # into multiple descriptors compiles but hangs real hardware (see - # design.py's comment above this check) -- so this must keep failing - # fast at construction, not silently emit a build that hangs. - with pytest.raises(ValueError, match="20-bit step field"): - FLMGEMM(M=M, K=K, N=N, context=aie_context).compile() +def test_flm_gemm_split_leg_windowing(aie_context): + # K or N > ~8191 at M > 256 makes the mega_row stride overflow the shim + # BD's 20-bit 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 iron.operators.flm_gemm.design import SHIM_BDS, SHIM_TASK_QUEUE + + worst = 1 + 2 * SHIM_TASK_QUEUE + assert worst <= SHIM_BDS, ( + f"a fully split block needs {worst} shim BDs of {SHIM_BDS}; " + "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. + FLMGEMM(M=512, K=10240, N=10240, context=aie_context).compile() From fadc2058b251b12d9e039a3226fc5becaa19d01a Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 14:06:57 -0600 Subject: [PATCH 27/31] flm_gemm/bench: record accuracy, and survive a competitor that cannot build Two gaps the 30-shape sweep hit once the E4B 10240-wide shapes stopped being skipped. err/mass was asserted against a budget and printed but never recorded, so a toolchain bump could move accuracy a long way inside that budget unnoticed. Add IronErr/FLMErr/GEMMErr to the metrics; the dev60->dev85 bump is bit-identical on all 30 shapes, which is only checkable because of this. IRON's generic GEMM still has the 20-bit mega_row stride limitation flm_gemm just fixed: at M>256 with N=10240 its C descriptor wants a 2621440-element stride and aiecc rejects the build. That is a property of that operator, not of the shape, and it is a clean compile-time error rather than a hang -- so report those shapes against the competitors that do build instead of losing flm_gemm's own numbers for two of thirty. Only that specific rejection is tolerated. Co-Authored-By: Claude --- iron/operators/flm_gemm/bench_vs_flm.py | 28 +++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/iron/operators/flm_gemm/bench_vs_flm.py b/iron/operators/flm_gemm/bench_vs_flm.py index e23b35673..6c4827973 100644 --- a/iron/operators/flm_gemm/bench_vs_flm.py +++ b/iron/operators/flm_gemm/bench_vs_flm.py @@ -17,6 +17,10 @@ flm : FastFlowLM's shipped ``mm.xclbin`` + its dumped TXN insts (order C, A, B) gemm : IRON's generic ``GEMM`` operator, same emulated-bfp16 numerics +``gemm`` drops out of the E4B gate/up shapes at M>256: its C descriptor needs a +mega_row stride past the shim's 20-bit step field there, which aiecc rejects. +Those shapes report iron against flm only. + The box is bimodal by ~6% (see the npu-bimodal-timing note), so the three are interleaved round-robin over several rounds and each is scored by the MINIMUM of its per-round medians. Running all of one competitor and then all of another @@ -316,6 +320,13 @@ def _artifacts_exist(op, ctx): GEMMLatency=r"gemm latency \(us\): (?P[\d\.]+)", SpeedupVsFLM=r"speedup vs flm: (?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. + IronErr=r"iron err/mass: (?P[\d\.e\+-]+)", + FLMErr=r"flm err/mass: (?P[\d\.e\+-]+)", + GEMMErr=r"gemm err/mass: (?P[\d\.e\+-]+)", IronThroughput=r"iron throughput: (?P[\d\.e\+-]+) GFLOP/s", IronJitterPct=r"iron jitter \(%\): (?P[\d\.]+)", IronXclbinKB=r"iron xclbin \(KB\): (?P[\d\.]+)", @@ -330,8 +341,20 @@ def test_flm_gemm_vs_flm(model, proj, M, K, N, aie_context): competitors = [ setup_iron(M, K, N, A, B, aie_context), setup_flm(M, K, N, A, B, aie_context), - setup_gemm(M, K, N, A, B, aie_context), ] + # IRON's generic GEMM still has the 20-bit mega_row stride limitation that + # flm_gemm fixed: at M>256 with a 10240-wide N its C descriptor wants a + # stride of 2621440 elements and aiecc rejects the build outright. That is + # a property of that operator, not of the shape, and it is caught at + # compile time rather than hanging -- so report the shape with the + # competitors that do build instead of losing flm_gemm's own numbers for + # it. Only that specific rejection is tolerated; anything else still fails. + try: + competitors.append(setup_gemm(M, K, N, A, B, aie_context)) + except RuntimeError as e: + if "aie.dma_bd" not in str(e) or "exceeds the" not in str(e): + raise + print("gemm unavailable: descriptor stride exceeds the shim's 20-bit step") bad = [c for c in competitors if not c.verify(M, N, expected, mass)] assert not bad, "; ".join( @@ -354,7 +377,8 @@ def test_flm_gemm_vs_flm(model, proj, M, K, N, aie_context): print(f"{c.name} latency (us): {c.us:.1f}") print(f"{c.name} err/mass: {c.err:.3e}") print(f"speedup vs flm: {by_name['flm'].us / iron.us:.3f}") - print(f"speedup vs gemm: {by_name['gemm'].us / iron.us:.3f}") + if "gemm" in by_name: + print(f"speedup vs gemm: {by_name['gemm'].us / iron.us:.3f}") print(f"iron throughput: {2.0 * M * K * N / (iron.us * 1e-6) / 1e9:.6e} GFLOP/s") print(f"iron jitter (%): {iron.jitter_pct:.2f}") print(f"iron xclbin (KB): {iron.xclbin.stat().st_size / 1024:.1f}") From aecb0dbbde7c0074a9562c575bf9eafc90565821 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 14:06:57 -0600 Subject: [PATCH 28/31] deps: bump mlir-aie to 1.4.3.dev85 and Peano to 2026090701 Takes the Peano pin from mlir-aie's own utils/peano-requirements.txt at df48abc, so the two stay in step. 20/20 flm_gemm tests and all 30 sweep shapes pass, and accuracy is BIT-IDENTICAL on every shape (IronErr unchanged to 4 significant figures). Latency regresses slightly: median +1.5% to +2.4% across the 30 shapes, 25-27 of 30 slower. That is small but real, not the box's bimodality -- running the same binaries twice on the same toolchain gives median +0.00% with 15/30 slower, so the median-of-30 is a reproducible statistic even though per-shape noise is +/-7%. Not bisected to a commit; aie-normalize-dma-bd-dims (#3688) is the most plausible suspect of the 25, since it rewrites BD dimensions. Co-Authored-By: Claude --- requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/requirements.txt b/requirements.txt index d0bd8880d..7f7c43b37 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.2026090201+a36c62b9 +mlir_aie==1.4.3.dev85+gdf48abc +llvm-aie==22.0.0.2026090701+3e93bf7b black reuse From 52a21abd4d00680c035a47308d8da36c4999e9bc Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 14:22:53 -0600 Subject: [PATCH 29/31] gemm: split C's drain when its row stride overflows the shim's 20-bit step C's descriptor walks current_tb_n_rows row-blocks with an outermost stride of mem_tile_m_C * N. That dimension lands in the shim BD's iteration field, whose step is 20 bits, so a wide N overflows it: M=1024 K=2560 N=10240 needs 2621440 and aiecc rejected the build outright with "Stride 3 exceeds the [1:1048576] range". Both of those shapes are real E4B gate/up projections. Issue one descriptor per row-block in that case, carrying the row jump in the offset -- which has no such limit -- and leaving the outer dimension degenerate. Same bytes, same order, same object count; only the descriptor is reshaped, and only when the stride actually overflows. The extra tasks stay inside the two shim limits that neither the verifier nor the toolchain models. BD ids: an iteration's tasks all stay live until tg.finish(), so they remain distinct, at 2 x (2 C + 2 A + 2 B) = 12 of 16. Channel task queue: C goes from 2 outstanding to 4, which is where A and B already sat -- and 4 is the measured ceiling (see flm_gemm's design.py). Verified on hardware; err/mass is identical to the unsplit path, including a shape that does not need the split: M=1024 K=2560 N=10240 16366 us 2.932e-04 (was: build failure) M=2048 K=2560 N=10240 29972 us 2.933e-04 (was: build failure) M=1024 K=2560 N=4096 5824 us 2.934e-04 (control, unsplit) Both new shapes join test.py's extensive params. 26/26 pass. Co-Authored-By: Claude --- iron/operators/gemm/design.py | 104 ++++++++++++++++++++++++---------- iron/operators/gemm/test.py | 6 ++ 2 files changed, 80 insertions(+), 30 deletions(-) diff --git a/iron/operators/gemm/design.py b/iron/operators/gemm/design.py index c8a4f6e7c..fe2d529a4 100644 --- a/iron/operators/gemm/design.py +++ b/iron/operators/gemm/design.py @@ -172,6 +172,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" @@ -597,39 +608,72 @@ 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, - ) + # This line does not change MLIR output at all - it's just for recording data movement + C_taps.append(C_tile) + + 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: diff --git a/iron/operators/gemm/test.py b/iron/operators/gemm/test.py index bbd41b00a..732942c1a 100755 --- a/iron/operators/gemm/test.py +++ b/iron/operators/gemm/test.py @@ -49,6 +49,12 @@ 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), + # C's row stride (mem_tile_m_C * N) overflows the shim BD's 20-bit + # iteration step at these widths, so the drain is issued as one + # descriptor per row-block instead. Both failed to build at all before + # that: aiecc "Stride 3 exceeds the [1:1048576] range". + (1024, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), + (2048, 2560, 10240, 8, False, False, 64, 64, 64, 0, 1), ] # fmt: on From 1504a9acc5ce0188603e791f4b7eedec0f9634e7 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 14:38:57 -0600 Subject: [PATCH 30/31] Address Code Review (#196) flm.gemm: NPU1 support, rename from flm_gemm, ship the prebuilt overlay Squashed for a clean rebase onto port-flm-gemm's stride-overflow fixes (this branch's original 18 commits are preserved on backup-gemm-broadcast-review-pre-rebase). - Port flm_gemm to run on NPU1 (Phoenix), picking tile_n per device; fix an s/t conflation in b_recv_dims along the way. - Drop the ablation knobs and environment variables that were only needed during tuning. - Rename iron.operators.flm_gemm -> iron.operators.flm.GEMM, give the kernels descriptive names, and unify mm_fused into one mmul entry point with geometry from one source. - Let the framework build artifact names instead of hand-narrating them. - Add iron.operators.flm.MMPrebuilt, running FastFlowLM's shipped mm.xclbin as an operator (RemoteFileArtifact, pinned by SHA-256), and delete bench_vs_flm.py in favor of comparing against it directly. Mark that comparison extensive so the default run never fetches the overlay. - Audit the README and document the prebuilt operator. - Re-measure the NPU1 numbers and fix a benchmark bug they exposed. - Satisfy clang-format on the touched kernel files. - Address remaining Copilot review comments: fix MMPrebuilt.pack_B's wrong B ordering (verified against the deleted, previously-correct flm_pack_B), add mm_prebuilt/test.py for actual correctness coverage in extensive CI, xfail the 4 known-unsupported E4B shapes in benchmark.py, validate tile_n before indexing CT_MAX_K_FOR_N, and fix NPU1's stale sweep-boundary test shapes. Co-authored-by: Claude --- aie_kernels/aie2p/flm_gemm.cc | 117 ----- aie_kernels/aie2p/flm_gemm_epilogue.cc | 96 ---- aie_kernels/aie2p/flm_gemm_geometry.h | 39 -- aie_kernels/aie2p/flm_gemm_mmul.h | 210 --------- aie_kernels/aie2p/nonlut_based_ops.h | 66 --- aie_kernels/generic/activations.h | 100 +++++ aie_kernels/generic/mm_fused.cc | 113 +++++ aie_kernels/generic/mm_fused_epilogue.cc | 97 ++++ aie_kernels/generic/mm_fused_mmul.h | 204 +++++++++ aie_kernels/generic/passThrough.cc | 9 + iron/common/__init__.py | 2 + iron/common/compilation/__init__.py | 2 + iron/common/compilation/base.py | 63 +++ iron/common/context.py | 1 + iron/operators/__init__.py | 10 +- iron/operators/flm/__init__.py | 33 ++ iron/operators/flm/gemm/README.md | 347 +++++++++++++++ iron/operators/flm/gemm/benchmark.py | 278 ++++++++++++ .../{flm_gemm => flm/gemm}/design.py | 369 ++++++++------- iron/operators/flm/gemm/op.py | 419 ++++++++++++++++++ .../{flm_gemm => flm/gemm}/reference.py | 28 +- iron/operators/flm/gemm/test.py | 218 +++++++++ iron/operators/flm/mm_prebuilt/README.md | 64 +++ iron/operators/flm/mm_prebuilt/design.py | 192 ++++++++ iron/operators/flm/mm_prebuilt/op.py | 180 ++++++++ iron/operators/flm/mm_prebuilt/test.py | 81 ++++ iron/operators/flm/packing.py | 130 ++++++ iron/operators/flm_gemm/README.md | 251 ----------- iron/operators/flm_gemm/bench_vs_flm.py | 387 ---------------- iron/operators/flm_gemm/op.py | 397 ----------------- iron/operators/flm_gemm/test.py | 163 ------- requirements.txt | 3 + 32 files changed, 2770 insertions(+), 1899 deletions(-) delete mode 100644 aie_kernels/aie2p/flm_gemm.cc delete mode 100644 aie_kernels/aie2p/flm_gemm_epilogue.cc delete mode 100644 aie_kernels/aie2p/flm_gemm_geometry.h delete mode 100644 aie_kernels/aie2p/flm_gemm_mmul.h delete mode 100644 aie_kernels/aie2p/nonlut_based_ops.h create mode 100644 aie_kernels/generic/activations.h create mode 100644 aie_kernels/generic/mm_fused.cc create mode 100644 aie_kernels/generic/mm_fused_epilogue.cc create mode 100644 aie_kernels/generic/mm_fused_mmul.h create mode 100644 iron/operators/flm/__init__.py create mode 100644 iron/operators/flm/gemm/README.md create mode 100644 iron/operators/flm/gemm/benchmark.py rename iron/operators/{flm_gemm => flm/gemm}/design.py (74%) create mode 100644 iron/operators/flm/gemm/op.py rename iron/operators/{flm_gemm => flm/gemm}/reference.py (74%) create mode 100644 iron/operators/flm/gemm/test.py create mode 100644 iron/operators/flm/mm_prebuilt/README.md create mode 100644 iron/operators/flm/mm_prebuilt/design.py create mode 100644 iron/operators/flm/mm_prebuilt/op.py create mode 100644 iron/operators/flm/mm_prebuilt/test.py create mode 100644 iron/operators/flm/packing.py delete mode 100644 iron/operators/flm_gemm/README.md delete mode 100644 iron/operators/flm_gemm/bench_vs_flm.py delete mode 100644 iron/operators/flm_gemm/op.py delete mode 100644 iron/operators/flm_gemm/test.py diff --git a/aie_kernels/aie2p/flm_gemm.cc b/aie_kernels/aie2p/flm_gemm.cc deleted file mode 100644 index 6cd655c7d..000000000 --- a/aie_kernels/aie2p/flm_gemm.cc +++ /dev/null @@ -1,117 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -// Compute kernel for the flm_gemm operator: a bf16 GEMM over a fixed 4x8 grid -// of compute tiles, where each tile owns an m x n slice of C and accumulates -// over K in f32. -// -// Split into per-step entry points rather than one entry point owning the whole -// loop nest. The loop nest lives in the design's core body instead -// (iron/operators/flm_gemm/design.py), which is what gives each level of it an -// ObjectFifo acquire point -- a fifo needs its consumer to acquire once per -// object, so a single entry point spanning a whole dispatch could not be fed by -// one. -// -// Tile geometry comes from the design as -D flags so design.py stays the single -// source of truth; the design's own buffer sizes and unroll factors are derived -// from the same constants. -#include "flm_gemm_mmul.h" -#include "zero.cc" -#include -#include - -#if !defined(FLM_GEMM_TILE_M) || !defined(FLM_GEMM_TILE_K) || \ - !defined(FLM_GEMM_TILE_N) -#error "design.py must pass -DFLM_GEMM_TILE_M / _TILE_K / _TILE_N" -#endif - -namespace { -constexpr int M = FLM_GEMM_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 and C must live across the -// whole K reduction, so sizing both to M pays the peak cost twice. Defaults -// to M, which is the symmetric design. -#ifndef FLM_GEMM_TILE_MA -#define FLM_GEMM_TILE_MA FLM_GEMM_TILE_M -#endif -constexpr int MA = FLM_GEMM_TILE_MA; -static_assert(M % MA == 0, "tile_m must be a whole number of A bands"); -constexpr int K = FLM_GEMM_TILE_K; -constexpr int N = FLM_GEMM_TILE_N; -constexpr int R = 8; // register tiling r/s/t -constexpr int S = 8; -constexpr int T = 8; - -// How much of K one compute tile holds at a time, given the n width. -#ifdef FLM_GEMM_CT_K -constexpr int CT_K = FLM_GEMM_CT_K; -#else -constexpr int CT_K = compute_CT_k_max_n(); -#endif -static_assert(CT_K > 0, "no K-blocking geometry for this tile_n"); - -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 with rounding_mode::floor. Truncation biases every -// operand conversion the same direction, so the error accumulates coherently -// over the K reduction instead of cancelling -- measured as a ~1% bias in the -// result, ~20x worse than round-to-nearest-even, which is far more than the -// bfp16 emulation itself costs. Set it explicitly in every entry point that -// converts (the mmul here, and the f32->bf16 store in the epilogue). -#ifdef FLM_GEMM_ROUND_FLOOR -constexpr aie::rounding_mode round_mode = aie::rounding_mode::floor; -#else -constexpr aie::rounding_mode round_mode = aie::rounding_mode::conv_even; -#endif -} // namespace - -extern "C" { - -// Zero the f32 accumulator. Called once per mega-block-row, before the k loop -// starts accumulating into it. -// -// ADD_BIAS is deliberately not supported: initialising the accumulator from a -// bias vector would mean consuming an extra object through the same handshake -// the B ObjectFifo now owns, which would desynchronise that fifo and hang -// rather than silently mis-compute. -void flm_gemm_acc_init(float *y_acc) { - // zero_vectorized brackets itself in event0/event1 for tracing. - zero_vectorized(y_acc); -} - -// One l-step of a k iteration: one B chunk multiplied against one A object, -// accumulated into y_acc. -// -// The l loop lives in the core body so that each B chunk gets its own acquire -// point. A is a single object spanning every z slice of the mmul, so this takes -// no locks -- the A and B fifos own that handshake. -#ifdef FLM_GEMM_BFP16_B -void flm_gemm_k_step(bfloat16 *a_buf, bfp16ebs8 *b_buf, float *y_acc, -#else -void flm_gemm_k_step(bfloat16 *a_buf, bfloat16 *b_buf, float *y_acc, -#endif - int32_t band) { -#ifdef FLM_GEMM_NULL_MMUL - // ABLATION ONLY: skip the multiply, keep every acquire, release and DMA. - // Output is garbage; never correctness-gate a build with this. - (void)a_buf; (void)b_buf; (void)y_acc; (void)band; - return; -#endif - ::aie::set_rounding(round_mode); - constexpr int NUM_ITER = K / CT_K; - // 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. -#ifdef FLM_GEMM_BFP16_B - flm_gemm_mmul_2x2_bfpb(a_buf, b_buf, y_acc + band * (MA * N)); -#else - flm_gemm_mmul_2x2( - a_buf, b_buf, y_acc + band * (MA * N)); -#endif -} -} diff --git a/aie_kernels/aie2p/flm_gemm_epilogue.cc b/aie_kernels/aie2p/flm_gemm_epilogue.cc deleted file mode 100644 index 5fee55c47..000000000 --- a/aie_kernels/aie2p/flm_gemm_epilogue.cc +++ /dev/null @@ -1,96 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -// flm_gemm's output stage: convert one chunk of the f32 accumulator into a bf16 -// output 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 every instantiation of this kernel -// has a branch-free inner loop. That differs from the design this was ported -// from, where one overlay served every activation and had to test a runtime -// mode word per chunk. -// -// This is a separate translation unit from flm_gemm.cc so it can be built with -// its own flags and, if the per-call overhead ever shows up in a trace, be -// switched to an inlined LLVM-IR kernel independently of the much larger mmul. -#include "../aie_kernel_utils.h" -#include "nonlut_based_ops.h" -#include -#include - -#if !defined(FLM_GEMM_OUT_CHUNK) || !defined(FLM_GEMM_C_DEPTH) -#error "design.py must pass -DFLM_GEMM_OUT_CHUNK / -DFLM_GEMM_C_DEPTH" -#endif - -// 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid -#ifndef FLM_GEMM_EPILOGUE_MODE -#define FLM_GEMM_EPILOGUE_MODE 0 -#endif -#ifndef FLM_GEMM_CLAMP -#define FLM_GEMM_CLAMP 0 -#endif -#ifndef FLM_GEMM_CLAMP_MIN -#define FLM_GEMM_CLAMP_MIN 0.0f -#endif -#ifndef FLM_GEMM_CLAMP_MAX -#define FLM_GEMM_CLAMP_MAX 0.0f -#endif - -namespace { -constexpr int CHUNK = FLM_GEMM_OUT_CHUNK; -constexpr int DEPTH = FLM_GEMM_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"); -} // namespace - -extern "C" { - -// Chunk (outer * DEPTH + half) of the accumulator -> one C object. -// -// 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 flm_gemm_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, - int32_t half) { - // The f32 -> bf16 store below is a conversion, so it depends on the rounding - // mode just as the mmul does; the core default is floor. See flm_gemm.cc. -#ifdef FLM_GEMM_ROUND_FLOOR - ::aie::set_rounding(aie::rounding_mode::floor); -#else - ::aie::set_rounding(aie::rounding_mode::conv_even); -#endif - const float *__restrict src = y_acc + (outer * DEPTH + half) * CHUNK; - -#if FLM_GEMM_CLAMP - const aie::vector lo = - aie::broadcast(static_cast(FLM_GEMM_CLAMP_MIN)); - const aie::vector hi = - aie::broadcast(static_cast(FLM_GEMM_CLAMP_MAX)); -#endif - - AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) - for (int j = 0; j < CHUNK / V; j++) { - aie::accum acc; - acc.from_vector(aie::load_v(src + j * V)); - // The assignment is the conversion: to_v16bfloat16 yields a raw - // v16bfloat16, not an aie::vector. - aie::vector v = to_v16bfloat16(acc); -#if FLM_GEMM_EPILOGUE_MODE == 1 - v = getGeluBf16_nonLUT(v); -#elif FLM_GEMM_EPILOGUE_MODE == 2 - v = getSiluBf16_nonLUT(v); -#elif FLM_GEMM_EPILOGUE_MODE == 3 - v = getSigmoidBf16_nonLUT(v); -#endif -#if FLM_GEMM_CLAMP - v = aie::clamp(v, lo, hi); -#endif - aie::store_v(y_out + j * V, v); - } -} -} diff --git a/aie_kernels/aie2p/flm_gemm_geometry.h b/aie_kernels/aie2p/flm_gemm_geometry.h deleted file mode 100644 index 3b0b46e21..000000000 --- a/aie_kernels/aie2p/flm_gemm_geometry.h +++ /dev/null @@ -1,39 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -// K-blocking geometry for the flm_gemm design: how much of K one compute tile -// holds at a time, as a function of the n tile width. flm_gemm.cc reads this to -// size its k loop; flm_gemm_mmul.h is the mmul that consumes the result. -// -// The values are a fixed L1 budget split two ways: a wider n tile leaves less -// room for the k slice of B, so the product stays roughly constant. n=128 (the -// only width flm_gemm currently builds) caps the core's k slice at 32. -#ifndef __FLM_GEMM_GEOMETRY_H__ -#define __FLM_GEMM_GEOMETRY_H__ -#include -#include - -constexpr int CT_k_max_n_16 = 16; -constexpr int CT_k_max_n_32 = 32; -constexpr int CT_k_max_n_64 = 128; -constexpr int CT_k_max_n_128 = 32; -constexpr int CT_k_max_n_256 = 16; - -template -constexpr int compute_CT_k_max_n() { - if constexpr (N == 16) { - return CT_k_max_n_16; - } else if constexpr (N == 32) { - return CT_k_max_n_32; - } else if constexpr (N == 64) { - return CT_k_max_n_64; - } else if constexpr (N == 128) { - return CT_k_max_n_128; - } else if constexpr (N == 256) { - return CT_k_max_n_256; - } else { - return -1; - } -} - -#endif // __FLM_GEMM_GEOMETRY_H__ diff --git a/aie_kernels/aie2p/flm_gemm_mmul.h b/aie_kernels/aie2p/flm_gemm_mmul.h deleted file mode 100644 index 44d7c4fa5..000000000 --- a/aie_kernels/aie2p/flm_gemm_mmul.h +++ /dev/null @@ -1,210 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -// The 2x2 mmul at the heart of flm_gemm, in the "one-buffer A" form its A -// ObjectFifo leg requires. -// -// A arrives as a SINGLE ObjectFifo object spanning every z slice, rather than a -// ping/pong pair the kernel locks for itself. The fifo owns that handshake, so -// there are no acquire/release pairs in here at all, and the core body acquires -// exactly one A object per call to this function. -// -// Each iteration of the j loop keeps four MMUL accumulators live (C00/C01/C10/ -// C11) across a 2x2 block of output tiles, so each pair of A loads and each -// pair of B loads feeds four macs. That is what keeps the vector unit busy; -// dropping to a 1x1 tile would halve the arithmetic per load. -#ifndef __FLM_GEMM_MMUL_H__ -#define __FLM_GEMM_MMUL_H__ -#include "../aie_kernel_utils.h" -#include "flm_gemm_geometry.h" -#include - -// rowA/colA/colB count r x s (A), s x t (B) and r x t (C) sub-tiles, not -// elements. b_row_maj / is_b_s_t_in_row_major select B's in-L1 layout; flm_gemm -// always instantiates (B column-major with row-major s x t -// sub-blocks), which is the layout its memtile forward() produces. The other -// combinations are kept because they are `if constexpr` and cost nothing, but -// they are untested here. -template -__aie_inline void -flm_gemm_mmul_2x2(const T_in *__restrict pA, const T_in *__restrict pB, - T_out *__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) { - T_out *__restrict pC1 = pC + (z * colB) * MMUL::size_C; - T_out *__restrict pC2 = pC + ((z + 1) * colB) * MMUL::size_C; - - // A is one object spanning every z slice; index it directly. - const T_in *__restrict pA_cur_buf = 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 T_in *__restrict pA1 = pA_cur_buf; - const T_in *__restrict pA2 = pA_cur_buf + colA * MMUL::size_A; - const T_in *__restrict pB1; - const T_in *__restrict pB2; - if constexpr (b_row_maj) { - pB1 = pB + (j)*MMUL::size_B; - pB2 = pB + (j + 1) * MMUL::size_B; - } else { - pB1 = pB + (j * colA) * MMUL::size_B; - 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, deliberately. An earlier 2x hand-unroll for the Peano path - // is now a pessimization: with colA/2 = 4 trip counts there are too - // few iterations to amortize the software pipeline's fill/drain. - // Measured cycles per mmul call (HW trace, event0/event1): rolled - // 1875, hand-unrolled 2x 2446, compiler unroll 4/8 3291/3221. Any - // extra live state across this loop loses more than it gains -- - // Peano's register allocator is pipelining-unaware and manufactures - // false loop-carried anti-deps (llvm-aie#1066), so keep the body - // minimal and let the pipeliner overlap the iterations. - 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; - - if constexpr (b_row_maj) { - B0 = aie::load_v(pB1); - pB1 += MMUL::size_B * colB; - B1 = aie::load_v(pB2); - pB2 += MMUL::size_B * colB; - } else { - if constexpr (is_b_s_t_in_row_major == false) { - B0 = aie::transpose(aie::load_v(pB1), t, s); - pB1 += MMUL::size_B; - B1 = aie::transpose(aie::load_v(pB2), t, s); - pB2 += MMUL::size_B; - } else { - 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); - } - aie::store_v(pC1, C00.template to_vector()); - pC1 += MMUL::size_C; - aie::store_v(pC1, C01.template to_vector()); - pC1 += MMUL::size_C; - aie::store_v(pC2, C10.template to_vector()); - pC2 += MMUL::size_C; - aie::store_v(pC2, C11.template to_vector()); - pC2 += MMUL::size_C; - } - } - - event1(); -} - -// The same 2x2 mmul, but B arrives ALREADY in bfp16ebs8 rather than bf16. -// -// The bf16 form converts B inside every mac -- transpose, widen, then -// to_v64bfp16ebs8 -- purely to feed hardware that only multiplies bfp16. 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, it -// does not add one. It also makes B 9 bytes per 8 elements instead of 16, -// which is why it is worth doing at all -- the operator is DMA-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, not -// blocks (llvm-aie#1232). The stream sidesteps both. -template -__aie_inline void flm_gemm_mmul_2x2_bfpb(const bfloat16 *__restrict pA, - const bfp16ebs8 *__restrict pB, - T_out *__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) { - T_out *__restrict pC1 = pC + (z * colB) * sizeC; - T_out *__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; - - // Rolled for the same reason as the bf16 form: extra live state across - // this loop loses more than it gains (llvm-aie#1066). - 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, exactly as - // mlir-aie's mm_bfp_mixed.cc 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 codegens. - 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); - } - aie::store_v(pC1, C00.template to_vector()); - pC1 += sizeC; - aie::store_v(pC1, C01.template to_vector()); - pC1 += sizeC; - aie::store_v(pC2, C10.template to_vector()); - pC2 += sizeC; - aie::store_v(pC2, C11.template to_vector()); - pC2 += sizeC; - } - } - event1(); -} - -#endif // __FLM_GEMM_MMUL_H__ diff --git a/aie_kernels/aie2p/nonlut_based_ops.h b/aie_kernels/aie2p/nonlut_based_ops.h deleted file mode 100644 index 5ec591513..000000000 --- a/aie_kernels/aie2p/nonlut_based_ops.h +++ /dev/null @@ -1,66 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -// Branch-free, LUT-free activations over an aie::vector. -// -// These are templated on the vector width and return a vector, so they compose -// inside an existing vector loop -- unlike aie_kernels/aie2p/{gelu,silu}.cc, -// which are whole-buffer entry points over a fixed 32-wide vector. flm_gemm's -// epilogue needs the former: it applies the activation to the same 16-wide -// vector it just converted from the f32 accumulator, without a second pass over -// L1. -// -// All three are built on tanh, which AIE2P has as a native vector op, 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. -#ifndef __NONLUT_BASED_OPS_H__ -#define __NONLUT_BASED_OPS_H__ -#include - -// sigmoid(x) = (tanh(x/2) + 1) / 2 -template -aie::vector -getSigmoidBf16_nonLUT(aie::vector x) { - const bfloat16 half = 0.5f; - const bfloat16 one = 1.0f; - aie::vector v_half = - aie::broadcast(half); - aie::vector v_one = - aie::broadcast(one); - aie::accum x_mul_half = aie::mul(x, v_half); - aie::vector tanh_x_half = - aie::tanh(x_mul_half.template to_vector()); - - aie::vector tanh_x_half_plus_one = - aie::add(tanh_x_half, v_one); - return aie::mul(tanh_x_half_plus_one, v_half); -} - -// silu(x) = x * sigmoid(x) -template -aie::vector -getSiluBf16_nonLUT(aie::vector x) { - aie::vector sigmoid_x = getSigmoidBf16_nonLUT(x); - return aie::mul(x, sigmoid_x); -} - -// gelu(x) ~= x * sigmoid(1.702x) -template -aie::vector -getGeluBf16_nonLUT(aie::vector x) { - constexpr bfloat16 x_scale = 1.702; - aie::vector v_x_scale = - aie::broadcast(x_scale); - aie::vector x_scaled = aie::mul(x, v_x_scale); - aie::vector sigmoid_x_scaled = - getSigmoidBf16_nonLUT(x_scaled); - return aie::mul(x, sigmoid_x_scaled); -} - -#endif // __NONLUT_BASED_OPS_H__ diff --git a/aie_kernels/generic/activations.h b/aie_kernels/generic/activations.h new file mode 100644 index 000000000..9421ac5ff --- /dev/null +++ b/aie_kernels/generic/activations.h @@ -0,0 +1,100 @@ +// 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 same 16-wide +// vector it just converted from the f32 accumulator, 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. +// +// This header was called nonlut_based_ops.h when mm_fused was AIE2P-only, which +// stopped being an accurate name once AIE2 brought in the LUT; it is mm_fused's +// only consumer, hence the rename rather than a second copy. +#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 + +// tanh of an f32 accumulator, on whichever path this architecture has. +template +__attribute__((always_inline)) aie::vector tanh_vec(aie::accum x) +{ +#if ACTIVATIONS_NATIVE_TANH + return aie::tanh(x.template to_vector()); +#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"); + return getTanhBf16(x.template to_vector()); +#endif +} + +// sigmoid(x) = (tanh(x/2) + 1) / 2 +template aie::vector sigmoid_vec(aie::vector x) +{ + const bfloat16 half = 0.5f; + const bfloat16 one = 1.0f; + aie::vector v_half = aie::broadcast(half); + aie::vector v_one = aie::broadcast(one); + aie::accum x_mul_half = aie::mul(x, v_half); + aie::vector tanh_x_half = tanh_vec(x_mul_half); + + aie::vector tanh_x_half_plus_one = aie::add(tanh_x_half, v_one); + return aie::mul(tanh_x_half_plus_one, v_half); +} + +// silu(x) = x * sigmoid(x) +template aie::vector silu_vec(aie::vector x) +{ + aie::vector sigmoid_x = sigmoid_vec(x); + return aie::mul(x, sigmoid_x); +} + +// gelu(x) ~= x * sigmoid(1.702x) +template aie::vector gelu_vec(aie::vector x) +{ + constexpr bfloat16 x_scale = 1.702; + aie::vector v_x_scale = aie::broadcast(x_scale); + aie::vector x_scaled = aie::mul(x, v_x_scale); + aie::vector sigmoid_x_scaled = sigmoid_vec(x_scaled); + return aie::mul(x, sigmoid_x_scaled); +} + +#endif // __ACTIVATIONS_H__ diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc new file mode 100644 index 000000000..44e7d3805 --- /dev/null +++ b/aie_kernels/generic/mm_fused.cc @@ -0,0 +1,113 @@ +// 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. +// +// Two 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 the accumulator +// +// 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. The +// matching output stage is mm_fused_epilogue.cc. +// +// 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 "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 + +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; + +// 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 below, and the f32->bf16 store in mm_fused_epilogue.cc. +// +// 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)); +} +} diff --git a/aie_kernels/generic/mm_fused_epilogue.cc b/aie_kernels/generic/mm_fused_epilogue.cc new file mode 100644 index 000000000..8a1f5f518 --- /dev/null +++ b/aie_kernels/generic/mm_fused_epilogue.cc @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// mm_fused's output stage: convert one chunk of the f32 accumulator into a bf16 +// output 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 every instantiation of this kernel +// has a branch-free inner loop. That differs from the design this was ported +// from, where one overlay served every activation and had to test a runtime +// mode word per chunk. +// +// This is a separate translation unit from mm_fused.cc so it can be built with +// its own flags and, if the per-call overhead ever shows up in a trace, be +// switched to an inlined LLVM-IR kernel independently of the much larger mmul. +#include "../aie_kernel_utils.h" +#include "activations.h" + +#include +#include + +#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 + +// 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid +#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 CHUNK = MM_FUSED_OUT_CHUNK; +constexpr int 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"); +} // namespace + +extern "C" { + +// Chunk (outer * DEPTH + half) of the accumulator -> one C object. +// +// 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 f32 -> bf16 store below is a conversion, so it depends on the rounding + // mode just as the mmul does, and must agree with it. Same flag, same + // polarity: see mm_fused.cc. +#ifdef ROUND_CONV_EVEN + ::aie::set_rounding(aie::rounding_mode::conv_even); +#else + ::aie::set_rounding(aie::rounding_mode::floor); +#endif + const float *__restrict src = y_acc + (outer * DEPTH + half) * CHUNK; + +#if MM_FUSED_CLAMP + const aie::vector lo = aie::broadcast(static_cast(MM_FUSED_CLAMP_MIN)); + const aie::vector hi = aie::broadcast(static_cast(MM_FUSED_CLAMP_MAX)); +#endif + + AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V) + for (int j = 0; j < CHUNK / V; j++) { + aie::accum acc; + acc.from_vector(aie::load_v(src + j * V)); + // The assignment is the conversion: to_v16bfloat16 yields a raw + // v16bfloat16, not an aie::vector. + aie::vector v = to_v16bfloat16(acc); +#if MM_FUSED_EPILOGUE_MODE == 1 + v = gelu_vec(v); +#elif MM_FUSED_EPILOGUE_MODE == 2 + v = silu_vec(v); +#elif MM_FUSED_EPILOGUE_MODE == 3 + v = sigmoid_vec(v); +#endif +#if MM_FUSED_CLAMP + v = aie::clamp(v, lo, hi); +#endif + 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 000000000..28645db65 --- /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 f4a784de5..e17cf5a62 100644 --- a/aie_kernels/generic/passThrough.cc +++ b/aie_kernels/generic/passThrough.cc @@ -10,6 +10,15 @@ #include #include +// Element width in bits, chosen by the caller with -DBIT_WIDTH. Only mha passes +// it (16); every other user wants the 32-bit form and used to reach it by +// leaving the macro undefined, which the preprocessor evaluates as 0 and so +// falls through to the #else below. Peano now compiles with -Werror=undef, so +// that default has to be written down rather than relied on. +#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 cb2ff31be..f448a6d65 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 d4e06c2e6..c1fb11855 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 af06dd128..487bb7583 100644 --- a/iron/common/compilation/base.py +++ b/iron/common/compilation/base.py @@ -37,8 +37,10 @@ from collections import deque from collections.abc import Iterator, Sequence from pathlib import Path +import hashlib import os.path import shutil +import urllib.request import zlib import logging import subprocess @@ -408,6 +410,33 @@ def __init__( super().__init__(filename, dependencies=[SourceArtifact(generator.source_path)]) +def _sha256_of(path: Path) -> str: + digest = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + digest.update(chunk) + return digest.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 +514,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 a7a5136ca..c2e4b2e84 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/operators/__init__.py b/iron/operators/__init__.py index 7555d1cb9..744278f69 100644 --- a/iron/operators/__init__.py +++ b/iron/operators/__init__.py @@ -13,7 +13,6 @@ _OPERATOR_MODULES = { "ElementwiseAdd": "elementwise_add", "ElementwiseMul": "elementwise_mul", - "FLMGEMM": "flm_gemm", "GEMM": "gemm", "GEMV": "gemv", "MHA": "mha", @@ -28,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}") @@ -40,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 000000000..3bc3767ae --- /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 000000000..83bdd63e8 --- /dev/null +++ b/iron/operators/flm/gemm/README.md @@ -0,0 +1,347 @@ + + +# `iron.operators.flm.GEMM` — bf16 GEMM with a fused epilogue + +```python +from iron.operators.flm import GEMM + +op = GEMM(M=1024, K=1536, N=6144, epilogue="silu", context=ctx) +op.compile() +op.get_callable()(A, op.pack_B(B), C_out) +``` + +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`. + +## 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.** bfp16 emulation drops mantissa bits; + native bf16 macs accumulating in f32 do not. Measured mean |err| against the + accumulated mass is under 1e-6 on NPU1 with `conv_even` versus 0.00042 on + NPU2, and 0.00015 versus 0.0099 with `floor`. `test.py` sets the budget per + architecture — inheriting NPU2's would leave ~70x of slack. +* **`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="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: + +| | err/mass | vs shipped | +|---|---|---| +| `rounding="floor"` | 0.009867 | **bit-identical, 6291456/6291456 elements** | +| `rounding="conv_even"` (default) | 0.000241 | differs everywhere | + +All four epilogues are bit-identical to the shipped kernel too, with `floor`. +The shipped kernel selects its activation from a runtime parameter, one overlay +serving every projection; this operator bakes it in at compile time instead, +with the same 0/1/2/3 mapping, which is what lets its inner loop be branch-free: + +| epilogue | vs shipped `output_mode` | +|---|---| +| `none` / `gelu` / `silu` / `sigmoid` | bit-identical, 1048576/1048576 each | + +**The default is `conv_even`, not `floor`.** 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 original. + +`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 expectations + +**Accuracy is architecture-dependent**, because the same 8x8x8 mmul lowers +differently: NPU2 uses two bfp16-emulated macs, which drop mantissa bits, while +NPU1 uses four native 4x8x4 bf16 macs, which do not. See +[Architectures](#architectures). `test.py` sets the budget per architecture. + +A pure elementwise *relative* tolerance is not meaningful here: 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 instead, +as `test.py` does. + +Reference points on NPU2, random signed A / non-negative B: + +| | mean err / mass | +|---|---| +| `flm.GEMM` (default) | 0.000241 | +| `iron.operators.GEMM`, same mode (`emulate=True, prio_accuracy=True`) | 0.000241 | +| `iron.operators.GEMM`, bf16 accumulator (`prio_accuracy=False`) | 0.000445 | +| `iron.operators.GEMM`, exact r=4 path (`emulate=False`) | 0.00007 | + +So on NPU2 this operator and `iron.operators.GEMM` in the same mode are +numerically **indistinguishable** — identical mean error, signed bias and +maximum, at both `tile_n` values. Same mmul shape, same emulation, same f32 +accumulation, same rounding. 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. Measured on NPU2 against the current design (rolled +mmul, resident B, ATB, bfp16 B), min of per-run medians over 6 rounds with the +two `tile_n` builds interleaved round-robin -- that box is bimodal ~6%, so +running all of one and then all of the other measures drift rather than design: + +| M / K / N | k_iters | `tile_n=64` | `tile_n=128` | +|---|---|---|---| +| 1024 / 512 / 4096 | 1 | 514 us | **498 us** | +| 1024 / 1024 / 4096 | 2 | **591 us** | 931 us | +| 1024 / 1536 / 6144 | 3 | **1141 us** | 1960 us | +| 1024 / 2560 / 4096 | 5 | **1239 us** | 1940 us | +| 2048 / 2048 / 2048 | 4 | **915 us** | 1581 us | +| 256 / 4096 / 1024 | 8 | **227 us** | 265 us | + +So `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 a wide margin. Measured on Phoenix, min of 8 interleaved rounds +of 20 dispatches: + +| M / K / N | k_iters | `tile_n=64` | `tile_n=128` | +|---|---|---|---| +| 512 / 512 / 1024 | 1 | **414 us** | 603 us | +| 512 / 1024 / 1024 | 2 | **618 us** | 1119 us | +| 512 / 1536 / 1536 | 3 | **1503 us** | 2066 us | + +## 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 — an earlier +split of wall time put it at roughly 40%, with a ceiling near 1.65x if compute +were free. So on NPU1 both a faster mmul and less traffic pay off, where on +NPU2 only the latter does. + +That split was measured on the previous L1 configuration and has not been +re-measured since; treat the 40% as indicative. It also cannot be reproduced as +written, because it relied on an ablation that nulled the mmul, and that knob +has been removed from the operator. + +### Measured dead ends on NPU1 + +Recorded so they are not retried. Both were plausible and both lost: + +* **Resident B does nothing here.** Measured off-versus-on at M=2048, the shape + where NPU2 gains 7.8%: 1259/1269 us against 1262/1268 (min/median) at K=512, + and 2238/2248 against 2221/2229 at K=1024 -- i.e. at best a no-op, marginally + negative at K=1024, against a 1.4-4.4% round spread. The `repeat_count` + BD-chain restart noted under [Resident B](#resident-b) eats ~77% of the win on + NPU2; on NPU1 it appears to eat all of it. Making B resident + *single*-buffered to reach K=2560 is worse + still, 5-12%, because it gives up the next column-block's prefetch. +* **The native 4x8x4 mmul shape is 22-30% slower.** See `register_tiling` in + `design.py`. Composing 8x8x8 out of native macs costs 2.5 `vshuffle` per + `vmac` and 4x8x4 costs zero, but removing every shuffle made it slower: the + kernel is load-port bound, not shuffle bound, and the narrower shape needs 53% + more loads per unit work because the 2x2 register block amortizes each load + over a quarter as much arithmetic. + +### 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; it grows with the height of the +problem, since B's re-reads scale with `m_row_blocks`. On NPU1 it is neither — +see [the dead ends above](#measured-dead-ends-on-npu1). At K=1024 N=4096 on +NPU2: + +| M | row-blocks | non-resident | resident | | +|---|---|---|---|---| +| 512 | 2 | 470.8 us | 468.5 us | 0.5% | +| 1024 | 4 | 860.4 us | 846.5 us | 1.6% | +| 2048 | 8 | 1760.1 us | **1622.8 us** | 7.8% | + +Do not evaluate this at small M: at M=512 the effect is inside the noise. + +Most of the available win is still on the table. With the mmul nulled, the +non-resident floor at M=2048 is 1739 us -- 118 MB at 68 GB/s, against a +memcpy-measured 63-70 GB/s roof for mixed read/write traffic -- and residency +drops that floor to 1135 us. Only 137 us of those 604 us reaches the full +build; the rest goes to `repeat_count` restarting the memtile BD chain at every +replay boundary. Closing that is the largest known remaining lever here. diff --git a/iron/operators/flm/gemm/benchmark.py b/iron/operators/flm/gemm/benchmark.py new file mode 100644 index 000000000..f01b5057f --- /dev/null +++ b/iron/operators/flm/gemm/benchmark.py @@ -0,0 +1,278 @@ +#!/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. + +Three implementations run per shape, on identical inputs: + + flm :class:`iron.operators.flm.GEMM`, the port + prebuilt :class:`iron.operators.flm.MMPrebuilt`, FastFlowLM's shipped + ``mm.xclbin``, downloaded and pinned by digest + 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 + +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. + +NPU2 only, because the shipped overlay is an 8-column NPU2 binary. + +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 +import time +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() +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, + ) + +# Every projection of both Gemma4 variants FastFlowLM ships, at three prefill +# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. +# 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 = [] + + already_built = self._is_built(op, ctx) + t0 = time.perf_counter() + op.compile() + # Only meaningful on a genuine miss; on a hit compile() returns in + # milliseconds and reporting that as a build time would be a lie. + self.compile_s = None if already_built else time.perf_counter() - t0 + + 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) + + @staticmethod + def _is_built(op, ctx): + if not op.artifacts: + op.set_up_artifacts() + return Path(op.xclbin_artifact.filename).is_file() + + 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): + ts = [] + for _ in range(ITERS): + t0 = time.perf_counter() + self.run() + ts.append((time.perf_counter() - t0) * 1e6) + 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\.]+)", + FLMCompileTime=r"flm compile \(s\): (?P[\d\.]+)", + GEMMCompileTime=r"gemm compile \(s\): (?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( + "prebuilt", + MMPrebuilt(M=M, K=K, N=N, context=aie_context), + A, + B, + M, + N, + BUDGET_FLOOR, + aie_context, + ), + Candidate( + "gemm", + IronGEMM(M=M, K=K, N=N, context=aie_context), + A, + B, + M, + N, + BUDGET_CONV_EVEN, + 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 c.compile_s is not None: + print(f"{c.name} compile (s): {c.compile_s:.1f}") + 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 similarity index 74% rename from iron/operators/flm_gemm/design.py rename to iron/operators/flm/gemm/design.py index eacb0b47c..d3c632eb9 100644 --- a/iron/operators/flm_gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -1,29 +1,39 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""bf16 GEMM over a fixed 4x8 compute-tile grid. - -This is a different design from ``iron.operators.gemm``, not a retuning of it. -The distinguishing choices, all of which the kernel's L1 layout depends on: - - * **A is broadcast along each compute row.** Four shim tiles (columns 0/2/4/6) - each feed one row of the grid, and every one of the 8 tiles in that row - consumes the same A object. B is broadcast down each column. So an A tile is - fetched once per row rather than once per tile. - * **C is joined at the memtile.** Each of the 4 tiles in a column writes its - own 64x128 slice into one memtile buffer, which drains to DDR as a single - 256x128 block, rather than each tile draining separately. - * **The mmul keeps A in a single ObjectFifo object** spanning every z slice - (``flm_gemm_mmul.h``), instead of a ping/pong pair the kernel locks itself. - * **The epilogue is fused**: 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. - -Geometry is fixed (m/k/n = 64/512/128, r/s/t = 8/8/8, 4x8 grid). 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. +"""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: + + * **Fixed tiling.** m/k/n = 64/512/128 and r/s/t = 8/8/8, rather than + parameterised tiles. 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. On NPU2 it is also + quantized to bfp16ebs8. + * **Asymmetric tile buffering**, so the A tile and the accumulator need not + share a height. + +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 + import numpy as np from ml_dtypes import bfloat16 @@ -40,7 +50,8 @@ Worker, ) from aie.iron.controlflow import range_ -from aie.iron.device import Tile +from aie.iron.device import NPU1, NPU2, Tile +from iron.common.device_utils import get_kernel_dir from iron.operators._trace import maybe_enable_trace # --- Fixed geometry ------------------------------------------------------- @@ -53,20 +64,63 @@ # README.md for the measured sweep, including the small-K shape where it # loses. N_TILE_DEFAULT = 64 -# k-slice per n width; must match compute_CT_k_max_n in flm_gemm_geometry.h +# 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. CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32, 256: 16} +# Register tiling. 8/8/8 on both architectures today; register_tiling() is the +# single source of truth and returns exactly these. R, S, T = 8, 8, 8 -ROWS, COLS = 4, 8 +ROWS = 4 +# Widest grid this design supports, i.e. NPU2. The actual width comes from the +# device (see grid_cols); this is only what op.py uses to name artifacts. +MAX_COLS = 8 + + +def grid_cols(dev): + """Grid width: 8 on NPU2 (Strix/Krackan), 4 on NPU1 (Phoenix).""" + return min(dev.cols, MAX_COLS) + + +def register_tiling(dev_name): + """The mmul's r/s/t, as (r, s, t). 8/8/8 on both architectures. + + A function rather than a bare constant because r/t is the natural thing to + retune per device, and because getting it wrong is silent: these set the + blocked L1 layout, so ``pack_B``, the four stream-dimension lists below and + ``gather_dims`` all key off them. + + Matching AIE2's native 4x8x4 mac shape here is a dead end -- it measures + 22-30% slower. The kernel is load-port bound, not 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. + """ + return (R, S, T) + + +def a_source_cols(cols): + """Which shim column sources the A broadcast for each compute row. + + On an 8-column grid the four A streams go to alternate columns, so each gets + its own shim MM2S path and never contends with a B fill; the existing gemm + operator pins A the same way. A 4-column grid has no such slack -- every + column must source one A row AND one B column AND drain C, which is 2 MM2S + + 1 S2MM, exactly saturating a shim tile's channels. + """ + return [2 * r for r in range(ROWS)] if cols >= 2 * ROWS else list(range(ROWS)) -# Which shim column sources the A broadcast for each compute row. Spreading -# them over alternate columns keeps four independent MM2S paths; the existing -# gemm operator pins A the same way in the 8-column case. -A_SOURCE_COL = [0, 2, 4, 6] 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 STACK_SIZE = 4096 # Usable L1 per compute tile: 64 KB less the stack and a little slack. @@ -75,7 +129,7 @@ EPILOGUE_MODES = {"none": 0, "gelu": 1, "silu": 2, "sigmoid": 3} # 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 = "flm_gemm_epilogue_chunk" +EPILOGUE_SYMBOL = "mm_fused_epilogue_chunk" # Minimum problem size, i.e. one pass of the whole grid. MIN_N depends on the # chosen n tile, so it is computed per call. @@ -83,10 +137,24 @@ MIN_K = K_TILE # 512 -def _bfp16_bytes(elems): - """bfp16ebs8 packs 8 values as 8 mantissa bytes plus one shared exponent.""" - assert elems % 8 == 0 - return elems // 8 * 9 +# B is bfp16ebs8 on AIE2P and bf16 on AIE2 -- the scalar BFP types exist only +# on AIE2P (see GEMM._bfp16_b). BFP16_GROUP is how many B values one element +# of the MLIR type holds: v8bfp16ebs8 holds 8, bf16 holds 1. +BFP16_GROUP = 8 + + +def bfp16_b_for(dev): + """Whether B is stored as bfp16ebs8 for this device. AIE2P only.""" + return get_kernel_dir(dev) == "aie2p" + + +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 * 9 # Shim-tile DMA BD step field is 20 bits wide (AIE2p; see mlir-aie's @@ -113,9 +181,12 @@ def _hw_stride_ok(stride_elems): return hw_stride <= (1 << _SHIM_STEP_BITS) - 1 -def _default_l1(n_tile, ct_max_k): +def _default_l1(n_tile, ct_max_k, b_elem_bytes): """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, so the + L1 budget below reflects what B actually costs on this device. + 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. @@ -130,7 +201,7 @@ def _default_l1(n_tile, ct_max_k): acc = M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH for b_depth in (B_DEPTH, 1): - b = _bfp16_bytes(ct_max_k * n_tile) * b_depth + 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 @@ -140,7 +211,7 @@ def _default_l1(n_tile, ct_max_k): raise ValueError(f"nothing fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}") -def flm_gemm( +def gemm( dev, M, K, @@ -149,37 +220,49 @@ def flm_gemm( tile_n=N_TILE_DEFAULT, tile_ma=None, overlap=None, - kernel_object="flm_gemm.o", - epilogue_object="flm_gemm_epilogue.o", + kernel_object="mm_fused.o", + epilogue_object="mm_fused_epilogue.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 - ``FLMGEMM.pack_B`` -- it emits B in the order the cores consume it, so both + ``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}" ) + # Grid width, and the shim columns feeding the A broadcast, both follow the + # device. Everything below is written against these rather than a constant, + # so the same dataflow covers NPU2's 4x8 and NPU1's 4x4. + COLS = grid_cols(dev) + A_SOURCE_COL = a_source_cols(COLS) + # r/t are the device's native mac shape; every blocked layout below is + # expressed in terms of them. + R, _S, T = register_tiling(dev.resolve().name) N_TILE = tile_n CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] + # B's storage format follows the device, and with it every B type and every + # B 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 = bfp16_b_for(dev) + B_GROUP = BFP16_GROUP if BFP16_B else 1 + b_elem_bytes = 9 / 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. - _t_ma_fit, L1_B_DEPTH = _default_l1(N_TILE, CT_MAX_K) + _t_ma_fit, L1_B_DEPTH = _default_l1(N_TILE, CT_MAX_K, b_elem_bytes) T_MA = _t_ma_fit if tile_ma is None else 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}" ) RHO = M_TILE // T_MA - import os as _os - - OVERLAP = int(_os.environ.get("FLM_OVERLAP", "2")) if overlap is None else overlap + 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 @@ -283,26 +366,32 @@ def flm_gemm( 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 // 8,), np.dtype[v8bfp16ebs8]] + 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 // 8,), np.dtype[v8bfp16ebs8]] - mt_b_bytes = _bfp16_bytes(K_TILE * N_TILE) + 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 // 8,), np.dtype[v8bfp16ebs8]] + b_l3_ty = np.ndarray[(K * N // B_GROUP,), b_elem_ty] c_l3_ty = np.ndarray[(M * N,), bf16_ty] - acc_init = Kernel("flm_gemm_acc_init", kernel_object, [ct_acc_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( - "flm_gemm_k_step", + "mm_fused_k_step", kernel_object, [ct_a_obj_ty, ct_b_ty, ct_acc_ty, np.int32], ) @@ -384,77 +473,33 @@ def _split_run(run): a_cons[(r, c)] = of_a.cons() # B: shim -> memtile -> broadcast down the compute column. - # Resident B: hold a whole column-block's B in the memtile as ONE object - # and re-walk it per row-block, so DDR sees it once instead of - # m_row_blocks times. B is the dominant DDR leg -- at M=1024 K=1536 N=6144 - # it is 75 MB of the 126 MB moved -- so this is ~43% less total traffic. # - # It saves power and time both, and the time is worth more the taller the - # problem is, because B's DDR re-reads scale with m_row_blocks. Measured at - # K=1024 N=4096 (full builds, min of interleaved rounds): + # 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. # - # M=512 (2 row-blocks) 470.8 -> 468.5 us 0.5% - # M=1024 (4 row-blocks) 860.4 -> 846.5 us 1.6% - # M=2048 (8 row-blocks) 1760.1 -> 1622.8 us 7.8% + # Three things here are load-bearing rather than tuning: # - # The operator IS DDR-bandwidth bound: with the mmul nulled out, the - # non-resident floor at M=2048 is 1739 us for 118 MB, i.e. 68 GB/s, against - # a measured 63-70 GB/s roof. Residency drops that floor to 1135 us. + # * 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. # - # Note the gap between that 604 us of floor and the 137 us actually - # captured: repeat_count restarts the memtile BD chain at every replay - # boundary, and ~77% of the win goes there. Closing it is the largest known - # remaining lever on this design. Do not measure it at small M -- at M=512 - # the effect is inside the noise, which is how it was first mistaken for a - # power-only optimisation. - # - # One object, not k_iters of them: iterating a pool replays each object in - # turn (k0,k0,k1,k1,...) rather than the sequence. Fitting k into the - # descriptors within the 4-dimension limit takes both hops -- inbound the - # outermost dim already steps by (N_TILE//T)*(K_TILE*T), exactly - # K_TILE*N_TILE, so widening its COUNT walks into the next k-block; - # outbound k becomes a new outermost dim, which also keeps one emitted - # object per CT_MAX_K slice. - # - # Replay is repeat_count, on the forward() below. iter_count cannot do this - # job: it only bounds how many times an end cycles through all its buffers - # (objects = iter_count * elemNumber * repeat_count), so it is in units of - # depth-cycles rather than objects, and getting it wrong hangs rather than - # mis-computes. - # - # Gated on the buffer fitting DOUBLE-buffered, so the next column-block - # still prefetches. B_DEPTH is load-bearing, not a safety margin: - # single-buffering would let the lowering use the DMA's native repeat field - # instead of a duplicated BD chain, but it also stops the next - # column-block's B prefetching behind this one's replay, and that costs - # more than it saves -- 1720 us against 1600 at M=2048 K=1024 N=4096, - # i.e. worse than not being resident at all. - # - # The budget counts only C's 64 KB, not A's 128 KB, so it admits k_iters - # <= 3 (K <= 1536 at tile_n=64) rather than 2. That deliberately overcommits - # the A-carrying memtiles by 64 KB and relies on aie-objectfifo-allocate - # spilling one buffer to the least-loaded adjacent memtile, which packs all - # eight to exactly 512 KB (C_L2L3_6 lands on mem_tile_7_1). There is ZERO - # slack: re-verify placement after any change to the A, B or C buffer - # sizes, or the build will fail address assignment rather than silently - # mis-run. - # - # Residency only pays once compute is off the critical path -- it was - # measured latency-neutral while the mmul was the wall, and worth 260 us - # immediately after the mmul loop was re-rolled. Larger K still falls back - # to non-resident, so check this gate before believing any measurement that - # claims to be testing residency. - # Count what A and C actually occupy rather than assuming C's 64 KB is the - # only other tenant. The old form ignored A entirely and hardcoded C's size, - # which at tile_n=128 (where C doubles and resident B is 432 KB) admitted a - # configuration that then failed address assignment outright. - # Prefer a double-buffered resident B so the next column-block prefetches - # behind this one's replay. Single-buffering costs that prefetch and was - # measured WORSE than not being resident at all -- at tile_n=64, where B is - # small enough that depth 2 fits anyway. At tile_n=128 B is twice the size - # and depth 2 does not fit, but depth 1 does; and there residency is worth - # far more, because it is also what keeps B's DDR from quadrupling. So take - # the deepest that fits rather than giving up on residency. + # 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 = 512 * 1024 - 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 @@ -466,7 +511,7 @@ def _split_run(run): # 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 // 8,), np.dtype[v8bfp16ebs8]] + mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE // B_GROUP,), b_elem_ty] b_l3l2_fifos = [] b_cons = {} @@ -616,20 +661,20 @@ def b_tap(mega_col, c): # 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 FLMGEMM.pack_B) so each k-block is one + # 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 // 8,), - offset=(mega_col * COLS + c) * N_TILE * K // 8, + 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 // 8] + [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 // 8] + 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 // 8, 0, 1] + [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // B_GROUP, 0, 1] ), ) @@ -637,9 +682,7 @@ 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. The - # ablation knobs below are unsplit-only; they write C to the wrong - # place by construction and exist purely for timing. + # the mega_row stride overflow the shim BD's iteration step. if c_split: return [ TensorAccessPattern( @@ -653,30 +696,6 @@ def c_taps(mega_col, c, mbs): return [_c_tap_unsplit(mega_col, c)] def _c_tap_unsplit(mega_col, c): - if _os.environ.get("FLM_C_LINEAR") == "1": - # ABLATION: same byte count, one contiguous run per column-block - # instead of ROWS*M_TILE runs of N_TILE. Writes C to the WRONG - # place -- only for measuring what the scattered write costs. - blk = m_row_blocks * ROWS * M_TILE * N_TILE - return TensorAccessPattern( - tensor_dims=(M * N,), - offset=(mega_col * COLS + c) * blk, - sizes=[1, 1, 1, blk], - strides=[0, 0, 0, 1], - ) - if _os.environ.get("FLM_C_RUN2") == "1": - # ABLATION: identical byte count and identical DDR footprint, but - # HALF as many runs each TWICE as long (N_TILE*2 = 256 B instead of - # 128 B), by walking every other row. Writes C to the WRONG place. - # This isolates exactly what the column-pair join would buy -- - # FLM_C_LINEAR above removes ALL scatter, so it is the ceiling for - # perfect linearisation, not for the 128 B -> 256 B step. - return TensorAccessPattern( - tensor_dims=(M * N,), - offset=(mega_col * COLS + c) * N_TILE, - sizes=[1, m_row_blocks, ROWS * M_TILE // 2, N_TILE * 2], - strides=[0, ROWS * M_TILE * N, N * 2, 1], - ) return TensorAccessPattern( tensor_dims=(M * N,), offset=(mega_col * COLS + c) * N_TILE, @@ -798,10 +817,7 @@ def emit_split(): a_l3_ty, b_l3_ty, c_l3_ty, - [ - f.prod(tile=Tile(A_SOURCE_COL[r], 0)) - for r, f in enumerate(a_l3l2_fifos) - ], + [f.prod(tile=Tile(A_SOURCE_COL[r], 0)) for r, f in enumerate(a_l3l2_fifos)], [f.prod(tile=Tile(c, 0)) for c, f in enumerate(b_l3l2_fifos)], [f.cons(tile=Tile(c, 0)) for c, f in enumerate(c_l2l3_fifos)], ], @@ -810,3 +826,48 @@ def emit_split(): 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=str, choices=sorted(EPILOGUE_MODES), default="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 000000000..0cd45e4f1 --- /dev/null +++ b/iron/operators/flm/gemm/op.py @@ -0,0 +1,419 @@ +# 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 iron.common.device_utils import get_kernel_dir +from iron.common.operator_bases import lut_based_ops_artifacts +import aie.utils as aie_utils + +from iron.operators.flm.packing import pack_b, packed_b_size +from iron.operators.flm.gemm.design import ( + CT_MAX_K_FOR_N, + C_DEPTH, + CT_OUT_LEN, + EPILOGUE_MODES, + K_TILE, + MIN_K, + MIN_M, + M_TILE, + _default_l1, + register_tiling, +) + + +@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 + # "none" | "gelu" | "silu" | "sigmoid", fused into the C drain. + epilogue: str = "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 _default_tile_n 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 + # "conv_even" (round to nearest even) or "floor" (truncate). The core powers + # up in floor and the overlay this was ported from never sets the mode, so + # "floor" reproduces its arithmetic exactly -- at ~40x the error, because + # truncation biases every conversion the same way and the bias accumulates + # over the K reduction instead of cancelling. + rounding: str = "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. + if self.tile_n is None: + self.tile_n = self._default_tile_n(self.K) + 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 + )[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, MIN_M), + ("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}") + if self.epilogue not in EPILOGUE_MODES: + raise ValueError( + f"epilogue must be one of {sorted(EPILOGUE_MODES)}, " + f"got {self.epilogue!r}" + ) + if self.clamp is not None: + lo, hi = self.clamp + if lo > hi: + raise ValueError(f"clamp min ({lo}) must be <= max ({hi})") + if self.rounding not in ("conv_even", "floor"): + raise ValueError( + f"rounding must be 'conv_even' or 'floor', got {self.rounding!r}" + ) + + MLIROperator.__init__(self, context=self.context) + + @staticmethod + def _default_tile_n(K: int, kernel_dir: str | None = None) -> int: + """Pick the n tile from the shape and the device. + + 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. + Measured ~20% for n=64 at K >= 1024 and ~9% the other way at K = 512. + + 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. Measured on Phoenix, n=64 wins + everywhere by 1.21x (M=256 K=512 N=512) to 1.38x (M=1024 K=2048 + N=1024) -- including at K=512, where NPU2's rule would pick 128. + """ + if kernel_dir is None: + kernel_dir = get_kernel_dir() + if kernel_dir == "aie2": + return 64 + return 128 if K // K_TILE <= 1 else 64 + + @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 _epilogue_artifact(self) -> str: + """Object name for the epilogue, over the flags that shape it.""" + clamp = "" + if self.clamp is not None: + clamp = "_clamp" + "_".join( + repr(float(v)).replace(".", "p").replace("-", "n").replace("+", "") + for v in self.clamp + ) + return f"mm_fused_epilogue_{self.epilogue}{clamp}_{self.rounding}.o" + + @property + def _needs_tanh_lut(self) -> bool: + """Whether the epilogue has to be linked against 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 converts and stores. Note the failure mode when this is + wrong is a LINK error for tanh_lut_ab/tanh_lut_cd, not a compile error, + so it surfaces late. + """ + return self.epilogue != "none" and get_kernel_dir() == "aie2" + + @property + def _epilogue_link_file(self) -> str: + """What the design should name as the epilogue kernel: the bare object, + or the archive bundling it with the LUT tables.""" + if self._needs_tanh_lut: + return f"{self.name}_epilogue.a" + return self._epilogue_artifact + + @property + def _rounding_flags(self) -> list[str]: + """Applies to both kernels: the mmul and the epilogue's f32->bf16 + store are both conversions and must agree. + + ROUND_CONV_EVEN is mm.cc's flag, reused here rather than inventing a + second spelling. Its polarity is mm.cc's too -- absent means the core's + power-up floor mode -- even though this operator defaults the other way. + """ + return ["-DROUND_CONV_EVEN"] if self.rounding == "conv_even" else [] + + @property + def _epilogue_source(self): + return ( + self.context.base_dir / "aie_kernels" / "generic" / "mm_fused_epilogue.cc" + ) + + @property + def _epilogue_flags(self) -> list[str]: + """Compile flags for the epilogue.""" + flags = [ + f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", + f"-DMM_FUSED_C_DEPTH={C_DEPTH}", + f"-DMM_FUSED_EPILOGUE_MODE={EPILOGUE_MODES[self.epilogue]}", + ] + 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", + ] + return flags + self._rounding_flags + + @property + def _bfp16_b(self) -> bool: + """Whether B is stored as bfp16ebs8 rather than bf16. + + AIE2P only. ``__AIE_API_SCALAR_BFP_TYPES__`` is defined solely in + ``aie_api/detail/aie2p/config.hpp`` and ``mmul_bfp16_bfp16.hpp`` exists + only under ``aie2p/``; on AIE2, ``aie_api/types.hpp`` gives + ``bfp16ebs8`` an empty placeholder struct. So AIE2 keeps B in bf16 and + uses the native 4x8x4-composed mmul, which is why both mmul templates + in the kernel header are live rather than one being dead code. + """ + return get_kernel_dir() == "aie2p" + + @property + def _b_elem_bytes(self) -> float: + """Bytes per B element in L1/L2: bfp16ebs8 packs 8 values into 9 bytes.""" + return 9 / 8 if self._bfp16_b else 2 + + @property + def _rst(self) -> tuple[int, int, int]: + """The mmul's r/s/t for this device. ``pack_B`` and the design's stream + dimensions must agree on these or the result is silently wrong.""" + return register_tiling(aie_utils.get_current_device().resolve().name) + + @property + def _kernel_object(self) -> str: + """Object name over every flag that changes the emitted code. + + r/t are included even though build dirs are already architecture-scoped, + because they set the blocked layout: an object built for one shape must + never satisfy a request for another. + """ + r, _s, t = self._rst + return ( + f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" + f"_r{r}t{t}_ma{self.tile_ma}_{self.rounding}.o" + ) + + 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._kernel_object, + "epilogue_object": self._epilogue_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. + emulate_flags = ( + ["-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", "-DMM_FUSED_BFP16_B"] + if self._bfp16_b + else [] + ) + + artifacts = [ + KernelObjectArtifact( + self._kernel_object, + dependencies=[ + SourceArtifact(generic / "mm_fused.cc"), + SourceArtifact(generic / "mm_fused_mmul.h"), + SourceArtifact(base_dir / "aie_kernels" / "aie_kernel_utils.h"), + SourceArtifact(base_dir / "aie_kernels" / kernel_dir / "zero.cc"), + ], + extra_flags=[ + 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={self._rst[0]}", + f"-DMM_FUSED_S={self._rst[1]}", + f"-DMM_FUSED_T={self._rst[2]}", + # The k slice. Passed rather than looked up in the kernel so + # that CT_MAX_K_FOR_N below is the only place it is chosen. + f"-DMM_FUSED_CT_K={CT_MAX_K_FOR_N[self.tile_n]}", + ] + + arch_include + + emulate_flags + + self._rounding_flags, + ), + ] + epilogue_obj = KernelObjectArtifact( + self._epilogue_artifact, + dependencies=[ + SourceArtifact(self._epilogue_source), + SourceArtifact(generic / "activations.h"), + SourceArtifact(base_dir / "aie_kernels" / "aie_kernel_utils.h"), + ], + extra_flags=self._epilogue_flags, + ) + if self._needs_tanh_lut: + # The LUT's coefficient tables live in their own translation unit in + # mlir-aie's runtime lib, so the epilogue object alone leaves + # tanh_lut_ab/tanh_lut_cd undefined at link time. + artifacts.append( + KernelArchiveArtifact( + self._epilogue_link_file, + dependencies=[epilogue_obj] + lut_based_ops_artifacts(kernel_dir), + ) + ) + else: + artifacts.append(epilogue_obj) + return artifacts + + 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. + """ + _r, s, t = self._rst + 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 == "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 similarity index 74% rename from iron/operators/flm_gemm/reference.py rename to iron/operators/flm/gemm/reference.py index 861580089..7872824b1 100644 --- a/iron/operators/flm_gemm/reference.py +++ b/iron/operators/flm/gemm/reference.py @@ -5,31 +5,27 @@ from iron.common.test_utils import torch_dtype_map -def _activation(x, epilogue): - if epilogue == "none": - return x - if epilogue == "gelu": - # Match the kernel, which uses the sigmoid approximation - # gelu(x) ~= x * sigmoid(1.702x) -- NOT torch's erf-exact gelu, and not - # the tanh approximation the standalone gelu operator uses. - return x * torch.sigmoid(1.702 * x) - if epilogue == "silu": - return x * torch.sigmoid(x) - if epilogue == "sigmoid": - return torch.sigmoid(x) - raise ValueError(f"unknown epilogue {epilogue!r}") - - def reference(input_a, input_b, epilogue="none", clamp=None): """CPU reference ``C = clamp(activation(A @ B))``. The matmul is accumulated in fp32 to mirror the kernel's f32 accumulator, then cast back to the input dtype at the end, which is where the kernel converts too. + + ``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. """ out_dtype = input_a.dtype C = torch.matmul(input_a.float(), input_b.float()) - C = _activation(C, epilogue) + if epilogue == "gelu": + C = C * torch.sigmoid(1.702 * C) + elif epilogue == "silu": + C = C * torch.sigmoid(C) + elif epilogue == "sigmoid": + C = torch.sigmoid(C) + elif epilogue != "none": + raise ValueError(f"unknown epilogue {epilogue!r}") if clamp is not None: C = torch.clamp(C, clamp[0], clamp[1]) return C.to(out_dtype) diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py new file mode 100644 index 000000000..c8155deed --- /dev/null +++ b/iron/operators/flm/gemm/test.py @@ -0,0 +1,218 @@ +#!/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 iron.operators import GEMM as GenericGEMM +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 + +# 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 + + +@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 == "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, + ) + + input_buffers = { + "A": golden_ref["input"].flatten(), + # B is consumed pre-packed; see GEMM.pack_B. + "B": operator.pack_B(golden_ref["input_b"]), + } + output_buffers = {"C": golden_ref["output"].flatten()} + + # The rule: bound 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, because with signed A the K-sum cancels by ~sqrt(K), so |C| is far + # smaller than the mass while the error tracks the mass -- which leaves + # near-zero outputs relatively uncheckable. + mass = ( + K + * golden_ref["input"].abs().float().mean() + * (golden_ref["input_b"].abs().float().mean()) + ) + # The fraction is per-architecture, because the two lower the same 8x8x8 + # mmul shape onto very different arithmetic: NPU2 emulates it with bfp16, + # which drops mantissa bits, while NPU1 has no bfp16 and lowers it onto + # four native bf16 macs accumulating in f32, which is exact up to the + # f32->bf16 store and so gets a ~20x tighter budget. + # + # floor rounding truncates rather than rounding to nearest, so its bias + # accumulates over the K reduction instead of cancelling, hence the + # separate, looser bound on both. + if aie_utils.get_current_device().resolve().name == "npu1": + budget = 0.002 if rounding == "floor" else 0.0002 + else: + budget = 0.05 if rounding == "floor" else 0.004 + errors, latency_us, bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=0.04, + abs_tol=float(budget * mass), + ) + + 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 iron.operators.flm.gemm.design import SHIM_BDS, SHIM_TASK_QUEUE + + worst = 1 + 2 * SHIM_TASK_QUEUE + assert worst <= SHIM_BDS, ( + f"a fully split block needs {worst} shim BDs of {SHIM_BDS}; " + "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() + + +@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 000000000..730fcb0b9 --- /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 000000000..b5e32d99e --- /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_MODES, + K_TILE, + MAX_COLS, + M_TILE, + ROWS, + a_source_cols, +) + +# The shipped overlay is built with n=128 on the full 8-column NPU2 grid; every +# other tiling knob matches flm.gemm, whose constants are imported above. +N_TILE = 128 +COLS = MAX_COLS +A_SOURCE_COL = a_source_cols(COLS) + +# 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="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`. + """ + if epilogue not in EPILOGUE_MODES: + raise ValueError( + f"epilogue must be one of {sorted(EPILOGUE_MODES)}, got {epilogue!r}" + ) + 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_MODES[epilogue]), + (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 000000000..a54a76bf5 --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/op.py @@ -0,0 +1,180 @@ +# 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_MODES, 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 + # "none" | "gelu" | "silu" | "sigmoid", selected through a runtime + # parameter rather than at compile time as in flm.GEMM. + epilogue: str = "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}") + if self.epilogue not in EPILOGUE_MODES: + raise ValueError( + f"epilogue must be one of {sorted(EPILOGUE_MODES)}, " + f"got {self.epilogue!r}" + ) + 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 000000000..243fb0ba8 --- /dev/null +++ b/iron/operators/flm/mm_prebuilt/test.py @@ -0,0 +1,81 @@ +#!/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 pytest + +import aie.utils as aie_utils + +from iron.common.test_utils import run_test +from iron.operators.flm.gemm.reference import generate_golden_reference +from iron.operators.flm.mm_prebuilt.op import MMPrebuilt + +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), + (512, 1024, 2048, "none", None), + (256, 512, 1024, "silu", None), + (256, 512, 1024, "gelu", None), + (256, 512, 1024, "sigmoid", None), + (256, 512, 1024, "none", (-2.0, 2.0)), + ], +) +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()} + + # Same absolute-mass bound as flm.gemm/test.py, for the same reason: with + # signed A the K-sum cancels by ~sqrt(K), leaving near-zero outputs + # relatively uncheckable under a plain relative tolerance. + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + errors, latency_us, bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=0.04, + abs_tol=float(BUDGET_FLOOR * mass), + ) + assert not errors, "Test failed" diff --git a/iron/operators/flm/packing.py b/iron/operators/flm/packing.py new file mode 100644 index 000000000..c02708ac5 --- /dev/null +++ b/iron/operators/flm/packing.py @@ -0,0 +1,130 @@ +# 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) + return blocked.permute(4, 0, 1, 5, 3, 2, 6).reshape(-1).contiguous() + # -> (cb, kb, kslice, tb, i, s_in, t_in) + # Row-major s x t within the block, which is what the plain mmul loads. + return blocked.permute(4, 0, 1, 5, 2, 3, 6).reshape(-1).contiguous() + # -> (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/flm_gemm/README.md b/iron/operators/flm_gemm/README.md deleted file mode 100644 index ed9598193..000000000 --- a/iron/operators/flm_gemm/README.md +++ /dev/null @@ -1,251 +0,0 @@ - - -# FLMGEMM — bf16 GEMM on a fixed 4x8 grid - -A second GEMM design, ported from FastFlowLM's `mm` overlay. It is a different -dataflow from [`GEMM`](../gemm), not a retuning of it: - -| | `GEMM` | `FLMGEMM` | -|---|---|---| -| geometry | parameterized tiles, 1–8 columns | fixed m=64 k=512, r/s/t 8/8/8, 4x8 grid; n selectable | -| A delivery | per column | broadcast along each compute row from 4 shim columns | -| C staging | full m x n tile in L1 | streamed out in 512-element chunks | -| C collection | per column | ObjectFifo `join` of 4 rows through the memtile | -| epilogue | separate `convert_copy` | fused f32→bf16 + activation + clamp | -| B layout | plain `(K, N)` | **pre-packed**, see below | - -On NPU2 (aie2p) only: the r=8 mmul shape exists solely on the bfp16-emulated -path, and the grid needs all 8 columns. - -## 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 * 8` stride: a -trailing group of fewer than 8 column-blocks 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 = FLMGEMM(M=M, K=K, N=N, context=ctx) -op.compile().get_callable()(A, op.pack_B(B), C_out) -``` - -`pack_B` reorders a row-major `(K, N)` matrix into the order the memtile -expects — each `K_TILE x tile_n` tile as the odometer `(n//T, k%S, k//S, n%T)`, -outermost first — so each fill is one contiguous read. - -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, and that was the original implementation, but its -innermost run is then `T=8` bf16 = 16 bytes: each 128 KB transfer becomes 8192 -scattered bursts. B is ~70% of the bytes a dispatch moves, so the whole -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 -FLMGEMM(M=M, K=K, N=N, rounding="floor", context=ctx) # matches shipped -FLMGEMM(M=M, K=K, N=N, context=ctx) # conv_even, default -``` - -Verified against FastFlowLM v1.0.4's -`xclbins/Gemma4-E2B-IT-NPU2/mm.xclbin`, driving it directly with the -instruction stream from that project's own TXN generator, on identical inputs: - -| | err/mass | vs shipped | -|---|---|---| -| `rounding="floor"` | 0.009867 | **bit-identical, 6291456/6291456 elements** | -| `rounding="conv_even"` (default) | 0.000241 | differs everywhere | - -All four epilogues are bit-identical to the shipped kernel too, with `floor` -(the shipped kernel selects its activation from RTP word 4; this operator -bakes it in at compile time, with the same 0/1/2/3 mapping): - -| epilogue | vs shipped `output_mode` | -|---|---| -| `none` / `gelu` / `silu` / `sigmoid` | bit-identical, 1048576/1048576 each | - -**The default is `conv_even`, not `floor`.** 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 original. - -`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 expectations - -The r=8 mmul exists only on the bfp16-emulated path, so the error budget is -that of an emulated GEMM. Do not compare against `GEMM`'s test tolerances, -which assert on the exact r=4 path (`emulate_bf16_mmul_with_bfp16=False`). - -A pure elementwise *relative* tolerance is not meaningful here: 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 instead, as `test.py` does. Reference points on random -signed A / non-negative B: - -| | mean err / mass | -|---|---| -| `FLMGEMM` (default) | 0.000241 | -| `GEMM`, same mode (`emulate=True, prio_accuracy=True`) | 0.000241 | -| `GEMM`, bf16 accumulator (`prio_accuracy=False`) | 0.000445 | -| `GEMM`, exact r=4 path (`emulate=False`) | 0.00007 | - -This operator and `GEMM` in the same mode are numerically **indistinguishable** --- identical mean error, signed bias and maximum, at both `tile_n` values. -Same mmul shape, same bfp16 emulation, same f32 accumulation, same rounding, -so there is no reason for them to differ and they do not. The only accuracy -difference worth knowing about is `conv_even` versus the shipped overlay's -`floor` (above). - -## Choosing `tile_n` - -`tile_n` defaults to `None`, which picks per shape: **128 when `K == 512`, -otherwise 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 1024 -columns of N per pass rather than 512. Which wins depends on whether compute -or data movement is the critical path, and 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. Measured 2026-09-09 against the current design -(rolled mmul, resident B, ATB, bfp16 B), min of per-run medians over 6 rounds -with the two `tile_n` builds interleaved round-robin -- this box is bimodal -~6%, so running all of one and then all of the other measures drift rather -than design: - -| M / K / N | k_iters | `tile_n=64` | `tile_n=128` | -|---|---|---|---| -| 1024 / 512 / 4096 | 1 | 514 us | **498 us** | -| 1024 / 1024 / 4096 | 2 | **591 us** | 931 us | -| 1024 / 1536 / 6144 | 3 | **1141 us** | 1960 us | -| 1024 / 2560 / 4096 | 5 | **1239 us** | 1940 us | -| 2048 / 2048 / 2048 | 4 | **915 us** | 1581 us | -| 256 / 4096 / 1024 | 8 | **227 us** | 265 us | - -The default rule is unchanged in sign: `tile_n=128` still wins only at -`k_iters=1`. But its margin there has narrowed to 3% (was 8%) and its penalty -everywhere else has grown -- at `k_iters>=2` it is now 1.2-1.7x slower where -it used to be 1.2-1.25x. 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): the more k -there is to reduce over, the more that costs. - -`pack_B` is bound to the operator because the packing layout depends on -`tile_n`; call `op.pack_B(B)`, not `FLMGEMM.pack_B(B)`. - -## Performance - -M=1024 K=1536 N=6144, min of per-run medians across separate processes: - -| | bytes moved | latency | DMA-only (compute nulled) | -|---|---|---|---| -| `FLMGEMM` (`tile_n=64`) | 47 MB | **1143 us** | -- | -| shipped `mm.xclbin` | 107 MB | 2175 us | -- | -| `GEMM` (`emulate=True, prio_accuracy=True`) | 126 MB | 3353 us | 3374 us | - -**1.90x the shipped overlay**, at err/mass 2.39e-04 against its 9.87e-03 -- -41x more accurate. Storing B in bfp16 is numerically free: the mmul only -multiplies bfp16, so quantizing on the host hoists a rounding that already -happened on every mac call. It has to reproduce the core's rounding MODE to -do so -- see ``pack_B``. The 69 MB is with B resident in the memtile; the 126 MB figure -this table used to quote was the non-resident fallback. At 1252 us against a -1231 us data-movement floor, this operator is now essentially DMA-bound: the -mmul is finally cheap enough to hide, so further gains have to come from -moving fewer bytes. - -Three changes compound to get there, and none of them works alone: - -* **The mmul's inner loop is rolled**, not hand-unrolled. 2446 -> 1875 cycles - per call. -* **`pack_B` emits the final consumption order**, so both B hops are linear - descriptors instead of blocked ones. Worth nothing by itself -- it is what - frees the descriptor dimensions the other two need. -* **Asymmetric tile buffering.** The A tile is 16 rows while the accumulator - is 64 (`rho = 4`), which pays for a 128-deep k slice. That halves the - accumulator traffic per mac and doubles the inner loop's trip count: - 3.67 -> 2.64 cycles per 8x8x8 mac. - -**Measure this carefully.** 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 -- that is how a change later shown to do nothing at all first -produced a convincing 5% "win". 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. - -Nulling the mmul out is what makes this legible. `GEMM` does not change at all -without it (3374 vs 3353 us), so it is entirely data-movement bound. - -This operator is **not**. An earlier revision of this section read the small -gap between the nulled floor and the full time (1692 of 1741 us) as proof that -`tile_n=64` was data-movement bound, and concluded that "further gains come -from moving fewer bytes, not from a faster kernel". That was backwards. The -floor sat just *above* the mmul, hiding it; the operator was compute bound the -whole time, and every attempt to move fewer bytes duly measured as worth -nothing. Re-rolling the mmul's inner loop cut it 23% (2446 -> 1875 cycles per -call, HW trace) and only then did keeping B resident pay -- together 1741 -> -1434 us. - -Two lessons worth keeping. When total latency is `max(compute, DMA)`, testing -levers **one at a time scores both as zero**: residency alone gained nothing -while the mmul was the wall, and the re-rolled mmul alone gained little while -the non-resident floor was. And a nulled-mmul floor close to the full time does -not by itself mean DMA-bound -- it equally means compute is hiding just -underneath. Compare against the *nulled floor*, not the full time, when judging -a data-movement change. - -Its transfers are cheaper mostly because B arrives pre-packed: the contiguous -run per transfer is 128 KB for B and 1 KB for A, against 128 bytes on every -leg for `GEMM`, which reorders in the descriptor instead. - -Two things dominate, and both are in the runtime sequence rather than the -kernel: B must be pre-packed (above), and each of A, B and C must go out as -**one transfer per column-block** rather than one per fifo object. A single -fill or drain may span many objects; issuing per object instead means a host -await per row-block, and a C await waits on the cores. Collapsing those is -also what makes overlapping column-blocks affordable — a block then costs 3 -buffer descriptors on a shim tile 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. - -This operator is DDR-bandwidth bound, so that is a latency win as well as a -power one, and it grows with the height of the problem because B's re-reads -scale with `m_row_blocks`. At K=1024 N=4096: - -| M | row-blocks | non-resident | resident | | -|---|---|---|---|---| -| 512 | 2 | 470.8 us | 468.5 us | 0.5% | -| 1024 | 4 | 860.4 us | 846.5 us | 1.6% | -| 2048 | 8 | 1760.1 us | **1622.8 us** | 7.8% | - -Do not evaluate this at small M: at M=512 the effect is inside the noise, which -is how it was first mistaken for a power-only optimisation. - -Most of the available win is still on the table. With the mmul nulled, the -non-resident floor at M=2048 is 1739 us -- 118 MB at 68 GB/s, against a -memcpy-measured 63-70 GB/s roof for mixed read/write traffic -- and residency -drops that floor to 1135 us. Only 137 us of those 604 us reaches the full -build; the rest goes to `repeat_count` restarting the memtile BD chain at every -replay boundary. Closing that is the largest known remaining lever here. diff --git a/iron/operators/flm_gemm/bench_vs_flm.py b/iron/operators/flm_gemm/bench_vs_flm.py deleted file mode 100644 index 6c4827973..000000000 --- a/iron/operators/flm_gemm/bench_vs_flm.py +++ /dev/null @@ -1,387 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Benchmark flm_gemm against the shipped FastFlowLM overlay and IRON's generic GEMM. - -Kept out of ``test.py`` on purpose: this file needs an external FastFlowLM -install, which the operator itself does not. ``test.py`` must stay runnable in -CI with nothing but this repo; this one skips the whole module when the install -is absent. - -Every shape is one test, covering the real Gemma4 projections (E2B and E4B x -{q, kv, o, gateup, down}) at three prefill lengths. Three competitors run per -shape: - - iron : the ``FLMGEMM`` operator (arg order A, B, C) - flm : FastFlowLM's shipped ``mm.xclbin`` + its dumped TXN insts (order C, A, B) - gemm : IRON's generic ``GEMM`` operator, same emulated-bfp16 numerics - -``gemm`` drops out of the E4B gate/up shapes at M>256: its C descriptor needs a -mega_row stride past the shim's 20-bit step field there, which aiecc rejects. -Those shapes report iron against flm only. - -The box is bimodal by ~6% (see the npu-bimodal-timing note), so the three are -interleaved round-robin over several rounds and each is scored by the MINIMUM -of its per-round medians. Running all of one competitor and then all of another -fabricates differences of about that size. Everything is compiled up front; -nothing is rebuilt between rounds. - -No time is reported for a dispatch whose output was not checked. - -Usage:: - - pytest iron/operators/flm_gemm/bench_vs_flm.py --iterations 1 - pytest iron/operators/flm_gemm/bench_vs_flm.py -k E2B --csv-output flm.csv -""" - -import os -import statistics -import subprocess -import time -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.flm_gemm.op import FLMGEMM -from iron.operators.gemm.op import GEMM - -# --------------------------------------------------------------------------- -# External FastFlowLM install. Both are overridable so this is not pinned to -# one machine's layout; absent either, the module skips rather than errors. -# --------------------------------------------------------------------------- -FLM_XCLBIN = Path( - os.environ.get( - "FLM_MM_XCLBIN", - "/scratch/ehunhoff/flm-release-1.0.4/extracted/opt/fastflowlm/share/flm" - "/xclbins/Gemma4-E2B-IT-NPU2/mm.xclbin", - ) -) -# Dumps the overlay's control instructions for one (M, K, N). The generator -# baked into it is gemma4_e2b_mm_txn, so it belongs to the xclbin above. -# Previously lived in /tmp, which is tmpfs here and did not survive reboots. -TXN_DUMP = Path( - os.environ.get("FLM_TXN_DUMP", "/scratch/ehunhoff/flm_gemm_bench/txn/dump") -) -TXN_CACHE = Path(os.environ.get("FLM_TXN_CACHE", TXN_DUMP.parent)) - -# The shipped overlay is a fixed n=128 design; its B layout is not negotiable. -FLM_N_TILE = 128 -# B's memtile odometer, shared by both packers. K_TILE is the operator's. -K_TILE, S, T = 512, 8, 8 - -# Interleaved rounds per test, and timed dispatches per competitor per round. -# 6 rounds is the floor at which min and median stopped disagreeing on this box. -ROUNDS = 6 -ITERS = 30 -WARMUP = 20 - -# err/mass budgets. iron and 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 _flm_available(): - return FLM_XCLBIN.is_file() and os.access(TXN_DUMP, os.X_OK) - - -if not _flm_available(): - pytest.skip( - f"FastFlowLM install not found (looked for {FLM_XCLBIN} and {TXN_DUMP}); " - "set FLM_MM_XCLBIN / FLM_TXN_DUMP to point at one", - allow_module_level=True, - ) - -_dev = aie_utils.get_current_device() -if _dev.cols < 8 or _dev.resolve().name != "npu2": - pytest.skip( - "flm_gemm is a fixed 4x8 npu2 design; this device cannot run it", - allow_module_level=True, - ) - - -# --------------------------------------------------------------------------- -# Shapes -# --------------------------------------------------------------------------- -# Every projection of both Gemma4 variants FastFlowLM ships, at three prefill -# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. -# 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] - -# All 30 shapes run. The four E4B projections with a 10240-wide dimension used -# to be skipped here: at M>256 their mega_row stride overflows the shim BD's -# 20-bit iteration step. design.py now issues that leg as one transfer per -# mega_row, retired in windows -- see its a_split comment. - - -def get_params(): - 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 - - -# --------------------------------------------------------------------------- -# Inputs and packing -# --------------------------------------------------------------------------- -def make_inputs(M, K, N): - """Identical data for all three competitors, and the reference to check.""" - 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, not 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, and near-zero outputs are - # relatively uncheckable. Same rationale as test.py's bound. - return A, B, Af @ Bf, float((Af.abs() @ Bf.abs()).mean()) - - -def flm_pack_B(Bt, n_tile=FLM_N_TILE): - """(K, N) row-major -> the shipped overlay's memtile order. - - Odometer (n//T, k%S, k//S, n%T), tiles ordered by column stripe then - k-block. Mirrors ``FLMGEMM.pack_B``'s tiling but stops at bf16: the - overlay consumes bf16, not the bfp16 blocks the IRON operator takes, and - is fixed at n=128 regardless of what the IRON operator chooses. - """ - b = Bt.float().numpy() - K, N = b.shape - out = [] - for cb in range(N // n_tile): - stripe = b[:, cb * n_tile : (cb + 1) * n_tile] - for kb in range(K // K_TILE): - tile = stripe[kb * K_TILE : (kb + 1) * K_TILE, :] - out.append( - tile.reshape(K_TILE // S, S, n_tile // T, T) - .transpose(2, 1, 0, 3) - .ravel() - ) - return torch.from_numpy(np.concatenate(out).astype(np.float32)).to(torch.bfloat16) - - -def flm_insts(M, K, N): - """Path to the overlay's TXN insts for one shape, dumping it if absent.""" - path = TXN_CACHE / f"txn_{M}_{K}_{N}.bin" - if not path.is_file(): - path.parent.mkdir(parents=True, exist_ok=True) - subprocess.run([str(TXN_DUMP), str(M), str(K), str(N), str(path)], check=True) - return path - - -# --------------------------------------------------------------------------- -# Competitors. Each returns (run, c_bo, label, xclbin_path) with the buffers -# already bound, so the timed section is nothing but the dispatch. -# --------------------------------------------------------------------------- -def _bind(xclbin, insts, args): - from aie.utils.npukernel import NPUKernel - - handle = aie_utils.DefaultNPURuntime.load(NPUKernel(str(xclbin), str(insts))) - return lambda: aie_utils.DefaultNPURuntime.run(handle, list(args)) - - -def setup_iron(M, K, N, A, B, ctx): - op = FLMGEMM(M=M, K=K, N=N, context=ctx) - built_before = _artifacts_exist(op, ctx) - t0 = time.perf_counter() - op.compile() - compile_s = time.perf_counter() - t0 - c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) - run = op.get_callable() - args = [ - XRTTensor.from_torch(A.flatten()), - XRTTensor.from_torch(op.pack_B(B).flatten()), - c_bo, - ] - return _Competitor( - "iron", - lambda: run(*args), - c_bo, - Path(op.xclbin_artifact.filename), - None if built_before else compile_s, - BUDGET_CONV_EVEN, - ) - - -def setup_gemm(M, K, N, A, B, ctx): - # Left at the operator's defaults, which are the same emulated-bfp16 mmul - # and conv_even rounding flm_gemm uses -- a like-for-like comparison, not - # flm_gemm against a more accurate and necessarily slower configuration. - op = GEMM(M=M, K=K, N=N, context=ctx) - built_before = _artifacts_exist(op, ctx) - t0 = time.perf_counter() - op.compile() - compile_s = time.perf_counter() - t0 - c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) - run = op.get_callable() - # b_col_maj defaults False, so B goes in as plain row-major (K, N). - args = [XRTTensor.from_torch(A.flatten()), XRTTensor.from_torch(B.flatten()), c_bo] - return _Competitor( - "gemm", - lambda: run(*args), - c_bo, - Path(op.xclbin_artifact.filename), - None if built_before else compile_s, - BUDGET_CONV_EVEN, - ) - - -def setup_flm(M, K, N, A, B, ctx): - c_bo = XRTTensor((M, N), dtype=np.dtype("bfloat16")) - # The shipped overlay's host contract is C, A, B -- not IRON's A, B, C. - args = [ - c_bo, - XRTTensor.from_torch(A.flatten()), - XRTTensor.from_torch(flm_pack_B(B).flatten()), - ] - run = _bind(FLM_XCLBIN, flm_insts(M, K, N), args) - # Prebuilt and shipped: there is no compile to time. - return _Competitor("flm", run, c_bo, FLM_XCLBIN, None, BUDGET_FLOOR) - - -class _Competitor: - def __init__(self, name, run, c_bo, xclbin, compile_s, budget): - self.name = name - self.run = run - self.c_bo = c_bo - self.xclbin = xclbin - self.compile_s = compile_s - self.budget = budget - self.round_medians = [] - - 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): - ts = [] - for _ in range(ITERS): - t0 = time.perf_counter() - self.run() - ts.append((time.perf_counter() - t0) * 1e6) - 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 - - -def _artifacts_exist(op, ctx): - """Whether this operator's xclbin is already built in ctx's build dir. - - Compile time is only meaningful on a genuine miss; on a hit ``compile()`` - returns in milliseconds and reporting that as a build time would be a lie. - """ - if not op.artifacts: - op.set_up_artifacts() - return (Path(ctx.build_dir) / f"{op.name}.xclbin").is_file() - - -# --------------------------------------------------------------------------- -# The benchmark -# --------------------------------------------------------------------------- -@pytest.mark.metrics( - IronLatency=r"iron latency \(us\): (?P[\d\.]+)", - FLMLatency=r"flm latency \(us\): (?P[\d\.]+)", - GEMMLatency=r"gemm latency \(us\): (?P[\d\.]+)", - SpeedupVsFLM=r"speedup vs flm: (?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. - IronErr=r"iron err/mass: (?P[\d\.e\+-]+)", - FLMErr=r"flm err/mass: (?P[\d\.e\+-]+)", - GEMMErr=r"gemm err/mass: (?P[\d\.e\+-]+)", - IronThroughput=r"iron throughput: (?P[\d\.e\+-]+) GFLOP/s", - IronJitterPct=r"iron jitter \(%\): (?P[\d\.]+)", - IronXclbinKB=r"iron xclbin \(KB\): (?P[\d\.]+)", - IronCompileTime=r"iron compile \(s\): (?P[\d\.]+)", -) -@pytest.mark.parametrize("model,proj,M,K,N", get_params()) -def test_flm_gemm_vs_flm(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. - competitors = [ - setup_iron(M, K, N, A, B, aie_context), - setup_flm(M, K, N, A, B, aie_context), - ] - # IRON's generic GEMM still has the 20-bit mega_row stride limitation that - # flm_gemm fixed: at M>256 with a 10240-wide N its C descriptor wants a - # stride of 2621440 elements and aiecc rejects the build outright. That is - # a property of that operator, not of the shape, and it is caught at - # compile time rather than hanging -- so report the shape with the - # competitors that do build instead of losing flm_gemm's own numbers for - # it. Only that specific rejection is tolerated; anything else still fails. - try: - competitors.append(setup_gemm(M, K, N, A, B, aie_context)) - except RuntimeError as e: - if "aie.dma_bd" not in str(e) or "exceeds the" not in str(e): - raise - print("gemm unavailable: descriptor stride exceeds the shim's 20-bit step") - - bad = [c for c in competitors 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 competitors: - for _ in range(WARMUP): - c.run() - # Round-robin, never all of one then all of another. - for _ in range(ROUNDS): - for c in competitors: - c.time_round() - - by_name = {c.name: c for c in competitors} - iron = by_name["iron"] - - print() - for c in competitors: - print(f"{c.name} latency (us): {c.us:.1f}") - print(f"{c.name} err/mass: {c.err:.3e}") - print(f"speedup vs flm: {by_name['flm'].us / iron.us:.3f}") - if "gemm" in by_name: - print(f"speedup vs gemm: {by_name['gemm'].us / iron.us:.3f}") - print(f"iron throughput: {2.0 * M * K * N / (iron.us * 1e-6) / 1e9:.6e} GFLOP/s") - print(f"iron jitter (%): {iron.jitter_pct:.2f}") - print(f"iron xclbin (KB): {iron.xclbin.stat().st_size / 1024:.1f}") - if iron.compile_s is not None: - print(f"iron compile (s): {iron.compile_s:.1f}") - print() diff --git a/iron/operators/flm_gemm/op.py b/iron/operators/flm_gemm/op.py deleted file mode 100644 index 109e3c6a5..000000000 --- a/iron/operators/flm_gemm/op.py +++ /dev/null @@ -1,397 +0,0 @@ -# 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 -import torch -from typing import ClassVar, Dict - -from iron.common import ( - MLIROperator, - AIERuntimeArgSpec, - KernelObjectArtifact, - SourceArtifact, - PythonGeneratedMLIRArtifact, - DesignGenerator, -) -from iron.common.device_utils import get_kernel_dir -import aie.utils as aie_utils - -from iron.operators.flm_gemm.design import ( - COLS, - C_DEPTH, - CT_OUT_LEN, - EPILOGUE_MODES, - K_TILE, - MIN_K, - MIN_M, - M_TILE, - N_TILE_DEFAULT, - S, - T, -) - - -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. flm_gemm calls ``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. - # TRUNCATION, not rounding -- that is what AIE2P does, and round-to-nearest - # here measures 6.18e-03 against truncation's 2.69e-04. - 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)) - - -@dataclass -class FLMGEMM(MLIROperator): - """AIE-accelerated bf16 GEMM on a fixed 4x8 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. - """ - - M: int - K: int - N: int - # "none" | "gelu" | "silu" | "sigmoid", fused into the C drain. - epilogue: str = field(default="none", repr=False) - # Optional (min, max) applied after the activation. - clamp: tuple[float, float] | None = field(default=None, repr=False) - # n tile width. 64 halves the mmul's accumulator traffic per mac; 128 - # halves A fetches instead and wins only when small K makes the operator - # DMA-bound. See README.md. - tile_n: int | None = field(default=None, repr=False) - # A-tile rows, decoupled from the accumulator's M_TILE (asymmetric tile - # buffering). None means symmetric (T_MA == M_TILE). - tile_ma: int | None = field(default=None, repr=False) - # "conv_even" (round to nearest even) or "floor" (truncate). The core - # powers up in floor, and the design this was ported from never sets the - # mode, so "floor" reproduces its arithmetic exactly -- at ~40x the error, - # because truncation biases every conversion the same way and the bias - # accumulates over the K reduction instead of cancelling. - rounding: str = field(default="conv_even", repr=False) - context: object = field(default=None, repr=False) - - _name_aliases: ClassVar[Dict[str, str]] = {**MLIROperator._name_aliases} - - def __post_init__(self): - if self.tile_n is None: - self.tile_n = self._default_tile_n(self.K) - # 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, MIN_M), - ("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}") - if self.epilogue not in EPILOGUE_MODES: - raise ValueError( - f"epilogue must be one of {sorted(EPILOGUE_MODES)}, " - f"got {self.epilogue!r}" - ) - if self.clamp is not None: - lo, hi = self.clamp - if lo > hi: - raise ValueError(f"clamp min ({lo}) must be <= max ({hi})") - if self.rounding not in ("conv_even", "floor"): - raise ValueError( - f"rounding must be 'conv_even' or 'floor', got {self.rounding!r}" - ) - - MLIROperator.__init__(self, context=self.context) - - @staticmethod - def _default_tile_n(K: int) -> int: - """Pick the n tile from the shape. - - 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, and that is set by how - much K there is to reduce over: with a single k iteration there is too - little compute to hide the extra A traffic. Measured ~20% for n=64 at - K >= 1024 and ~9% the other way at K = 512. - """ - return 128 if K // K_TILE <= 1 else 64 - - @property - def name(self) -> str: - # epilogue/clamp are repr=False so the plain path keeps a stable name, - # but they change the emitted kernel, so the variants must not share an - # artifact name: in a shared build dir a cached plain build would - # otherwise satisfy a fused op and silently skip the activation. - base = super().name - if self.epilogue != "none": - base = f"{base}_epi{self.epilogue}" - if self.clamp is not None: - base = f"{base}_clamp{self._clamp_tag}" - if self.rounding != "conv_even": - base = f"{base}_{self.rounding}" - if self.tile_n != self._default_tile_n(self.K): - base = f"{base}_tn{self.tile_n}" - # The RESOLVED height, not just an explicit override: it changes the - # emitted MLIR and the kernel object, so a build dir holding another - # value's artifacts must not satisfy this one. - if self._tile_ma != M_TILE: - base = f"{base}_ma{self._tile_ma}" - return base - - @property - def _clamp_tag(self) -> str: - lo, hi = self.clamp - return f"{lo:g}_{hi:g}".replace("-", "m").replace(".", "p") - - @property - def _epilogue_artifact(self) -> str: - obj = f"flm_gemm_epilogue_{self.epilogue}" - if self.clamp is not None: - obj = f"{obj}_clamp{self._clamp_tag}" - if self.rounding != "conv_even": - obj = f"{obj}_{self.rounding}" - return f"{obj}.o" - - @property - def _rounding_flags(self) -> list[str]: - """Applies to both kernels: the mmul and the epilogue's f32->bf16 - store are both conversions and must agree.""" - return ["-DFLM_GEMM_ROUND_FLOOR"] if self.rounding == "floor" else [] - - @property - def _epilogue_source(self): - return self.context.base_dir / "aie_kernels" / "aie2p" / "flm_gemm_epilogue.cc" - - @property - def _epilogue_flags(self) -> list[str]: - """Compile flags for the epilogue.""" - flags = [ - f"-DFLM_GEMM_OUT_CHUNK={CT_OUT_LEN}", - f"-DFLM_GEMM_C_DEPTH={C_DEPTH}", - f"-DFLM_GEMM_EPILOGUE_MODE={EPILOGUE_MODES[self.epilogue]}", - ] - 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 += [ - "-DFLM_GEMM_CLAMP=1", - f"-DFLM_GEMM_CLAMP_MIN={float(lo)!r}f", - f"-DFLM_GEMM_CLAMP_MAX={float(hi)!r}f", - ] - return flags + self._rounding_flags - - @property - def _ablate_mmul(self) -> bool: - """ABLATION: FLM_NULL_MMUL=1 nulls the multiply, leaving all data - movement. Threaded into the object AND operator names because the - build cache is keyed on filename.""" - import os - - return os.environ.get("FLM_NULL_MMUL", "") == "1" - - @property - def _tile_ma(self) -> int: - """Resolved A-tile height. design.py picks the default, and it MUST be - the same value the kernel is compiled with -- the design sizes the A - object from it while the kernel derives the mmul's rowA from it, so a - mismatch reads past the buffer and produces garbage rather than a build - error.""" - from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N, _default_l1 - - if self.tile_ma is not None: - return self.tile_ma - return _default_l1(self.tile_n, CT_MAX_K_FOR_N[self.tile_n])[0] - - @property - def _kernel_object(self) -> str: - rnd = "" if self.rounding == "conv_even" else f"_{self.rounding}" - ma = f"_ma{self._tile_ma}" + ("_nomm" if self._ablate_mmul else "") - return f"flm_gemm_{M_TILE}x{K_TILE}x{self.tile_n}{rnd}{ma}.o" - - def get_mlir_artifact(self): - return PythonGeneratedMLIRArtifact( - f"{self.name}.mlir", - DesignGenerator( - self.operator_dir / "design.py", - "flm_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._kernel_object, - "epilogue_object": self._epilogue_artifact, - "trace_size": 0, - }, - ), - ) - - def get_kernel_artifacts(self): - # The mmul and the epilogue are both aie2p-only: the mmul relies on the - # bf16 emulation path and the grid needs 8 columns. - kernel_dir = get_kernel_dir() - if kernel_dir != "aie2p": - raise NotImplementedError( - f"flm_gemm is only available on NPU2 (aie2p); got {kernel_dir!r}" - ) - base_dir = self.context.base_dir - aie2p = base_dir / "aie_kernels" / "aie2p" - - artifacts = [ - KernelObjectArtifact( - self._kernel_object, - dependencies=[SourceArtifact(aie2p / "flm_gemm.cc")], - extra_flags=[ - f"-DFLM_GEMM_TILE_M={M_TILE}", - f"-DFLM_GEMM_TILE_K={K_TILE}", - f"-DFLM_GEMM_TILE_N={self.tile_n}", - f"-DFLM_GEMM_TILE_MA={self._tile_ma}", - "-DFLM_GEMM_BFP16_B", - *(["-DFLM_GEMM_NULL_MMUL"] if self._ablate_mmul else []), - # The r=8 mmul shape this design uses only exists on the - # bfp16-emulated path; without this the kernel will not - # compile. - "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", - ] - + self._rounding_flags, - ), - ] - artifacts.append( - KernelObjectArtifact( - self._epilogue_artifact, - dependencies=[SourceArtifact(self._epilogue_source)], - extra_flags=self._epilogue_flags, - ) - ) - return artifacts - - def pack_B(self, B): - """Reorder and quantize a row-major ``(K, N)`` weight matrix into the - layout the B fill expects. Returns a flat uint8 tensor of bfp16ebs8 - blocks, NOT a bf16 tensor. - - The quantization is not a loss this adds. The mmul only multiplies - bfp16, so the bf16 path converts B inside every mac call; doing it here - hoists a rounding that already happened and leaves the arithmetic - bit-identical. It also makes B 9 bytes per 8 values instead of 16, - which is the point -- this operator is data-movement bound. - - Each ``K_TILE x N_TILE`` tile is emitted in t-block-major order -- the - odometer ``(n//T, k%S, k//S, n%T)``, outermost first -- with tiles - ordered by column stripe and then by k-block, so each fill is one - contiguous read. - - This is deliberately the caller's job. The same reorder is expressible - as a strided descriptor over an unpacked B, but its innermost run is - then T=8 bf16 = 16 bytes, turning each 128 KB transfer into 8192 - scattered bursts -- measured 5.4x slower end to end, and the whole of - this operator's gap against the design it was ported from, which packs - its weights on the host for the same reason. Weights are packed once - and reused across dispatches, so the cost belongs here. - """ - K, N = B.shape - N_TILE = self.tile_n - if K % K_TILE or N % N_TILE: - raise ValueError( - f"B ({K}, {N}) must tile to ({K_TILE}, {N_TILE}) to be packed" - ) - # Emit the FINAL consumption order, not an intermediate one. The old - # layout left a 4-dimension scatter for the memtile's dims_from_stream - # to finish, which put the n-block index outside the k-slice index and - # so cost two descriptor dimensions on the way back out. Packing all - # the way here makes both B hops linear, which is what leaves room for - # a k-slice deep enough to halve the accumulator traffic (CT_MAX_K=128) - # while B is also memtile-resident. - from iron.operators.flm_gemm.design import CT_MAX_K_FOR_N - - CT_K = CT_MAX_K_FOR_N[N_TILE] - col_a = CT_K // S - t = 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) - # -> (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 below match the kernel's. - # Grouping the shared exponent over 8 consecutive k (for one n) is - # verified: grouping over n instead measures 1.95e-02 against this - # layout's 2.69e-04. - t = t.permute(4, 0, 1, 5, 2, 6, 3).reshape(-1, 8).contiguous() - return _f32_to_bfp16ebs8( - t.float().numpy(), round_conv_even=self.rounding == "conv_even" - ) - - @staticmethod - def unpack_B_size(K, N): - """Bytes ``pack_B`` returns for a ``(K, N)`` weight matrix.""" - return K * N // 8 * 9 - - def get_arg_spec(self): - return [ - AIERuntimeArgSpec("in", (self.M, self.K)), # A - # B arrives pre-packed AND quantized by pack_B: bfp16ebs8, which is - # 9 bytes per 8 values rather than bf16's 16. Declared in bytes so - # the buffer is sized from what pack_B actually returns -- a - # (K, N) bf16 spec would over-allocate the largest buffer by 1.78x. - AIERuntimeArgSpec( - "in", (self.unpack_B_size(self.K, self.N),), dtype=np.uint8 - ), # 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/test.py b/iron/operators/flm_gemm/test.py deleted file mode 100644 index f9a10ce87..000000000 --- a/iron/operators/flm_gemm/test.py +++ /dev/null @@ -1,163 +0,0 @@ -#!/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 iron.operators.flm_gemm.op import FLMGEMM -from iron.operators.flm_gemm.reference import generate_golden_reference -from iron.common.test_utils import run_test - - -def get_params(): - dev = aie_utils.get_current_device() - # The design is a fixed 4x8 grid, so it needs all 8 columns. - if dev.cols < 8 or dev.resolve().name != "npu2": - return [] - - # N values that are NOT a multiple of N_TILE*COLS=1024 exercise the - # trailing partial column-block, where some columns compute it and the - # rest only drain the A broadcast. Real transformer o/down projections - # have N = model dim, so they always land here: 1536 leaves 4 active - # columns, 2560 leaves 4, and 128 leaves just 1. - # fmt: off - # 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 bit for bit; 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 - ] - # 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 - - -@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_flm_gemm(M, K, N, epilogue, clamp, rounding, aie_context): - # Keep the activation tests in the range where the curve is not flat. - scale = 4.0 if epilogue == "none" else 0.5 - golden_ref = generate_golden_reference( - M=M, K=K, N=N, epilogue=epilogue, clamp=clamp, scale=scale - ) - - - operator = FLMGEMM( - M=M, - K=K, - N=N, - epilogue=epilogue, - clamp=clamp, - rounding=rounding, - context=aie_context, - ) - - input_buffers = { - "A": golden_ref["input"].flatten(), - # B is consumed pre-packed; see FLMGEMM.pack_B. - "B": operator.pack_B(golden_ref["input_b"]), - } - output_buffers = {"C": golden_ref["output"].flatten()} - - # This design's r=8 mmul shape exists only on the bfp16-emulated path, so - # its error budget is that of an emulated GEMM, not of the exact one the - # GEMM operator's test asserts on (that test opts into r=4 via - # emulate_bf16_mmul_with_bfp16=False, which is not available here). - # - # A pure relative tolerance cannot work: 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. So the error is bounded in ABSOLUTE terms - # against the accumulated mass, which is what bfp16 error actually scales - # with. Measured on this data: mean |err| is 0.00042 of the mass and the - # worst element 0.0025 -- both marginally better than the GEMM operator run - # in the same emulated mode (0.00044 / 0.0031), so the budget below is not - # papering over a regression in this port. The bound is tight enough to - # have caught a real bug: leaving the core in its default floor rounding - # mode pushes mean error to 0.0099 of mass, ~7x over. - mass = K * golden_ref["input"].abs().float().mean() * ( - golden_ref["input_b"].abs().float().mean() - ) - # - # floor rounding truncates rather than rounding to nearest, so its bias - # accumulates over the K reduction instead of cancelling: ~0.0099 of mass - # rather than ~0.00042, measured, and bit-identical to the shipped overlay. - # It gets a bound to match; holding it to the conv_even budget would just - # fail. - budget = 0.05 if rounding == "floor" else 0.004 - errors, latency_us, bandwidth_gbps = run_test( - operator, - input_buffers, - output_buffers, - rel_tol=0.04, - abs_tol=float(budget * mass), - ) - - 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_flm_gemm_split_leg_windowing(aie_context): - # K or N > ~8191 at M > 256 makes the mega_row stride overflow the shim - # BD's 20-bit 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 iron.operators.flm_gemm.design import SHIM_BDS, SHIM_TASK_QUEUE - - worst = 1 + 2 * SHIM_TASK_QUEUE - assert worst <= SHIM_BDS, ( - f"a fully split block needs {worst} shim BDs of {SHIM_BDS}; " - "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. - FLMGEMM(M=512, K=10240, N=10240, context=aie_context).compile() diff --git a/requirements.txt b/requirements.txt index 7f7c43b37..622f716a0 100755 --- a/requirements.txt +++ b/requirements.txt @@ -13,6 +13,9 @@ --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple +# Keep these two in step: each mlir-aie wheel is built against one Peano +# nightly, pinned in that revision's utils/peano-requirements.txt. The pair +# below is the one dev85 declares. mlir_aie==1.4.3.dev85+gdf48abc llvm-aie==22.0.0.2026090701+3e93bf7b From 81b923d1ee48ecee1a4fcf7a47e0df59a8c462e6 Mon Sep 17 00:00:00 2001 From: Erika Hunhoff Date: Thu, 10 Sep 2026 15:08:18 -0600 Subject: [PATCH 31/31] flm_gemm: address Copilot review on PR #195 - design.py: drop tile_n=256 from CT_MAX_K_FOR_N -- it always overflows L1 (the accumulator alone exceeds the budget), so it was advertised as valid but unbuildable. - design.py: re-validate L1 B depth against a caller-overridden tile_ma instead of reusing the depth picked for the default fit, which could silently overflow L1. - packing.py: pack_b's non-bfp16 branches now cast to bf16 explicitly, matching the documented return type instead of preserving whatever dtype the caller passed in. - gemm/test.py: add an extensive execution test for the two-sided split path (M=512, K=N=10240) -- the existing test only compiled it, which can't exercise the runtime hang the split guards against. - mm_prebuilt/test.py: add remainder-only and full+remainder N shapes to cover the trailing-column path, and use a domain-scaled tolerance for sigmoid/clamp so an all-zero result can no longer pass. - gemm/reference.py: fix the docstring's description of operation order to match the kernel (bf16 cast happens before the epilogue, not after). Co-Authored-By: Claude --- iron/operators/flm/gemm/design.py | 41 +++++++++++++++++++++----- iron/operators/flm/gemm/reference.py | 11 +++++-- iron/operators/flm/gemm/test.py | 41 ++++++++++++++++++++++++++ iron/operators/flm/mm_prebuilt/test.py | 29 +++++++++++++----- iron/operators/flm/packing.py | 15 +++++++--- 5 files changed, 115 insertions(+), 22 deletions(-) diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index d3c632eb9..8e371d7db 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -68,7 +68,10 @@ # 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. -CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32, 256: 16} +# n=256 is deliberately absent: at that width the f32 accumulator alone +# (M_TILE * 256 * 4 = 65536 bytes) already exceeds L1_BUDGET, 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. 8/8/8 on both architectures today; register_tiling() is the # single source of truth and returns exactly these. R, S, T = 8, 8, 8 @@ -211,6 +214,27 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes): 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): + """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. + """ + 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 <= L1_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, @@ -255,12 +279,15 @@ def gemm( # 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. - _t_ma_fit, L1_B_DEPTH = _default_l1(N_TILE, CT_MAX_K, b_elem_bytes) - T_MA = _t_ma_fit if tile_ma is None else 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}" - ) + if tile_ma is None: + T_MA, L1_B_DEPTH = _default_l1(N_TILE, CT_MAX_K, b_elem_bytes) + 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) RHO = M_TILE // T_MA OVERLAP = OVERLAP_DEFAULT if overlap is None else overlap K_DIV_CT_K_MAX = K_TILE // CT_MAX_K diff --git a/iron/operators/flm/gemm/reference.py b/iron/operators/flm/gemm/reference.py index 7872824b1..3dae7f21b 100644 --- a/iron/operators/flm/gemm/reference.py +++ b/iron/operators/flm/gemm/reference.py @@ -8,9 +8,14 @@ def reference(input_a, input_b, epilogue="none", clamp=None): """CPU reference ``C = clamp(activation(A @ B))``. - The matmul is accumulated in fp32 to mirror the kernel's f32 accumulator, - then cast back to the input dtype at the end, which is where the kernel - converts too. + The matmul is accumulated in fp32 to mirror the kernel's f32 accumulator. + This is an idealized reference, not operation-for-operation matching: it + applies the activation and clamp in fp32 and casts to the input dtype only + at the end, whereas the kernel (``mm_fused_epilogue_chunk`` in + ``mm_fused_epilogue.cc``) converts the accumulator to bf16 first and then + applies the activation and clamp in bf16. The two are close enough that the + per-test tolerances absorb the difference, but do not expect a bit-exact + match. ``gelu`` is the sigmoid approximation ``x * sigmoid(1.702x)``, matching the kernel -- NOT torch's erf-exact gelu, and not the tanh approximation the diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index c8155deed..a330795f4 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -203,6 +203,47 @@ def test_gemm_split_leg_windowing(aie_context): GEMM(M=512, K=10240, N=10240, context=aie_context).compile() +@pytest.mark.extensive +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 can't exercise. This dispatches the same shape on + hardware and checks the result. + """ + 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) + + input_buffers = { + "A": golden_ref["input"].flatten(), + "B": operator.pack_B(golden_ref["input_b"]), + } + output_buffers = {"C": golden_ref["output"].flatten()} + + # Same mass-based bound as test_gemm, at the non-floor budget for whichever + # architecture this runs on. + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + budget = ( + 0.0002 if aie_utils.get_current_device().resolve().name == "npu1" else 0.004 + ) + errors, _latency_us, _bandwidth_gbps = run_test( + operator, + input_buffers, + output_buffers, + rel_tol=0.04, + abs_tol=float(budget * mass), + ) + 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``. diff --git a/iron/operators/flm/mm_prebuilt/test.py b/iron/operators/flm/mm_prebuilt/test.py index 243fb0ba8..dfcd72c51 100644 --- a/iron/operators/flm/mm_prebuilt/test.py +++ b/iron/operators/flm/mm_prebuilt/test.py @@ -39,8 +39,10 @@ @pytest.mark.parametrize( "M,K,N,epilogue,clamp", [ - (256, 512, 1024, "none", None), - (512, 1024, 2048, "none", None), + (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), (256, 512, 1024, "sigmoid", None), @@ -66,16 +68,27 @@ def test_mm_prebuilt(M, K, N, epilogue, clamp, aie_context): # Same absolute-mass bound as flm.gemm/test.py, for the same reason: with # signed A the K-sum cancels by ~sqrt(K), leaving near-zero outputs # relatively uncheckable under a plain relative tolerance. - mass = ( - K - * golden_ref["input"].abs().float().mean() - * golden_ref["input_b"].abs().float().mean() - ) + # + # sigmoid and clamp are the exception: their output is bounded to a known, + # narrow range (sigmoid to (0, 1), this clamp to (-2, 2)), far smaller than + # the mass-based bound above -- which would then pass even an all-zero + # result. Scale the tolerance to the actual output domain for those instead. + if epilogue == "sigmoid": + abs_tol = BUDGET_FLOOR + elif clamp is not None: + abs_tol = BUDGET_FLOOR * (clamp[1] - clamp[0]) + else: + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + abs_tol = float(BUDGET_FLOOR * mass) errors, latency_us, bandwidth_gbps = run_test( operator, input_buffers, output_buffers, rel_tol=0.04, - abs_tol=float(BUDGET_FLOOR * mass), + abs_tol=abs_tol, ) assert not errors, "Test failed" diff --git a/iron/operators/flm/packing.py b/iron/operators/flm/packing.py index c02708ac5..538ecfeb3 100644 --- a/iron/operators/flm/packing.py +++ b/iron/operators/flm/packing.py @@ -110,10 +110,17 @@ def pack_b( if not bfp16: if overlay_order: # -> (cb, kb, kslice, tb, s_in, i, t_in) - return blocked.permute(4, 0, 1, 5, 3, 2, 6).reshape(-1).contiguous() - # -> (cb, kb, kslice, tb, i, s_in, t_in) - # Row-major s x t within the block, which is what the plain mmul loads. - return blocked.permute(4, 0, 1, 5, 2, 3, 6).reshape(-1).contiguous() + 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