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

tt-model push meteor-p150 (container)

e91e989 verified 5 days ago

1.45 kB

// SPDX-License-Identifier: Apache-2.0
// TGATE writer (meteor-p150 OPT2 round 3 R3.2, tt/tgate.py): the 3 bf16 output tiles of each 32-row block (CB 16) to
// the interleaved [R, 96] bf16 TILE output (tiles b * 3 + k); work split as the reader.
// CT args: [0] NB, [1] NC, [2] GX, 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 auto out_args = TensorAccessorArgs<3>();
constexpr uint32_t TB16 = 2048;
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);
    const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), TB16);
    for (uint32_t b = b0; b < b0 + n; ++b) {
        cb_wait_front(cb_out, 3);
        const uint32_t rp = get_read_ptr(cb_out);
        for (uint32_t k = 0; k < 3; ++k) {
            noc_async_write(rp + k * TB16, out.get_noc_addr(b * 3 + k), TB16);
        }
        noc_async_writes_flushed();
        cb_pop_front(cb_out, 3);
    }
    noc_async_write_barrier();
}

	// SPDX-License-Identifier: Apache-2.0
	// TGATE writer (meteor-p150 OPT2 round 3 R3.2, tt/tgate.py): the 3 bf16 output tiles of each 32-row block (CB 16) to
	// the interleaved [R, 96] bf16 TILE output (tiles b * 3 + k); work split as the reader.
	// CT args: [0] NB, [1] NC, [2] GX, 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 auto out_args = TensorAccessorArgs<3>();
	constexpr uint32_t TB16 = 2048;
	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);
	const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), TB16);
	for (uint32_t b = b0; b < b0 + n; ++b) {
	cb_wait_front(cb_out, 3);
	const uint32_t rp = get_read_ptr(cb_out);
	for (uint32_t k = 0; k < 3; ++k) {
	noc_async_write(rp + k * TB16, out.get_noc_addr(b * 3 + k), TB16);
	}
	noc_async_writes_flushed();
	cb_pop_front(cb_out, 3);
	}
	noc_async_write_barrier();
	}