File size: 3,598 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 | // 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);
}
}
|