wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge upstream RDNA4 layouts and validate native Windows K integration

Browse Source
This commit is contained in:
Codex committed 2026-09-30 14:54:38 +03:00
2 parents 2dea066d6a 4334a0c3c4
commit 64472e5a61
47 files changed
+3250 -576
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/k/dltss_pwin_dec0_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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_pos.h"
POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)
kernels/k/dltss_pwin_dec1_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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)
kernels/k/dltss_pwin_dec2_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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)
kernels/k/dltss_pwin_enc0_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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_pos.h"
POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)
kernels/k/dltss_pwin_enc1_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)
kernels/k/dltss_pwin_enc2_layer.hip
+2
View File
@@ -1,3 +1,5 @@
// 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)
kernels/k/pwin_common.h
+68 -74
View File
@@ -1,10 +1,11 @@
// Native RDNA3 (gfx11 wave32 WMMA) kernels for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
// 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, 16 K values)
// and the activations as the B operand (lane l: token l & 15, 16 K values). The f32 result then has
// lane l = token l & 15 and VGPR i = output channel 2i + (l >> 4) of the 16-channel tile ("D^T").
// 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"
@@ -12,7 +13,7 @@
#include <hip/hip_runtime.h>
#endif
#include <stdint.h>
#include "../common/wmma_layout.h"
#pragma clang fp contract(off)
typedef _Float16 half_t;
@@ -21,23 +22,19 @@ typedef _Float16 h16 __attribute__((ext_vector_type(16)));
typedef float f8v __attribute__((ext_vector_type(8)));
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
#include "../common/wmma_backend.h"
__device__ __forceinline__ uint32_t lane_id()
{
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
}
__device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
{
return d4r_wmma_legacy_layout(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
}
typedef wm_op op_t; // WMMA operand: u8v (gfx11) or u4v (gfx12)
// 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__ f8v mma16(u8v a, u8v b, f8v c)
__device__ __forceinline__ f8v mma16(const op_t& a, const op_t& b, f8v c)
{
f8v d = wmma(a, b, c);
f8v d = wm_mma(a, b, c);
#ifndef PWIN_F32ACC
#pragma unroll
for (int i = 0; i < 8; ++i)
@@ -70,50 +67,19 @@ __device__ __forceinline__ half_t hi16(uint32_t v)
return __builtin_bit_cast(hv2, v)[1];
}
// D^T f16 values (own[i] = channel 2i + hf of token l & 15) -> B operand of the same token (all 16
// channels in order): the other parity comes from lane l ^ 16
__device__ __forceinline__ u8v operand_from_dt(const half_t own[8])
// 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])
{
const uint32_t hf = lane_id() >> 4;
uint32_t mine[4], theirs[4];
#pragma unroll
for (int j = 0; j < 4; ++j)
{
mine[j] = pack2(own[2 * j], own[2 * j + 1]);
theirs[j] = other_half(mine[j]);
}
#ifndef PWIN_NO_PERM
// v_perm_b32(theirs, mine): bytes 0-3 = mine, 4-7 = theirs. Even j: the low halves, odd j: the high
// halves, in the order (mine, theirs) for hf = 0 and (theirs, mine) for hf = 1.
const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
return r;
#else
u8v r;
#pragma unroll
for (int j = 0; j < 8; ++j)
{
const half_t a = (j & 1) ? hi16(mine[j >> 1]) : lo16(mine[j >> 1]);
const half_t b = (j & 1) ? hi16(theirs[j >> 1]) : lo16(theirs[j >> 1]);
r[j] = hf ? pack2(b, a) : pack2(a, b);
}
return r;
#endif
return wm_op_from_acc(own);
}
__device__ __forceinline__ u8v operand_from_f8(f8v d)
__device__ __forceinline__ op_t operand_from_f8(f8v d)
{
half_t h[8];
#pragma unroll
for (int i = 0; i < 8; ++i)
h[i] = (half_t)d[i];
return operand_from_dt(h);
return wm_op_from_f8(d);
}
// element c (0..15) of a B operand (16 halves)
// 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]);
@@ -150,6 +116,58 @@ __device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
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();
@@ -167,27 +185,3 @@ __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);
}
// A-operand image of a K x N weight: tile (kt, nt) = 16 lanes x 8 dwords (lane = output channel)
struct WDesc
{
int src; // byte offset of the fragments in the weight buffer
int kt_stride; // bytes between k tiles
int nt_stride; // bytes between n tiles
int KT, NT; // tiles
int dst; // u8v index of the image in the prep buffer
};
// weight operand tile `tile` of an A-operand image (16 lanes x 32 bytes): lanes l and l + 16 need the same
// row, so each half-wave loads one 16-byte half and the halves are swapped across (halves L2 traffic)
__device__ __forceinline__ u8v wtile(const u8v* __restrict__ img, int tile)
{
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
const u4v mine = ((const u4v*)&img[tile * 16 + m])[hf];
u4v other;
#pragma unroll
for (int j = 0; j < 4; ++j)
other[j] = other_half(mine[j]);
const u4v lo = hf ? other : mine, hi = hf ? mine : other;
return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
}
kernels/k/pwin_layer.h
+90 -80
View File
@@ -72,14 +72,10 @@ __device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
return *(const half_t*)(w + byte);
}
// channels 16 nt + 2i + hf (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, uint32_t hf, half_t out[8])
// 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])
{
const u4v a = *(const u4v*)(w + byte), b = *(const u4v*)(w + byte + 16);
const u8v v = (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
#pragma unroll
for (int i = 0; i < 8; ++i)
out[i] = hf ? hi16(v[i]) : lo16(v[i]);
wm_acc_vec_load((const half_t*)(w + byte), out);
}
__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
@@ -131,7 +127,7 @@ template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u
#pragma unroll
for (int i = 0; i < 8; ++i)
{
const int key = 16 * kt + 2 * i + (lane >> 4), r = q & 15, c = key & 15;
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);
}
@@ -209,6 +205,45 @@ __device__ uint64_t g_stamp[16384][16];
__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>
@@ -265,7 +300,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
half_t vv_[8];
dvec8(w, L::BO + 32 * nt, hf, vv_);
dvec8(w, L::BO + 32 * nt, vv_);
#pragma unroll
@@ -277,7 +312,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
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)
u8v vtop[2][4];
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
@@ -287,28 +322,26 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
{
f8v d = splat8(0.0f);
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * wv + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
const u8v vo = operand_from_f8(d);
if (!hf)
vsh[h][nt][wv][m] = vo;
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] = vsh[h][nt][kt][m];
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
}
u8v pop[4];
op_t pop[4];
if constexpr (L::POS)
{
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = ((const u8v*)bias)[((h * 4 + wv) * 4 + kt) * 16 + m];
pop[kt] = op_img((const u8v*)bias, ((h * 4 + wv) * 4 + kt) * 16 + m);
}
else
{
u8v qop[2];
op_t qop[2];
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
f8v dq[6];
#pragma unroll
@@ -316,10 +349,10 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const u8v row = lds_row16(hrow + 16 * kt);
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[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)
@@ -331,15 +364,13 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
const u8v kop = operand_from_f8(d);
if (!hf)
store_row16(&kl[16 * wv + m][16 * nt], kop);
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 + 2 * i + hf][16 * wv + m] = (half_t)d[i];
vt[16 * nt + wm_acc_row(i)][16 * wv + m] = (half_t)d[i];
}
}
block_sync();
@@ -358,15 +389,16 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(lds_row16(&kl[16 * kt + m][16 * ds]), qop[ds], d);
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]);
rs += (float)e[kt][i];
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
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)
@@ -378,7 +410,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
pop[kt] = operand_from_dt(pv);
}
}
u8v oop[2];
op_t oop[2];
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
@@ -389,20 +421,20 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
if constexpr (L::POS)
d = mma16(vtop[nt][kt], pop[kt], d);
else
d = mma16(lds_row16(&vt[16 * nt + m][16 * kt]), pop[kt], d);
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 + 2 * i + hf] = (float)(half_t)d[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(img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m], oop[ks], acc[nt]);
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
}
@@ -411,34 +443,32 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
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 + 2 * i + hf] = acc[nt][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 u8v xv = lds_row16(arow + 16 * nt);
const op_t xv = op_lds(arow + 16 * nt);
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
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] + op_get(xv, 2 * i + hf);
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]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(hrow + 16 * nt, mo);
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)
{
u8v gop[2];
op_t gop[2];
#pragma unroll
for (int hn = 0; hn < 2; ++hn)
{
@@ -447,7 +477,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
half_t vv_[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv_);
dvec8(w, L::B1 + 64 * c + 32 * hn, vv_);
#pragma unroll
@@ -457,7 +487,7 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
for (int kt = 0; kt < KT; ++kt)
d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
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);
@@ -466,16 +496,14 @@ __device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int nt = 0; nt < KT; ++nt)
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
acc[nt] = mma16(img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m], gop[ks], acc[nt]);
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)
{
const u8v yo = operand_from_f8(acc[nt]);
if (!hf)
store_row16(arow + 16 * nt, yo);
op_store(arow + 16 * nt, operand_from_f8(acc[nt]));
}
__builtin_amdgcn_wave_barrier();
}
@@ -530,31 +558,29 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
u8v xe[L::KE];
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
xe[ks] = gload_row16(xin + 16 * ks);
xe[ks] = op_gload(xin + 16 * ks);
#pragma unroll
for (int nt = 0; nt < C / 16; ++nt)
{
f8v d;
{
half_t vv_[8];
dvec8(p0.w, 32 * nt, hf, vv_);
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(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
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);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tok][16 * nt], xv);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
@@ -592,7 +618,7 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
half_t vv_[8];
dvec8(w, L::PMB + 32 * vt, hf, vv_);
dvec8(w, L::PMB + 32 * vt, vv_);
#pragma unroll
@@ -606,18 +632,11 @@ __device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img,
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(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;
half_t* dst = mout + L::NPW * g + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
}
@@ -749,7 +768,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
half_t vv_[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv_);
dvec8(p.w, L::EXPB + 32 * gnt, vv_);
#pragma unroll
@@ -759,15 +778,13 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
}
for (int ks = 0; ks < L::KL; ++ks)
d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
const u8v sv = gload_row16(sk + 16 * nt);
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] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[t][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
block_sync();
@@ -788,7 +805,7 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
half_t vv_[8];
dvec8(q.w, L::HB + 32 * nt, hf, vv_);
dvec8(q.w, L::HB + 32 * nt, vv_);
#pragma unroll
@@ -799,15 +816,8 @@ __device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, c
}
#pragma unroll
for (int kt = 0; kt < L::KT; ++kt)
d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
const u8v ho = operand_from_f8(d);
const int valid = L::NOUT - 16 * nt;
if (!hf && inb)
{
gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
if (valid >= 16)
gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
}
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)
kernels/k/pwin_pos.h
+39 -63
View File
@@ -53,7 +53,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int nt = 0; nt < KT; ++nt)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, hf, vv);
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
#pragma unroll
@@ -67,7 +67,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#endif
for (int h = 0; h < H; ++h)
{
u8v vtop[2][4];
op_t vtop[2][4];
#ifdef POS_VSHARE
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
@@ -78,17 +78,15 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
const u8v vo = operand_from_f8(d);
if (!hf)
vsh[h][nt][kt][m] = vo;
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] = vsh[h][nt][kt][m];
vtop[nt][kt] = op_img(&vsh[h][nt][0][0], kt * 16 + m);
#else
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
@@ -98,19 +96,19 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
f8v d = splat8(0.0f);
#pragma unroll
for (int ks = 0; ks < KT; ++ks)
d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
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
u8v oop[MT][2];
op_t oop[MT][2];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
{
const int qt = wv * MT + mi;
u8v pop[4];
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
pop[kt] = ((const u8v*)bias)[((h * 4 + qt) * 4 + kt) * 16 + m];
pop[kt] = op_img((const u8v*)bias, ((h * 4 + qt) * 4 + kt) * 16 + m);
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
{
@@ -126,7 +124,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const u8v wo = img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m];
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]);
@@ -139,24 +137,22 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
for (int nt = 0; nt < KT; ++nt)
{
half_t g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
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 u8v xv = lds_row16(&act[tok][16 * nt]);
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] + op_get(xv, 2 * i + hf);
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]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(&hb[tok][16 * nt], mo);
op_store(&hb[tok][16 * nt], operand_from_dt(mv));
}
}
__builtin_amdgcn_wave_barrier();
@@ -165,12 +161,12 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
// ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
for (int c = 0; c < L::NMLP; ++c)
{
u8v gop[MT][2];
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, hf, bb);
dvec8(w, L::B1 + 64 * c + 32 * hn, bb);
f8v d[MT];
#pragma unroll
for (int mi = 0; mi < MT; ++mi)
@@ -180,10 +176,10 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int kt = 0; kt < KT; ++kt)
{
const u8v w1 = img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m];
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, lds_row16(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[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)
@@ -198,7 +194,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int ks = 0; ks < 2; ++ks)
{
const u8v w2 = img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m];
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]);
@@ -211,9 +207,7 @@ __device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const
#pragma unroll
for (int nt = 0; nt < KT; ++nt)
{
const u8v yo = operand_from_f8(acc[mi][nt]);
if (!hf)
store_row16(&act[16 * (wv * MT + mi) + m][16 * nt], yo);
op_store(&act[16 * (wv * MT + mi) + m][16 * nt], operand_from_f8(acc[mi][nt]));
}
block_sync();
}
@@ -239,29 +233,27 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
if constexpr (L::CIN != 0)
{
const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
u8v xe[L::KE];
op_t xe[L::KE];
#pragma unroll
for (int ks = 0; ks < L::KE; ++ks)
xe[ks] = gload_row16(xin + 16 * 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, hf, vv);
dvec8(p0.w, 32 * nt, vv);
f8v 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(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
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);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tok][16 * nt], xv);
op_store(&act[tok][16 * nt], operand_from_dt(xo));
}
}
else
@@ -296,7 +288,7 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
for (int vt = wv; vt < L::PMT; vt += NW)
{
half_t vv[8];
dvec8(p.w, L::PMB + 32 * vt, hf, vv);
dvec8(p.w, L::PMB + 32 * vt, vv);
f8v d;
#pragma unroll
for (int i = 0; i < 8; ++i)
@@ -305,17 +297,10 @@ __device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, c
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&act[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(d);
const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
half_t* dst = mout + L::NPW * g + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * g + 16 * u, operand_from_f8(d), minb, valid);
}
STAMP(8);
}
@@ -333,10 +318,10 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
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;
u8v xls[L::KL];
op_t xls[L::KL];
#pragma unroll
for (int ks = 0; ks < L::KL; ++ks)
xls[ks] = gload_row16(xl + 16 * 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;
@@ -347,22 +332,20 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
{
const int gnt = q * (C / 16) + nt;
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
dvec8(p.w, L::EXPB + 32 * gnt, vv);
f8v 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(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], xls[ks], d);
const u8v sv = gload_row16(sk + 16 * nt);
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] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[t][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[t][16 * nt], operand_from_dt(xo));
}
}
}
@@ -380,22 +363,15 @@ __device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, co
for (int nt = 0; nt < L::NOUTA / 16; ++nt)
{
half_t vv[8];
dvec8(q.w, L::HB + 32 * nt, hf, vv);
dvec8(q.w, L::HB + 32 * nt, vv);
f8v 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(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
const u8v ho = operand_from_f8(d);
const int valid = L::NOUT - 16 * nt;
if (!hf && inb)
{
gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
if (valid >= 16)
gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
}
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);
}
}
}
kernels/k/pwin_wide.h
+39 -55
View File
@@ -61,7 +61,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
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 + 2 i + hf)
// 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)
@@ -69,10 +69,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
const u8v xv = lds_row16(arow + 16 * nt);
const op_t xv = op_lds(arow + 16 * nt);
#pragma unroll
for (int i = 0; i < 8; ++i)
x0[j][i] = op_get(xv, 2 * i + hf);
x0[j][i] = wm_op_acc_elem(xv, i);
}
}
block_sync(); // hb complete, R free
@@ -86,7 +86,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (nt < KT)
{
half_t vv[8];
dvec8(w, L::BO + 32 * nt, hf, vv);
dvec8(w, L::BO + 32 * nt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
acc[j][i] = (float)vv[i];
@@ -99,7 +99,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
for (int r = 0; r < D::ROUNDS; ++r)
{
const int h = g < NA ? r * NA + g : H;
u8v qop[2], oop[2];
op_t qop[2], oop[2];
if (h < H)
{
// six independent chains (Q, K, V x 2 n-tiles), each over kt in order
@@ -109,10 +109,10 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
dq[u] = splat8(0.0f);
for (int kt = 0; kt < KT; ++kt)
{
const u8v row = lds_row16(hrow + 16 * kt);
const op_t row = op_lds(hrow + 16 * kt);
#pragma unroll
for (int u = 0; u < 6; ++u)
dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[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)
@@ -124,15 +124,13 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
qop[nt] = operand_from_f8(d);
else if (j == 1)
{
const u8v kop = operand_from_f8(d);
if (!hf)
store_row16(&kl[g][tok][16 * nt], kop);
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 + 2 * i + hf][tok] = (half_t)d[i];
vt[g][16 * nt + wm_acc_row(i)][tok] = (half_t)d[i];
}
}
}
@@ -154,17 +152,18 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
}
#pragma unroll
for (int ds = 0; ds < 2; ++ds)
d = mma16(lds_row16(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
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]);
rs += (float)e[kt][i];
if constexpr (!kShimRowSum)
rs += (float)e[kt][i];
}
}
rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
rs = row_sum_total(rs, e);
const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
u8v pop[4];
op_t pop[4];
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
{
@@ -180,15 +179,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v d = splat8(0.0f);
#pragma unroll
for (int kt = 0; kt < 4; ++kt)
d = mma16(lds_row16(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
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 && !hf)
if (h < H)
{
store_row16(&ob[g][tok][0], oop[0]);
store_row16(&ob[g][tok][16], oop[1]);
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)
@@ -196,15 +195,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int hh = r * NA + jh;
if (hh >= H)
break;
const u8v o0 = lds_row16(&ob[jh][tok][0]), o1 = lds_row16(&ob[jh][tok][16]);
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(img[(L::T_WO + (2 * hh) * KT + nt) * 16 + m], o0, acc[j]);
acc[j] = mma16(img[(L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m], o1, acc[j]);
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]);
}
}
}
@@ -220,8 +219,8 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
if (nt < KT)
{
half_t mv[8], g2[8], b2[8];
dvec8(w, L::G2 + 32 * nt, hf, g2);
dvec8(w, L::B2 + 32 * nt, hf, b2);
dvec8(w, L::G2 + 32 * nt, g2);
dvec8(w, L::B2 + 32 * nt, b2);
#pragma unroll
for (int i = 0; i < 8; ++i)
{
@@ -229,9 +228,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
mv[i] = x1 * g2[i];
acc[j][i] = (float)(half_t)(x1 + b2[i]);
}
const u8v mo = operand_from_dt(mv);
if (!hf)
store_row16(hrow + 16 * nt, mo);
op_store(hrow + 16 * nt, operand_from_dt(mv));
}
}
block_sync();
@@ -249,18 +246,16 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
f8v d;
{
half_t vv[8];
dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv);
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(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
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);
const u8v gop = operand_from_dt(gg);
if (!hf)
store_row16(&hid[rr % D::NB][g][tok][16 * hn], gop);
op_store(&hid[rr % D::NB][g][tok][16 * hn], operand_from_dt(gg));
}
block_sync();
for (int jc = 0; jc < NG; ++jc)
@@ -268,15 +263,15 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int cc = rr * NG + jc;
if (cc >= L::NMLP)
break;
const u8v h0 = lds_row16(&hid[rr % D::NB][jc][tok][0]), h1 = lds_row16(&hid[rr % D::NB][jc][tok][16]);
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(img[(L::T_W2 + (cc * 2) * KT + nt) * 16 + m], h0, acc[j]);
acc[j] = mma16(img[(L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m], h1, acc[j]);
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]);
}
}
}
@@ -290,9 +285,7 @@ __device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, cons
const int nt = g + NG * j;
if (nt < KT)
{
const u8v yo = operand_from_f8(acc[j]);
if (!hf)
store_row16(hrow + 16 * nt, yo);
op_store(hrow + 16 * nt, operand_from_f8(acc[j]));
}
}
block_sync();
@@ -349,7 +342,7 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
f8v d;
{
half_t vv[8];
dvec8(w, L::PMB + 32 * vt, hf, vv);
dvec8(w, L::PMB + 32 * vt, vv);
#pragma unroll
for (int i = 0; i < 8; ++i)
d[i] = (float)vv[i];
@@ -358,17 +351,10 @@ __device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, c
{
const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
const int src = 8 * (2 * my + dy) + 2 * mx + dx;
d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&hb[src][16 * (ks % L::KT)]), d);
d = mma16(op_img(img, (L::T_PM + ks * L::PMT + vt) * 16 + m), op_lds(&hb[src][16 * (ks % L::KT)]), d);
}
const u8v mo = operand_from_f8(d);
const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
half_t* dst = mout + L::NPW * gw + 16 * u;
if (!hf && minb)
{
gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
if (valid >= 16)
gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
}
op_gstore(mout + L::NPW * gw + 16 * u, operand_from_f8(d), minb, valid);
}
}
@@ -409,21 +395,19 @@ __device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, c
f8v d;
{
half_t vv[8];
dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
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(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
const u8v sv = gload_row16(sk + 16 * nt);
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] + op_get(sv, 2 * i + hf);
const u8v xv = operand_from_dt(xo);
if (!hf)
store_row16(&act[tt][16 * nt], xv);
xo[i] = (half_t)d[i] + wm_op_acc_elem(sv, i);
op_store(&act[tt][16 * nt], operand_from_dt(xo));
}
}
block_sync();