frnet-p150 / code /tt_frnet /tt /kernels /rsz_writer.cpp

tt-model push frnet-p150 (container)

7a24154 verified 16 days ago

4.77 kB

// SPDX-License-Identifier: Apache-2.0
// RSZ writer (frnet-p150 OPT round 2, tt/resize.py): per unit (ct, y) it waits for the whole H-blended row V (NWI
// fp32 tiles in CB 5), then per output tile column wt gathers the rows of the 32 output pixels' two input columns
// j0(w), j1(w) (column table, resident in CB 8: [W][2] u32) from V into g0 / g1 (CBs 6 / 7: fp32, row i = V row of
// pixel j; 2 face rows of 64 B per tile row), and writes the bf16 output tile (CB 16) to tile (y * NWT + wt) * CT + ct.
// CT args: [0] U, [1] NC, [2] GX, [3] H, [4] CT, [5] NWI, [6] NWT, [7] W, then the TensorAccessorArgs of the output
// and the column table (ROW_MAJOR, one page). Common RT args: [out_addr, coltab_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 uint32_t W = get_compile_time_arg_val(7);
constexpr auto o_args = TensorAccessorArgs<8>();
constexpr auto c_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096, TB16 = 2048;
constexpr uint32_t cb_v = 5, cb_g0 = 6, cb_g1 = 7, cb_ct = 8, cb_out = 16;

// the fp32 tile row of pixel j (in V) -> row i of the tile at dst: two 64-byte face rows. RSZ_CPU_GATHER: plain
// unrolled word copies on this RISC-V; default: two local L1 -> L1 NoC reads (the NoC moves the bytes, the RISC-V
// only issues; one read barrier per gathered tile pair).
FORCE_INLINE void copy_row(uint32_t vbase, uint32_t j, uint32_t dst, uint32_t i) {
    const uint32_t r = j & 31;
    const uint32_t s = vbase + (j >> 5) * TB + (r >> 4) * 2048 + (r & 15) * 64;
    const uint32_t d = dst + (i >> 4) * 2048 + (i & 15) * 64;
#ifdef RSZ_CPU_GATHER
    const uint32_t* ps = reinterpret_cast<const uint32_t*>(s);
    uint32_t* pd = reinterpret_cast<uint32_t*>(d);
#pragma GCC unroll 16
    for (uint32_t k = 0; k < 16; ++k) {
        pd[k] = ps[k];
    }
#pragma GCC unroll 16
    for (uint32_t k = 0; k < 16; ++k) {
        pd[256 + k] = ps[256 + k];
    }
#else
    noc_async_read(get_noc_addr(s), d, 64);
    noc_async_read(get_noc_addr(s + 1024), d + 1024, 64);
#endif
}

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 ao = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), TB16);
    const auto ac = TensorAccessor(c_args, get_common_arg_val<uint32_t>(1), W * 8);
    cb_reserve_back(cb_ct, 1);
    const uint32_t ctab = get_write_ptr(cb_ct);
    noc_async_read(ac.get_noc_addr(0), ctab, W * 8);
    noc_async_read_barrier();
    const uint32_t* col = reinterpret_cast<const uint32_t*>(ctab);    // read after the barrier above
    for (uint32_t u = u0; u < u0 + n; ++u) {
        const uint32_t ct = u / H, y = u % H;
        cb_wait_front(cb_v, NWI);
        asm volatile("" ::: "memory");                      // no V load above the wait
        const uint32_t vbase = get_read_ptr(cb_v);
        // gather tile wt while the compute kernel blends tile wt - 1 (g0 / g1 hold 2 tiles), then write tile wt - 1
        for (uint32_t wt = 0; wt <= NWT; ++wt) {
            if (wt < NWT) {
                cb_reserve_back(cb_g0, 1);
                cb_reserve_back(cb_g1, 1);
                asm volatile("" ::: "memory");              // no g store above the reserve
                const uint32_t g0 = get_write_ptr(cb_g0), g1 = get_write_ptr(cb_g1);
#ifndef RSZ_SKIP_GATHER
                for (uint32_t i = 0; i < 32; ++i) {
                    const uint32_t w = wt * 32 + i;
                    copy_row(vbase, col[2 * w], g0, i);
                    copy_row(vbase, col[2 * w + 1], g1, i);
                }
#endif
#ifndef RSZ_CPU_GATHER
                noc_async_read_barrier();
#endif
                asm volatile("" ::: "memory");
                cb_push_back(cb_g0, 1);
                cb_push_back(cb_g1, 1);
            }
            if (wt > 0) {
                cb_wait_front(cb_out, 1);
                noc_async_write(get_read_ptr(cb_out), ao.get_noc_addr((y * NWT + wt - 1) * CT + ct), TB16);
                noc_async_writes_flushed();
                cb_pop_front(cb_out, 1);
            }
        }
        cb_pop_front(cb_v, NWI);
    }
    noc_async_write_barrier();
}

	// SPDX-License-Identifier: Apache-2.0
	// RSZ writer (frnet-p150 OPT round 2, tt/resize.py): per unit (ct, y) it waits for the whole H-blended row V (NWI
	// fp32 tiles in CB 5), then per output tile column wt gathers the rows of the 32 output pixels' two input columns
	// j0(w), j1(w) (column table, resident in CB 8: [W][2] u32) from V into g0 / g1 (CBs 6 / 7: fp32, row i = V row of
	// pixel j; 2 face rows of 64 B per tile row), and writes the bf16 output tile (CB 16) to tile (y * NWT + wt) * CT + ct.
	// CT args: [0] U, [1] NC, [2] GX, [3] H, [4] CT, [5] NWI, [6] NWT, [7] W, then the TensorAccessorArgs of the output
	// and the column table (ROW_MAJOR, one page). Common RT args: [out_addr, coltab_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 uint32_t W = get_compile_time_arg_val(7);
	constexpr auto o_args = TensorAccessorArgs<8>();
	constexpr auto c_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
	constexpr uint32_t TB = 4096, TB16 = 2048;
	constexpr uint32_t cb_v = 5, cb_g0 = 6, cb_g1 = 7, cb_ct = 8, cb_out = 16;

	// the fp32 tile row of pixel j (in V) -> row i of the tile at dst: two 64-byte face rows. RSZ_CPU_GATHER: plain
	// unrolled word copies on this RISC-V; default: two local L1 -> L1 NoC reads (the NoC moves the bytes, the RISC-V
	// only issues; one read barrier per gathered tile pair).
	FORCE_INLINE void copy_row(uint32_t vbase, uint32_t j, uint32_t dst, uint32_t i) {
	const uint32_t r = j & 31;
	const uint32_t s = vbase + (j >> 5) * TB + (r >> 4) * 2048 + (r & 15) * 64;
	const uint32_t d = dst + (i >> 4) * 2048 + (i & 15) * 64;
	#ifdef RSZ_CPU_GATHER
	const uint32_t* ps = reinterpret_cast<const uint32_t*>(s);
	uint32_t* pd = reinterpret_cast<uint32_t*>(d);
	#pragma GCC unroll 16
	for (uint32_t k = 0; k < 16; ++k) {
	pd[k] = ps[k];
	}
	#pragma GCC unroll 16
	for (uint32_t k = 0; k < 16; ++k) {
	pd[256 + k] = ps[256 + k];
	}
	#else
	noc_async_read(get_noc_addr(s), d, 64);
	noc_async_read(get_noc_addr(s + 1024), d + 1024, 64);
	#endif
	}

	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 ao = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), TB16);
	const auto ac = TensorAccessor(c_args, get_common_arg_val<uint32_t>(1), W * 8);
	cb_reserve_back(cb_ct, 1);
	const uint32_t ctab = get_write_ptr(cb_ct);
	noc_async_read(ac.get_noc_addr(0), ctab, W * 8);
	noc_async_read_barrier();
	const uint32_t* col = reinterpret_cast<const uint32_t*>(ctab); // read after the barrier above
	for (uint32_t u = u0; u < u0 + n; ++u) {
	const uint32_t ct = u / H, y = u % H;
	cb_wait_front(cb_v, NWI);
	asm volatile("" ::: "memory"); // no V load above the wait
	const uint32_t vbase = get_read_ptr(cb_v);
	// gather tile wt while the compute kernel blends tile wt - 1 (g0 / g1 hold 2 tiles), then write tile wt - 1
	for (uint32_t wt = 0; wt <= NWT; ++wt) {
	if (wt < NWT) {
	cb_reserve_back(cb_g0, 1);
	cb_reserve_back(cb_g1, 1);
	asm volatile("" ::: "memory"); // no g store above the reserve
	const uint32_t g0 = get_write_ptr(cb_g0), g1 = get_write_ptr(cb_g1);
	#ifndef RSZ_SKIP_GATHER
	for (uint32_t i = 0; i < 32; ++i) {
	const uint32_t w = wt * 32 + i;
	copy_row(vbase, col[2 * w], g0, i);
	copy_row(vbase, col[2 * w + 1], g1, i);
	}
	#endif
	#ifndef RSZ_CPU_GATHER
	noc_async_read_barrier();
	#endif
	asm volatile("" ::: "memory");
	cb_push_back(cb_g0, 1);
	cb_push_back(cb_g1, 1);
	}
	if (wt > 0) {
	cb_wait_front(cb_out, 1);
	noc_async_write(get_read_ptr(cb_out), ao.get_noc_addr((y * NWT + wt - 1) * CT + ct), TB16);
	noc_async_writes_flushed();
	cb_pop_front(cb_out, 1);
	}
	}
	cb_pop_front(cb_v, NWI);
	}
	noc_async_write_barrier();
	}