// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Cell-tile 2x2 pool, vertical half (models/tt/conv_cell.py PoolCell), data movement. Input: the
// half-pooled TILE map [cells, QTI tiles per tile row] of this core (CB_A bound to it), 80 cells per
// image row. Pooled row Y (Y % 2 == PROC builds its operands) reads the even image row (cells
// 160Y + 32k .., k = 0..2) and the odd row below (+80 cells): CB_B gets, per k, the QTI tiles of the
// odd row = rows 16..31 of tile row 5Y + k + 2 (faces 2, 3 -> 0, 1) and rows 0..15 of tile row
// 5Y + k + 3 (faces 0, 1 -> 2, 3), the latter only when it exists (k = 2 of the last pooled row
// reads past the shard: those output rows are not used).
// PROC 1 also moves each result block of CB_U (compute order: Y, k) into the output (CB_O bound):
//   MODE 0 (rm):   untilized [32 cells, QTI * 32] block, valid rows (32, or 16 for k = 2) at cell
//                  Y*80 + 32k of the ROW_MAJOR output (QTI * 64 B per cell)
//   MODE 1 (cell): QTI pooled tiles, 16-row halves copied to pooled cell Y*80 + 32k (+16) of the
//                  TILE output (QTI tiles per tile row)
// CT args: PROC, cb_a, cb_b, cb_u, cb_o, QTI, TRS, PROW, MODE
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr uint32_t PROC = get_compile_time_arg_val(0);
    constexpr uint32_t cb_a = get_compile_time_arg_val(1);
    constexpr uint32_t cb_b = get_compile_time_arg_val(2);
    constexpr uint32_t cb_u = get_compile_time_arg_val(3);
    constexpr uint32_t cb_o = get_compile_time_arg_val(4);
    constexpr uint32_t QTI = get_compile_time_arg_val(5);
    constexpr uint32_t TRS = get_compile_time_arg_val(6);
    constexpr uint32_t PROW = get_compile_time_arg_val(7);
    constexpr uint32_t MODE = get_compile_time_arg_val(8);
    constexpr uint32_t TILE = 2048;
    if constexpr (PROC == 0) {
        cb_reserve_back(cb_a, TRS * QTI);
        cb_push_back(cb_a, TRS * QTI);
    }
    const uint64_t a_noc = get_noc_addr(my_x[noc_index], my_y[noc_index], get_write_ptr(cb_a));
    for (uint32_t Y = PROC; Y < PROW; Y += 2) {
        cb_reserve_back(cb_b, 3 * QTI);
        const uint32_t b_l1 = get_write_ptr(cb_b);
        for (uint32_t k = 0; k < 3; ++k) {
            const uint32_t t0 = 5 * Y + k + 2;
            for (uint32_t q = 0; q < QTI; ++q) {
                const uint32_t dst = b_l1 + (k * QTI + q) * TILE;
                noc_async_read(a_noc + (t0 * QTI + q) * TILE + 1024, dst, 1024);
                if (t0 + 1 < TRS) {
                    noc_async_read(a_noc + ((t0 + 1) * QTI + q) * TILE, dst + 1024, 1024);
                }
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_b, 3 * QTI);
    }
    if constexpr (PROC == 1) {
        const uint32_t o_l1 = get_write_ptr(cb_o);
        const uint64_t me = get_noc_addr(my_x[noc_index], my_y[noc_index], 0);
        for (uint32_t Y = 0; Y < PROW; ++Y) {
            for (uint32_t k = 0; k < 3; ++k) {
                cb_wait_front(cb_u, QTI);
                const uint32_t u = get_read_ptr(cb_u);
                const uint32_t rows = k == 2 ? 16 : 32;
                const uint32_t c0 = Y * 80 + 32 * k;
                if constexpr (MODE == 0) {
                    noc_async_read(me | u, o_l1 + c0 * QTI * 64, rows * QTI * 64);
                } else {
                    for (uint32_t hs = 0; hs < rows / 16; ++hs) {
                        const uint32_t dc = c0 + 16 * hs;
                        for (uint32_t q = 0; q < QTI; ++q) {
                            noc_async_read(me | (u + q * TILE + hs * 1024),
                                           o_l1 + ((dc >> 5) * QTI + q) * TILE + ((dc >> 4) & 1) * 1024, 1024);
                        }
                    }
                }
                noc_async_read_barrier();
                cb_pop_front(cb_u, QTI);
            }
        }
    }
}