[WIP] Alternative GEMM Design

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<bfloat16, N>, 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/<arch>/{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 <aie_api/aie.hpp>
+
+#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 <int vec_size>
+__attribute__((always_inline)) aie::vector<bfloat16, vec_size> tanh_vec(aie::accum<accfloat, vec_size> x)
+{
+#if ACTIVATIONS_NATIVE_TANH
+    return aie::tanh<bfloat16>(x.template to_vector<float>());
+#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<bfloat16>());
+#endif
+}
+
+// sigmoid(x) = (tanh(x/2) + 1) / 2
+template <int vec_size> aie::vector<bfloat16, vec_size> sigmoid_vec(aie::vector<bfloat16, vec_size> x)
+{
+    const bfloat16 half = 0.5f;
+    const bfloat16 one = 1.0f;
+    aie::vector<bfloat16, vec_size> v_half = aie::broadcast<bfloat16, vec_size>(half);
+    aie::vector<bfloat16, vec_size> v_one = aie::broadcast<bfloat16, vec_size>(one);
+    aie::accum<accfloat, vec_size> x_mul_half = aie::mul(x, v_half);
+    aie::vector<bfloat16, vec_size> tanh_x_half = tanh_vec<vec_size>(x_mul_half);
+
+    aie::vector<bfloat16, vec_size> 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 <int vec_size> aie::vector<bfloat16, vec_size> silu_vec(aie::vector<bfloat16, vec_size> x)
+{
+    aie::vector<bfloat16, vec_size> sigmoid_x = sigmoid_vec<vec_size>(x);
+    return aie::mul(x, sigmoid_x);
+}
+
+// gelu(x) ~= x * sigmoid(1.702x)
+template <int vec_size> aie::vector<bfloat16, vec_size> gelu_vec(aie::vector<bfloat16, vec_size> x)
+{
+    constexpr bfloat16 x_scale = 1.702;
+    aie::vector<bfloat16, vec_size> v_x_scale = aie::broadcast<bfloat16, vec_size>(x_scale);
+    aie::vector<bfloat16, vec_size> x_scaled = aie::mul(x, v_x_scale);
+    aie::vector<bfloat16, vec_size> sigmoid_x_scaled = sigmoid_vec<vec_size>(x_scaled);
+    return aie::mul(x, sigmoid_x_scaled);
+}
+
+#endif // __ACTIVATIONS_H__

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 <aie_api/aie.hpp>
+#include <stdint.h>
+
+#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<float, M, N>(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));
+}
+}

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 <aie_api/aie.hpp>
+#include <stdint.h>
+
+#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<bfloat16, V> lo = aie::broadcast<bfloat16, V>(static_cast<bfloat16>(MM_FUSED_CLAMP_MIN));
+    const aie::vector<bfloat16, V> hi = aie::broadcast<bfloat16, V>(static_cast<bfloat16>(MM_FUSED_CLAMP_MAX));
+#endif
+
+    AIE_LOOP_MAX_ITERATION_COUNT(CHUNK / V)
+    for (int j = 0; j < CHUNK / V; j++) {
+        aie::accum<accfloat, V> acc;
+        acc.from_vector(aie::load_v<V>(src + j * V));
+        // The assignment is the conversion: to_v16bfloat16 yields a raw
+        // v16bfloat16, not an aie::vector.
+        aie::vector<bfloat16, V> v = to_v16bfloat16(acc);
+#if MM_FUSED_EPILOGUE_MODE == 1
+        v = gelu_vec<V>(v);
+#elif MM_FUSED_EPILOGUE_MODE == 2
+        v = silu_vec<V>(v);
+#elif MM_FUSED_EPILOGUE_MODE == 3
+        v = sigmoid_vec<V>(v);
+#endif
+#if MM_FUSED_CLAMP
+        v = aie::clamp(v, lo, hi);
+#endif
+        aie::store_v(y_out + j * V, v);
+    }
+}
+}