File size: 5,560 Bytes
58ba6df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 | // SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Round 11: cmm compute for the k2 K-split layout with a WIDE output block (NT > 4 tile columns per core, e.g. qkv:
// 2 token tile rows x 12 tile columns). The 2 x NT fp32 accumulators do not fit DST, so the matmul runs in subblocks of
// SB = 4 columns (2 x 4 = 8 fp32 DST tiles, full-sync DST) and accumulates over the K blocks with packer L1
// accumulation into c_part (fp32, [row i][col n] = tile i * NT + n; block 0 overwrites). Then:
// - the partner row's partial (c_part tiles NT .. 2 NT - 1) is written by the BRISC into the partner's c_int;
// - this core's own partial (tiles 0 .. NT - 1) is copied into DST exactly (c_part is unpacked to DST in fp32,
// UnpackToDestFp32) and packed with L1 accumulation onto the partner's partial in c_int (exact fp32 sum) once the
// NCRISC released the c_gate token (partner partial landed);
// - + bias (row broadcast) -> c_out (the output dtype, e.g. bfp8), in chunks of 4 tiles.
// c_part is never popped (the BRISC reads its partner half after the compute is done with it; CB state resets per launch).
// compile-time: NB (K blocks of this core), KB, NT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/cb_api.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#ifdef CMM_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define KZONE(n) DeviceZoneScopedN(n)
#else
#define KZONE(n)
#endif
void kernel_main() {
constexpr uint32_t NB = get_compile_time_arg_val(0);
constexpr uint32_t KB = get_compile_time_arg_val(1);
constexpr uint32_t NT = get_compile_time_arg_val(2);
constexpr uint32_t MPC = CMM_MPC;
#ifdef CMM_K2W_SB
constexpr uint32_t SB = CMM_K2W_SB; // 2: half-sync DST (4 fp32 tiles, math / pack overlap); 4: full-sync DST
#else
constexpr uint32_t SB = 4;
#endif
constexpr uint32_t DSTC = MPC * SB; // fp32 DST tiles per acquire
constexpr uint32_t CH = 4; // tiles per DST acquire in the sum / bias passes (fits half- and full-sync fp32 DST)
static_assert(NT % SB == 0 && NT % CH == 0 && DSTC <= 8, "subblocks");
constexpr uint32_t NSB = NT / SB;
constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24, c_part = 26, c_gate = 22;
compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_part);
cb_reserve_back(c_gate, 1); // arm the gate (1-page CB, full until the NCRISC pops it)
cb_push_back(c_gate, 1);
matmul_block_init(c_a, c_w, 0, SB, MPC, KB);
cb_reserve_back(c_part, MPC * NT);
for (uint32_t b = 0; b < NB; ++b) {
KZONE("CBLK");
const uint32_t ca = (b & 1) ? c_a1 : c_a, cw = (b & 1) ? c_w1 : c_w;
{
KZONE("CWA");
cb_wait_front(ca, KB * MPC);
}
{
KZONE("CWW");
cb_wait_front(cw, KB * NT);
}
for (uint32_t sb = 0; sb < NSB; ++sb) {
tile_regs_acquire();
#ifndef CMM_NO_MM
for (uint32_t k = 0; k < KB; ++k) {
matmul_block(ca, cw, k, k * NT + sb * SB, 0, 0, SB, MPC, KB);
}
#endif
tile_regs_commit();
tile_regs_wait();
if (b == 1 && sb == 0) {
PACK((llk_pack_reconfig_l1_acc(1))); // block 0 overwrote, every later block accumulates
}
for (uint32_t i = 0; i < MPC; ++i) {
for (uint32_t c = 0; c < SB; ++c) {
pack_tile<true>(i * SB + c, c_part, i * NT + sb * SB + c);
}
}
tile_regs_release();
}
cb_pop_front(ca, KB * MPC);
cb_pop_front(cw, KB * NT);
}
PACK((llk_pack_reconfig_l1_acc(0)));
cb_push_back(c_part, MPC * NT);
// own partial + the partner's partial (exact fp32: unpack-to-dest copy, packer L1 accumulation)
cb_wait_front(c_part, MPC * NT);
{
KZONE("CSUM");
reconfig_data_format_srca(c_part);
copy_tile_to_dst_init_short(c_part);
pack_reconfig_data_format(c_int);
{
KZONE("CXWAIT");
cb_reserve_back(c_gate, 1); // the partner's partial is in c_int
}
cb_reserve_back(c_int, NT);
PACK((llk_pack_reconfig_l1_acc(1)));
for (uint32_t q = 0; q < NT; q += CH) {
constexpr uint32_t n = CH;
tile_regs_acquire();
for (uint32_t t = 0; t < n; ++t) {
copy_tile(c_part, q + t, t);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t t = 0; t < n; ++t) {
pack_tile<true>(t, c_int, q + t);
}
tile_regs_release();
}
PACK((llk_pack_reconfig_l1_acc(0)));
cb_push_back(c_int, NT);
}
// + bias -> output
KZONE("CBIAS");
cb_wait_front(c_int, NT);
cb_wait_front(c_bias, NT);
reconfig_data_format(c_int, c_bias);
pack_reconfig_data_format(c_out);
add_bcast_rows_init(c_int, c_bias);
cb_reserve_back(c_out, NT);
for (uint32_t q = 0; q < NT; q += CH) {
tile_regs_acquire();
for (uint32_t i = 0; i < CH; ++i) {
add_tiles_bcast_rows(c_int, c_bias, q + i, q + i, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < CH; ++i) {
pack_tile(i, c_out);
}
tile_regs_release();
}
cb_push_back(c_out, NT);
cb_pop_front(c_int, NT);
cb_pop_front(c_bias, NT);
}
|