// SPDX-License-Identifier: Apache-2.0 // msda_mk (tt/msda_mk.py, BEVFORMER_MSDA_MK; OPT_PLAN section 4 MK4a / MK4b, D19): the MSDA block module megakernel, // data movement. One program per encoder MSDA call does what six (TSA) / four (SCA) programs did: // the query concat [A | B] (TSA: [v0 | query + pos]) -> the compact sampling + attention projections in ONE // K-blocked matmul (fp32 weights resident in L1) -> + bias -> loc = int32(cols + base [+ shift row]) and the fp32 // logits -> (the msda_k steps) in-kernel softmax, pixel-pair gathers, fp32 corner weights, packer-L1 sum -> the // merged-head output. // The same source runs on both data-movement RISC-Vs: PART 0 (NCRISC) owns the points [0, P/2), the per-block inputs // (the K tiles of the query rows from A / B, the LOC_TW base tiles) and the reduce scalers / masks / E tiles; PART 1 // (BRISC) owns the points [P/2, P), loads the weights / bias / shift once and writes the output tiles. // // Work: the items (b, j) (32-query block b, head block j) in BLOCK-MAJOR order w = b * H + j, contiguous ranges per // core (core index from the logical coordinates; common RT args only, probe P3). A core computes the front (the // projections and the glue) of every block its range touches (at most two partial blocks per core are recomputed by // a neighbour: identical values) and the MSDA of its own items. // Front -> gathers hand-off: the compute kernel packs each block's LOC_TW int32 loc tiles into cb_loc0 and cb_loc1 // (one copy per RISC) and its AT_TW fp32 logit tiles into cb_lgc (PART 0); the readers wait for them at the block's // first item and pop them after its last. The compute kernel runs the front one block ahead (CB depth 2 blocks). #include #include "api/dataflow/dataflow_api.h" #include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp" constexpr uint32_t PART = get_compile_time_arg_val(0); constexpr int32_t HV = get_compile_time_arg_val(1); constexpr int32_t WV = get_compile_time_arg_val(2); constexpr uint32_t NQ = get_compile_time_arg_val(3); constexpr uint32_t NB = get_compile_time_arg_val(4); constexpr uint32_t H = get_compile_time_arg_val(5); constexpr uint32_t SUB = get_compile_time_arg_val(6); constexpr uint32_t P = get_compile_time_arg_val(7); constexpr uint32_t NC = get_compile_time_arg_val(8); constexpr uint32_t GX = get_compile_time_arg_val(9); constexpr uint32_t KA = get_compile_time_arg_val(10); // K tiles of the query rows from A (per block) constexpr uint32_t KBT = get_compile_time_arg_val(11); // K tiles from B (0: A only) constexpr uint32_t LOC_TW = get_compile_time_arg_val(12); // loc tiles per block (H * 2P / 32) constexpr uint32_t AT_TW = get_compile_time_arg_val(13); // logit tiles per block (H * P / 32) constexpr uint32_t PAD_BITS = get_compile_time_arg_val(14); constexpr uint32_t HAS_SHIFT = get_compile_time_arg_val(15); // ACT (BEVFORMER_SCA_ACTIVE, SCA only): the per-calibration runtime tensor `nblk` (uint32 page, entries 0..SUB-1 = // the active 32-row blocks of each camera's rebatched rows, valid rows first) restricts the work to the active // blocks: item w = a * H + j over the NA = sum(nblk) active blocks a, block(a) = camera c * BPS + (a - prefix c). // The padding blocks behind each camera's rows are never computed (their output rows are never gathered). constexpr uint32_t ACT = get_compile_time_arg_val(16); // RAW (BEVFORMER_SCA_BASE=kernel): `base` is the raw grid base [Q, 32] fp32 (one tile per block, the 2P point columns // in 0..2P-1); PART 0 reads each block's tile with faces 0 / 2 also written into faces 1 / 3 (both column halves = // the point columns: the compact layout of two heads per tile), PART 1 loads the `aff` tiles (S, O rows replicated) constexpr uint32_t RAW = get_compile_time_arg_val(17); constexpr auto v_args = TensorAccessorArgs<18>(); constexpr auto a_args = TensorAccessorArgs(); constexpr auto b_args = TensorAccessorArgs(); constexpr auto w_args = TensorAccessorArgs(); constexpr auto bi_args = TensorAccessorArgs(); constexpr auto ba_args = TensorAccessorArgs(); constexpr auto s_args = TensorAccessorArgs(); constexpr auto o_args = TensorAccessorArgs(); constexpr auto n_args = TensorAccessorArgs(); constexpr auto f_args = TensorAccessorArgs(); constexpr uint32_t KT = KA + KBT; constexpr uint32_t NOUT = LOC_TW + AT_TW; constexpr uint32_t LOC_HS = 2 * P; constexpr uint32_t cb_in = PART ? 1 : 0, cb_geo = PART ? 3 : 2, cb_e = 4, cb_out = 16; constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12; constexpr uint32_t cb_x = 21, cb_wsa = 22, cb_bias = 23, cb_base = 25, cb_shift = 26; constexpr uint32_t cb_loc = PART ? 29 : 28, cb_lgc = 30; constexpr uint32_t cb_cnt = PART ? 8 : 7; constexpr uint32_t cb_aff = 31; constexpr uint32_t BASE_TILES = RAW ? 1 : LOC_TW; // ACT: the nblk page (PART 0's copy is pushed to the compute kernel) constexpr uint32_t PH = P / 2; constexpr uint32_t P0 = PART * PH; constexpr uint32_t NS = PH * 4; constexpr uint32_t TB16 = 2048, TB32 = 4096; constexpr uint32_t NWORK = NB * H; constexpr uint32_t BPS = NB / SUB; constexpr int32_t LOC_LIM = 1 << 28; static_assert(P % 2 == 0 && P <= 8, "msda_mk: P even, <= 8"); static_assert(NB % SUB == 0, "msda_mk: blocks split evenly over the value maps of a head"); static_assert(!RAW || LOC_HS == 16, "msda_mk: the raw base needs 16 point columns (two heads per tile)"); static_assert(32 % LOC_HS == 0 && (H * LOC_HS) % 32 == 0 && (H * P) % 32 == 0, "msda_mk: compact layouts"); uint32_t nblk[SUB > 0 ? SUB : 1]; // active block index -> absolute 32-row block (identity without ACT) static inline uint32_t blk(uint32_t a) { if constexpr (ACT) { uint32_t c = 0; while (c + 1 < SUB && a >= nblk[c]) { a -= nblk[c]; ++c; } return c * BPS + a; } else { return a; } } static inline uint32_t tile_elem(uint32_t r, uint32_t c) { return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15); } constexpr uint32_t FRAC16_ONE = 0x3F800000; // fp32 1.0 (MSDA_MK_FRAC16: 1 + q / 2^16 has q << 7 as mantissa) static inline uint32_t q16_to_f32(uint32_t q) { if (q == 0) { return 0; } const uint32_t n = 31 - __builtin_clz(q); return ((127 - 16 + n) << 23) | ((q << (23 - n)) & 0x7FFFFF); } static inline void zero_l1(uint32_t addr, uint32_t bytes) { const uint64_t z = get_noc_addr(MEM_ZEROS_BASE); for (uint32_t o = 0; o < bytes; o += MEM_ZEROS_SIZE) { noc_async_read(z, addr + o, MEM_ZEROS_SIZE); } } // PART 0: the KT query tiles of block b (A then B) and its LOC_TW base tiles template static inline void push_block_inputs(uint32_t b, const TA& A, const TB_& B, const TC& BA) { cb_reserve_back(cb_x, KT); cb_reserve_back(cb_base, BASE_TILES); const uint32_t xa = get_write_ptr(cb_x); const uint32_t ba = get_write_ptr(cb_base); for (uint32_t k = 0; k < KA; ++k) { noc_async_read_tile(b * KA + k, A, xa + k * TB16); } if constexpr (KBT > 0) { for (uint32_t k = 0; k < KBT; ++k) { noc_async_read_tile(b * KBT + k, B, xa + (KA + k) * TB16); } } if constexpr (RAW) { // fp32 faces are 1 KB: face f of the source tile -> faces f and f + 1 (f = 0, 2) for (uint32_t f = 0; f < 4; f += 2) { const uint64_t src = BA.get_noc_addr(b, f * 1024); noc_async_read(src, ba + f * 1024, 1024); noc_async_read(src, ba + (f + 1) * 1024, 1024); } } else { for (uint32_t t = 0; t < LOC_TW; ++t) { noc_async_read_tile(b * LOC_TW + t, BA, ba + t * TB32); } } noc_async_read_barrier(); cb_push_back(cb_x, KT); cb_push_back(cb_base, BASE_TILES); } #ifdef MSDA_MK_MCAST // BEVFORMER_MSDA_MK_MCAST (D20 round 2): PART 1's resident weights are read from DRAM once per grid row by the row's // first core (logical x 0) and multicast to the GX - 1 other cores of the row (mmres_k's / bevfusion K6's handshake: // each receiver zeroes its 'valid' semaphore, increments the sender's 'ready' semaphore and waits for valid = 1; the // sender waits for GX - 1, multicasts, flushes, multicasts valid = 1; the rectangle never contains the sender). // Every core of the grid takes part, also one without work (NC = the whole grid). constexpr uint32_t MC_VYS[] = {MC_VY}; #endif template inline void load_weights(const WA& W, uint32_t wa) { #ifdef MSDA_MK_MCAST const bool sender = get_absolute_logical_x() == 0; const uint32_t y = MC_VYS[get_absolute_logical_y()]; const uint32_t ready_addr = get_semaphore(0); const uint32_t valid_addr = get_semaphore(1); volatile tt_l1_ptr uint32_t* ready_ptr = reinterpret_cast(ready_addr); volatile tt_l1_ptr uint32_t* valid_ptr = reinterpret_cast(valid_addr); constexpr uint32_t NRECV = GX - 1; if (sender) { for (uint32_t t = 0; t < KT * NOUT; ++t) { noc_async_read_tile(t, W, wa + t * TB32); } noc_async_read_barrier(); const bool n1 = noc_index == 1; const uint32_t x0 = n1 ? MC_VXE : MC_VX1, x1 = n1 ? MC_VX1 : MC_VXE; noc_semaphore_wait(ready_ptr, NRECV); noc_semaphore_set(ready_ptr, 0); noc_async_write_multicast(wa, get_noc_multicast_addr(x0, y, x1, y, wa), KT * NOUT * TB32, NRECV, true); noc_async_writes_flushed(); *valid_ptr = 1; noc_semaphore_set_multicast(valid_addr, get_noc_multicast_addr(x0, y, x1, y, valid_addr), NRECV); noc_async_write_barrier(); } else { noc_semaphore_set(valid_ptr, 0); noc_semaphore_inc(get_noc_addr(MC_VX0, y, ready_addr), 1); noc_semaphore_wait(valid_ptr, 1); } #else #ifndef MSDA_MK_NOWLOAD for (uint32_t t = 0; t < KT * NOUT; ++t) { noc_async_read_tile(t, W, wa + t * TB32); } #endif #endif } void kernel_main() { const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x(); if (ci >= NC) { return; } uint32_t nwork = NWORK; if constexpr (ACT) { const auto n_acc = TensorAccessor(n_args, get_common_arg_val(8), 128); cb_reserve_back(cb_cnt, 1); const uint32_t na = get_write_ptr(cb_cnt); noc_async_read(n_acc.get_noc_addr(0), na, 128); noc_async_read_barrier(); volatile tt_l1_ptr uint32_t* nv = reinterpret_cast(na); uint32_t tot = 0; for (uint32_t c = 0; c < SUB; ++c) { const uint32_t v = nv[c] > BPS ? BPS : nv[c]; nv[c] = v; // the compute kernel reads the clamped counts nblk[c] = v; tot += v; } if constexpr (PART == 0) { cb_push_back(cb_cnt, 1); } nwork = tot * H; } const uint32_t w0 = ci * nwork / NC; const uint32_t w1 = (ci + 1) * nwork / NC; if (w0 >= w1) { #ifdef MSDA_MK_MCAST if constexpr (PART == 1) { // no work here, but the row's weight multicast needs every core const auto W = TensorAccessor(w_args, get_common_arg_val(3), TB32); cb_reserve_back(cb_wsa, KT * NOUT); load_weights(W, get_write_ptr(cb_wsa)); } #endif return; } const uint32_t bf = w0 / H; // active block indices (absolute blocks without ACT) const uint32_t bl = (w1 - 1) / H; const auto v_acc = TensorAccessor(v_args, get_common_arg_val(0), 128); const auto A = TensorAccessor(a_args, get_common_arg_val(1), TB16); const auto B = TensorAccessor(b_args, get_common_arg_val(2), TB16); const auto W = TensorAccessor(w_args, get_common_arg_val(3), TB32); const auto BI = TensorAccessor(bi_args, get_common_arg_val(4), TB32); const auto BA = TensorAccessor(ba_args, get_common_arg_val(5), TB32); const auto S = TensorAccessor(s_args, get_common_arg_val(6), TB32); const auto o_acc = TensorAccessor(o_args, get_common_arg_val(7), TB16); const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE); if constexpr (PART == 1) { // the projection weights [KT, NOUT] tiles (row-major tile order k * NOUT + n), the bias row and the shift // row (row 0 replicated over the 32 rows: the value binary_ng's row broadcast adds to every row) cb_reserve_back(cb_wsa, KT * NOUT); const uint32_t wa = get_write_ptr(cb_wsa); load_weights(W, wa); cb_reserve_back(cb_bias, NOUT); const uint32_t bia = get_write_ptr(cb_bias); for (uint32_t t = 0; t < NOUT; ++t) { noc_async_read_tile(t, BI, bia + t * TB32); } if constexpr (RAW) { const auto F = TensorAccessor(f_args, get_common_arg_val(9), TB32); cb_reserve_back(cb_aff, 2); const uint32_t fa = get_write_ptr(cb_aff); noc_async_read_tile(0, F, fa); noc_async_read_tile(1, F, fa + TB32); } noc_async_read_barrier(); cb_push_back(cb_wsa, KT * NOUT); cb_push_back(cb_bias, NOUT); if constexpr (RAW) { cb_push_back(cb_aff, 2); } if constexpr (HAS_SHIFT) { cb_reserve_back(cb_shift, LOC_TW); const uint32_t sa = get_write_ptr(cb_shift); for (uint32_t j = 0; j < LOC_TW; ++j) { noc_async_read_tile(j, S, sa + j * TB32); } noc_async_read_barrier(); for (uint32_t j = 0; j < LOC_TW; ++j) { volatile tt_l1_ptr uint32_t* t = reinterpret_cast(sa + j * TB32); for (uint32_t f = 0; f < 4; ++f) { const uint32_t src = (f & 1) * 256; for (uint32_t r = (f < 2 ? 1 : 0); r < 16; ++r) { for (uint32_t c = 0; c < 16; ++c) { t[f * 256 + r * 16 + c] = t[src + c]; } } } } cb_push_back(cb_shift, LOC_TW); } } else { // E_p (bf16): row p all ones; the softmax's reduce scalers; the two half masks; the logit pages pre-filled // with PAD_LOGIT (the msda_k PART-0 setup) cb_reserve_back(cb_e, P); const uint32_t e_addr = get_write_ptr(cb_e); zero_l1(e_addr, P * TB16); noc_async_read_barrier(); volatile tt_l1_ptr uint16_t* e = reinterpret_cast(e_addr); for (uint32_t p = 0; p < P; ++p) { for (uint32_t c = 0; c < 32; ++c) { e[p * (TB16 / 2) + tile_elem(p, c)] = 0x3F80; } } cb_push_back(cb_e, P); dataflow_kernel_lib::calculate_and_prepare_reduce_scaler(); dataflow_kernel_lib::calculate_and_prepare_reduce_scaler(); cb_reserve_back(cb_msk, 2); volatile tt_l1_ptr uint32_t* m = reinterpret_cast(get_write_ptr(cb_msk)); for (uint32_t h = 0; h < 2; ++h) { for (uint32_t r = 0; r < 32; ++r) { for (uint32_t c = 0; c < 32; ++c) { m[h * 1024 + tile_elem(r, c)] = (c >= h * PH && c < (h + 1) * PH) ? 0x3F800000u : 0u; } } } cb_push_back(cb_msk, 2); cb_reserve_back(cb_lg, 2); volatile tt_l1_ptr uint32_t* lg = reinterpret_cast(get_write_ptr(cb_lg)); for (uint32_t i = 0; i < 2048; ++i) { lg[i] = PAD_BITS; } // the first two blocks' inputs (the compute kernel runs the front one block ahead) push_block_inputs(blk(bf), A, B, BA); if (bf + 1 <= bl) { push_block_inputs(blk(bf + 1), A, B, BA); } } cb_reserve_back(cb_geo, 6); zero_l1(get_write_ptr(cb_geo), 6 * TB32); noc_async_read_barrier(); uint32_t prev = 0xFFFFFFFF; for (uint32_t a = bf; a <= bl; ++a) { const uint32_t b = blk(a); const uint32_t j0 = a == bf ? w0 - bf * H : 0; const uint32_t j1 = a == bl ? w1 - bl * H : H; cb_wait_front(cb_loc, LOC_TW); const uint32_t lbuf = get_read_ptr(cb_loc); uint32_t lgbuf = 0; if constexpr (PART == 0) { cb_wait_front(cb_lgc, AT_TW); lgbuf = get_read_ptr(cb_lgc); } for (uint32_t j = j0; j < j1; ++j) { const uint32_t vbase = (j * SUB + b / BPS) * HV * (WV + 1); cb_reserve_back(cb_geo, 3); cb_reserve_back(cb_in, NS); const uint32_t g_addr = get_write_ptr(cb_geo); const uint32_t in_base = get_write_ptr(cb_in); volatile tt_l1_ptr int32_t* loc = reinterpret_cast(lbuf + ((j * LOC_HS) / 32) * TB32); const uint32_t lcol = (j * LOC_HS) % 32; if constexpr (PART == 0) { // the head's logit tile: columns 0..P-1 <- compact columns (j P) % 32 + p; columns >= P keep PAD cb_reserve_back(cb_lg, 1); volatile tt_l1_ptr uint32_t* lg = reinterpret_cast(get_write_ptr(cb_lg)); volatile tt_l1_ptr uint32_t* src = reinterpret_cast(lgbuf + ((j * P) / 32) * TB32); const uint32_t acol = (j * P) % 32; for (uint32_t r = 0; r < 32; ++r) { const bool ok = b * 32 + r < NQ; for (uint32_t p = 0; p < P; ++p) { lg[tile_elem(r, p)] = ok ? src[tile_elem(r, acol + p)] : 0u; } } cb_push_back(cb_lg, 1); } volatile tt_l1_ptr uint32_t* fx = reinterpret_cast(g_addr); volatile tt_l1_ptr uint32_t* fy = reinterpret_cast(g_addr + TB32); #ifdef MSDA_MK_NOPTS for (uint32_t r = 0; r < 0; ++r) { #else for (uint32_t r = 0; r < 32; ++r) { #endif const bool row_ok = b * 32 + r < NQ; for (uint32_t p = P0; p < P0 + PH; ++p) { const uint32_t ew = tile_elem(r, p); const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128; if (!row_ok) { #ifdef MSDA_MK_FRAC16 fx[ew] = FRAC16_ONE; // decodes to 0.0 fy[ew] = FRAC16_ONE; #else fx[ew] = 0; fy[ew] = 0; #endif noc_async_read(zeros, blk, 128); noc_async_read(zeros, blk + 2 * TB16, 128); continue; } int32_t vx = loc[tile_elem(r, lcol + 2 * p)]; int32_t vy = loc[tile_elem(r, lcol + 2 * p + 1)]; vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx); vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy); const int32_t x0 = vx >> 16; const int32_t y0 = vy >> 16; #ifdef MSDA_MK_FRAC16 // BEVFORMER_MSDA_FRAC16 (D20 round 3): the fraction q (16 bits) as the fp32 bits of 1 + q / 2^16 // (exponent of 1.0, q in the top mantissa bits): one shift and or; the compute kernel's SFPU // weight step takes 1 + q / 2^16 - 1 = q / 2^16 exactly (= q16_to_f32(q), bit for bit) fx[ew] = FRAC16_ONE | ((static_cast(vx) & 0xFFFF) << 7); fy[ew] = FRAC16_ONE | ((static_cast(vy) & 0xFFFF) << 7); #elif defined(MSDA_MK_RAWFRAC) fx[ew] = static_cast(vx) & 0xFFFF; // (anatomy probe: no fp32 conversion) fy[ew] = static_cast(vy) & 0xFFFF; #elif defined(MSDA_MK_NOGEO) #else fx[ew] = q16_to_f32(static_cast(vx) & 0xFFFF); fy[ew] = q16_to_f32(static_cast(vy) & 0xFFFF); #endif const bool yv0 = y0 >= 0 && y0 < HV, yv1 = y0 + 1 >= 0 && y0 + 1 < HV; const bool xp = x0 >= -1 && x0 < WV; // page i = x0 + 1 in [0, WV] const uint32_t pg = vbase + static_cast(y0 * (WV + 1) + x0 + 1); #ifdef MSDA_MK_NOREQ if (pg == 0xFFFFFFFF) { noc_async_read(zeros, blk, 128); } continue; #endif #ifdef MSDA_MK_NOREAD if (false) { #else if (xp && yv0) { #endif noc_async_read(v_acc.get_noc_addr(pg), blk, 128); } else { noc_async_read(zeros, blk, 128); } #ifdef MSDA_MK_NOREAD if (false) { #else if (xp && yv1) { #endif noc_async_read(v_acc.get_noc_addr(pg + WV + 1), blk + 2 * TB16, 128); } else { noc_async_read(zeros, blk + 2 * TB16, 128); } } } noc_async_read_barrier(); cb_push_back(cb_geo, 3); cb_push_back(cb_in, NS); if constexpr (PART == 1) { if (prev != 0xFFFFFFFF) { cb_wait_front(cb_out, 1); noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16); noc_async_writes_flushed(); cb_pop_front(cb_out, 1); } prev = b * H + j; // output tile (b, j) of [Q, H * 32] } } cb_pop_front(cb_loc, LOC_TW); if constexpr (PART == 0) { cb_pop_front(cb_lgc, AT_TW); if (a + 2 <= bl) { push_block_inputs(blk(a + 2), A, B, BA); } } } if constexpr (PART == 1) { cb_wait_front(cb_out, 1); noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(prev), TB16); noc_async_write_barrier(); cb_pop_front(cb_out, 1); } }