File size: 9,754 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 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 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | // SPDX-License-Identifier: Apache-2.0
// MK-H (OPT round 5, FRNET_POINT_MK=head; OPT_PLAN 3.1 / D19): FRHead as ONE row-streamed module megakernel, compute.
// Per tile row (32 points) of this core's range, every intermediate stays in L1 / DEST:
// g = tilize(gathered fuse2 rows) (tilize_block: exact bf16 datacopy, as P2P)
// layer 0: y = act([g | g] @ [W_hi; W_lo] + b) (bf16 input: x_lo = 0)
// layer l: y = act([x_hi | x_hi | x_lo] @ [W_hi; W_lo; W_hi] + b) x = y_{l-1} + skip_{l-1}
// with exactly the arithmetic of the stock path it replaces (KCAT operand kernel + ttnn.linear 1D program, OPT rounds
// 1-4): the K-concatenated matmul accumulated over K' in fp32 DEST in K' order (matmul_block, HiFi4; the stock
// matmul's cross-block reload is exact, UnpackToDestFp32), packed fp32, the bias added by the FPU row broadcast of
// the packed partials (add_tiles_bcast_rows: the stock fused-bias epilogue), ReLU, then in DEST: + skip (copy_tile
// with UnpackToDestFp32 + add_binary_tile<NearestEven>, the HEAD-SKIP add) and the two-term split of the sum
// (typecast_tile RNE, sub_binary_tile, typecast_tile: kcat_compute.cpp), packed bf16 as the next layer's x_hi / x_lo.
// The last layer (logits) is packed fp32 to CB 20 for the writer.
// CT args: [0] per-core RT-arg count (P3), [1] NL, [2] KT0, [3] (unused), [4] (unused), [5 + 3 l ..] per layer:
// K' tiles, Nt, SKIP; [20 + l] RELU per layer; [25] SB (output tiles per DEST batch of the matmul, <= 4).
// Per-core RT args: [nr, 0].
#include <cstdint>
#include "api/compute/bcast.h"
#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/relu.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"
#include "api/compute/tilize.h"
using namespace ckernel;
constexpr uint32_t MAXL = 5;
constexpr uint32_t NL = get_compile_time_arg_val(1);
constexpr uint32_t KT0 = get_compile_time_arg_val(2);
constexpr uint32_t KP[MAXL] = {get_compile_time_arg_val(5), get_compile_time_arg_val(8), get_compile_time_arg_val(11),
get_compile_time_arg_val(14), get_compile_time_arg_val(17)};
constexpr uint32_t NT[MAXL] = {get_compile_time_arg_val(6), get_compile_time_arg_val(9), get_compile_time_arg_val(12),
get_compile_time_arg_val(15), get_compile_time_arg_val(18)};
constexpr uint32_t SK[MAXL] = {get_compile_time_arg_val(7), get_compile_time_arg_val(10), get_compile_time_arg_val(13),
get_compile_time_arg_val(16), get_compile_time_arg_val(19)};
constexpr uint32_t RL[MAXL] = {get_compile_time_arg_val(20), get_compile_time_arg_val(21), get_compile_time_arg_val(22),
get_compile_time_arg_val(23), get_compile_time_arg_val(24)};
constexpr uint32_t SB = get_compile_time_arg_val(25);
constexpr uint32_t EB = get_compile_time_arg_val(26); // epilogue tiles per DEST acquire (2 slots each; fp32 DEST: <= 2)
constexpr uint32_t cb_rm = 0, cb_w = 2, cb_b = 3, cb_s = 4, cb_g = 5;
constexpr uint32_t cb_hi[2] = {16, 18}, cb_lo[2] = {17, 19};
constexpr uint32_t cb_out = 20, cb_i = 24;
void kernel_main() {
const uint32_t nr = get_arg_val<uint32_t>(0);
if (nr == 0) {
return;
}
uint32_t woff[MAXL], boff[MAXL];
{
uint32_t w = 0, b = 0;
for (uint32_t l = 0; l < NL; ++l) {
woff[l] = w;
boff[l] = b;
w += KP[l] * NT[l];
b += NT[l];
}
compute_kernel_hw_startup<SrcOrder::Reverse>(cb_g, cb_w, cb_i);
cb_wait_front(cb_w, w);
cb_wait_front(cb_b, b);
}
for (uint32_t r = 0; r < nr; ++r) {
// pixel2point(fuse2): the gathered row-major block -> KT0 bf16 tiles
reconfig_data_format_srca(cb_rm);
pack_reconfig_data_format(cb_g);
tilize_init(cb_rm, KT0, cb_g);
cb_wait_front(cb_rm, KT0);
cb_reserve_back(cb_g, KT0);
tilize_block(cb_rm, KT0, cb_g);
cb_push_back(cb_g, KT0);
cb_pop_front(cb_rm, KT0);
tilize_uninit(cb_rm, cb_g);
for (uint32_t l = 0; l < NL; ++l) {
const uint32_t nt = NT[l], kp = KP[l];
const uint32_t a = l == 0 ? cb_g : cb_hi[(l - 1) & 1];
const uint32_t lo = l == 0 ? cb_g : cb_lo[(l - 1) & 1];
const uint32_t kt = l == 0 ? kp / 2 : kp / 3;
const bool last = l + 1 == NL;
const uint32_t sb = nt < SB ? nt : SB;
// y = X' @ W' (fp32 DEST over all K' tiles) -> fp32 partials in CB 24
cb_wait_front(a, kt);
if (l > 0) {
cb_wait_front(lo, kt);
}
reconfig_data_format<SrcOrder::Reverse>(a, cb_w);
pack_reconfig_data_format(cb_i);
matmul_block_init(a, cb_w, 0, sb, 1, kp);
for (uint32_t n0 = 0; n0 < nt; n0 += sb) {
tile_regs_acquire();
for (uint32_t k = 0; k < kp; ++k) {
uint32_t cb = a, it = k;
if (k >= 2 * kt) {
cb = lo;
it = k - 2 * kt;
} else if (k >= kt) {
it = k - kt;
}
matmul_block(cb, cb_w, it, woff[l] + k * nt + n0, 0, 0, sb, 1, kp);
}
tile_regs_commit();
cb_reserve_back(cb_i, sb);
tile_regs_wait();
for (uint32_t i = 0; i < sb; ++i) {
pack_tile(i, cb_i);
}
tile_regs_release();
cb_push_back(cb_i, sb);
}
cb_pop_front(a, kt);
if (l > 0) {
cb_pop_front(lo, kt);
}
// epilogue per output tile: + bias (FPU row broadcast), ReLU, then + skip and the split, or fp32 out
const uint32_t ohi = cb_hi[l & 1], olo = cb_lo[l & 1];
if (SK[l]) {
cb_wait_front(cb_s, nt);
}
for (uint32_t n = 0; n < nt; n += EB) {
const uint32_t ge = (nt - n) < EB ? (nt - n) : EB; // tiles of this DEST group (2 slots each)
cb_wait_front(cb_i, ge);
reconfig_data_format(cb_i, cb_b);
add_bcast_rows_init(cb_i, cb_b);
tile_regs_acquire();
for (uint32_t e = 0; e < ge; ++e) {
add_tiles_bcast_rows(cb_i, cb_b, e, boff[l] + n + e, 2 * e);
}
if (RL[l]) {
relu_tile_init();
for (uint32_t e = 0; e < ge; ++e) {
relu_tile(2 * e);
}
}
if (last) {
tile_regs_commit();
cb_reserve_back(cb_out, ge);
pack_reconfig_data_format(cb_out);
tile_regs_wait();
for (uint32_t e = 0; e < ge; ++e) {
pack_tile(2 * e, cb_out);
}
tile_regs_release();
cb_push_back(cb_out, ge);
} else {
if (SK[l]) {
reconfig_data_format_srca(cb_s);
copy_init(cb_s);
for (uint32_t e = 0; e < ge; ++e) {
copy_tile(cb_s, n + e, 2 * e + 1);
}
add_binary_tile_init();
for (uint32_t e = 0; e < ge; ++e) {
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(2 * e, 2 * e + 1, 2 * e);
}
}
copy_dest_values_init();
for (uint32_t e = 0; e < ge; ++e) {
copy_dest_values<DataFormat::Float32>(2 * e, 2 * e + 1);
}
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
for (uint32_t e = 0; e < ge; ++e) {
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2 * e);
}
sub_binary_tile_init();
for (uint32_t e = 0; e < ge; ++e) {
sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2 * e + 1, 2 * e, 2 * e + 1);
}
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
for (uint32_t e = 0; e < ge; ++e) {
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2 * e + 1);
}
tile_regs_commit();
cb_reserve_back(ohi, ge);
cb_reserve_back(olo, ge);
pack_reconfig_data_format(ohi);
tile_regs_wait();
for (uint32_t e = 0; e < ge; ++e) {
pack_tile(2 * e, ohi);
pack_tile(2 * e + 1, olo);
}
tile_regs_release();
cb_push_back(ohi, ge);
cb_push_back(olo, ge);
}
cb_pop_front(cb_i, ge);
}
if (SK[l]) {
cb_pop_front(cb_s, nt);
}
}
}
}
|