File size: 4,945 Bytes
58ba6df | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | // SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Face-granular (16x16) head split / concat for the ViT with UNPADDED heads (hamer/tt_kernels.py).
// HaMeR ViT-H heads are 80 wide = 5 faces (2.5 tiles). The qkv matmul produces the unpadded
// [192, 3*16*80] activation (120 tile columns = one per core on a 12x10 grid); SDPA needs head
// dims that are tile multiples, so every head is re-laid out to 96 = 3 tiles (6 faces, the 6th
// face zero). A bf16 TILE is 4 contiguous 512-byte faces: f0 rows 0-15 x cols 0-15, f1 rows 0-15
// x cols 16-31, f2 rows 16-31 x cols 0-15, f3 rows 16-31 x cols 16-31, so moving a 16-column
// slice of a tile row means moving the 2 faces (fr = 0, 1) of that face column.
//
// MODE_SPLIT src [R tile rows x W tile cols], W*2 = T*H*F face cols (F = faces per head)
// -> T outputs [H, R, DP] tiles (DP*2 >= F face cols per head, extra faces zero)
// out index i = ((t*H + h)*R + r)*DP + c ; out face (fr, fc): j = 2c + fc
// j < F: src face col g = (t*H + h)*F + j -> src page r*W + g/2, face fr*2 + g%2
// MODE_CONCAT src [H, R, DP] tiles -> 1 output [R x W] tiles, W*2 = H*F
// out index i = r*W + oc ; out face (fr, fc): g = 2*oc + fc, h = g / F, j = g % F
// -> src page (h*R + r)*DP + j/2, face fr*2 + j%2
//
// compile-time args: T, H, R, W, DP, F, CHUNK, then TensorAccessorArgs(src), (dst)
// runtime args: src_addr, start, n, dst0_addr .. dst{T-1}_addr
#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 W = get_compile_time_arg_val(3);
constexpr uint32_t DP = get_compile_time_arg_val(4);
constexpr uint32_t F = get_compile_time_arg_val(5);
constexpr uint32_t CHUNK = get_compile_time_arg_val(6);
constexpr auto src_args = TensorAccessorArgs<7>();
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 face = page / 4;
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 jt = 0; jt < m; ++jt) {
const uint32_t i = start + base + jt;
const uint32_t out_tile = scratch + jt * page;
for (uint32_t fc = 0; fc < 2; ++fc) {
#ifdef MODE_SPLIT
const uint32_t c = i % DP;
const uint32_t r = (i / DP) % R;
const uint32_t th = i / (DP * R); // t*H + h
const uint32_t j = 2 * c + fc;
if (j >= F) {
// zero pad faces (rows 0-15 and 16-31 of this face column)
for (uint32_t fr = 0; fr < 2; ++fr) {
volatile tt_l1_ptr uint32_t* p =
reinterpret_cast<volatile tt_l1_ptr uint32_t*>(out_tile + (fr * 2 + fc) * face);
for (uint32_t w = 0; w < face / 4; ++w) {
p[w] = 0;
}
}
continue;
}
const uint32_t g = th * F + j;
const uint32_t src_page = r * W + g / 2;
const uint32_t sfc = g % 2;
#else
const uint32_t oc = i % W;
const uint32_t r = i / W;
const uint32_t g = 2 * oc + fc;
const uint32_t h = g / F;
const uint32_t j = g % F;
const uint32_t src_page = (h * R + r) * DP + j / 2;
const uint32_t sfc = j % 2;
#endif
for (uint32_t fr = 0; fr < 2; ++fr) {
noc_async_read(
src.get_noc_addr(src_page, (fr * 2 + sfc) * face), out_tile + (fr * 2 + fc) * face, face);
}
}
}
noc_async_read_barrier();
for (uint32_t jt = 0; jt < m; ++jt) {
const uint32_t i = start + base + jt;
#ifdef MODE_SPLIT
const uint32_t t = i / (H * R * DP);
const uint32_t dst_page = i % (H * R * DP);
#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);
noc_async_write(scratch + jt * page, dst.get_noc_addr(dst_page), page);
}
noc_async_write_barrier();
}
}
|