File size: 3,997 Bytes
e91e989 | 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 | // SPDX-License-Identifier: Apache-2.0
// BRES reader (meteor-p150 OPT2 round 2, tt/bresize.py): a separable resize of a bf16 TILE map written as ONE
// block-sparse matmul out = M @ x, M = kron(A_h, A_w) over the rows of the (flat or 4-D) input and output maps.
// Output row block b (32 output rows) = sum over its KM entries e of Wt[w_e] @ x[k_e] (KM fixed per program; unused
// entries point at the all-zero weight tile 0 and repeat a real input block, so every value read is finite).
// Per output row block b: its table page [k_0, w_0, k_1, w_1, ...] (uint32, PW words, 64 B multiple) into the
// private scratch CB 8, the KM weight tiles (bf16, CB 1) and the KM x CT input tiles (bf16, CB 0, entry-major:
// e * CT + ct), one barrier per block.
// Work split (no per-core RT args, P3): core i of the first NC cores (row-major over a GX-wide grid) owns blocks
// [i B + min(i, E), + B + (i < E)) of NB.
// NX = 2 (an fp32 map as its bf16 terms x = x_hi + x_lo): the KM x CT tiles of x_lo follow those of x_hi in CB 0.
// CT args: [0] NB, [1] NC, [2] GX, [3] KM, [4] CT, [5] PW, [6] NX, then the TensorAccessorArgs of x (x_hi), of the
// weight tiles, of the table and of x_lo (x again when NX = 1). Common RT args: [x_addr, w_addr, table_addr, lo_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NB = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t KM = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr uint32_t PW = get_compile_time_arg_val(5);
constexpr uint32_t NX = get_compile_time_arg_val(6);
constexpr auto x_args = TensorAccessorArgs<7>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto t_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto l_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048;
constexpr uint32_t cb_x = 0, cb_w = 1, cb_tab = 8;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NB / NC, E = NB % NC;
const uint32_t b0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
const auto ax = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB16);
const auto aw = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB16);
const auto at = TensorAccessor(t_args, get_common_arg_val<uint32_t>(2), PW * 4);
const auto al = TensorAccessor(l_args, get_common_arg_val<uint32_t>(3), TB16);
const uint32_t tab = get_write_ptr(cb_tab);
volatile tt_l1_ptr uint32_t* ent = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(tab);
for (uint32_t b = b0; b < b0 + n; ++b) {
noc_async_read(at.get_noc_addr(b), tab, PW * 4);
noc_async_read_barrier();
cb_reserve_back(cb_w, KM);
cb_reserve_back(cb_x, NX * KM * CT);
uint32_t dw = get_write_ptr(cb_w), dx = get_write_ptr(cb_x);
for (uint32_t e = 0; e < KM; ++e) {
const uint32_t k = ent[2 * e], w = ent[2 * e + 1];
noc_async_read(aw.get_noc_addr(w), dw, TB16);
dw += TB16;
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_read(ax.get_noc_addr(k * CT + ct), dx, TB16);
dx += TB16;
}
}
if constexpr (NX > 1) {
for (uint32_t e = 0; e < KM; ++e) {
const uint32_t k = ent[2 * e];
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_read(al.get_noc_addr(k * CT + ct), dx, TB16);
dx += TB16;
}
}
}
noc_async_read_barrier();
cb_push_back(cb_w, KM);
cb_push_back(cb_x, NX * KM * CT);
}
}
|