File size: 2,880 Bytes
e91e989 | 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 | // SPDX-License-Identifier: Apache-2.0
// BRES compute (meteor-p150 OPT2 round 2, tt/bresize.py): per output row block b and channel tile ct
// D = sum_e Wt[e] @ x[e, ct] (FPU matmuls into one fp32 DEST tile; bf16 x bf16 products, BRES_FID fidelity)
// out = bf16_rne(D) (SFPU typecast, BRES_OUT_BF16) or the fp32 D (packed to an fp32 CB)
// CB 1 holds the block's KM weight tiles (in0), CB 0 its KM x CT input tiles (in1, entry-major).
// NX = 2: D also sums Wt[e] @ x_lo[e, ct] (an fp32 map as its two bf16 terms; tiles KM x CT on in CB 0).
// CT args: [0] NB, [1] NC, [2] GX, [3] KM, [4] CT, [5] NX. Define BRES_OUT_BF16 (1: bf16 output).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#ifndef BRES_OUT_BF16
#define BRES_OUT_BF16 1
#endif
using namespace ckernel;
constexpr uint32_t NB = 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 KM = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr uint32_t NX = get_compile_time_arg_val(5);
constexpr uint32_t cb_x = 0, cb_w = 1, cb_out = 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 = NB / NC, E = NB % NC;
const uint32_t b0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
compute_kernel_hw_startup<SrcOrder::Reverse>(cb_w, cb_x, cb_out);
for (uint32_t b = b0; b < b0 + n; ++b) {
cb_wait_front(cb_w, KM);
cb_wait_front(cb_x, NX * KM * CT);
for (uint32_t ct = 0; ct < CT; ++ct) {
tile_regs_acquire();
matmul_init(cb_w, cb_x);
for (uint32_t e = 0; e < KM; ++e) {
matmul_tiles(cb_w, cb_x, e, e * CT + ct, 0);
}
if constexpr (NX > 1) {
for (uint32_t e = 0; e < KM; ++e) {
matmul_tiles(cb_w, cb_x, e, (KM + e) * CT + ct, 0);
}
}
#if BRES_OUT_BF16
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
#endif
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_x, NX * KM * CT);
cb_pop_front(cb_w, KM);
}
}
|