File size: 4,146 Bytes
7a24154 | 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 | // SPDX-License-Identifier: Apache-2.0
// RSZ reader (frnet-p150 OPT round 2, tt/resize.py): the bilinear align-corners resize of an fp32 NHWC TILE map
// [1, 1, h * w, C] to a bf16 [1, 1, H * W, C] map, one work unit = (channel tile ct, output row y).
// Per unit: the two input rows h0(y), h1(y) at channel tile ct as whole tiles (w / 32 tiles each, CBs 0 / 1, in
// step), and the two row-weight tiles A0 = 1 - a(y), A1 = a(y) (every element, CB 2) filled from the row table.
// Once per program: the W-direction weight tiles B0 / B1 of every output tile column (CB 3, resident: tile 2 wt = B0,
// 2 wt + 1 = B1; row i of tile wt = weight of output pixel 32 wt + i) and the row table [H][4] u32 (h0, h1, bits(1 -
// a), bits(a)) into CB 4.
// Work split (no per-core RT args, P3): core i of the first NC cores (row-major) owns units [i*B + min(i, E), + B +
// (i < E)) of U = CT * H units (u -> ct = u / H, y = u % H).
// CT args: [0] U, [1] NC, [2] GX, [3] H, [4] CT, [5] NWI (= w / 32), [6] NWT (= W / 32), then the TensorAccessorArgs
// of the input map, the B tensor [W, 64] fp32 TILE and the row table (ROW_MAJOR, one page). Common RT args:
// [x_addr, b_addr, rowtab_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t U = 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 uint32_t H = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr uint32_t NWI = get_compile_time_arg_val(5);
constexpr uint32_t NWT = get_compile_time_arg_val(6);
constexpr auto x_args = TensorAccessorArgs<7>();
constexpr auto b_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto r_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_x0 = 0, cb_x1 = 1, cb_a = 2, cb_b = 3, cb_rt = 4;
FORCE_INLINE void fill(uint32_t addr, uint32_t bits) {
volatile tt_l1_ptr uint32_t* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(addr);
for (uint32_t k = 0; k < 1024; ++k) {
p[k] = bits;
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = U / NC, E = U % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
const auto ax = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB);
const auto ab = TensorAccessor(b_args, get_common_arg_val<uint32_t>(1), TB);
const auto ar = TensorAccessor(r_args, get_common_arg_val<uint32_t>(2), H * 16);
// resident: row table, then the B weight tiles
cb_reserve_back(cb_rt, 1);
const uint32_t rt = get_write_ptr(cb_rt);
noc_async_read(ar.get_noc_addr(0), rt, H * 16);
cb_reserve_back(cb_b, 2 * NWT);
{
uint32_t w = get_write_ptr(cb_b);
for (uint32_t k = 0; k < 2 * NWT; ++k) {
noc_async_read(ab.get_noc_addr(k), w, TB); // tile k of [W, 64]: (k / 2, k % 2) = (wt, B0 | B1)
w += TB;
}
}
noc_async_read_barrier();
cb_push_back(cb_b, 2 * NWT);
volatile tt_l1_ptr uint32_t* tab = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(rt);
for (uint32_t u = u0; u < u0 + n; ++u) {
const uint32_t ct = u / H, y = u % H;
const uint32_t h0 = tab[4 * y], h1 = tab[4 * y + 1];
cb_reserve_back(cb_a, 2);
const uint32_t wa = get_write_ptr(cb_a);
fill(wa, tab[4 * y + 2]);
fill(wa + TB, tab[4 * y + 3]);
cb_push_back(cb_a, 2);
for (uint32_t k = 0; k < NWI; ++k) {
cb_reserve_back(cb_x0, 1);
cb_reserve_back(cb_x1, 1);
noc_async_read(ax.get_noc_addr((h0 * NWI + k) * CT + ct), get_write_ptr(cb_x0), TB);
noc_async_read(ax.get_noc_addr((h1 * NWI + k) * CT + ct), get_write_ptr(cb_x1), TB);
noc_async_read_barrier();
cb_push_back(cb_x0, 1);
cb_push_back(cb_x1, 1);
}
}
}
|