code/tt_sam3/tt/kernels/tile_stream_dm.cpp

// SPDX-License-Identifier: Apache-2.0
// Interleaved tile streaming for element-wise custom ops: IS_WRITER=0 reads tiles [start, start+n) of the source
// into cb (BLK per batch, one barrier per batch); IS_WRITER=1 writes cb to the same tile ids of the destination.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

constexpr uint32_t cb = get_compile_time_arg_val(0);
constexpr uint32_t BLK = get_compile_time_arg_val(1);
constexpr uint32_t IS_WRITER = get_compile_time_arg_val(2);
constexpr uint32_t TILE_BYTES = get_compile_time_arg_val(3);
constexpr auto t_args = TensorAccessorArgs<4>();

void kernel_main() {
    const uint32_t addr = get_arg_val<uint32_t>(0);
    const uint32_t n = get_arg_val<uint32_t>(1);
    const uint32_t start = get_arg_val<uint32_t>(2);
    const auto acc = TensorAccessor(t_args, addr, TILE_BYTES);
    for (uint32_t t = 0; t < n; t += BLK) {
        const uint32_t b = (n - t) < BLK ? (n - t) : BLK;
        if constexpr (IS_WRITER == 0) {
            cb_reserve_back(cb, b);
            uint32_t l1 = get_write_ptr(cb);
            for (uint32_t j = 0; j < b; ++j) {
                noc_async_read(acc.get_noc_addr(start + t + j), l1 + j * TILE_BYTES, TILE_BYTES);
            }
            noc_async_read_barrier();
            cb_push_back(cb, b);
        } else {
            cb_wait_front(cb, b);
            uint32_t l1 = get_read_ptr(cb);
            for (uint32_t j = 0; j < b; ++j) {
                noc_async_write(l1 + j * TILE_BYTES, acc.get_noc_addr(start + t + j), TILE_BYTES);
            }
            noc_async_writes_flushed();
            cb_pop_front(cb, b);
        }
    }
    if constexpr (IS_WRITER == 1) {
        noc_async_write_barrier();
    }
}