Clear windows branch; preserve previous contents in win-old
This commit is contained in:
ayo committed
1 parent
f0c881985b
commit
73994d4634
311 files changed
-57018
No files matched your search
@@ -1,10 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#ifndef D4R_K_POS_MT
|
||||
#define D4R_K_POS_MT 2
|
||||
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
|
||||
#error "experimental position token tiling requires gfx12 strict K baseline"
|
||||
#endif
|
||||
#include "pwin_pos.h"
|
||||
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, D4R_K_POS_MT)
|
||||
@@ -1,5 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)
|
||||
@@ -1,5 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)
|
||||
@@ -1,2 +0,0 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)
|
||||
@@ -1,2 +0,0 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 4)
|
||||
@@ -1,2 +0,0 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 4)
|
||||
@@ -1,19 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#ifndef D4R_K_POS_MT
|
||||
#define D4R_K_POS_MT 2
|
||||
#elif !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
|
||||
#error "experimental position token tiling requires gfx12 strict K baseline"
|
||||
#endif
|
||||
// Reuse V across the window's waves. Paired gfx1201 replays retain every half
|
||||
// code and reduce enc0 time by about 5%; gfx11 and arithmetic experiments keep
|
||||
// their previous path. D4R_K_PRIVATE_V supplies the original A/B control.
|
||||
#if defined(__GFX12__) && defined(D4R_K_FP16_BASELINE) && !defined(D4R_K_PRIVATE_V) && \
|
||||
!defined(D4R_K_PACKED_ACC) && !defined(D4R_K_F16_WMMA)
|
||||
#ifndef POS_VSHARE
|
||||
#define POS_VSHARE
|
||||
#endif
|
||||
#endif
|
||||
#include "pwin_pos.h"
|
||||
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, D4R_K_POS_MT)
|
||||
@@ -1,16 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#ifdef D4R_K_ENC1_POS_MT
|
||||
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE)
|
||||
#error "experimental enc1 token tiling requires gfx12 strict K baseline"
|
||||
#endif
|
||||
#if D4R_K_ENC1_POS_MT != 1 && D4R_K_ENC1_POS_MT != 2 && D4R_K_ENC1_POS_MT != 4
|
||||
#error "enc1 token tiles per wave must be 1, 2 or 4"
|
||||
#endif
|
||||
#include "pwin_pos.h"
|
||||
POS_ENCODER(dltss_pwin_enc1_layer, 2, 64, 64, 0, D4R_K_ENC1_POS_MT)
|
||||
#else
|
||||
#include "pwin_layer.h"
|
||||
PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
|
||||
#endif
|
||||
@@ -1,5 +0,0 @@
|
||||
// CU mode: the waves of a window share one L0 cache (1-6% faster here; slower in the wide layers)
|
||||
// d4r-build-flags: -mcumode
|
||||
#define PWIN_F32ACC
|
||||
#include "pwin_layer.h"
|
||||
PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)
|
||||
@@ -1,2 +0,0 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_ENCODER(dltss_pwin_enc3_layer, 4, 96, 128, 2)
|
||||
@@ -1,2 +0,0 @@
|
||||
#include "pwin_wide.h"
|
||||
WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 4)
|
||||
@@ -1,264 +0,0 @@
|
||||
// Native kernels (gfx11 / gfx12 wave32 WMMA) for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
|
||||
// Semantics: kernels/tools/pwin_model.py (verified against the PTX interpreter, kernels/tools/ptxsim.py).
|
||||
//
|
||||
// Layout conventions (transposed WMMA, as in the preset M kernels): a GEMM Y = X W is computed as
|
||||
// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15) and the activations
|
||||
// as the B operand (lane l: token l & 15). The f32 result then has lane l = token l & 15 and VGPR i =
|
||||
// output channel wm_acc_row(i) of the 16-channel tile ("D^T"); ../common/wmma_layout.h has the per-target
|
||||
// register layouts (gfx11: channel 2i + (l >> 4); gfx12: channel i + 8 (l >> 4)).
|
||||
#pragma once
|
||||
#ifdef D4R_DEVICE_ONLY_MINIMAL
|
||||
#include "../common/hip_device_minimal.h"
|
||||
#else
|
||||
#include <hip/hip_runtime.h>
|
||||
#endif
|
||||
#include <stdint.h>
|
||||
#include "../common/wmma_layout.h"
|
||||
#if defined(D4R_ACCURACY) && defined(PWIN_F32ACC)
|
||||
#undef PWIN_F32ACC
|
||||
#endif
|
||||
#pragma clang fp contract(off)
|
||||
|
||||
typedef _Float16 half_t;
|
||||
typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
|
||||
typedef _Float16 h16 __attribute__((ext_vector_type(16)));
|
||||
typedef float f8v __attribute__((ext_vector_type(8)));
|
||||
#ifdef D4R_K_PACKED_ACC
|
||||
#if !defined(__GFX12__) || !defined(D4R_K_FP16_BASELINE) || defined(D4R_K_F16_WMMA)
|
||||
#error "packed accumulator storage requires gfx12 strict F32 WMMA baseline"
|
||||
#endif
|
||||
// Keep already-rounded values packed between steps. Every matrix instruction
|
||||
// still uses F32 C/D and rounds once to half afterwards, like the baseline.
|
||||
typedef _Float16 acc8v __attribute__((ext_vector_type(8)));
|
||||
#else
|
||||
typedef f8v acc8v;
|
||||
#endif
|
||||
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
|
||||
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
|
||||
|
||||
__device__ __forceinline__ uint32_t lane_id()
|
||||
{
|
||||
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
|
||||
}
|
||||
|
||||
typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
|
||||
|
||||
// The CUDA fragment reduction rounds products and every tree node to f16.
|
||||
// A float sum of the squared halves changes rsqrt and can amplify differences
|
||||
// through attention. The strict baseline preserves its channel grouping.
|
||||
template <int C>
|
||||
__device__ __forceinline__ half_t l2_sum(const half_t* row)
|
||||
{
|
||||
#ifdef D4R_K_FP16_BASELINE
|
||||
static_assert(C % 32 == 0 && C <= 160, "audited K channel counts");
|
||||
hv2 partial[4];
|
||||
#pragma unroll
|
||||
for (int t = 0; t < 4; ++t)
|
||||
{
|
||||
hv2 halves[2];
|
||||
#pragma unroll
|
||||
for (int parity = 0; parity < 2; ++parity)
|
||||
{
|
||||
hv2 pairs[C / 32];
|
||||
#pragma unroll
|
||||
for (int group = 0; group < C / 32; ++group)
|
||||
{
|
||||
const hv2 a = *(const hv2*)(row + 2 * t + 8 * parity + 32 * group);
|
||||
const hv2 b = *(const hv2*)(row + 2 * t + 8 * parity + 32 * group + 16);
|
||||
pairs[group] = (a * a) + (b * b);
|
||||
}
|
||||
hv2 value = pairs[0];
|
||||
if constexpr (C >= 64) value = value + pairs[1];
|
||||
if constexpr (C == 96) value = value + pairs[2];
|
||||
if constexpr (C >= 128) value = value + (pairs[2] + pairs[3]);
|
||||
if constexpr (C == 160) value = value + pairs[4];
|
||||
halves[parity] = value;
|
||||
}
|
||||
partial[t] = halves[0] + halves[1];
|
||||
}
|
||||
const hv2 sum = (partial[0] + partial[1]) + (partial[2] + partial[3]);
|
||||
return sum[0] + sum[1];
|
||||
#else
|
||||
float sum = 0;
|
||||
#pragma unroll
|
||||
for (int c = 0; c < C; ++c) sum += (float)(half_t)(row[c] * row[c]);
|
||||
return (half_t)sum;
|
||||
#endif
|
||||
}
|
||||
|
||||
// one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)
|
||||
// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
|
||||
__device__ __forceinline__ acc8v mma16(const op_t& a, const op_t& b, acc8v c)
|
||||
{
|
||||
#ifdef D4R_K_PACKED_ACC
|
||||
const f8v fc = __builtin_convertvector(c, f8v);
|
||||
const f8v d = wm_mma(a, b, fc);
|
||||
return __builtin_convertvector(d, acc8v);
|
||||
#elif defined(D4R_K_F16_WMMA)
|
||||
#if !defined(__GFX12__) || D4R_WMMA_LAYOUT != 12 || !defined(D4R_K_FP16_BASELINE)
|
||||
#error "experimental packed F16 WMMA requires gfx12 native layout and strict K baseline"
|
||||
#endif
|
||||
return wm_mma_f16_step(a, b, c);
|
||||
#else
|
||||
f8v d = wm_mma(a, b, c);
|
||||
#if !defined(PWIN_F32ACC) || defined(D4R_K_FP16_BASELINE)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)(half_t)d[i];
|
||||
#endif
|
||||
return d;
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ acc8v splat8(float v)
|
||||
{
|
||||
#ifdef D4R_K_PACKED_ACC
|
||||
const half_t h = (half_t)v;
|
||||
return (acc8v){h, h, h, h, h, h, h, h};
|
||||
#else
|
||||
return (acc8v){v, v, v, v, v, v, v, v};
|
||||
#endif
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t other_half(uint32_t v)
|
||||
{
|
||||
return wm_other_half(v);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t pack2(half_t lo, half_t hi)
|
||||
{
|
||||
return __builtin_bit_cast(uint32_t, (hv2){lo, hi});
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t lo16(uint32_t v)
|
||||
{
|
||||
return __builtin_bit_cast(hv2, v)[0];
|
||||
}
|
||||
__device__ __forceinline__ half_t hi16(uint32_t v)
|
||||
{
|
||||
return __builtin_bit_cast(hv2, v)[1];
|
||||
}
|
||||
|
||||
// D^T f16 values (own[i] = channel wm_acc_row(i) of token l & 15) -> operand of the same token (the
|
||||
// channels become K)
|
||||
__device__ __forceinline__ op_t operand_from_dt(const half_t own[8])
|
||||
{
|
||||
return wm_op_from_acc(own);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ op_t operand_from_f8(acc8v d)
|
||||
{
|
||||
#ifdef D4R_K_PACKED_ACC
|
||||
const half_t h[8] = {d[0], d[1], d[2], d[3], d[4], d[5], d[6], d[7]};
|
||||
return wm_op_from_acc(h);
|
||||
#else
|
||||
return wm_op_from_f8(d);
|
||||
#endif
|
||||
}
|
||||
|
||||
// element c (0..15) of a row of 16 halves (VALU code; not a WMMA operand)
|
||||
__device__ __forceinline__ half_t op_get(const u8v& v, int c)
|
||||
{
|
||||
return (c & 1) ? hi16(v[c >> 1]) : lo16(v[c >> 1]);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ u8v lds_row16(const half_t* p)
|
||||
{
|
||||
const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
|
||||
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
|
||||
}
|
||||
|
||||
// Global activation traffic (tokens in, rows out): PWIN_NT_LOAD / PWIN_NT_STORE mark it non-temporal so
|
||||
// it does not evict the weight images from the caches. Same data either way.
|
||||
__device__ __forceinline__ u8v gload_row16(const half_t* p)
|
||||
{
|
||||
#ifdef PWIN_NT_LOAD
|
||||
const u4v a = __builtin_nontemporal_load((const u4v*)p), b = __builtin_nontemporal_load((const u4v*)(p + 8));
|
||||
#else
|
||||
const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
|
||||
#endif
|
||||
return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
|
||||
}
|
||||
__device__ __forceinline__ void gstore16(half_t* p, u4v v)
|
||||
{
|
||||
#ifdef PWIN_NT_STORE
|
||||
__builtin_nontemporal_store(v, (u4v*)p);
|
||||
#else
|
||||
*(u4v*)p = v;
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
|
||||
{
|
||||
gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
|
||||
gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
|
||||
}
|
||||
|
||||
// WMMA operands from rows of 16 channels: LDS, global (PWIN_NT_LOAD as above), prep images (32-byte slots)
|
||||
__device__ __forceinline__ op_t op_lds(const half_t* p)
|
||||
{
|
||||
return wm_op_load(p);
|
||||
}
|
||||
__device__ __forceinline__ op_t op_gload(const half_t* p)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
#ifdef PWIN_NT_LOAD
|
||||
return __builtin_nontemporal_load((const u4v*)(p + 8 * wm_half()));
|
||||
#else
|
||||
return *(const u4v*)(p + 8 * wm_half());
|
||||
#endif
|
||||
#else
|
||||
return gload_row16(p);
|
||||
#endif
|
||||
}
|
||||
__device__ __forceinline__ op_t op_img(const u8v* __restrict__ img, int idx)
|
||||
{
|
||||
return wm_op_image(img, idx);
|
||||
}
|
||||
// operand -> slot idx of an operand image in LDS (read back with op_img)
|
||||
__device__ __forceinline__ void op_img_store(u8v* img, int idx, const op_t& v)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
((u4v*)img)[2 * idx + wm_half()] = v;
|
||||
#else
|
||||
if (!wm_half())
|
||||
img[idx] = v;
|
||||
#endif
|
||||
}
|
||||
// operand (16 channels of this lane's token) -> row of `count` (8 or 16) channels in global memory, if ok
|
||||
__device__ __forceinline__ void op_gstore(half_t* dst, const op_t& v, bool ok, int count = 16)
|
||||
{
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
if (ok && (wm_half() == 0 || count >= 16))
|
||||
gstore16(dst + 8 * wm_half(), v);
|
||||
#else
|
||||
if (!wm_half() && ok)
|
||||
{
|
||||
gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
|
||||
if (count >= 16)
|
||||
gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
|
||||
}
|
||||
#endif
|
||||
}
|
||||
// operand -> row of 16 channels in LDS
|
||||
__device__ __forceinline__ void op_store(half_t* dst, const op_t& v)
|
||||
{
|
||||
wm_op_store(dst, v);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void block_sync()
|
||||
{
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
__device__ __forceinline__ int mirror(int v, int n)
|
||||
{
|
||||
v = v < 0 ? -v : v;
|
||||
return v < 2 * n - 2 - v ? v : 2 * n - 2 - v;
|
||||
}
|
||||
|
||||
// ------------------------------------------------ pre-swizzled f16 weight fragments (NVIDIA)
|
||||
// logical B[k][n] of one 16x16 fragment (512 bytes)
|
||||
__device__ __forceinline__ int frag_offset(int k, int n)
|
||||
{
|
||||
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
|
||||
}
|
||||
@@ -1,837 +0,0 @@
|
||||
// Native DltssPaddedWinLayer (DLSS 4, preset K) template: 8x8 windows, H heads x 32, C channels.
|
||||
// One block per window, 4 waves (one 16-token tile each). Semantics: kernels/tools/pwin_model.py (swin_core,
|
||||
// pwin_encoder); weight layout: pwin_model.Layout.
|
||||
#pragma once
|
||||
#include "pwin_common.h"
|
||||
#ifndef PWIN_VGPR
|
||||
#define PWIN_VGPR
|
||||
#endif
|
||||
|
||||
struct PwinParams
|
||||
{
|
||||
int W, H; // 0: token grid
|
||||
const uint8_t* in; // 8
|
||||
const uint8_t* skip; // 16 (decoders)
|
||||
uint8_t* out24; // 24: encoders: merged output; decoders: output
|
||||
uint8_t* out32; // 32: encoders: full-resolution output
|
||||
uint64_t pad40, pad48; // 40
|
||||
int sx, sy; // 56
|
||||
const uint8_t* w; // 64
|
||||
uint8_t rest[176 - 72];
|
||||
};
|
||||
static_assert(sizeof(PwinParams) == 176, "param block");
|
||||
|
||||
template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
|
||||
{
|
||||
static constexpr bool POS = POS_; // position-only attention (trait flag 2): the table is P itself
|
||||
static constexpr int H = H_, C = C_, COUT = COUT_, KT = C / 16, NCH = C / 32, NMLP = C / 8;
|
||||
static constexpr int G1 = 0, QKV = 2 * C, HEAD = 192 * C, BIAS = QKV + HEAD * H, BO = BIAS + 8192 * H,
|
||||
WO = BO + 2 * C, G2 = WO + 64 * C * H, B2 = G2 + 2 * C, W1 = B2 + 2 * C, B1 = W1 + 8 * C * C,
|
||||
W2 = B1 + 8 * C, PM = W2 + 8 * C * C;
|
||||
// patch merge: NVIDIA's warps own COUT / H channels each, allocated as whole 16-column tiles
|
||||
static constexpr int NPW = COUT / H, NPA = (NPW + 15) / 16 * 16, PMT = H * NPA / 16, PMB = PM + 8 * C * NPA * H;
|
||||
// prep image tiles (16 lanes of u8v each)
|
||||
static constexpr int T_QKV = 0; // [h][j][kt KT][nt 2]
|
||||
static constexpr int T_WO = T_QKV + H * 3 * KT * 2; // [ks 2H][nt KT]
|
||||
static constexpr int T_W1 = T_WO + 2 * H * KT; // [c][kt KT][nt 2]
|
||||
static constexpr int T_W2 = T_W1 + NMLP * KT * 2; // [c][ks 2][nt KT]
|
||||
static constexpr int T_PM = T_W2 + NMLP * 2 * KT; // [ks C/4][nt COUT/16]
|
||||
static constexpr int T_END = T_PM + (C / 4) * PMT;
|
||||
static constexpr int NBIAS = H * 4 * 4 * 32; // u4v per lane
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + NBIAS;
|
||||
|
||||
__device__ static int src_of(int tile)
|
||||
{
|
||||
if (tile < T_WO)
|
||||
{
|
||||
const int nt = tile & 1, kt = (tile >> 1) % KT, j = (tile / (2 * KT)) % 3, h = tile / (6 * KT);
|
||||
return QKV + HEAD * h + 2048 * (3 * (kt >> 1) + j) + 512 * (kt & 1) + 1024 * nt;
|
||||
}
|
||||
if (tile < T_W1)
|
||||
{
|
||||
const int t = tile - T_WO, nt = t % KT, ks = t / KT;
|
||||
return WO + 64 * C * (ks >> 1) + 512 * (ks & 1) + 1024 * nt;
|
||||
}
|
||||
if (tile < T_W2)
|
||||
{
|
||||
const int t = tile - T_W1, nt = t & 1, kt = (t >> 1) % KT, c = t / (2 * KT);
|
||||
return W1 + 64 * C * c + 512 * kt + 32 * C * nt;
|
||||
}
|
||||
if (tile < T_PM)
|
||||
{
|
||||
const int t = tile - T_W2, nt = t % KT, ks = (t / KT) & 1, c = t / (2 * KT);
|
||||
return W2 + 64 * C * c + 512 * ks + 1024 * nt;
|
||||
}
|
||||
const int t = tile - T_PM, nt = t % PMT, ks = t / PMT;
|
||||
return PM + 512 * ks + 128 * C * nt; // allocated tiles are contiguous
|
||||
}
|
||||
};
|
||||
|
||||
__device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
|
||||
{
|
||||
return *(const half_t*)(w + byte);
|
||||
}
|
||||
|
||||
// channels 16 nt + wm_acc_row(i) (i = 0..7) of an f16 vector at byte offset `byte` (16 nt already added)
|
||||
__device__ __forceinline__ void dvec8(const uint8_t* w, int byte, half_t out[8])
|
||||
{
|
||||
wm_acc_vec_load((const half_t*)(w + byte), out);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
|
||||
{
|
||||
u8v r;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
r[j] = pack2(wload(w, src + frag_offset(2 * j, n)), wload(w, src + frag_offset(2 * j + 1, n)));
|
||||
return r;
|
||||
}
|
||||
|
||||
template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
|
||||
{
|
||||
const uint8_t* w = p.w;
|
||||
if (idx < L::T_END * 16)
|
||||
{
|
||||
img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
|
||||
return;
|
||||
}
|
||||
const int b = idx - L::T_END * 16;
|
||||
if (b >= L::NBIAS)
|
||||
return;
|
||||
if constexpr (L::POS)
|
||||
{
|
||||
// B operand of P V: lane = query 16 qt + l16, 16 keys 16 kt .. 16 kt + 15 (u8v per lane; 16 lanes)
|
||||
if (b >= L::H * 4 * 4 * 16)
|
||||
return;
|
||||
const int l16 = b & 15, kt = (b >> 4) & 3, qt = (b >> 6) & 3, h = b >> 8, r = l16;
|
||||
u8v v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
{
|
||||
half_t e[2];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 2; ++u)
|
||||
{
|
||||
const int c = 2 * j + u;
|
||||
const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
|
||||
e[u] = wload(w, L::BIAS + 8192 * h + off);
|
||||
}
|
||||
v[j] = pack2(e[0], e[1]);
|
||||
}
|
||||
((u8v*)bias)[b] = v;
|
||||
return;
|
||||
}
|
||||
const int lane = b & 31, kt = (b >> 5) & 3, qt = (b >> 7) & 3, h = b >> 9;
|
||||
const int q = 16 * qt + (lane & 15);
|
||||
half_t v[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int key = 16 * kt + wm_acc_row_h(i, lane >> 4), r = q & 15, c = key & 15;
|
||||
const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
|
||||
v[i] = wload(w, L::BIAS + 8192 * h + off);
|
||||
}
|
||||
bias[b] = (u4v){pack2(v[0], v[1]), pack2(v[2], v[3]), pack2(v[4], v[5]), pack2(v[6], v[7])};
|
||||
}
|
||||
|
||||
template <class L> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
|
||||
__device__ __forceinline__ half_t hmax(half_t a, half_t b) { return a > b ? a : b; }
|
||||
|
||||
__device__ __forceinline__ half_t softmax_e(half_t s)
|
||||
{
|
||||
const half_t hi = __builtin_bit_cast(half_t, (uint16_t)0x4D69), lo = __builtin_bit_cast(half_t, (uint16_t)0xCD69);
|
||||
const half_t a = __builtin_bit_cast(half_t, (uint16_t)0x91D5), c = __builtin_bit_cast(half_t, (uint16_t)0xC50D);
|
||||
const half_t t = hmax(hmin(s, hi), lo);
|
||||
const half_t u = t * t;
|
||||
const half_t v = __builtin_fmaf16(u, a, (half_t)1.0f);
|
||||
const half_t x = __builtin_fmaf16(t, v, c);
|
||||
return (half_t)__builtin_amdgcn_exp2f((float)x);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ half_t gelu(half_t x)
|
||||
{
|
||||
const half_t k1 = (half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
|
||||
const half_t c = hmin(hmax(x, (half_t)-2.0f), (half_t)2.0f);
|
||||
const half_t ac = c < (half_t)0.0f ? -c : c;
|
||||
const half_t u = k1 - k2 * ac;
|
||||
return x * ((half_t)0.5f + c * u);
|
||||
}
|
||||
|
||||
// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
|
||||
__device__ __forceinline__ void gelu8(const acc8v& d, half_t g[8])
|
||||
{
|
||||
const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]};
|
||||
const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
|
||||
const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
|
||||
const hv2 r = x * ((hv2)(half_t)0.5f + c * u);
|
||||
g[2 * j] = r[0];
|
||||
g[2 * j + 1] = r[1];
|
||||
}
|
||||
}
|
||||
|
||||
// store the D^T tile of token m as 16 contiguous channels (lanes of one half-wave write)
|
||||
__device__ __forceinline__ void store_row16(half_t* dst, const u8v& v)
|
||||
{
|
||||
*(u4v*)dst = (u4v){v[0], v[1], v[2], v[3]};
|
||||
*(u4v*)(dst + 8) = (u4v){v[4], v[5], v[6], v[7]};
|
||||
}
|
||||
|
||||
#ifdef PWIN_STAMP
|
||||
// per-block stage timestamps of wave 0 (s_memtime), slot k
|
||||
__device__ uint64_t g_stamp[16384][16];
|
||||
#define STAMP(k) \
|
||||
do \
|
||||
{ \
|
||||
if (threadIdx.z == 0 && lane_id() == 0) \
|
||||
{ \
|
||||
const int bid = blockIdx.y * gridDim.x + blockIdx.x; \
|
||||
if (bid < 16384) \
|
||||
g_stamp[bid][k] = __builtin_readcyclecounter(); \
|
||||
} \
|
||||
} while (0)
|
||||
#else
|
||||
#define STAMP(k) do {} while (0)
|
||||
#endif
|
||||
#ifdef PWIN_DBG
|
||||
__device__ float g_dbg[12][64][160];
|
||||
#define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
|
||||
#endif
|
||||
// Softmax row sum over the 64 keys of this lane's query: this lane's 32 accumulator rows (summed into rs as
|
||||
// the scores are computed), plus the other half's. The layout shim (gfx12 layout on gfx11) instead redoes the
|
||||
// sum in gfx11's order (keys 16 kt + 2 i + h, then the other parity), which the f32 result depends on, so
|
||||
// that its output stays bit-identical to the gfx11 build.
|
||||
#ifdef D4R_WMMA_SHIM
|
||||
constexpr bool kShimRowSum = true;
|
||||
#else
|
||||
constexpr bool kShimRowSum = false;
|
||||
#endif
|
||||
__device__ __forceinline__ float row_sum_total(float rs, const half_t (&e)[4][8])
|
||||
{
|
||||
if constexpr (!kShimRowSum)
|
||||
return rs + __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
|
||||
// this lane: keys 16 kt + 8 h + j; the other half: keys 16 kt + 8 (1 - h) + j
|
||||
const bool hf = wm_half() != 0;
|
||||
float s[2] = {0.0f, 0.0f};
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
half_t lo[8], hi[8];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; j += 2)
|
||||
{
|
||||
const uint32_t own = pack2(e[kt][j], e[kt][j + 1]), other = other_half(own);
|
||||
const uint32_t l = hf ? other : own, u = hf ? own : other;
|
||||
lo[j] = lo16(l), lo[j + 1] = hi16(l), hi[j] = lo16(u), hi[j + 1] = hi16(u);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int par = 0; par < 2; ++par)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const int key = 2 * i + par;
|
||||
s[par] += (float)(key < 8 ? lo[key] : hi[key - 8]);
|
||||
}
|
||||
}
|
||||
return s[0] + s[1];
|
||||
}
|
||||
|
||||
// The Swin core for one window: x0 rows of this wave's 16 tokens are in act (f16, C per token).
|
||||
// On return act holds the block output y. hb: normalised input rows; K/V of one head at a time in kl/vt.
|
||||
template <class L>
|
||||
__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*act)[L::C], half_t (*hb)[L::C], half_t (*kl)[32], half_t (*vt)[64])
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z;
|
||||
const uint8_t* w = p.w;
|
||||
half_t* arow = act[16 * wv + m];
|
||||
half_t* hrow = hb[16 * wv + m];
|
||||
|
||||
STAMP(1);
|
||||
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
|
||||
const half_t ss = l2_sum<C>(arow);
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
u8v hv;
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt); // 16 contiguous gammas
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
#ifdef PWIN_DBG
|
||||
if (DBG_ON && !hf)
|
||||
for (int c = 0; c < C; ++c)
|
||||
g_dbg[0][16 * wv + m][c] = (float)hrow[c];
|
||||
#endif
|
||||
|
||||
if constexpr (L::POS)
|
||||
block_sync(); // V of all 64 tokens is computed by every wave from hb
|
||||
STAMP(2);
|
||||
// ---- attention, one head at a time; the output projection accumulates head by head
|
||||
acc8v acc[KT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::BO + 32 * nt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
acc[nt][i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
for (int h = 0; h < H; ++h)
|
||||
{
|
||||
// position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
|
||||
op_t vtop[2][4];
|
||||
if constexpr (L::POS)
|
||||
{
|
||||
// each wave computes the V^T operands of its own key tile; shared through LDS, one slot per head
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
acc8v d = splat8(0.0f);
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * wv + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
|
||||
op_img_store(&vsh[h][nt][0][0], wv * 16 + m, operand_from_f8(d));
|
||||
}
|
||||
block_sync();
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
|
||||
}
|
||||
op_t pop[4];
|
||||
if constexpr (L::POS)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
|
||||
}
|
||||
else
|
||||
{
|
||||
op_t qop[2];
|
||||
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
|
||||
acc8v dq[6];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = splat8(0.0f);
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 3; ++j)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const acc8v d = dq[2 * j + nt];
|
||||
if (j == 0)
|
||||
qop[nt] = operand_from_f8(d);
|
||||
else if (j == 1)
|
||||
{
|
||||
op_store(&kl[16 * wv + m][16 * nt], operand_from_f8(d));
|
||||
}
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
half_t e[4][8];
|
||||
float rs = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
d[2 * i] = (float)lo16(bi[i]);
|
||||
d[2 * i + 1] = (float)hi16(bi[i]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ds = 0; ds < 2; ++ds)
|
||||
d = mma16(op_lds(&kl[16 * kt + m][16 * ds]), qop[ds], d);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
e[kt][i] = softmax_e((half_t)d[i]);
|
||||
if constexpr (!kShimRowSum)
|
||||
rs += (float)e[kt][i];
|
||||
}
|
||||
}
|
||||
rs = row_sum_total(rs, e);
|
||||
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
half_t pv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
pv[i] = e[kt][i] * rc;
|
||||
pop[kt] = operand_from_dt(pv);
|
||||
}
|
||||
}
|
||||
op_t oop[2];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
acc8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
if constexpr (L::POS)
|
||||
d = mma16(vtop[nt][kt], pop[kt], d);
|
||||
else
|
||||
d = mma16(op_lds(&vt[16 * nt + m][16 * kt]), pop[kt], d);
|
||||
}
|
||||
oop[nt] = operand_from_f8(d);
|
||||
#ifdef PWIN_DBG
|
||||
if (DBG_ON)
|
||||
for (int i = 0; i < 8; ++i)
|
||||
g_dbg[2 + h][16 * wv + m][16 * nt + wm_acc_row(i)] = (float)(half_t)d[i];
|
||||
#endif
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
acc[nt] = mma16(op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m), oop[ks], acc[nt]);
|
||||
if constexpr (!L::POS)
|
||||
block_sync(); // kl / vt are reused by the next head
|
||||
}
|
||||
|
||||
#ifdef PWIN_DBG
|
||||
if (DBG_ON)
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
for (int i = 0; i < 8; ++i)
|
||||
g_dbg[1][16 * wv + m][16 * nt + wm_acc_row(i)] = acc[nt][i];
|
||||
#endif
|
||||
STAMP(3);
|
||||
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
const op_t xv = op_lds(arow + 16 * nt);
|
||||
half_t mv[8], g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const half_t x1 = (half_t)acc[nt][i] + wm_op_acc_elem(xv, i);
|
||||
mv[i] = x1 * g2[i];
|
||||
acc[nt][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(hrow + 16 * nt, operand_from_dt(mv));
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
|
||||
// ---- MLP: C/8 chunks of 32 hidden
|
||||
for (int c = 0; c < L::NMLP; ++c)
|
||||
{
|
||||
op_t gop[2];
|
||||
#pragma unroll
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
acc8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
|
||||
half_t g[8];
|
||||
gelu8(d, g);
|
||||
gop[hn] = operand_from_dt(g);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
acc[nt] = mma16(op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m), gop[ks], acc[nt]);
|
||||
}
|
||||
STAMP(4);
|
||||
// y -> act (rows of this wave)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
}
|
||||
|
||||
// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
|
||||
template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout : PwinLayout<H_, C_, COUT_, POS_>
|
||||
{
|
||||
using B = PwinLayout<H_, C_, COUT_, POS_>;
|
||||
static constexpr int CIN = CIN_, KE = CIN / 16, CB = CIN ? 2 * C_ + 2 * CIN * C_ : 0;
|
||||
static constexpr int T_EMB = B::T_END, T_END = T_EMB + KE * (C_ / 16);
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
|
||||
};
|
||||
|
||||
template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x;
|
||||
if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
|
||||
{
|
||||
const int t = (idx >> 4) - L::T_EMB, nt = t % (L::C / 16), ks = t / (L::C / 16);
|
||||
img[idx] = img_tile(p.w, 2 * L::C + 512 * ks + 512 * L::KE * nt, idx & 15);
|
||||
return;
|
||||
}
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
// core tiles and the bias / P images (indices after T_END are shifted down to the base layout's)
|
||||
const int base_idx = idx < L::T_EMB * 16 ? idx : idx - (L::T_END - L::T_EMB) * 16;
|
||||
if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
|
||||
return;
|
||||
pwin_prep_at<typename L::B>(q, img, bias, base_idx);
|
||||
}
|
||||
|
||||
// encoder: core + full-resolution output + 2x2 patch merge
|
||||
template <class L>
|
||||
__device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
PwinParams p = p0;
|
||||
if constexpr (L::CB != 0)
|
||||
p.w = p0.w + L::CB;
|
||||
constexpr int C = L::C, COUT = L::COUT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int bx = blockIdx.x, by = blockIdx.y;
|
||||
|
||||
STAMP(0);
|
||||
// input rows (mirror padded) -> act (first layer: embedding + ReLU)
|
||||
{
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
if constexpr (L::CIN != 0)
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
op_t xe[L::KE];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
xe[ks] = op_gload(xin + 16 * ks);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
acc8v d;
|
||||
{
|
||||
half_t vv_[8];
|
||||
dvec8(p0.w, 32 * nt, vv_);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv_[i];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
|
||||
op_store(&act[tok][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
pwin_core<L>(p, img, bias, act, hb, kl, vt);
|
||||
|
||||
STAMP(5);
|
||||
// full-resolution output
|
||||
const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
block_sync();
|
||||
|
||||
STAMP(6);
|
||||
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, COUT/16 n-tiles over the 4 waves
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
|
||||
const uint8_t* w = p.w;
|
||||
for (int vt = wv; vt < L::PMT; vt += 4)
|
||||
{
|
||||
acc8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(w, L::PMB + 32 * vt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
STAMP(7);
|
||||
}
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
// allocated column 16 vt + c belongs to original warp g, local column 16 u + c
|
||||
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
|
||||
}
|
||||
}
|
||||
|
||||
#define PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, CIN) \
|
||||
using NAME##_L = PwinEmbLayout<H, C, COUT, POS, CIN>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
pwin_encoder<NAME##_L>(p, g_img, g_bias); \
|
||||
}
|
||||
#define PWIN_ENCODER(NAME, H, C, COUT) PWIN_ENCODER_EMB(NAME, H, C, COUT, false, 0)
|
||||
#define PWIN_ENCODER_POS(NAME, H, C, COUT, POS) PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, 0)
|
||||
|
||||
// bottleneck (dec5): core only, output at the same resolution to +24
|
||||
template <class L>
|
||||
__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[64][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[32][64];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int bx = blockIdx.x, by = blockIdx.y;
|
||||
{
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
pwin_core<L>(p, img, bias, act, hb, kl, vt);
|
||||
const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
|
||||
#define PWIN_KERNEL(NAME, H, C, COUT, BODY) \
|
||||
using NAME##_L = PwinLayout<H, C, COUT>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p) \
|
||||
{ \
|
||||
BODY<NAME##_L>(p, g_img, g_bias); \
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- decoders (patch expand + skip)
|
||||
template <int H_, int C_, int CL_, bool POS_ = false, int NOUT_ = 0, int NOUTA_ = 0>
|
||||
struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
|
||||
{
|
||||
using B = PwinLayout<H_, C_, 16, POS_>;
|
||||
// last layer: output head out = y W_h + b_h (NOUTA computed, NOUT stored); b_h at B::PM, W_h after it
|
||||
static constexpr int NOUT = NOUT_, NOUTA = NOUTA_, HB = B::PM, HW = B::PM + 2 * NOUTA_;
|
||||
static constexpr int CL = CL_, KL = CL / 16, EXPN = 4 * C_ / 16; // expand: K CL, N 4C
|
||||
static constexpr int EXPB = 2 * CL * 4 * C_, CB = EXPB + 8 * C_; // expand bias, core weights base
|
||||
static constexpr int T_EXP = B::T_PM; // expand tiles replace the merge tiles
|
||||
static constexpr int T_HEAD = T_EXP + KL * EXPN;
|
||||
static constexpr int T_END = T_HEAD + B::KT * (NOUTA_ / 16);
|
||||
static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
|
||||
};
|
||||
|
||||
template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
|
||||
{
|
||||
const int idx = blockIdx.x * 128 + threadIdx.x;
|
||||
if (idx >= L::T_EXP * 16 && idx < L::T_HEAD * 16)
|
||||
{
|
||||
const int t = (idx >> 4) - L::T_EXP, nt = t % L::EXPN, ks = t / L::EXPN;
|
||||
img[idx] = img_tile(p.w, 512 * ks + 32 * L::CL * nt, idx & 15);
|
||||
return;
|
||||
}
|
||||
if (idx >= L::T_HEAD * 16 && idx < L::T_END * 16)
|
||||
{
|
||||
const int t = (idx >> 4) - L::T_HEAD, nt = t % (L::NOUTA / 16), ks = t / (L::NOUTA / 16);
|
||||
img[idx] = img_tile(p.w + L::CB, L::HW + 512 * ks + 512 * L::KT * nt, idx & 15);
|
||||
return;
|
||||
}
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
// core tiles lie below T_EXP; bias / P images follow the base layout's (unused) merge tiles
|
||||
pwin_prep_at<typename L::B>(q, img, bias, idx < L::T_EXP * 16 ? idx : L::B::T_END * 16 + (idx - L::T_END * 16));
|
||||
}
|
||||
|
||||
template <class L>
|
||||
__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
|
||||
__shared__ __attribute__((aligned(16))) half_t kl[KVR][32];
|
||||
__shared__ __attribute__((aligned(16))) half_t vt[KVC][64];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
|
||||
// ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
|
||||
{
|
||||
const int W2 = p.W / 2, H2 = p.H / 2;
|
||||
const int ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
|
||||
const int dy = wv >> 1, dx = wv & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
const int gnt = wv * (C / 16) + nt; // expand output column tile
|
||||
acc8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_store(&act[t][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
pwin_core<L>(q, img, bias, act, hb, kl, vt);
|
||||
|
||||
const int tok = 16 * wv + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
|
||||
if constexpr (L::NOUTA != 0)
|
||||
{
|
||||
half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
|
||||
{
|
||||
acc8v d;
|
||||
{
|
||||
|
||||
half_t vv_[8];
|
||||
|
||||
dvec8(q.w, L::HB + 32 * nt, vv_);
|
||||
|
||||
#pragma unroll
|
||||
|
||||
for (int i = 0; i < 8; ++i)
|
||||
|
||||
d[i] = (float)vv_[i];
|
||||
|
||||
}
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < L::KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
|
||||
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
|
||||
}
|
||||
}
|
||||
else if (inb)
|
||||
{
|
||||
half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
|
||||
#define PWIN_DECODER_EX(NAME, H, C, CL, POS, NOUT, NOUTA) \
|
||||
using NAME##_L = PwinDecLayout<H, C, CL, POS, NOUT, NOUTA>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
pwin_decoder<NAME##_L>(p, g_img, g_bias); \
|
||||
}
|
||||
|
||||
#define PWIN_DECODER(NAME, H, C, CL) PWIN_DECODER_EX(NAME, H, C, CL, false, 0, 0)
|
||||
@@ -1,399 +0,0 @@
|
||||
// Position-only-attention layers (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
|
||||
// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
|
||||
#pragma once
|
||||
#include "pwin_layer.h"
|
||||
|
||||
template <class L, int MT>
|
||||
__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*act)[L::C], half_t (*hb)[L::C])
|
||||
{
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
|
||||
static_assert(MT == 1 || MT == 2 || MT == 4, "position token tiles per wave");
|
||||
static_assert(L::POS, "position-only attention layers");
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z;
|
||||
const uint8_t* w = p.w;
|
||||
|
||||
// ---- L2 norm of every tile of this wave -> hb
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const half_t* arow = act[16 * (wv * MT + mi) + m];
|
||||
half_t* hrow = hb[16 * (wv * MT + mi) + m];
|
||||
const half_t ss = l2_sum<C>(arow);
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
for (int kt = (int)hf; kt < KT; kt += 2)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
|
||||
u8v hv;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
}
|
||||
block_sync(); // V of all 64 tokens is computed from hb
|
||||
STAMP(2);
|
||||
|
||||
// ---- attention: O = P V per head (P from the table), output projection accumulated head by head
|
||||
acc8v acc[MT][KT];
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::BO + 32 * nt, vv);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[mi][nt][i] = (float)vv[i];
|
||||
}
|
||||
#ifdef POS_VSHARE
|
||||
// V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
|
||||
// (one slot per head, so a head's slot is never rewritten while it is read)
|
||||
__shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
|
||||
#endif
|
||||
for (int h = 0; h < H; ++h)
|
||||
{
|
||||
op_t vtop[2][4];
|
||||
#ifdef POS_VSHARE
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int kt = wv * MT + mi;
|
||||
acc8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
|
||||
op_img_store(&vsh[h][nt][0][0], kt * 16 + m, operand_from_f8(d));
|
||||
}
|
||||
block_sync();
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
|
||||
#else
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
acc8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < KT; ++ks)
|
||||
d = mma16(op_lds(&hb[16 * kt + m][16 * ks]), op_img(img, (L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m), d);
|
||||
vtop[nt][kt] = operand_from_f8(d);
|
||||
}
|
||||
#endif
|
||||
op_t oop[MT][2];
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int qt = wv * MT + mi;
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
acc8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
d = mma16(vtop[nt][kt], pop[kt], d);
|
||||
oop[mi][nt] = operand_from_f8(d);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t wo = op_img(img, (L::T_WO + (2 * h + ks) * KT + nt) * 16 + m);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
|
||||
STAMP(3);
|
||||
// ---- residual, MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
half_t g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m;
|
||||
const op_t xv = op_lds(&act[tok][16 * nt]);
|
||||
half_t mv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const half_t x1 = (half_t)acc[mi][nt][i] + wm_op_acc_elem(xv, i);
|
||||
mv[i] = x1 * g2[i];
|
||||
acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
|
||||
STAMP(4);
|
||||
// ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
|
||||
for (int c = 0; c < L::NMLP; ++c)
|
||||
{
|
||||
op_t gop[MT][2];
|
||||
#pragma unroll
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
half_t bb[8];
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
|
||||
acc8v d[MT];
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[mi][i] = (float)bb[i];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t w1 = op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
d[mi] = mma16(w1, op_lds(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
half_t g[8];
|
||||
gelu8(d[mi], g);
|
||||
gop[mi][hn] = operand_from_dt(g);
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < 2; ++ks)
|
||||
{
|
||||
const op_t w2 = op_img(img, (L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
|
||||
}
|
||||
}
|
||||
STAMP(5);
|
||||
// y -> act
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < KT; ++nt)
|
||||
{
|
||||
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
|
||||
}
|
||||
block_sync();
|
||||
}
|
||||
|
||||
// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
|
||||
template <class L, int MT>
|
||||
__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, NW = 4 / MT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
PwinParams p = p0;
|
||||
if constexpr (L::CB != 0)
|
||||
p.w = p0.w + L::CB;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
STAMP(0);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m, ty = tok >> 3, tx = tok & 7;
|
||||
const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
|
||||
if constexpr (L::CIN != 0)
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
|
||||
op_t xe[L::KE];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
xe[ks] = op_gload(xin + 16 * ks);
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p0.w, 32 * nt, vv);
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KE; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EMB + ks * (C / 16) + nt) * 16 + m), xe[ks], d);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = hmax((half_t)d[i], (half_t)0.0f);
|
||||
op_store(&act[tok][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
|
||||
}
|
||||
}
|
||||
__builtin_amdgcn_wave_barrier();
|
||||
STAMP(1);
|
||||
pos_core<L, MT>(p, img, bias, act, hb);
|
||||
STAMP(6);
|
||||
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m;
|
||||
const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
|
||||
for (int kt = (int)hf; kt < L::KT; kt += 2)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
STAMP(7);
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
|
||||
for (int vt = wv; vt < L::PMT; vt += NW)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::PMB + 32 * vt, vv);
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
|
||||
}
|
||||
STAMP(8);
|
||||
}
|
||||
|
||||
// dec0: patch expand + skip, core, output head
|
||||
template <class L, int MT>
|
||||
__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, NW = 4 / MT;
|
||||
__shared__ __attribute__((aligned(16))) half_t act[64][C];
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
|
||||
{
|
||||
const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
|
||||
op_t xls[L::KL];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
xls[ks] = op_gload(xl + 16 * ks);
|
||||
for (int q = wv; q < 4; q += NW)
|
||||
{
|
||||
const int dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < C / 16; ++nt)
|
||||
{
|
||||
const int gnt = q * (C / 16) + nt;
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv);
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
#pragma unroll
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), xls[ks], d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_store(&act[t][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PwinParams q = p;
|
||||
q.w = p.w + L::CB;
|
||||
pos_core<L, MT>(q, img, bias, act, hb);
|
||||
#pragma unroll
|
||||
for (int mi = 0; mi < MT; ++mi)
|
||||
{
|
||||
const int tok = 16 * (wv * MT + mi) + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
|
||||
const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
|
||||
half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(q.w, L::HB + 32 * nt, vv);
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < L::KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m), op_lds(&act[tok][16 * kt]), d);
|
||||
op_gstore(hout + 16 * nt, operand_from_f8(d), inb, L::NOUT - 16 * nt);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#define POS_ENCODER(NAME, H, C, COUT, CIN, MT) \
|
||||
using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_emb_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
pos_encoder<NAME##_L, MT>(p, g_img, g_bias); \
|
||||
}
|
||||
|
||||
#define POS_DECODER(NAME, H, C, CL, NOUT, NOUTA, MT) \
|
||||
using NAME##_L = PwinDecLayout<H, C, CL, true, NOUT, NOUTA>; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
pwin_dec_prep<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
pos_decoder<NAME##_L, MT>(p, g_img, g_bias); \
|
||||
}
|
||||
@@ -1,429 +0,0 @@
|
||||
// Wide variant of the DltssPaddedWinLayer template for the deep, latency-bound levels (few windows):
|
||||
// 4 * NG waves per window. Wave wv = threadIdx.z works on token tile t = wv & 3 and channel group
|
||||
// g = wv >> 2: heads g, g + NG, ... (one head per group and round), output tiles nt = g + NG j of the
|
||||
// output projection and MLP, and MLP hidden chunks c = NG rr + g. Every output element goes through
|
||||
// the same f16 operations in the same order as pwin_layer.h (bit-identical results).
|
||||
//
|
||||
// Shared memory: hb (64 x C) = normalised input, then MLP input, then the block output y; R = union of
|
||||
// x0 (64 x C), the K / V images of the NG heads of an attention round, their outputs O, and the double
|
||||
// buffered MLP hidden chunks of a round.
|
||||
#pragma once
|
||||
#include "pwin_layer.h"
|
||||
|
||||
template <class L, int NG> struct WideDims
|
||||
{
|
||||
// NA: heads per attention round (their K / V images share R)
|
||||
static constexpr int NA = NG < 4 ? NG : 4;
|
||||
static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
|
||||
static constexpr int X0_BYTES = 64 * C * 2, KV_BYTES = NA * 2 * 64 * 32 * 2;
|
||||
// MLP hidden chunks: one per group and round, double-buffered while that fits
|
||||
static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * 32 * 2;
|
||||
static constexpr int R_BYTES = X0_BYTES > KV_BYTES ? (X0_BYTES > HID_BYTES ? X0_BYTES : HID_BYTES)
|
||||
: (KV_BYTES > HID_BYTES ? KV_BYTES : HID_BYTES);
|
||||
static_assert(!L::POS, "position-only attention layers use pwin_pos.h");
|
||||
};
|
||||
|
||||
// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
|
||||
template <class L, int NG>
|
||||
__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
|
||||
half_t (*hb)[L::C], uint8_t* R)
|
||||
{
|
||||
using D = WideDims<L, NG>;
|
||||
constexpr int C = L::C, KT = L::KT, H = L::H, NTL = D::NTL;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const uint8_t* w = p.w;
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
const half_t* arow = act[tok];
|
||||
half_t* hrow = hb[tok];
|
||||
|
||||
// ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
|
||||
const half_t ss = l2_sum<C>(arow);
|
||||
const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
|
||||
for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
|
||||
{
|
||||
const u8v xv = lds_row16(arow + 16 * kt);
|
||||
u8v hv;
|
||||
const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 8; ++j)
|
||||
hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
|
||||
op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
|
||||
store_row16(hrow + 16 * kt, hv);
|
||||
}
|
||||
// residual input of the owned output tiles (D^T order: channel 16 nt + wm_acc_row(i))
|
||||
half_t x0[NTL][8];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
const op_t xv = op_lds(arow + 16 * nt);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
x0[j][i] = wm_op_acc_elem(xv, i);
|
||||
}
|
||||
}
|
||||
block_sync(); // hb complete, R free
|
||||
|
||||
// ---- attention: NG heads per round; the output projection accumulates head by head
|
||||
acc8v acc[NTL];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::BO + 32 * nt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[j][i] = (float)vv[i];
|
||||
}
|
||||
}
|
||||
half_t(*kl)[64][32] = (half_t(*)[64][32])R; // [group][token][dim]
|
||||
constexpr int NA = D::NA;
|
||||
half_t(*vt)[32][64] = (half_t(*)[32][64])(R + NA * 64 * 32 * 2); // [group][dim][token]
|
||||
half_t(*ob)[64][32] = (half_t(*)[64][32])R; // [group][token][dim] (after K / V)
|
||||
for (int r = 0; r < D::ROUNDS; ++r)
|
||||
{
|
||||
const int h = g < NA ? r * NA + g : H;
|
||||
op_t qop[2], oop[2];
|
||||
if (h < H)
|
||||
{
|
||||
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
|
||||
acc8v dq[6];
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = splat8(0.0f);
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
{
|
||||
const op_t row = op_lds(hrow + 16 * kt);
|
||||
#pragma unroll
|
||||
for (int u = 0; u < 6; ++u)
|
||||
dq[u] = mma16(op_img(img, (L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m), row, dq[u]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 3; ++j)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
const acc8v d = dq[2 * j + nt];
|
||||
if (j == 0)
|
||||
qop[nt] = operand_from_f8(d);
|
||||
else if (j == 1)
|
||||
{
|
||||
op_store(&kl[g][tok][16 * nt], operand_from_f8(d));
|
||||
}
|
||||
else
|
||||
{
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
vt[g][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
if (h < H)
|
||||
{
|
||||
half_t e[4][8];
|
||||
float rs = 0.0f;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
|
||||
acc8v d;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
{
|
||||
d[2 * i] = (float)lo16(bi[i]);
|
||||
d[2 * i + 1] = (float)hi16(bi[i]);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int ds = 0; ds < 2; ++ds)
|
||||
d = mma16(op_lds(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
e[kt][i] = softmax_e((half_t)d[i]);
|
||||
if constexpr (!kShimRowSum)
|
||||
rs += (float)e[kt][i];
|
||||
}
|
||||
}
|
||||
rs = row_sum_total(rs, e);
|
||||
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
|
||||
op_t pop[4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
{
|
||||
half_t pv[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
pv[i] = e[kt][i] * rc;
|
||||
pop[kt] = operand_from_dt(pv);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
{
|
||||
acc8v d = splat8(0.0f);
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < 4; ++kt)
|
||||
d = mma16(op_lds(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
|
||||
oop[nt] = operand_from_f8(d);
|
||||
}
|
||||
}
|
||||
block_sync(); // K / V of the round consumed
|
||||
if (h < H)
|
||||
{
|
||||
op_store(&ob[g][tok][0], oop[0]);
|
||||
op_store(&ob[g][tok][16], oop[1]);
|
||||
}
|
||||
block_sync();
|
||||
for (int jh = 0; jh < NA; ++jh)
|
||||
{
|
||||
const int hh = r * NA + jh;
|
||||
if (hh >= H)
|
||||
break;
|
||||
const op_t o0 = op_lds(&ob[jh][tok][0]), o1 = op_lds(&ob[jh][tok][16]);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh) * KT + nt) * 16 + m), o0, acc[j]);
|
||||
acc[j] = mma16(op_img(img, (L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m), o1, acc[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (r + 1 < D::ROUNDS)
|
||||
block_sync(); // ob is overwritten by the next round's K / V
|
||||
}
|
||||
|
||||
// ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
half_t mv[8], g2[8], b2[8];
|
||||
dvec8(w, L::G2 + 32 * nt, g2);
|
||||
dvec8(w, L::B2 + 32 * nt, b2);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
{
|
||||
const half_t x1 = (half_t)acc[j][i] + x0[j][i];
|
||||
mv[i] = x1 * g2[i];
|
||||
acc[j][i] = (float)(half_t)(x1 + b2[i]);
|
||||
}
|
||||
op_store(hrow + 16 * nt, operand_from_dt(mv));
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
|
||||
// ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
|
||||
half_t(*hid)[NG][64][32] = (half_t(*)[NG][64][32])R;
|
||||
for (int rr = 0; rr < D::MROUNDS; ++rr)
|
||||
{
|
||||
const int c = rr * NG + g;
|
||||
#pragma unroll
|
||||
for (int hn = 0; hn < 2; ++hn)
|
||||
{
|
||||
if (c >= L::NMLP)
|
||||
break;
|
||||
acc8v d;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::B1 + 64 * c + 32 * hn, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
for (int kt = 0; kt < KT; ++kt)
|
||||
d = mma16(op_img(img, (L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m), op_lds(hrow + 16 * kt), d);
|
||||
half_t gg[8];
|
||||
gelu8(d, gg);
|
||||
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
|
||||
}
|
||||
block_sync();
|
||||
for (int jc = 0; jc < NG; ++jc)
|
||||
{
|
||||
const int cc = rr * NG + jc;
|
||||
if (cc >= L::NMLP)
|
||||
break;
|
||||
const op_t h0 = op_lds(&hid[rr % D::NB][jc][tok][0]), h1 = op_lds(&hid[rr % D::NB][jc][tok][16]);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2) * KT + nt) * 16 + m), h0, acc[j]);
|
||||
acc[j] = mma16(op_img(img, (L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m), h1, acc[j]);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (D::NB == 1 && rr + 1 < D::MROUNDS)
|
||||
block_sync();
|
||||
}
|
||||
// y -> hb (the last reads of hb were the W1 operands before the final barrier)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < NTL; ++j)
|
||||
{
|
||||
const int nt = g + NG * j;
|
||||
if (nt < KT)
|
||||
{
|
||||
op_store(hrow + 16 * nt, operand_from_f8(acc[j]));
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
}
|
||||
|
||||
// x0 rows of a window (mirror padded) -> R
|
||||
template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
|
||||
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
|
||||
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
|
||||
store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
|
||||
}
|
||||
|
||||
// y rows of the window's own tokens -> dst (C per token)
|
||||
template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst)
|
||||
{
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
|
||||
const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
|
||||
if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
|
||||
{
|
||||
half_t* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * L::C;
|
||||
for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
|
||||
gstore_row16(yout + 16 * kt, lds_row16(&hb[tok][16 * kt]));
|
||||
}
|
||||
}
|
||||
|
||||
template <class L, int NG>
|
||||
__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C, COUT = L::COUT;
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
wide_load_x0<L, NG>(p, R);
|
||||
block_sync();
|
||||
wide_core<L, NG>(p, img, bias, hb, R);
|
||||
wide_store_y<L, NG>(p, hb, p.out32);
|
||||
|
||||
// patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z;
|
||||
const int my = m >> 2, mx = m & 3;
|
||||
const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
|
||||
const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
|
||||
half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
|
||||
const uint8_t* w = p.w;
|
||||
for (int vt = wv; vt < L::PMT; vt += 4 * NG)
|
||||
{
|
||||
acc8v d;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(w, L::PMB + 32 * vt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
for (int ks = 0; ks < C / 4; ++ks)
|
||||
{
|
||||
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
|
||||
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
|
||||
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&hb[src][16 * (ks % L::KT)]), d);
|
||||
}
|
||||
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
|
||||
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
|
||||
}
|
||||
}
|
||||
|
||||
// bottleneck (dec5): core only, output at the same resolution to +24
|
||||
template <class L, int NG>
|
||||
__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][L::C];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
wide_load_x0<L, NG>(p, R);
|
||||
block_sync();
|
||||
wide_core<L, NG>(p, img, bias, hb, R);
|
||||
wide_store_y<L, NG>(p, hb, p.out24);
|
||||
}
|
||||
|
||||
template <class L, int NG>
|
||||
__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
|
||||
{
|
||||
constexpr int C = L::C;
|
||||
__shared__ __attribute__((aligned(16))) half_t hb[64][C];
|
||||
__shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
|
||||
half_t(*act)[C] = (half_t(*)[C])R;
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const int wv = threadIdx.z, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
|
||||
|
||||
// ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
|
||||
{
|
||||
const int W2 = p.W / 2, H2 = p.H / 2;
|
||||
const int ly = m >> 2, lx = m & 3;
|
||||
const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
|
||||
const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
|
||||
const int dy = q >> 1, dx = q & 1, tt = 8 * (2 * ly + dy) + 2 * lx + dx;
|
||||
const int SX = mirror(8 * bx - p.sx + (tt & 7), p.W), SY = mirror(8 * by - p.sy + (tt >> 3), p.H);
|
||||
const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
|
||||
for (int nt = g; nt < C / 16; nt += NG)
|
||||
{
|
||||
const int gnt = q * (C / 16) + nt;
|
||||
acc8v d;
|
||||
{
|
||||
half_t vv[8];
|
||||
dvec8(p.w, L::EXPB + 32 * gnt, vv);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
d[i] = (float)vv[i];
|
||||
}
|
||||
for (int ks = 0; ks < L::KL; ++ks)
|
||||
d = mma16(op_img(img, (L::T_EXP + ks * L::EXPN + gnt) * 16 + m), op_gload(xl + 16 * ks), d);
|
||||
const op_t sv = op_gload(sk + 16 * nt);
|
||||
half_t xo[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
|
||||
op_store(&act[tt][16 * nt], operand_from_dt(xo));
|
||||
}
|
||||
}
|
||||
block_sync();
|
||||
PwinParams qp = p;
|
||||
qp.w = p.w + L::CB;
|
||||
wide_core<L, NG>(qp, img, bias, hb, R);
|
||||
static_assert(L::NOUTA == 0, "output head: use pwin_layer.h");
|
||||
wide_store_y<L, NG>(p, hb, p.out24);
|
||||
}
|
||||
|
||||
#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...) \
|
||||
using NAME##_L = __VA_ARGS__; \
|
||||
__device__ u8v g_img[NAME##_L::T_END * 16]; \
|
||||
__device__ u4v g_bias[NAME##_L::NBIAS]; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128; \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 * (NG); \
|
||||
extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */ \
|
||||
extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p) \
|
||||
{ \
|
||||
PREP<NAME##_L>(p, g_img, g_bias); \
|
||||
} \
|
||||
extern "C" __global__ void __launch_bounds__(128 * (NG)) PWIN_VGPR NAME(PwinParams p) \
|
||||
{ \
|
||||
BODY<NAME##_L, NG>(p, g_img, g_bias); \
|
||||
}
|
||||
|
||||
#define WIDE_ENCODER(NAME, H, C, COUT, NG) WIDE_KERNEL(NAME, pwin_prep, wide_encoder, NG, PwinLayout<H, C, COUT>)
|
||||
#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout<H, C, 16>)
|
||||
#define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout<H, C, CL>)
|
||||
Reference in new issue
Block a user