File size: 2,160 Bytes
8eeb550 | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | // SPDX-License-Identifier: Apache-2.0
// STEM_POOL compute (meteor-p150 OPT round 3 R3.2, tt/pool.py): the vertical half of the 3x3 / s2 max pool, per tile
// of a unit's three input rows (CBs 0 / 1 / 2, WT tiles each): v = max(max(r0, r1), r2) with binary_max_tile, the
// fp32 SFPU max of ttnn.maximum (binary_ng), operands unpacked to DEST as fp32 (UnpackToDestFp32 on CBs 0..2), packed
// fp32 to CB 16. The writer does the horizontal half. Max is exact; the order of the operands does not matter.
// CT args: [0] NU, [1] NC, [2] GX, [3] WT.
#include <cstdint>
#include "api/compute/binary_max_min.h"
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t NU = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t WT = get_compile_time_arg_val(3);
constexpr uint32_t cb_v = 16;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NU / NC, E = NU % NC;
const uint32_t nu = B + (core < E ? 1 : 0);
if (nu == 0) {
return;
}
compute_kernel_hw_startup(0, cb_v);
for (uint32_t u = 0; u < nu; ++u) {
cb_wait_front(0, WT);
cb_wait_front(1, WT);
cb_wait_front(2, WT);
for (uint32_t tx = 0; tx < WT; ++tx) {
tile_regs_acquire();
copy_init(0);
copy_tile(0, tx, 0);
copy_init(1);
copy_tile(1, tx, 1);
binary_max_tile_init();
binary_max_tile(0, 1, 0);
copy_init(2);
copy_tile(2, tx, 1);
binary_max_tile_init();
binary_max_tile(0, 1, 0);
tile_regs_commit();
cb_reserve_back(cb_v, 1);
tile_regs_wait();
pack_tile(0, cb_v);
tile_regs_release();
cb_push_back(cb_v, 1);
}
cb_pop_front(0, WT);
cb_pop_front(1, WT);
cb_pop_front(2, WT);
}
}
|