d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton
This commit is contained in:
ayo committed
commit
118173ecd7
128 files changed
+31333
No files matched your search
Executable
+26
@@ -0,0 +1,26 @@
|
||||
#!/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_PTXLOAD and
|
||||
# needs D4R_ZLUDA_BUILD.
|
||||
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}"
|
||||
LLVM="$R/lib/llvm/bin"
|
||||
B="$(mktemp -d)"
|
||||
trap 'rm -rf "$B"' EXIT
|
||||
# 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 -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" -
|
||||
python3 "$D/make_ptx.py" "$K" "$B/$K.ptx" >/dev/null
|
||||
env LD_LIBRARY_PATH="$D4R_ZLUDA_BUILD:$R/lib" XDG_CACHE_HOME="$B/jit" D4R_ZLUDA_DUMP_DIR="$B/dump" \
|
||||
D4R_ZLUDA_EXTRA_BC="$B/extra.bc" D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 D4R_ZLUDA_IGNORE_DENORMAL=1 \
|
||||
D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 "$D4R_PTXLOAD" "$B/$K.ptx" >/dev/null
|
||||
cp "$B/dump/module.hsaco" "$OUT/$K.hsaco"
|
||||
printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
|
||||
@@ -0,0 +1,112 @@
|
||||
// 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
|
||||
__attribute__((device)) u8v g_dec0_w16[8 * 2 * 2 * 16];
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
|
||||
|
||||
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));
|
||||
g_dec0_w16[idx] = operand16(d);
|
||||
}
|
||||
|
||||
__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;
|
||||
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;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
// Native tail of rrlite_enc0_4x4_mvhi_hdr_folded: 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.
|
||||
__attribute__((device)) u8v g_enc0_w16[4 * 4 * 2 * 16];
|
||||
DEV __attribute__((used)) uint32_t d4r_prep_blocks = 4;
|
||||
|
||||
struct Enc0Params
|
||||
{
|
||||
uint64_t w;
|
||||
uint8_t rest[352];
|
||||
};
|
||||
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4x4_mvhi_hdr_folded_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));
|
||||
g_enc0_w16[idx] = operand16(d);
|
||||
}
|
||||
|
||||
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;
|
||||
|
||||
// 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;
|
||||
}
|
||||
}
|
||||
Executable
+253
@@ -0,0 +1,253 @@
|
||||
#!/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"))
|
||||
|
||||
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="""{
|
||||
.param .b32 q0;
|
||||
st.param.b32 [q0+0], %r259;
|
||||
.param .b64 q1;
|
||||
st.param.b64 [q1+0], %rd106;
|
||||
.param .b64 q2;
|
||||
st.param.b64 [q2+0], %rd78;
|
||||
.param .b32 q3;
|
||||
st.param.b32 [q3+0], %r1066;
|
||||
.param .b32 q4;
|
||||
st.param.b32 [q4+0], %r1067;
|
||||
.param .b32 q5;
|
||||
st.param.b32 [q5+0], %r1230;
|
||||
.param .b32 q6;
|
||||
st.param.b32 [q6+0], %r1233;
|
||||
call.uni d4r_enc0_tail, (q0, q1, q2, q3, q4, q5, q6);
|
||||
}
|
||||
ret;
|
||||
|
||||
}
|
||||
""",
|
||||
),
|
||||
# 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
|
||||
|
||||
|
||||
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 known.exists() 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 = KERNELS[name]
|
||||
lines = module_file(name, k).read_text().replace("\r\n", "\n").split("\n")
|
||||
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,375 @@
|
||||
// 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>
|
||||
|
||||
#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).
|
||||
#ifdef D4R_SUST_OUTLINE
|
||||
#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);
|
||||
}
|
||||
|
||||
// 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