File size: 3,001 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 | // SPDX-License-Identifier: Apache-2.0
// Stem compute (sceneseg-p150 OPT_PLAN K-IN kernel form / MK0, tt/stem.py): per output tile of the stem conv,
// y = A @ W_hi matmul_tiles into DEST 0 (bf16 x bf16 products, fp32 DEST, HiFi4)
// y = y + A @ W_lo matmul_tiles into DEST 1, then add_binary_tile<NearestEven> (the EPI term sum)
// y = y + bias the exact fp32 bias-map tile (stem_fold), add_binary_tile<NearestEven>
// y = silu(y) silu_tile (ttnn.silu), fp32 out
// CT args: [0] T, [1] NC, [2] GX (work split as the reader).
#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/matmul.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
#ifndef STEM_DBG
#define STEM_DBG 0 // probe only: 1 = A @ W_hi, 2 = the bias tile, 3 = the term sum (no bias / SiLU)
#endif
constexpr uint32_t T = 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 cb_a = 0, cb_w = 1, cb_b = 2, 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 = T / NC, E = T % NC;
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
compute_kernel_hw_startup(cb_a, cb_w, cb_out);
cb_wait_front(cb_w, 2);
for (uint32_t t = 0; t < n; ++t) {
cb_wait_front(cb_a, 1);
cb_wait_front(cb_b, 1);
tile_regs_acquire();
#if STEM_DBG == 2
reconfig_data_format_srca(cb_b);
copy_init(cb_b);
copy_tile(cb_b, 0, 0);
reconfig_data_format_srca(cb_w);
#else
matmul_init(cb_a, cb_w);
matmul_tiles(cb_a, cb_w, 0, 0, 0);
#if STEM_DBG != 1
matmul_tiles(cb_a, cb_w, 0, 1, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#if STEM_DBG != 3
// the fp32 bias tile goes to DEST through the unpacker's fp32 path: srcA reconfigured from the bf16 weights
// (matmul: srcA = in1) to fp32 and back (bmm_large_block_zm_fused_bias_activation.cpp does the same)
reconfig_data_format_srca(cb_b);
copy_init(cb_b);
copy_tile(cb_b, 0, 1);
reconfig_data_format_srca(cb_w);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
silu_tile_init();
silu_tile(0);
#endif
#endif
#endif
tile_regs_commit();
cb_reserve_back(cb_out, 1);
tile_regs_wait();
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
cb_pop_front(cb_a, 1);
cb_pop_front(cb_b, 1);
}
cb_pop_front(cb_w, 2);
}
|