Hugging Face's logo

changh95
/
meteor-p150

Robotics
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
meteor-p150 / code /tt_meteor /tt /kernels /bres_reader2.cpp
changh95's picture
changh95
tt-model push meteor-p150 (container)
e91e989 verified
Raw History Blame Contribute Delete
6.19 kB
// SPDX-License-Identifier: Apache-2.0
// BRES reader v2 (meteor-p150 OPT2 round 5, tt/bresize.py, METEOR_BRES=cache): bres_reader.cpp with reader-private
// L1 caches of the input blocks and the weight tiles. v1 read every entry's KM x CT input tiles and its weight tile
// from DRAM: a x2 resize reads each input tile ~10 times (the output blocks of one core share their input rows), so
// the kernel was bound by its DRAM reads. v2 keeps the input blocks (all NX x CT tiles of block k) in a direct-mapped
// cache of S slots (tag k) and the weight tiles in one of SW slots (tag w): a block's misses are read
// from DRAM, then its KM x CT input tiles and KM weight tiles are copied from the caches into CB 0 / CB 1 by NOC
// reads of this core's own L1, in the same order as v1. Compute and writer are v1's: the same tiles in the same order,
// so the output is bit-identical to v1.
// The v2 table (bresize.cache_plan) gives per output block and entry e the words [k_e, w_e, sx_e, sw_e]: the input
// block, the weight tile and their cache slots (k % S / w % SW, probed on the host so that no two distinct input
// blocks (weight tiles) of one output block share a slot: the misses of a block never overwrite a tile the same block
// still copies). A slot whose tag is not the key is refilled from DRAM. One barrier after the misses, one after the
// copies (the copies of block b complete before block b + 1 refills a slot).
// CT args: [0] NB, [1] NC, [2] GX, [3] KM, [4] CT, [5] PW, [6] NX, [7] S, [8] SW, 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]. CBs: 0 x tiles, 1 weight tiles, 8 table page (scratch),
// 9 input cache (S * NX * CT tiles, scratch), 10 weight cache (SW tiles, scratch), 11 tags (S + SW words, scratch).
#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 uint32_t S = get_compile_time_arg_val(7);
constexpr uint32_t SW = get_compile_time_arg_val(8);
constexpr auto x_args = TensorAccessorArgs<9>();
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 SLOT = NX * CT * TB16; // bytes of one cached input block
constexpr uint32_t cb_x = 0, cb_w = 1, cb_tab = 8, cb_xc = 9, cb_wc = 10, cb_tags = 11;
constexpr uint32_t NONE = 0xFFFFFFFFu;
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);
const uint32_t xc = get_write_ptr(cb_xc);
const uint32_t wc = get_write_ptr(cb_wc);
volatile tt_l1_ptr uint32_t* ent = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(tab);
volatile tt_l1_ptr uint32_t* xtag = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_tags));
volatile tt_l1_ptr uint32_t* wtag = xtag + S;
for (uint32_t s = 0; s < S; ++s) {
xtag[s] = NONE;
}
for (uint32_t s = 0; s < SW; ++s) {
wtag[s] = NONE;
}
for (uint32_t b = b0; b < b0 + n; ++b) {
noc_async_read(at.get_noc_addr(b), tab, PW * 4);
noc_async_read_barrier();
// misses: DRAM -> the caches
for (uint32_t e = 0; e < KM; ++e) {
const uint32_t k = ent[4 * e], w = ent[4 * e + 1], s = ent[4 * e + 2], sw = ent[4 * e + 3];
if (xtag[s] != k) {
xtag[s] = k;
uint32_t d = xc + s * SLOT;
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_read(ax.get_noc_addr(k * CT + ct), d, TB16);
d += TB16;
}
if constexpr (NX > 1) {
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_read(al.get_noc_addr(k * CT + ct), d, TB16);
d += TB16;
}
}
}
if (wtag[sw] != w) {
wtag[sw] = w;
noc_async_read(aw.get_noc_addr(w), wc + sw * TB16, TB16);
}
}
noc_async_read_barrier();
// copies: the caches -> CB 1 (KM weight tiles) and CB 0 (entry-major x_hi tiles, then x_lo), v1's order
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 s = ent[4 * e + 2], sw = ent[4 * e + 3];
noc_async_read(get_noc_addr(wc + sw * TB16), dw, TB16);
dw += TB16;
noc_async_read(get_noc_addr(xc + s * SLOT), dx, CT * TB16);
dx += CT * TB16;
}
if constexpr (NX > 1) {
for (uint32_t e = 0; e < KM; ++e) {
const uint32_t s = ent[4 * e + 2];
noc_async_read(get_noc_addr(xc + s * SLOT + CT * TB16), dx, CT * TB16);
dx += CT * TB16;
}
}
noc_async_read_barrier();
cb_push_back(cb_w, KM);
cb_push_back(cb_x, NX * KM * CT);
}
}