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

Clear windows branch; preserve previous contents in win-old

Browse Source
This commit is contained in:
ayo committed 2026-10-16 05:20:08 -04:00
1 parent f0c881985b
commit 73994d4634
311 files changed
-57018
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

kernels/tex/build_tex.sh
-37
View File
@@ -1,37 +0,0 @@
#!/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' ')"
kernels/tex/dec0_head.hip
-175
View File
@@ -1,175 +0,0 @@
// 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
}
kernels/tex/enc0_tail.hip
-204
View File
@@ -1,204 +0,0 @@
// 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
}
kernels/tex/make_ptx.py
-329
View File
@@ -1,329 +0,0 @@
#!/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()
kernels/tex/sust_only.hip
-2
View File
@@ -1,2 +0,0 @@
// Native surface stores only (see make_ptx.py: post, downsample); the kernels are ZLUDA's compile.
#include "tex_common.h"
kernels/tex/tex_common.h
-448
View File
@@ -1,448 +0,0 @@
// 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);
}