File size: 1,446 Bytes
e91e989 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | // 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();
}
|