Add RDNA4 (gfx12) support and a 0.1.2 optimization round
RDNA4 (experimental; untested on RDNA4 hardware) - kernels/common/wmma_layout.h picks the WMMA operand/accumulator layout per target; K (pwin_*) and M (swin_*) index through it, so the same sources build for gfx11 and gfx12. D4R_WMMA_LAYOUT=12 on gfx11 is a test shim that runs the gfx12 indexing on gfx11 WMMA. - M can use native e4m3 FP8 WMMA on gfx12 (D4R_NATIVE_FP8=1). Texture tails (enc0_tail, dec0_head) get gfx12 f16 and FP8 paths that match ZLUDA's gfx12 lowering. - ZLUDA patch 0007: gfx12 WMMA helpers, optional native FP8 MMA (D4R_ZLUDA_WMMA_FP8_NATIVE), a layout override for testing, and an architecture argument for d4r_emit. Texture kernels are now compiled offline per target through d4r_emit (D4R_ZLUDA_EMIT) instead of on the local GPU. - New d4r.ini option [Kernels] NativeFp8 (default on), wired through the shim and d4r_config.py. On gfx12 the bridge serves kernels/<arch>-fp8 when it is on. - The bridge picks the KFD GPU with the most SIMDs (so an iGPU is not chosen over a discrete GPU); D4R_GPU_ARCH overrides it. - kernels/build.sh and package_release.sh build gfx1100-gfx1103 and gfx1200/gfx1201 (plus -fp8 variants), and D4R_BUNDLE_TEX takes per-target folders. Performance (RDNA3; harness byte-identical unless noted) - ZLUDA patch 0006: m16n8k8 f16 MMAs on WMMA (E; not bit-exact, 52-55 dB per layer, D4R_ZLUDA_WMMA_K8=0 disables it), per-key weight-image slots for native prep kernels (d4r_prep_key_at, d4r_prep_key_slots), and wave64 builds of MMA-free kernels (D4R_ZLUDA_WAVE64). - M: packed FP8 encoding, conflict-free patch-merge reads, scalar weight loads in enc1, and weight prep once per weight set for the six enc3 tube blocks. - K: CU mode for the six non-wide pwin layers. - Texture kernels: every DLSS flag variant is built natively (24 K output, 8 M post, 4 M downsample, 4 M enc0), not only Townfall's; post and K output run as wave64. - M direct output: rrlite_downsample_kernel writes the shim's output buffer; the shim's default direct-output presets are now 11,13. Tools - dump_runner takes the prep grid from d4r_prep_blocks and refuses texture/surface dumps. - New pwin_check.py, swin_check.py, model_enc3.py and tail_check.hip. Docs: README GPU table and RDNA4 notes, building (patches 0002-0007, d4r_emit), native-kernels (RDNA4 section), performance (new numbers), packaging README. Version 0.1.2.
This commit is contained in:
ayo committed
1 parent
f310d52b69
commit
4334a0c3c4
42 files changed
+3068
-529
No files matched your search
@@ -2,25 +2,32 @@
|
||||
# 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.
|
||||
# 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}"
|
||||
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)
|
||||
# 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"
|
||||
--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" -
|
||||
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
|
||||
# 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 D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 \
|
||||
D4R_ZLUDA_WAVE64="${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):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
|
||||
printf '%-52s %s\n' "$K" "$(grep -E '^\s+\.(vgpr_count|private_segment_fixed_size|wavefront_size):' "$B/dump/asm.s" | tr -s ' ' | paste -sd' ')"
|
||||
@@ -6,8 +6,14 @@
|
||||
|
||||
// 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
|
||||
{
|
||||
@@ -27,7 +33,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_dec0_4
|
||||
#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)
|
||||
@@ -47,6 +57,58 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
|
||||
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)
|
||||
@@ -109,4 +171,5 @@ DEV void d4r_dec0_head(uint32_t smem, uint64_t wp, uint64_t inp, uint32_t dimx,
|
||||
__builtin_amdgcn_perm(b23, a23, 0x05040100u), __builtin_amdgcn_perm(b23, a23, 0x07060302u)};
|
||||
*(__attribute__((address_space(3))) u4v*)(s + base + goff[G]) = words;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
// Native tail of rrlite_enc0_4x4_mvhi_hdr_folded: the 2x2-pixel patch embedding GEMM (16 tokens x
|
||||
// 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
|
||||
@@ -7,8 +7,14 @@
|
||||
|
||||
// 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
|
||||
{
|
||||
@@ -16,7 +22,13 @@ struct Enc0Params
|
||||
uint8_t rest[352];
|
||||
};
|
||||
|
||||
KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4x4_mvhi_hdr_folded_prep(Enc0Params prm)
|
||||
// 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)
|
||||
@@ -29,7 +41,11 @@ KERNEL __attribute__((amdgpu_flat_work_group_size(128, 128))) void rrlite_enc0_4
|
||||
#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)
|
||||
@@ -38,6 +54,69 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
|
||||
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);
|
||||
@@ -121,4 +200,5 @@ DEV void d4r_enc0_tail(uint32_t smem, uint64_t wp, uint64_t outp, uint32_t dimx,
|
||||
uint32_t* out = (uint32_t*)outp + rowstride * (uint32_t)Y + 8 * (uint32_t)Xc + plane * hf;
|
||||
*(u8v*)out = words;
|
||||
}
|
||||
#endif
|
||||
}
|
||||
+97
-24
@@ -11,6 +11,28 @@ 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(
|
||||
@@ -25,27 +47,7 @@ KERNELS = {
|
||||
)
|
||||
;
|
||||
""",
|
||||
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;
|
||||
|
||||
}
|
||||
""",
|
||||
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.
|
||||
@@ -206,12 +208,83 @@ def replace_rz_round(lines):
|
||||
return out, n
|
||||
|
||||
|
||||
ENC0_REF = "rrlite_enc0_4x4_mvhi_hdr_folded"
|
||||
# the other flag combinations of the kernels whose only native part is the surface-store rewrite
|
||||
SUST_ONLY_RE = re.compile(r"^(hiluma_engine_output_depth(inv|reg)_mv(hi|lo)_(hdr|ldr)(_max)?_v[12]_rel|"
|
||||
r"rrlite_post_3_[12]_mv(hi|lo)_(hdr|ldr)_folded|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_output"))
|
||||
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 known.exists() and entry.search(known.read_text()):
|
||||
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()):
|
||||
@@ -221,8 +294,8 @@ def module_file(name, k):
|
||||
|
||||
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")
|
||||
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"]:
|
||||
|
||||
@@ -3,6 +3,10 @@
|
||||
// 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)
|
||||
|
||||
@@ -65,6 +69,75 @@ __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)
|
||||
{
|
||||
|
||||
Reference in new issue
Block a user