File size: 3,850 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | // SPDX-License-Identifier: Apache-2.0
// RSZ compute (frnet-p150 OPT round 2, tt/resize.py): per unit (ct, y)
// V_k = x0_k * A0 + x1_k * A1 (the H-direction lerp of input tile k, fp32 SFPU, k < NWI) -> CB 5 (fp32)
// out_wt = g0_wt * B0_wt + g1_wt * B1_wt (the W-direction lerp; g0 / g1: the V rows of each output pixel's
// two input columns, gathered by the writer) -> bf16 (typecast) -> CB 16
// fp32 DEST, every operand unpacked to DEST as fp32 (no TF32 rounding). CT args: [0] U, [1] NC, [2] GX, [3] NWI,
// [4] NWT (work split as the reader).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t U = 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 NWI = get_compile_time_arg_val(3);
constexpr uint32_t NWT = get_compile_time_arg_val(4);
constexpr uint32_t cb_x0 = 0, cb_x1 = 1, cb_a = 2, cb_b = 3, cb_v = 5, cb_g0 = 6, cb_g1 = 7, cb_out = 16;
// D[d] = cb[i] * w[j]; then D[0] += D[d] (d = 1 for the second product)
FORCE_INLINE void lerp(uint32_t ca, uint32_t ia, uint32_t cw0, uint32_t iw0, uint32_t cb, uint32_t ib, uint32_t cw1,
uint32_t iw1) {
copy_init(ca);
copy_tile(ca, ia, 0);
copy_init(cw0);
copy_tile(cw0, iw0, 1);
mul_binary_tile_init();
mul_binary_tile(0, 1, 0);
copy_init(cb);
copy_tile(cb, ib, 1);
copy_init(cw1);
copy_tile(cw1, iw1, 2);
mul_binary_tile_init();
mul_binary_tile(1, 2, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = U / NC, E = U % NC;
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
compute_kernel_hw_startup(cb_x0, cb_v);
cb_wait_front(cb_b, 2 * NWT);
for (uint32_t u = 0; u < n; ++u) {
cb_wait_front(cb_a, 2);
pack_reconfig_data_format(cb_v);
for (uint32_t k = 0; k < NWI; ++k) {
cb_wait_front(cb_x0, 1);
cb_wait_front(cb_x1, 1);
tile_regs_acquire();
lerp(cb_x0, 0, cb_a, 0, cb_x1, 0, cb_a, 1);
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(cb_x0, 1);
cb_pop_front(cb_x1, 1);
}
cb_pop_front(cb_a, 2);
pack_reconfig_data_format(cb_out);
for (uint32_t wt = 0; wt < NWT; ++wt) {
cb_wait_front(cb_g0, 1);
cb_wait_front(cb_g1, 1);
tile_regs_acquire();
#ifndef RSZ_SKIP_MATH
lerp(cb_g0, 0, cb_b, 2 * wt, cb_g1, 0, cb_b, 2 * wt + 1);
#else
copy_init(cb_g0);
copy_tile(cb_g0, 0, 0);
#endif
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
tile_regs_commit();
cb_reserve_back(cb_out, 1);
tile_regs_wait();
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
cb_pop_front(cb_g0, 1);
cb_pop_front(cb_g1, 1);
}
}
cb_pop_front(cb_b, 2 * NWT);
}
|