File size: 6,990 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 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 | // 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
}
|