meteor-p150 / code /tt_meteor /tt /kernels /bres_writer.cpp

tt-model push meteor-p150 (container)

e91e989 verified 5 days ago

1.68 kB

// SPDX-License-Identifier: Apache-2.0
// BRES writer (meteor-p150 OPT2 round 2, tt/bresize.py): the CT output tiles of each output row block b of this
// core's range (CB 16, ct order) -> tiles b * CT + ct of the output map (flat [1, 1, R, C] or 4-D [N, H, W, C] TILE,
// W padded to tiles: its row blocks never cross an image row). Work split as the reader.
// CT args: [0] NB, [1] NC, [2] GX, [3] CT, [4] TB (output tile bytes), then the TensorAccessorArgs of the output.
// Common RT args: [out_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NB = 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 CT = get_compile_time_arg_val(3);
constexpr uint32_t TB = get_compile_time_arg_val(4);
constexpr auto o_args = TensorAccessorArgs<5>();
constexpr uint32_t cb_out = 16;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NB / NC, E = NB % NC;
    const uint32_t b0 = 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), TB);
    for (uint32_t b = b0; b < b0 + n; ++b) {
        for (uint32_t ct = 0; ct < CT; ++ct) {
            cb_wait_front(cb_out, 1);
            noc_async_write(get_read_ptr(cb_out), ao.get_noc_addr(b * CT + ct), TB);
            noc_async_writes_flushed();
            cb_pop_front(cb_out, 1);
        }
    }
    noc_async_write_barrier();
}

	// SPDX-License-Identifier: Apache-2.0
	// BRES writer (meteor-p150 OPT2 round 2, tt/bresize.py): the CT output tiles of each output row block b of this
	// core's range (CB 16, ct order) -> tiles b * CT + ct of the output map (flat [1, 1, R, C] or 4-D [N, H, W, C] TILE,
	// W padded to tiles: its row blocks never cross an image row). Work split as the reader.
	// CT args: [0] NB, [1] NC, [2] GX, [3] CT, [4] TB (output tile bytes), then the TensorAccessorArgs of the output.
	// Common RT args: [out_addr].
	#include <cstdint>

	#include "api/dataflow/dataflow_api.h"

	constexpr uint32_t NB = 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 CT = get_compile_time_arg_val(3);
	constexpr uint32_t TB = get_compile_time_arg_val(4);
	constexpr auto o_args = TensorAccessorArgs<5>();
	constexpr uint32_t cb_out = 16;

	void kernel_main() {
	const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
	if (core >= NC) {
	return;
	}
	constexpr uint32_t B = NB / NC, E = NB % NC;
	const uint32_t b0 = 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), TB);
	for (uint32_t b = b0; b < b0 + n; ++b) {
	for (uint32_t ct = 0; ct < CT; ++ct) {
	cb_wait_front(cb_out, 1);
	noc_async_write(get_read_ptr(cb_out), ao.get_noc_addr(b * CT + ct), TB);
	noc_async_writes_flushed();
	cb_pop_front(cb_out, 1);
	}
	}
	noc_async_write_barrier();
	}