Hugging Face's logo

changh95
/
sceneseg-p150

Image Segmentation
ONNX
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
sceneseg-p150 / code /tt_sceneseg /tt /kernels /se_compute.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
6.13 kB
// 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);
}