| // SPDX-License-Identifier: Apache-2.0 | |
| // Stem writer (sceneseg-p150 OPT_PLAN K-IN kernel form / MK0, tt/stem.py): the fp32 output tiles of this core's | |
| // range (work split as the reader) to the [1, 1, HOUT * WOUT, 32] TILE map, flushed per tile. | |
| // CT args: [0] T, [1] NC, [2] GX, then the output's TensorAccessorArgs. Common RT args: [out]. | |
| 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 auto out_args = TensorAccessorArgs<3>(); | |
| 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 = 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), 4096); | |
| for (uint32_t t = t0; t < t0 + n; ++t) { | |
| cb_wait_front(cb_out, 1); | |
| noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(t), 4096); | |
| noc_async_writes_flushed(); | |
| cb_pop_front(cb_out, 1); | |
| } | |
| noc_async_write_barrier(); | |
| } | |