File size: 10,554 Bytes
8eeb550 | 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 | // SPDX-License-Identifier: Apache-2.0
// PACK (meteor-p150 OPT round 4 R4.2, tt/pack.py; OPT_PLAN #8 MK7 output tail): the frame's fp32 outputs into the
// packed readback rows in ONE data-movement program on one RISC / NOC (reads and writes on the same NOC: no two-NOC
// small DRAM writes, SYS-1419). Replaces the stock per-output untilize / pad / RM reshape / concat / slice chain.
//
// Entry e (an output, logical [R, C] rows of fp32) lands as its R * C elements in row-major order at element offset
// O_e of the packed stream (O_e a multiple of 1024: whole 4 KB packed rows); the packed stream is cut into NS segment
// tensors of SEG_PAGES rows of 1024 fp32 (page p -> segment p / SEG_PAGES). The tail of the last row of an entry is
// padding (never read by the host).
// Work unit = 32 rows of one entry (u = U_e + g: rows [32 g, min(R, 32 g + 32))):
// KIND 0 (TILE fp32, interleaved, padded width Wt tiles): the Wt tiles of tile row g are read, the 32 x C valid
// elements reordered from the tile faces (face f = (row / 16) * 2 + (col % 32) / 16, row-major 16 x 16)
// into the staging buffer by the RISC;
// KIND 1 (ROW_MAJOR fp32, C * 4 % 64 == 0): the 32 row pages are read straight into the staging buffer.
// Then the 32 * C * 4 B (a multiple of 128 B; the last unit rounded up to 64 B) go out in pieces split at the 4 KB
// packed-row boundaries.
// CM entries (channel-major, OPT round 4 R4.3: the host's NHWC -> NCHW transpose done here; TILE only, R % 32 == 0,
// HW % 16 == 0): row r = n * HW + hw, column c lands at element (n * C + c) * HW + hw, so the unit is reordered into
// [C][32] and each channel's 32 values go out as one 128 B piece (two 64 B pieces when the unit crosses a plane). Reads of unit k + 1 are issued before unit k is reordered (two input / staging buffers).
// Work split (no per-core RT args, probe P3): core i = y * GX + x owns units [BND[i], BND[i + 1]) (CT table, balanced
// by the bytes moved on the host).
// CT args: [0] NC, [1] GX, [2] NE, [3] SEG_PAGES, [4] NS, [5] WMAX (tiles per unit buffer), [6] STG (staging bytes),
// [7 .. 7 + NC] BND, then per entry 8 words (KIND, R, C, Wt, U_e, O_e, CM, HW), then the TensorAccessorArgs of the
// TILE inputs, of the ROW_MAJOR inputs (one set per kind: every input is DRAM interleaved) and of the
// segments. Common RT args: NE input addresses,
// then NS segment addresses.
// CB 0: 2 pages of WMAX x 4 KB (input buffers); CB 1: 2 pages of STG bytes (staging buffers); both reserved once and
// used as plain L1 scratch, never pushed. (The first bring-up reserved 2 x WMAX pages of the 2-page CB 0 and hung in
// cb_reserve_back: OPT_REPORT "Known hangs" R4.2.)
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NC = get_compile_time_arg_val(0);
constexpr uint32_t GX = get_compile_time_arg_val(1);
constexpr uint32_t NE = get_compile_time_arg_val(2);
constexpr uint32_t SEG_PAGES = get_compile_time_arg_val(3);
constexpr uint32_t NS = get_compile_time_arg_val(4);
constexpr uint32_t WMAX = get_compile_time_arg_val(5);
constexpr uint32_t STG = get_compile_time_arg_val(6);
constexpr uint32_t BND0 = 7;
constexpr uint32_t META = BND0 + NC + 1;
constexpr uint32_t MW = 8; // meta words per entry
constexpr auto tile_args = TensorAccessorArgs<META + MW * NE>();
constexpr auto rm_args = TensorAccessorArgs<tile_args.next_compile_time_args_offset()>();
constexpr auto out_args = TensorAccessorArgs<rm_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096; // fp32 tile bytes == packed row bytes
constexpr uint32_t ROWB = 4096;
struct Unit {
uint32_t e, g, kind, rows, c, wt, o, cm, hw;
};
FORCE_INLINE Unit unit_of(uint32_t u) {
uint32_t e = 0;
while (e + 1 < NE && kernel_compile_time_args[META + MW * (e + 1) + 4] <= u) {
++e;
}
const uint32_t m = META + MW * e;
Unit r;
r.e = e;
r.kind = kernel_compile_time_args[m];
const uint32_t R = kernel_compile_time_args[m + 1];
r.c = kernel_compile_time_args[m + 2];
r.wt = kernel_compile_time_args[m + 3];
r.g = u - kernel_compile_time_args[m + 4];
r.o = kernel_compile_time_args[m + 5];
r.cm = kernel_compile_time_args[m + 6];
r.hw = kernel_compile_time_args[m + 7];
const uint32_t r0 = r.g * 32;
r.rows = (R - r0 < 32) ? R - r0 : 32;
return r;
}
FORCE_INLINE void issue_reads(const Unit& un, uint32_t buf) {
const uint32_t addr = get_common_arg_val<uint32_t>(un.e);
if (un.kind == 0) {
const auto src = TensorAccessor(tile_args, addr, TB);
for (uint32_t k = 0; k < un.wt; ++k) {
noc_async_read(src.get_noc_addr(un.g * un.wt + k), buf + k * TB, TB);
}
} else {
const uint32_t pb = un.c * 4;
const auto src = TensorAccessor(rm_args, addr, pb);
for (uint32_t i = 0; i < un.rows; ++i) {
noc_async_read(src.get_noc_addr(un.g * 32 + i), buf + i * pb, pb);
}
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
const uint32_t u0 = kernel_compile_time_args[BND0 + core];
const uint32_t u1 = kernel_compile_time_args[BND0 + core + 1];
if (u0 >= u1) {
return;
}
cb_reserve_back(0, 2); // CB 0: 2 pages of WMAX x 4 KB (the two input buffers)
const uint32_t inb = get_write_ptr(0);
cb_reserve_back(1, 2); // CB 1: 2 pages of STG bytes (the two staging buffers)
const uint32_t stb = get_write_ptr(1);
Unit cur = unit_of(u0);
issue_reads(cur, inb);
uint32_t slot = 0;
for (uint32_t u = u0; u < u1; ++u) {
noc_async_read_barrier(); // unit u is in buffer `slot`
// the other input buffer and staging buffer `slot` were write sources of earlier units: their writes must
// have left L1 before they are refilled
noc_async_writes_flushed();
Unit nxt = cur;
if (u + 1 < u1) {
nxt = unit_of(u + 1);
issue_reads(nxt, inb + (slot ^ 1) * WMAX * TB);
}
const uint32_t ib = inb + slot * WMAX * TB;
const uint32_t sb = stb + slot * STG;
const uint32_t n = cur.rows * cur.c; // elements of this unit
uint32_t src = ib;
if (cur.cm) {
// [C][32] from the tile faces: column c, rows i: tile c / 32, face (i / 16) * 2 + (c % 32) / 16, word
// (i % 16) * 16 + c % 16
volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sb);
for (uint32_t c = 0; c < cur.c; ++c) {
const uint32_t base = ib + (c >> 5) * TB + ((((c >> 4) & 1) * 256) + (c & 15)) * 4;
volatile tt_l1_ptr uint32_t* s0 = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base);
volatile tt_l1_ptr uint32_t* s1 = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base + 512 * 4);
for (uint32_t i = 0; i < 16; ++i) {
dst[i] = s0[i * 16];
dst[16 + i] = s1[i * 16];
}
dst += 32;
}
const uint32_t r0 = cur.g * 32;
const uint32_t n0 = r0 / cur.hw;
const uint32_t hw0 = r0 - n0 * cur.hw;
const uint32_t k = (cur.hw - hw0 < 32) ? cur.hw - hw0 : 32; // rows in plane n0 (16 or 32)
for (uint32_t c = 0; c < cur.c; ++c) {
const uint32_t sc = sb + c * 128;
for (uint32_t part = 0; part < 2; ++part) {
const uint32_t rows = part == 0 ? k : 32 - k;
if (rows == 0) {
continue;
}
uint32_t src2 = sc + (part == 0 ? 0 : k * 4);
uint32_t off = (cur.o + ((n0 + part) * cur.c + c) * cur.hw + (part == 0 ? hw0 : 0)) * 4;
uint32_t bytes = rows * 4;
while (bytes) {
const uint32_t page = off / ROWB;
const uint32_t inpg = off - page * ROWB;
const uint32_t len = (ROWB - inpg < bytes) ? ROWB - inpg : bytes;
const uint32_t sg = page / SEG_PAGES;
const auto dst2 = TensorAccessor(out_args, get_common_arg_val<uint32_t>(NE + sg), ROWB);
noc_async_write(src2, dst2.get_noc_addr(page - sg * SEG_PAGES, inpg), len);
src2 += len;
off += len;
bytes -= len;
}
}
}
cur = nxt;
slot ^= 1;
continue;
}
if (cur.kind == 0) {
volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(sb);
for (uint32_t i = 0; i < cur.rows; ++i) {
const uint32_t fr = (i >> 4) * 2; // face row
const uint32_t ro = (i & 15) * 16; // row offset inside the face
for (uint32_t c0 = 0; c0 < cur.c; c0 += 16) {
const uint32_t t = c0 >> 5;
const uint32_t f = fr + ((c0 >> 4) & 1);
volatile tt_l1_ptr uint32_t* s =
reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ib + t * TB + (f * 256 + ro) * 4);
const uint32_t m = (cur.c - c0 < 16) ? cur.c - c0 : 16;
for (uint32_t j = 0; j < m; ++j) {
dst[j] = s[j];
}
dst += m;
}
}
src = sb;
}
// out: n * 4 bytes (rounded up to 64 B) at byte offset (O_e + g * 32 * C) * 4, split at 4 KB rows
uint32_t bytes = (n * 4 + 63) & ~63u;
uint32_t off = (cur.o + cur.g * 32 * cur.c) * 4;
while (bytes) {
const uint32_t page = off / ROWB;
const uint32_t inpg = off - page * ROWB;
const uint32_t len = (ROWB - inpg < bytes) ? ROWB - inpg : bytes;
const uint32_t s = page / SEG_PAGES;
const auto dst = TensorAccessor(out_args, get_common_arg_val<uint32_t>(NE + s), ROWB);
noc_async_write(src, dst.get_noc_addr(page - s * SEG_PAGES, inpg), len);
src += len;
off += len;
bytes -= len;
}
cur = nxt;
slot ^= 1;
}
noc_async_write_barrier();
}
|