// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// 2x2 / stride-2 max pool of an L1 height-sharded ROW_MAJOR activation holding ONE image row per
// core (W pixels x C channels), output core c = half h = c % 2 of output row r = c / 2. RISC PROC
// (0 / 1) gathers input row 2r + PROC (on core c - h + PROC, read over the NoC; local for one of
// them) pixels [h * W/2, (h + 1) * W/2): even pixels into CB_E, odd pixels into CB_O, 2 KB pages
// (pseudo tiles of 1024 bf16: the elementwise max does not care about the element order).
// RT args: src_noc_x, src_noc_y, x_addr, px0 (first input pixel), CT args: cb_e, cb_o, NPX (output
// pixels per core), STICK (bytes per pixel)
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    const uint32_t nx = get_arg_val<uint32_t>(0);
    const uint32_t ny = get_arg_val<uint32_t>(1);
    const uint32_t x_addr = get_arg_val<uint32_t>(2);
    const uint32_t px0 = get_arg_val<uint32_t>(3);
    constexpr uint32_t cb_e = get_compile_time_arg_val(0);
    constexpr uint32_t cb_o = get_compile_time_arg_val(1);
    constexpr uint32_t NPX = get_compile_time_arg_val(2);
    constexpr uint32_t STICK = get_compile_time_arg_val(3);
    constexpr uint32_t PER_PAGE = 2048 / STICK;
    constexpr uint32_t PAGES = NPX / PER_PAGE;
    static_assert(PAGES * PER_PAGE == NPX, "whole pages per core");
    uint64_t src = get_noc_addr(nx, ny, x_addr + px0 * STICK);
    for (uint32_t t = 0; t < PAGES; ++t) {
        cb_reserve_back(cb_e, 1);
        cb_reserve_back(cb_o, 1);
        uint32_t we = get_write_ptr(cb_e), wo = get_write_ptr(cb_o);
        for (uint32_t k = 0; k < PER_PAGE; ++k) {
            noc_async_read(src, we, STICK);
            noc_async_read(src + STICK, wo, STICK);
            src += 2 * STICK;
            we += STICK;
            wo += STICK;
        }
        noc_async_read_barrier();
        cb_push_back(cb_e, 1);
        cb_push_back(cb_o, 1);
    }
}