code/tt_frnet/tt/kernels/rsz_reader.cpp

// 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);
        }
    }
}