File size: 12,563 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 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 | // SPDX-License-Identifier: Apache-2.0
// K11 depthwise reader (sceneseg-p150 OPT_PLAN K11, tt/dw.py): the float32 TILE input map [1, 1, Hin*Win, C] (DRAM
// interleaved) of one depthwise kxk conv, gathered into one TILE page per (output pixel tile, channel tile, tap).
//
// Work: the NW = CT * PT items (channel tile ct, output pixel tile pt; ct-major) split over the first NC cores in
// row-major order, core i owning [i * B + min(i, E), + B + (i < E)) (no per-core RT args: probe P3). A core walks
// its items in chunks of at most CH pixel tiles of one ct. Per chunk:
// 1. the input rows the chunk needs are loaded into an L1 band: two planes per row (channel half h = 0 / 1, 64 B
// per pixel = 16 fp32 channels), P zero columns on each side; for stride 2 each plane is stored column-parity
// major (even padded columns, then odd), so every tap reads consecutive output pixels from consecutive band
// pixels. A 16-pixel face half of an input tile is one contiguous DRAM run: stride 1 reads it straight into
// the band; stride 2 reads the row into a staging row (CB 5, two rows ping-pong) and scatters its pixels to
// the parity planes with 64 B local copies (64 B DRAM reads were latency-bound: R2 profile).
// 2. per output tile, the KK TILE pages (one per tap) are assembled from the band: face (rg, h) = 16 output
// pixels of channel half h, copied in runs (split at output-row ends); out-of-image rows and pixels past
// Hout*Wout read a zero block. Every copy is a local NoC read; ONE barrier per output tile (KK pages pushed
// together, CB 0 holds 2 * KK pages), not per page.
// When ct changes the NWT weight tiles (tile ct * NWT + j; NWT = KK: every row of tile j = w[ct*32 .. +32, tap j];
// R7.1 DW_WPACK: NWT = ceil(KK / 8), rows 4l .. 4l+3 of tile j = tap 8j + l) and the bias tile ct are read into
// CBs 1 / 2 (resident until the next ct).
// CT args: [0] HIN, [1] WIN, [2] HOUT, [3] WOUT, [4] K, [5] S, [6] P, [7] CT, [8] PT, [9] NC, [10] GX, [11] CH,
// [12] NRMAX (band rows), [13] STAGE (stride 2: 1 = staging rows, 0 = 64 B DRAM reads), [14] NWT (weight
// tiles per channel tile), then the TensorAccessorArgs of x, w, b.
// Common RT args: [x_addr, w_addr, b_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t HIN = get_compile_time_arg_val(0);
constexpr uint32_t WIN = get_compile_time_arg_val(1);
constexpr uint32_t HOUT = get_compile_time_arg_val(2);
constexpr uint32_t WOUT = get_compile_time_arg_val(3);
constexpr uint32_t K = get_compile_time_arg_val(4);
constexpr uint32_t S = get_compile_time_arg_val(5);
constexpr uint32_t P = get_compile_time_arg_val(6);
constexpr uint32_t CT = get_compile_time_arg_val(7);
constexpr uint32_t PT = get_compile_time_arg_val(8);
constexpr uint32_t NC = get_compile_time_arg_val(9);
constexpr uint32_t GX = get_compile_time_arg_val(10);
constexpr uint32_t CH = get_compile_time_arg_val(11);
constexpr uint32_t NRMAX = get_compile_time_arg_val(12);
constexpr uint32_t STAGE = get_compile_time_arg_val(13);
constexpr uint32_t NWT = get_compile_time_arg_val(14);
constexpr auto x_args = TensorAccessorArgs<15>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr uint32_t KK = K * K;
constexpr uint32_t HWO = HOUT * WOUT;
constexpr uint32_t NW = CT * PT;
constexpr uint32_t TB = 4096;
constexpr uint32_t WP = WIN + 2 * P; // padded input width (pixels)
constexpr uint32_t PH = (WP + 1) / 2; // stride 2: pixels of the even-column half of a plane row
constexpr uint32_t RS = (S == 1) ? WP : 2 * PH; // pixels per plane row
constexpr uint32_t cb_x = 0, cb_w = 1, cb_b = 2, cb_band = 3, cb_zero = 4, cb_stage = 5;
constexpr uint32_t STAGE_ROW = WIN * 128; // one staging row: two planes of WIN pixels x 64 B
// L1 byte address of padded column j of band row r, channel half h
FORCE_INLINE uint32_t band_addr(uint32_t base, uint32_t r, uint32_t h, uint32_t j) {
uint32_t pos;
if constexpr (S == 1) {
pos = j;
} else {
pos = (j & 1) * PH + (j >> 1);
}
return base + (((r * 2 + h) * RS) + pos) * 64;
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NW / NC, E = NW % NC;
const uint32_t w0 = core * B + (core < E ? core : E);
const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
const auto xa = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB);
const auto wa = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB);
const auto ba = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB);
// the 1 KB zero block (16 pixels x 64 B)
const uint32_t zero = get_write_ptr(cb_zero);
{
volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
for (uint32_t i = 0; i < 256; ++i) {
z[i] = 0;
}
}
const uint64_t zero_noc = get_noc_addr(zero);
const uint32_t band = get_write_ptr(cb_band);
const uint32_t stage = (S == 2 && STAGE) ? get_write_ptr(cb_stage) : 0;
uint32_t cur_ct = 0xFFFFFFFF;
uint32_t w = w0;
while (w < w1) {
const uint32_t ct = w / PT;
const uint32_t pt0 = w - ct * PT;
uint32_t pt1 = pt0 + CH;
if (pt1 > PT) {
pt1 = PT;
}
if (pt1 - pt0 > w1 - w) {
pt1 = pt0 + (w1 - w);
}
if (ct != cur_ct) {
cb_reserve_back(cb_w, NWT);
uint32_t l1 = get_write_ptr(cb_w);
for (uint32_t t = 0; t < NWT; ++t) {
noc_async_read(wa.get_noc_addr(ct * NWT + t), l1, TB);
l1 += TB;
}
cb_reserve_back(cb_b, 1);
noc_async_read(ba.get_noc_addr(ct), get_write_ptr(cb_b), TB);
noc_async_read_barrier();
cb_push_back(cb_w, NWT);
cb_push_back(cb_b, 1);
cur_ct = ct;
}
// ---- 1. the band: padded input rows R0 .. R0 + NR - 1 (input row iy = R0 + r - P)
const uint32_t p_first = pt0 * 32;
uint32_t p_last = pt1 * 32 - 1;
if (p_last >= HWO) {
p_last = HWO - 1;
}
const uint32_t ya = p_first / WOUT, yb = p_last / WOUT;
const uint32_t R0 = ya * S;
const uint32_t NR = (yb - ya) * S + K; // <= NRMAX by construction (host)
for (uint32_t r = 0; r < NR; ++r) {
const int32_t iy = (int32_t)(R0 + r) - (int32_t)P;
if (iy < 0 || iy >= (int32_t)HIN) {
continue;
}
for (uint32_t h = 0; h < 2; ++h) {
for (uint32_t j = 0; j < P; ++j) {
noc_async_read(zero_noc, band_addr(band, r, h, j), 64);
noc_async_read(zero_noc, band_addr(band, r, h, P + WIN + j), 64);
}
}
const uint32_t qa = (uint32_t)iy * WIN;
const uint32_t qb = qa + WIN;
if constexpr (S == 1) {
uint32_t q = qa;
while (q < qb) {
uint32_t ge = (q | 15) + 1;
if (ge > qb) {
ge = qb;
}
const uint32_t n = ge - q;
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa),
n * 64);
}
q = ge;
}
} else if constexpr (STAGE == 0) {
// stride 2 without staging (R2.1 form, knob DW_STAGE=0): one 64 B DRAM read per pixel and half
for (uint32_t q = qa; q < qb; ++q) {
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa), 64);
}
}
} else {
const uint32_t st = stage + (r & 1) * STAGE_ROW;
uint32_t q = qa;
while (q < qb) {
uint32_t ge = (q | 15) + 1;
if (ge > qb) {
ge = qb;
}
const uint32_t n = ge - q;
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), st + h * (WIN * 64) + (q - qa) * 64,
n * 64);
}
q = ge;
}
// the staging row is complete (and the copies out of the other staging row, issued one row ago,
// are done before that row is refilled: this barrier covers them too)
noc_async_read_barrier();
noc_async_read_set_state(zero_noc);
for (uint32_t h = 0; h < 2; ++h) {
for (uint32_t px = 0; px < WIN; ++px) {
noc_async_read_with_state(st + h * (WIN * 64) + px * 64, band_addr(band, r, h, P + px), 64);
}
}
}
}
noc_async_read_barrier();
// ---- 2. KK TILE pages per output tile. The output pixels of the tile split into at most 8 segments (per
// 16-pixel row group: runs inside one output row, or past Hout*Wout); a tap only shifts each segment's band
// source, so the segments are computed once per tile. All copies are local: NoC state set once.
noc_async_read_set_state(zero_noc);
for (uint32_t pt = pt0; pt < pt1; ++pt) {
uint32_t seg_d[8], seg_n[8], seg_y[8], seg_x[8], ns = 0;
for (uint32_t rg = 0; rg < 2; ++rg) {
const uint32_t p = pt * 32 + rg * 16;
uint32_t y = p / WOUT;
uint32_t x = p - y * WOUT;
uint32_t i = 0;
while (i < 16) {
seg_d[ns] = rg * 2048 + i * 64;
if (p + i >= HWO) {
seg_n[ns] = (16 - i) * 64;
seg_y[ns] = 0xFFFF; // zero segment
seg_x[ns] = 0;
++ns;
break;
}
uint32_t n = WOUT - x;
if (n > 16 - i) {
n = 16 - i;
}
seg_n[ns] = n * 64;
seg_y[ns] = y;
seg_x[ns] = x;
++ns;
i += n;
x += n;
if (x == WOUT) {
x = 0;
++y;
}
}
}
cb_reserve_back(cb_x, KK);
uint32_t dst = get_write_ptr(cb_x);
for (uint32_t dy = 0; dy < K; ++dy) {
for (uint32_t dx = 0; dx < K; ++dx, dst += TB) {
for (uint32_t g = 0; g < ns; ++g) {
const uint32_t d0 = dst + seg_d[g];
const uint32_t rr = seg_y[g] * S + dy; // padded row (huge for a zero segment)
if (rr >= P && rr < P + HIN) {
const uint32_t src = band_addr(band, rr - R0, 0, seg_x[g] * S + dx);
noc_async_read_with_state(src, d0, seg_n[g]);
noc_async_read_with_state(src + RS * 64, d0 + 1024, seg_n[g]);
} else {
noc_async_read_with_state(zero, d0, seg_n[g]);
noc_async_read_with_state(zero, d0 + 1024, seg_n[g]);
}
}
}
}
noc_async_read_barrier();
cb_push_back(cb_x, KK);
}
w += pt1 - pt0;
}
}
|