Hugging Face's logo

changh95
/
hamer-p150

Keypoint Detection
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-container
tenstorrent
ttnn
tt-metal
tt-nn
hand-pose
mano
3d-mesh
hamer
tt-model-catalog
Model card Files
xet
 Community
1
hamer-p150 / code /hamer /kernels /tile_permute.cpp
changh95's picture
changh95
Optimized build (2026-10-03): model call 3.49 ms, device 3.36 ms
58ba6df verified
Raw History Blame Contribute Delete
3.35 kB
// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Tile-granular head split / head concat for the HaMeR ViT and MANO head (hamer/tt_kernels.py).
// One data-movement kernel per core: it owns a contiguous range [start, start+n) of OUTPUT tiles,
// gathers them from the interleaved source with batched async reads into a local scratch CB,
// then scatters them to the interleaved destination(s) with batched async writes.
//
// MODE_SPLIT src [R rows x T*H*D tile cols] -> T outputs (T <= 12), each [H, R, D] tiles
// out index i = ((t*H + h)*R + r)*D + c
// src page = r*(T*H*D) + t*H*D + h*D + c ; dst page (tensor t) = (h*R + r)*D + c
// MODE_CONCAT src [H, R, D] tiles -> 1 output [R x H*D] tiles
// out index j = r*(H*D) + h*D + c ; src page = (h*R + r)*D + c
//
// compile-time args: T, H, R, D, CHUNK (tiles per batch), then TensorAccessorArgs(src), (dst)
// runtime args: src_addr, start, n, dst0_addr .. dst{T-1}_addr
// Source / destinations may be interleaved or sharded (TensorAccessor resolves page -> bank/core).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
void kernel_main() {
const uint32_t src_addr = get_arg_val<uint32_t>(0);
const uint32_t start = get_arg_val<uint32_t>(1);
const uint32_t n = get_arg_val<uint32_t>(2);
constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t H = get_compile_time_arg_val(1);
constexpr uint32_t R = get_compile_time_arg_val(2);
constexpr uint32_t D = get_compile_time_arg_val(3);
constexpr uint32_t CHUNK = get_compile_time_arg_val(4);
constexpr auto src_args = TensorAccessorArgs<5>();
constexpr auto dst_args = TensorAccessorArgs<src_args.next_compile_time_args_offset()>();
constexpr uint32_t cb = 0;
const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
const uint32_t scratch = get_write_ptr(cb);
const auto src = TensorAccessor(src_args, src_addr, page);
for (uint32_t base = 0; base < n; base += CHUNK) {
const uint32_t m = (n - base) < CHUNK ? (n - base) : CHUNK;
for (uint32_t j = 0; j < m; ++j) {
const uint32_t i = start + base + j;
const uint32_t c = i % D;
#ifdef MODE_SPLIT
const uint32_t r = (i / D) % R;
const uint32_t th = i / (D * R); // t*H + h
const uint32_t src_page = r * (T * H * D) + th * D + c;
#else
const uint32_t hc = i % (H * D); // h*D + c
const uint32_t r = i / (H * D);
const uint32_t h = hc / D;
const uint32_t src_page = (h * R + r) * D + c;
#endif
noc_async_read(src.get_noc_addr(src_page), scratch + j * page, page);
}
noc_async_read_barrier();
for (uint32_t j = 0; j < m; ++j) {
const uint32_t i = start + base + j;
#ifdef MODE_SPLIT
const uint32_t t = i / (H * R * D);
const uint32_t dst_page = i % (H * R * D);
#else
const uint32_t t = 0;
const uint32_t dst_page = i;
#endif
const auto dst = TensorAccessor(dst_args, get_arg_val<uint32_t>(3 + t), page);
const uint64_t a = dst.get_noc_addr(dst_page);
noc_async_write(scratch + j * page, a, page);
}
noc_async_write_barrier();
}
}