File size: 10,172 Bytes
e0adbab | 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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | // SPDX-License-Identifier: Apache-2.0
// MK-T reader (sceneseg-p150 OPT_PLAN 3.4, the D19 seg-tail module megakernel, tt/mkt.py), NCRISC.
//
// Core c (row-major logical index, c < NC) owns the n = RB + (c < RE) image rows ya .. ya + n - 1; band row r
// (0 .. n + 1) is image row ya - 1 + r (rows outside the image are zeros: the convs' zero padding).
// 0. the LOAD-time tensors once: decode_layer_9 weights (72 bf16 tiles, K tile (dy * 3 + dx) * 4 + k, N tiles 0 / 1)
// and bias K tile (rows 0 / 1 = bf16 hi / lo, 2 N tiles), margin weights (18 tiles) and bias K tile (1), and
// the 3 identical ones tiles (columns 0 / 1 = 1.0) the bias products use as activations.
// 1. phase 1, per tile column x (32 pixels): the column band, per band row and channel group g (16 channels):
// 34 pixel sticks of 32 B (pixel x0 - 1 .. x0 + 32; a TILE face of the bf16 map is 16 consecutive sticks of one
// group, so the band is 2 face reads + 2 halo-stick reads per (row, group) from DRAM; matched 64-byte
// alignment of source and destination), double buffered (the next column's DRAM reads are issued before this
// column's tiles are built); then the column's activation tiles A[(r * 3 + dx) * 4 + k] (pixels x0 + dx - 1 ..
// + 31, channels 32 k .. 32 k + 31) by four local 512-byte face copies each.
// 2. phase 2 (after the writer's token: the D store holds the 64-channel map of rows ya - 1 .. ya + n), per column:
// the margin conv's activation tiles A[(r * 3 + dx) * 2 + k] from the D store, four local face copies each.
// Every NoC read is waited on before the push that hands it to compute. No semaphores here.
// CT args: [0] H, [1] W, [2] NC, [3] GX, [4] RB, [5] RE, then the TensorAccessorArgs of the input map, w9, b9, wm, bm, ones.
// Common RT args: [map, w9, b9, wm, bm, ones].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t H = get_compile_time_arg_val(0);
constexpr uint32_t W = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t RB = get_compile_time_arg_val(4);
constexpr uint32_t RE = get_compile_time_arg_val(5);
constexpr auto in_args = TensorAccessorArgs<6>();
constexpr auto w9_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr auto b9_args = TensorAccessorArgs<w9_args.next_compile_time_args_offset()>();
constexpr auto wm_args = TensorAccessorArgs<b9_args.next_compile_time_args_offset()>();
constexpr auto bm_args = TensorAccessorArgs<wm_args.next_compile_time_args_offset()>();
constexpr auto one_args = TensorAccessorArgs<bm_args.next_compile_time_args_offset()>();
constexpr uint32_t XT = W / 32; // tile columns
constexpr uint32_t NMAX = RB + (RE ? 1 : 0);
constexpr uint32_t NRMAX = NMAX + 2;
constexpr uint32_t TB16 = 2048, TB32 = 4096, FACE = 512, STICK = 32;
constexpr uint32_t BS = 1152; // band region per (row, group)
constexpr uint32_t BUFB = NRMAX * 8 * BS; // one column band
constexpr uint32_t RSB = (W + 2) * STICK; // D store: bytes per (row slot, group)
constexpr uint32_t SLOTB = 4 * RSB; // D store: one image row (4 groups of 16 channels)
constexpr uint32_t cb_a = 0, cb_w9 = 1, cb_b9 = 2, cb_wm = 3, cb_bm = 4, cb_band = 5, cb_d = 6, cb_sync = 7,
cb_zero = 8, cb_one = 10;
static_assert(W % 32 == 0, "the map width must be whole tiles");
static_assert(NMAX >= 1 && NMAX <= 3, "1 to 3 rows per core (DEST: 2 n fp32 tiles + 2 bias tiles <= 8)");
#ifndef MKT_DBG
#define MKT_DBG 0
#endif
template <typename Acc>
FORCE_INLINE void load_const(const Acc& acc, uint32_t cb, uint32_t ntiles, uint32_t tb) {
cb_reserve_back(cb, ntiles);
const uint32_t l1 = get_write_ptr(cb);
for (uint32_t i = 0; i < ntiles; ++i) {
noc_async_read(acc.get_noc_addr(i), l1 + i * tb, tb);
}
noc_async_read_barrier();
cb_push_back(cb, ntiles);
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
const uint32_t n = RB + (core < RE ? 1 : 0);
const uint32_t ya = core * RB + (core < RE ? core : RE);
const uint32_t nr = n + 2;
const auto ain = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB16);
// zero face (64-aligned) for padding sticks / out-of-image rows
const uint32_t zero = (get_write_ptr(cb_zero) + 63u) & ~63u;
{
volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
for (uint32_t i = 0; i < FACE / 4; ++i) {
z[i] = 0;
}
asm volatile("" ::: "memory");
invalidate_l1_cache();
}
const uint64_t zero_noc = get_noc_addr(zero);
const uint32_t band = (get_write_ptr(cb_band) + 63u) & ~63u;
// band rows outside the image: zero once in both buffers (never written by the column loads)
for (uint32_t b = 0; b < 2; ++b) {
for (uint32_t r = 0; r < nr; ++r) {
const int32_t yy = static_cast<int32_t>(ya + r) - 1;
if (yy >= 0 && yy < static_cast<int32_t>(H)) {
continue;
}
for (uint32_t g = 0; g < 8; ++g) {
const uint32_t reg = band + b * BUFB + (r * 8 + g) * BS;
noc_async_read(zero_noc, reg, FACE);
noc_async_read(zero_noc, reg + FACE, FACE);
noc_async_read(zero_noc, reg + 2 * FACE, BS - 2 * FACE);
}
}
}
load_const(TensorAccessor(w9_args, get_common_arg_val<uint32_t>(1), TB16), cb_w9, 72, TB16);
load_const(TensorAccessor(b9_args, get_common_arg_val<uint32_t>(2), TB16), cb_b9, 2, TB16);
load_const(TensorAccessor(wm_args, get_common_arg_val<uint32_t>(3), TB16), cb_wm, 18, TB16);
load_const(TensorAccessor(bm_args, get_common_arg_val<uint32_t>(4), TB16), cb_bm, 1, TB16);
load_const(TensorAccessor(one_args, get_common_arg_val<uint32_t>(5), TB16), cb_one, 3, TB16);
// column band loads: region (r, g) of buffer b at band + b * BUFB + (r * 8 + g) * BS; stick s (pixel x0 + s - 1)
// at region + 32 + 32 s: the left halo stick at +32, pixels x0 .. at +64 (64-aligned), the right halo at +1088
auto load_band = [&](uint32_t x, uint32_t b) {
const uint32_t buf = band + b * BUFB;
for (uint32_t r = 0; r < nr; ++r) {
const int32_t yy = static_cast<int32_t>(ya + r) - 1;
if (yy < 0 || yy >= static_cast<int32_t>(H)) {
continue;
}
const uint32_t trow = static_cast<uint32_t>(yy) * XT;
for (uint32_t g = 0; g < 8; ++g) {
const uint32_t reg = buf + (r * 8 + g) * BS;
const uint32_t k = g >> 1, par = g & 1;
const uint32_t tile = (trow + x) * 4 + k;
noc_async_read(ain.get_noc_addr(tile, par * FACE), reg + 64, FACE);
noc_async_read(ain.get_noc_addr(tile, (2 + par) * FACE), reg + 64 + FACE, FACE);
if (x > 0) {
noc_async_read(ain.get_noc_addr(tile - 4, (2 + par) * FACE + 15 * STICK), reg + 32, STICK);
} else {
noc_async_read(zero_noc, reg + 32, STICK);
}
if (x + 1 < XT) {
noc_async_read(ain.get_noc_addr(tile + 4, par * FACE), reg + 64 + 2 * FACE, STICK);
} else {
noc_async_read(zero_noc, reg + 64 + 2 * FACE, STICK);
}
}
}
};
// ---------------- phase 1: decode_layer_9 activation tiles
constexpr uint32_t NA1 = 12; // tiles per band row
load_band(0, 0);
noc_async_read_barrier();
for (uint32_t x = 0; x < XT; ++x) {
if (x + 1 < XT) {
load_band(x + 1, (x + 1) & 1);
}
const uint32_t buf = band + (x & 1) * BUFB;
cb_reserve_back(cb_a, NRMAX * NA1); // fixed push size: a CB push never wraps
const uint32_t base = get_write_ptr(cb_a);
for (uint32_t r = 0; r < nr; ++r) {
for (uint32_t dx = 0; dx < 3; ++dx) {
for (uint32_t k = 0; k < 4; ++k) {
const uint32_t dst = base + ((r * 3 + dx) * 4 + k) * TB16;
const uint32_t s0 = buf + (r * 8 + 2 * k) * BS + 32 + dx * STICK;
const uint32_t s1 = s0 + BS;
noc_async_read(get_noc_addr(s0), dst, FACE);
noc_async_read(get_noc_addr(s1), dst + FACE, FACE);
noc_async_read(get_noc_addr(s0 + FACE), dst + 2 * FACE, FACE);
noc_async_read(get_noc_addr(s1 + FACE), dst + 3 * FACE, FACE);
}
}
}
noc_async_read_barrier(); // this column's tiles and the next column's band
cb_push_back(cb_a, NRMAX * NA1);
}
#if MKT_DBG != 1
// ---------------- phase 2: margin conv activation tiles from the D store
cb_wait_front(cb_sync, 1);
cb_pop_front(cb_sync, 1);
const uint32_t dst0 = get_write_ptr(cb_d);
constexpr uint32_t NA2 = 6;
for (uint32_t x = 0; x < XT; ++x) {
cb_reserve_back(cb_a, NRMAX * NA2);
const uint32_t base = get_write_ptr(cb_a);
for (uint32_t r = 0; r < nr; ++r) {
for (uint32_t dx = 0; dx < 3; ++dx) {
for (uint32_t k = 0; k < 2; ++k) {
const uint32_t dst = base + ((r * 3 + dx) * 2 + k) * TB16;
const uint32_t s0 = dst0 + r * SLOTB + (2 * k) * RSB + (x * 32 + dx) * STICK;
const uint32_t s1 = s0 + RSB;
noc_async_read(get_noc_addr(s0), dst, FACE);
noc_async_read(get_noc_addr(s1), dst + FACE, FACE);
noc_async_read(get_noc_addr(s0 + FACE), dst + 2 * FACE, FACE);
noc_async_read(get_noc_addr(s1 + FACE), dst + 3 * FACE, FACE);
}
}
}
noc_async_read_barrier();
cb_push_back(cb_a, NRMAX * NA2);
}
#endif
}
|