feat: add d4r inference engines and Windows parity tooling
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
This commit is contained in:
ayo committed
1 parent
25493b3198
commit
7c5d3c9157
430 files changed
+76191
No files matched your search
@@ -0,0 +1,37 @@
|
||||
#!/usr/bin/env bash
|
||||
# build_tex.sh KERNEL SRC OUT_DIR: one texture kernel. NVIDIA's PTX for KERNEL, edited by make_ptx.py so
|
||||
# part of it calls functions from SRC.hip, is compiled by ZLUDA with SRC's bitcode linked in
|
||||
# (D4R_ZLUDA_EXTRA_BC) and the resulting code object is saved as OUT_DIR/KERNEL.hsaco.
|
||||
# Called by kernels/build.sh, which sets D4R_ROCM_DIR, D4R_GPU_ARCH, D4R_DLSS_PTX_DIR, D4R_ZLUDA_EMIT,
|
||||
# D4R_ZLUDA_WAVE64 (1 = wave64 code object) and D4R_TEX_FP8 (1 = the gfx12 native-FP8 variant).
|
||||
# ZLUDA compiles offline for D4R_GPU_ARCH with its d4r_emit tool, so any target builds on any machine.
|
||||
set -euo pipefail
|
||||
K=$1 SRC=$2 OUT=$3
|
||||
D="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
R="${D4R_ROCM_DIR:-/opt/rocm}"
|
||||
ARCH="${D4R_GPU_ARCH:-gfx1101}"
|
||||
FP8="${D4R_TEX_FP8:-0}"
|
||||
ACCURACY="${D4R_PREFER_ACCURACY:-0}"
|
||||
LLVM="$R/lib/llvm/bin"
|
||||
B="$(mktemp -d)"
|
||||
trap 'rm -rf "$B"' EXIT
|
||||
# D4R_TEX_CFLAGS: extra flags for SRC.hip (tests, e.g. -DD4R_WMMA_LAYOUT=12 for the gfx12-layout shim on gfx11)
|
||||
read -r -a FLAGS <<< "${D4R_TEX_CFLAGS:-}"
|
||||
[[ "$FP8" == 1 ]] && FLAGS+=(-DD4R_TEX_FP8)
|
||||
[[ "$ACCURACY" == 1 ]] && FLAGS+=(-DD4R_ACCURACY)
|
||||
# SRC.hip -> bitcode without target attributes, so ZLUDA's own target settings apply after linking
|
||||
"$LLVM/clang++" -x hip -std=c++20 -nogpuinc -nogpulib -O3 -mno-wavefrontsize64 --offload-device-only \
|
||||
--offload-arch="$ARCH" -fgpu-rdc -emit-llvm -c -Xclang -fdenormal-fp-math=dynamic -DD4R_KERNEL_NAME="$K" \
|
||||
"${FLAGS[@]}" -o "$B/raw.bc" "$D/$SRC.hip"
|
||||
"$LLVM/llvm-dis" "$B/raw.bc" -o "$B/raw.ll"
|
||||
sed -E -e '/@llvm.used/d' -e '/wchar_size/d' -e '/llvm.module.flags/d' -e '/__hip_cuid/d' -e 's/optnone//g' \
|
||||
-e "s/\"target-cpu\"=\"$ARCH\"//g" -e 's/"target-features"="[^"]+"//g' "$B/raw.ll" | "$LLVM/llvm-as" -o "$B/extra.bc" -
|
||||
D4R_PREFER_ACCURACY="$ACCURACY" python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
|
||||
# the code-generation switches of the runtime (d4r.ini [Kernels]), which ZLUDA's module cache keys on
|
||||
env LD_LIBRARY_PATH="$R/lib" D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 \
|
||||
D4R_ZLUDA_WMMA_FP8_NATIVE="$FP8" D4R_ZLUDA_IGNORE_DENORMAL="$((1 - ACCURACY))" \
|
||||
D4R_ZLUDA_WMMA_F32ACC=0 D4R_ZLUDA_FAST_MATH=0 D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 \
|
||||
D4R_ZLUDA_WAVE64="$([[ "$ACCURACY" == 1 ]] && echo 0 || echo "${D4R_ZLUDA_WAVE64:-0}")" \
|
||||
"$D4R_ZLUDA_EMIT" "$B/$K.ptx" "$B/dump" "$ARCH" >/dev/null
|
||||
cp "$B/dump/module.hsaco" "$OUT/$K.hsaco"
|
||||
printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size|wavefront_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
|
||||
@@ -0,0 +1,175 @@
|
||||
// Native head of rrlite_dec0_4x4_folded: the token GEMM (16 half-resolution tokens x K 64 (two
|
||||
// e4m3 planes) x N 128, f16 bias, f16 accumulation per k32 step) and the e4m3 staging of its result
|
||||
// in shared memory, which the rest of the kernel (ZLUDA's compile of NVIDIA's PTX, dec0_head.ptx)
|
||||
// reads. Arithmetic as ZLUDA's m16n8k32 lowering (see enc0_tail.hip); tokens are mirror padded.
|
||||
#include "tex_common.h"
|
||||
|
||||
// weights as WMMA A images: index ((p * 2 + c) * 2 + h) * 16 + n16 for channel pair p (16p..16p+15),
|
||||
// k32 chunk c (input plane), k16 half h
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
__attribute__((device)) kslot g_dec0_w16[8 * 2 * 2 * 16];
|
||||
#else
|
||||
__attribute__((device)) u8v g_dec0_w16[8 * 2 * 2 * 16];
|
||||
#endif
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
|
||||
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
|
||||
|
||||
struct Dec0Params
|
||||
{
|
||||
uint64_t w;
|
||||
uint8_t rest[160]; // 168-byte parameter block
|
||||
};
|
||||
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4x4_folded_prep(Dec0Params prm)
|
||||
{
|
||||
const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
|
||||
if (idx >= 512)
|
||||
return;
|
||||
const uint8_t* w = (const uint8_t*)prm.w;
|
||||
const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 1, p = idx >> 6;
|
||||
const int g = n16 & 7, second = n16 >> 3;
|
||||
u4v d;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
d[j] = *(const uint32_t*)(w + 512 * (2 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
g_dec0_w16[idx] = kslot_from16(d);
|
||||
#else
|
||||
g_dec0_w16[idx] = operand16(d);
|
||||
#endif
|
||||
}
|
||||
|
||||
__attribute__((device)) static inline uint32_t mirror_clamp(int32_t v, uint32_t full, uint32_t half)
|
||||
{
|
||||
const int32_t a = v < 0 ? -v : v;
|
||||
const int32_t m = __builtin_elementwise_min(a, (int32_t)(full & ~1u) - a - 2);
|
||||
return __builtin_elementwise_min(half - 1u, (uint32_t)m);
|
||||
}
|
||||
|
||||
DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
|
||||
{
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const uint8_t* w = (const uint8_t*)wp;
|
||||
lds_u8* s = (lds_u8*)(uintptr_t)smem;
|
||||
|
||||
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
|
||||
const uint32_t Y = mirror_clamp(oy / 2 + (int32_t)(m >> 2), dimy, H2);
|
||||
const uint32_t Xc = mirror_clamp(ox / 2 + (int32_t)(m & 3), dimx, W2);
|
||||
const uint32_t* in = (const uint32_t*)inp + rowstride * Y + 8 * Xc;
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
kop X[2][2];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c)
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 2; ++h)
|
||||
X[c][h] = kop_from8(*(const uint2_t*)(in + plane * c + 4 * h + 2 * hf));
|
||||
|
||||
// E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7
|
||||
uint32_t E[8][2];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 8; ++p)
|
||||
{
|
||||
f8v cf;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)*(const half_t*)(w + 8192 + 2 * (16 * p + i + 8 * hf));
|
||||
half_t acc[8];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c)
|
||||
{
|
||||
const int slot = ((p * 2 + c) * 2) * 16 + m;
|
||||
cf = k32_e4m3(cf, kop_image(g_dec0_w16, slot), X[c][0], kop_image(g_dec0_w16, slot + 16), X[c][1]);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (half_t)cf[i];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)acc[i];
|
||||
}
|
||||
uint32_t codes[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
codes[k] = e4m3x2(pack2(acc[2 * k], acc[2 * k + 1]));
|
||||
E[p][0] = codes[0] | (codes[1] << 16);
|
||||
E[p][1] = codes[2] | (codes[3] << 16);
|
||||
}
|
||||
|
||||
// staging (layout as below): this lane holds tiles 4G + hf and 4G + 2 + hf (channels 16p + 8hf ..), so it
|
||||
// writes H = hf itself
|
||||
const uint32_t base = 1024 * hf + 32 * (m & 3) + 256 * (m >> 2);
|
||||
const uint32_t goff[4] = {0, 16, 128, 144};
|
||||
#pragma unroll
|
||||
for (int G = 0; G < 4; ++G)
|
||||
{
|
||||
const uint32_t* p0 = E[2 * G];
|
||||
const uint32_t* p1 = E[2 * G + 1];
|
||||
u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
|
||||
__builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
|
||||
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
|
||||
}
|
||||
#else
|
||||
u8v X[2][2];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c)
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 2; ++h)
|
||||
X[c][h] = operand16(*(const u4v*)(in + plane * c + 4 * h));
|
||||
|
||||
uint32_t E[8][2];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 8; ++p)
|
||||
{
|
||||
half_t acc[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = *(const half_t*)(w + 8192 + 2 * (16 * p + 2 * i + hf));
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 2; ++c)
|
||||
{
|
||||
f8v cf;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)acc[i];
|
||||
cf *= 0x1p-16f;
|
||||
const u8v* wv = &g_dec0_w16[((p * 2 + c) * 2) * 16 + m];
|
||||
cf = wmma(wv[0], X[c][0], cf);
|
||||
cf = wmma(wv[16], X[c][1], cf);
|
||||
cf = cf * 0x1p16f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (half_t)cf[i];
|
||||
}
|
||||
uint32_t codes[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
codes[k] = e4m3x2(pack2(acc[2 * k], acc[2 * k + 1]));
|
||||
E[p][0] = codes[0] | (codes[1] << 16);
|
||||
E[p][1] = codes[2] | (codes[3] << 16);
|
||||
}
|
||||
|
||||
// staging: word t of token m at 1024 * H + 32 * (m & 3) + 256 * (m >> 2) + {0, 16, 128, 144}[G] + 4t
|
||||
// holds columns 2t, 2t + 1 of tiles 4G + H and 4G + 2 + H; lanes 0..15 write H = 0, 16..31 H = 1
|
||||
uint32_t EV[8], OD[8];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 8; ++p)
|
||||
{
|
||||
const uint32_t theirs = other_half(hf ? E[p][0] : E[p][1]);
|
||||
EV[p] = hf ? theirs : E[p][0];
|
||||
OD[p] = hf ? E[p][1] : theirs;
|
||||
}
|
||||
const uint32_t base = 1024 * hf + 32 * (m & 3) + 256 * (m >> 2);
|
||||
const uint32_t goff[4] = {0, 16, 128, 144};
|
||||
#pragma unroll
|
||||
for (int G = 0; G < 4; ++G)
|
||||
{
|
||||
const uint32_t a01 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x05010400u);
|
||||
const uint32_t a23 = __builtin_amdgcn_perm(OD[2 * G], EV[2 * G], 0x07030602u);
|
||||
const uint32_t b01 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x05010400u);
|
||||
const uint32_t b23 = __builtin_amdgcn_perm(OD[2 * G + 1], EV[2 * G + 1], 0x07030602u);
|
||||
u4v words = {__builtin_amdgcn_perm(b01, a01, 0x05040100u), __builtin_amdgcn_perm(b01, a01, 0x07060302u),
|
||||
__builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
|
||||
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,204 @@
|
||||
// Native tail of rrlite_enc0_4x4_*_folded (all flag combinations): the 2x2-pixel patch embedding GEMM (16 tokens x
|
||||
// K 128 x N 64, e4m3 inputs, f16 accumulation per k32 step, f16 bias), clamp to +-2pi and e4m3 store.
|
||||
// The feature code before it is ZLUDA's compile of the original PTX (enc0_tail.ptx calls this after
|
||||
// the kernel's own shared-memory staging); the arithmetic reproduces ZLUDA's m16n8k32 lowering
|
||||
// (scaled e4m3 widening, same k slot order, C * 2^-16 inside the f32 WMMA, f16 rounding per k32).
|
||||
#include "tex_common.h"
|
||||
|
||||
// B operands (weights) as WMMA A images: index ((p * 4 + c) * 2 + h) * 16 + n16 for channel pair p
|
||||
// (channels 16p..16p+15), k32 chunk c, k16 half h.
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
__attribute__((device)) kslot g_enc0_w16[4 * 4 * 2 * 16];
|
||||
#else
|
||||
__attribute__((device)) u8v g_enc0_w16[4 * 4 * 2 * 16];
|
||||
#endif
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
|
||||
// the prep reads only prm.w (offset 0): the hook skips it while that pointer repeats
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_key_at = 0 + 1;
|
||||
|
||||
struct Enc0Params
|
||||
{
|
||||
uint64_t w;
|
||||
uint8_t rest[352];
|
||||
};
|
||||
|
||||
// every flag combination (mvhi/mvlo, hdr/ldr) has this tail; build_tex.sh names the prep after the kernel
|
||||
#ifndef D4R_KERNEL_NAME
|
||||
#define D4R_KERNEL_NAME rrlite_enc0_4x4_mvhi_hdr_folded
|
||||
#endif
|
||||
#define D4R_CAT2(a, b) a##b
|
||||
#define D4R_CAT(a, b) D4R_CAT2(a, b)
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void D4R_CAT(D4R_KERNEL_NAME, _prep)(Enc0Params prm)
|
||||
{
|
||||
const uint32_t idx = __builtin_amdgcn_workgroup_id_x() * 128 + __builtin_amdgcn_workitem_id_x();
|
||||
if (idx >= 512)
|
||||
return;
|
||||
const uint8_t* w = (const uint8_t*)prm.w;
|
||||
const int n16 = idx & 15, h = (idx >> 4) & 1, c = (idx >> 5) & 3, p = idx >> 7;
|
||||
// NVIDIA B fragments: v4 load 4p + c of lane 4 * column + j holds {b0, b1} of tile 2p then 2p + 1
|
||||
const int g = n16 & 7, second = n16 >> 3;
|
||||
u4v d;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
d[j] = *(const uint32_t*)(w + 4096 + 512 * (4 * p + c) + 16 * (4 * g + j) + 4 * (2 * second + h));
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
g_enc0_w16[idx] = kslot_from16(d);
|
||||
#else
|
||||
g_enc0_w16[idx] = operand16(d);
|
||||
#endif
|
||||
}
|
||||
|
||||
DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx, uint32_t dimy, int32_t ox, int32_t oy)
|
||||
{
|
||||
const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
|
||||
const uint8_t* w = (const uint8_t*)wp;
|
||||
lds_u8* s = (lds_u8*)(uintptr_t)smem;
|
||||
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
// activations: token m's 128 e4m3 features (addresses as below); this half's 8 bytes of each k16 part
|
||||
const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
|
||||
const uint32_t choff[4] = {0, 16, 128, 144};
|
||||
kop X[4][2];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 4; ++c)
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 2; ++h)
|
||||
X[c][h] = kop_from8(*(const __attribute__((address_space(3))) uint2_t*)(s + 1024 * h + rowoff + choff[c] + 8 * hf));
|
||||
|
||||
const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
|
||||
// E[p]: e4m3 codes of channels 16p + 8hf .. 16p + 8hf + 7 (D^T element i = channel 16p + i + 8hf)
|
||||
uint32_t E[4][2];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 4; ++p)
|
||||
{
|
||||
f8v cf;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)*(const half_t*)(w + 12288 + 2 * (16 * p + i + 8 * hf));
|
||||
half_t acc[8];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 4; ++c)
|
||||
{
|
||||
const int slot = ((p * 4 + c) * 2) * 16 + m;
|
||||
cf = k32_e4m3(cf, kop_image(g_enc0_w16, slot), X[c][0], kop_image(g_enc0_w16, slot + 16), X[c][1]);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (half_t)cf[i];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)acc[i];
|
||||
}
|
||||
uint32_t codes[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
hv2 v = {acc[2 * k], acc[2 * k + 1]};
|
||||
v = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, lo), hi);
|
||||
codes[k] = e4m3x2(__builtin_bit_cast(uint32_t, v));
|
||||
}
|
||||
E[p][0] = codes[0] | (codes[1] << 16);
|
||||
E[p][1] = codes[2] | (codes[3] << 16);
|
||||
}
|
||||
|
||||
// plane pl, word 4 j + q = columns 2q, 2q + 1 of tile (2 pl, j) then of tile (2 pl + 1, j) (tile (p, j) =
|
||||
// channels 16p + 8j ..): this lane holds tiles j = hf, i.e. words 4 hf .. 4 hf + 3 of both planes
|
||||
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
|
||||
const int32_t Y = oy / 2 + (int32_t)(m >> 2), Xc = ox / 2 + (int32_t)(m & 3);
|
||||
if ((uint32_t)Y < H2 && (uint32_t)Xc < W2)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int pl = 0; pl < 2; ++pl)
|
||||
{
|
||||
const uint32_t* p0 = E[2 * pl];
|
||||
const uint32_t* p1 = E[2 * pl + 1];
|
||||
const u4v words = {__builtin_amdgcn_perm(p1[0], p0[0], 0x05040100u), __builtin_amdgcn_perm(p1[0], p0[0], 0x07060302u),
|
||||
__builtin_amdgcn_perm(p1[1], p0[1], 0x05040100u), __builtin_amdgcn_perm(p1[1], p0[1], 0x07060302u)};
|
||||
*(u4v*)((uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * pl + 4 * hf) = words;
|
||||
}
|
||||
}
|
||||
#else
|
||||
// activations: token m's 128 e4m3 features; A[m][k] of chunk c at
|
||||
// 1024 * (k >> 4) + 32 * (m & 3) + 256 * (m >> 2) + (k & 15) + {0, 16, 128, 144}[c]
|
||||
const uint32_t rowoff = 32 * (m & 3) + 256 * (m >> 2);
|
||||
const uint32_t choff[4] = {0, 16, 128, 144};
|
||||
u8v X[4][2];
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 4; ++c)
|
||||
#pragma unroll
|
||||
for (int h = 0; h < 2; ++h)
|
||||
X[c][h] = operand16(*(const __attribute__((address_space(3))) u4v*)(s + 1024 * h + rowoff + choff[c]));
|
||||
|
||||
const hv2 lo = {(half_t)-6.2831855f, (half_t)-6.2831855f}, hi = {(half_t)6.2831855f, (half_t)6.2831855f};
|
||||
uint32_t E[4][2];
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 4; ++p)
|
||||
{
|
||||
// D^T layout: element i is channel 16p + 2i + hf of token m
|
||||
half_t acc[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = *(const half_t*)(w + 12288 + 2 * (16 * p + 2 * i + hf));
|
||||
#pragma unroll
|
||||
for (int c = 0; c < 4; ++c)
|
||||
{
|
||||
f8v cf;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
cf[i] = (float)acc[i];
|
||||
cf *= 0x1p-16f;
|
||||
const u8v* wv = &g_enc0_w16[((p * 4 + c) * 2) * 16 + m];
|
||||
cf = wmma(wv[0], X[c][0], cf);
|
||||
cf = wmma(wv[16], X[c][1], cf);
|
||||
cf = cf * 0x1p16f;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i)
|
||||
acc[i] = (half_t)cf[i];
|
||||
}
|
||||
uint32_t codes[4];
|
||||
#pragma unroll
|
||||
for (int k = 0; k < 4; ++k)
|
||||
{
|
||||
hv2 v = {acc[2 * k], acc[2 * k + 1]};
|
||||
v = __builtin_elementwise_minimumnum(__builtin_elementwise_maximumnum(v, lo), hi);
|
||||
codes[k] = e4m3x2(__builtin_bit_cast(uint32_t, v));
|
||||
}
|
||||
// byte i of E[p][0] (E[p][1]) is column 2i + hf (2(i - 4) + hf) of tile 2p (2p + 1)
|
||||
E[p][0] = codes[0] | (codes[1] << 16);
|
||||
E[p][1] = codes[2] | (codes[3] << 16);
|
||||
}
|
||||
|
||||
// lanes 0..15 write plane 0 (tiles 0..3), lanes 16..31 plane 1 (tiles 4..7): swap the other
|
||||
// plane's codes with the partner lane holding the other column parity
|
||||
uint32_t mine[2][2], theirs[2][2];
|
||||
#pragma unroll
|
||||
for (int a = 0; a < 2; ++a)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 2; ++j)
|
||||
{
|
||||
const uint32_t send = hf ? E[a][j] : E[2 + a][j];
|
||||
theirs[a][j] = other_half(send);
|
||||
mine[a][j] = hf ? E[2 + a][j] : E[a][j];
|
||||
}
|
||||
u8v words;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 2; ++j)
|
||||
{
|
||||
const uint32_t ev0 = hf ? theirs[0][j] : mine[0][j], od0 = hf ? mine[0][j] : theirs[0][j];
|
||||
const uint32_t ev1 = hf ? theirs[1][j] : mine[1][j], od1 = hf ? mine[1][j] : theirs[1][j];
|
||||
const uint32_t a01 = __builtin_amdgcn_perm(od0, ev0, 0x05010400u), a23 = __builtin_amdgcn_perm(od0, ev0, 0x07030602u);
|
||||
const uint32_t b01 = __builtin_amdgcn_perm(od1, ev1, 0x05010400u), b23 = __builtin_amdgcn_perm(od1, ev1, 0x07030602u);
|
||||
words[4 * j + 0] = __builtin_amdgcn_perm(b01, a01, 0x05040100u);
|
||||
words[4 * j + 1] = __builtin_amdgcn_perm(b01, a01, 0x07060302u);
|
||||
words[4 * j + 2] = __builtin_amdgcn_perm(b23, a23, 0x05040100u);
|
||||
words[4 * j + 3] = __builtin_amdgcn_perm(b23, a23, 0x07060302u);
|
||||
}
|
||||
|
||||
const uint32_t H2 = dimy >> 1, W2 = dimx >> 1, rowstride = W2 * 8, plane = H2 * rowstride;
|
||||
const int32_t Y = oy / 2 + (int32_t)(m >> 2), Xc = ox / 2 + (int32_t)(m & 3);
|
||||
if ((uint32_t)Y < H2 && (uint32_t)Xc < W2)
|
||||
{
|
||||
uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
|
||||
*(u8v*)out = words;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -0,0 +1,329 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Hand-edited PTX for the texture kernels: keep NVIDIA's code up to a cut line, then call a native tail.
|
||||
|
||||
usage: make_ptx.py KERNEL OUT.ptx
|
||||
"""
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
# PTX extracted from nvngx_dlss.dll by kernels/tools/extract_dlss_ptx.py (kernels/build.sh does this)
|
||||
PTX_DIR = Path(os.environ.get("D4R_DLSS_PTX_DIR", Path(__file__).resolve().parent.parent / "extracted" / "ptx"))
|
||||
|
||||
ENC0_TAIL = """{{
|
||||
.param .b32 q0;
|
||||
st.param.b32 [q0+0], {smem};
|
||||
.param .b64 q1;
|
||||
st.param.b64 [q1+0], {w};
|
||||
.param .b64 q2;
|
||||
st.param.b64 [q2+0], {out};
|
||||
.param .b32 q3;
|
||||
st.param.b32 [q3+0], {dimx};
|
||||
.param .b32 q4;
|
||||
st.param.b32 [q4+0], {dimy};
|
||||
.param .b32 q5;
|
||||
st.param.b32 [q5+0], {ox};
|
||||
.param .b32 q6;
|
||||
st.param.b32 [q6+0], {oy};
|
||||
call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
|
||||
}}
|
||||
ret;
|
||||
|
||||
}}
|
||||
"""
|
||||
|
||||
KERNELS = {
|
||||
# enc0: cut after the bar.warp.sync that ends the smem staging of the GEMM A operand
|
||||
"rrlite_enc0_4x4_mvhi_hdr_folded": dict(
|
||||
file="dlss-0185-00.ptx",
|
||||
cut_after=2548,
|
||||
cut_check="bar.warp.sync -1;",
|
||||
sust=True,
|
||||
extern=""".extern .func d4r_enc0_tail
|
||||
(
|
||||
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
|
||||
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
|
||||
)
|
||||
;
|
||||
""",
|
||||
tail=ENC0_TAIL.format(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233"),
|
||||
),
|
||||
# dec0: the GEMM (from the B fragment loads to the e4m3 staging stores) becomes a native call; the
|
||||
# A loads before it are left dead. %r560 (shared base) is the only register of the range used later.
|
||||
"rrlite_dec0_4x4_folded": dict(
|
||||
file="dlss-0195-00.ptx",
|
||||
delete=(139, 587),
|
||||
delete_check=("mov.u32 %r17, %laneid;", "st.shared.v2.u16 [%r570+1168], {%rs63, %rs64};"),
|
||||
sust=True,
|
||||
extern=""".extern .func d4r_dec0_head
|
||||
(
|
||||
.param .b32 d4r_a0, .param .b64 d4r_a1, .param .b64 d4r_a2, .param .b32 d4r_a3,
|
||||
.param .b32 d4r_a4, .param .b32 d4r_a5, .param .b32 d4r_a6
|
||||
)
|
||||
;
|
||||
""",
|
||||
insert="""mov.u32 %r560, _ZZ22rrlite_dec0_4x4_foldedN6RRLite18Decoder0ParametersEE4smem;
|
||||
{
|
||||
.param .b32 q0;
|
||||
st.param.b32 [q0+0], %r560;
|
||||
.param .b64 q1;
|
||||
st.param.b64 [q1+0], %rd22;
|
||||
.param .b64 q2;
|
||||
st.param.b64 [q2+0], %rd24;
|
||||
.param .b32 q3;
|
||||
st.param.b32 [q3+0], %r475;
|
||||
.param .b32 q4;
|
||||
st.param.b32 [q4+0], %r476;
|
||||
.param .b32 q5;
|
||||
st.param.b32 [q5+0], %r3;
|
||||
.param .b32 q6;
|
||||
st.param.b32 [q6+0], %r4;
|
||||
call.uni d4r_dec0_head, (q0, q1, q2, q3, q4, q5, q6);
|
||||
}""",
|
||||
),
|
||||
# post and downsample: only the formatted surface stores become native (inline format conversion
|
||||
# instead of ZLUDA's per-channel surface_formatted_store_bits calls)
|
||||
"rrlite_post_3_2_mvhi_hdr_folded": dict(file="dlss-0176-00.ptx", sust=True, extern=""),
|
||||
"rrlite_downsample_kernel_static_hdr": dict(file="dlss-0196-00.ptx", sust=True, extern=""),
|
||||
# preset K (DLSS 4): the output kernel's 17 surface stores (13 formatted, 4 raw 32-bit)
|
||||
"hiluma_engine_output_depthinv_mvhi_hdr_max_v2_rel": dict(file="dlss-0078-00.ptx", sust=True, rz_round=True, extern=""),
|
||||
"hiluma_engine_input_depthinv_mvhi_hdr_v2_rel": dict(file="dlss-0062-00.ptx", sust=True, rz_round=True, extern=""),
|
||||
}
|
||||
|
||||
|
||||
SUST_EXTERN = """.extern .func d4r_sust_v2b16
|
||||
(
|
||||
.param .b64 d4r_s0, .param .b32 d4r_s1, .param .b32 d4r_s2, .param .b16 d4r_s3, .param .b16 d4r_s4
|
||||
)
|
||||
;
|
||||
"""
|
||||
SUSTP_EXTERN = """.extern .func d4r_sust_p_v4b32
|
||||
(
|
||||
.param .b64 d4r_t0, .param .b32 d4r_t1, .param .b32 d4r_t2, .param .b32 d4r_t3, .param .b32 d4r_t4,
|
||||
.param .b32 d4r_t5, .param .b32 d4r_t6
|
||||
)
|
||||
;
|
||||
"""
|
||||
SUSTP_RE = re.compile(r"^sust\.p\.2d\.v4\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%f\d+),(%f\d+),(%f\d+),(%f\d+)\};$")
|
||||
SUSTB32_EXTERN = """.extern .func d4r_sust_b32
|
||||
(
|
||||
.param .b64 d4r_w0, .param .b32 d4r_w1, .param .b32 d4r_w2, .param .b32 d4r_w3
|
||||
)
|
||||
;
|
||||
"""
|
||||
SUSTB32_RE = re.compile(r"^sust\.b\.2d\.b32\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%r\d+)\};$")
|
||||
SUST_RE = re.compile(r"^sust\.b\.2d\.v2\.b16\.zero \[(%rd\d+), \{(%r\d+),(%r\d+)\}\], \{(%rs\d+),(%rs\d+)\};$")
|
||||
|
||||
|
||||
SUST_KINDS = set(os.environ.get("D4R_SUST_KINDS", "b16,p,b32").split(","))
|
||||
|
||||
|
||||
def replace_sust(lines):
|
||||
out, n = [], {"b16": 0, "p": 0, "b32": 0}
|
||||
for line in lines:
|
||||
mb = SUSTB32_RE.match(line.strip()) if "b32" in SUST_KINDS else None
|
||||
if mb:
|
||||
s, x, y, v = mb.groups()
|
||||
out.append(f"""{{
|
||||
.param .b64 u0;
|
||||
st.param.b64 [u0+0], {s};
|
||||
.param .b32 u1;
|
||||
st.param.b32 [u1+0], {x};
|
||||
.param .b32 u2;
|
||||
st.param.b32 [u2+0], {y};
|
||||
.param .b32 u3;
|
||||
st.param.b32 [u3+0], {v};
|
||||
call d4r_sust_b32, (u0, u1, u2, u3);
|
||||
}}""")
|
||||
n["b32"] += 1
|
||||
continue
|
||||
mp = SUSTP_RE.match(line.strip()) if "p" in SUST_KINDS else None
|
||||
if mp:
|
||||
s, x, y, *v = mp.groups()
|
||||
st = "\n".join(f"\t.param .b32 v{i};\n\tst.param.b32 [v{i}+0], {r};" for i, r in enumerate(v))
|
||||
out.append(f"""{{
|
||||
.param .b64 u0;
|
||||
st.param.b64 [u0+0], {s};
|
||||
.param .b32 u1;
|
||||
st.param.b32 [u1+0], {x};
|
||||
.param .b32 u2;
|
||||
st.param.b32 [u2+0], {y};
|
||||
{st}
|
||||
call d4r_sust_p_v4b32, (u0, u1, u2, v0, v1, v2, v3);
|
||||
}}""")
|
||||
n["p"] += 1
|
||||
continue
|
||||
m = SUST_RE.match(line.strip())
|
||||
if not m:
|
||||
out.append(line)
|
||||
continue
|
||||
s, x, y, a, b = m.groups()
|
||||
out.append(f"""{{
|
||||
.param .b64 u0;
|
||||
st.param.b64 [u0+0], {s};
|
||||
.param .b32 u1;
|
||||
st.param.b32 [u1+0], {x};
|
||||
.param .b32 u2;
|
||||
st.param.b32 [u2+0], {y};
|
||||
.param .b16 u3;
|
||||
st.param.b16 [u3+0], {a};
|
||||
.param .b16 u4;
|
||||
st.param.b16 [u4+0], {b};
|
||||
call d4r_sust_v2b16, (u0, u1, u2, u3, u4);
|
||||
}}""")
|
||||
n["b16"] += 1
|
||||
return out, n
|
||||
|
||||
|
||||
RZ_RE = re.compile(r"^add\.rz\.ftz\.f32 (%f\d+), (%f\d+), (%f\d+);$")
|
||||
CVT_RE = re.compile(r"^cvt\.rzi\.f32\.f32 (%f\d+), (%f\d+);$")
|
||||
|
||||
|
||||
def replace_rz_round(lines):
|
||||
"""roundf idiom add.rz(x, copysign(0.5, x)) + cvt.rzi -> exact round-half-away in default rounding:
|
||||
t = trunc(x); r = |x - t| >= 0.5 ? t + 2 * copysign(0.5, x) : t (x - t is exact). Lets ZLUDA compile the
|
||||
kernel without the strict-FP (constrained) mode that one explicit rounding mode forces."""
|
||||
out, n, i = [], 0, 0
|
||||
while i < len(lines):
|
||||
m = RZ_RE.match(lines[i].strip())
|
||||
if m:
|
||||
d, x, sgn = m.groups()
|
||||
j = i + 1
|
||||
while j < len(lines) and not lines[j].strip():
|
||||
j += 1
|
||||
c = CVT_RE.match(lines[j].strip()) if j < len(lines) else None
|
||||
if c and c.group(2) == d:
|
||||
r = c.group(1)
|
||||
out += [f"cvt.rzi.f32.f32 %rq{4*n}, {x};", f"sub.f32 %rq{4*n+1}, {x}, %rq{4*n};",
|
||||
f"abs.f32 %rq{4*n+1}, %rq{4*n+1};", f"setp.ge.f32 %pq{n}, %rq{4*n+1}, 0f3F000000;",
|
||||
f"add.f32 %rq{4*n+2}, {sgn}, {sgn};", f"add.f32 %rq{4*n+3}, %rq{4*n}, %rq{4*n+2};",
|
||||
f"selp.f32 {r}, %rq{4*n+3}, %rq{4*n}, %pq{n};"]
|
||||
n += 1
|
||||
i = j + 1
|
||||
continue
|
||||
raise SystemExit(f"add.rz at line {i} not followed by its cvt.rzi")
|
||||
out.append(lines[i])
|
||||
i += 1
|
||||
return out, n
|
||||
|
||||
|
||||
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
|
||||
# Surface-store and round-idiom recipes for the other flag combinations.
|
||||
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
|
||||
r"hiluma_engine_input_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)_v2_rel|"
|
||||
r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)(_folded)?|"
|
||||
r"rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)|rrlite_dec0_4x4|"
|
||||
r"rrlite_downsample_kernel_(static|dynamic)_(hdr|ldr))$")
|
||||
ENC0_RE = re.compile(r"^rrlite_enc0_4x4_mv(hi|lo)_(hdr|ldr)_folded$")
|
||||
|
||||
|
||||
def read_lines(path):
|
||||
return path.read_text().replace("\r\n", "\n").split("\n")
|
||||
|
||||
|
||||
def enc0_registers(lines):
|
||||
"""Cut line (1-based) and the tail call's registers of an enc0 variant. The code after the third
|
||||
bar.warp.sync is the same GEMM in every flag combination, only register numbers differ: the shared
|
||||
base and weights pointer are read off by aligning that block with the reference variant; the output
|
||||
pointer and grid size are the kernel's loads of parameters +40 and +8; the block offsets are the two
|
||||
values halved (shr 31 / add) by the epilogue."""
|
||||
syncs = [i for i, l in enumerate(lines) if l.strip() == "bar.warp.sync -1;"]
|
||||
if len(syncs) != 3:
|
||||
sys.exit(f"enc0: expected 3 bar.warp.sync, found {len(syncs)}")
|
||||
cut = syncs[2] + 1
|
||||
before, after = lines[:cut], lines[cut:]
|
||||
out = [m.group(1) for m in (re.match(r"^ld\.param\.u64 (%rd\d+), \[%rd\d+\+40\];$", l.strip()) for l in before) if m]
|
||||
dims = [m.groups() for m in (re.match(r"^ld\.param\.v2\.u32 \{(%r\d+), (%r\d+)\}, \[%rd\d+\+8\];$", l.strip()) for l in before) if m]
|
||||
halves = []
|
||||
for a, b in zip(after, after[1:]):
|
||||
m = re.match(r"^shr\.u32 (%r\d+), (%r\d+), 31;$", a.strip())
|
||||
if m and re.match(rf"^add\.s32 %r\d+, {re.escape(m.group(2))}, {re.escape(m.group(1))};$", b.strip()):
|
||||
halves.append(m.group(2))
|
||||
if len(out) != 1 or len(dims) != 1 or len(halves) < 2:
|
||||
sys.exit(f"enc0: parameter loads or epilogue not found ({out}, {dims}, {halves[:2]})")
|
||||
return cut, dict(out=out[0], dimx=dims[0][0], dimy=dims[0][1], oy=halves[0], ox=halves[1])
|
||||
|
||||
|
||||
REG_RE = re.compile(r"%(?:rd|rs|r|fd|f|p)\d+")
|
||||
|
||||
|
||||
def enc0_variant(name):
|
||||
ref = read_lines(module_file(ENC0_REF, KERNELS[ENC0_REF]))
|
||||
path = module_file(name, {"file": ""})
|
||||
lines = read_lines(path)
|
||||
rcut, rregs = enc0_registers(ref)
|
||||
k = KERNELS[ENC0_REF]
|
||||
known = dict(smem="%r259", w="%rd106", out="%rd78", dimx="%r1066", dimy="%r1067", ox="%r1230", oy="%r1233")
|
||||
if rcut != k["cut_after"] or any(rregs[key] != known[key] for key in rregs):
|
||||
sys.exit(f"enc0: the register search does not reproduce {ENC0_REF} ({rcut}, {rregs})")
|
||||
cut, regs = enc0_registers(lines)
|
||||
# align the GEMM block (up to the reference's first closing brace) and map registers
|
||||
n = next(i for i, l in enumerate(ref[rcut:]) if l.startswith("}"))
|
||||
mapping = {}
|
||||
for a, b in zip(ref[rcut:rcut + n], lines[cut:cut + n]):
|
||||
if REG_RE.sub("R", a) != REG_RE.sub("R", b):
|
||||
sys.exit(f"enc0: {name} differs from {ENC0_REF} after the cut: {a!r} / {b!r}")
|
||||
for x, y in zip(REG_RE.findall(a), REG_RE.findall(b)):
|
||||
if mapping.setdefault(x, y) != y:
|
||||
sys.exit(f"enc0: inconsistent register mapping {x} -> {mapping[x]} / {y}")
|
||||
regs.update(smem=mapping[known["smem"]], w=mapping[known["w"]])
|
||||
return dict(file=path.name, cut_after=cut, cut_check="bar.warp.sync -1;", sust=True, extern=k["extern"],
|
||||
tail=ENC0_TAIL.format(**regs))
|
||||
|
||||
|
||||
def kernel_spec(name):
|
||||
if name in KERNELS:
|
||||
return KERNELS[name]
|
||||
if ENC0_RE.match(name):
|
||||
return enc0_variant(name)
|
||||
if SUST_ONLY_RE.match(name):
|
||||
return dict(file="", sust=True, extern="", rz_round=name.startswith("hiluma_engine_"))
|
||||
sys.exit(f"make_ptx: no recipe for {name}")
|
||||
|
||||
|
||||
def module_file(name, k):
|
||||
"""The extracted PTX module that defines NAME: the 310.7 file number, else any module that does
|
||||
(other DLSS versions number their modules differently)."""
|
||||
entry = re.compile(rf"\.entry\s+{re.escape(name)}\s*\(")
|
||||
known = PTX_DIR / k["file"]
|
||||
if k["file"] and known.is_file() and entry.search(known.read_text()):
|
||||
return known
|
||||
for path in sorted(PTX_DIR.glob("*.ptx")):
|
||||
if entry.search(path.read_text()):
|
||||
return path
|
||||
sys.exit(f"no PTX module in {PTX_DIR} defines {name}")
|
||||
|
||||
|
||||
def main():
|
||||
name, out = sys.argv[1], Path(sys.argv[2])
|
||||
k = kernel_spec(name)
|
||||
lines = read_lines(module_file(name, k))
|
||||
if "cut_after" in k:
|
||||
cut = k["cut_after"]
|
||||
if lines[cut - 1].strip() != k["cut_check"]:
|
||||
sys.exit(f"line {cut} is {lines[cut - 1]!r}, expected {k['cut_check']!r}")
|
||||
head, tail = lines[:cut], k["tail"].split("\n")
|
||||
elif "delete" in k:
|
||||
a, b = k["delete"]
|
||||
if (lines[a - 1].strip(), lines[b - 1].strip()) != k["delete_check"]:
|
||||
sys.exit(f"lines {a}/{b} are {lines[a - 1]!r}/{lines[b - 1]!r}, expected {k['delete_check']}")
|
||||
head, tail = lines[:a - 1] + k["insert"].split("\n") + lines[b:], []
|
||||
else:
|
||||
head, tail = lines, []
|
||||
entry = next(i for i, l in enumerate(head) if l.startswith(".visible .entry"))
|
||||
body, n = replace_sust(head[entry:]) if k.get("sust") else (head[entry:], {"b16": 0, "p": 0, "b32": 0})
|
||||
if k.get("rz_round"):
|
||||
body, nrz = replace_rz_round(body)
|
||||
# temporaries for the rewrite, declared after the kernel's own .reg lines
|
||||
at = next(i for i, l in enumerate(body) if l.startswith(".reg"))
|
||||
body = body[:at] + [f".reg .f32 %rq<{4 * nrz + 1}>;", f".reg .pred %pq<{nrz + 1}>;"] + body[at:]
|
||||
n["rz"] = nrz
|
||||
extern = k["extern"] + (SUST_EXTERN if n["b16"] else "") + (SUSTP_EXTERN if n["p"] else "") + (SUSTB32_EXTERN if n["b32"] else "")
|
||||
text = "\n".join(head[:entry]) + "\n" + extern + "\n".join(body + tail)
|
||||
out.write_text(text)
|
||||
print(f"{out}: {len(text.splitlines())} lines, surface stores replaced: {n}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,2 @@
|
||||
// Native surface stores only (see make_ptx.py: post, downsample); the kernels are ZLUDA's compile.
|
||||
#include "tex_common.h"
|
||||
@@ -0,0 +1,448 @@
|
||||
// Shared device code for native tails of the rrlite texture kernels (DLSS 310.7.0, preset M).
|
||||
// These are linked into ZLUDA's own compile of a hand-edited PTX kernel (D4R_ZLUDA_EXTRA_BC), so
|
||||
// only plain clang builtins are used (compiled with -nogpuinc -nogpulib).
|
||||
#pragma once
|
||||
#include <stdint.h>
|
||||
#include "../common/wmma_layout.h"
|
||||
#if defined(D4R_TEX_FP8) && D4R_WMMA_LAYOUT != 12 && defined(__HIP_DEVICE_COMPILE__)
|
||||
#error "D4R_TEX_FP8 needs the gfx12 WMMA layout"
|
||||
#endif
|
||||
|
||||
#pragma clang fp contract(off)
|
||||
|
||||
#define DEV extern "C" __attribute__((device))
|
||||
// Kernels that ZLUDA compiles in strict-FP mode (constrained intrinsics) cannot take inlined FP code from
|
||||
// here (the AMDGPU backend fails on some constrained operations): build them with -DD4R_SUST_OUTLINE,
|
||||
// which keeps the FP bodies of the surface stores in noinline leaf functions (no calls, no stack).
|
||||
#if defined(D4R_SUST_OUTLINE) || defined(D4R_ACCURACY)
|
||||
#define FP_BODY __attribute__((device, noinline)) static
|
||||
#else
|
||||
#define FP_BODY __attribute__((device, always_inline)) static inline
|
||||
#endif
|
||||
#define KERNEL extern "C" __attribute__((global))
|
||||
|
||||
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)));
|
||||
typedef uint32_t u8v __attribute__((ext_vector_type(8)));
|
||||
typedef uint32_t u4v __attribute__((ext_vector_type(4)));
|
||||
typedef uint32_t uint2_t __attribute__((ext_vector_type(2)));
|
||||
typedef uint16_t u16x2 __attribute__((ext_vector_type(2)));
|
||||
typedef int16_t i16x2 __attribute__((ext_vector_type(2)));
|
||||
typedef __attribute__((address_space(3))) uint8_t lds_u8;
|
||||
|
||||
__attribute__((device)) static inline uint32_t lane_id()
|
||||
{
|
||||
return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
|
||||
}
|
||||
|
||||
// ZLUDA's e4m3x4_to_scaled_f16x2x2 (zluda_ptx_impl.cpp): e4m3 bytes -> f16 pairs holding value * 2^-8,
|
||||
// {bytes 0,1} and {bytes 2,3}.
|
||||
__attribute__((device)) static inline void e4m3x4_scaled(uint32_t packed, uint32_t& low, uint32_t& high)
|
||||
{
|
||||
uint32_t magnitude = packed & 0x7f7f7f7fu;
|
||||
magnitude |= (magnitude + 0x01010101u) & 0x80808080u;
|
||||
uint32_t sign = packed & 0x80808080u;
|
||||
low = (__builtin_amdgcn_perm(0u, magnitude, 0x0c010c00u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x010c000cu);
|
||||
high = (__builtin_amdgcn_perm(0u, magnitude, 0x0c030c02u) << 7) | __builtin_amdgcn_perm(0u, sign, 0x030c020cu);
|
||||
}
|
||||
|
||||
// A k16 WMMA operand from 16 consecutive e4m3 bytes of NVIDIA's k order (4 dwords): register j holds
|
||||
// k 4j, 4j+1 and register 4 + j holds k 4j+2, 4j+3 -- ZLUDA's slot order (e4m3_a_k16_fragment).
|
||||
__attribute__((device)) static inline u8v operand16(u4v d)
|
||||
{
|
||||
u8v v;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < 4; ++j)
|
||||
{
|
||||
uint32_t lo, hi;
|
||||
e4m3x4_scaled(d[j], lo, hi);
|
||||
v[j] = lo;
|
||||
v[4 + j] = hi;
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__attribute__((device)) static inline f8v wmma(u8v a, u8v b, f8v c)
|
||||
{
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- gfx12 m16n8k32 e4m3 steps
|
||||
// On gfx12 the GEMM tails do a k32 step as ZLUDA's gfx12 lowering of NVIDIA's m16n8k32 e4m3 MMA does, so they
|
||||
// stay bit-identical to ZLUDA's compile of the kernel they replace: a lane's k16 operand holds this half's
|
||||
// e4m3 k 8h .. 8h + 7 (ZLUDA's gfx12 slot order), either widened as above with C * 2^-16 inside the f32
|
||||
// WMMAs (d4r_wmma12_*), or with D4R_TEX_FP8 (ZLUDA's D4R_ZLUDA_WMMA_FP8_NATIVE) as the bytes of a native FP8
|
||||
// WMMA with no scaling (d4r_wmma12f8_*). (The gfx11 tails use operand16 and wmma above.)
|
||||
#if D4R_WMMA_LAYOUT == 12
|
||||
// 8 e4m3 bytes (k in order) widened to value * 2^-8 f16 pairs, k in order
|
||||
__attribute__((device)) static inline u4v widen8(uint2_t d)
|
||||
{
|
||||
uint32_t a0, a1, b0, b1;
|
||||
e4m3x4_scaled(d.x, a0, a1);
|
||||
e4m3x4_scaled(d.y, b0, b1);
|
||||
return (u4v){a0, a1, b0, b1};
|
||||
}
|
||||
#ifdef D4R_TEX_FP8
|
||||
typedef uint2_t kop;
|
||||
typedef u4v kslot; // weight image slot: the 16 e4m3 bytes of one k16 operand row
|
||||
__attribute__((device)) static inline kop kop_from8(uint2_t d)
|
||||
{
|
||||
return d;
|
||||
}
|
||||
__attribute__((device)) static inline kslot kslot_from16(u4v d)
|
||||
{
|
||||
return d;
|
||||
}
|
||||
#else
|
||||
typedef u4v kop;
|
||||
typedef u8v kslot; // weight image slot: the 16 widened values of one k16 operand row, k in order
|
||||
__attribute__((device)) static inline kop kop_from8(uint2_t d)
|
||||
{
|
||||
return widen8(d);
|
||||
}
|
||||
__attribute__((device)) static inline kslot kslot_from16(u4v d)
|
||||
{
|
||||
const u4v lo = widen8((uint2_t){d[0], d[1]}), hi = widen8((uint2_t){d[2], d[3]});
|
||||
return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
|
||||
}
|
||||
#endif
|
||||
// this half's k16 operand from 16 consecutive e4m3 bytes of NVIDIA's k order
|
||||
__attribute__((device)) static inline kop kop_from16(u4v d)
|
||||
{
|
||||
return kop_from8(lane_id() >= 16 ? (uint2_t){d[2], d[3]} : (uint2_t){d[0], d[1]});
|
||||
}
|
||||
// this half's operand from slot idx of a weight image
|
||||
__attribute__((device)) static inline kop kop_image(const kslot* img, int idx)
|
||||
{
|
||||
return ((const kop*)img)[2 * idx + (lane_id() >> 4)];
|
||||
}
|
||||
// f32 C (true scale) + one k32 step
|
||||
__attribute__((device)) static inline f8v k32_e4m3(f8v c, const kop& a0, const kop& b0, const kop& a1, const kop& b1)
|
||||
{
|
||||
#if defined(D4R_TEX_FP8) && defined(__GFX12__)
|
||||
typedef int i2v __attribute__((ext_vector_type(2)));
|
||||
c = __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a0), __builtin_bit_cast(i2v, b0), c);
|
||||
return __builtin_amdgcn_wmma_f32_16x16x16_fp8_fp8_w32_gfx12(__builtin_bit_cast(i2v, a1), __builtin_bit_cast(i2v, b1), c);
|
||||
#elif defined(D4R_TEX_FP8)
|
||||
// layout shim, like ZLUDA's: each FP8 WMMA as the scaled widening (exact powers of two)
|
||||
c = wm_mma_zluda(widen8(a0), widen8(b0), c * 0x1p-16f) * 0x1p16f;
|
||||
return wm_mma_zluda(widen8(a1), widen8(b1), c * 0x1p-16f) * 0x1p16f;
|
||||
#else
|
||||
c = c * 0x1p-16f;
|
||||
c = wm_mma_zluda(a0, b0, c);
|
||||
c = wm_mma_zluda(a1, b1, c);
|
||||
return c * 0x1p16f;
|
||||
#endif
|
||||
}
|
||||
#endif
|
||||
|
||||
// ZLUDA's f16x2_to_e4m3x2_satfinite_bits: RNE satfinite e4m3 of both halves, low half -> low byte.
|
||||
__attribute__((device)) static inline uint32_t e4m3x2(uint32_t bits)
|
||||
{
|
||||
u16x2 half = __builtin_bit_cast(u16x2, bits);
|
||||
u16x2 magnitude = half & (uint16_t)0x7fffu;
|
||||
u16x2 normal = (magnitude + ((magnitude >> (uint16_t)7) & (uint16_t)1) + (uint16_t)0xe03fu) >> (uint16_t)7;
|
||||
u16x2 exponent = magnitude >> (uint16_t)10;
|
||||
u16x2 clamped = __builtin_elementwise_max(__builtin_elementwise_min(exponent, (u16x2)8), (u16x2)1);
|
||||
u16x2 shift = (uint16_t)16 - clamped;
|
||||
u16x2 significand = (magnitude & (uint16_t)0x3ffu) |
|
||||
(__builtin_elementwise_min(magnitude, (u16x2)0x400) & (uint16_t)0x400u);
|
||||
u16x2 subnormal = (significand + (((u16x2)1 << (shift - (uint16_t)1)) - (uint16_t)1) +
|
||||
((significand >> shift) & (uint16_t)1)) >> shift;
|
||||
u16x2 below = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)(magnitude - (uint16_t)0x2400u)) >> (int16_t)15));
|
||||
u16x2 code = (subnormal & below) | (normal & ~below);
|
||||
code = __builtin_elementwise_min(code, (u16x2)0x7e);
|
||||
u16x2 nan = __builtin_bit_cast(u16x2, (i16x2)(__builtin_bit_cast(i16x2, (u16x2)((uint16_t)0x7c00u - magnitude)) >> (int16_t)15));
|
||||
code |= nan & (uint16_t)0x7fu;
|
||||
code |= (half >> (uint16_t)8) & (uint16_t)0x80u;
|
||||
return uint32_t(code.x) | (uint32_t(code.y) << 8);
|
||||
}
|
||||
|
||||
__attribute__((device)) static inline uint32_t pack2(half_t a, half_t b)
|
||||
{
|
||||
return __builtin_bit_cast(uint32_t, (hv2){a, b});
|
||||
}
|
||||
|
||||
// value from the lane with the same index in the other half of the wave
|
||||
__attribute__((device)) static inline uint32_t other_half(uint32_t v)
|
||||
{
|
||||
return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- surface stores
|
||||
typedef int32_t v2i __attribute__((ext_vector_type(2)));
|
||||
typedef int16_t v2s __attribute__((ext_vector_type(2)));
|
||||
typedef float v4f __attribute__((ext_vector_type(4)));
|
||||
typedef int32_t v4i_t __attribute__((ext_vector_type(4)));
|
||||
typedef __attribute__((address_space(4))) unsigned int tsharp_t;
|
||||
extern "C" __attribute__((device)) int __ockl_image_channel_data_type_2D(tsharp_t* image);
|
||||
extern "C" __attribute__((device)) int __ockl_image_channel_order_2D(tsharp_t* image);
|
||||
extern "C" __attribute__((device)) void __ockl_image_store_2D(tsharp_t* image, v2i coord, v4f value);
|
||||
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t surface, v2i coord, v2s value);
|
||||
|
||||
// ZLUDA's surface_half_bits_to_float (exact, independent of the f16 denormal mode)
|
||||
__attribute__((device)) static inline float half_bits_to_float(uint32_t bits)
|
||||
{
|
||||
uint32_t sign = (bits & 0x8000u) << 16;
|
||||
uint32_t exponent = (bits >> 10) & 0x1fu;
|
||||
uint32_t mantissa = bits & 0x3ffu;
|
||||
if (exponent == 0)
|
||||
{
|
||||
float magnitude = float(mantissa) * 0x1p-24f;
|
||||
return __builtin_bit_cast(float, __builtin_bit_cast(uint32_t, magnitude) | sign);
|
||||
}
|
||||
if (exponent == 0x1f)
|
||||
return __builtin_bit_cast(float, sign | 0x7f800000u | (mantissa << 13));
|
||||
return __builtin_bit_cast(float, sign | ((exponent + 112u) << 23) | (mantissa << 13));
|
||||
}
|
||||
|
||||
FP_BODY void sust_rg16f_body(tsharp_t* image, int32_t px, int32_t y, uint16_t a, uint16_t b)
|
||||
{
|
||||
__ockl_image_store_2D(image, (v2i){px, y}, (v4f){half_bits_to_float(a), half_bits_to_float(b), 0.0f, 0.0f});
|
||||
}
|
||||
|
||||
// sust.b.2d.v2.b16.zero: a whole RG16F pixel is one image store of the decoded halves (what ZLUDA's
|
||||
// surface_store_raw_2d does for that format); anything else goes through ZLUDA's generic helper.
|
||||
__attribute__((device, noinline)) static void sust_v2b16_generic(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
|
||||
{
|
||||
__zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(surface, (v2i){x, y}, (v2s){(int16_t)a, (int16_t)b});
|
||||
}
|
||||
|
||||
DEV __attribute__((always_inline)) void d4r_sust_v2b16(uint64_t surface, int32_t x, int32_t y, uint16_t a, uint16_t b)
|
||||
{
|
||||
tsharp_t* image = (tsharp_t*)surface;
|
||||
if (__ockl_image_channel_data_type_2D(image) == 14 && __ockl_image_channel_order_2D(image) == 3 && (x & 3) == 0)
|
||||
{
|
||||
if (x >= 0)
|
||||
sust_rg16f_body(image, x >> 2, y, a, b);
|
||||
}
|
||||
else
|
||||
sust_v2b16_generic(surface, x, y, a, b);
|
||||
}
|
||||
|
||||
|
||||
// f32 comparisons on the bits (no fcmp: kernels ZLUDA compiles in strict-FP mode cannot lower
|
||||
// constrained fcmp); same results as the float compares, -0.0 included
|
||||
__attribute__((device)) static inline bool f_isnan(uint32_t u) { return (u & 0x7fffffffu) > 0x7f800000u; }
|
||||
__attribute__((device)) static inline bool f_lt0(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) != 0 && !f_isnan(u); }
|
||||
// value > 1.0f (value not NaN)
|
||||
__attribute__((device)) static inline bool f_gt1(uint32_t u) { return (int32_t)u > 0x3f800000; }
|
||||
// value < -1.0f (value not NaN)
|
||||
__attribute__((device)) static inline bool f_ltm1(uint32_t u) { return (int32_t)u < 0 && (u & 0x7fffffffu) > 0x3f800000u; }
|
||||
|
||||
// ZLUDA's surface_formatted_store_bits (zluda_ptx_impl.cpp), inlined: normalized CUDA formats convert
|
||||
// the f32 value, everything else stores the bits unchanged.
|
||||
__attribute__((device)) static inline uint32_t formatted_store_bits(uint32_t bits, uint32_t cuda_format, int channel)
|
||||
{
|
||||
if (cuda_format == 80)
|
||||
{
|
||||
uint32_t u = bits;
|
||||
if (f_isnan(u) || f_lt0(u))
|
||||
u = 0u;
|
||||
if (f_gt1(u))
|
||||
u = 0x3f800000u;
|
||||
const float value = __builtin_bit_cast(float, u);
|
||||
const float scale = channel == 3 ? 3.0f : 1023.0f;
|
||||
// float(int(floor(x))) == floor(x) here (0 <= x <= 1023.5): no int <-> float conversions, which
|
||||
// strict-FP kernels cannot lower as constrained signed conversions
|
||||
return __builtin_bit_cast(uint32_t, __builtin_floorf(value * scale + 0.5f) / scale);
|
||||
}
|
||||
float scale = 0.0f;
|
||||
bool signed_normalized = false;
|
||||
if (cuda_format >= 192 && cuda_format <= 194)
|
||||
scale = 255.0f;
|
||||
else if (cuda_format >= 195 && cuda_format <= 197)
|
||||
scale = 65535.0f;
|
||||
else if (cuda_format >= 198 && cuda_format <= 200)
|
||||
{
|
||||
scale = 127.0f;
|
||||
signed_normalized = true;
|
||||
}
|
||||
else if (cuda_format >= 201 && cuda_format <= 203)
|
||||
{
|
||||
scale = 32767.0f;
|
||||
signed_normalized = true;
|
||||
}
|
||||
if (!(cuda_format >= 192 && cuda_format <= 203))
|
||||
return bits;
|
||||
uint32_t u = bits;
|
||||
if (f_isnan(u))
|
||||
u = 0u;
|
||||
if (signed_normalized ? f_ltm1(u) : f_lt0(u))
|
||||
u = signed_normalized ? 0xbf800000u : 0u;
|
||||
if (f_gt1(u))
|
||||
u = 0x3f800000u;
|
||||
const float scaled = __builtin_bit_cast(float, u) * scale;
|
||||
const uint32_t su = __builtin_bit_cast(uint32_t, scaled);
|
||||
int32_t integer = f_lt0(su) ? -int32_t(__builtin_floorf(-scaled + 0.5f)) : int32_t(__builtin_floorf(scaled + 0.5f));
|
||||
return uint32_t(integer);
|
||||
}
|
||||
|
||||
// sust.p.2d.v4.b32.zero: ZLUDA stores the raw bits unless the original CUDA format was normalized
|
||||
// (surface_formatted_store_bits; format word after the image and sampler descriptors).
|
||||
FP_BODY void sust_p_body(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b, uint32_t c, uint32_t d)
|
||||
{
|
||||
const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
|
||||
__ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
|
||||
(v4f){__builtin_bit_cast(float, formatted_store_bits(a, format, 0)),
|
||||
__builtin_bit_cast(float, formatted_store_bits(b, format, 1)),
|
||||
__builtin_bit_cast(float, formatted_store_bits(c, format, 2)),
|
||||
__builtin_bit_cast(float, formatted_store_bits(d, format, 3))});
|
||||
}
|
||||
|
||||
// d4r output redirect (wine_nvcuda_bridge.c d4rSetArrayRedirect): a tagged row pitch at byte 84 of the
|
||||
// ZLUDA surface object (zero padding otherwise) sends stores of that surface to the linear memory whose
|
||||
// address is at byte 88, rows `pitch` bytes apart. Out-of-range coordinates are dropped as by .zero stores.
|
||||
extern "C" __attribute__((device)) int __ockl_image_width_2D(tsharp_t* image);
|
||||
extern "C" __attribute__((device)) int __ockl_image_height_2D(tsharp_t* image);
|
||||
__attribute__((device)) static inline uint32_t redirect_pitch(uint64_t surface)
|
||||
{
|
||||
// 'R2' tag in the high half (other handles may have data there), row pitch / 8 in the low half
|
||||
const uint32_t word = ((const __attribute__((address_space(4))) uint32_t*)surface)[21];
|
||||
return (word >> 16) == 0x5232u ? (word & 0xffffu) << 3 : 0u;
|
||||
}
|
||||
__attribute__((device)) static inline uint8_t* redirect_row(uint64_t surface, int32_t y)
|
||||
{
|
||||
const uint64_t base = *(const __attribute__((address_space(4))) uint64_t*)(surface + 88);
|
||||
return (uint8_t*)(base + (uint64_t)(uint32_t)y * redirect_pitch(surface));
|
||||
}
|
||||
__attribute__((device)) static inline bool redirect_inside(uint64_t surface, int32_t x, int32_t y)
|
||||
{
|
||||
tsharp_t* image = (tsharp_t*)surface;
|
||||
return x >= 0 && y >= 0 && x < __ockl_image_width_2D(image) && y < __ockl_image_height_2D(image);
|
||||
}
|
||||
|
||||
DEV __attribute__((always_inline)) void d4r_sust_p_v4b32(uint64_t surface, int32_t x, int32_t y, uint32_t a, uint32_t b,
|
||||
uint32_t c, uint32_t d)
|
||||
{
|
||||
const uint32_t format0 = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
|
||||
if (redirect_pitch(surface) != 0 && format0 != 80 && !(format0 >= 192 && format0 <= 203))
|
||||
{
|
||||
// RGBA16F: the f32 values as halves, rounded toward zero like the image store's format conversion
|
||||
if (redirect_inside(surface, x, y))
|
||||
{
|
||||
const uint32_t lo = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, a), __builtin_bit_cast(float, b)));
|
||||
const uint32_t hi = __builtin_bit_cast(uint32_t, __builtin_amdgcn_cvt_pkrtz(__builtin_bit_cast(float, c), __builtin_bit_cast(float, d)));
|
||||
*(uint2_t*)(redirect_row(surface, y) + 8 * x) = (uint2_t){lo, hi};
|
||||
}
|
||||
return;
|
||||
}
|
||||
// formatted_store_bits keeps the bits unless the format is normalized: one uniform test per store
|
||||
const uint32_t format = ((const __attribute__((address_space(4))) uint32_t*)surface)[20];
|
||||
if (format != 80 && !(format >= 192 && format <= 203))
|
||||
__ockl_image_store_2D((tsharp_t*)surface, (v2i){x, y},
|
||||
(v4f){__builtin_bit_cast(float, a), __builtin_bit_cast(float, b), __builtin_bit_cast(float, c),
|
||||
__builtin_bit_cast(float, d)});
|
||||
else
|
||||
sust_p_body(surface, x, y, a, b, c, d);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- raw 32-bit surface stores
|
||||
extern "C" __attribute__((device)) void __zluda_ptx_impl_surfobj_b_2d_b32_zero(uint64_t surface, v2i coord, uint32_t value);
|
||||
|
||||
// exact int -> float for |v| < 2^24 without a signed conversion (see formatted_store_bits)
|
||||
__attribute__((device)) static inline float i2f(int32_t v)
|
||||
{
|
||||
return v < 0 ? -(float)(uint32_t)(-v) : (float)(uint32_t)v;
|
||||
}
|
||||
|
||||
// ZLUDA's surface_decode_channel (bits of one memory channel -> the float4 lane value the store converts back)
|
||||
__attribute__((device)) static inline float decode_channel(uint32_t bits, int data_type)
|
||||
{
|
||||
switch (data_type)
|
||||
{
|
||||
case 0:
|
||||
return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int8_t(bits)), 127), -127)) / 127.0f;
|
||||
case 1:
|
||||
return i2f(__builtin_elementwise_max(__builtin_elementwise_min(int32_t(int16_t(bits)), 32767), -32767)) / 32767.0f;
|
||||
case 2:
|
||||
return (float)(bits & 0xffu) / 255.0f;
|
||||
case 3:
|
||||
return (float)(bits & 0xffffu) / 65535.0f;
|
||||
case 8:
|
||||
return __builtin_bit_cast(float, int32_t(int8_t(bits)));
|
||||
case 9:
|
||||
return __builtin_bit_cast(float, int32_t(int16_t(bits)));
|
||||
case 14:
|
||||
return half_bits_to_float(bits);
|
||||
default:
|
||||
return __builtin_bit_cast(float, bits);
|
||||
}
|
||||
}
|
||||
|
||||
// out of line: inlining ZLUDA's generic raw store into large kernels crashes AMDGPU instruction
|
||||
// selection (see zluda_ptx_impl.cpp); it makes no calls, so the wrapper needs no stack
|
||||
__attribute__((device, noinline)) static void sust_b32_generic(uint64_t surface, int32_t x, int32_t y, uint32_t data)
|
||||
{
|
||||
__zluda_ptx_impl_surfobj_b_2d_b32_zero(surface, (v2i){x, y}, data);
|
||||
}
|
||||
|
||||
// sust.b.2d.b32.zero (x is a byte offset): whole 4-byte pixels and pairs of 2-byte pixels in R/RG/RGBA
|
||||
// channel order need no read-modify-write (ZLUDA's surface_store_raw_2d, same conversions); anything
|
||||
// else goes through ZLUDA's generic helper.
|
||||
FP_BODY void sust_b32_fast(tsharp_t* image, int32_t x, int32_t y, uint32_t data, int dt, int cb, int pb)
|
||||
{
|
||||
if (pb == 4)
|
||||
{
|
||||
v4f v = {0.0f, 0.0f, 0.0f, 0.0f};
|
||||
if (cb == 4)
|
||||
v[0] = decode_channel(data, dt);
|
||||
else if (cb == 2)
|
||||
{
|
||||
v[0] = decode_channel(data & 0xffffu, dt);
|
||||
v[1] = decode_channel(data >> 16, dt);
|
||||
}
|
||||
else
|
||||
{
|
||||
v[0] = decode_channel(data & 0xffu, dt);
|
||||
v[1] = decode_channel((data >> 8) & 0xffu, dt);
|
||||
v[2] = decode_channel((data >> 16) & 0xffu, dt);
|
||||
v[3] = decode_channel(data >> 24, dt);
|
||||
}
|
||||
__ockl_image_store_2D(image, (v2i){x >> 2, y}, v);
|
||||
}
|
||||
else
|
||||
{
|
||||
v4f v0 = {0.0f, 0.0f, 0.0f, 0.0f}, v1 = {0.0f, 0.0f, 0.0f, 0.0f};
|
||||
if (cb == 2)
|
||||
{
|
||||
v0[0] = decode_channel(data & 0xffffu, dt);
|
||||
v1[0] = decode_channel(data >> 16, dt);
|
||||
}
|
||||
else
|
||||
{
|
||||
v0[0] = decode_channel(data & 0xffu, dt);
|
||||
v0[1] = decode_channel((data >> 8) & 0xffu, dt);
|
||||
v1[0] = decode_channel((data >> 16) & 0xffu, dt);
|
||||
v1[1] = decode_channel(data >> 24, dt);
|
||||
}
|
||||
__ockl_image_store_2D(image, (v2i){x >> 1, y}, v0);
|
||||
__ockl_image_store_2D(image, (v2i){(x >> 1) + 1, y}, v1);
|
||||
}
|
||||
}
|
||||
|
||||
DEV __attribute__((always_inline)) void d4r_sust_b32(uint64_t surface, int32_t x, int32_t y, uint32_t data)
|
||||
{
|
||||
tsharp_t* image = (tsharp_t*)surface;
|
||||
if (redirect_pitch(surface) != 0)
|
||||
{
|
||||
// raw bits at byte offset x (RGBA16F rows are 8 bytes per pixel)
|
||||
if (x >= 0 && (x & 3) == 0 && redirect_inside(surface, x >> 3, y))
|
||||
*(uint32_t*)(redirect_row(surface, y) + x) = data;
|
||||
return;
|
||||
}
|
||||
const int dt = __ockl_image_channel_data_type_2D(image), order = __ockl_image_channel_order_2D(image);
|
||||
const int cb = (dt == 0 || dt == 2 || dt == 8 || dt == 11) ? 1 : (dt == 1 || dt == 3 || dt == 9 || dt == 12 || dt == 14) ? 2
|
||||
: (dt == 10 || dt == 13 || dt == 15) ? 4 : 0;
|
||||
const int n = order == 1 ? 1 : order == 3 ? 2 : order == 8 ? 4 : 0; // R, RG, RGBA (identity lane order)
|
||||
const int pb = cb * n;
|
||||
if ((pb == 4 && (x & 3) == 0) || (pb == 2 && (x & 1) == 0))
|
||||
{
|
||||
if (x >= 0)
|
||||
sust_b32_fast(image, x, y, data, dt, cb, pb);
|
||||
}
|
||||
else
|
||||
sust_b32_generic(surface, x, y, data);
|
||||
}
|
||||
Reference in new issue
Block a user