File size: 1,628 Bytes
7a24154 | 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 | // SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, compute: each gathered 32-row
// block of CB 0 (row-major, Kt pages) tilized into Kt bf16 tiles of CB 16 (tilize_block: an exact bf16 datacopy
// through DEST, the LLK of the embedding op's own tilized output).
// CT args: [0] per-core RT-arg count (P3), [1] Kt, [2] TZ (0 = tilize_block, 1 = fast_tilize_block, 9 = no tilize:
// timing diagnostic only, the output is garbage). Per-core RT args: [nb, 0].
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/tilize.h"
constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t TZ = get_compile_time_arg_val(2);
constexpr uint32_t cb_rm = 0, cb_t = 16;
void kernel_main() {
const uint32_t nb = get_arg_val<uint32_t>(0);
if (nb == 0) {
return;
}
compute_kernel_hw_startup(cb_rm, cb_t);
if constexpr (TZ == 1) {
fast_tilize_init(cb_rm, Kt, cb_t);
} else if constexpr (TZ == 0) {
tilize_init(cb_rm, Kt, cb_t);
}
for (uint32_t b = 0; b < nb; ++b) {
cb_wait_front(cb_rm, Kt);
cb_reserve_back(cb_t, Kt);
if constexpr (TZ == 1) {
fast_tilize_block(cb_rm, Kt, cb_t);
} else if constexpr (TZ == 0) {
tilize_block(cb_rm, Kt, cb_t);
}
cb_push_back(cb_t, Kt);
cb_pop_front(cb_rm, Kt);
}
if constexpr (TZ == 1) {
fast_tilize_uninit(cb_rm, cb_t, Kt);
} else if constexpr (TZ == 0) {
tilize_uninit(cb_rm, cb_t);
}
}
|