File size: 1,540 Bytes
e0adbab | 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 37 | // SPDX-License-Identifier: Apache-2.0
// K11 depthwise writer (sceneseg-p150 OPT_PLAN K11, tt/dw.py): the output tiles of this core's items (work split and
// order as dw_reader.cpp: item w = ct * PT + pt) to the float32 TILE output [1, 1, Hout*Wout, C], tile pt * CT + ct.
// CT args: [0] CT, [1] PT, [2] NC, [3] GX, then the output's TensorAccessorArgs. Common RT args: [out_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t CT = get_compile_time_arg_val(0);
constexpr uint32_t PT = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr auto out_args = TensorAccessorArgs<4>();
constexpr uint32_t NW = CT * PT;
constexpr uint32_t TB = 4096;
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 = NW / NC, E = NW % NC;
const uint32_t w0 = core * B + (core < E ? core : E);
const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), TB);
for (uint32_t w = w0; w < w1; ++w) {
const uint32_t ct = w / PT;
const uint32_t pt = w - ct * PT;
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(pt * CT + ct), TB);
noc_async_writes_flushed();
cb_pop_front(cb_out, 1);
}
noc_async_write_barrier();
}
|