File size: 4,679 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 | // SPDX-License-Identifier: Apache-2.0
// TGATE compute (meteor-p150 OPT2 round 3 R3.2, tt/tgate.py): the temporal fuse's gate glue in one pass per 32-row
// block. The merged 1x1 conv writes t = W_t x (96 columns, CB 0) and d_j = l_j - l_0 (j = 1..3, columns 0..2 of the
// fourth tile, CB 1; its columns 3..31 are the zero padding of the conv output):
// E = exp(d) SFPU, fp32 DEST (exp(0) = 1 in the padding columns)
// S = E @ ONES FPU matmul with the ones tile (rows 0..2 = 1): every column = sum_j exp(d_j)
// G = SCALE / (1 + S) SFPU (= SCALE * softmax(l)[0], identical in every column of a row)
// out_k = bf16_rne(relu(t_k * G + b_k)) fp32 SFPU mul / add, relu, RNE typecast (k = 0..2)
// Every CB reaches DEST through SrcA / SrcB (TF32; tt/tgate.py: fp32 unpack-to-DEST before the matmul broke it).
// CBs: 0 t (3 tiles), 1 d, 2 ones, 3 bias (3 tiles, every row the bias), 4 E (loop-back), 5 G (loop-back), 16 out.
// 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/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_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(cb_l, cb_e);
cb_wait_front(cb_one, 1);
cb_wait_front(cb_bias, 3);
for (uint32_t b = b0; b < b0 + n; ++b) {
// E = exp(d)
cb_wait_front(cb_l, 1);
tile_regs_acquire();
copy_init(cb_l);
copy_tile(cb_l, 0, 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);
}
}
|