File size: 1,241 Bytes
7a24154 | 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 | // SPDX-License-Identifier: Apache-2.0
// MK-H (OPT round 5, FRNET_POINT_MK=head; OPT_PLAN 3.1 / D19): FRHead as ONE row-streamed module megakernel, writer
// (RISCV_1): the NTO fp32 output tiles (logits) of each tile row r of this core's range to the output [.., N, 32 NTO].
// CT args: [0] per-core RT-arg count (P3), [1] NTO, then the TensorAccessorArgs of the output.
// Common RT args: [out_addr]. Per-core RT args: [r0, nr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NTO = get_compile_time_arg_val(1);
constexpr auto o_args = TensorAccessorArgs<2>();
constexpr uint32_t cb_out = 20;
constexpr uint32_t FTB = 4096;
void kernel_main() {
const uint32_t r0 = get_arg_val<uint32_t>(0);
const uint32_t nr = get_arg_val<uint32_t>(1);
const auto o = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), FTB);
for (uint32_t r = r0; r < r0 + nr; ++r) {
cb_wait_front(cb_out, NTO);
uint32_t rp = get_read_ptr(cb_out);
for (uint32_t j = 0; j < NTO; ++j) {
noc_async_write(rp, o.get_noc_addr(r * NTO + j), FTB);
rp += FTB;
}
noc_async_writes_flushed();
cb_pop_front(cb_out, NTO);
}
noc_async_write_barrier();
}
|