File size: 3,356 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 | // SPDX-License-Identifier: Apache-2.0
// Tile-permutation writer for the attention head split / merge (moge-2 opt kernel, ttnn.generic_op).
// Pages arrive in c_0 in source order (reader_batched.cpp); each is written to its destination page:
// MODE 0 (create heads): src [S, 3*H*2 tiles] row-major tiles, page p = r*(6H) + c, c = which*2H + h*2 + j
// -> dst[which] [H, S, 64]: page (h*RT + r)*2 + j (which = q, k, v)
// 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: dst0_addr, dst1_addr, dst2_addr, num_pages, start_id.
// Compile args: BLOCK, MODE, RT (tile rows of S), H (heads), TensorAccessorArgs dst0 [, dst1, dst2].
#include "api/dataflow/dataflow_api.h"
#include "api/dataflow/noc.h"
#include "api/dataflow/dataflow_buffer.h"
#include "api/tensor/noc_traits.h"
void kernel_main() {
const uint32_t a0 = get_arg_val<uint32_t>(0);
const uint32_t a1 = get_arg_val<uint32_t>(1);
const uint32_t a2 = get_arg_val<uint32_t>(2);
const uint32_t num_pages = get_arg_val<uint32_t>(3);
const uint32_t start_id = get_arg_val<uint32_t>(4);
constexpr uint32_t BLOCK = get_compile_time_arg_val(0);
constexpr uint32_t MODE = get_compile_time_arg_val(1);
constexpr uint32_t RT = get_compile_time_arg_val(2);
constexpr uint32_t H = get_compile_time_arg_val(3);
constexpr auto d0_args = TensorAccessorArgs<4>();
constexpr auto d1_args = TensorAccessorArgs<d0_args.next_compile_time_args_offset()>();
constexpr auto d2_args = TensorAccessorArgs<d1_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_id = 0;
const uint32_t page_bytes = get_local_cb_interface(cb_id).fifo_page_size;
const auto s0 = TensorAccessor(d0_args, a0);
const auto s1 = TensorAccessor(d1_args, a1);
const auto s2 = TensorAccessor(d2_args, a2);
Noc noc;
DataflowBuffer dfb(cb_id);
uint32_t i = start_id;
const uint32_t end_id = start_id + num_pages;
while (i < end_id) {
const uint32_t n = (end_id - i) < BLOCK ? (end_id - i) : BLOCK;
dfb.wait_front(n);
for (uint32_t t = 0; t < n; ++t) {
const uint32_t p = i + t;
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);
if (which == 0) {
noc.async_write(dfb, s0, page_bytes, {.offset_bytes = t * page_bytes}, {.page_id = dst});
} else if (which == 1) {
noc.async_write(dfb, s1, page_bytes, {.offset_bytes = t * page_bytes}, {.page_id = dst});
} else {
noc.async_write(dfb, s2, page_bytes, {.offset_bytes = t * page_bytes}, {.page_id = dst});
}
} else {
const uint32_t j = p & 1, hr = p >> 1;
const uint32_t h = hr / RT, r = hr % RT;
const uint32_t dst = r * 2 * H + h * 2 + j;
noc.async_write(dfb, s0, page_bytes, {.offset_bytes = t * page_bytes}, {.page_id = dst});
}
}
noc.async_writes_flushed();
dfb.pop_front(n);
i += n;
}
noc.async_write_barrier();
}
|