Skip to content

flm.GEMM: take M, K, N and the activation as runtime parameters #200

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
hunhoffe wants to merge 10 commits into
base: devel
Choose a base branch
from
Draft

flm.GEMM: take M, K, N and the activation as runtime parameters #200

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f94d169
flm.GEMM: take M, K, N and the activation as runtime parameters
andrej Sep 11, 2026
f424116
flm_gemm: reconcile docs after merging the RTP work
hunhoffe Sep 11, 2026
92350c4
flm_gemm: make the clamp bounds runtime parameters
hunhoffe Sep 11, 2026
801fea7
flm_gemm: add m_chunk, folding row-blocks into one B fetch
hunhoffe Sep 11, 2026
6b2af06
deps: bump mlir-aie to the v1.4.3 release
hunhoffe Sep 14, 2026
4fbb53b
flm_gemm: name ct_max_k in the artifact stems
hunhoffe Sep 14, 2026
c378b27
flm_gemm: send four runtime-parameter words instead of ten
hunhoffe Sep 14, 2026
0103a74
flm_gemm: fix benchmark.py's usage docstring
hunhoffe Sep 14, 2026
b2379c0
flm_gemm: retire split-leg transfers rolling, not in windows
hunhoffe Sep 14, 2026
560b0b8
flm_gemm: cut the comment volume roughly in half
hunhoffe Sep 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 120 additions & 95 deletions aie_kernels/generic/mm_fused.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,14 @@
// SPDX-License-Identifier: Apache-2.0

// bf16 GEMM compute kernel. Each compute tile owns an m x n slice of C and
// accumulates over K into an f32 accumulator that stays in L1 for the whole
// reduction.
//
// Three entry points, each called once per iteration of a loop nest that lives
// in the design (iron/operators/flm/gemm/design.py) rather than here:
// accumulates over K into an f32 accumulator that stays in L1.
//
// mm_fused_acc_init zero the accumulator, once per output tile
// mm_fused_k_step multiply one A band by one B chunk into it
// mm_fused_epilogue_chunk drain one chunk of it to a bf16 C object
//
// The nest lives in the design so that every level of it has an ObjectFifo
// acquire point, a fifo consumer having to acquire once per object.
//
// Tile geometry arrives as -D flags from design.py, which is the single source
// of truth for it: the same constants size the design's buffers and set its
// unroll factors.
// The loop nest is in design.py, not here, so every level has an ObjectFifo
// acquire point. Tile geometry arrives from there as -D flags.
#include "../aie_kernel_utils.h"
#include "activations.h"
#include "mm_fused_mmul.h"
Expand All @@ -35,43 +27,32 @@

// Epilogue selection. 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid, matching
// Epilogue.mode in design.py.
#ifndef MM_FUSED_EPILOGUE_MODE
#define MM_FUSED_EPILOGUE_MODE 0
#ifndef MM_FUSED_EPILOGUE_MODE_MASK
#define MM_FUSED_EPILOGUE_MODE_MASK 0xF
#endif
// Whether a clamp is compiled in at all: a capability, not a selection, since
// the clamped path costs program memory. The bounds are runtime.
#ifndef MM_FUSED_CLAMP
#define MM_FUSED_CLAMP 0
#endif
#ifndef MM_FUSED_CLAMP_MIN
#define MM_FUSED_CLAMP_MIN 0.0f
#endif
#ifndef MM_FUSED_CLAMP_MAX
#define MM_FUSED_CLAMP_MAX 0.0f
#endif

namespace
{
constexpr int M = MM_FUSED_TILE_M;
// Asymmetric tile buffering: the A tile spans MA rows while the accumulator
// spans M, so the core folds RHO = M / MA A bands into one C tile before
// releasing it. A dies as soon as it is consumed while C must live across the
// whole K reduction, so sizing both to M would pay the peak L1 cost twice.
// MA == M is the symmetric case.
// Asymmetric tile buffering: A spans MA rows and the accumulator M, so the
// core folds RHO = M / MA A bands into one C tile before releasing it. A dies
// on consumption while C lives across the 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.
// From arXiv:2511.16041, "Can Asymmetric Tile Buffering Be Beneficial?";
// reference AIE implementation in Xilinx/mlir-aie PR #3076. Those configs
// accumulate in bf16/bfp16, affording larger C tiles; this one keeps an f32
// accumulator, so the freed L1 buys a deeper k slice instead.
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.
// Register tiling, and how much of K a tile holds at a time. Both are
// design.py's to choose; CT_K trades against the n width for a fixed budget.
constexpr int R = MM_FUSED_R;
constexpr int S = MM_FUSED_S;
constexpr int T = MM_FUSED_T;
Expand All @@ -91,44 +72,86 @@ 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
// The core powers up in rounding_mode::floor, so a converting kernel must
// choose explicitly. Truncation biases every conversion the same direction, so
// the error accumulates over the K reduction instead of cancelling -- ~1% of
// the result, against ~0.02% for round-to-nearest-even, which is far more than
// the bfp16 emulation itself costs. Every entry point that converts sets it:
// the mmul and the epilogue's f32->bf16 store, both below.
//
// Flag name and polarity follow mm.cc, so the two kernels are configured the
// same way; the operator passes -DROUND_CONV_EVEN by default.
// the error accumulates over the K reduction: ~1% of the result against ~0.02%
// for round-to-nearest-even, far more than the bfp16 emulation costs. Flag
// name and polarity follow mm.cc.
#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
// One activation's inner loop. Templated so each mode compiles branch-free;
// mm_fused_epilogue_chunk selects between them once per chunk.
template <int MODE, bool CLAMP>
static inline void
epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float clamp_min, float clamp_max)
{
aie::vector<float, V> lo, hi;
if constexpr (CLAMP) {
lo = aie::broadcast<float, V>(clamp_min);
hi = aie::broadcast<float, V>(clamp_max);
}

AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V)
for (int j = 0; j < CHUNK / V; j++) {
// f32 through the activation and clamp, converted exactly once on the
// store. Converting first would round twice and let the activation's
// slope amplify the first rounding.
aie::vector<float, V> f = aie::load_v<V>(src + j * V);
if constexpr (MODE == 1)
f = gelu_vec<V>(f);
else if constexpr (MODE == 2)
f = silu_vec<V>(f);
else if constexpr (MODE == 3)
f = sigmoid_vec<V>(f);
if constexpr (CLAMP)
f = aie::max(aie::min(f, hi), lo);
aie::accum<accfloat, V> out;
out.from_vector(f);
// The assignment is the conversion: to_v16bfloat16 yields a raw
// v16bfloat16, not an aie::vector.
aie::vector<bfloat16, V> v = to_v16bfloat16(out);
aie::store_v(y_out + j * V, v);
}
}

// Pick the clamped or unclamped instantiation. Only one exists unless
// MM_FUSED_CLAMP compiled the clamp in.
template <int MODE>
static inline void epilogue_dispatch(bfloat16 *__restrict y_out,
const float *__restrict src,
int32_t clamp_enabled,
float clamp_min,
float clamp_max)
{
#if MM_FUSED_CLAMP
if (clamp_enabled) {
epilogue_body<MODE, true>(y_out, src, clamp_min, clamp_max);
return;
}
#endif
epilogue_body<MODE, false>(y_out, src, clamp_min, clamp_max);
}
} // 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.
// Zero the f32 accumulator before the k loop accumulates into it. A bias is
// deliberately unsupported: initialising from one would consume an extra
// object through the B fifo's handshake, desynchronising it into a hang.
void mm_fused_acc_init(float *y_acc)
{
// zero_vectorized brackets itself in event0/event1 for tracing.
zero_vectorized<float, M, N>(y_acc);
}

// One step of the k loop: one B chunk multiplied against one A band,
// accumulated into y_acc.
// One step of the k loop: one B chunk against one A band, 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.
// Takes no locks -- the A and B fifos own the handshake, so the core body
// acquires around this call. mm_fused_b_elem_t is bfp16ebs8 or bfloat16
// depending on the architecture, but the signature is the same either way.
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);
Expand All @@ -137,54 +160,56 @@ void mm_fused_k_step(bfloat16 *a_buf, mm_fused_b_elem_t *b_buf, float *y_acc, in
mm_fused_mmul_2x2<(MA / R), (CT_K / S), (N / T), R, S, T>(a_buf, b_buf, y_acc + band * (MA * N));
}

// The output stage: convert chunk (outer * C_DEPTH + half) of the f32
// accumulator into a bf16 C object the core body has already acquired from the
// C ObjectFifo, optionally applying an activation and a clamp on the way out.
//
// Fusing the activation here is the point: the values are already in registers
// after the f32 -> bf16 conversion, so gelu/silu/sigmoid costs one more vector
// op per 16 elements instead of a separate pass over L1 (which is what chaining
// a standalone activation operator after a GEMM would cost). The mode and clamp
// are compile-time, so the inner loop below is branch-free.
// Convert chunk (outer * C_DEPTH + half) of the f32 accumulator into a bf16 C
// object the core body already acquired, applying an activation and clamp on
// the way out.
//
// 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)
// Fusing the activation is the point: the values are already in registers, so
// gelu/silu/sigmoid costs one more vector op per 16 elements rather than a
// separate pass over L1. The mode is runtime, tested once per chunk so the
// inner loop stays branch-free; the cost is program memory, since every mode
// in the mask is compiled in. The clamp splits the same way, its bounds
// arriving as raw int32 because npu_write_rtp only writes i32 words. The chunk
// index comes in two parts because the core body unrolls the drain.
void mm_fused_epilogue_chunk(bfloat16 *y_out,
float *y_acc,
int32_t outer,
int32_t half,
int32_t mode,
int32_t clamp_enabled,
int32_t clamp_min_bits,
int32_t clamp_max_bits)
{
// The store below is a conversion, so it obeys the same rounding mode the
// mmul does and must agree with it.
::aie::set_rounding(round_mode);
const float *__restrict src = y_acc + (outer * C_DEPTH + half) * CHUNK;
// __builtin_bit_cast, not memcpy: memcpy leaves an unresolved external
// call here rather than folding to a register move.
const float clamp_min = __builtin_bit_cast(float, clamp_min_bits);
const float clamp_max = __builtin_bit_cast(float, clamp_max_bits);

#if MM_FUSED_CLAMP
const aie::vector<float, V> lo = aie::broadcast<float, V>(MM_FUSED_CLAMP_MIN);
const aie::vector<float, V> hi = aie::broadcast<float, V>(MM_FUSED_CLAMP_MAX);
switch (mode) {
#if MM_FUSED_EPILOGUE_MODE_MASK & 2
case 1:
epilogue_dispatch<1>(y_out, src, clamp_enabled, clamp_min, clamp_max);
return;
#endif

AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V)
for (int j = 0; j < CHUNK / V; j++) {
// The accumulator stays f32 through the activation and the clamp, and
// is converted to bf16 exactly once, on the store. Converting first
// would round twice and let the activation's slope amplify the first
// rounding -- see activations.h.
aie::vector<float, V> f = aie::load_v<V>(src + j * V);
#if MM_FUSED_EPILOGUE_MODE == 1
f = gelu_vec<V>(f);
#elif MM_FUSED_EPILOGUE_MODE == 2
f = silu_vec<V>(f);
#elif MM_FUSED_EPILOGUE_MODE == 3
f = sigmoid_vec<V>(f);
#if MM_FUSED_EPILOGUE_MODE_MASK & 4
case 2:
epilogue_dispatch<2>(y_out, src, clamp_enabled, clamp_min, clamp_max);
return;
#endif
#if MM_FUSED_CLAMP
f = aie::max(aie::min(f, hi), lo);
#if MM_FUSED_EPILOGUE_MODE_MASK & 8
case 3:
epilogue_dispatch<3>(y_out, src, clamp_enabled, clamp_min, clamp_max);
return;
#endif
aie::accum<accfloat, V> out;
out.from_vector(f);
// The assignment is the conversion: to_v16bfloat16 yields a raw
// v16bfloat16, not an aie::vector.
aie::vector<bfloat16, V> v = to_v16bfloat16(out);
aie::store_v(y_out + j * V, v);
// Mode 0 is always compiled, so a mode the mask leaves out yields an
// unactivated result rather than an unwritten buffer.
default:
epilogue_dispatch<0>(y_out, src, clamp_enabled, clamp_min, clamp_max);
return;
}
}
}
Loading
Loading