Hugging Face's logo

changh95
/
moge-2-p150

Depth Estimation
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-container
tenstorrent
ttnn
tt-metal
tt-nn
monocular-geometry
surface-normals
moge
tt-model-catalog
Model card Files
xet
 Community
1
moge-2-p150 / code /tt_moge /kernels /out_bands_writer.cpp
changh95's picture
changh95
Optimized build (2026-10-03): model call 19.0 ms, trace 17.7 ms
13b4736 verified
Raw History Blame Contribute Delete
2.02 kB
// SPDX-License-Identifier: Apache-2.0
// Output tail writer (moge-2 opt kernel): c_16 holds per block a 32 x (B*32) row-major block (row c =
// channel c of B*32 consecutive pixels, out_bands_compute.cpp). Pixel tile t0 of the block lies in band
// t0 / TPB; channel rows c < ROWS are written as one B*64-byte segment into page c (= channel row) of that
// band's ROW_MAJOR [ROWS, TPB*32] tensor (rows ROW0 + c with the ROW0 define). Blocks never straddle a band (TPB % B == 0, host check).
// Runtime args: band addresses [8], num_blocks, start_block. Compile args: B, ROWS, TPB, TensorAccessorArgs.
#include "api/dataflow/dataflow_api.h"
#include "api/dataflow/noc.h"
#include "api/dataflow/dataflow_buffer.h"
#include "api/tensor/noc_traits.h"
#ifndef ROW0
#define ROW0 0 // first channel row written (MOGE_OUT_GROUPS: one output tensor per head's channel rows)
#endif
void kernel_main() {
uint32_t addr[8];
for (uint32_t i = 0; i < 8; ++i) {
addr[i] = get_arg_val<uint32_t>(i);
}
const uint32_t nblocks = get_arg_val<uint32_t>(8);
const uint32_t start_block = get_arg_val<uint32_t>(9);
constexpr uint32_t B = get_compile_time_arg_val(0);
constexpr uint32_t ROWS = get_compile_time_arg_val(1);
constexpr uint32_t TPB = get_compile_time_arg_val(2);
constexpr auto d_args = TensorAccessorArgs<3>();
constexpr uint32_t cb_id = 16;
constexpr uint32_t seg = B * 32 * 2;
Noc noc;
DataflowBuffer dfb(cb_id);
for (uint32_t b = 0; b < nblocks; ++b) {
const uint32_t t0 = (start_block + b) * B;
const uint32_t band = t0 / TPB;
const uint32_t off = (t0 - band * TPB) * 32 * 2;
const auto s = TensorAccessor(d_args, addr[band]);
dfb.wait_front(B);
for (uint32_t c = 0; c < ROWS; ++c) {
noc.async_write(dfb, s, seg, {.offset_bytes = (ROW0 + c) * seg}, {.page_id = c, .offset_bytes = off});
}
noc.async_writes_flushed();
dfb.pop_front(B);
}
noc.async_write_barrier();
}