File size: 3,851 Bytes
13b4736 | 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 | // SPDX-License-Identifier: Apache-2.0
// Head split / merge as a tile permutation with LOCAL reads (moge-2 opt kernel, round 7, ttnn.generic_op).
// The source is L1-interleaved: each core handles exactly the source pages that live in its own L1 bank and
// writes them straight from L1 to their destination page (no reader, no CB, no read round trip). Two RISCs
// per core (PART 0 / 1 on their own NoC) take alternate pages of the bank.
// MODE 0 (create heads): src [S, 3*H*2 tiles], page p = r*(6H) + c, c = which*2H + h*2 + j
// -> dst[which] [H, S, 64]: page (h*RT + r)*2 + j
// MODE 1 (concat heads): src [H, S, 64]: page p = (h*RT + r)*2 + j -> dst [S, H*64]: page r*2H + h*2 + j
// Runtime args: src_addr, dst0_addr, dst1_addr, dst2_addr, num_pages.
// Compile args: MODE, RT, H, PART, TensorAccessorArgs src, dst0, dst1, dst2.
#include "api/dataflow/dataflow_api.h"
#include "api/tensor/noc_traits.h"
void kernel_main() {
const uint32_t src_addr = get_arg_val<uint32_t>(0);
const uint32_t a0 = get_arg_val<uint32_t>(1);
const uint32_t a1 = get_arg_val<uint32_t>(2);
const uint32_t a2 = get_arg_val<uint32_t>(3);
const uint32_t num_pages = get_arg_val<uint32_t>(4);
constexpr uint32_t MODE = get_compile_time_arg_val(0);
constexpr uint32_t RT = get_compile_time_arg_val(1);
constexpr uint32_t H = get_compile_time_arg_val(2);
constexpr uint32_t PART = get_compile_time_arg_val(3);
constexpr uint32_t PAGE = 2048;
constexpr auto src_args = TensorAccessorArgs<4>();
constexpr auto d0_args = TensorAccessorArgs<src_args.next_compile_time_args_offset()>();
constexpr auto d1_args = TensorAccessorArgs<d0_args.next_compile_time_args_offset()>();
constexpr auto d2_args = TensorAccessorArgs<d1_args.next_compile_time_args_offset()>();
const auto s = TensorAccessor(src_args, src_addr);
const auto s0 = TensorAccessor(d0_args, a0);
const auto s1 = TensorAccessor(d1_args, a1);
const auto s2 = TensorAccessor(d2_args, a2);
const uint32_t mx = my_x[noc_index], my = my_y[noc_index];
auto is_local = [&](uint32_t p) -> bool {
const uint64_t a = s.get_noc_addr(p, 0, noc_index);
return NOC_UNICAST_ADDR_X(a) == mx && NOC_UNICAST_ADDR_Y(a) == my;
};
// find this core's bank: first local page b, bank stride nb = distance to the next local page
uint32_t b = 0xFFFFFFFF, nb = 0;
for (uint32_t p = 0; p < num_pages && p < 512; ++p) {
if (is_local(p)) {
if (b == 0xFFFFFFFF) {
b = p;
} else {
nb = p - b;
break;
}
}
}
if (b == 0xFFFFFFFF) {
return; // no source page in this core's bank
}
if (nb == 0) {
nb = num_pages; // only one page lives here
}
uint32_t k = 0;
for (uint32_t p = b; p < num_pages; p += nb, ++k) {
if ((k & 1) != PART) {
continue;
}
const uint32_t src_l1 = (uint32_t)(s.get_noc_addr(p, 0, noc_index) & NOC_LOCAL_ADDR_MASK);
if constexpr (MODE == 0) {
const uint32_t r = p / (6 * H), c = p % (6 * H);
const uint32_t which = c / (2 * H), hj = c % (2 * H);
const uint32_t dst = ((hj >> 1) * RT + r) * 2 + (hj & 1);
const uint64_t d = which == 0 ? s0.get_noc_addr(dst, 0, noc_index)
: which == 1 ? s1.get_noc_addr(dst, 0, noc_index)
: s2.get_noc_addr(dst, 0, noc_index);
noc_async_write(src_l1, d, PAGE);
} else {
const uint32_t j = p & 1, hr = p >> 1;
const uint32_t h = hr / RT, r = hr % RT;
noc_async_write(src_l1, s0.get_noc_addr(r * 2 * H + h * 2 + j, 0, noc_index), PAGE);
}
}
noc_async_write_barrier();
}
|