| // SPDX-License-Identifier: Apache-2.0 | |
| // SE-K writer (sceneseg-p150 OPT_PLAN SE-K, tt/sek.py): this core's fp32 gate tiles (work split as the reader) to | |
| // the [1, 1, 1, C] TILE gate (row tile 0, column tile n), flushed per tile. | |
| // CT args: [0] NT2, [1] NC, [2] GX, then the output's TensorAccessorArgs. Common RT args: [out]. | |
| constexpr uint32_t NT2 = 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 = NT2 / NC, E = NT2 % NC; | |
| const uint32_t n0 = core * B + (core < E ? core : E); | |
| const uint32_t nn = B + (core < E ? 1 : 0); | |
| const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), 4096); | |
| for (uint32_t j = 0; j < nn; ++j) { | |
| cb_wait_front(cb_out, 1); | |
| noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(n0 + j), 4096); | |
| noc_async_writes_flushed(); | |
| cb_pop_front(cb_out, 1); | |
| } | |
| noc_async_write_barrier(); | |
| } | |