File size: 5,198 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 | // SPDX-License-Identifier: Apache-2.0
// MK-T compute (sceneseg-p150 OPT_PLAN 3.4, the D19 seg-tail module megakernel, tt/mkt.py).
//
// Phase 1 (decode_layer_9), per tile column (the reader's (n + 2) x 12 activation tiles, A[(r * 3 + dx) * 4 + k]):
// y[i][c] = sum over taps (dy, dx) and K tiles k of A[i + dy][dx][k] @ W9[(dy * 3 + dx) * 4 + k][c]
// matmul_block(rt = n rows, ct = 2 column tiles, row stride kt = 12), bf16 x bf16, fp32 DEST
// y = y + b9[c] one more matmul_block: the ones tile x [b_hi; b_lo] (the fp32 bias as two bf16 terms, 16
// significant bits; 1.0 x bf16 is exact), accumulated in DEST after the conv K loop. No
// unpack-to-DEST copy: an fp32 bias copy_tile after the K loop hung in DEST full-sync mode
// (logs/sceneseg/opt_mk/HANG_1.txt)
// y = gelu(y) gelu_tile<false> (ttnn.gelu, exact erf form; EPI_ACT 2)
// y = bf16(y) typecast_tile<Float32, Float16_b> (EPI_OUT_BF16), packed to CB 16 (the writer's D store)
// Phase 2 (margins), per tile column (A[(r * 3 + dx) * 2 + k] from the D store):
// m[i] = sum A[i + dy][dx][k] @ WM[(dy * 3 + dx) * 2 + k] (rt = n, ct = 1, row stride 6) + ones x [bm_hi; bm_lo]
// -> fp32 margin tiles to CB 17 (the writer decides the classes, K8).
// DEST full sync (8 fp32 tiles): phase 1 uses tiles 0 .. 2n - 1 (<= 6), phase 2 0 .. n - 1.
// CT args: [0] H, [1] W, [2] NC, [3] GX, [4] RB, [5] RE.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.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;
#ifndef MKT_DBG
#define MKT_DBG 0
#endif
constexpr uint32_t W = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t RB = get_compile_time_arg_val(4);
constexpr uint32_t RE = get_compile_time_arg_val(5);
constexpr uint32_t XT = W / 32;
constexpr uint32_t NRMAX = RB + (RE ? 1 : 0) + 2;
constexpr uint32_t cb_a = 0, cb_w9 = 1, cb_b9 = 2, cb_wm = 3, cb_bm = 4, cb_d9 = 16, cb_m = 17, cb_one = 10;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
const uint32_t n = RB + (core < RE ? 1 : 0);
compute_kernel_hw_startup(cb_a, cb_w9, cb_d9);
cb_wait_front(cb_w9, 72);
cb_wait_front(cb_b9, 2);
cb_wait_front(cb_wm, 18);
cb_wait_front(cb_bm, 1);
cb_wait_front(cb_one, 3);
// ---------------- phase 1
const uint32_t no = 2 * n;
for (uint32_t x = 0; x < XT; ++x) {
cb_wait_front(cb_a, NRMAX * 12); // the reader's fixed push size (rows nr .. NRMAX - 1 unused)
tile_regs_acquire();
matmul_block_init(cb_a, cb_w9, 0, 2, n, 12);
for (uint32_t t = 0; t < 9; ++t) {
for (uint32_t k = 0; k < 4; ++k) {
const uint32_t kt = t * 4 + k;
// tap t = dy * 3 + dx: A tile of row i is (i + dy) * 12 + dx * 4 + k = i * 12 + kt
matmul_block(cb_a, cb_w9, kt, kt * 2, 0, 0, 2, n, 12);
}
}
// + bias: ones tile (columns 0 / 1 = 1.0, one per output row: row stride 1) x [b_hi; b_lo]
matmul_block_init(cb_one, cb_b9, 0, 2, n, 1);
matmul_block(cb_one, cb_b9, 0, 0, 0, 0, 2, n, 1);
gelu_tile_init<false>();
for (uint32_t j = 0; j < no; ++j) {
gelu_tile<false>(j);
}
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
for (uint32_t j = 0; j < no; ++j) {
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(j);
}
tile_regs_commit();
cb_reserve_back(cb_d9, no);
tile_regs_wait();
for (uint32_t j = 0; j < no; ++j) {
pack_tile(j, cb_d9);
}
tile_regs_release();
cb_push_back(cb_d9, no);
cb_pop_front(cb_a, NRMAX * 12);
}
#if MKT_DBG != 1
// ---------------- phase 2
pack_reconfig_data_format(cb_m);
for (uint32_t x = 0; x < XT; ++x) {
cb_wait_front(cb_a, NRMAX * 6);
tile_regs_acquire();
matmul_block_init(cb_a, cb_wm, 0, 1, n, 6);
for (uint32_t kt = 0; kt < 18; ++kt) {
matmul_block(cb_a, cb_wm, kt, kt, 0, 0, 1, n, 6);
}
matmul_block_init(cb_one, cb_bm, 0, 1, n, 1);
matmul_block(cb_one, cb_bm, 0, 0, 0, 0, 1, n, 1);
tile_regs_commit();
cb_reserve_back(cb_m, n);
tile_regs_wait();
for (uint32_t i = 0; i < n; ++i) {
pack_tile(i, cb_m);
}
tile_regs_release();
cb_push_back(cb_m, n);
cb_pop_front(cb_a, NRMAX * 6);
}
#endif
cb_pop_front(cb_w9, 72);
cb_pop_front(cb_b9, 2);
cb_pop_front(cb_wm, 18);
cb_pop_front(cb_bm, 1);
cb_pop_front(cb_one, 3);
}
|