diff --git a/aie_kernels/generic/mm_fused.cc b/aie_kernels/generic/mm_fused.cc index 719c09361..46a2173b0 100644 --- a/aie_kernels/generic/mm_fused.cc +++ b/aie_kernels/generic/mm_fused.cc @@ -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" @@ -35,43 +27,27 @@ // Epilogue selection. 0 = none, 1 = gelu, 2 = silu, 3 = sigmoid, matching // Epilogue.mode in design.py. -#ifndef MM_FUSED_EPILOGUE_MODE -#define MM_FUSED_EPILOGUE_MODE 0 -#endif -#ifndef MM_FUSED_CLAMP -#define MM_FUSED_CLAMP 0 -#endif -#ifndef MM_FUSED_CLAMP_MIN -#define MM_FUSED_CLAMP_MIN 0.0f -#endif -#ifndef MM_FUSED_CLAMP_MAX -#define MM_FUSED_CLAMP_MAX 0.0f +#ifndef MM_FUSED_EPILOGUE_MODE_MASK +#define MM_FUSED_EPILOGUE_MODE_MASK 0xF #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; @@ -91,44 +67,68 @@ 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. +// +// The clamp is unconditional. An unclamped caller sends (-inf, +inf), which +// leaves every finite value bit-identical, so there is no unclamped +// instantiation to compile and no clamped-versus-not fork in the build. +template +static inline void +epilogue_body(bfloat16 *__restrict y_out, const float *__restrict src, float clamp_min, float clamp_max) +{ + const aie::vector lo = aie::broadcast(clamp_min); + const aie::vector hi = aie::broadcast(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 f = aie::load_v(src + j * V); + if constexpr (MODE == 1) + f = gelu_vec(f); + else if constexpr (MODE == 2) + f = silu_vec(f); + else if constexpr (MODE == 3) + f = sigmoid_vec(f); + f = aie::max(aie::min(f, hi), lo); + aie::accum out; + out.from_vector(f); + // The assignment is the conversion: to_v16bfloat16 yields a raw + // v16bfloat16, not an aie::vector. + aie::vector v = to_v16bfloat16(out); + aie::store_v(y_out + j * V, v); + } +} } // 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(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); @@ -137,54 +137,55 @@ 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. +// 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. Fusing them costs one more vector op per 16 elements instead of +// a separate pass over L1. // -// Fusing the activation here is the point: the values are already in registers -// after the f32 -> bf16 conversion, so gelu/silu/sigmoid costs one more vector -// op per 16 elements instead of a separate pass over L1 (which is what chaining -// a standalone activation operator after a GEMM would cost). The mode and clamp -// are compile-time, so the inner loop below is branch-free. -// -// The chunk index is split in two because the core body unrolls the drain by -// the C fifo depth to keep the acquired buffer index a compile-time constant; -// passing both parts avoids doing that arithmetic up there. -void mm_fused_epilogue_chunk(bfloat16 *y_out, float *y_acc, int32_t outer, int32_t half) +// The mode is runtime, tested once per chunk so the inner loops stay +// branch-free; the cost is program memory, since every mode in the mask is +// compiled in. The bounds arrive as raw int32 because npu_write_rtp only +// writes i32 words, and 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_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 lo = aie::broadcast(MM_FUSED_CLAMP_MIN); - const aie::vector hi = aie::broadcast(MM_FUSED_CLAMP_MAX); + switch (mode) { +#if MM_FUSED_EPILOGUE_MODE_MASK & 2 + case 1: + epilogue_body<1>(y_out, src, 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 f = aie::load_v(src + j * V); -#if MM_FUSED_EPILOGUE_MODE == 1 - f = gelu_vec(f); -#elif MM_FUSED_EPILOGUE_MODE == 2 - f = silu_vec(f); -#elif MM_FUSED_EPILOGUE_MODE == 3 - f = sigmoid_vec(f); +#if MM_FUSED_EPILOGUE_MODE_MASK & 4 + case 2: + epilogue_body<2>(y_out, src, 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_body<3>(y_out, src, clamp_min, clamp_max); + return; #endif - aie::accum out; - out.from_vector(f); - // The assignment is the conversion: to_v16bfloat16 yields a raw - // v16bfloat16, not an aie::vector. - aie::vector v = to_v16bfloat16(out); - aie::store_v(y_out + j * V, v); + // Mode 0 is always compiled, so a mode the mask leaves out yields an + // unactivated result rather than an unwritten buffer. op.py rejects that + // combination up front; this is the backstop. + default: + epilogue_body<0>(y_out, src, clamp_min, clamp_max); + return; } } } diff --git a/iron/operators/flm/gemm/README.md b/iron/operators/flm/gemm/README.md index e185873f7..28375662b 100644 --- a/iron/operators/flm/gemm/README.md +++ b/iron/operators/flm/gemm/README.md @@ -21,6 +21,12 @@ A second GEMM implementation alongside [`iron.operators.GEMM`](../../gemm), specialised for transformer projection shapes and ported from FastFlowLM's `mm` overlay. +**M, K, N and the activation are runtime parameters**, so one xclbin serves +every shape and only the instruction stream is rebuilt per shape. The +FastFlowLM harness needs that: it registers one `mm.xclbin` per model and swaps +instruction streams, against a budget of 16 xclbins for a whole model. See +[Runtime parameters](#runtime-parameters). + 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: @@ -40,17 +46,6 @@ The shipped overlay itself is available as [`iron.operators.flm.MMPrebuilt`](../mm_prebuilt) for comparison; `benchmark.py` measures the two against each other and against `iron.operators.GEMM`. -**Not yet a drop-in replacement for the shipped overlay in FastFlowLM itself.** -FLM's runtime selects matrix shape and activation per call via runtime -parameters (RTPs) on one compiled xclbin. `M`/`K`/`N`/`epilogue`/`rounding` here -are `GEMM(...)` constructor arguments instead -- baked into the MLIR and the -kernel's `-D` flags at compile time (see -[Matching the shipped overlay](#matching-the-shipped-fastflowlm-overlay)) -- so -each shape+epilogue combination is its own compiled kernel object, not one -kernel switchable at runtime. Using this operator inside FLM today means -precompiling and swapping between kernels per combination; making shape and -epilogue RTP-selectable is follow-up work. - ## Architectures Runs on both NPU2 (aie2p — Strix/Krackan) and NPU1 (aie2 — Phoenix/Hawk Point). @@ -81,6 +76,91 @@ Two consequences of the native-vs-emulated split are worth knowing: NPU2 only.** NPU1 sums the K reduction in a different order, so it matches the rounding *mode* but not the exact results. +## Runtime parameters + +**Six words** in an L1 buffer per core, written by the runtime sequence and +read by the core once its barrier opens: + +| word | value | +|---|---| +| `N` | the raw N; the core derives its own `n_work` / `n_drain` from it | +| `m_row_blocks` | `M / 256` | +| `k_iters` | `K / 512` | +| epilogue | the `Epilogue` mode | +| `clamp_min` / `clamp_max` | the bounds, as raw `int32` bit patterns | + +A word is not free: each costs ~66 ns per core and the sequence writes +`ROWS * COLS` = 32 of them, so **every word is ~2 us of dispatch latency** +(measured by padding the buffer at a fixed core count). Against a ~107 us +floor that is most of a short-prefill dispatch, so `rtp_layout()` in design.py +sizes the buffer per configuration rather than sending words a build cannot +use. Two groups are therefore conditional or gone: + +* `n_chunks` / `n_units` are sent **only when `m_chunk > 1`**. They are + `m_row_blocks // M_CHUNK` and each other, so at the shipped `M_CHUNK = 1` + the core just reads `m_row_blocks`. +* `n_work` / `n_drain` are **not sent at all**. The core derives them from `N` + and its own column, which it reads from a per-tile buffer initialised at + build time. Branch-free, and both divisors are powers of two, so it costs + shifts rather than a `__divsi3` call: + + n_tiles = N // N_TILE + n_work = (n_tiles - my_col + COLS - 1) // COLS + n_drain = ((n_tiles + COLS - 1) // COLS) - n_work + +The clamp bounds are the one group that is unconditional despite most callers +not clamping. The kernel has no unclamped instantiation to compile out: an +unclamped dispatch sends `(-inf, +inf)`, which leaves every finite value +bit-identical. Two always-sent words buy one xclbin for clamped and unclamped +callers alike, which is the whole point of the runtime parameters, and a +clamping caller now sends one word fewer than the old `clamp_enabled` trio did. +The bounds stay raw `int32` because `npu_write_rtp` writes i32 only; the kernel +casts back with `__builtin_bit_cast`, since `memcpy` leaves an unresolved +external call rather than folding to a register move. + + The column index is per-tile **static data**, deliberately not a constant + folded into the program: the 32 core programs today differ only in symbol + names, and baking it into code would make them differ in instructions, + foreclosing a future one-program xclbin. + +Measured on the 30-shape suite at **four** words: **-3.5% median at M=256** +(best -11.3%, E2B/kv), and within noise at M >= 1024 -- the saving is a +constant ~12 us, so it is a short-prefill and decode lever, not a prefill one. +Making the clamp bounds unconditional put two words back, which the ~2 us per +word above prices at **~4 us of that ~12**; the shape of the result is +unchanged but the median has not been re-measured since. + +All columns are always built. One with no work for a shape gets `n_work = 0` +and still drains its share of the A broadcast, because the memtile will not +release an A object until every consumer has taken it. + +The two artifacts therefore carry different stems: the xclbin's `config_name` +covers tile_n, ct_max_k, tile_ma, the compiled activation set, rounding and the +device, while `name` adds every runtime parameter -- M, K, N, the activation +and the clamp bounds. It has to: the sequence writes those as immediates and +the build cache keys on filename and mtime, so a stem that omits one serves the +first caller's instruction stream to the second. The xclbin is built from a +module emitted at a reference shape, whose runtime sequence is discarded. + +One thing stays build-time, because it costs program memory: which activations +the epilogue can *select between* (`epilogue_modes`). It lands in the xclbin's +name. The clamp does not -- every build compiles it, so `clamp=(-2, 2)`, +`clamp=(-4, 4)` and no clamp at all share one xclbin. + +The core releases its barrier straight after reading the parameters. +`wait_for_value` emits `LockAction.Acquire`, which does not leave the lock +consumed, so without the release a core that runs twice does not wait the +second time and reads the previous dispatch's parameters. Releasing before the +work is safe, because the sequence cannot set the barrier again until it has +drained this dispatch's C. + +The repo's other barrier users never hit this, because neither waits twice: +`mha` puts its infinite loop *inside* the wait, and `softmax` writes the same +parameters every dispatch. `test_one_xclbin_serves_every_shape` is the +regression test -- without the release it hangs the device on the second +shape, and the parametrised tests cannot catch it, because the `aie_context` +fixture reconfigures the array between cases. + ## Shape constraints `M % 256 == 0`, `K % 512 == 0`, `N % tile_n == 0` (so 64 by default). @@ -145,10 +225,11 @@ improves and gelu's worst case drops 5.5%. Measured perf-neutral (0.993-1.006x, inside the run-to-run spread). The shipped kernel selects its activation -- and its shape -- from runtime -parameters, one overlay serving every projection; this operator bakes both in -at compile time instead (activation keeps the shipped 0/1/2/3 mapping), which -is what lets its inner loop be branch-free. See the FLM-compatibility note near -the top of this file for what that means for using this operator inside FLM. +parameters, one overlay serving every projection. This operator does the same +(activation keeps the shipped 0/1/2/3 mapping); see +[Runtime parameters](#runtime-parameters). Which activations the epilogue can +*select between* is still a build-time choice, because each one compiled in +costs program memory. `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 @@ -201,9 +282,8 @@ to reduce over -- with a single k iteration there is not enough compute to hide the extra A traffic. On NPU2, `tile_n=128` wins only at `k_iters=1`, and by ~3%; at `k_iters>=2` it -is 1.2-1.7x slower. Both follow from `tile_n=128` giving up resident B — its -`mt_b` is 128 KB, so `k_iters` copies do not fit the memtile — and the more k -there is to reduce over, the more that costs. +is 1.2-1.7x slower, because its `CT_MAX_K` falls to 32 and the compute cost of +the shorter k slice swamps the A traffic it saves. 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 @@ -279,7 +359,6 @@ which is why the gap grows with the problem size rather than being flat. Unlike NPU2, compute here is **not** hidden behind the transfers, so on NPU1 both a faster mmul and less traffic pay off, where on NPU2 only the latter does. -Resident B is also a no-op on NPU1 — see [Resident B](#resident-b). ### Why the transfers are cheap @@ -295,29 +374,47 @@ Two things, both in the runtime sequence rather than the kernel: block then costs 3 shim buffer descriptors instead of `1 + 2*k_iters`, so two can be in flight without exhausting the 16 available. -### Resident B - -Where a whole column-block's B fits in the memtile double-buffered -(`k_iters <= 2`, i.e. K <= 1024 at `tile_n=64`) it is held there and replayed -per row-block, so DDR reads it once instead of `m_row_blocks` times -- about -43% less traffic. Larger K falls back to re-reading it, unchanged. - -On NPU2 this is a latency win as well as a power one, because there the -operator is close to DDR-bandwidth bound, and it grows with the height of the -problem since B's re-reads scale with `m_row_blocks`. At K=1024 N=4096, min of -per-round medians over 10 interleaved rounds of 20 dispatches, `npu_time`, -power mode `turbo`: - -| M | row-blocks | non-resident | resident | | -|---|---|---|---|---| -| 512 | 2 | 527.0 us | **461.2 us** | 12.5% | -| 1024 | 4 | 1025.8 us | **857.3 us** | 16.4% | -| 2048 | 8 | 1958.0 us | **1579.5 us** | 19.3% | - -Most of the available win is still on the table: `repeat_count` restarts the -memtile BD chain at every replay boundary, which costs part of the traffic -saving back. Closing that is the largest known remaining lever here. - -On NPU1 residency is neither a latency nor a power win — measured off-versus-on -at M=2048 it is at best a no-op and marginally negative at K=1024, well inside -the 1.4-4.4% round spread. +### B is re-fetched per row-block + +DDR reads B `m_row_blocks` times rather than once. Holding a whole column-block +in the memtile and replaying it would size that buffer from `k_iters` and set +the replay from `m_row_blocks`, putting **both K and M into the device +configuration** — and the configuration is what one xclbin has to share across +every shape. That is the standing cost of M, K and N being runtime parameters, +and it is why B is the dominant DDR leg here. + +**M is the tractable half.** `aiex.npu.push_queue` takes both `bd_id` and +`repeat_count` as SSA operands, and `aiex.dma_channel_reset_for(@fifo)` expands +into the whole re-arm trio a resident fifo needs — channel reset, `aiex.set_lock` +per bound lock, START_QUEUE re-push — inside the **runtime sequence**, which this +operator regenerates per shape. So a per-shape replay count does not have to +reach the xclbin. All of it is reachable from Python and has an npu2 device test +(1000 dispatches on one hardware context). + +**K is the part still in the way.** Correct ordering needs one memtile object +spanning every k-block, so the buffer is sized from `k_iters` and that sizing is +device configuration. Selecting among several pre-programmed BD chains via +`push_queue`'s runtime `bd_id` is the obvious line of attack and has not been +tried. + +### Split legs retire rolling, not in windows + +Where K or N is 10240, the row-block stride overflows the shim BD's 20-bit +iteration step and that leg is issued as one transfer per row-block. Two shim +resources bound how many may be outstanding, and neither is modelled by the +toolchain: BD ids (16/tile, freed without a completion check) and the channel +task queue (4 deep, pushed unconditionally). + +The sequence retires the **oldest** transfer as it issues the next, which +bounds both resources directly while keeping the channel full. + +**Do not "simplify" this into windowing** — issue four, await the whole window, +issue the next four. That bounds the same two resources and reads more simply, +but it drains the channel to *empty* at every window boundary and again at every +column-block boundary, and on a DDR-rate-bound design those bubbles are the +entire cost of the split path. Measured at up to **-12.4%** on the shapes that +take this path (E4B/gateup M1024), for no change in what is in flight. + +`m_chunk` takes this path too, since its only structural effect is to force the +split on for A. It is off by default regardless — see `M_CHUNK_FOR_N` in +design.py, which would fork the xclbin. diff --git a/iron/operators/flm/gemm/benchmark.py b/iron/operators/flm/gemm/benchmark.py index 00279cf81..507c17681 100644 --- a/iron/operators/flm/gemm/benchmark.py +++ b/iron/operators/flm/gemm/benchmark.py @@ -7,41 +7,36 @@ Up to three implementations run per shape, on identical inputs: flm :class:`iron.operators.flm.GEMM`, the port - gemm :class:`iron.operators.GEMM`, left at its defaults, which are the - same emulated-bfp16 mmul and conv_even rounding -- a like-for-like - comparison rather than one against a more accurate, slower build + gemm :class:`iron.operators.GEMM` at its defaults, which are the same + emulated-bfp16 mmul and conv_even rounding, so the comparison is + like-for-like rather than against a more accurate, slower build prebuilt :class:`iron.operators.flm.MMPrebuilt`, FastFlowLM's shipped - ``mm.xclbin``, downloaded and pinned by digest. NPU2 only, because - that binary is a fixed 8-column NPU2 overlay -- on any other device - it is dropped and the flm-vs-gemm comparison still runs. - -Nothing here needs an external install or a host-specific path: the overlay is -a ``RemoteFileArtifact``, so it is fetched into the (gitignored) build dir like -any other artifact, and its instruction stream is generated by IRON. The binary -is never checked in -- it is pinned by SHA-256 against an immutable FastFlowLM -commit and downloaded on demand. - -Marked ``extensive`` for documentation, but pytest never collects it either way: -``pytest.ini`` sets ``python_files = test.py``, and this file is named -``benchmark.py`` so no CI job -- extensive or otherwise -- runs it. That is -deliberate: it is a timing comparison, meant to be invoked directly, not a -correctness gate. The correctness half now lives in -``iron/operators/flm/mm_prebuilt/test.py``, which IS a collected ``test.py`` -and so IS reached by the extensive job. - -Timing is the runtime's own ``npu_time`` (device-side), the same source -``iron.common.test_utils.run_test`` reports, rather than a host wall clock: it -excludes host dispatch and so compares the designs rather than the driver. - -The shapes below are the ones this operator exists to serve, so they overlap -with ``test.py``'s by construction. They are not redundant with it: ``test.py`` -asserts correctness on one implementation, this compares latency across three -frozen binaries, and neither can stand in for the other. + ``mm.xclbin``, pinned by digest. NPU2 only, since that binary is a + fixed 8-column overlay; elsewhere it is dropped and the flm-vs-gemm + comparison still runs. + +The overlay is a ``RemoteFileArtifact`` pinned by SHA-256 against an immutable +FastFlowLM commit and fetched into the gitignored build dir, so nothing here +needs an external install or a host-specific path. + +pytest never collects this: ``pytest.ini`` sets ``python_files = test.py``. It +is a timing comparison meant to be invoked directly, not a correctness gate; +the correctness half lives in ``iron/operators/flm/mm_prebuilt/test.py``. + +Timing is the runtime's device-side ``npu_time`` rather than a host wall clock, +so it compares the designs rather than the driver. + +Pass ``--iterations 1``. Each test already averages ``ITERS`` dispatches over +``ROUNDS`` interleaved rounds, so conftest's default of 5 repeats the whole +matrix five times for nothing. 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 + pytest iron/operators/flm/gemm/benchmark.py --iterations 1 + pytest iron/operators/flm/gemm/benchmark.py --iterations 1 -k E2B --csv-output flm.csv + +Do not pass ``-s`` when you want the CSV: conftest's reporter parses the +captured stdout, so disabling capture yields a CSV with no metric columns. """ import statistics @@ -64,22 +59,15 @@ _dev = aie_utils.get_current_device() # The shipped overlay is a fixed 8-column NPU2 binary. Where that does not -# match the device, drop that ONE candidate rather than skipping the module: -# flm vs iron.operators.GEMM is measurable on every supported device. +# match the device, drop that one candidate rather than skipping the module, +# since flm vs iron.operators.GEMM is measurable on every supported device. HAVE_PREBUILT = _dev is not None and _dev.resolve().name == "npu2" and _dev.cols >= 8 # Every projection of both Gemma4 variants FastFlowLM ships, at three prefill -# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. -# -# These two are the right coverage for the shipped mm.xclbin. Checked against -# FastFlowLM f81eba71: Gemma4-E2B-IT-NPU2, Gemma4-E4B-IT-NPU2 and Gemma3-4B-NPU2 -# all ship the SAME mm.xclbin blob (git 4727df98, 512220 bytes) -- one overlay -# serving several models -- so E2B and E4B between them already exercise it. -# -# Gemma4-12B-IT-NPU2 does not ship an mm.xclbin at all. Its overlays are a -# different set (attn_global, attn_sliding, audio_image_mm, dequant_mm, layer, -# lm_head), so its projections go through a quantized matmul rather than this -# bf16 one and cannot be compared against MMPrebuilt. +# lengths. E2B is dim 1536 / ffn 6144; E4B is dim 2560 / ffn 10240. Both ship +# the same mm.xclbin blob (checked at FastFlowLM f81eba71), so between them +# they cover it. Gemma4-12B ships no mm.xclbin at all -- its projections go +# through a quantized matmul -- so it cannot be compared against MMPrebuilt. # proj, K, N E2B_PROJ = [ ("q", 1536, 4096), @@ -111,11 +99,10 @@ 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). + # 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 split that leg into per-mega_row + # transfers now (design.py's a_split/c_split, test_gemm_split_leg_bounds). params = [] for model, projections in (("E2B", E2B_PROJ), ("E4B", E4B_PROJ)): for M in PREFILL_LENGTHS: @@ -132,10 +119,7 @@ def make_inputs(M, K, N): 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. + # Bounded against accumulated mass rather than relatively; see test.py. return A, B, Af @ Bf, float((Af.abs() @ Bf.abs()).mean()) @@ -152,9 +136,8 @@ def __init__(self, name, op, A, B, M, N, budget, ctx): 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. + # Only the flm operators take B pre-packed. iron.operators.GEMM + # reorders in the descriptor, so it wants plain row-major (K, N). packed_b = op.pack_B(B) if hasattr(op, "pack_B") else B args = [ XRTTensor.from_torch(A.flatten()), @@ -192,10 +175,8 @@ def jitter_pct(self): 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. + # The budget below is loose enough that a toolchain change could move the + # error a long way inside it unnoticed, so record the numbers too. FLMErr=r"flm err/mass: (?P[\d\.e\+-]+)", PrebuiltErr=r"prebuilt err/mass: (?P[\d\.e\+-]+)", GEMMErr=r"gemm err/mass: (?P[\d\.e\+-]+)", @@ -254,11 +235,10 @@ def test_gemm_vs_prebuilt(model, proj, M, K, N, aie_context): 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". + # Round-robin, never all of one then all of another. Dispatch latency here + # is bimodal with modes about 6% apart, so a batch landing wholly in one + # mode turns min-of-medians into a mode selector. That is how a change + # later shown to do nothing once produced a convincing 5% "win". for _ in range(ROUNDS): for c in candidates: c.time_round() diff --git a/iron/operators/flm/gemm/design.py b/iron/operators/flm/gemm/design.py index 864e55f62..01b140e15 100644 --- a/iron/operators/flm/gemm/design.py +++ b/iron/operators/flm/gemm/design.py @@ -4,53 +4,17 @@ """bf16 GEMM over a 4-row compute-tile grid, as wide as the device. A second GEMM design alongside ``iron.operators.gemm``, specialised for -transformer projection shapes. The overall dataflow is the same whole-array -shape as that operator's -- A broadcast along each compute row, B down each -column, C joined through the memtile -- so those are NOT what distinguishes it. -What does: - - * **B is quantized to bfp16ebs8 on NPU2** by ``GEMM.pack_B``, not bf16. - ``iron.operators.GEMM`` only ever moves bf16. This is not primarily a DMA - saving: it is what makes NPU2's fast mmul lowering available at all -- - ``aie::mmul<8,8,8>`` needs bfp16 operands to decompose into two emulated - macs instead of four, which is most of the NPU2 speedup (see NPU1 in - README.md's Performance section, where B stays bf16 and the margin over - ``iron.operators.GEMM`` is correspondingly smaller). Quantizing is - numerically free -- the mmul only multiplies bfp16 regardless, so this - hoists a rounding that already happened on every mac -- provided it - reproduces the core's rounding mode; see ``packing.py``. - * **The tile shape is fixed, not parameterised** -- but fixedness alone is - not the advantage: ``iron.operators.GEMM`` is equally fixed once compiled - with a choice of tile args. What differs is *which* shape is fixed. r/s/t - stays 8/8/8 on both architectures for the reason below. m/k = 64/512 (n - defaults to 64) is chosen for the L1-budget tradeoff documented next to - ``CT_MAX_K_FOR_N`` below: n=64 gives the mmul a colA of 8 rather than 4, - which wins whenever compute is the critical path, at the cost of A being - re-read more often. Only the grid WIDTH varies with the device: 8 columns - on NPU2, 4 on NPU1. - * **A fused epilogue.** The f32->bf16 conversion, an optional activation and - an optional clamp all happen while the values are still in registers, on the - way into the C object, instead of a separate pass over L1. - * **B arrives pre-packed** by ``GEMM.pack_B``, in the order the cores consume - it, so both B hops are plain linear descriptors instead of the 128-byte - scattered bursts ``iron.operators.GEMM`` reorders in the descriptor. - * **Asymmetric tile buffering (ATB)**, so the A tile and the accumulator need - not share a height -- this is what buys the deep k slice (K_TILE=512) - within the L1 budget; see README.md's ATB reference. - -None of these four helps alone -- see README.md's Performance section for the -measured, per-choice breakdown of the gap against both the shipped FastFlowLM -overlay and ``iron.operators.GEMM``. +transformer projection shapes. Same dataflow -- A broadcast along each compute +row, B down each column, C joined through the memtile -- but with a fixed tile +shape, B quantized to bfp16ebs8 and pre-packed into consumption order, an +activation and clamp fused into the C drain, and 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. +README.md has the per-choice breakdown against both the shipped FastFlowLM +overlay and ``iron.operators.GEMM``. -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. +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 argparse @@ -71,6 +35,7 @@ Runtime, TaskGroup, Worker, + WorkerRuntimeBarrier, ) from aie.iron.controlflow import range_ from aie.dialects.aie import get_target_model @@ -82,32 +47,21 @@ # --- Fixed geometry ------------------------------------------------------- # GEMM tiling per compute tile, and the register tiling inside it. M_TILE, K_TILE = 64, 512 -# Default n tile. 64 gives the mmul a colA of 8 rather than 4, halving the -# accumulator traffic per mac, at the cost of doubling A fetches (the grid -# then covers 512 columns of N per pass instead of 1024). That trade wins -# whenever compute is the critical path, which is the usual case; see -# README.md for the measured sweep, including the small-K shape where it -# loses. +# Default n tile. 64 doubles A fetches but gives the mmul colA=8 instead of 4, +# which wins when compute is the critical path. op.py picks per shape. N_TILE_DEFAULT = 64 -# How much of K one compute tile holds at a time, per n width. This is a fixed -# L1 budget split two ways, so a wider n tile leaves less room for B's k slice -# and the product stays roughly constant. op.py passes the chosen value to the -# kernel as -DMM_FUSED_CT_K, making this table the only place it is decided. -# n=256 is deliberately absent: at that width the f32 accumulator alone -# (M_TILE * 256 * 4 = 65536 bytes) already fills the whole of L1, before A, B -# or C are even counted, so no ct_max_k could ever make it fit. +# How much of K one compute tile holds at a time, per n width. It is a fixed +# L1 budget split two ways, passed to the kernel as -DMM_FUSED_CT_K. n=256 is +# absent because its f32 accumulator alone (M_TILE*256*4) fills all of L1. CT_MAX_K_FOR_N = {16: 16, 32: 32, 64: 128, 128: 32} -# Register tiling, shared by both architectures. These set the blocked L1 -# layout, so ``pack_B``, the four stream-dimension lists below and -# ``gather_dims`` all key off them; changing one without the others is silently -# wrong rather than a build error. -# -# Matching AIE2's native 4x8x4 mac shape instead is a measured dead end, at -# 22-30% slower: the kernel is load-port bound rather than shuffle bound, and -# 4/8/4 needs 1.62 loads per mac against 8/8/8's 1.06, because the 2x2 register -# block amortizes each load over four macs either way but over far less work. -# Retrying it needs a wider register block on the native shape, i.e. a 4x4 mmul -# kernel, not just a different r/s/t here. +# (tile_n, ct_max_k) pairs verified on hardware. The table above looks tunable +# but is not, and a wrong value fails SILENTLY: ct_k=64 at tile_n=64 gives +# err/mass 3.45e-02 against 2.42e-04, and tile_n=128 NaNs. Cause unknown, so +# refuse rather than miscompute. +_VERIFIED_CT_K = {(16, 16), (32, 32), (64, 128), (128, 32)} +# Register tiling, shared by both architectures. pack_B, the stream-dimension +# lists and gather_dims all key off these; changing one alone is silently +# wrong. AIE2's native 4x8x4 shape measured 22-30% slower (load-port bound). R, S, T = 8, 8, 8 @@ -121,16 +75,28 @@ def compute_rows(dev): 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 +# L1 bytes reserved for the core's stack, which the buffer budget below must +# not hand out. The device default is 1024 and aiecc measures what a build +# actually needs: 1088 on NPU1, which is the activation LUT path plus the +# epilogue's clamp vectors, so the default fails the build outright. 2048 +# leaves headroom; aiecc names the exact requirement if a change outgrows it. +STACK_SIZE = 2048 +# Row-blocks a core folds into one B fetch, cutting B's DDR reads by M_CHUNK +# at the cost of that many L1 accumulators and forcing a_split. Off everywhere +# for a contractual reason: it must divide m_row_blocks (M % 512 == 0) while +# the overlay this replaces takes any multiple of 256, so a shape that cannot +# use it forks config_name. See README.md. +M_CHUNK_FOR_N = {16: 1, 32: 1, 64: 1, 128: 1} # 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. +# shim buffer descriptors on a column (A + B + C) of the 16 available, so the +# ceiling is 5. 2 is enough to keep the fills ahead of the cores. OVERLAP_DEFAULT = 2 class Epilogue(StrEnum): """Activation folded into the C drain. - Declaration order is the wire format -- it is both the kernel's + Declaration order is the wire format: it is both the kernel's ``-DMM_FUSED_EPILOGUE_MODE`` and the shipped overlay's ``output_mode``. """ @@ -145,12 +111,45 @@ def mode(self) -> int: return list(Epilogue).index(self) +# The parameter buffer each core reads once its barrier opens. These six are +# always present; conditional words follow at offsets rtp_layout() computes, +# so a word a build cannot use is never allocated. +# +# The clamp bounds are unconditional even though most callers do not clamp, +# because the alternative is a second xclbin: the kernel's clamp is not +# compiled out, it is neutralised by sending (-inf, +inf). Two words is the +# price of that, and a clamping caller now pays one word less than it did. +( + RTP_N_VAL, + RTP_M_ROW_BLOCKS, + RTP_K_ITERS, + RTP_EPILOGUE, + RTP_CLAMP_MIN, + RTP_CLAMP_MAX, +) = range(6) + + +def rtp_layout(m_chunk): + """Slot index for each optional parameter, and the total word count. + + A word is not free: the sequence writes one per core, costing ~2 us of + dispatch latency against a ~107 us floor. So optional groups are omitted + rather than defaulted. + """ + slots = {} + n = 6 + if m_chunk > 1: + slots["n_chunks"] = n + slots["n_units"] = n + 1 + n += 2 + return slots, n + + class Rounding(StrEnum): """Rounding for every f32->bf16 conversion. - The core powers up in floor; conv_even is the default because truncation - biases every conversion the same way and the error then accumulates over - the K reduction. floor reproduces the shipped overlay. + conv_even by default: truncation biases every conversion the same way, so + the error accumulates over the K reduction. floor matches the overlay. """ CONV_EVEN = "conv_even" @@ -183,20 +182,16 @@ def _b_bytes(elems, bfp16_b): # --- Shim DMA limits ------------------------------------------------------ # -# Hardware facts the Python bindings do not expose: gemm() reads AIETargetModel -# directly for the L1 ceiling, BD count and grid, but neither getDmaBdStepBits -# nor getDmaBdWrapSizeBits is bound, and nothing models the channel task queue. -# gemv/design.py and repeat/design.py hardcode the same fields. -# -# Step field width. An IR-level bf16-element stride S is re-expressed as -# (S - 1) * 2 bytes / 4-byte granularity before AIEXDialect.cpp checks it. +# Hardware facts the Python bindings do not expose: getDmaBdStepBits and +# getDmaBdWrapSizeBits are unbound, and nothing models the channel task queue. +# gemv/design.py and repeat/design.py hardcode the same fields. An IR-level +# bf16 stride S is re-expressed as (S-1)*2 bytes / 4-byte granularity before +# AIEXDialect.cpp checks it. _SHIM_STEP_BITS = 20 _BF16_BYTES = 2 _ADDR_GRANULARITY_BYTES = 4 -# Entries in a shim DMA channel's task queue. AIEDmaToNpu's NpuPushQueueOp -# pushes unconditionally, so overrunning this is a silent device hang rather -# than a diagnostic. Measured at K=10240 M=1024: 4 outstanding tasks on one -# channel run, 8 hang. +# Entries in a shim DMA channel's task queue. NpuPushQueueOp pushes +# unconditionally, so overrunning this hangs silently. Measured: 4 run, 8 hang. SHIM_TASK_QUEUE = 4 @@ -205,30 +200,17 @@ def _hw_stride_ok(stride_elems): return hw_stride <= (1 << _SHIM_STEP_BITS) - 1 -def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): - """Pick (A-tile height, L1 B depth) -- the largest working set that fits. +def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): + """Pick the largest working set that fits: (A-tile height, L1 B depth). - ``b_elem_bytes`` is 9/8 where B is bfp16ebs8 and 2 where it is bf16, and - ``budget`` is the core's data memory, so the search below reflects what B - actually costs on this device. - - No stack is reserved out of ``budget``: the cores leave ``stack_size`` - unset and aiecc measures each core's requirement and fails the build if it - does not fit, so the stack is the toolchain's to enforce. This kernel - measures 192 bytes against the >=6 KB the search leaves unused anyway. - - A dies as soon as it is consumed while the accumulator lives across the - whole K reduction, so they need not share a height; shrinking A is what - pays for a k slice deep enough to halve the accumulator traffic per mac. - B's depth is searched too because at n=128 the k=128 slice makes the B - object 18 KB, and a double-buffered pair simply does not fit -- giving that - up is what buys colA=16 there, and colA is worth far more than B's L1 - prefetch (3.67 -> 2.28 cycles per mac, measured). - - Deeper B first, then the tallest A that still fits, so the n=64 default is - unchanged at (32, 2). + Deeper B first, then the tallest A that still fits, since colA is worth + far more than B's L1 prefetch. ``budget`` is the whole local memory; the + stack comes off it here so callers can keep passing the raw size. """ - acc = M_TILE * n_tile * 4 + budget -= STACK_SIZE + # m_chunk accumulators, since the core holds a B chunk across that many + # row-blocks. The only term that scales with it. + acc = m_chunk * M_TILE * n_tile * 4 cout = CT_OUT_LEN * 2 * C_DEPTH for b_depth in (B_DEPTH, 1): b = int(ct_max_k * n_tile * b_elem_bytes) * b_depth @@ -241,22 +223,15 @@ def _default_l1(n_tile, ct_max_k, b_elem_bytes, budget): raise ValueError(f"nothing fits L1 for tile_n={n_tile}, ct_max_k={ct_max_k}") -def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget): +def _b_depth_for(t_ma, n_tile, ct_max_k, b_elem_bytes, budget, m_chunk=1): """Deepest B fifo depth that fits L1 alongside an explicit A-tile height. - ``_default_l1`` picks L1_B_DEPTH together with the t_ma IT chooses; that - pairing need not fit a caller-overridden t_ma; a taller A tile leaves less - L1 for B, and can push a working set that fit at the default t_ma over - budget. Raise rather than silently reusing a depth that doesn't fit. - - How much the depth is worth, measured on npu2 (turbo, 12 interleaved rounds - of 20 dispatches, min of per-round medians) by forcing depth 1 against the - default: 1.2% at M=1024 K=1536 N=6144, and within noise at K=1024 N=4096 and - K=512 N=1024. So the prefetch earns its L1 at the largest shapes and is - close to free elsewhere -- worth keeping, but not worth contorting the - search for. + ``_default_l1``'s depth is chosen with its own t_ma, which need not fit a + caller-overridden one. Raise rather than reuse a depth that does not fit. """ - acc = M_TILE * n_tile * 4 + budget -= STACK_SIZE + # Same terms as _default_l1; acc is the only one that scales with m_chunk. + acc = m_chunk * 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): @@ -275,47 +250,49 @@ def gemm( K, N, epilogue=Epilogue.NONE, + clamp=None, tile_n=N_TILE_DEFAULT, + m_chunk=None, tile_ma=None, - overlap=None, kernel_object="mm_fused.o", trace_size=0, ): """Emit the MLIR module for an M x K @ K x N bf16 GEMM. - A is (M, K) row-major, B is (K, N) row-major and C is (M, N) row-major, all - bf16 and all plain dense tensors, except that B must arrive pre-packed by - ``GEMM.pack_B`` -- it emits B in the order the cores consume it, so both - B hops are plain linear descriptors. + A, B and C are row-major bf16 dense tensors, except that B must arrive + pre-packed by ``GEMM.pack_B`` in the order the cores consume it. """ if tile_n not in CT_MAX_K_FOR_N: raise ValueError( f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {tile_n}" ) - # Everything shape-related below comes from the device rather than a - # constant, so the same dataflow covers NPU2's 4x8 and NPU1's 4x4. + # From the device, not constants, so one dataflow covers 4x8 and 4x4. tm = get_target_model(dev.resolve()) COLS, ROWS = dev.cols, compute_rows(dev) MIN_M = M_TILE * ROWS SHIM_BDS = tm.get_num_bds(0, 0) N_TILE = tile_n CT_MAX_K = CT_MAX_K_FOR_N[N_TILE] - # B is bfp16ebs8 on AIE2P and bf16 on AIE2: the scalar BFP types are gated - # on __AIE_API_SCALAR_BFP_TYPES__, which only aie_api/detail/aie2p/config.hpp - # defines, so on AIE2 B stays bf16 and the mmul lowers onto four native - # 4x8x4 macs. That choice drives every B type and extent below. B_GROUP is - # the number of B values per element of the MLIR type, so a length in values - # becomes a length in elements by dividing. + if (N_TILE, CT_MAX_K) not in _VERIFIED_CT_K: + raise ValueError( + f"tile_n={N_TILE} with ct_max_k={CT_MAX_K} is not a verified " + f"combination (verified: {sorted(_VERIFIED_CT_K)}). It would build, " + f"run, and compute the WRONG ANSWER -- see _VERIFIED_CT_K. If you " + f"are retuning CT_MAX_K_FOR_N, fix that coupling first and add the " + f"pair here once a hardware test passes." + ) + M_CHUNK = M_CHUNK_FOR_N[N_TILE] if m_chunk is None else m_chunk + # B is bfp16ebs8 on AIE2P and bf16 on AIE2. B_GROUP is the B values per + # element of the MLIR type, so a length in values divides to elements. BFP16_B = dev.arch == AIEArch.AIE2p B_GROUP = BFP16_GROUP if BFP16_B else 1 b_elem_bytes = BFP16_GROUP_BYTES / BFP16_GROUP if BFP16_B else 2 - # Asymmetric tile buffering: the A tile spans T_MA rows while the - # accumulator spans M_TILE, so the core folds RHO bands into one C tile. - # A is dead the moment it is consumed while C lives across the whole K - # reduction, so sizing both to M_TILE pays the peak L1 cost twice. + # Asymmetric tile buffering: A spans T_MA rows and the accumulator M_TILE, + # so the core folds RHO bands into one C tile. A dies on consumption while + # C lives across the K reduction, so sizing both to M_TILE pays twice. if tile_ma is None: T_MA, L1_B_DEPTH = _default_l1( - N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size(), M_CHUNK ) else: T_MA = tile_ma @@ -324,10 +301,15 @@ def gemm( f"tile_ma ({T_MA}) must divide {M_TILE} and be a multiple of {2 * R}" ) L1_B_DEPTH = _b_depth_for( - T_MA, N_TILE, CT_MAX_K, b_elem_bytes, tm.get_local_memory_size() + T_MA, + N_TILE, + CT_MAX_K, + b_elem_bytes, + tm.get_local_memory_size(), + M_CHUNK, ) RHO = M_TILE // T_MA - OVERLAP = OVERLAP_DEFAULT if overlap is None else overlap + OVERLAP = OVERLAP_DEFAULT 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 @@ -337,12 +319,18 @@ def gemm( MIN_N = N_TILE * COLS epilogue = Epilogue(epilogue) - # A compute tile does a whole m x n block or nothing, so M and K must tile - # exactly. N need only be a multiple of N_TILE: a trailing group of fewer - # than COLS blocks is handled by giving the columns different trip counts - # (see col_work / col_drain below). That matters in practice -- for a - # transformer the o and down projections have N = model dim, which is - # essentially never a multiple of N_TILE*COLS. + # Clamp bounds ride the RTP buffer as raw int32 bit patterns, since + # npu_write_rtp writes i32 only. No clamp means the identity bounds rather + # than a different build: min(x, +inf) and max(x, -inf) leave every finite + # value bit-identical, so an unclamped dispatch is numerically unchanged. + rtp_slots, rtp_words = rtp_layout(M_CHUNK) + clamp_lo, clamp_hi = clamp if clamp is not None else (-np.inf, np.inf) + clamp_min_bits = int(np.float32(clamp_lo).view(np.int32)) + clamp_max_bits = int(np.float32(clamp_hi).view(np.int32)) + # A 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 short trailing group is + # handled by per-column trip counts, which matters because o and down + # have N = model dim. 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}") @@ -350,82 +338,63 @@ def gemm( bf16_ty = np.dtype[bfloat16] f32 = np.dtype[np.float32] - # How many times the whole grid sweeps, in each dimension. m_row_blocks = M // MIN_M k_iters = K // K_TILE - # A mega_row dimension with stride ROWS*M_TILE*{K,N} lands in the shim - # BD's ITERATION field, whose step is 20 bits, so it overflows once K or N - # crosses ~8191 elements -- E4B's FFN width (10240) does, E2B's max (6144) - # does not. Only relevant once M walks more than one mega_row; at M=256 the - # dimension is degenerate (size 1) and is stripped before the stride is - # ever encoded, so it never fails there regardless of K/N. - # - # Such a leg is instead issued as m_row_blocks separate transfers, each - # carrying the mega_row jump in its OFFSET (unbounded) rather than a shared - # STRIDE. The two legs are independent -- E4B's down-proj overflows on K - # (A only) and its gate/up on N (C only) -- so neither shape pays for both. + # A "unit" is one group of M_CHUNK row-blocks. Every leg is issued per + # unit, so A, B and C stay aligned with each other and the core's nest. + n_chunks, n_rem = divmod(m_row_blocks, M_CHUNK) + if n_rem: + # op.py resolves m_chunk, so this cannot fire. A partial group is + # inexpressible: the object is M_CHUNK tiles wide and the forward + # always drains that much, and no way of padding it lowers correctly. + raise ValueError( + f"m_row_blocks ({m_row_blocks}) must be a multiple of m_chunk " + f"({M_CHUNK}); op.py should have resolved m_chunk to 1 here" + ) + n_units = n_chunks + + def unit_rows(u): + """(first row-block, how many) for unit ``u``; always a full group.""" + return u * M_CHUNK, M_CHUNK + + # A mega_row stride lands in the shim BD's 20-bit iteration step, so it + # overflows once K or N passes ~8191 elements; only E4B's 10240 does. Such + # a leg goes out as one transfer per mega_row, carrying the jump in its + # unbounded offset. M_CHUNK > 1 forces the same path for A. # - # Those transfers must stay live in their TaskGroup until awaited. - # TaskGroup.finish() emits dma_free_task, which returns the buffer - # descriptor id to a COMPILE-TIME allocator that does not check the - # transfer finished (mlir-aie AIEAssignRuntimeSequenceBDIDs::recycle, - # isAwait=false); ids are per shim TILE, shared across channels and - # directions, so retiring one early lets the next task reprogram a live - # descriptor. Verify with aie-opt --aie-substitute-shim-dma-allocations - # --aie-assign-runtime-sequence-bd-ids: the ids on a shim tile must be - # distinct. - a_split = m_row_blocks > 1 and not _hw_stride_ok(ROWS * M_TILE * K) + # Those transfers must stay live in their TaskGroup until awaited: the + # BD-id allocator is compile-time and does not check that a transfer + # finished, so freeing one early lets the next task reprogram a live + # descriptor and corrupt silently. + a_split = n_units > 1 and (M_CHUNK > 1 or 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: + # Split legs share one channel, whose task queue is 4 deep and modelled + # nowhere; overrunning it hangs (4 outstanding run, 8 hang). emit_split() + # bounds it by retiring the oldest as it issues the next, which also keeps + # the channel full -- do not simplify that to awaiting a whole batch, which + # drains the channel at every boundary and costs up to 12.4%. The bound + # counts transfers, not units: under c_split a unit drains M_CHUNK of them. + _per_unit = M_CHUNK if c_split else 1 + # Live descriptors on a shim tile: SHIM_TASK_QUEUE from the rolling window, + # plus B and the unsplit leg for each of the two blocks a boundary spans. + bds_per_block = SHIM_TASK_QUEUE + 2 + 2 + if (a_split or c_split) and bds_per_block > SHIM_BDS: raise ValueError( f"M={M} K={K} N={N} needs {bds_per_block} shim buffer descriptors " - f"per window (1 B + {MB_WINDOW if a_split else 1} A + " - f"{MB_WINDOW if c_split else 1} C) but a shim tile has only " - f"{SHIM_BDS}." + f"for the split path but a shim tile has only {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)) + # The unsplit path pipelines whole column-blocks instead, at 3 descriptors + # each (A + B + C). The split path does its own bounding above and ignores + # this. + OVERLAP = max(1, min(OVERLAP, SHIM_BDS // 3)) # Sweeps where all COLS columns have work, plus a trailing group of # rem_blocks columns (0 <= rem_blocks < COLS) that do one block more. n_full = N // MIN_N rem_blocks = (N % MIN_N) // N_TILE - # Columns that participate at all. With no full sweep (N below the grid's - # COLS*N_TILE stride) only the first rem_blocks columns do, and the rest - # are not instantiated -- giving them fifos that nothing ever drains builds - # dead dataflow, which newer mlir-aie rejects outright with - # "objectfifo.pool op segment 0 has no drainer". - n_active_cols = COLS if n_full else rem_blocks - # Per column: how many column-blocks it computes, and whether it sits out a - # trailing one while still draining the A broadcast for its row. That drain - # only arises for a column that exists and skips the trailing block, which - # requires at least one full sweep. - col_work = [n_full + (1 if c < rem_blocks else 0) for c in range(n_active_cols)] - col_drain = [ - 1 if (rem_blocks and n_full and c >= rem_blocks) else 0 - for c in range(n_active_cols) - ] - + # Every column is instantiated for every shape. Which columns exist is + # configuration, and this design has only one. A column with no work for + # the current shape gets n_work = 0 and drains instead. + n_active_cols = 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 @@ -434,14 +403,12 @@ def gemm( 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 + # L2 (per memtile). M_CHUNK stacked row-block tiles, so the forward below + # can interleave them on the way out; see a_send_dims. + mt_a_ty = np.ndarray[(M_CHUNK * M_TILE * K_TILE,), bf16_ty] 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. + # L3 (DDR), flat; the taps below index them linearly. a_l3_ty = np.ndarray[(M * K,), bf16_ty] b_l3_ty = np.ndarray[(K * N // B_GROUP,), b_elem_ty] c_l3_ty = np.ndarray[(M * N,), bf16_ty] @@ -460,34 +427,36 @@ def gemm( epilogue_chunk = Kernel( EPILOGUE_SYMBOL, kernel_object, - [ct_out_ty, ct_acc_ty, np.int32, np.int32], + # outer, half, mode, clamp_min_bits, clamp_max_bits + [ct_out_ty, ct_acc_ty] + [np.int32] * 5, ) # --- 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. + # These turn a row-major DDR tile into the blocked layout the mmul + # indexes. A mismatch is silently wrong, not a build error. # C: de-block each core's r x t tiled output back into row-major within its # 64x128 slice, on the way into the memtile. gather_dims = [(M_TILE // R, R * N_TILE), (N_TILE // T, T), (R, N_TILE), (T, 1)] - # B: DDR row-major (k x n) -> s x t blocks (recv), then split into the - # CT_MAX_K-deep chunks a single mmul call consumes (send). - # B needs no reblocking on either hop: pack_B already emits it in the - # order the cores consume, so the memtile just streams it through. That - # frees every descriptor dimension B used to spend -- which is what lets - # CT_MAX_K reach 128 (the innermost run would otherwise overflow the BD's - # 10-bit size field and need a split dimension) at the same time as - # residency (which spends one on its outer k walk). + # B needs no reblocking on either hop: pack_B emits it in consume order. + # That frees the descriptor dimensions that let CT_MAX_K reach 128. 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: same idea, r x s blocks. The outermost row-group dimension spans + # M_CHUNK tiles. mc's stride is exactly this dimension's size*stride, so + # the two merge and the walk stays within the memtile BD's four dims. + a_recv_dims = [ + (M_CHUNK * M_TILE // R, R * K_TILE), + (R, S), + (K_TILE // S, R * S), + (S, 1), + ] + # Emits (b_iter, mc, band): the order the core acquires A in while holding + # a B chunk across the group. a_send_dims = [ (K_DIV_CT_K_MAX, R * CT_MAX_K), - (M_TILE // R, R * K_TILE), + (M_CHUNK * M_TILE // R, R * K_TILE), ] + split_run(R * CT_MAX_K) # C: one join per column. Each of the ROWS cores in the column drops its @@ -507,9 +476,9 @@ def gemm( 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: shim -> memtile -> broadcast along the compute row, reblocking on + # the forward(). One fifo per row even at M_CHUNK > 1: a second would + # want a third core input DMA channel, and a tile has two. a_l3l2_fifos = [] a_cons = {} for r in range(ROWS): @@ -521,185 +490,189 @@ def gemm( name=f"A_L2L1_{r}", dims_to_stream=a_send_dims, ) - # One cons() handle per active column; every tile in the row sees - # this object, so inactive columns must not be consumers at all. + # Every tile in the row sees this object, so inactive columns must + # not be consumers at all. for c in range(n_active_cols): a_cons[(r, c)] = of_a.cons() - # B: shim -> memtile -> broadcast down the compute column. - # - # Where it fits, a whole column-block's B is held in the memtile as ONE - # object and re-walked per row-block, so DDR reads it once instead of - # m_row_blocks times. B is the dominant DDR leg, so that is roughly 43% less - # total traffic; the latency it buys grows with M, because B's re-reads - # scale with m_row_blocks. Larger K does not fit and falls back to - # re-reading. See README.md for the measured effect. - # - # Three things here are load-bearing rather than tuning: - # - # * ONE object spanning every k-block, not a pool of k_iters objects. - # Iterating a pool replays each object in turn (k0,k0,k1,k1,...) rather - # than the k0..kn sequence the cores accumulate in. - # * repeat_count on the forward() below is what re-sends an object. - # iter_count only bounds how many times an end cycles through all its - # buffers, so it is in units of depth-cycles; getting it wrong hangs - # rather than mis-computing. - # * The depth search takes the deepest that fits, not depth 1. Single - # buffering stops the next column-block prefetching behind this one's - # replay, which measures worse than not being resident at all. - # - # The budget must count what A and C actually occupy: at tile_n=128 C - # doubles and a resident B is 432 KB, and a budget that assumed C's size - # admitted a configuration that then failed address assignment. Placement - # has zero slack -- the eight memtiles pack to exactly 512 KB, relying on - # aie-objectfifo-allocate spilling one buffer to an adjacent tile -- so - # re-verify it after any change to the A, B or C buffer sizes. - mt_free = tm.get_mem_tile_size() - mt_a_bytes * A_DEPTH - mt_out_bytes * C_DEPTH - MT_B_DEPTH = next( - (d for d in (B_DEPTH, 1) if k_iters * mt_b_bytes * d <= mt_free), 0 - ) - b_resident = MT_B_DEPTH > 0 - if b_resident: - # Just a bigger buffer. With B packed in consumption order the walk is - # linear, so spanning every k-block needs no extra descriptor - # dimension -- the objects simply come out in k order. (The previous - # blocked layout had to widen one dim inbound and add an outermost k - # dim outbound, which is what collided with CT_MAX_K=128.) - mt_b_ty = np.ndarray[(k_iters * K_TILE * N_TILE // B_GROUP,), b_elem_ty] - + # B: shim -> memtile -> broadcast down the compute column, one k-block per + # object and re-fetched per row-block. Holding a whole column-block would + # put both K and M in the device configuration (README.md). Placement has + # zero slack -- the memtiles pack to exactly 512 KB, one buffer spilled to + # a neighbour -- so re-verify after any A/B/C size change. b_l3l2_fifos = [] b_cons = {} for c in range(n_active_cols): - of_b_in = ObjectFifo( - mt_b_ty, name=f"B_L3L2_{c}", depth=MT_B_DEPTH if b_resident else B_DEPTH - ) + 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( - # The one placement pin this design keeps. Everything else -- the - # workers, the accumulator buffers, the C join, the A forward and - # the shim ends -- is left to the placer, and measures the same. - # - # Without it, aie-place-tiles merges the 20 logical memtiles (4 A - # relays + 8 B relays + 8 C joins) onto the 8 physical ones in a way - # that aie-objectFifo-stateful-transform then rejects with "number - # of input DMA channel exceeded". Spreading B one-per-column is - # enough to steer it to a legal assignment; see the mlir-aie issue - # referenced in README.md. Reproduces at M=1024 K=2048 N=2048, which - # test.py covers. + # The one placement pin; everything else is left to the placer. + # Without it the 20 logical memtiles merge onto the 8 physical + # ones in a way rejected with "number of input DMA channel + # exceeded". Reproduces at M=1024 K=2048 N=2048. tile=Tile(c, 1), obj_type=ct_b_ty, depth=L1_B_DEPTH, name=f"B_L2L1_{c}", dims_to_stream=b_send_dims, - # Replay the resident memtile object once per row-block. This is - # the mechanism that actually re-sends an object; iter_count only - # bounds how many chain iterations happen in total. Correct - # ordering depends on the memtile holding ONE object spanning every - # k-block: replicating a pool of k_iters smaller objects would - # emit k0,k0,k1,k1,... rather than the k0..kn sequence the cores - # accumulate in. - repeat_count=m_row_blocks if b_resident else None, ) for r in range(ROWS): b_cons[(r, c)] = of_b.cons() + # Data, not an immediate folded into the program: the core programs differ + # only in symbol names, and baking this in as code would foreclose a + # one-program xclbin. Written once, so it costs nothing per dispatch. + my_cols = [ + [ + Buffer( + np.ndarray[(1,), np.dtype[np.int32]], + name=f"my_col_{r}_{c}", + initial_value=np.array([c], dtype=np.int32), + ) + for c in range(n_active_cols) + ] + for r in range(ROWS) + ] + + # --- Runtime parameters ----------------------------------------------- + rtps = [ + [ + Buffer( + np.ndarray[(rtp_words,), np.dtype[np.int32]], + name=f"rtp_{r}_{c}", + initial_value=np.zeros(rtp_words, dtype=np.int32), + use_write_rtp=True, + ) + for c in range(n_active_cols) + ] + for r in range(ROWS) + ] + barriers = [ + [WorkerRuntimeBarrier() for _ in range(n_active_cols)] for _ in range(ROWS) + ] + # --- Compute ---------------------------------------------------------- - def core_fn(n_work, n_drain, acc, o_h, b_h, a_h, init_k, kstep_k, epi_k): - """Core body for a column that computes ``n_work`` column-blocks and - then drains A for ``n_drain`` more (0 or 1). - - n_work/n_drain are bound per column via functools.partial below; they - are compile-time constants, so the trip counts below fold away. - """ - # The loop nest lives here rather than inside the kernel so that - # every level has an ObjectFifo acquire point. With M/K/N known at - # compile time all the trip counts are constants. - if n_work: - for _ in range_(n_work): - for _ in range_(m_row_blocks): - init_k(acc) - for _ in range_(k_iters): - # The l loop is unrolled by the B fifo depth so the - # acquired buffer index stays a compile-time - # constant. - for _ in range_(B_ITERS // B_DEPTH): - for _ in range(B_DEPTH): - # One B chunk feeds every A band, so B is - # acquired once around the band loop. - b = b_h.acquire(1) - for band in range(RHO): - a = a_h.acquire(1) - kstep_k(a, b, acc, band) - a_h.release(1) - b_h.release(1) - # Drain the accumulator. Unrolled by C_DEPTH for the - # same reason; a full O_CHUNKS unroll overflows program - # memory. - for chunk in range_(O_CHUNKS // C_DEPTH): - for half in range(C_DEPTH): - o = o_h.acquire(1) - epi_k(o, acc, chunk, half) - o_h.release(1) - if n_drain: - # The trailing partial column-block, for a column that sits it - # out. A is broadcast along the whole compute row, so this - # column must still consume its share or the columns that DO - # have work stall waiting for the fifo to advance. No B and no - # C here -- the runtime sequence issues neither for it. - for _ in range_(n_drain): - for _ in range_(m_row_blocks): - for _ in range_(k_iters): - for _ in range_(B_ITERS // B_DEPTH): - for _ in range(B_DEPTH): - for _ in range(RHO): - a_h.acquire(1) - a_h.release(1) + def core_fn(accs, o_h, b_h, a_h, init_k, kstep_k, epi_k, my_rtp, my_col, barrier): + """Core body. Every trip count and the activation come from the + runtime parameter buffer, so one core program serves every shape.""" + # The nest is here, not in the kernel, so every level has an acquire. + barrier.wait_for_value(1) + # Derived rather than sent, saving an RTP word: column c has work in + # block j iff (j*COLS + c)*N_TILE < N. Both divisors are powers of two, + # so this must leave no __divsi3 -- check the .o, not the .ll. Use // + # rather than >>; ScalarValue has no shift operators. + n_tiles = my_rtp[RTP_N_VAL] // N_TILE + n_work = (n_tiles - my_col[0] + COLS - 1) // COLS + n_drain = ((n_tiles + COLS - 1) // COLS) - n_work + n_row_blocks = my_rtp[RTP_M_ROW_BLOCKS] + n_k_iters = my_rtp[RTP_K_ITERS] + epi_mode = my_rtp[RTP_EPILOGUE] + clamp_min_bits = my_rtp[RTP_CLAMP_MIN] + clamp_max_bits = my_rtp[RTP_CLAMP_MAX] + # An absent slot becomes a compile-time constant rather than a load: at + # M_CHUNK == 1 both chunk counts are just n_row_blocks. + if "n_chunks" in rtp_slots: + n_chunks = my_rtp[rtp_slots["n_chunks"]] + n_units_rt = my_rtp[rtp_slots["n_units"]] + else: + n_chunks = n_row_blocks + n_units_rt = n_row_blocks + # Acquire does not consume the barrier, so take it back to zero or the + # next dispatch re-reads these instead of waiting. Safe before the + # work: the sequence cannot re-set it until this dispatch's C drains. + barrier.release_with_value(1) + + def sweep(group): + """One k reduction feeding ``group`` accumulators off a shared B. + + ``group`` is a Python list, so the mc loops unroll. b_h is acquired + outside them, so DDR reads B once per len(group) row-blocks. + """ + for a_acc in group: + init_k(a_acc) + for _ in range_(n_k_iters): + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + b = b_h.acquire(1) + for a_acc in group: + for band in range(RHO): + a = a_h.acquire(1) + kstep_k(a, b, a_acc, band) + a_h.release(1) + b_h.release(1) + # Unrolled by C_DEPTH; a full O_CHUNKS unroll overflows program + # memory. + for a_acc in group: + for chunk in range_(O_CHUNKS // C_DEPTH): + for half in range(C_DEPTH): + o = o_h.acquire(1) + epi_k( + o, + a_acc, + chunk, + half, + epi_mode, + clamp_min_bits, + clamp_max_bits, + ) + o_h.release(1) + + for _ in range_(n_work): + for _ in range_(n_chunks): + sweep(accs) + + # Column-blocks this column sits out. A is broadcast along the row, so + # it must still consume its share or the columns that do have work will + # stall on the fifo. No B and no C; the sequence issues neither. + for _ in range_(n_drain): + # Every unit delivers a full M_CHUNK tiles of A, leftover or not, + # so an idle column drains that much per unit. + for _ in range_(n_units_rt): + for _ in range_(n_k_iters): + for _ in range_(B_ITERS // B_DEPTH): + for _ in range(B_DEPTH): + for _ in range(M_CHUNK * RHO): + a_h.acquire(1) + a_h.release(1) workers = [] for r in range(ROWS): for c in range(n_active_cols): - acc = Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}") + # Worker flattens nested fn_args, so the length stays + # compile-time in core_fn. + accs = [ + Buffer(type=ct_acc_ty, name=f"c_acc_{r}_{c}_{mc}") + for mc in range(M_CHUNK) + ] workers.append( Worker( - partial(core_fn, col_work[c], col_drain[c]), + core_fn, [ - acc, + accs, c_prod[(r, c)].prod(), b_cons[(r, c)], a_cons[(r, c)], acc_init, k_step, epilogue_chunk, + rtps[r][c], + my_cols[r][c], + barriers[r][c], ], + stack_size=STACK_SIZE, ) ) # --- Runtime ---------------------------------------------------------- # - # Every wrap below stays under the shim's wrap/size field (DMA_BD_MAX_WRAP): the - # largest are K_TILE=512 and ROWS*M_TILE=256. - # One transfer per (column-block, leg) instead of one per object. - # - # A single fill/drain may span MANY fifo objects -- the descriptor just - # walks them in the order the cores consume -- so a whole column-block's - # worth of A, B and C each go out as one task. Issuing per object instead - # meant a host-side await for every sweep, and those awaits were the - # serialisation: the next sweep could not start until the previous sweep's - # C had come all the way back. - # - # Dimension order must match the core loop nest exactly: for each - # column-block it walks mega_row, then k. Every wrap stays under the shim's - # wrap/size field (largest are K_TILE=512 and ROWS*M_TILE=256). - def a_taps(mega_col, r, mbs): - # Every (mega_row, k) block this compute row consumes for one - # column-block. A does not depend on mega_col; it is re-fetched per - # column-block because the cores re-consume it. - # - # Returns a LIST: one 4D descriptor normally, or one 3D descriptor per - # mega_row when the mega_row stride would overflow the shim BD's 20-bit - # iteration step (see a_split above). The split form carries the - # mega_row jump in the offset, which has no such limit. - if not a_split: + # One transfer per (column-block, leg), not one per object: a descriptor + # walks many fifo objects in consume order, and per-object issue meant a + # host await per sweep. Dimension order must match the core's nest. + def a_taps(mega_col, r, units): + # Every (row-block, k) block this row consumes for one column-block, + # k outermost. A does not depend on mega_col; it is re-fetched because + # the cores re-consume it. + if M_CHUNK == 1 and not a_split: return [ TensorAccessPattern( tensor_dims=(M * K,), @@ -708,53 +681,53 @@ def a_taps(mega_col, r, mbs): 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], + taps = [] + for u in units: + first, count = unit_rows(u) + taps.append( + TensorAccessPattern( + tensor_dims=(M * K,), + offset=first * ROWS * M_TILE * K + r * M_TILE * K, + sizes=[k_iters, M_CHUNK, M_TILE, K_TILE], + strides=[K_TILE, ROWS * M_TILE * K, K, 1], + ) ) - for mb in mbs - ] + return taps def b_tap(mega_col, c): # Every (mega_row, k) chunk this column consumes. B does not depend on - # mega_row, hence the 0 stride: the same k-blocks are replayed for each - # row-block, which is what the cores expect. - # - # B must arrive PRE-PACKED (see GEMM.pack_B) so each k-block is one - # contiguous run. Expressing that reorder in the descriptor instead - # gives an innermost run of T=8 bf16, turning each 128 KB transfer into - # 8192 scattered bursts -- measured 5.4x slower end to end. + # mega_row, hence the 0 stride. It must arrive pre-packed so each + # k-block is one contiguous run; reordering in the descriptor instead + # gives an innermost run of T=8 bf16 and measured 5.4x slower. return TensorAccessPattern( tensor_dims=(K * N // B_GROUP,), offset=(mega_col * COLS + c) * N_TILE * K // B_GROUP, - sizes=( - [1, 1, 1, k_iters * K_TILE * N_TILE // B_GROUP] - if b_resident - else [m_row_blocks, k_iters, 1, K_TILE * N_TILE // B_GROUP] - ), - strides=( - [0, 0, 0, 1] if b_resident else [0, K_TILE * N_TILE // B_GROUP, 0, 1] - ), + # One k sweep per unit, not per row-block: the cores hold each B + # chunk across a group. The unit dimension keeps stride 0. + sizes=[n_units, k_iters, 1, K_TILE * N_TILE // B_GROUP], + strides=[0, K_TILE * N_TILE // B_GROUP, 0, 1], ) - def c_taps(mega_col, c, mbs): - # Every joined block this column produces for one column-block: one - # ROWS*M_TILE x N_TILE block per row-block. Returns a LIST, for the - # same reason a_taps does -- one descriptor per mega_row when N makes - # the mega_row stride overflow the shim BD's iteration step. + def c_taps(mega_col, c, units): + # Every joined block this column produces: one ROWS*M_TILE x N_TILE + # per row-block, in plain row-block order even under M_CHUNK. 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 - ] + taps = [] + for u in units: + first, count = unit_rows(u) + # One descriptor per row-block: grouping them would put the + # ROWS*M_TILE*N stride back in, which c_split exists to avoid. + for i in range(count): + taps.append( + TensorAccessPattern( + tensor_dims=(M * N,), + offset=(mega_col * COLS + c) * N_TILE + + (first + i) * ROWS * M_TILE * N, + sizes=[1, 1, ROWS * M_TILE, N_TILE], + strides=[0, 0, N, 1], + ) + ) + return taps return [_c_tap_unsplit(mega_col, c)] def _c_tap_unsplit(mega_col, c): @@ -766,43 +739,49 @@ def _c_tap_unsplit(mega_col, c): ) def sequence(A, B, C, a_prods, b_prods, c_conses): - # Column-blocks 0..n_full-1 use every column; the trailing one (when N - # is not a multiple of N_TILE*COLS) uses only the first rem_blocks. A - # is always issued for every row, because the columns sitting the - # trailing block out still drain their share of the broadcast. + # Write every core's parameters, then open every barrier. Both loops + # run to completion before the first fill is issued, so no core can + # read a half-written buffer. + for r in range(ROWS): + for c in range(n_active_cols): + rtps[r][c][RTP_N_VAL] = N + rtps[r][c][RTP_M_ROW_BLOCKS] = m_row_blocks + rtps[r][c][RTP_K_ITERS] = k_iters + rtps[r][c][RTP_EPILOGUE] = epilogue.mode + rtps[r][c][RTP_CLAMP_MIN] = clamp_min_bits + rtps[r][c][RTP_CLAMP_MAX] = clamp_max_bits + # Only what this configuration actually reads; see rtp_layout. + if "n_chunks" in rtp_slots: + rtps[r][c][rtp_slots["n_chunks"]] = n_chunks + rtps[r][c][rtp_slots["n_units"]] = n_units + for r in range(ROWS): + for c in range(n_active_cols): + barriers[r][c].set(1) + + # A trailing block uses only the first rem_blocks columns. A is still + # issued for every row, since the sitting-out columns drain it. blocks = [(mc, COLS) for mc in range(n_full)] if rem_blocks: blocks.append((n_full, rem_blocks)) - # One task per (column-block, leg): three per column instead of one per - # object, so a whole column-block retires on a single await rather than - # one per row-block. The C drain is issued first and retired last -- it - # is an S2MM that simply waits for the cores, so keeping it outstanding - # is what overlaps compute with write-back, and it must not share a - # group with the fills it depends on. - # Depth-2: issue column-block i+1 before retiring i, so its transfers - # are already moving while i computes. Retiring a block before issuing - # the next serialises on the C await, which waits for the cores. - # - # This is affordable only because each leg is a single task per - # column-block. Per-object tasks need 1 + 2*k_iters and cannot be - # overlapped at all. - # Keep OVERLAP column-blocks in flight, against SHIM_BDS buffer - # descriptors per column; the operator is DDR-rate bound rather than - # byte bound, so how deeply the fills are pipelined is what decides the - # rate. - # - # Every task of a block stays LIVE in its TaskGroup until the block is - # retired here. That is load-bearing, not tidiness -- see the - # dma_free_task note on a_split above. - all_mb = list(range(m_row_blocks)) - - # One emitter per leg, so the two paths below differ only in HOW they - # group and retire, not in how a leg is issued. + # C is issued first and retired last: keeping that S2MM outstanding + # overlaps compute with write-back, and it must not share a group with + # the fills it depends on. Tasks stay live until retired here. + all_mb = list(range(n_units)) + + # One emitter per leg, so the paths below differ only in how they + # group and retire. def issue_a(mega_col, mbs, group, wait=False): for r in range(ROWS): - for tap in a_taps(mega_col, r, mbs): - a_prods[r].fill(A, tap, group=group, wait=wait) + taps = a_taps(mega_col, r, mbs) + for i, tap in enumerate(taps): + # A leftover unit emits many fills back to back on one + # channel, so await every SHIM_TASK_QUEUE-th to stay inside + # the queue depth. + bounded = ( + len(taps) > SHIM_TASK_QUEUE and (i + 1) % SHIM_TASK_QUEUE == 0 + ) + a_prods[r].fill(A, tap, group=group, wait=wait or bounded) def issue_b(mega_col, active_cols, group): for c in range(active_cols): @@ -816,10 +795,8 @@ def issue_c(mega_col, active_cols, mbs, group): def emit_unsplit(): pending = [] for mega_col, active_cols in blocks: - # C in its own group, issued first and retired last: it is an - # S2MM that waits on the cores, so keeping it outstanding is - # what overlaps compute with write-back, and it must not share - # a group with the fills it depends on. + # C in its own group so it does not share one with the fills + # it depends on; see above. tg_c = TaskGroup() issue_c(mega_col, active_cols, all_mb, tg_c) tg_f = TaskGroup() @@ -836,43 +813,56 @@ def emit_unsplit(): 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. + """Issue split legs one unit at a time, retiring the oldest. + + One TaskGroup per unit, retired only when a new one would exceed + SHIM_TASK_QUEUE outstanding. ``pending`` is retired in append + order, which keeps a block's B and unsplit-leg descriptors alive + until its units have been awaited. + """ + # What a unit costs on the busiest channel. Under c_split it + # drains M_CHUNK C descriptors onto one, so counting units instead + # would overrun the queue by that factor. + unit_cost = _per_unit if c_split else 1 + pending = [] # (group, queue cost), oldest first + + def retire(limit): + while sum(q for _, q in pending) > limit: + pending.pop(0)[0].finish() + for mega_col, active_cols in blocks: tg_b = TaskGroup() issue_b(mega_col, active_cols, tg_b) - # The leg that did NOT split is still one task for the whole - # block -- its single descriptor already spans every mega_row, - # so re-issuing it per window would transfer the block twice. - # It stays live alongside the windows and retires with them. tg_whole = TaskGroup() if not c_split: issue_c(mega_col, active_cols, all_mb, tg_whole) if not a_split: issue_a(mega_col, all_mb, tg_whole) - for w in range(0, m_row_blocks, MB_WINDOW): - mbs = all_mb[w : w + MB_WINDOW] - tg_w = TaskGroup() + for u in all_mb: + # Before issuing, not after: fill/drain pushes the task + # immediately while TaskGroup.finish() emits the await, so + # retiring afterwards would leave the queue transiently one + # over. Await down to where this unit's transfers fit. + retire(SHIM_TASK_QUEUE - unit_cost) + tg_u = TaskGroup() if c_split: - issue_c(mega_col, active_cols, mbs, tg_w) + issue_c(mega_col, active_cols, [u], tg_u) if a_split: - # wait=True: the await is what makes this window's - # descriptors reusable by the next. - issue_a(mega_col, mbs, tg_w, wait=True) - tg_w.finish() - - tg_whole.finish() - tg_b.finish() + issue_a(mega_col, [u], tg_u, wait=True) + pending.append((tg_u, unit_cost)) + + # Not queue-counted: B and the unsplit leg ride channels the + # units do not contend for. Still retired in order. + pending.append((tg_whole, 0)) + pending.append((tg_b, 0)) + + # Drain everything, not retire(0): the tail groups are weighted 0, + # so a count-driven loop stops with them still open and the build + # fails with "Failed to close task groups". + for tg, _ in pending: + tg.finish() emit_split() if (a_split or c_split) else emit_unsplit() diff --git a/iron/operators/flm/gemm/op.py b/iron/operators/flm/gemm/op.py index f7c3be3ca..b87063b75 100644 --- a/iron/operators/flm/gemm/op.py +++ b/iron/operators/flm/gemm/op.py @@ -18,8 +18,8 @@ from aie.dialects.aie import get_target_model from aie.dialects._aie_enum_gen import AIEArch from iron.common.device_utils import get_kernel_dir +from iron.common.compilation import InstsBinArtifact, XclbinArtifact from iron.common.operator_bases import lut_based_ops_artifacts -from iron.common.utils import float_to_name import aie.utils as aie_utils from iron.operators.flm.packing import pack_b, packed_b_size @@ -27,6 +27,7 @@ BFP16_GROUP, BFP16_GROUP_BYTES, CT_MAX_K_FOR_N, + M_CHUNK_FOR_N, C_DEPTH, compute_rows, CT_OUT_LEN, @@ -39,6 +40,7 @@ S, T, _default_l1, + _hw_stride_ok, ) @@ -46,26 +48,28 @@ 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. + Fixed 64/512/128 tiling, no tiling knobs, and an activation plus optional + clamp folded into the output stage. The grid is as wide as the device: 8 + columns on NPU2, 4 on NPU1. See ``design.py``. """ # 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. + # stem from all of them. The build cache keys on filename, not on source or + # flags, so any field that changes the emitted MLIR or the kernel object + # must reach the stem or a stale build silently satisfies the request. M: int K: int N: int # Activation fused into the C drain. epilogue: Epilogue = Epilogue.NONE - # Optional (min, max) applied after the activation. + # The activations the epilogue can select between at run time. Each one + # compiled in costs program memory, so a deployment that dispatches two + # should compile two. Unlike `epilogue`, this is part of the + # configuration. + epilogue_modes: tuple[Epilogue, ...] = tuple(Epilogue) + # Optional (min, max) applied after the activation. The bounds are runtime + # parameters and the kernel always clamps, so this changes the instruction + # stream only -- clamped and unclamped callers share one xclbin. 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 @@ -74,6 +78,9 @@ class GEMM(MLIROperator): # 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 + # Row-blocks folded into one B fetch. __post_init__ resolves None from + # tile_n, falling back to 1 when it would not divide m_row_blocks. + m_chunk: int | None = None # Rounding for every f32->bf16 conversion; see Rounding in design.py. rounding: Rounding = Rounding.CONV_EVEN context: object = field(default=None, repr=False) @@ -83,44 +90,48 @@ class GEMM(MLIROperator): "epilogue": "epi", "tile_n": "tn", "tile_ma": "ma", + "m_chunk": "mc", "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. + # Resolve both tile knobs here, so the fields hold what the build + # actually uses. tile_ma especially must reach the artifact name: the + # design sizes the A object from it and the kernel derives the mmul's + # rowA from it. dev = aie_utils.get_current_device() if self.tile_n is None: - # n=64 gives the mmul colA=8 instead of 4, halving accumulator - # traffic per mac; n=128 halves A fetches instead. Which wins - # depends on whether compute or data movement is the critical path. - # - # On NPU2 that flips with K: with a single k iteration there is too - # little compute to hide the extra A traffic, so n=128 wins there - # (~9%), while n=64 wins by ~20% at K >= 1024. - # - # NPU1 never reaches that crossover. It has half the columns AND a - # quarter of the per-tile bf16 mac throughput (four native 4x8x4 - # macs per 8x8x8 shape, against NPU2's two bfp16-emulated ones), so - # it stays compute-bound at every K, and n=128's 32 KB f32 - # accumulator also overflows bank-aware L1 allocation. n=64 wins - # everywhere there by 1.21-1.38x, including at K=512 where NPU2's - # rule would pick 128. + # The trade flips with K on NPU2: one k iteration has too little + # compute to hide n=64's extra A traffic, so n=128 wins there by + # ~9% and n=64 by ~20% at K >= 1024. NPU1 has half the columns and + # a quarter of the per-tile bf16 throughput, so it stays + # compute-bound and n=64 wins at every K by 1.21-1.38x. single_k_iter = self.K // K_TILE <= 1 self.tile_n = 128 if (dev.arch == AIEArch.AIE2p and single_k_iter) else 64 elif self.tile_n not in CT_MAX_K_FOR_N: raise ValueError( f"tile_n must be one of {sorted(CT_MAX_K_FOR_N)}, got {self.tile_n}" ) + # m_chunk falls back to 1 unless both hold: it divides m_row_blocks (a + # partial group is inexpressible, see design.py), and the group's + # row-blocks sit ROWS*M_TILE*K apart inside the A descriptor, a stride + # that must fit the shim BD's 20-bit step. K=10240 overflows it where + # m_chunk=1 would not. + if self.m_chunk is None: + want = M_CHUNK_FOR_N[self.tile_n] + rows = M_TILE * compute_rows(dev) + m_row_blocks = self.M // rows if self.M % rows == 0 else 0 + fits = m_row_blocks and m_row_blocks % want == 0 + if fits and not _hw_stride_ok(compute_rows(dev) * M_TILE * self.K): + fits = False + self.m_chunk = want if fits else 1 if self.tile_ma is None: self.tile_ma = _default_l1( self.tile_n, CT_MAX_K_FOR_N[self.tile_n], self._b_elem_bytes, get_target_model(dev.resolve()).get_local_memory_size(), + self.m_chunk, )[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 @@ -136,6 +147,23 @@ def __post_init__(self): # the resolved fields still serialize into artifact names unchanged. self.epilogue = Epilogue(self.epilogue) self.rounding = Rounding(self.rounding) + # Deduplicated, since the mask ORs one bit per mode and a repeat would + # otherwise have to be tolerated by every consumer of the tuple. + self.epilogue_modes = tuple( + dict.fromkeys(Epilogue(m) for m in self.epilogue_modes) + ) + # A mode the mask leaves out reaches the kernel's default arm, which is + # NONE -- an unactivated result rather than a build or dispatch error. + # Refuse instead: this is the caller contradicting itself. + if ( + self.epilogue is not Epilogue.NONE + and self.epilogue not in self.epilogue_modes + ): + raise ValueError( + f"epilogue {self.epilogue} is not in epilogue_modes " + f"{tuple(str(m) for m in self.epilogue_modes)}, so it would not " + "be compiled in and the kernel would silently apply none" + ) if self.clamp is not None: lo, hi = self.clamp if lo > hi: @@ -144,24 +172,72 @@ def __post_init__(self): MLIROperator.__init__(self, context=self.context) @property - def name(self) -> str: - """Artifact stem, prefixed to disambiguate from ``iron.operators.GEMM``. + def _epilogue_mask(self) -> int: + """Bitmask of the modes compiled into the epilogue. Mode 0 is always + present -- the kernel falls back to it. - ``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. + OR rather than sum: ``__post_init__`` deduplicates, but a sum would + make that a correctness requirement rather than tidiness, since two + copies of a mode carry into the neighbouring mode's bit. """ - return f"FLM_{super().name}" + mask = 1 + for m in self.epilogue_modes: + mask |= 1 << Epilogue(m).mode + return mask + + @property + def config_name(self) -> str: + """Stem of the artifacts that do not depend on the shape. + + Everything here shapes the device configuration, and so the xclbin. M, + K, N, the activation and the clamp bounds are absent: they are runtime + parameters, so they reach the instruction stream instead -- see + ``name``. + + ``ck`` needs naming separately because retuning CT_MAX_K_FOR_N moves it + while tn is unmoved, and tile_ma is caller-overridable. Omitting it + once served an xclbin built at one ck to a request for another. + """ + dev = aie_utils.get_current_device().resolve().name + return ( + f"FLM_GEMM_tn{self.tile_n}_ck{CT_MAX_K_FOR_N[self.tile_n]}" + f"_ma{self.tile_ma}_mc{self.m_chunk}" + f"_em{self._epilogue_mask:x}_{self.rounding}_{dev}" + ) + + @property + def name(self) -> str: + """Artifact stem for the instruction stream, which does depend on it. + + The configuration it runs on, then the runtime parameters on top. That + also inherits ``config_name``'s prefix, which disambiguates from + ``iron.operators.GEMM`` -- that class would otherwise share a stem and + satisfy this operator's cache lookups. + + Every runtime parameter has to appear, because the sequence writes them + as immediates and the build cache keys on filename and mtime: a stem + that omits one serves the first caller's instruction stream to the + second and silently applies the first caller's values. The clamp bounds + go in as raw bit patterns, so bounds that differ only below the printed + precision still get their own stem. + """ + base = f"{self.config_name}_M{self.M}_K{self.K}_N{self.N}" + if self.epilogue != Epilogue.NONE: + base = f"{base}_epi{self.epilogue}" + if self.clamp is not None: + lo, hi = ( + int(np.float32(v).view(np.int32)) & 0xFFFFFFFF for v in self.clamp + ) + base = f"{base}_cl{lo:08x}{hi:08x}" + return base @property def _bfp16_b(self) -> bool: """Whether B is stored as bfp16ebs8 rather than bf16. - AIE2P only, and the reason both mmul templates in the kernel header are - live rather than one being dead code: on AIE2 the scalar BFP types do - not exist, so B stays bf16 and the mmul lowers onto four native 4x8x4 - macs. + AIE2P only, which is why both mmul templates in the kernel header are + live: on AIE2 the scalar BFP types do not exist, so B stays bf16 and + the mmul lowers onto four native 4x8x4 macs. """ return aie_utils.get_current_device().arch == AIEArch.AIE2p @@ -174,20 +250,16 @@ def _b_elem_bytes(self) -> float: def _kernel_object(self) -> str: """Object name over every flag that changes the emitted code. - Everything the -D flags in ``get_kernel_artifacts`` carry has to appear - here: this repo's build cache keys on filename and mtime rather than on - source or flags, so an object built for one configuration would - otherwise silently satisfy a request for another. That includes r/t, - which set the blocked layout, and the epilogue flags, which since the - epilogue was folded into this translation unit shape the same object. + Every -D flag from ``get_kernel_artifacts`` has to appear, for the + cache reason above. ``ck`` looks derivable from tile_n, but that is a + tuning table: naming it means retuning an entry does not also require + wiping the build dir. """ - clamp = "" - if self.clamp is not None: - clamp = "_clamp" + "_".join(float_to_name(float(v)) for v in self.clamp) return ( f"mm_fused_{M_TILE}x{K_TILE}x{self.tile_n}" + f"_ck{CT_MAX_K_FOR_N[self.tile_n]}" f"_r{R}t{T}_ma{self.tile_ma}_{self.rounding}" - f"_epi{self.epilogue}{clamp}.o" + f"_em{self._epilogue_mask:x}.o" ) @property @@ -195,57 +267,106 @@ def _link_file(self) -> str: """What the design names as its kernel: the bare object, or the archive bundling it with the tanh LUT tables. - Only AIE2 evaluates the activations through a LUT (AIE2P has a native - vector tanh), and only an activation references tanh at all -- the plain - epilogue just converts and stores. When this is wrong the failure is a - LINK error for tanh_lut_ab/tanh_lut_cd rather than a compile error, so - it surfaces late. + Only AIE2 evaluates activations through a LUT, and only an activation + references tanh. Getting this wrong is a link error, so it surfaces + late. """ - if self.epilogue is not Epilogue.NONE and get_kernel_dir() == "aie2": - return f"{self.name}_kernels.a" + if ( + any(Epilogue(m) is not Epilogue.NONE for m in self.epilogue_modes) + and get_kernel_dir() == "aie2" + ): + # config_name, not name: this string reaches the design's + # link_with, so a shape in it would put the shape in the device + # configuration. + return f"{self.config_name}_kernels.a" return self._kernel_object - def get_mlir_artifact(self): + @property + def _reference_shape(self) -> tuple[int, int, int]: + """The shape the configuration-only module is emitted at. + + Its runtime sequence is discarded; only the device body reaches the + xclbin. The smallest valid shape keeps it cheap and makes the + shape-independence explicit. + """ + dev = aie_utils.get_current_device() + # M must be at least m_chunk row-blocks: a partial group is + # inexpressible (see design.py), and this module must build. + return ( + M_TILE * compute_rows(dev) * self.m_chunk, + MIN_K, + self.tile_n * dev.cols, + ) + + def _mlir_artifact(self, filename, M, K, N, epilogue, clamp): return PythonGeneratedMLIRArtifact( - f"{self.name}.mlir", + filename, DesignGenerator( self.operator_dir / "design.py", "gemm", (), { "dev": aie_utils.get_current_device(), - "M": self.M, - "K": self.K, - "N": self.N, + "M": M, + "K": K, + "N": N, "tile_n": self.tile_n, "tile_ma": self.tile_ma, - "epilogue": self.epilogue, + "m_chunk": self.m_chunk, + "epilogue": epilogue, + "clamp": clamp, "kernel_object": self._link_file, "trace_size": 0, }, ), ) + def get_mlir_artifact(self): + return self._mlir_artifact( + f"{self.name}.mlir", self.M, self.K, self.N, self.epilogue, self.clamp + ) + + def set_up_artifacts(self) -> None: + kernels = self.get_kernel_artifacts() + + # Emitted at a reference shape and activation, so every shape sharing + # this configuration reuses it. No clamp, not this instance's bounds: + # they reach only the discarded runtime sequence. + config_mlir = self._mlir_artifact( + f"{self.config_name}.mlir", + *self._reference_shape, + Epilogue.NONE, + None, + ) + self.xclbin_artifact = XclbinArtifact( + f"{self.config_name}.xclbin", + mlir_input=config_mlir, + dependencies=[config_mlir] + kernels, + ) + shape_mlir = self.get_mlir_artifact() + self.insts_artifact = InstsBinArtifact( + f"{self.name}.bin", + mlir_input=shape_mlir, + # aiecc compiles the cores on the way to an instruction stream, so + # this needs the kernel objects too. + dependencies=[shape_mlir] + kernels, + ) + self.add_artifacts([self.xclbin_artifact, self.insts_artifact]) + 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. + # mm_fused.cc includes zero.cc, which is per-architecture. A quoted + # include searches generic/ first, so the arch directory must 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. + # AIE2P lowers the 8x8x8 mmul onto two bfp16-emulated macs, which this + # selects; AIE2 lowers it onto four native bf16 macs and ignores it. + # MM_FUSED_BFP16_B rides along, since bfp16ebs8 storage needs the + # scalar BFP types. flags = [ # Tile geometry and register tiling, for the mmul. f"-DMM_FUSED_TILE_M={M_TILE}", @@ -261,28 +382,17 @@ def get_kernel_artifacts(self): # Output stage. f"-DMM_FUSED_OUT_CHUNK={CT_OUT_LEN}", f"-DMM_FUSED_C_DEPTH={C_DEPTH}", - f"-DMM_FUSED_EPILOGUE_MODE={self.epilogue.mode}", + f"-DMM_FUSED_EPILOGUE_MODE_MASK={self._epilogue_mask}", ] + arch_include if self._bfp16_b: flags += [ "-DAIE_API_EMULATE_BFLOAT16_MMUL_WITH_BFP16", "-DMM_FUSED_BFP16_B", ] - if self.clamp is not None: - lo, hi = self.clamp - # repr() rather than :g -- the latter renders -4.0 as "-4", and - # "-4f" is not a valid C float literal. - flags += [ - "-DMM_FUSED_CLAMP=1", - f"-DMM_FUSED_CLAMP_MIN={float(lo)!r}f", - f"-DMM_FUSED_CLAMP_MAX={float(hi)!r}f", - ] if self.rounding is Rounding.CONV_EVEN: - # ROUND_CONV_EVEN is mm.cc's flag, reused rather than inventing a - # second spelling, and its polarity is mm.cc's too: absent means the - # core's power-up floor mode, even though this operator defaults the - # other way. It covers both conversions in the kernel -- the mmul - # and the epilogue's f32->bf16 store -- which must agree. + # mm.cc's flag and polarity, reused: absent means the core's + # power-up floor mode, though this operator defaults the other way. + # Covers both conversions in the kernel, which must agree. flags.append("-DROUND_CONV_EVEN") kernel_obj = KernelObjectArtifact( @@ -298,9 +408,8 @@ def get_kernel_artifacts(self): ) if self._link_file == self._kernel_object: return [kernel_obj] - # The tanh LUT's coefficient tables live in their own translation unit - # in mlir-aie's runtime lib, so on AIE2 the kernel object alone leaves - # tanh_lut_ab/tanh_lut_cd undefined at link time. See _link_file. + # The tanh LUT tables live in their own translation unit, so on AIE2 + # the kernel object alone leaves them undefined at link time. return [ KernelArchiveArtifact( self._link_file, @@ -311,16 +420,11 @@ def get_kernel_artifacts(self): 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. + Flat uint8 bfp16ebs8 blocks on NPU2, flat bf16 on NPU1. Bound to the + operator because the layout depends on the resolved ``tile_n`` and the + device. Packing to consumption order is what makes both B hops linear + descriptors, freeing the dimensions a deep k slice needs. See + :mod:`iron.operators.flm.packing`. """ return pack_b( B, @@ -340,12 +444,9 @@ def packed_B_size(self, K, N): 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. + # On AIE2P B is quantized to bfp16ebs8, so it is declared in + # bytes and sized from what pack_B returns; a (K, N) bf16 spec + # would over-allocate by 1.78x. On AIE2 it is an element count. ( AIERuntimeArgSpec( "in", (self.packed_B_size(self.K, self.N),), dtype=np.uint8 diff --git a/iron/operators/flm/gemm/test.py b/iron/operators/flm/gemm/test.py index 0dd4312e1..88a64306b 100644 --- a/iron/operators/flm/gemm/test.py +++ b/iron/operators/flm/gemm/test.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import os + import pytest import aie.utils as aie_utils @@ -13,10 +15,12 @@ BFP16_GROUP, BFP16_GROUP_BYTES, CT_MAX_K_FOR_N, + M_CHUNK_FOR_N, Epilogue, M_TILE, R, Rounding, + SHIM_TASK_QUEUE, _b_depth_for, _default_l1, ) @@ -44,17 +48,11 @@ def get_params(): 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. + # One full sweep is N_TILE * COLS wide, so both it and the N values that + # leave a trailing partial column-block differ per device. The trailing + # case is the interesting one: some columns compute the block while the + # rest only drain the A broadcast, and real o/down projections always land + # there. At K = 512 tile_n defaults to 128, halving at K >= 1024. # fmt: off if dev_name == "npu2": # M, K, N, epilogue, clamp, rounding @@ -80,15 +78,14 @@ def get_params(): ( 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. + # mega_row iteration step, so that leg goes out as one transfer + # per mega_row against a bounded outstanding count. These are the + # real E4B FFN projections, unsupported until that landed, and + # M=2048 is what pushes past the bound. ( 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 + ( 2048, 10240, 2560, NONE, None, CONV_EVEN), # E4B down, 2x + ( 2048, 2560, 10240, NONE, None, CONV_EVEN), # E4B gateup, 2x ] 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 @@ -127,18 +124,14 @@ def get_params(): def check_on_device(operator, golden_ref, K, rounding=CONV_EVEN): """Run ``operator`` against its golden reference and return run_test's result. - Bounds the error in ABSOLUTE terms as a fraction of the accumulated mass, - i.e. the expected size of the K reduction before cancellation, - K * mean|a| * mean|b|. A plain relative tolerance cannot work: with signed A - the K-sum cancels by ~sqrt(K), so |C| ends up far smaller than the mass - while the error tracks the mass, leaving near-zero outputs uncheckable. - - The fraction is per-architecture, because the two lower the same 8x8x8 mmul - onto very different arithmetic: NPU2 emulates it with bfp16, which drops - mantissa bits, while NPU1 has no bfp16 and lowers onto four native bf16 macs - accumulating in f32 -- exact up to the f32->bf16 store, so ~20x tighter. - floor truncates rather than rounding to nearest, so its bias accumulates - over the K reduction instead of cancelling and gets a looser bound on both. + Bounds the error absolutely, as a fraction of the accumulated mass + K * mean|a| * mean|b|. A relative tolerance cannot work: with signed A the + K-sum cancels by ~sqrt(K), so |C| ends up far smaller than the mass the + error tracks, leaving near-zero outputs uncheckable. + + The fraction is per-architecture, since NPU2 emulates the mmul with bfp16 + while NPU1 accumulates four native bf16 macs in f32 (~20x tighter). floor + truncates, so its bias accumulates and gets a looser bound on both. """ mass = ( K @@ -196,51 +189,33 @@ def test_gemm(M, K, N, epilogue, clamp, rounding, aie_context): 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. +def test_gemm_split_leg_bounds(aie_context): + """K or N = 10240 overflows the shim BD's 20-bit mega_row step, so that leg + goes out one transfer per mega_row. Two unmodelled shim resources bound how + many may be live -- BD ids and the channel task queue -- and overrunning + either hangs silently. The live set is 4 + 2 + 2 = 8 of 16 descriptors; + assert that here, since retuning SHIM_TASK_QUEUE could break it silently. """ - from aie.dialects.aie import get_target_model - from iron.operators.flm.gemm.design import SHIM_TASK_QUEUE - dev = aie_utils.get_current_device() available = get_target_model(dev.resolve()).get_num_bds(0, 0) - worst = 1 + 2 * SHIM_TASK_QUEUE + worst = SHIM_TASK_QUEUE + 2 + 2 assert worst <= available, ( f"a fully split block needs {worst} shim BDs of {available}; " - "windowing no longer fits and the split shapes will hang" + "the 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. + # The square case splits both legs, which the real Gemma shapes never do + # (E4B's down overflows on K and its gate/up on N, never both), so it is + # the only cover for the two-sided path. GEMM(M=512, K=10240, N=10240, context=aie_context).compile() -def test_gemm_split_leg_windowing_runs(aie_context): +def test_gemm_split_leg_bounds_runs(aie_context): """Execute the two-sided split path, not just compile it. - test_gemm_split_leg_windowing above only compiles this shape: the failure - mode it guards against -- BD-id aliasing and shim task-queue overrun (see - that test's docstring) -- is a runtime device hang or silent corruption, - which compiling the MLIR cannot exercise. This dispatches the same shape on - hardware and checks the result. - - Regular rather than extensive despite being the largest shape here. What it - catches is a hang or silently wrong output, not a wrong number, and its - compile-only sibling is already regular, so leaving the executing half out - of the default run is the wrong side to err on. Costs ~8s against the - regular suite's ~13s. + The failure the sibling test guards against is a runtime hang or silent + corruption, which compiling cannot exercise. Regular rather than extensive + despite the size: ~8s against the suite's ~13s. """ M, K, N = 512, 10240, 10240 golden_ref = generate_golden_reference(M=M, K=K, N=N) @@ -254,16 +229,12 @@ def test_gemm_split_leg_windowing_runs(aie_context): def tile_option_params(): """Every (tile_n, tile_ma) the design accepts on this device. - The shape parameters above exercise only the DEFAULT tile geometry, because - __post_init__ resolves both knobs from the shape and the device. These cover - the knobs themselves, which change the blocked L1 layout: tile_n selects - CT_MAX_K and the B object width, tile_ma sets the mmul's rowA and the A - object height, and pack_B, the four stream-dimension lists and gather_dims - all key off them. A mismatch is silently wrong output rather than a build - error, so each combination has to actually run on hardware. - - The default tile_ma per tile_n stays in the regular suite; the overrides are - extensive, since each is its own kernel object and xclbin. + The shape parameters above only exercise the default geometry, since + __post_init__ resolves both knobs. These cover the knobs themselves, which + change the blocked L1 layout, and a mismatch is silently wrong output + rather than a build error, so each has to run on hardware. The defaults + stay in the regular suite; the overrides are extensive, since each is its + own kernel object and xclbin. """ dev = aie_utils.get_current_device() if dev is None or dev.resolve().name not in ("npu1", "npu2"): @@ -273,14 +244,17 @@ def tile_option_params(): params = [] for tile_n, ct_k in sorted(CT_MAX_K_FOR_N.items()): - default_ma = _default_l1(tile_n, ct_k, b_elem, l1)[0] + # m_chunk matters: the core holds a B chunk across that many + # accumulators, so it is what decides which A heights still fit. + m_chunk = M_CHUNK_FOR_N[tile_n] + default_ma = _default_l1(tile_n, ct_k, b_elem, l1, m_chunk)[0] # One full sweep of the grid at this tile_n, so every column has work. M, K, N = 256, 512, tile_n * dev.cols for tile_ma in (16, 32, 64): if M_TILE % tile_ma or tile_ma % (2 * R): continue try: - _b_depth_for(tile_ma, tile_n, ct_k, b_elem, l1) + _b_depth_for(tile_ma, tile_n, ct_k, b_elem, l1, m_chunk) except ValueError: continue # this A height leaves no room for B at this width marks = [] if tile_ma == default_ma else [pytest.mark.extensive] @@ -312,12 +286,101 @@ def test_gemm_tile_options(M, K, N, tile_n, tile_ma, aie_context): 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. + Both classes are named ``GEMM`` and MLIROperator.name derives the stem from + the class name, so with the cache keyed on filename the two operators would + 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 ) + + +def test_one_xclbin_serves_every_shape(aie_context): + """Several shapes back to back on one loaded xclbin. + + The parametrised tests cannot cover this: each gets a fresh context, so + the array is reconfigured between cases. Here the shapes share one, they + disagree on every parameter, and none may rebuild the xclbin. + """ + # Every shape must resolve to the same m_chunk, which shapes the core + # program and so the xclbin. These are all even in m_row_blocks and exclude + # the K that would overflow the A descriptor's step. + shapes = [ + (512, 1536, 2048, "none"), + (512, 1536, 256, "none"), # only 4 of 8 columns compute + (1024, 2048, 1536, "none"), + (512, 1536, 6144, "gelu"), + (512, 1536, 2048, "none"), # back to the first, after the rest + ] + xclbin = None + for M, K, N, epilogue in shapes: + operator = GEMM(M=M, K=K, N=N, epilogue=epilogue, context=aie_context) + golden_ref = generate_golden_reference( + M=M, K=K, N=N, epilogue=epilogue, scale=4.0 if epilogue == "none" else 0.5 + ) + mass = ( + K + * golden_ref["input"].abs().float().mean() + * golden_ref["input_b"].abs().float().mean() + ) + errors, _, _ = run_test( + operator, + { + "A": golden_ref["input"].flatten(), + "B": operator.pack_B(golden_ref["input_b"]), + }, + {"C": golden_ref["output"].flatten()}, + rel_tol=0.04, + abs_tol=float(0.004 * mass), + ) + assert not errors, f"{M}x{K}x{N} {epilogue} failed" + + stamp = ( + operator.xclbin_artifact.filename, + os.path.getmtime(operator.xclbin_artifact.filename), + ) + if xclbin is None: + xclbin = stamp + assert stamp == xclbin, f"{M}x{K}x{N} rebuilt the xclbin" + + +def test_one_xclbin_serves_every_clamp_bound(aie_context): + """Different clamp bounds back to back on one loaded xclbin. + + The bounds are runtime parameters, so they must not rebuild anything. + Separate from test_one_xclbin_serves_every_shape, which never clamps and so + cannot catch bounds leaking back into the configuration. + """ + M, K, N = 256, 512, 1024 + bounds = [(-2.0, 2.0), (-4.0, 4.0), (-0.5, 0.5)] + xclbin = None + for clamp in bounds: + operator = GEMM(M=M, K=K, N=N, clamp=clamp, context=aie_context) + golden_ref = generate_golden_reference( + M=M, K=K, N=N, clamp=clamp, scale=INPUT_SCALE + ) + errors, _, _ = check_on_device(operator, golden_ref, K) + assert not errors, f"clamp={clamp} produced wrong output" + + stamp = ( + operator.xclbin_artifact.filename, + os.path.getmtime(operator.xclbin_artifact.filename), + ) + if xclbin is None: + xclbin = stamp + assert stamp == xclbin, f"clamp={clamp} rebuilt the xclbin" + + # ...and neither does dropping the clamp: the kernel always clamps, and an + # unclamped caller neutralises it with (-inf, +inf) rather than compiling + # a second build. config_name rather than xclbin_artifact, which only + # exists once compile() has run. + clamped = GEMM(M=M, K=K, N=N, clamp=bounds[0], context=aie_context) + unclamped = GEMM(M=M, K=K, N=N, context=aie_context) + assert unclamped.config_name == clamped.config_name + # The bounds do reach the instruction stream, though, so they must reach + # its stem or the build cache serves one caller's stream to another. + assert unclamped.name != clamped.name + assert ( + clamped.name != GEMM(M=M, K=K, N=N, clamp=bounds[1], context=aie_context).name + ) diff --git a/requirements.txt b/requirements.txt index 7f7c43b37..fdedee171 100755 --- a/requirements.txt +++ b/requirements.txt @@ -9,11 +9,12 @@ # CUDA build served from PyPI. We therefore also pin torch to the "+cpu" local # version below, which is only available from the PyTorch CPU index. --index-url https://download.pytorch.org/whl/cpu +--find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/v1.4.3 --find-links https://github.com/Xilinx/mlir-aie/releases/expanded_assets/latest-wheels-4 --find-links https://github.com/Xilinx/llvm-aie/releases/expanded_assets/nightly --extra-index-url https://pypi.org/simple -mlir_aie==1.4.3.dev85+gdf48abc +mlir_aie==1.4.3 llvm-aie==22.0.0.2026090701+3e93bf7b black