Hugging Face's logo

changh95
/
meteor-p150

Robotics
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
meteor-p150 / code /tt_meteor /tt /kernels /terms_epi_writer.cpp
changh95's picture
changh95
tt-model push meteor-p150 (container)
8eeb550 verified
Raw History Blame Contribute Delete
2.52 kB
// SPDX-License-Identifier: Apache-2.0
// TERMS_EPI writer (meteor-p150 OPT_PLAN #2, tt/epi.py): the output tiles of this core's contiguous range (work split
// as the reader) to the interleaved outputs, flushed per tile, one barrier at the end: the fp32 main output (CB 16,
// HAS_MAIN) and, with SPLIT, the bf16 hi / lo terms (CBs 17 / 18).
// CT args: [0] T, [1] NC, [2] GX, [3] HAS_MAIN, [4] SPLIT, then the TensorAccessorArgs of the main, hi and lo outputs
// (repeats of the first present one when absent). Common RT args: [main_addr, hi_addr, lo_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t T = 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 HAS_MAIN = get_compile_time_arg_val(3);
constexpr uint32_t SPLIT = get_compile_time_arg_val(4);
constexpr auto out_args = TensorAccessorArgs<5>();
constexpr auto hi_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr auto lo_args = TensorAccessorArgs<hi_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_out = 16, cb_hi = 17, cb_lo = 18;
constexpr uint32_t TB32 = 4096, TB16 = 2048;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / NC, E = T % NC;
const uint32_t t0 = 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), TB32);
const auto hi = TensorAccessor(hi_args, get_common_arg_val<uint32_t>(1), TB16);
const auto lo = TensorAccessor(lo_args, get_common_arg_val<uint32_t>(2), TB16);
for (uint32_t t = t0; t < t0 + n; ++t) {
if constexpr (HAS_MAIN) {
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(t), TB32);
}
if constexpr (SPLIT) {
cb_wait_front(cb_hi, 1);
cb_wait_front(cb_lo, 1);
noc_async_write(get_read_ptr(cb_hi), hi.get_noc_addr(t), TB16);
noc_async_write(get_read_ptr(cb_lo), lo.get_noc_addr(t), TB16);
}
noc_async_writes_flushed();
if constexpr (HAS_MAIN) {
cb_pop_front(cb_out, 1);
}
if constexpr (SPLIT) {
cb_pop_front(cb_hi, 1);
cb_pop_front(cb_lo, 1);
}
}
noc_async_write_barrier();
}