Hugging Face's logo

changh95
/
sceneseg-p150

Image Segmentation
ONNX
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
sceneseg-p150 / code /tt_sceneseg /tt /kernels /s1_writer.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
1.57 kB
// SPDX-License-Identifier: Apache-2.0
// S1 writer (sceneseg-p150 OPT_PLAN S1, tt/subpixel.py): the output tiles of this core's units (u = c * R + r, work
// split as the reader) to tile r * CT + c of the 2H x 2W map (DRAM interleaved TILE), flushed per tile, one barrier
// at the end.
// CT args: [0] R, [1] CT, [2] NC, [3] GX, [4] OTB (output tile bytes), then the output's TensorAccessorArgs.
// Common RT args: [out].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t R = get_compile_time_arg_val(0);
constexpr uint32_t CT = 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 uint32_t OTB = get_compile_time_arg_val(4);
constexpr auto out_args = TensorAccessorArgs<5>();
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 U = R * CT;
constexpr uint32_t B = U / NC, E = U % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t u1 = u0 + B + (core < E ? 1 : 0);
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), OTB);
for (uint32_t u = u0; u < u1; ++u) {
const uint32_t c = u / R, r = u % R;
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(r * CT + c), OTB);
noc_async_writes_flushed();
cb_pop_front(cb_out, 1);
}
noc_async_write_barrier();
}