File size: 6,134 Bytes
e0adbab | 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 | // SPDX-License-Identifier: Apache-2.0
// SE-K compute (sceneseg-p150 OPT_PLAN SE-K, tt/sek.py): one squeeze-excitation MLP on the fp32 row sum x
// (the 1 / (H W) of the mean is folded into W1), with the port's precision scheme (tt/layers.py Linear: bf16 terms
// of the activation and of the weights, the products hi x W_hi, hi x W_lo, lo x W_hi, fp32 bias, fp32 activation)
// but all products of a layer accumulated in one fp32 DEST tile (the stock path rounded each product's matmul to
// fp32 separately and summed them in the epilogue):
// x -> (x_hi, x_lo) typecast_tile<Float32, Float16_b> + sub_binary_tile (layers.split_activation)
// h = silu(sum_k x_hi W1hi + x_hi W1lo + x_lo W1hi + b1) (every core, all NT1 tiles)
// h -> (h_hi, h_lo)
// g = sigmoid(sum_k h_hi W2hi + h_hi W2lo + h_lo W2hi + b2) (this core's fc2 output tiles)
// CT args: [0] KT1, [1] NT1, [2] NT2, [3] NC, [4] GX.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/copy_dest_values.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/matmul.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t KT1 = get_compile_time_arg_val(0);
constexpr uint32_t NT1 = get_compile_time_arg_val(1);
constexpr uint32_t NT2 = get_compile_time_arg_val(2);
constexpr uint32_t NC = get_compile_time_arg_val(3);
constexpr uint32_t GX = get_compile_time_arg_val(4);
constexpr uint32_t cb_x = 0, cb_w1 = 1, cb_b1 = 2, cb_w2 = 3, cb_b2 = 4, cb_out = 16;
constexpr uint32_t cb_xhi = 24, cb_xlo = 25, cb_hhi = 26, cb_hlo = 27;
// DEST 0 holds y (fp32): DEST 1 = bf16(y), DEST 2 = bf16(y - DEST 1), packed to the bf16 CBs hi / lo
FORCE_INLINE void split_pack(uint32_t cb_hi, uint32_t cb_lo) {
copy_dest_values_init();
copy_dest_values<DataFormat::Float32>(0, 1);
copy_dest_values<DataFormat::Float32>(0, 2);
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
sub_binary_tile_init();
sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2, 1, 2);
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2);
tile_regs_commit();
cb_reserve_back(cb_hi, 1);
cb_reserve_back(cb_lo, 1);
tile_regs_wait();
pack_reconfig_data_format(cb_hi);
pack_tile(1, cb_hi);
pack_tile(2, cb_lo);
tile_regs_release();
cb_push_back(cb_hi, 1);
cb_push_back(cb_lo, 1);
}
// srcA: the bf16 weights for matmul (srcA = in1), the fp32 CBs for the copies to DEST: reconfigured explicitly
// around each copy (as bmm_large_block_zm_fused_bias_activation.cpp does around its partials reload)
FORCE_INLINE void add_bias(uint32_t cb, uint32_t idx) {
reconfig_data_format_srca(cb);
copy_init(cb);
copy_tile(cb, idx, 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 = NT2 / NC, E = NT2 % NC;
const uint32_t nn = B + (core < E ? 1 : 0);
compute_kernel_hw_startup(cb_xhi, cb_w1, cb_out);
// ---- x -> bf16 terms
cb_wait_front(cb_x, KT1);
for (uint32_t k = 0; k < KT1; ++k) {
tile_regs_acquire();
reconfig_data_format_srca(cb_x);
copy_init(cb_x);
copy_tile(cb_x, k, 0);
split_pack(cb_xhi, cb_xlo);
}
cb_pop_front(cb_x, KT1);
// ---- fc1 + SiLU -> h terms
cb_wait_front(cb_xhi, KT1);
cb_wait_front(cb_xlo, KT1);
cb_wait_front(cb_w1, 2 * KT1 * NT1);
cb_wait_front(cb_b1, NT1);
for (uint32_t n = 0; n < NT1; ++n) {
tile_regs_acquire();
reconfig_data_format(cb_xhi, cb_w1);
matmul_init(cb_xhi, cb_w1);
for (uint32_t k = 0; k < KT1; ++k) {
matmul_tiles(cb_xhi, cb_w1, k, k * NT1 + n, 0); // x_hi W1_hi
matmul_tiles(cb_xhi, cb_w1, k, (KT1 + k) * NT1 + n, 0); // x_hi W1_lo
}
matmul_init(cb_xlo, cb_w1);
for (uint32_t k = 0; k < KT1; ++k) {
matmul_tiles(cb_xlo, cb_w1, k, k * NT1 + n, 0); // x_lo W1_hi
}
add_bias(cb_b1, n);
silu_tile_init();
silu_tile(0);
split_pack(cb_hhi, cb_hlo);
}
cb_pop_front(cb_xhi, KT1);
cb_pop_front(cb_xlo, KT1);
cb_pop_front(cb_w1, 2 * KT1 * NT1);
cb_pop_front(cb_b1, NT1);
// ---- fc2 + sigmoid on this core's output tiles
cb_wait_front(cb_hhi, NT1);
cb_wait_front(cb_hlo, NT1);
cb_wait_front(cb_w2, 2 * NT1 * nn);
cb_wait_front(cb_b2, nn);
for (uint32_t j = 0; j < nn; ++j) {
tile_regs_acquire();
reconfig_data_format(cb_hhi, cb_w2);
matmul_init(cb_hhi, cb_w2);
for (uint32_t k = 0; k < NT1; ++k) {
matmul_tiles(cb_hhi, cb_w2, k, k * nn + j, 0); // h_hi W2_hi
matmul_tiles(cb_hhi, cb_w2, k, (NT1 + k) * nn + j, 0); // h_hi W2_lo
}
matmul_init(cb_hlo, cb_w2);
for (uint32_t k = 0; k < NT1; ++k) {
matmul_tiles(cb_hlo, cb_w2, k, k * nn + j, 0); // h_lo W2_hi
}
add_bias(cb_b2, j);
sigmoid_tile_init<false>();
sigmoid_tile<VectorMode::RC, false>(0);
tile_regs_commit();
cb_reserve_back(cb_out, 1);
tile_regs_wait();
pack_reconfig_data_format(cb_out);
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
}
cb_pop_front(cb_hhi, NT1);
cb_pop_front(cb_hlo, NT1);
cb_pop_front(cb_w2, 2 * NT1 * nn);
cb_pop_front(cb_b2, nn);
}
|