File size: 1,996 Bytes
c699c4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 | // 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);
}
}
|