code/tt_moge/kernels/out_bands_compute.cpp

// SPDX-License-Identifier: Apache-2.0
// Output tail (moge-2 opt kernel, ttnn.generic_op): TILE [HW, 32] (pixels x channels) -> channel-major
// ROW_MAJOR rows. Per block of B pixel tiles: WH-transpose each tile into DEST, then pack-untilize the B
// DEST tiles as one 32 x (B*32) row-major block into c_16 (row c = channel c, B*32 consecutive pixels).
// Replaces ttnn.transpose + per-band slice + untilize_with_unpadding. Runtime arg 0: blocks for this core.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/transpose.h"
#include "api/compute/pack_untilize.h"
#include "api/dataflow/dataflow_buffer.h"

void kernel_main() {
    const uint32_t nblocks = get_arg_val<uint32_t>(0);
    constexpr uint32_t B = get_compile_time_arg_val(0);
    constexpr auto cb_in = tt::CBIndex::c_0;
    constexpr auto cb_out = tt::CBIndex::c_16;
    DataflowBuffer dfb_in(cb_in);
    DataflowBuffer dfb_out(cb_out);
    unary_op_init_common(cb_in, cb_out);
    transpose_init(cb_in);
    pack_untilize_dest_init<B>(cb_out);
    for (uint32_t b = 0; b < nblocks; ++b) {
        dfb_in.wait_front(B);
        tile_regs_acquire();
        for (uint32_t j = 0; j < B; ++j) {
            transpose_tile(cb_in, j, j);
        }
        tile_regs_commit();
        dfb_in.pop_front(B);
        dfb_out.reserve_back(B);
        tile_regs_wait();
        pack_untilize_dest<B>(cb_out);
        tile_regs_release();
        dfb_out.push_back(B);
    }
    pack_untilize_uninit(cb_out);
}