File size: 1,738 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 44 45 46 | // SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
// pool_rows: out page t = max(max(E0, O0), max(E1, O1)) page t (SFPU binary max on bf16: exact),
// packed into the output shard (CB_OUT bound to it) in the same element order (pseudo tiles).
#include <cstdint>
#include "api/compute/compute_kernel_api.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/binary_max_min.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
void kernel_main() {
constexpr uint32_t cb_e0 = get_compile_time_arg_val(0);
constexpr uint32_t cb_o0 = get_compile_time_arg_val(1);
constexpr uint32_t cb_e1 = get_compile_time_arg_val(2);
constexpr uint32_t cb_o1 = get_compile_time_arg_val(3);
constexpr uint32_t cb_out = get_compile_time_arg_val(4);
constexpr uint32_t PAGES = get_compile_time_arg_val(5);
unary_op_init_common(cb_e0, cb_out);
copy_tile_to_dst_init_short(cb_e0);
binary_max_tile_init();
cb_reserve_back(cb_out, PAGES);
for (uint32_t t = 0; t < PAGES; ++t) {
cb_wait_front(cb_e0, 1);
cb_wait_front(cb_o0, 1);
cb_wait_front(cb_e1, 1);
cb_wait_front(cb_o1, 1);
tile_regs_acquire();
copy_tile(cb_e0, 0, 0);
copy_tile(cb_o0, 0, 1);
copy_tile(cb_e1, 0, 2);
copy_tile(cb_o1, 0, 3);
binary_max_tile(0, 1, 0);
binary_max_tile(2, 3, 2);
binary_max_tile(0, 2, 0);
tile_regs_commit();
tile_regs_wait();
pack_tile<true>(0, cb_out, t);
tile_regs_release();
cb_pop_front(cb_e0, 1);
cb_pop_front(cb_o0, 1);
cb_pop_front(cb_e1, 1);
cb_pop_front(cb_o1, 1);
}
cb_push_back(cb_out, PAGES);
}
|