File size: 7,531 Bytes
55474c5 | 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 | // SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py), compute. The C14 x0.8 INTER_LINEAR resize of the normalised node image on the SFPU, in the
// exact float32 steps of the host twin (host.fast; each lerp is fma(x1 - x0, w, x0): the difference rounded to
// fp32, one fused multiply-add rounded once; probe logs/bevformer/opt_r4/probe_sfpu: SFPU sub / SFPMAD are the
// correctly rounded fp32 results on this domain, 0 of 262144 differ):
// H row r (4 fp32 pseudo-tiles): H = (P1 - P0) * WX + P0 (P0 / P1 / WX from the reader)
// output row y (4 bf16 pseudo-tiles): V = (H[r0 + 1] - H[r0]) * WY[y % 4] + H[r0], rounded to bf16 (RNE, as
// torch's float32 -> bfloat16 cast of the host path) on the SFPU, then packed (exact)
// cb_h keeps all H rows of the core (no wrap of the 2-row window).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "kpre_common.h"
#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"
namespace ckernel::sfpu {
// DEST tiles 0 / 1 / 2 = x0, x1, w -> tile 0 = fma(x1 - x0, w, x0)
inline void kpre_lerp() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vFloat x0 = sfpi::dst_reg[0];
sfpi::vFloat x1 = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat dx = x1 - x0;
sfpi::dst_reg[0] = dx * w + x0;
sfpi::dst_reg++;
}
}
// U16 rows: DEST tiles 0 / 1 = the uint16 bytes of planes p / p + 1 (low 16 bits of each 32-bit word), 2 = the
// per-lane -mean_c, 3 = w -> tile 0 = fma(x1 - x0, w, x0) with x = f32(u) + (-mean_c) (exact int -> fp32 cast, then
// the single fp32 add of the host's normalisation)
inline void kpre_lerp_u16() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vInt u0 = sfpi::dst_reg[0];
sfpi::vInt u1 = sfpi::dst_reg[32];
sfpi::vFloat beta = sfpi::dst_reg[64];
sfpi::vFloat w = sfpi::dst_reg[96];
sfpi::vFloat x0 = sfpi::convert<sfpi::vFloat>(u0 & 0xFFFF) + beta;
sfpi::vFloat x1 = sfpi::convert<sfpi::vFloat>(u1 & 0xFFFF) + beta;
sfpi::vFloat dx = x1 - x0;
sfpi::dst_reg[0] = dx * w + x0;
sfpi::dst_reg++;
}
}
// the same, then round to bf16 (RNE: + 0x7fff + lsb, low 16 bits cleared; the bf16 pack is then exact)
inline void kpre_lerp_bf16() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vFloat x0 = sfpi::dst_reg[0];
sfpi::vFloat x1 = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat dx = x1 - x0;
sfpi::vFloat v = dx * w + x0;
sfpi::vUInt bits = sfpi::as<sfpi::vUInt>(v);
sfpi::vUInt lsb = (bits >> 16) & 1;
bits = bits + 0x7fffU + lsb;
bits = bits & 0xFFFF0000U;
sfpi::dst_reg[0] = sfpi::as<sfpi::vFloat>(bits);
sfpi::dst_reg++;
}
}
} // namespace ckernel::sfpu
#endif
constexpr uint32_t N = get_compile_time_arg_val(0);
constexpr uint32_t S = get_compile_time_arg_val(1);
constexpr uint32_t W = get_compile_time_arg_val(2);
constexpr uint32_t RH = get_compile_time_arg_val(3);
constexpr uint32_t NC = get_compile_time_arg_val(4);
constexpr uint32_t GX = get_compile_time_arg_val(5);
constexpr uint32_t RA = get_compile_time_arg_val(6);
constexpr uint32_t RB = get_compile_time_arg_val(7);
constexpr uint32_t RD = get_compile_time_arg_val(8);
constexpr uint32_t SPLIT = get_compile_time_arg_val(9); // 1: odd rows (from rlo) come from the writer RISC in cb_p2
constexpr uint32_t U16 = get_compile_time_arg_val(10); // 1: the P rows are raw uint16 planes (+ cb_beta)
constexpr uint32_t cb_p = 0, cb_wx = 1, cb_wy = 2, cb_h = 3, cb_p2 = 5, cb_beta = 7, cb_v = 16;
constexpr uint32_t QT = 4; // pseudo-tiles per row
template <bool BF16>
inline void lerp_rows(uint32_t cb_x, uint32_t i0, uint32_t i1, uint32_t cb_w, uint32_t iw, uint32_t cb_o) {
tile_regs_acquire();
#ifndef KPRE_NOCOPY
copy_tile(cb_x, i0, 0);
copy_tile(cb_x, i1, 1);
copy_tile(cb_w, iw, 2);
#endif
#ifndef KPRE_NOSFPU
if constexpr (BF16) {
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_bf16, 0, VectorMode::RC)));
} else {
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp, 0, VectorMode::RC)));
}
#endif
tile_regs_commit();
tile_regs_wait();
pack_tile(0, cb_o);
tile_regs_release();
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
const KpreGeom g = kpre_geom<N, S, W, RH, NC, RA, RB, RD>(core);
if (!g.rows) {
return;
}
compute_kernel_hw_startup(cb_p, cb_h);
copy_init(cb_p);
MATH((SFPU_UNARY_INIT(unused)));
cb_wait_front(cb_wx, QT);
cb_wait_front(cb_wy, QT);
if constexpr (U16) {
cb_wait_front(cb_beta, 1);
}
// cb_h holds every H row of the core (rlo .. rhi, sized on the host): rows are never popped before the end, so
// the two rows an output row reads are always contiguous pages (a FIFO window could wrap the ring)
uint32_t produced = g.rlo; // next H row to compute
for (uint32_t y = g.ya; y <= g.yv; ++y) {
const uint32_t r0 = kpre_r0<RA, RB, RD>(y);
while (produced <= r0 + 1) {
// H row `produced`: (P1 - P0) * WX + P0, fp32
// phase planes s = 0..4 of the node row: H plane p = lerp(plane p, plane p + 1, (2p + 1) / 8)
const uint32_t cbp = (SPLIT && ((produced - g.rlo) & 1)) ? cb_p2 : cb_p;
cb_wait_front(cbp, QT + 1);
cb_reserve_back(cb_h, QT);
pack_reconfig_data_format(cb_h);
if constexpr (U16) {
for (uint32_t t = 0; t < QT; ++t) {
tile_regs_acquire();
copy_tile_to_dst_init_short_with_dt(cb_wx, cbp);
copy_tile(cbp, t, 0);
copy_tile(cbp, t + 1, 1);
copy_tile_to_dst_init_short_with_dt(cbp, cb_wx);
copy_tile(cb_beta, 0, 2);
copy_tile(cb_wx, t, 3);
#ifndef KPRE_NOSFPU
MATH((SFPU_UNARY_INIT(unused)));
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_u16, 0,
VectorMode::RC)));
#endif
tile_regs_commit();
tile_regs_wait();
pack_tile(0, cb_h);
tile_regs_release();
}
} else {
for (uint32_t t = 0; t < QT; ++t) {
lerp_rows<false>(cbp, t, t + 1, cb_wx, t, cb_h);
}
}
cb_push_back(cb_h, QT);
cb_pop_front(cbp, QT + 1);
++produced;
}
// output row y: (H[r0 + 1] - H[r0]) * WY + H[r0] -> bf16
const uint32_t h0 = (r0 - g.rlo) * QT;
cb_wait_front(cb_h, h0 + 2 * QT);
cb_reserve_back(cb_v, QT);
pack_reconfig_data_format(cb_v);
for (uint32_t t = 0; t < QT; ++t) {
lerp_rows<true>(cb_h, h0 + t, h0 + QT + t, cb_wy, y % 4, cb_v);
}
cb_push_back(cb_v, QT);
}
cb_pop_front(cb_h, (produced - g.rlo) * QT);
if constexpr (U16) {
cb_pop_front(cb_beta, 1);
}
}
|