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 /tfuse_reader.cpp
changh95's picture
changh95
tt-model push meteor-p150 (container)
e91e989 verified
Raw History Blame Contribute Delete
3.6 kB
// SPDX-License-Identifier: Apache-2.0
// TFUSE reader (meteor-p150 OPT2 MK-TF, tt/tgate.py fuse_mm): the temporal fuse's merged 1x1 conv AND its gate glue
// in one program. Per 32-row block b of the bf16 raw map x (flat [R, 96] or the 4-D [1, H, W, 96] TILE form; both
// are 3 tiles per 32-row block): tiles b * 3 + 0..2 into CB 8. Once: the 12 fp32 weight tiles of W^T [96, 128]
// (tile k * 4 + j: input channels 32k.., output columns 32j..; columns 0..95 tfuse.0, 96..98 the gate logit
// differences, the rest zero) into CB 6, the ones tile (CB 2), the 3 tfuse.0 bias tiles (CB 3) and the two gate
// bias tiles (hi / lo bf16 terms of b_j - b_0, every row; CB 7).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous block
// range [i * B + min(i, E), + B + (i < E)) of the NB blocks, B = NB / NC, E = NB % NC.
// CT args: [0] NB, [1] NC, [2] GX, then the TensorAccessorArgs of x, w, ones, bias, gbias.
// Common RT args: [x_addr, w_addr, ones_addr, bias_addr, gbias_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 auto x_args = TensorAccessorArgs<3>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr auto gb_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr uint32_t TB32 = 4096, TB16 = 2048;
constexpr uint32_t cb_one = 2, cb_bias = 3, cb_w = 6, cb_gb = 7, cb_x = 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 x = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB16);
const auto w = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB32);
const auto ones = TensorAccessor(o_args, get_common_arg_val<uint32_t>(2), TB32);
const auto bias = TensorAccessor(b_args, get_common_arg_val<uint32_t>(3), TB32);
const auto gbias = TensorAccessor(gb_args, get_common_arg_val<uint32_t>(4), TB32);
cb_reserve_back(cb_w, 12);
for (uint32_t t = 0; t < 12; ++t) {
noc_async_read(w.get_noc_addr(t), get_write_ptr(cb_w) + t * TB32, TB32);
}
cb_reserve_back(cb_one, 1);
noc_async_read(ones.get_noc_addr(0), get_write_ptr(cb_one), TB32);
cb_reserve_back(cb_bias, 3);
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read(bias.get_noc_addr(k), get_write_ptr(cb_bias) + k * TB32, TB32);
}
cb_reserve_back(cb_gb, 2);
for (uint32_t k = 0; k < 2; ++k) {
noc_async_read(gbias.get_noc_addr(k), get_write_ptr(cb_gb) + k * TB32, TB32);
}
noc_async_read_barrier();
cb_push_back(cb_w, 12);
cb_push_back(cb_one, 1);
cb_push_back(cb_bias, 3);
cb_push_back(cb_gb, 2);
for (uint32_t b = b0; b < b0 + n; ++b) {
cb_reserve_back(cb_x, 3);
const uint32_t wx = get_write_ptr(cb_x);
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read(x.get_noc_addr(b * 3 + k), wx + k * TB16, TB16);
}
noc_async_read_barrier();
cb_push_back(cb_x, 3);
}
}