code/tt_moge/kernels/heads_writer.cpp

// 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();
}