File size: 6,190 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 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | // SPDX-License-Identifier: Apache-2.0
// TFUSE compute (meteor-p150 OPT2 MK-TF, tt/tgate.py fuse_mm): the temporal fuse's merged 1x1 conv and its gate glue
// in one pass per 32-row block (the TGATE kernel, tgate_compute.cpp, with the conv as its first phase):
// phase 1, the conv: g_j = sum_k x_k @ W[k, j] (j = 0..3, k = 0..2; FPU matmuls, bf16 x (SrcB) times fp32 W
// (SrcA, TF32) into fp32 DEST, HiFi4: the stock fp32-weight matmul's operands), packed fp32:
// j = 0..2 -> CB 0 (t), j = 3 -> CB 1 (the logit differences without their bias)
// phase 2, the glue: E = exp(d + gb_hi + gb_lo) (SFPU, fp32; the gate bias as two bf16 terms, CB 7)
// S = E @ ONES, G = SCALE / (1 + S), out_k = bf16_rne(relu(t_k * G + b_k)) (as TGATE)
// SrcA stays fp32 throughout (W, every copy source, the ones tile); SrcB is bf16 for phase 1 (x) and fp32 for the
// E @ ONES matmul: reconfigured twice per block. Every CB reaches DEST through SrcA / SrcB (TF32, as TGATE).
// CBs: 8 x (3 tiles bf16), 6 W (12 fp32), 0 t (3, loop-back), 1 d (loop-back), 2 ones, 3 bias (3), 7 gate bias (2),
// 4 E (loop-back), 5 G (loop-back), 16 out (bf16).
// CT args: [0] NB, [1] NC, [2] GX, [3] SCALE (fp32 bits).
#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_binary_sfpu.h"
#include "api/compute/eltwise_unary/binop_with_scalar.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/exp.h"
#include "api/compute/eltwise_unary/recip.h"
#include "api/compute/eltwise_unary/relu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
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 SCALE = get_compile_time_arg_val(3);
constexpr uint32_t ONE = 0x3F800000u;
constexpr uint32_t cb_t = 0, cb_l = 1, cb_one = 2, cb_bias = 3, cb_e = 4, cb_g = 5, cb_w = 6, cb_gb = 7, cb_x = 8,
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;
}
// matmul operand order: in0 (x) -> SrcB, in1 (W) -> SrcA
compute_kernel_hw_startup<SrcOrder::Reverse>(cb_x, cb_w, cb_t);
cb_wait_front(cb_w, 12);
cb_wait_front(cb_one, 1);
cb_wait_front(cb_bias, 3);
cb_wait_front(cb_gb, 2);
for (uint32_t b = b0; b < b0 + n; ++b) {
// phase 1: g_j = x @ W[:, j]
reconfig_data_format_srcb(cb_x);
cb_wait_front(cb_x, 3);
for (uint32_t j = 0; j < 4; ++j) {
const uint32_t dst = j < 3 ? cb_t : cb_l;
tile_regs_acquire();
matmul_init(cb_x, cb_w);
for (uint32_t k = 0; k < 3; ++k) {
matmul_tiles(cb_x, cb_w, k, k * 4 + j, 0);
}
tile_regs_commit();
cb_reserve_back(dst, 1);
tile_regs_wait();
pack_reconfig_data_format(dst);
pack_tile(0, dst);
tile_regs_release();
cb_push_back(dst, 1);
}
cb_pop_front(cb_x, 3);
reconfig_data_format_srcb(cb_e);
// E = exp(d + gb_hi + gb_lo)
cb_wait_front(cb_l, 1);
tile_regs_acquire();
copy_init(cb_l);
copy_tile(cb_l, 0, 0);
copy_init(cb_gb);
copy_tile(cb_gb, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
copy_init(cb_gb);
copy_tile(cb_gb, 1, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
exp_tile_init();
exp_tile(0);
tile_regs_commit();
cb_reserve_back(cb_e, 1);
tile_regs_wait();
pack_reconfig_data_format(cb_e);
pack_tile(0, cb_e);
tile_regs_release();
cb_push_back(cb_e, 1);
cb_pop_front(cb_l, 1);
// G = SCALE / (1 + E @ ONES)
cb_wait_front(cb_e, 1);
tile_regs_acquire();
matmul_init(cb_e, cb_one);
matmul_tiles(cb_e, cb_one, 0, 0, 0);
binop_with_scalar_tile_init();
add_unary_tile(0, ONE);
recip_tile_init();
recip_tile(0);
binop_with_scalar_tile_init();
mul_unary_tile(0, SCALE);
tile_regs_commit();
cb_reserve_back(cb_g, 1);
tile_regs_wait();
pack_reconfig_data_format(cb_g);
pack_tile(0, cb_g);
tile_regs_release();
cb_push_back(cb_g, 1);
cb_pop_front(cb_e, 1);
// out_k = bf16(relu(t_k * G + b_k))
cb_wait_front(cb_g, 1);
cb_wait_front(cb_t, 3);
cb_reserve_back(cb_out, 3);
for (uint32_t k = 0; k < 3; ++k) {
tile_regs_acquire();
copy_init(cb_t);
copy_tile(cb_t, k, 0);
copy_init(cb_g);
copy_tile(cb_g, 0, 1);
mul_binary_tile_init();
mul_binary_tile(0, 1, 0);
copy_init(cb_bias);
copy_tile(cb_bias, k, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
relu_tile_init();
relu_tile(0);
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();
tile_regs_wait();
pack_reconfig_data_format(cb_out);
pack_tile(0, cb_out);
tile_regs_release();
}
cb_push_back(cb_out, 3);
cb_pop_front(cb_t, 3);
cb_pop_front(cb_g, 1);
}
}
|