File size: 7,618 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 | // SPDX-License-Identifier: Apache-2.0
// K11 depthwise compute (sceneseg-p150 OPT_PLAN K11, tt/dw.py): per output tile (32 pixels x 32 channels, float32)
// acc = x_0 * w_0; acc = x_t * w_t + acc (t = 1 .. KK-1, SFPMAD: one fp32 rounding per tap)
// acc = acc + bias (add_binary_tile, fp32 SFPU add, as EPI)
// acc = silu(acc) (silu_tile, ttnn.silu's LLK, as EPI) [DW_ACT 1; 0 = none]
// x_t: the gathered input page of tap t (CB 0, unpacked to DEST as fp32), w_t: weight tile t of the current channel
// tile (CB 1, every row = the 32 channel weights of tap t), bias: CB 2. DEST tiles: 0 = acc, 1 = x, 2 = w.
// Work split and order as dw_reader.cpp; the weight / bias CBs are popped when the channel tile changes.
// CT args: [0] KK, [1] CT, [2] PT, [3] NC, [4] GX, [5] CH, [6] NWT (weight tiles per channel tile). Defines: DW_ACT,
// DW_WPACK.
// R7.1 DW_WPACK=1: the NWT = ceil(KK / 8) weight tiles hold 8 taps each (rows 4l .. 4l+3 = tap 8j + l); per output
// tile they are unpacked to DEST tiles 2 / 3 once (5x5: again for taps 16 .. 24) and tap t's SFPMADs read their weight
// vectors from its 4-row group (k11_tap_tile). An SFPU vector is 4 DEST rows x 8 columns of one parity (sfpi unit
// 2k = rows 4k' .. +3 even columns, 2k + 1 the odd ones), so the x vector of unit d takes the group's unit of the same
// parity: the same operands per SFPMAD as the one-tap-per-tile form (bit-identical).
#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/tile_move_copy.h"
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// one face (8 row groups) of DEST tile 0 = tile 1 * tile 2 (+ tile 0)
inline void k11_mul_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::dst_reg[0] = x * w;
sfpi::dst_reg++;
}
}
inline void k11_mad_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat a = sfpi::dst_reg[0];
sfpi::dst_reg[0] = x * w + a;
sfpi::dst_reg++;
}
}
// R7.1: the whole tile in one call (VectorMode::None, base = DEST tile 0): tile 0 = tile 1 * w (+ tile 0), w = the
// tap's weight group at sfpi unit WR relative to the face base for faces 0 / 1, WR - 16 for faces 2 / 3.
template <bool FIRST>
inline void k11_tap_tile(const std::uint32_t, const std::uint32_t, const std::uint32_t, const std::uint32_t wr) {
#pragma GCC unroll 4
for (int f = 0; f < 4; f++) {
const int wo = (int)wr - (f >= 2 ? 16 : 0);
sfpi::vFloat we = sfpi::dst_reg[wo];
sfpi::vFloat wd = sfpi::dst_reg[wo + 1];
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
if constexpr (FIRST) {
sfpi::dst_reg[0] = x * ((d & 1) ? wd : we);
} else {
sfpi::vFloat a = sfpi::dst_reg[0];
sfpi::dst_reg[0] = x * ((d & 1) ? wd : we) + a;
}
sfpi::dst_reg++;
}
}
}
} // namespace ckernel::sfpu
#endif
using namespace ckernel;
constexpr uint32_t KK = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t PT = 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 CH = get_compile_time_arg_val(5);
constexpr uint32_t NWT = get_compile_time_arg_val(6);
constexpr uint32_t NW = CT * PT;
constexpr uint32_t cb_x = 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 = NW / NC, E = NW % NC;
const uint32_t w0 = core * B + (core < E ? core : E);
const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
compute_kernel_hw_startup(cb_x, cb_out);
uint32_t cur_ct = 0xFFFFFFFF;
uint32_t w = w0;
while (w < w1) {
const uint32_t ct = w / PT;
const uint32_t pt0 = w - ct * PT;
uint32_t pt1 = pt0 + CH;
if (pt1 > PT) {
pt1 = PT;
}
if (pt1 - pt0 > w1 - w) {
pt1 = pt0 + (w1 - w);
}
if (ct != cur_ct) {
if (cur_ct != 0xFFFFFFFF) {
cb_pop_front(cb_w, NWT);
cb_pop_front(cb_b, 1);
}
cb_wait_front(cb_w, NWT);
cb_wait_front(cb_b, 1);
cur_ct = ct;
}
for (uint32_t pt = pt0; pt < pt1; ++pt) {
tile_regs_acquire();
#if DW_WPACK == 1
copy_init(cb_w);
copy_tile(cb_w, 0, 2);
if constexpr (NWT > 1) {
copy_tile(cb_w, 1, 3);
}
copy_init(cb_x);
add_binary_tile_init();
for (uint32_t t = 0; t < KK; ++t) {
if (t == 16) {
copy_init(cb_w);
copy_tile(cb_w, 2, 2);
if constexpr (NWT > 3) {
copy_tile(cb_w, 3, 3);
}
copy_init(cb_x);
add_binary_tile_init();
}
cb_wait_front(cb_x, 1);
copy_tile(cb_x, 0, 1);
const uint32_t wr = 64 + 32 * ((t >> 3) & 1) + 16 * ((t >> 2) & 1) + 2 * (t & 3);
if (t == 0) {
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (true), 0, 1, 0,
VectorMode::None, wr)));
} else {
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (false), 0, 1, 0,
VectorMode::None, wr)));
}
cb_pop_front(cb_x, 1);
}
#else
for (uint32_t t = 0; t < KK; ++t) {
cb_wait_front(cb_x, 1);
copy_init(cb_x);
copy_tile(cb_x, 0, 1);
copy_init(cb_w);
copy_tile(cb_w, t, 2);
add_binary_tile_init();
if (t == 0) {
MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mul_face, 0, 1, 0,
VectorMode::RC)));
} else {
MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mad_face, 0, 1, 0,
VectorMode::RC)));
}
cb_pop_front(cb_x, 1);
}
#endif
copy_init(cb_b);
copy_tile(cb_b, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#if DW_ACT == 1
silu_tile_init();
silu_tile(0);
#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);
}
w += pt1 - pt0;
}
if (cur_ct != 0xFFFFFFFF) {
cb_pop_front(cb_w, NWT);
cb_pop_front(cb_b, 1);
}
}
|