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 /epi_compute.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
6.99 kB
// SPDX-License-Identifier: Apache-2.0
// EPI compute (sceneseg-p150 OPT_PLAN EPI, tt/epi.py): the conv / linear epilogue of tt/layers.py finish() in one
// pass per tile, with the LLKs and the order of the stock programs it replaces, so the output is bit-identical:
// y = p0 + p1 (+ p2) add_binary_tile<NearestEven> (binary_ng's fp32 SFPU add, operands unpacked to DEST
// as fp32: UnpackToDestFp32 on CBs 0..5)
// y = y + bias[c] add_binary_tile<NearestEven> (binary_ng's row broadcast: the reader-filled bias tile)
// y = act(y) silu_tile (EPI_ACT 1, ttnn.silu) / gelu_tile<false> (EPI_ACT 2, ttnn.gelu exact erf
// form, fast_and_approximate_mode=False) / sigmoid_tile<RC, false> (EPI_ACT 3, ttnn.sigmoid's
// default: vector mode RC, SigmoidMode.Accurate); math_approx_mode false as the unary factory
// EPI-split (OPT_REPORT R2.3), each the stock op it replaces:
// y = y + res add_binary_tile (EPI_RES; the MBConv residual ttnn.add(out, x))
// y = y * row[c] mul_binary_tile (EPI_ROW; the SE gate ttnn.multiply(y, s), row-broadcast like binary_ng)
// main = y fp32, or bf16 via typecast_tile<Float32, Float16_b> (EPI_OUT_BF16; ttnn.typecast)
// hi = bf16(y), lo = bf16(y - hi) (EPI_SPLIT; layers.split_activation: typecast, typecast back (exact),
// subtract (fp32 SFPU), typecast), hi / lo packed to CBs 17 / 18 as bf16
// DEST tiles: 0 = y, 1 = operand / hi, 2 = lo. EPI_MAIN 0 skips the main output (split only).
// CT args: [0] T, [1] NC, [2] GX, [3] CT, [4] NP, [5] HAS_BIAS, [6] UNIT (work split as the reader).
// Defines: EPI_ACT, EPI_OUT_BF16, EPI_RES, EPI_ROW, EPI_SPLIT, EPI_MAIN.
#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/gelu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#ifndef EPI_RES
#define EPI_RES 0
#endif
#ifndef EPI_ROW
#define EPI_ROW 0
#endif
#ifndef EPI_SPLIT
#define EPI_SPLIT 0
#endif
#ifndef EPI_MAIN
#define EPI_MAIN 1
#endif
using namespace ckernel;
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 CT = get_compile_time_arg_val(3);
constexpr uint32_t NP = get_compile_time_arg_val(4);
constexpr uint32_t HAS_BIAS = get_compile_time_arg_val(5);
constexpr uint32_t UNIT = get_compile_time_arg_val(6);
constexpr uint32_t cb_bias = 3, cb_res = 4, cb_row = 5, cb_out = 16, cb_hi = 17, cb_lo = 18;
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 / UNIT / NC, E = T / UNIT % NC;
const uint32_t t0 = (core * B + (core < E ? core : E)) * UNIT;
const uint32_t n = (B + (core < E ? 1 : 0)) * UNIT;
if (n == 0) {
return;
}
#if EPI_MAIN
compute_kernel_hw_startup(0, cb_out);
#else
compute_kernel_hw_startup(0, cb_hi);
#endif
if constexpr (HAS_BIAS) {
cb_wait_front(cb_bias, CT);
}
#if EPI_ROW
cb_wait_front(cb_row, CT);
#endif
for (uint32_t t = t0; t < t0 + n; ++t) {
cb_wait_front(0, 1);
if constexpr (NP > 1) {
cb_wait_front(1, 1);
}
if constexpr (NP > 2) {
cb_wait_front(2, 1);
}
#if EPI_RES
cb_wait_front(cb_res, 1);
#endif
tile_regs_acquire();
copy_init(0);
copy_tile(0, 0, 0);
if constexpr (NP > 1) {
copy_init(1);
copy_tile(1, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
if constexpr (NP > 2) {
copy_init(2);
copy_tile(2, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
if constexpr (HAS_BIAS) {
copy_init(cb_bias);
copy_tile(cb_bias, t % CT, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
#if EPI_ACT == 1
silu_tile_init();
silu_tile(0);
#elif EPI_ACT == 2
gelu_tile_init<false>();
gelu_tile<false>(0);
#elif EPI_ACT == 3
sigmoid_tile_init<false>();
sigmoid_tile<VectorMode::RC, false>(0);
#endif
#if EPI_RES
copy_init(cb_res);
copy_tile(cb_res, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#endif
#if EPI_ROW
copy_init(cb_row);
copy_tile(cb_row, t % CT, 1);
mul_binary_tile_init();
mul_binary_tile(0, 1, 0);
#endif
#if EPI_SPLIT
// hi (tile 1) = bf16(y); lo (tile 2) = bf16(y - hi); y stays in tile 0
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);
#endif
#if EPI_MAIN && EPI_OUT_BF16
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
#endif
tile_regs_commit();
#if EPI_MAIN
cb_reserve_back(cb_out, 1);
#endif
#if EPI_SPLIT
cb_reserve_back(cb_hi, 1);
cb_reserve_back(cb_lo, 1);
#endif
tile_regs_wait();
#if EPI_MAIN
#if EPI_SPLIT
pack_reconfig_data_format(cb_out);
#endif
pack_tile(0, cb_out);
#endif
#if EPI_SPLIT
pack_reconfig_data_format(cb_hi);
pack_tile(1, cb_hi);
pack_tile(2, cb_lo);
#endif
tile_regs_release();
#if EPI_MAIN
cb_push_back(cb_out, 1);
#endif
#if EPI_SPLIT
cb_push_back(cb_hi, 1);
cb_push_back(cb_lo, 1);
#endif
cb_pop_front(0, 1);
if constexpr (NP > 1) {
cb_pop_front(1, 1);
}
if constexpr (NP > 2) {
cb_pop_front(2, 1);
}
#if EPI_RES
cb_pop_front(cb_res, 1);
#endif
}
if constexpr (HAS_BIAS) {
cb_pop_front(cb_bias, CT);
}
#if EPI_ROW
cb_pop_front(cb_row, CT);
#endif
}