feat: add d4r inference engines and Windows parity tooling
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
This commit is contained in:
ayo committed
1 parent
25493b3198
commit
7c5d3c9157
430 files changed
+76191
No files matched your search
@@ -0,0 +1,185 @@
|
||||
// Launch a native kernel on a replay dump (manifest.txt, args.bin, alloc-N.bin), save every allocation
|
||||
// after the first launch, then time repeated launches.
|
||||
// usage: dump_runner <hsaco> <kernel> <dump dir> <out dir> [iterations]
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <algorithm>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#define CHECK(x) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (x); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
fprintf(stderr, "%s failed: %s\n", #x, hipGetErrorString(e_)); \
|
||||
return 1; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static std::vector<char> read_file(const std::string& path)
|
||||
{
|
||||
std::ifstream f(path, std::ios::binary);
|
||||
return std::vector<char>((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc < 5)
|
||||
{
|
||||
fprintf(stderr, "usage: %s hsaco kernel dump out [iters]\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
const std::string dump = argv[3], out = argv[4];
|
||||
const int iters = argc > 5 ? atoi(argv[5]) : 20;
|
||||
hipModule_t mod;
|
||||
hipFunction_t fn;
|
||||
CHECK(hipModuleLoad(&mod, argv[1]));
|
||||
CHECK(hipModuleGetFunction(&fn, mod, argv[2]));
|
||||
hipFunction_t prep = nullptr;
|
||||
if (hipModuleGetFunction(&prep, mod, (std::string(argv[2]) + "_prep").c_str()) != hipSuccess)
|
||||
prep = nullptr;
|
||||
// the prep grid: PREP_GRID, else the module's d4r_prep_blocks (what the ZLUDA hook launches), else 88
|
||||
unsigned prepGrid = 88;
|
||||
{
|
||||
hipDeviceptr_t pb;
|
||||
size_t pbs = 0;
|
||||
uint32_t v = 0;
|
||||
if (hipModuleGetGlobal(&pb, &pbs, mod, "d4r_prep_blocks") == hipSuccess && pbs == 4 &&
|
||||
hipMemcpyDtoH(&v, pb, 4) == hipSuccess && v != 0)
|
||||
prepGrid = v;
|
||||
}
|
||||
if (getenv("PREP_GRID"))
|
||||
prepGrid = atoi(getenv("PREP_GRID"));
|
||||
|
||||
std::vector<char> args = read_file(dump + "/args.bin");
|
||||
std::ifstream man(dump + "/manifest.txt");
|
||||
std::string line;
|
||||
unsigned grid[3] = {1, 1, 1}, block[3] = {1, 1, 1}, shared = 0;
|
||||
std::vector<int> ids;
|
||||
std::vector<void*> dev;
|
||||
std::vector<size_t> sizes;
|
||||
while (std::getline(man, line))
|
||||
{
|
||||
std::istringstream s(line);
|
||||
std::string kind;
|
||||
s >> kind;
|
||||
if (kind == "texture" || kind == "surface")
|
||||
{
|
||||
// texture/surface object handles in the arguments would point at nothing here: the kernel faults
|
||||
fprintf(stderr, "%s: the dump uses texture or surface objects, which dump_runner cannot recreate\n",
|
||||
dump.c_str());
|
||||
return 3;
|
||||
}
|
||||
if (kind == "launch")
|
||||
s >> grid[0] >> grid[1] >> grid[2] >> block[0] >> block[1] >> block[2] >> shared;
|
||||
else if (kind == "alloc")
|
||||
{
|
||||
int id;
|
||||
std::string base;
|
||||
size_t size;
|
||||
s >> id >> base >> size;
|
||||
std::vector<char> data = read_file(dump + "/alloc-" + std::to_string(id) + ".bin");
|
||||
if (data.size() != size)
|
||||
{
|
||||
fprintf(stderr, "alloc %d size mismatch\n", id);
|
||||
return 1;
|
||||
}
|
||||
void* d;
|
||||
CHECK(hipMalloc(&d, size));
|
||||
CHECK(hipMemcpy(d, data.data(), size, hipMemcpyHostToDevice));
|
||||
ids.push_back(id);
|
||||
dev.push_back(d);
|
||||
sizes.push_back(size);
|
||||
}
|
||||
else if (kind == "pointer")
|
||||
{
|
||||
size_t off, delta;
|
||||
int id;
|
||||
s >> off >> id >> delta;
|
||||
size_t i = std::find(ids.begin(), ids.end(), id) - ids.begin();
|
||||
uint64_t v = (uint64_t)dev[i] + delta;
|
||||
memcpy(args.data() + off, &v, 8);
|
||||
}
|
||||
}
|
||||
{
|
||||
hipDeviceptr_t bz;
|
||||
size_t bzs = 0;
|
||||
uint32_t v = 0;
|
||||
if (hipModuleGetGlobal(&bz, &bzs, mod, "d4r_block_z") == hipSuccess && bzs == 4 &&
|
||||
hipMemcpyDtoH(&v, bz, 4) == hipSuccess && v != 0)
|
||||
block[2] = v;
|
||||
}
|
||||
{
|
||||
hipDeviceptr_t gp;
|
||||
size_t gs = 0;
|
||||
uint32_t v = 0;
|
||||
if (hipModuleGetGlobal(&gp, &gs, mod, "d4r_grid_x") == hipSuccess && gs == 4 &&
|
||||
hipMemcpyDtoH(&v, gp, 4) == hipSuccess && v != 0)
|
||||
{
|
||||
grid[0] = v;
|
||||
grid[1] = grid[2] = 1;
|
||||
}
|
||||
}
|
||||
if (getenv("GRID_X"))
|
||||
grid[0] = atoi(getenv("GRID_X"));
|
||||
if (getenv("GRID_Y"))
|
||||
grid[1] = atoi(getenv("GRID_Y"));
|
||||
size_t argSize = args.size();
|
||||
void* cfg[] = {HIP_LAUNCH_PARAM_BUFFER_POINTER, args.data(), HIP_LAUNCH_PARAM_BUFFER_SIZE, &argSize,
|
||||
HIP_LAUNCH_PARAM_END};
|
||||
const bool prepOnce = getenv("PREP_ONCE") != nullptr;
|
||||
int launches = 0;
|
||||
auto launch = [&]() {
|
||||
if (prep != nullptr && !(prepOnce && launches++ > 0))
|
||||
{
|
||||
hipError_t e = hipModuleLaunchKernel(prep, prepGrid, 1, 1, 128, 1, 1, 0, nullptr, nullptr, cfg);
|
||||
if (e != hipSuccess)
|
||||
return e;
|
||||
}
|
||||
return hipModuleLaunchKernel(fn, grid[0], grid[1], grid[2], block[0], block[1], block[2], shared, nullptr,
|
||||
nullptr, cfg);
|
||||
};
|
||||
CHECK(launch());
|
||||
CHECK(hipDeviceSynchronize());
|
||||
for (size_t i = 0; i < ids.size(); ++i)
|
||||
{
|
||||
std::vector<char> h(sizes[i]);
|
||||
CHECK(hipMemcpy(h.data(), dev[i], sizes[i], hipMemcpyDeviceToHost));
|
||||
std::ofstream(out + "/alloc-" + std::to_string(ids[i]) + ".bin", std::ios::binary).write(h.data(), h.size());
|
||||
}
|
||||
if (const char* dbg = getenv("DBG_GLOBAL"))
|
||||
{
|
||||
hipDeviceptr_t ptr;
|
||||
size_t size = 0;
|
||||
if (hipModuleGetGlobal(&ptr, &size, mod, dbg) == hipSuccess)
|
||||
{
|
||||
std::vector<char> h(size);
|
||||
CHECK(hipMemcpyDtoH(h.data(), ptr, size));
|
||||
std::ofstream(out + "/dbg.bin", std::ios::binary).write(h.data(), h.size());
|
||||
}
|
||||
}
|
||||
hipEvent_t a, b;
|
||||
CHECK(hipEventCreate(&a));
|
||||
CHECK(hipEventCreate(&b));
|
||||
std::vector<float> ms;
|
||||
for (int i = 0; i < iters; ++i)
|
||||
{
|
||||
CHECK(hipEventRecord(a));
|
||||
CHECK(launch());
|
||||
CHECK(hipEventRecord(b));
|
||||
CHECK(hipEventSynchronize(b));
|
||||
float t;
|
||||
CHECK(hipEventElapsedTime(&t, a, b));
|
||||
ms.push_back(t);
|
||||
}
|
||||
std::sort(ms.begin(), ms.end());
|
||||
if (!ms.empty())
|
||||
printf("%s: median %.4f ms, min %.4f ms over %d launches\n", argv[2], ms[ms.size() / 2], ms[0], iters);
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract the PTX modules embedded in NVIDIA's nvngx_dlss.dll.
|
||||
|
||||
usage: extract_dlss_ptx.py NVNGX_DLSS_DLL OUT_DIR
|
||||
|
||||
Writes OUT_DIR/dlss-NNNN-MM.ptx for PTX entry MM of the NNNN-th CUDA fatbin found in the DLL (in file
|
||||
order). The texture-kernel build (kernels/tex) edits some of these modules; the DLL is NVIDIA's and is
|
||||
never part of this repository, and neither are the extracted files.
|
||||
"""
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
|
||||
FATBIN_MAGIC = 0xBA55ED50
|
||||
|
||||
|
||||
def lz4_block(src: bytes, size: int) -> bytes:
|
||||
dst = bytearray()
|
||||
i = 0
|
||||
while i < len(src):
|
||||
token = src[i]
|
||||
i += 1
|
||||
literals = token >> 4
|
||||
if literals == 15:
|
||||
while True:
|
||||
b = src[i]
|
||||
i += 1
|
||||
literals += b
|
||||
if b != 255:
|
||||
break
|
||||
dst += src[i:i + literals]
|
||||
i += literals
|
||||
if i >= len(src) or len(dst) >= size:
|
||||
break
|
||||
offset = src[i] | (src[i + 1] << 8)
|
||||
i += 2
|
||||
match = token & 15
|
||||
if match == 15:
|
||||
while True:
|
||||
b = src[i]
|
||||
i += 1
|
||||
match += b
|
||||
if b != 255:
|
||||
break
|
||||
match += 4
|
||||
for _ in range(match):
|
||||
dst.append(dst[-offset])
|
||||
return bytes(dst[:size])
|
||||
|
||||
|
||||
def fatbins(data: bytes):
|
||||
"""Yields (offset, bytes) of every fatbin container (magic, version 1) in file order."""
|
||||
magic = struct.pack('<I', FATBIN_MAGIC)
|
||||
start = data.find(magic)
|
||||
while start >= 0:
|
||||
_, version, header_size, size = struct.unpack_from('<IHHQ', data, start)
|
||||
if version == 1 and header_size == 16 and start + header_size + size <= len(data):
|
||||
yield start, data[start:start + header_size + size]
|
||||
start = data.find(magic, start + header_size + size)
|
||||
else:
|
||||
start = data.find(magic, start + 4)
|
||||
|
||||
|
||||
def ptx_entries(fatbin: bytes):
|
||||
"""Yields the (decompressed) PTX text of each PTX entry of one fatbin."""
|
||||
_, _, header_size, size = struct.unpack_from('<IHHQ', fatbin, 0)
|
||||
offset = header_size
|
||||
while offset < header_size + size:
|
||||
kind, _, entry_header, entry_size = struct.unpack_from('<HHIQ', fatbin, offset)
|
||||
flags = struct.unpack_from('<Q', fatbin, offset + 40)[0]
|
||||
decompressed = struct.unpack_from('<Q', fatbin, offset + 56)[0] if entry_header >= 64 else 0
|
||||
payload = fatbin[offset + entry_header:offset + entry_header + entry_size]
|
||||
if kind == 1: # PTX
|
||||
yield lz4_block(payload, decompressed) if (flags & 0x2000 and decompressed) else payload
|
||||
offset += entry_header + entry_size
|
||||
|
||||
|
||||
def main(argv):
|
||||
if len(argv) != 3:
|
||||
sys.exit(__doc__)
|
||||
data = open(argv[1], 'rb').read()
|
||||
os.makedirs(argv[2], exist_ok=True)
|
||||
total = 0
|
||||
for index, (_, fatbin) in enumerate(fatbins(data)):
|
||||
for entry, ptx in enumerate(ptx_entries(fatbin)):
|
||||
ptx = ptx.rstrip(b'\0')
|
||||
with open(os.path.join(argv[2], f'dlss-{index:04d}-{entry:02d}.ptx'), 'wb') as out:
|
||||
out.write(ptx)
|
||||
total += 1
|
||||
print(f'{total} PTX modules written to {argv[2]}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(sys.argv)
|
||||
@@ -0,0 +1,69 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Write the d4r-kernels.txt manifest of a native kernel directory.
|
||||
|
||||
usage: kernel_manifest.py KERNEL_DIR NVNGX_DLSS_DLL [NVNGX_DLSS_DLL...]
|
||||
|
||||
For every NAME.hsaco in KERNEL_DIR, finds the PTX module of each DLL that defines `.entry NAME` and
|
||||
lists the FNV-1a 64 hash of the module's text (trailing NUL bytes removed), one "NAME HASH" line per
|
||||
distinct hash. The Wine CUDA bridge serves a native kernel only while DLSS loads a module with a
|
||||
listed hash, so a DLSS version that changed the kernel runs ZLUDA's own compile of it instead.
|
||||
The manifest is written to KERNEL_DIR/d4r-kernels.txt; it holds hashes, not NVIDIA code.
|
||||
"""
|
||||
import hashlib
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from extract_dlss_ptx import fatbins, ptx_entries # noqa: E402
|
||||
|
||||
ENTRY = re.compile(rb'\.entry[ \t\n]+([A-Za-z0-9_$]+)[ \t\n\r]*\(')
|
||||
|
||||
|
||||
def fnv1a64(data: bytes) -> int:
|
||||
value = 0xcbf29ce484222325
|
||||
for byte in data:
|
||||
value = ((value ^ byte) * 0x100000001b3) & 0xffffffffffffffff
|
||||
return value
|
||||
|
||||
|
||||
def module_hashes(dll: str) -> dict:
|
||||
"""NAME -> set of module hashes for every PTX entry of one DLL."""
|
||||
data = open(dll, 'rb').read()
|
||||
found = {}
|
||||
for _, fatbin in fatbins(data):
|
||||
for ptx in ptx_entries(fatbin):
|
||||
text = ptx.rstrip(b'\0')
|
||||
value = fnv1a64(text)
|
||||
for match in ENTRY.finditer(text):
|
||||
found.setdefault(match.group(1).decode(), set()).add(value)
|
||||
return found
|
||||
|
||||
|
||||
def main(argv):
|
||||
if len(argv) < 3:
|
||||
sys.exit(__doc__)
|
||||
directory = argv[1]
|
||||
names = sorted(name[:-len('.hsaco')] for name in os.listdir(directory) if name.endswith('.hsaco'))
|
||||
lines = ['# d4r native kernels: NAME and the FNV-1a 64 hash of the DLSS PTX module each was written for',
|
||||
'# (written by kernels/tools/kernel_manifest.py from:']
|
||||
listed = {name: set() for name in names}
|
||||
for dll in argv[2:]:
|
||||
digest = hashlib.sha256(open(dll, 'rb').read()).hexdigest()[:16]
|
||||
lines.append(f'# {os.path.basename(dll)} sha256 {digest}...)')
|
||||
hashes = module_hashes(dll)
|
||||
for name in names:
|
||||
listed[name] |= hashes.get(name, set())
|
||||
missing = [name for name in names if not listed[name]]
|
||||
if missing:
|
||||
sys.exit(f'no DLL defines {", ".join(missing)}; not writing a manifest')
|
||||
for name in names:
|
||||
for value in sorted(listed[name]):
|
||||
lines.append(f'{name} {value:016x}')
|
||||
with open(os.path.join(directory, 'd4r-kernels.txt'), 'w') as out:
|
||||
out.write('\n'.join(lines) + '\n')
|
||||
print(f'{len(names)} kernels listed in {os.path.join(directory, "d4r-kernels.txt")}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main(sys.argv)
|
||||
@@ -0,0 +1,13 @@
|
||||
#!/usr/bin/env python3
|
||||
# usage: kprof.py RUN.log FRAMES [WARM] : per-kernel median GPU ms and per-frame total from KERNEL_PROFILE lines
|
||||
import sys,re,statistics as st
|
||||
lines=[l for l in open(sys.argv[1],errors='replace') if 'KERNEL_PROFILE' in l]
|
||||
frames=int(sys.argv[2]); warm=int(sys.argv[3]) if len(sys.argv)>3 else frames//2
|
||||
rows=[(int(m.group(1)),float(m.group(2)),m.group(3)) for m in (re.search(r'launch=(\d+).*gpu_ms=([0-9.]+) name=(\S+)',l) for l in lines) if m]
|
||||
rows.sort(); n=len(rows); per=n/frames; cut=int(per*warm); rows=rows[cut:]; f=frames-warm
|
||||
agg={}
|
||||
for _,ms,name in rows: agg.setdefault(name,[]).append(ms)
|
||||
tot=sum(sum(v) for v in agg.values())/f
|
||||
print(f"launches/frame={per:.1f} total_gpu_ms/frame={tot:.3f} (frames {warm+1}-{frames})")
|
||||
for name,v in sorted(agg.items(),key=lambda kv:-sum(kv[1])):
|
||||
print(f" {sum(v)/f:8.3f} ms/frame {len(v)/f:5.1f}x med {st.median(v):7.3f} {name}")
|
||||
@@ -0,0 +1,141 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Numeric primitives and replay-dump access for swin_model.py (the rrlite Swin blocks of DLSS 4.5, preset M).
|
||||
|
||||
Everything here is NVIDIA's arithmetic as the native kernels reproduce it (kernels/m/swin_common.h,
|
||||
swin_block.h): f16 rounding after every operation, e4m3 (OCP, satfinite, RNE) re-quantisation, one f16 rounding
|
||||
per k32 MMA step, and the softmax's bit-level exponent trick.
|
||||
"""
|
||||
import struct
|
||||
|
||||
import numpy as np
|
||||
|
||||
F16, F32, F64 = np.float16, np.float32, np.float64
|
||||
|
||||
|
||||
def f16(x):
|
||||
"""round to f16; the result stays float64 (on the f16 grid), so products and sums of such values are exact
|
||||
until the next f16() (numpy would otherwise do f16 arithmetic, e.g. in matmul)"""
|
||||
return np.asarray(x, F64).astype(F16).astype(F64)
|
||||
|
||||
|
||||
def _e4m3_table():
|
||||
t = np.zeros(256, F64)
|
||||
for c in range(256):
|
||||
s, e, m = c >> 7, (c >> 3) & 15, c & 7
|
||||
if e == 15 and m == 7:
|
||||
v = np.nan
|
||||
elif e == 0:
|
||||
v = m * 2.0 ** -9
|
||||
else:
|
||||
v = (1 + m / 8) * 2.0 ** (e - 7)
|
||||
t[c] = -v if s else v
|
||||
return t
|
||||
|
||||
|
||||
E4M3 = _e4m3_table()
|
||||
_POS = E4M3[:127] # codes 0..126: the finite non-negative values, ascending
|
||||
|
||||
|
||||
def q8(x):
|
||||
"""e4m3 RNE satfinite round trip of f16 values (NaN stays NaN), as f16"""
|
||||
h = f16(x).astype(F64)
|
||||
a = np.abs(h)
|
||||
i = np.clip(np.searchsorted(_POS, a), 1, 126)
|
||||
lo, hi = _POS[i - 1], _POS[i]
|
||||
# nearest; ties to the even code
|
||||
pick_hi = (hi - a < a - lo) | ((hi - a == a - lo) & ((i & 1) == 0))
|
||||
v = np.where(pick_hi, hi, lo)
|
||||
v = np.where(a >= _POS[126], _POS[126], v)
|
||||
v = np.where(a <= 0, 0.0, v)
|
||||
v = np.where(np.isnan(h), np.nan, np.copysign(v, h))
|
||||
return f16(v)
|
||||
|
||||
|
||||
# softmax constants (swin_block.h stage 2)
|
||||
C_SCALE = f16(0.0972222164273262 / 5.656854152679443)
|
||||
CLAMP = F16(0.55615234375)
|
||||
P1 = f16(0.92730712890625)
|
||||
P0 = F16(1.375)
|
||||
|
||||
|
||||
def exp_trick(poly):
|
||||
"""per row, columns (2j, 2j+1): the 32-bit word even | odd << 16 of the f16 bits, << 5, + 0x7FF88000; the even
|
||||
column takes its low half, the odd column its high half (as f16 bits)"""
|
||||
b = poly.astype(F16).view(np.uint16).astype(np.uint64)
|
||||
out = np.empty_like(b)
|
||||
even, odd = b[:, 0::2], b[:, 1::2]
|
||||
word = ((even | (odd << np.uint64(16))) << np.uint64(5)) + np.uint64(0x7FF88000)
|
||||
word &= np.uint64(0xFFFFFFFF)
|
||||
out[:, 0::2] = word & np.uint64(0xFFFF)
|
||||
out[:, 1::2] = word >> np.uint64(16)
|
||||
return out.astype(np.uint16).view(F16)
|
||||
|
||||
|
||||
def row_sum16(w):
|
||||
"""f16 butterfly over the 16 columns: partner c ^ 8, then ^ 2, ^ 4, ^ 1 (every column ends with the same sum)"""
|
||||
s = w.astype(F16)
|
||||
c = np.arange(16)
|
||||
for m in (8, 2, 4, 1):
|
||||
s = f16(s.astype(F64) + s[:, c ^ m].astype(F64))
|
||||
return s
|
||||
|
||||
|
||||
def gelu_poly(z):
|
||||
z = f16(z)
|
||||
cz = np.clip(z, F16(-2.0), F16(2.0))
|
||||
inner = f16(f16(0.41216981).astype(F64) - f16(f16(0.08108133).astype(F64) * np.abs(cz).astype(F64)).astype(F64))
|
||||
return f16(z.astype(F64) * f16(0.5 + f16(cz.astype(F64) * inner.astype(F64)).astype(F64)).astype(F64))
|
||||
|
||||
|
||||
def mma_chain(A, W, C0=None):
|
||||
"""chained k32 MMA steps with f16 accumulation: D = f16(A[:, 32k:32k+32] W[32k:32k+32] + D), exact sums"""
|
||||
M, K = A.shape
|
||||
D = np.zeros((M, W.shape[1]), F64) if C0 is None else f16(C0)
|
||||
A = np.asarray(A, F64)
|
||||
W = np.asarray(W, F64)
|
||||
for k in range(0, K, 32):
|
||||
D = f16(A[:, k:k + 32] @ W[k:k + 32] + D.astype(F64))
|
||||
return D
|
||||
|
||||
|
||||
class Params:
|
||||
"""one rrlite Swin-block launch from a replay dump (CommonParams: w 0, sx sy 8, tw th 16, in 24, p32 32,
|
||||
out 40, p48 48)"""
|
||||
|
||||
def __init__(self, d):
|
||||
man = open(d + '/manifest.txt').read().split('\n')
|
||||
self.kernel = man[0].split()[1]
|
||||
la = man[1].split()
|
||||
self.grid = tuple(map(int, la[1:4]))
|
||||
self.args = open(d + '/args.bin', 'rb').read()
|
||||
self.allocs = {}
|
||||
for line in man:
|
||||
p = line.split()
|
||||
if p and p[0] == 'alloc':
|
||||
self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
|
||||
self.w_ptr, = struct.unpack_from('<Q', self.args, 0)
|
||||
self.sx, self.sy, self.W, self.H = struct.unpack_from('<iiii', self.args, 8)
|
||||
self.inp, self.p32, self.out, self.p48 = struct.unpack_from('<QQQQ', self.args, 24)
|
||||
wbase, wbuf = self.alloc_of(self.w_ptr)
|
||||
self.wbuf = wbuf[self.w_ptr - wbase:]
|
||||
self.abase, self.arena = self.alloc_of(self.inp)
|
||||
|
||||
def alloc_of(self, a):
|
||||
for base, buf in self.allocs.values():
|
||||
if base <= a < base + buf.size:
|
||||
return base, buf
|
||||
raise ValueError(hex(a))
|
||||
|
||||
def f16vec(self, off, n):
|
||||
return self.wbuf[off:off + 2 * n].view(F16).copy()
|
||||
|
||||
def token_xy(self, bx, by):
|
||||
T = np.arange(64)
|
||||
tx = (T & 3) + 4 * ((T >> 4) & 1)
|
||||
ty = ((T >> 2) & 3) + 4 * (T >> 5)
|
||||
return 8 * bx - self.sx + tx, 8 * by - self.sy + ty
|
||||
|
||||
@staticmethod
|
||||
def mirror(v, n):
|
||||
v = np.abs(v)
|
||||
return np.minimum(v, 2 * n - 2 - v)
|
||||
@@ -0,0 +1,15 @@
|
||||
#!/usr/bin/env python3
|
||||
"""psnr.py DIR_A DIR_B: per-frame PSNR of rgba16f outputs (watermark box excluded, values clipped to [0, 1] after tonemap-free clamp)."""
|
||||
import sys, glob, os, numpy as np
|
||||
a, b = sys.argv[1:3]
|
||||
vals = []
|
||||
for f in sorted(glob.glob(a + '/frame-*.rgba16f')):
|
||||
g = os.path.join(b, os.path.basename(f))
|
||||
if not os.path.exists(g): continue
|
||||
x = np.fromfile(f, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
|
||||
y = np.fromfile(g, np.float16).reshape(1440, 2560, 4)[..., :3].astype(np.float64)
|
||||
m = np.ones(x.shape[:2], bool); m[1370:1392, 1850:2460] = False
|
||||
x, y = np.clip(x[m], 0, 1), np.clip(y[m], 0, 1)
|
||||
mse = np.mean((x - y) ** 2)
|
||||
vals.append(10 * np.log10(1.0 / mse) if mse > 0 else 99.0)
|
||||
print('frames', len(vals), 'PSNR dB min %.2f mean %.2f' % (min(vals), np.mean(vals)), ' '.join('%.1f' % v for v in vals))
|
||||
@@ -0,0 +1,20 @@
|
||||
// usage: ptxload FILE.ptx... : cuModuleLoadData + cuModuleGetFunction for every .entry
|
||||
#include <dlfcn.h>
|
||||
#include <stdio.h>
|
||||
#include <stdlib.h>
|
||||
#include <string.h>
|
||||
typedef int (*I)(unsigned); typedef int (*DG)(int*,int); typedef int (*CC)(void**,unsigned,int);
|
||||
typedef int (*ML)(void**,const void*); typedef int (*GF)(void**,void*,const char*); typedef int (*MU)(void*);
|
||||
int main(int argc,char**argv){
|
||||
void*l=dlopen("libcuda.so",RTLD_NOW); if(!l){fprintf(stderr,"%s\n",dlerror());return 1;}
|
||||
I init=dlsym(l,"cuInit"); DG dg=dlsym(l,"cuDeviceGet"); CC cc=dlsym(l,"cuCtxCreate_v2");
|
||||
ML ml=dlsym(l,"cuModuleLoadData"); GF gf=dlsym(l,"cuModuleGetFunction"); MU mu=dlsym(l,"cuModuleUnload");
|
||||
int d; void*ctx; init(0); dg(&d,0); cc(&ctx,0,d); int bad=0;
|
||||
for(int a=1;a<argc;a++){
|
||||
FILE*f=fopen(argv[a],"rb"); fseek(f,0,2); long n=ftell(f); rewind(f); char*s=malloc(n+1); fread(s,1,n,f); s[n]=0; fclose(f);
|
||||
void*m=0; int r=ml(&m,s); printf("%s load=%d",argv[a],r); if(r){bad++;puts("");free(s);continue;}
|
||||
for(char*p=strstr(s,".entry ");p;p=strstr(p+1,".entry ")){
|
||||
char name[256]; sscanf(p+7,"%255[A-Za-z0-9_$]",name); void*fn=0; int g=gf(&fn,m,name);
|
||||
printf(" %s=%d",name,g); if(g) bad++; }
|
||||
puts(""); mu(m); free(s);}
|
||||
return bad?1:0;}
|
||||
@@ -0,0 +1,883 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Minimal vectorised PTX interpreter for DLSS rrlite kernels.
|
||||
|
||||
Runs one thread block of a fully unrolled kernel on the CPU (all threads of the
|
||||
block as numpy lanes), against memory restored from a bridge replay dump
|
||||
(D4R_CUDA_REPLAY_DUMP_DIR). It is a reference and an instrument: every
|
||||
statement can be observed (hooks), which is how the logical computation
|
||||
(GEMMs, weight layouts, normalisations) of a kernel is recovered.
|
||||
|
||||
Only the instruction forms the rrlite kernels use are implemented; anything
|
||||
else raises NotImplementedError naming the statement.
|
||||
"""
|
||||
import re
|
||||
import struct
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
F16, F32, F64 = np.float16, np.float32, np.float64
|
||||
U64 = np.uint64
|
||||
PARAM_BASE = 1 << 60 # virtual address of the kernel parameter block
|
||||
MASK = {8: 0xff, 16: 0xffff, 32: 0xffffffff, 64: 0xffffffffffffffff}
|
||||
|
||||
|
||||
def _e4m3_decode_table():
|
||||
t = np.zeros(256, F64)
|
||||
for c in range(256):
|
||||
s = -1.0 if c & 0x80 else 1.0
|
||||
e, m = (c >> 3) & 0xf, c & 7
|
||||
if (c & 0x7f) == 0x7f:
|
||||
t[c] = np.nan
|
||||
elif e == 0:
|
||||
t[c] = s * m * 2.0 ** -9
|
||||
else:
|
||||
t[c] = s * (1 + m / 8) * 2.0 ** (e - 7)
|
||||
return t
|
||||
|
||||
|
||||
E4M3 = _e4m3_decode_table()
|
||||
|
||||
|
||||
def _f16_to_e4m3_table():
|
||||
"""Exact RNE + satfinite f16 -> e4m3 for all 65536 f16 codes (ZLUDA's helper)."""
|
||||
h = np.arange(65536, dtype=np.uint32)
|
||||
mag = h & 0x7fff
|
||||
normal = (mag - 0x2000 + 0x3f + ((mag >> 7) & 1)) >> 7
|
||||
exp = mag >> 10
|
||||
sig = (mag & 0x3ff) | np.where(exp != 0, 0x400, 0)
|
||||
cl = np.clip(exp, 1, 8)
|
||||
sh = 16 - cl
|
||||
sub = (sig + (np.left_shift(1, sh - 1)) - 1 + ((sig >> sh) & 1)) >> sh
|
||||
code = np.where(mag >= 0x2400, normal, sub)
|
||||
code = np.where(mag > 0x5f00, 0x7e, code)
|
||||
code = np.where(mag > 0x7c00, 0x7f, code)
|
||||
return (((h >> 8) & 0x80) | code).astype(np.uint32)
|
||||
|
||||
|
||||
F16_TO_E4M3 = _f16_to_e4m3_table()
|
||||
|
||||
|
||||
def bits_to_f16(x):
|
||||
return (x & 0xffff).astype(np.uint16).view(F16)
|
||||
|
||||
|
||||
def f16_to_bits(v):
|
||||
return np.asarray(v, F16).view(np.uint16).astype(U64)
|
||||
|
||||
|
||||
def bits_to_f32(x):
|
||||
return (x & 0xffffffff).astype(np.uint32).view(F32)
|
||||
|
||||
|
||||
def f32_to_bits(v):
|
||||
return np.asarray(v, F32).view(np.uint32).astype(U64)
|
||||
|
||||
|
||||
class Stmt:
|
||||
__slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index')
|
||||
|
||||
|
||||
def split_args(s):
|
||||
out, depth, cur = [], 0, ''
|
||||
for ch in s:
|
||||
if ch in '{[':
|
||||
depth += 1
|
||||
elif ch in '}]':
|
||||
depth -= 1
|
||||
if ch == ',' and depth == 0:
|
||||
out.append(cur.strip())
|
||||
cur = ''
|
||||
else:
|
||||
cur += ch
|
||||
if cur.strip():
|
||||
out.append(cur.strip())
|
||||
return out
|
||||
|
||||
|
||||
def parse_kernel(path, name):
|
||||
with open(path) as file:
|
||||
src = file.read()
|
||||
i = src.index('.entry ' + name + '(')
|
||||
j = src.find('\n.visible', i + 10)
|
||||
j = len(src) if j < 0 else j
|
||||
body = src[i:j]
|
||||
body = body[body.index(')') + 1:]
|
||||
# module-scope .shared arrays (declared before the entry) are visible to it too
|
||||
module_shared = {m.group(1): int(m.group(2)) for m in re.finditer(r'^\.shared[^\n;]*?(\w+)\[(\d+)\]', src[:i], re.M)}
|
||||
body = re.sub(r'//[^\n]*', '', body)
|
||||
body = re.sub(r'\.(maxntid|reqntid|minnctapersm|maxnreg)[^\n{]*', '', body)
|
||||
stmts, types, shared = [], {}, {}
|
||||
pending_labels = []
|
||||
for raw in re.sub(r'\s*\n\s*', ' ', body).split(';'):
|
||||
s = raw.strip()
|
||||
while True:
|
||||
s = s.strip()
|
||||
if s.startswith('{') and not s.startswith('{%'):
|
||||
s = s[1:]
|
||||
continue
|
||||
if s.startswith('}'):
|
||||
s = s[1:]
|
||||
continue
|
||||
m = re.match(r'(\$[\w]+):\s*', s)
|
||||
if m:
|
||||
pending_labels.append(m.group(1))
|
||||
s = s[m.end():]
|
||||
continue
|
||||
break
|
||||
if not s:
|
||||
continue
|
||||
if s.startswith('.reg'):
|
||||
m = re.match(r'\.reg\s*\.(\w+)\s+(.*)', s)
|
||||
ty = m.group(1)
|
||||
for decl in m.group(2).split(','):
|
||||
decl = decl.strip()
|
||||
m2 = re.match(r'(%[\w]+)<(\d+)>', decl)
|
||||
if m2:
|
||||
types[('prefix', m2.group(1))] = ty
|
||||
else:
|
||||
types[decl] = ty
|
||||
continue
|
||||
if s.startswith('.shared'):
|
||||
m = re.search(r'(\w+)\[(\d+)\]', s)
|
||||
shared[m.group(1)] = int(m.group(2))
|
||||
continue
|
||||
if s.startswith('.'):
|
||||
continue
|
||||
st = Stmt()
|
||||
st.labels, pending_labels = pending_labels, []
|
||||
st.pred, st.neg = None, False
|
||||
m = re.match(r'@(!?)(%\w+)\s+', s)
|
||||
if m:
|
||||
st.neg, st.pred = m.group(1) == '!', m.group(2)
|
||||
s = s[m.end():]
|
||||
parts = s.split(None, 1)
|
||||
st.op = parts[0]
|
||||
st.args = split_args(parts[1]) if len(parts) > 1 else []
|
||||
st.text = s
|
||||
st.index = len(stmts)
|
||||
stmts.append(st)
|
||||
if pending_labels:
|
||||
st = Stmt()
|
||||
st.labels, st.pred, st.neg, st.op, st.args, st.text, st.index = pending_labels, None, False, 'nop', [], 'nop', len(stmts)
|
||||
stmts.append(st)
|
||||
for k, v in module_shared.items():
|
||||
shared.setdefault(k, v)
|
||||
return stmts, types, shared
|
||||
|
||||
|
||||
def reg_width(types, name):
|
||||
ty = types.get(name)
|
||||
if ty is None:
|
||||
m = re.match(r'(%[a-z]+)', name)
|
||||
ty = types.get(('prefix', m.group(1)))
|
||||
return {'pred': 1, 'b16': 16, 'u16': 16, 's16': 16, 'f16': 16, 'b32': 32, 'u32': 32, 's32': 32, 'f32': 32,
|
||||
'b64': 64, 'u64': 64, 's64': 64, 'f64': 64, 'b8': 8}[ty]
|
||||
|
||||
|
||||
class Memory:
|
||||
"""Global memory: the replay dump's allocations at their captured addresses."""
|
||||
|
||||
def __init__(self, dump_dir):
|
||||
manifest = open(dump_dir + '/manifest.txt').read().split('\n')
|
||||
self.args = open(dump_dir + '/args.bin', 'rb').read()
|
||||
self.bases, self.data, self.ids = [], [], []
|
||||
for line in manifest:
|
||||
p = line.split()
|
||||
if p and p[0] == 'alloc':
|
||||
idx, base, size = int(p[1]), int(p[2], 16), int(p[3])
|
||||
buf = np.fromfile(f'{dump_dir}/alloc-{idx}.bin', np.uint8)
|
||||
assert buf.size == size
|
||||
self.bases.append(base)
|
||||
self.data.append(buf)
|
||||
self.ids.append(idx)
|
||||
order = np.argsort(self.bases)
|
||||
self.bases = [self.bases[i] for i in order]
|
||||
self.data = [self.data[i] for i in order]
|
||||
self.ids = [self.ids[i] for i in order]
|
||||
self.base_arr = np.array(self.bases, U64)
|
||||
self.stores = [] # (addresses, bytes) of every global store
|
||||
|
||||
def locate(self, addr):
|
||||
idx = np.searchsorted(self.base_arr, addr, side='right') - 1
|
||||
if np.any(idx < 0):
|
||||
raise ValueError('address below every allocation')
|
||||
return idx
|
||||
|
||||
def load(self, addr, nbytes, mask):
|
||||
out = np.zeros((addr.size, nbytes), np.uint8)
|
||||
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
|
||||
for a in np.unique(idx[mask]):
|
||||
sel = mask & (idx == a)
|
||||
off = (addr[sel] - self.bases[a]).astype(np.int64)
|
||||
if np.any(off + nbytes > self.data[a].size):
|
||||
raise ValueError('load past the end of an allocation')
|
||||
out[sel] = self.data[a][off[:, None] + np.arange(nbytes)]
|
||||
return out
|
||||
|
||||
def store(self, addr, raw, mask):
|
||||
idx = self.locate(np.where(mask, addr, self.base_arr[0]))
|
||||
n = raw.shape[1]
|
||||
for a in np.unique(idx[mask]):
|
||||
sel = mask & (idx == a)
|
||||
off = (addr[sel] - self.bases[a]).astype(np.int64)
|
||||
self.data[a][off[:, None] + np.arange(n)] = raw[sel]
|
||||
self.stores.append((self.ids[a], off.copy(), raw[sel].copy()))
|
||||
|
||||
|
||||
class Block:
|
||||
def __init__(self, stmts, types, shared, mem, param_name, block_dim, grid_dim, ctaid):
|
||||
self.stmts, self.types, self.mem = stmts, types, mem
|
||||
self.param_name = param_name
|
||||
bx, by, bz = block_dim
|
||||
self.T = bx * by * bz
|
||||
t = np.arange(self.T)
|
||||
self.tid = {'x': t % bx, 'y': (t // bx) % by, 'z': t // (bx * by)}
|
||||
self.lane = t % 32
|
||||
self.warp = t // 32
|
||||
self.ctaid = ctaid
|
||||
self.grid = grid_dim
|
||||
self.ntid = block_dim
|
||||
self.regs = {}
|
||||
self.shared = {k: np.zeros(v, np.uint8) for k, v in shared.items()}
|
||||
self.shared_base = {}
|
||||
base = 0
|
||||
for k, v in shared.items():
|
||||
self.shared_base[k] = base
|
||||
base += (v + 15) & ~15
|
||||
self.shared_mem = np.zeros(max(base, 16), np.uint8)
|
||||
self.labels = {}
|
||||
for st in stmts:
|
||||
for l in st.labels:
|
||||
self.labels[l] = st.index
|
||||
self.hooks = []
|
||||
self.retired = np.zeros(self.T, bool)
|
||||
self.prov = {}
|
||||
|
||||
# --- operands -----------------------------------------------------------
|
||||
def width(self, name):
|
||||
return reg_width(self.types, name)
|
||||
|
||||
def special(self, name):
|
||||
m = re.match(r'%(tid|ntid|ctaid|nctaid)\.([xyz])', name)
|
||||
if m:
|
||||
kind, ax = m.groups()
|
||||
k = 'xyz'.index(ax)
|
||||
if kind == 'tid':
|
||||
return self.tid[ax].astype(U64)
|
||||
if kind == 'ntid':
|
||||
return np.full(self.T, self.ntid[k], U64)
|
||||
if kind == 'ctaid':
|
||||
return np.full(self.T, self.ctaid[k], U64)
|
||||
return np.full(self.T, self.grid[k], U64)
|
||||
if name == 'WARP_SZ':
|
||||
return np.full(self.T, 32, U64)
|
||||
if name == '%laneid':
|
||||
return self.lane.astype(U64)
|
||||
if name == '%warpid':
|
||||
return self.warp.astype(U64)
|
||||
return None
|
||||
|
||||
def imm(self, s, width=64):
|
||||
if s.startswith('0f'):
|
||||
return np.full(self.T, int(s[2:], 16), U64)
|
||||
if s.startswith('0d'):
|
||||
return np.full(self.T, int(s[2:], 16), U64)
|
||||
v = int(s, 0)
|
||||
return np.full(self.T, v & MASK[64], U64)
|
||||
|
||||
def get(self, s, width=None):
|
||||
s = s.strip()
|
||||
if s.startswith('%') or s in self.types or s == 'WARP_SZ':
|
||||
sp = self.special(s) if (s.startswith('%') or s == 'WARP_SZ') else None
|
||||
if sp is not None:
|
||||
return sp
|
||||
if s not in self.regs:
|
||||
w = self.width(s)
|
||||
self.regs[s] = np.zeros(self.T, bool if w == 1 else U64)
|
||||
return self.regs[s]
|
||||
if s in self.shared_base:
|
||||
return np.full(self.T, self.shared_base[s], U64)
|
||||
if s == self.param_name:
|
||||
return np.full(self.T, PARAM_BASE, U64)
|
||||
if s[0].isdigit() or s[0] == '-':
|
||||
v = self.imm(s)
|
||||
if width:
|
||||
v = v & np.uint64(MASK[width])
|
||||
return v
|
||||
raise NotImplementedError('operand ' + s)
|
||||
|
||||
def put(self, name, value, mask, keep_prov=False):
|
||||
name = name.strip()
|
||||
if not keep_prov:
|
||||
self.prov.pop(name, None)
|
||||
w = self.width(name)
|
||||
old = self.get(name)
|
||||
if w == 1:
|
||||
self.regs[name] = np.where(mask, value.astype(bool), old)
|
||||
else:
|
||||
self.regs[name] = np.where(mask, value.astype(U64) & np.uint64(MASK[w]), old)
|
||||
|
||||
def vec(self, s):
|
||||
return [x.strip() for x in s.strip()[1:-1].split(',')]
|
||||
|
||||
def address(self, s):
|
||||
"""[reg+off] or [symbol+off] -> (space, address array)."""
|
||||
inner = s.strip()[1:-1]
|
||||
m = re.match(r'([%\w$.]+)\s*(\+\s*-?\d+)?', inner)
|
||||
base, off = m.group(1), int(m.group(2).replace('+', '')) if m.group(2) else 0
|
||||
if base == self.param_name:
|
||||
return 'param', np.full(self.T, off, U64)
|
||||
if base in self.shared_base:
|
||||
return 'shared', np.full(self.T, self.shared_base[base] + off, U64)
|
||||
return 'reg', (self.get(base) + np.uint64(off & MASK[64])) & np.uint64(MASK[64])
|
||||
|
||||
# --- execution ------------------------------------------------------------
|
||||
def run(self, limit=None):
|
||||
stmts = self.stmts
|
||||
active = np.ones(self.T, bool)
|
||||
pending = {}
|
||||
pc = 0
|
||||
steps = 0
|
||||
while pc < len(stmts):
|
||||
st = stmts[pc]
|
||||
if pc in pending:
|
||||
active |= pending.pop(pc)
|
||||
if not active.any():
|
||||
# Reconverge at either a forward target or a loop fallthrough.
|
||||
targets = sorted(pending)
|
||||
if not targets:
|
||||
break
|
||||
pc = targets[0]
|
||||
continue
|
||||
active &= ~self.retired
|
||||
mask = active
|
||||
if st.pred is not None:
|
||||
p = self.get(st.pred)
|
||||
mask = active & (~p if st.neg else p)
|
||||
if st.op == 'bra' or st.op.startswith('bra.'):
|
||||
target = self.labels[st.args[0]]
|
||||
taking = mask
|
||||
if target <= pc:
|
||||
if (taking == active).all():
|
||||
pc = target
|
||||
continue
|
||||
if not taking.any():
|
||||
pc += 1
|
||||
continue
|
||||
# Lanes finishing a loop wait at its fallthrough while the
|
||||
# other lanes execute their remaining iterations. Their
|
||||
# registers and memory writes stay masked until reconvergence.
|
||||
pending[pc + 1] = pending.get(pc + 1, np.zeros(self.T, bool)) | (active & ~taking)
|
||||
active = taking.copy()
|
||||
pc = target
|
||||
continue
|
||||
active = active & ~taking
|
||||
if taking.any():
|
||||
pending[target] = pending.get(target, np.zeros(self.T, bool)) | taking
|
||||
pc += 1
|
||||
continue
|
||||
if mask.any() or st.op.startswith(('mma', 'shfl', 'movmatrix', 'bar')):
|
||||
self.execute(st, mask)
|
||||
for hook in self.hooks:
|
||||
hook(self, st, mask)
|
||||
pc += 1
|
||||
steps += 1
|
||||
if limit and steps >= limit:
|
||||
break
|
||||
return steps
|
||||
|
||||
def execute(self, st, mask):
|
||||
op, a = st.op, st.args
|
||||
parts = op.split('.')
|
||||
base = parts[0]
|
||||
fn = getattr(self, 'op_' + base, None)
|
||||
if fn is None:
|
||||
raise NotImplementedError(st.text)
|
||||
try:
|
||||
fn(st, parts, a, mask)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f'{type(e).__name__} {e} in: {st.text}') from e
|
||||
|
||||
# --- helpers --------------------------------------------------------------
|
||||
@staticmethod
|
||||
def sext(v, w):
|
||||
v = v.astype(np.int64) & ((1 << w) - 1) if w < 64 else v.view(np.int64)
|
||||
if w < 64:
|
||||
v = np.where(v >= (1 << (w - 1)), v - (1 << w), v)
|
||||
return v
|
||||
|
||||
def typed(self, s, ty):
|
||||
w = int(re.sub(r'\D', '', ty) or 32) if ty not in ('pred',) else 1
|
||||
v = self.get(s, w)
|
||||
if ty.startswith('s'):
|
||||
return self.sext(v, w)
|
||||
if ty.startswith('f'):
|
||||
if ty == 'f16':
|
||||
return bits_to_f16(v).astype(F64)
|
||||
if ty == 'f32':
|
||||
return bits_to_f32(v).astype(F64)
|
||||
return v.view(F64)
|
||||
return (v & np.uint64(MASK[w])).astype(np.int64) if w < 64 else v
|
||||
|
||||
def result(self, ty, value):
|
||||
if ty == 'f16':
|
||||
return f16_to_bits(np.asarray(value).astype(F16))
|
||||
if ty == 'f32':
|
||||
return f32_to_bits(np.asarray(value).astype(F32))
|
||||
if ty == 'f64':
|
||||
return np.asarray(value, F64).view(U64)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
return (np.asarray(value).astype(np.int64).astype(U64)) & np.uint64(MASK[w])
|
||||
|
||||
# --- instructions -----------------------------------------------------------
|
||||
def op_nop(self, st, parts, a, mask):
|
||||
pass
|
||||
|
||||
def op_bar(self, st, parts, a, mask):
|
||||
pass
|
||||
|
||||
def op_ret(self, st, parts, a, mask):
|
||||
self.retired |= mask
|
||||
|
||||
def op_neg(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80008000), mask)
|
||||
elif ty == 'f16':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x8000), mask)
|
||||
elif ty == 'f32':
|
||||
self.put(a[0], self.get(a[1]) ^ np.uint64(0x80000000), mask)
|
||||
else:
|
||||
self.put(a[0], self.result('u' + ty[1:], -self.typed(a[1], ty)), mask)
|
||||
|
||||
def op_cvta(self, st, parts, a, mask):
|
||||
self.put(a[0], self.get(a[1]), mask)
|
||||
|
||||
def op_mov(self, st, parts, a, mask):
|
||||
d, s = a
|
||||
if d.startswith('{'):
|
||||
names = self.vec(d)
|
||||
v = self.get(s)
|
||||
w = self.width(names[0])
|
||||
for i, n in enumerate(names):
|
||||
self.put(n, (v >> np.uint64(i * w)) & np.uint64(MASK[w]), mask)
|
||||
return
|
||||
if s.startswith('{'):
|
||||
names = self.vec(s)
|
||||
w = self.width(names[0])
|
||||
v = np.zeros(self.T, U64)
|
||||
for i, n in enumerate(names):
|
||||
v |= (self.get(n) & np.uint64(MASK[w])) << np.uint64(i * w)
|
||||
self.put(d, v, mask)
|
||||
return
|
||||
self.put(d, self.get(s), mask)
|
||||
if s in self.prov:
|
||||
self.prov[d] = self.prov[s]
|
||||
|
||||
def op_ld(self, st, parts, a, mask):
|
||||
space = 'global' if 'global' in parts else 'shared' if 'shared' in parts else 'param' if 'param' in parts else None
|
||||
ty = parts[-1]
|
||||
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
nbytes = w // 8
|
||||
kind, addr = self.address(a[1])
|
||||
total = nbytes * vec
|
||||
if kind == 'param':
|
||||
raw = np.frombuffer(self.mem.args, np.uint8)
|
||||
off = int(addr[0])
|
||||
data = np.tile(raw[off:off + total], (self.T, 1))
|
||||
elif space == 'param':
|
||||
# ld.param through a computed address (mov.b64 %rd, <param>; add; ld.param [%rd])
|
||||
raw = np.frombuffer(self.mem.args, np.uint8)
|
||||
off = np.where(mask, addr - np.uint64(PARAM_BASE), 0).astype(np.int64)
|
||||
data = raw[off[:, None] + np.arange(total)]
|
||||
elif kind == 'shared' or space == 'shared':
|
||||
off = addr.astype(np.int64)
|
||||
data = self.shared_mem[off[:, None] + np.arange(total)]
|
||||
else:
|
||||
data = self.mem.load(addr, total, mask)
|
||||
names = self.vec(a[0]) if a[0].startswith('{') else [a[0]]
|
||||
for i, n in enumerate(names):
|
||||
chunk = data[:, i * nbytes:(i + 1) * nbytes]
|
||||
v = np.zeros(self.T, U64)
|
||||
for b in range(nbytes):
|
||||
v |= chunk[:, b].astype(U64) << np.uint64(8 * b)
|
||||
if ty.startswith('s') and w < self.width(n):
|
||||
v = (self.sext(v, w).astype(np.int64)).astype(U64)
|
||||
self.put(n, v, mask)
|
||||
self.prov[n] = (st.index, kind if kind != 'reg' else space, addr + np.uint64(i * nbytes), nbytes)
|
||||
|
||||
def op_st(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
vec = next((int(p[1:]) for p in parts if re.fullmatch(r'v\d', p)), 1)
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
nbytes = w // 8
|
||||
kind, addr = self.address(a[0])
|
||||
names = self.vec(a[1]) if a[1].startswith('{') else [a[1]]
|
||||
raw = np.zeros((self.T, nbytes * len(names)), np.uint8)
|
||||
for i, n in enumerate(names):
|
||||
v = self.get(n, w)
|
||||
for b in range(nbytes):
|
||||
raw[:, i * nbytes + b] = ((v >> np.uint64(8 * b)) & np.uint64(0xff)).astype(np.uint8)
|
||||
if kind == 'shared' or 'shared' in parts:
|
||||
off = addr.astype(np.int64)[mask]
|
||||
self.shared_mem[off[:, None] + np.arange(raw.shape[1])] = raw[mask]
|
||||
else:
|
||||
self.mem.store(addr, raw, mask)
|
||||
|
||||
def int_binop(self, parts, a, mask, fn):
|
||||
ty = parts[-1]
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
self.put(a[0], self.result(ty if not ty.startswith('b') else 'u' + ty[1:], fn(x, y)), mask)
|
||||
|
||||
def op_add(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) + self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x + y)
|
||||
|
||||
def op_sub(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) - self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x - y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x - y)
|
||||
|
||||
def op_mul(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if 'wide' in parts:
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
self.put(a[0], (x * y).astype(np.int64).astype(U64), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty)), mask)
|
||||
elif ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y), mask)
|
||||
else:
|
||||
self.int_binop(parts, a, mask, lambda x, y: x * y)
|
||||
|
||||
def op_mad(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
x, y, z = self.typed(a[1], ty), self.typed(a[2], ty), self.typed(a[3], ty)
|
||||
self.put(a[0], self.result('u' + ty[1:], x * y + z), mask)
|
||||
|
||||
def h2(self, sa, sb, fn, sc=None):
|
||||
va, vb = self.get(sa), self.get(sb)
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
x = bits_to_f16(va >> np.uint64(16 * i)).astype(F64)
|
||||
y = bits_to_f16(vb >> np.uint64(16 * i)).astype(F64)
|
||||
if sc is not None:
|
||||
z = bits_to_f16(self.get(sc) >> np.uint64(16 * i)).astype(F64)
|
||||
r = fn(x, y, z)
|
||||
else:
|
||||
r = fn(x, y)
|
||||
out |= f16_to_bits(np.asarray(r).astype(F16)) << np.uint64(16 * i)
|
||||
return out
|
||||
|
||||
def op_fma(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3]), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty) + self.typed(a[3], ty)), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def minmax(self, parts, a, mask, is_max):
|
||||
ty = parts[-1]
|
||||
if ty == 'f16x2':
|
||||
f = np.fmax if is_max else np.fmin
|
||||
self.put(a[0], self.h2(a[1], a[2], f), mask)
|
||||
elif ty in ('f16', 'f32'):
|
||||
f = np.fmax if is_max else np.fmin
|
||||
self.put(a[0], self.result(ty, f(self.typed(a[1], ty), self.typed(a[2], ty))), mask)
|
||||
else:
|
||||
f = np.maximum if is_max else np.minimum
|
||||
self.int_binop(parts, a, mask, f)
|
||||
|
||||
def op_min(self, st, parts, a, mask):
|
||||
self.minmax(parts, a, mask, False)
|
||||
|
||||
def op_max(self, st, parts, a, mask):
|
||||
self.minmax(parts, a, mask, True)
|
||||
|
||||
def op_abs(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
if ty in ('f16x2', 'f16', 'f32'):
|
||||
m = {'f16x2': 0x7fff7fff, 'f16': 0x7fff, 'f32': 0x7fffffff}[ty]
|
||||
self.put(a[0], self.get(a[1]) & np.uint64(m), mask)
|
||||
return
|
||||
self.put(a[0], self.result('u' + ty[1:], np.abs(self.typed(a[1], ty))), mask)
|
||||
|
||||
def logic(self, a, mask, fn):
|
||||
if self.width(a[0]) == 1:
|
||||
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
|
||||
else:
|
||||
self.put(a[0], fn(self.get(a[1]), self.get(a[2])), mask)
|
||||
|
||||
def op_and(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x & y)
|
||||
|
||||
def op_or(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x | y)
|
||||
|
||||
def op_xor(self, st, parts, a, mask):
|
||||
self.logic(a, mask, lambda x, y: x ^ y)
|
||||
|
||||
def op_not(self, st, parts, a, mask):
|
||||
self.put(a[0], ~self.get(a[1]), mask)
|
||||
|
||||
def op_shl(self, st, parts, a, mask):
|
||||
w = int(re.sub(r'\D', '', parts[-1]))
|
||||
n = self.get(a[2]) & np.uint64(0xff)
|
||||
v = np.where(n >= w, 0, (self.get(a[1]) << np.minimum(n, np.uint64(63))))
|
||||
self.put(a[0], v.astype(U64), mask)
|
||||
|
||||
def op_shr(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
w = int(re.sub(r'\D', '', ty))
|
||||
n = (self.get(a[2]) & np.uint64(0xff)).astype(np.int64)
|
||||
if ty.startswith('s'):
|
||||
v = self.sext(self.get(a[1]), w)
|
||||
r = v >> np.minimum(n, w - 1)
|
||||
self.put(a[0], r.astype(np.int64).astype(U64), mask)
|
||||
else:
|
||||
v = self.get(a[1]) & np.uint64(MASK[w])
|
||||
r = np.where(n >= w, 0, v >> np.minimum(n, 63).astype(U64))
|
||||
self.put(a[0], r.astype(U64), mask)
|
||||
|
||||
def op_bfi(self, st, parts, a, mask):
|
||||
f, b = self.get(a[1]), self.get(a[2])
|
||||
pos = self.get(a[3]).astype(np.int64) & 0xff
|
||||
ln = self.get(a[4]).astype(np.int64) & 0xff
|
||||
w = int(re.sub(r'\D', '', parts[-1]))
|
||||
out = np.array(b, U64)
|
||||
for i in range(self.T):
|
||||
p, l = int(pos[i]), int(ln[i])
|
||||
if l == 0 or p >= w:
|
||||
continue
|
||||
l = min(l, w - p)
|
||||
m = ((1 << l) - 1) << p
|
||||
out[i] = (int(b[i]) & ~m & MASK[w]) | ((int(f[i]) << p) & m)
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_setp(self, st, parts, a, mask):
|
||||
cmp, ty = parts[1], parts[-1]
|
||||
x, y = self.typed(a[1], ty), self.typed(a[2], ty)
|
||||
r = {'lt': x < y, 'le': x <= y, 'gt': x > y, 'ge': x >= y, 'eq': x == y, 'ne': x != y,
|
||||
'lo': x < y, 'ls': x <= y, 'hi': x > y, 'hs': x >= y}[cmp]
|
||||
dests = a[0].split('|')
|
||||
self.put(dests[0], r, mask)
|
||||
|
||||
def op_selp(self, st, parts, a, mask):
|
||||
p = self.get(a[3])
|
||||
self.put(a[0], np.where(p, self.get(a[1]), self.get(a[2])), mask)
|
||||
|
||||
def op_ex2(self, st, parts, a, mask):
|
||||
ty = parts[-1]
|
||||
v = self.get(a[1])
|
||||
if ty == 'f16x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
h = bits_to_f16((v >> np.uint64(16 * i)) & np.uint64(0xffff)).astype(F64)
|
||||
out |= f16_to_bits(np.exp2(h).astype(F16)) << np.uint64(16 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif ty == 'f16':
|
||||
self.put(a[0], f16_to_bits(np.exp2(bits_to_f16(v).astype(F64)).astype(F16)), mask)
|
||||
elif ty == 'f32':
|
||||
self.put(a[0], f32_to_bits(np.exp2(bits_to_f32(v).astype(F64)).astype(F32)), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def op_rsqrt(self, st, parts, a, mask):
|
||||
x = self.typed(a[1], 'f32').astype(F32)
|
||||
with np.errstate(divide='ignore', invalid='ignore'):
|
||||
self.put(a[0], f32_to_bits((1.0 / np.sqrt(x.astype(F64))).astype(F32)), mask)
|
||||
|
||||
def op_rcp(self, st, parts, a, mask):
|
||||
x = self.typed(a[1], 'f32')
|
||||
with np.errstate(divide='ignore'):
|
||||
self.put(a[0], f32_to_bits((1.0 / x).astype(F32)), mask)
|
||||
|
||||
def op_div(self, st, parts, a, mask):
|
||||
x, y = self.typed(a[1], 'f32'), self.typed(a[2], 'f32')
|
||||
with np.errstate(divide='ignore', invalid='ignore'):
|
||||
self.put(a[0], f32_to_bits((x / y).astype(F32)), mask)
|
||||
|
||||
def op_cvt(self, st, parts, a, mask):
|
||||
types = [p for p in parts[1:] if p not in ('rn', 'rz', 'rm', 'rp', 'rni', 'rzi', 'sat', 'satfinite', 'relu', 'ftz')]
|
||||
dt, stp = types[0], types[1]
|
||||
relu = 'relu' in parts
|
||||
v = self.get(a[1])
|
||||
if dt == 'e4m3x2' and stp == 'f16x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
h = (v >> np.uint64(16 * i)) & np.uint64(0xffff)
|
||||
if relu:
|
||||
f = bits_to_f16(h)
|
||||
h = np.where((f < 0) & ~np.isnan(f.astype(F32)), 0, h).astype(U64)
|
||||
out |= F16_TO_E4M3[h.astype(np.int64)].astype(U64) << np.uint64(8 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif dt == 'f16x2' and stp == 'e4m3x2':
|
||||
out = np.zeros(self.T, U64)
|
||||
for i in range(2):
|
||||
c = ((v >> np.uint64(8 * i)) & np.uint64(0xff)).astype(np.int64)
|
||||
out |= f16_to_bits(E4M3[c].astype(F16)) << np.uint64(16 * i)
|
||||
self.put(a[0], out, mask)
|
||||
elif dt == 'f16' and stp == 'f32':
|
||||
self.put(a[0], f16_to_bits(bits_to_f32(v).astype(F16)), mask)
|
||||
elif dt == 'f16' and stp == 'f64':
|
||||
self.put(a[0], f16_to_bits(v.view(F64).astype(F16)), mask)
|
||||
elif dt == 'f32' and stp == 'f16':
|
||||
self.put(a[0], f32_to_bits(bits_to_f16(v).astype(F32)), mask)
|
||||
elif dt in ('f16', 'f32') and stp[0] in 'us':
|
||||
ws = int(re.sub(r'\D', '', stp))
|
||||
iv = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
|
||||
if dt == 'f16':
|
||||
self.put(a[0], f16_to_bits(iv.astype(F64).astype(F16)), mask)
|
||||
else:
|
||||
self.put(a[0], f32_to_bits(iv.astype(F64).astype(F32)), mask)
|
||||
elif dt[0] in 'us' and stp[0] in 'us':
|
||||
ws = int(re.sub(r'\D', '', stp))
|
||||
val = self.sext(v, ws) if stp[0] == 's' else (v & np.uint64(MASK[ws])).astype(np.int64)
|
||||
self.put(a[0], val.astype(np.int64).astype(U64), mask)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
|
||||
def op_shfl(self, st, parts, a, mask):
|
||||
mode = parts[2]
|
||||
v = self.get(a[1])
|
||||
b = self.get(a[2]).astype(np.int64)
|
||||
out = np.array(v)
|
||||
for w in range(self.T // 32):
|
||||
sl = slice(32 * w, 32 * w + 32)
|
||||
lane = np.arange(32)
|
||||
if mode == 'bfly':
|
||||
src = lane ^ b[sl]
|
||||
elif mode == 'down':
|
||||
src = np.where(lane + b[sl] < 32, lane + b[sl], lane)
|
||||
elif mode == 'up':
|
||||
src = np.where(lane - b[sl] >= 0, lane - b[sl], lane)
|
||||
elif mode == 'idx':
|
||||
src = b[sl] & 31
|
||||
out[sl] = v[sl][src]
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_movmatrix(self, st, parts, a, mask):
|
||||
v = self.get(a[1])
|
||||
out = np.zeros(self.T, U64)
|
||||
for w in range(self.T // 32):
|
||||
m = np.zeros((8, 8), U64)
|
||||
for lane in range(32):
|
||||
x = int(v[32 * w + lane])
|
||||
r, c = lane // 4, (lane % 4) * 2
|
||||
m[r, c], m[r, c + 1] = x & 0xffff, (x >> 16) & 0xffff
|
||||
mt = m.T
|
||||
for lane in range(32):
|
||||
r, c = lane // 4, (lane % 4) * 2
|
||||
out[32 * w + lane] = int(mt[r, c]) | (int(mt[r, c + 1]) << 16)
|
||||
self.put(a[0], out, mask)
|
||||
|
||||
def op_mma(self, st, parts, a, mask):
|
||||
shape = parts[3]
|
||||
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
|
||||
A = [self.get(n) for n in a_names]
|
||||
Bv = [self.get(n) for n in b_names]
|
||||
C = [self.get(n) for n in c_names]
|
||||
dout = [np.array(self.get(n)) for n in d_names]
|
||||
for w in range(self.T // 32):
|
||||
base = 32 * w
|
||||
if shape == 'm16n8k32':
|
||||
Am = np.zeros((16, 32)); Bm = np.zeros((32, 8)); Cm = np.zeros((16, 8))
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(16):
|
||||
byte = (int(A[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
|
||||
row = g if (i < 4 or 8 <= i < 12) else g + 8
|
||||
col = t * 4 + (i & 3) + (16 if i >= 8 else 0)
|
||||
Am[row, col] = E4M3[byte]
|
||||
for i in range(8):
|
||||
byte = (int(Bv[i // 4][base + lane]) >> (8 * (i % 4))) & 0xff
|
||||
Bm[t * 4 + (i & 3) + (16 if i >= 4 else 0), g] = E4M3[byte]
|
||||
elif shape == 'm16n8k16':
|
||||
Am = np.zeros((16, 16)); Bm = np.zeros((16, 8)); Cm = np.zeros((16, 8))
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(8):
|
||||
h = (int(A[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
row = g if (i < 2 or 4 <= i < 6) else g + 8
|
||||
col = t * 2 + (i & 1) + (8 if i >= 4 else 0)
|
||||
Am[row, col] = np.uint16(h).view(F16)
|
||||
for i in range(4):
|
||||
h = (int(Bv[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
Bm[t * 2 + (i & 1) + (8 if i >= 2 else 0), g] = np.uint16(h).view(F16)
|
||||
else:
|
||||
raise NotImplementedError(st.text)
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for i in range(4):
|
||||
h = (int(C[i // 2][base + lane]) >> (16 * (i % 2))) & 0xffff
|
||||
Cm[g if i < 2 else g + 8, t * 2 + (i & 1)] = np.uint16(h).view(F16)
|
||||
Dm = (Am @ Bm + Cm).astype(F16)
|
||||
if self.mma_hook:
|
||||
self.mma_hook(self, st, w, Am, Bm, Cm, Dm, [self.prov.get(n) for n in a_names], [self.prov.get(n) for n in b_names])
|
||||
for lane in range(32):
|
||||
g, t = lane >> 2, lane & 3
|
||||
for r in range(2):
|
||||
lo = int(Dm[g + 8 * r, t * 2].view(np.uint16))
|
||||
hi = int(Dm[g + 8 * r, t * 2 + 1].view(np.uint16))
|
||||
dout[r][base + lane] = lo | (hi << 16)
|
||||
for n, v in zip(d_names, dout):
|
||||
self.put(n, v, mask)
|
||||
|
||||
mma_hook = None
|
||||
|
||||
# wmma.mma.sync.aligned.row.col.m16n16k16.f16.f16 with register fragments: two m16n8k16 MMAs
|
||||
# sharing A registers 0..3, B registers 0,1 / 2,3 and C/D registers 0,1 / 2,3 (the upper A and B
|
||||
# registers of the fragment ABI are duplicates; ZLUDA's LLVMZludaSplitWmmaM16n16k16 mapping)
|
||||
def op_wmma(self, st, parts, a, mask):
|
||||
if 'm16n16k16' not in parts or 'mma' not in parts:
|
||||
raise NotImplementedError(st.text)
|
||||
d_names, a_names, b_names, c_names = (self.vec(x) for x in a)
|
||||
for half in range(2):
|
||||
sub = Stmt_like(st, 'mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16',
|
||||
['{%s}' % ', '.join(d_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(a_names[0:4]),
|
||||
'{%s}' % ', '.join(b_names[2 * half:2 * half + 2]), '{%s}' % ', '.join(c_names[2 * half:2 * half + 2])])
|
||||
self.op_mma(sub, sub.op.split('.'), sub.args, mask)
|
||||
|
||||
|
||||
class Stmt_like:
|
||||
def __init__(self, st, op, args):
|
||||
self.op, self.args, self.text = op, args, st.text + f' [as {op}]'
|
||||
for k in Stmt.__slots__:
|
||||
if k not in ('op', 'args', 'text'):
|
||||
setattr(self, k, getattr(st, k, None))
|
||||
|
||||
|
||||
def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None):
|
||||
stmts, types, shared = parse_kernel(ptx, kernel)
|
||||
mem = Memory(dump)
|
||||
launch = [l for l in open(dump + '/manifest.txt') if l.startswith('launch')][0].split()
|
||||
grid = tuple(int(x) for x in launch[1:4])
|
||||
block = tuple(int(x) for x in launch[4:7])
|
||||
b = Block(stmts, types, shared, mem, kernel + '_param_0', block, grid, ctaid)
|
||||
b.hooks = list(hooks)
|
||||
b.mma_hook = mma_hook
|
||||
steps = b.run(limit)
|
||||
return b, mem, steps
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
ptx, kernel, dump = sys.argv[1:4]
|
||||
cx, cy = (int(v) for v in sys.argv[4].split(',')) if len(sys.argv) > 4 else (0, 0)
|
||||
import time
|
||||
t0 = time.time()
|
||||
b, mem, steps = run_block(ptx, kernel, dump, (cx, cy, 0))
|
||||
print(f'block ({cx},{cy}): {steps} statements in {time.time() - t0:.1f}s, {len(mem.stores)} store groups')
|
||||
@@ -0,0 +1,127 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compares native preset-K layer outputs with the numpy model (pwin_model.py) on sampled windows of a replay dump.
|
||||
|
||||
usage: pwin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--windows N]
|
||||
DUMP_DIR replay dump of one dltss_pwin_* launch (inputs)
|
||||
OUT_DIR dump_runner output of a native kernel on that dump (alloc-N.bin after the launch)
|
||||
For every sampled window (the four corners, border windows and random interior ones) the model output and the
|
||||
kernel's stored values are compared per f16 element: exact matches, and the error in f16 ulps at the model
|
||||
value's magnitude (values below 2^-10 count with the ulp of 2^-10, so near-zero sign flips do not dominate).
|
||||
The kernels round as NVIDIA's f16 wmma does except in the PWIN_F32ACC layers, and gfx12 WMMA groups its f32
|
||||
sums differently from gfx11, so neither build is expected to match the model exactly.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import pwin_model as pm # noqa: E402
|
||||
|
||||
F16 = np.float16
|
||||
|
||||
# layer -> (kind, H, C, COUT/CL, extra)
|
||||
LAYERS = {
|
||||
'enc0': ('enc0', 2, 32, 64, None), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
|
||||
'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 16, None),
|
||||
'dec4': ('dec', 4, 128, 160, None), 'dec3': ('dec', 4, 96, 128, None), 'dec2': ('dec', 2, 64, 96, None),
|
||||
'dec1': ('dec', 2, 64, 64, None), 'dec0': ('dec0', 2, 32, 64, None),
|
||||
}
|
||||
|
||||
|
||||
class Outputs:
|
||||
"""the dump's pointer map over the allocations a native run saved"""
|
||||
def __init__(self, dump, out_dir):
|
||||
self.d = dump
|
||||
self.allocs = {i: (base, np.fromfile(f'{out_dir}/alloc-{i}.bin', np.uint8)) for i, (base, _) in dump.allocs.items()}
|
||||
|
||||
def rows(self, off_param, index, width):
|
||||
a = self.d.u64(off_param)
|
||||
for base, buf in self.allocs.values():
|
||||
if base <= a < base + buf.size:
|
||||
o = a - base + 2 * index * width
|
||||
return buf[o:o + 2 * width].view(F16)
|
||||
raise ValueError(hex(a))
|
||||
|
||||
|
||||
def ulps(model, got):
|
||||
m = model.astype(np.float64)
|
||||
g = got.astype(np.float64)
|
||||
scale = np.maximum(np.abs(m), 2.0 ** -10)
|
||||
ulp = 2.0 ** (np.floor(np.log2(scale)) - 10)
|
||||
return np.abs(m - g) / ulp
|
||||
|
||||
|
||||
def windows(grid, n, seed=1):
|
||||
gx, gy = grid[0], grid[1]
|
||||
w = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
|
||||
r = random.Random(seed)
|
||||
while len(w) < min(n, gx * gy):
|
||||
w.add((r.randrange(gx), r.randrange(gy)))
|
||||
return sorted(w)
|
||||
|
||||
|
||||
def model_window(P, layer, bx, by):
|
||||
kind, H, C, X, extra = LAYERS[layer]
|
||||
if kind == 'enc0':
|
||||
return pm.pwin_enc0(P, bx, by)
|
||||
if kind == 'enc':
|
||||
return pm.pwin_encoder(P, bx, by, H, C, X, posattn=extra)
|
||||
if kind == 'plain':
|
||||
return pm.swin_core(P, pm.Layout(H, C, C), pm.load_window(P, bx, by, C)), None
|
||||
if kind == 'dec':
|
||||
return pm.pwin_decoder(P, bx, by, H, C, X), None
|
||||
return pm.pwin_dec0(P, bx, by)
|
||||
|
||||
|
||||
def stored(P, O, layer, bx, by, y, extra_out):
|
||||
"""(model values, kernel values) of everything the kernel stores for window (bx, by)"""
|
||||
kind, H, C, X, _ = LAYERS[layer]
|
||||
W, Hh = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
ms, gs = [], []
|
||||
full_off = 32 if kind in ('enc0', 'enc') else 24
|
||||
for t in range(64):
|
||||
Y0, X0 = 8 * by - sy + (t >> 3), 8 * bx - sx + (t & 7)
|
||||
if not (0 <= Y0 < Hh and 0 <= X0 < W):
|
||||
continue
|
||||
if kind == 'dec0':
|
||||
ms.append(extra_out[t]); gs.append(O.rows(24, Y0 * W + X0, 40))
|
||||
else:
|
||||
ms.append(y[t]); gs.append(O.rows(full_off, Y0 * W + X0, C))
|
||||
if kind in ('enc0', 'enc'):
|
||||
cout = extra_out.shape[1]
|
||||
for r in range(16):
|
||||
Ym, Xm = int((8 * by - sy) / 2) + (r >> 2), int((8 * bx - sx) / 2) + (r & 3)
|
||||
if 0 <= Ym < Hh // 2 and 0 <= Xm < W // 2:
|
||||
ms.append(extra_out[r]); gs.append(O.rows(24, Ym * (W // 2) + Xm, cout))
|
||||
return np.concatenate(ms), np.concatenate(gs)
|
||||
|
||||
|
||||
def main():
|
||||
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
||||
n = 12
|
||||
if '--windows' in sys.argv:
|
||||
n = int(sys.argv[sys.argv.index('--windows') + 1])
|
||||
args.remove(str(n))
|
||||
dump_dir, outs = args[0], args[1:]
|
||||
P = pm.Dump(dump_dir)
|
||||
layer = P.kernel.replace('dltss_pwin_', '').replace('_layer', '')
|
||||
wins = windows(P.grid, n)
|
||||
models = {w: model_window(P, layer, *w) for w in wins}
|
||||
for out in outs:
|
||||
O = Outputs(P, out)
|
||||
m, g = zip(*(stored(P, O, layer, bx, by, *models[(bx, by)]) for bx, by in wins))
|
||||
m, g = np.concatenate(m), np.concatenate(g)
|
||||
u = ulps(m, g)
|
||||
exact = (m.view(np.uint16) == g.view(np.uint16)).mean()
|
||||
md, gd = m.astype(np.float64), g.astype(np.float64)
|
||||
psnr = 10 * np.log10(np.abs(md).max() ** 2 / max(((md - gd) ** 2).mean(), 1e-30))
|
||||
print(f'{layer:5s} {os.path.basename(os.path.normpath(out)):10s} {len(wins)} windows {m.size} values: exact {100 * exact:5.1f}% '
|
||||
f'<=1ulp {100 * (u <= 1).mean():5.1f}% <=4ulp {100 * (u <= 4).mean():6.2f}% max {u.max():7.1f} ulp '
|
||||
f'mean {u.mean():.3f} ulp psnr {psnr:.1f} dB finite {np.isfinite(g.astype(np.float32)).mean() * 100:.0f}%')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,440 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Numpy reference of DLSS 4 (preset K) DltssPaddedWinLayer blocks, built stage by stage against the
|
||||
PTX interpreter (kernels/tools/ptxsim.py). f16 arithmetic in the kernel's operation order."""
|
||||
import struct
|
||||
import numpy as np
|
||||
|
||||
F16, F32, F64 = np.float16, np.float32, np.float64
|
||||
|
||||
# Legacy replay reports retain their original approximate norm by default.
|
||||
# Windows strict K validation selects the PTX fragment's rounded reduction.
|
||||
NORM_FRAGMENT_ORDER = False
|
||||
|
||||
|
||||
def norm_rows(x):
|
||||
if not NORM_FRAGMENT_ORDER:
|
||||
return f16(np.sum(x.astype(F32).astype(F64) ** 2, axis=1))
|
||||
tokens, channels = x.shape
|
||||
assert channels in (32, 64, 96, 128, 160)
|
||||
square = f16(x.astype(F64) ** 2)
|
||||
partial = []
|
||||
for lane in range(4):
|
||||
halves = []
|
||||
for parity in range(2):
|
||||
pairs = [add16(square[:, 2 * lane + 8 * parity + 32 * g:2 * lane + 8 * parity + 32 * g + 2],
|
||||
square[:, 2 * lane + 8 * parity + 32 * g + 16:2 * lane + 8 * parity + 32 * g + 18])
|
||||
for g in range(channels // 32)]
|
||||
value = pairs[0]
|
||||
if channels >= 64:
|
||||
value = add16(value, pairs[1])
|
||||
if channels == 96:
|
||||
value = add16(value, pairs[2])
|
||||
if channels >= 128:
|
||||
value = add16(value, add16(pairs[2], pairs[3]))
|
||||
if channels == 160:
|
||||
value = add16(value, pairs[4])
|
||||
halves.append(value)
|
||||
partial.append(add16(halves[0], halves[1]))
|
||||
total = add16(add16(partial[0], partial[1]), add16(partial[2], partial[3]))
|
||||
return add16(total[:, 0], total[:, 1])
|
||||
|
||||
|
||||
class Dump:
|
||||
def __init__(self, d):
|
||||
self.dir = d
|
||||
man = open(d + '/manifest.txt').read().split('\n')
|
||||
self.kernel = man[0].split()[1]
|
||||
la = man[1].split()
|
||||
self.grid, self.block = tuple(map(int, la[1:4])), tuple(map(int, la[4:7]))
|
||||
self.args = open(d + '/args.bin', 'rb').read()
|
||||
self.allocs = {}
|
||||
for l in man:
|
||||
p = l.split()
|
||||
if p and p[0] == 'alloc':
|
||||
self.allocs[int(p[1])] = (int(p[2], 16), np.fromfile(f'{d}/alloc-{p[1]}.bin', np.uint8))
|
||||
|
||||
def u64(self, off):
|
||||
return struct.unpack_from('<Q', self.args, off)[0]
|
||||
|
||||
def i32x2(self, off):
|
||||
return struct.unpack_from('<ii', self.args, off)
|
||||
|
||||
def ptr(self, off):
|
||||
a = self.u64(off)
|
||||
for i, (base, buf) in self.allocs.items():
|
||||
if base <= a < base + buf.size:
|
||||
return buf, a - base
|
||||
raise ValueError(hex(a))
|
||||
|
||||
def f16(self, off_param, byte, n):
|
||||
buf, o = self.ptr(off_param)
|
||||
return buf[o + byte:o + byte + 2 * n].view(F16).copy()
|
||||
|
||||
|
||||
def f16(x):
|
||||
return np.asarray(x, F64).astype(F16)
|
||||
|
||||
|
||||
def mirror(v, n):
|
||||
v = abs(v)
|
||||
return min(v, 2 * n - 2 - v)
|
||||
|
||||
|
||||
def load_window(P, bx, by, C):
|
||||
"""64 tokens (raster order of the 8x8 window) x C channels, f16, mirror padded."""
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(8)
|
||||
x = np.zeros((64, C), F16)
|
||||
for r in range(64):
|
||||
ty, tx = r >> 3, r & 7
|
||||
Y = mirror(8 * by - sy + ty, H)
|
||||
X = mirror(8 * bx - sx + tx, W)
|
||||
a = o + 2 * ((Y * W + X) * C)
|
||||
x[r] = buf[a:a + 2 * C].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
# --- reductions in the kernel's order ---------------------------------------------------------
|
||||
def sumsq_rows(x):
|
||||
"""f16 sum of squares per token, reproducing the kernel's tree for C = 64 (per lane: 16 products
|
||||
over 8 words (t + 4k), pairwise f16x2 adds, butterfly over the 4 lanes of a row, final hi+lo)."""
|
||||
T, C = x.shape
|
||||
assert C == 64
|
||||
sq = f16(x.astype(F64) * x.astype(F64)) # f16 products
|
||||
out = np.zeros(T, F16)
|
||||
for r in range(T):
|
||||
lanes = []
|
||||
for t in range(4):
|
||||
# word k holds channels 2(t+4k), 2(t+4k)+1 as an f16 pair
|
||||
w = [np.array([sq[r, 2 * (t + 4 * k)], sq[r, 2 * (t + 4 * k) + 1]], F16) for k in range(8)]
|
||||
a = [f16(w[2 * i].astype(F64) + w[2 * i + 1].astype(F64)) for i in range(4)] # (0+1),(2+3),(4+5),(6+7)
|
||||
b = [f16(a[0].astype(F64) + a[2].astype(F64)), f16(a[1].astype(F64) + a[3].astype(F64))]
|
||||
lanes.append(f16(b[0].astype(F64) + b[1].astype(F64)))
|
||||
s1 = [f16(lanes[t].astype(F64) + lanes[t ^ 1].astype(F64)) for t in range(4)]
|
||||
s2 = f16(s1[0].astype(F64) + s1[2].astype(F64))
|
||||
out[r] = f16(s2[0].astype(F64) + s2[1].astype(F64))
|
||||
return out
|
||||
|
||||
|
||||
def rsqrt16(s):
|
||||
return f16(1.0 / np.sqrt(s.astype(F32).astype(F64)))
|
||||
|
||||
|
||||
def mma(A, B, C):
|
||||
"""one wmma k16 step: D = f16(A @ B + C) with exact sums (the interpreter's semantics)."""
|
||||
return f16(A.astype(F64) @ B.astype(F64) + C.astype(F64))
|
||||
|
||||
|
||||
def frag_offset(k, n):
|
||||
"""byte offset of logical B[k][n] (k, n < 16) inside one 512-byte pre-swizzled f16 fragment"""
|
||||
return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1)
|
||||
|
||||
|
||||
FRAG = np.array([[frag_offset(k, n) for n in range(16)] for k in range(16)])
|
||||
|
||||
|
||||
def weights(P, byte, K, N, nt_stride=None):
|
||||
"""logical W[K][N] from fragments at weights + byte + 512 * kt + nt_stride * nt"""
|
||||
buf, o = P.ptr(64)
|
||||
nt_stride = nt_stride or 512 * (K // 16)
|
||||
W = np.zeros((K, N), F16)
|
||||
for kt in range(K // 16):
|
||||
for nt in range(N // 16):
|
||||
base = o + byte + 512 * kt + nt_stride * nt
|
||||
vals = buf[(base + FRAG)[..., None] + np.arange(2)].reshape(16, 16, 2).view(F16)[..., 0]
|
||||
W[16 * kt:16 * kt + 16, 16 * nt:16 * nt + 16] = vals
|
||||
return W
|
||||
|
||||
|
||||
def gemm(A, W, C0=None):
|
||||
"""chained wmma k16 steps (f16 rounding after each step), per 16x16 output tile"""
|
||||
M, K = A.shape
|
||||
N = W.shape[1]
|
||||
D = np.zeros((M, N), F16) if C0 is None else C0.astype(F16).copy()
|
||||
for kt in range(K // 16):
|
||||
D = mma(A[:, 16 * kt:16 * kt + 16], W[16 * kt:16 * kt + 16], D)
|
||||
return D
|
||||
|
||||
|
||||
def hv(bits):
|
||||
return np.array(bits, np.uint16).view(F16)
|
||||
|
||||
|
||||
def fma16(a, b, c):
|
||||
return f16(a.astype(F64) * b.astype(F64) + c.astype(F64))
|
||||
|
||||
|
||||
def add16(a, b):
|
||||
return f16(a.astype(F64) + b.astype(F64))
|
||||
|
||||
|
||||
def mul16(a, b):
|
||||
return f16(a.astype(F64) * b.astype(F64))
|
||||
|
||||
|
||||
def softmax_exp(S):
|
||||
"""e = ex2(t * (1 + a t^2) + c), t = clamp(S, +-21.64) (f16x2 ops of the kernel)"""
|
||||
hi, lo = hv(0x4D69), hv(0xCD69)
|
||||
a, one, c = hv(0x91D5), hv(0x3C00), hv(0xC50D)
|
||||
t = np.maximum(np.minimum(S, hi), lo)
|
||||
u = mul16(t, t)
|
||||
v = fma16(u, a, one)
|
||||
w = fma16(t, v, c)
|
||||
return f16(np.exp2(w.astype(F64)))
|
||||
|
||||
|
||||
def gelu16(x):
|
||||
"""x * (0.5 + c (0.41216 - 0.081081 |c|)), c = clamp(x, +-2) (f16x2 ops of the kernel)"""
|
||||
k1, k2 = f16(np.float32(np.uint32(0x3ED306EB).view(F32))), f16(np.float32(np.uint32(0x3DA60DD6).view(F32)))
|
||||
half, two = F16(0.5), F16(2.0)
|
||||
c = np.minimum(np.maximum(x, -two), two)
|
||||
u = f16(k1.astype(F64) - mul16(np.full_like(c, k2), np.abs(c)).astype(F64))
|
||||
return mul16(x, add16(np.full_like(c, half), mul16(c, u)))
|
||||
|
||||
|
||||
def bias_table(P, byte):
|
||||
"""64x64 f16 table stored in C-fragment order, 512 bytes per 16x16 tile (tile 4 mt + nt)"""
|
||||
buf, o = P.ptr(64)
|
||||
B = np.zeros((64, 64), F16)
|
||||
for i in range(64):
|
||||
for j in range(64):
|
||||
mt, nt, r, c = i >> 4, j >> 4, i & 15, j & 15
|
||||
off = 512 * (4 * nt + mt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3)
|
||||
B[i, j] = buf[o + byte + off:o + byte + off + 2].view(F16)[0]
|
||||
return B
|
||||
|
||||
|
||||
def vec(P, byte, n):
|
||||
return P.f16(64, byte, n)
|
||||
|
||||
|
||||
def pwin_enc2(P, bx, by, pm_order=((0, 0), (0, 1), (1, 0), (1, 1)), b2_in_c0=True, dbg=None):
|
||||
"""logical enc2 layer (H 2, C 64 -> merged 96) for one window; returns (y [64 x 64], merged [16 x 96])"""
|
||||
C, H, Dh = 64, 2, 32
|
||||
x0 = load_window(P, bx, by, C)
|
||||
s = norm_rows(x0)
|
||||
r = rsqrt16(s)
|
||||
h1 = mul16(x0, mul16(r[:, None], vec(P, 0, C)[None, :]))
|
||||
Os = []
|
||||
for h in range(H):
|
||||
hb = 128 + 12288 * h
|
||||
q, k, v = [gemm(h1[:, 32:64], weights(P, hb + 2048 * (j + 3), 32, 32, 1024),
|
||||
gemm(h1[:, 0:32], weights(P, hb + 2048 * j, 32, 32, 1024))) for j in range(3)]
|
||||
S = gemm(q, k.T.copy(), bias_table(P, 24704 + 8192 * h))
|
||||
E = softmax_exp(S)
|
||||
rs = f16(np.sum(E.astype(F64), axis=1))
|
||||
rc = f16(1.0 / rs.astype(F32).astype(F64))
|
||||
Pm = mul16(E, rc[:, None])
|
||||
Os.append(gemm(Pm, v))
|
||||
if dbg is not None:
|
||||
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = q, k, v, S, Pm, Os[-1]
|
||||
O = np.concatenate(Os, axis=1)
|
||||
bo = vec(P, 41088, C)
|
||||
Y = gemm(O[:, 32:64], weights(P, 45312, 32, 64, 1024), gemm(O[:, 0:32], weights(P, 41216, 32, 64, 1024), np.tile(bo, (64, 1))))
|
||||
x1 = add16(Y, x0)
|
||||
if dbg is not None:
|
||||
dbg['Y'], dbg['x1'] = Y, x1
|
||||
m = mul16(x1, vec(P, 49408, C)[None, :])
|
||||
b2 = vec(P, 49536, C)
|
||||
acc = add16(x1, np.tile(b2, (64, 1))) if b2_in_c0 else x1
|
||||
for c in range(8):
|
||||
Hc = gemm(m, weights(P, 49664 + 4096 * c, 64, 32, 2048), np.tile(vec(P, 82432 + 64 * c, 32), (64, 1)))
|
||||
acc = gemm(gelu16(Hc), weights(P, 82944 + 4096 * c, 32, 64, 1024), acc)
|
||||
if dbg is not None:
|
||||
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
|
||||
y = acc
|
||||
# patch merge: 16 merged tokens (4x4), K = 4 sub-tokens x 64
|
||||
A = np.zeros((16, 256), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(pm_order):
|
||||
A[mt, 64 * q:64 * q + 64] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
merged = np.concatenate([gemm(A, weights(P, 115712 + 24576 * w, 256, 48, 8192), np.tile(vec(P, 164864 + 96 * w, 48), (16, 1)))
|
||||
for w in range(2)], axis=1)
|
||||
return y, merged
|
||||
|
||||
|
||||
class Layout:
|
||||
"""weight-buffer layout of a plain DltssPaddedWinLayer (encoder with patch merge), bytes"""
|
||||
def __init__(self, H, C, COUT):
|
||||
self.H, self.C, self.COUT = H, C, COUT
|
||||
self.G1 = 0
|
||||
self.QKV = 2 * C
|
||||
self.HEAD = 192 * C
|
||||
self.BIAS = self.QKV + self.HEAD * H
|
||||
self.BO = self.BIAS + 8192 * H
|
||||
self.WO = self.BO + 2 * C
|
||||
self.G2 = self.WO + 64 * C * H
|
||||
self.B2 = self.G2 + 2 * C
|
||||
self.W1 = self.B2 + 2 * C
|
||||
self.B1 = self.W1 + 8 * C * C
|
||||
self.W2 = self.B1 + 8 * C
|
||||
self.PM = self.W2 + 8 * C * C
|
||||
self.NPW = COUT // H # merged channels per original warp
|
||||
self.NPA = -(-self.NPW // 16) * 16 # allocated (padded to whole 16-column tiles)
|
||||
self.PMB = self.PM + 8 * C * self.NPA * H
|
||||
|
||||
|
||||
def swin_core(P, L, x0, dbg=None, posattn=False):
|
||||
"""norm -> attention -> Wo + residual -> MLP (+ residual): the block output y [64 x C]"""
|
||||
C, H = L.C, L.H
|
||||
NCH = C // 32
|
||||
s = norm_rows(x0)
|
||||
r = rsqrt16(s)
|
||||
h1 = mul16(x0, mul16(r[:, None], vec(P, L.G1, C)[None, :]))
|
||||
Os = []
|
||||
for h in range(H):
|
||||
hb = L.QKV + L.HEAD * h
|
||||
proj = []
|
||||
for j in range(3):
|
||||
acc = None
|
||||
for ch in range(NCH):
|
||||
acc = gemm(h1[:, 32 * ch:32 * ch + 32], weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024), acc)
|
||||
proj.append(acc)
|
||||
q, k, v = proj
|
||||
if posattn:
|
||||
# position-only attention (trait flag 2): the table is the probability matrix itself
|
||||
Os.append(gemm(bias_table(P, L.BIAS + 8192 * h), v))
|
||||
if dbg is not None:
|
||||
dbg[f'v{h}'], dbg[f'O{h}'] = v, Os[-1]
|
||||
continue
|
||||
S = gemm(q, k.T.copy(), bias_table(P, L.BIAS + 8192 * h))
|
||||
E = softmax_exp(S)
|
||||
rs = f16(np.sum(E.astype(F64), axis=1))
|
||||
rc = f16(1.0 / rs.astype(F32).astype(F64))
|
||||
Os.append(gemm(mul16(E, rc[:, None]), v))
|
||||
if dbg is not None:
|
||||
dbg[f'q{h}'], dbg[f'k{h}'], dbg[f'v{h}'], dbg[f'O{h}'], dbg[f'S{h}'], dbg[f'E{h}'], dbg[f'rs{h}'] = q, k, v, Os[-1], S, E, rs
|
||||
Y = np.tile(vec(P, L.BO, C), (64, 1))
|
||||
for h in range(H):
|
||||
Y = gemm(Os[h], weights(P, L.WO + 64 * C * h, 32, C, 1024), Y)
|
||||
x1 = add16(Y, x0)
|
||||
m = mul16(x1, vec(P, L.G2, C)[None, :])
|
||||
acc = add16(x1, np.tile(vec(P, L.B2, C), (64, 1)))
|
||||
if dbg is not None:
|
||||
dbg['Y'], dbg['x1'] = Y, x1
|
||||
for c in range(C // 8):
|
||||
Hc = gemm(m, weights(P, L.W1 + 64 * C * c, C, 32, 32 * C), np.tile(vec(P, L.B1 + 64 * c, 32), (64, 1)))
|
||||
acc = gemm(gelu16(Hc), weights(P, L.W2 + 64 * C * c, 32, C, 1024), acc)
|
||||
if dbg is not None:
|
||||
dbg[f'H{c}'], dbg[f'acc{c}'] = Hc, acc
|
||||
return acc
|
||||
|
||||
|
||||
def pwin_encoder(P, bx, by, H, C, COUT, pm_bias_padded=True, posattn=False):
|
||||
L = Layout(H, C, COUT)
|
||||
x0 = load_window(P, bx, by, C)
|
||||
y = swin_core(P, L, x0, posattn=posattn)
|
||||
A = np.zeros((16, 4 * C), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
Wp = weights(P, L.PM, 4 * C, L.NPA * H, 128 * C)
|
||||
bias = vec(P, L.PMB, L.NPA * H) if pm_bias_padded else None
|
||||
cols = []
|
||||
for w in range(H):
|
||||
sl = slice(L.NPA * w, L.NPA * w + L.NPW)
|
||||
b = bias[sl] if pm_bias_padded else vec(P, L.PMB + 2 * L.NPW * w, L.NPW)
|
||||
cols.append(gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(np.concatenate([b, np.zeros(L.NPA - L.NPW, F16)]), (16, 1)))[:, :L.NPW])
|
||||
merged = np.concatenate(cols, axis=1)
|
||||
return y, merged
|
||||
|
||||
|
||||
class Shifted:
|
||||
"""view of the weight buffer starting at a byte offset (decoders: core weights follow the expand)"""
|
||||
def __init__(self, P, base):
|
||||
self.P, self.base = P, base
|
||||
|
||||
def ptr(self, off):
|
||||
buf, o = self.P.ptr(off)
|
||||
return buf, o + (self.base if off == 64 else 0)
|
||||
|
||||
def f16(self, off_param, byte, n):
|
||||
return self.P.f16(off_param, byte + (self.base if off_param == 64 else 0), n)
|
||||
|
||||
def i32x2(self, off):
|
||||
return self.P.i32x2(off)
|
||||
|
||||
|
||||
def load_lowres(P, bx, by, CL):
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(8)
|
||||
W2, H2 = W // 2, H // 2
|
||||
lx0, ly0 = int((8 * bx - sx) / 2), int((8 * by - sy) / 2)
|
||||
x = np.zeros((16, CL), F16)
|
||||
for r in range(16):
|
||||
Y = mirror(ly0 + (r >> 2), H2); X = mirror(lx0 + (r & 3), W2)
|
||||
a = o + 2 * ((Y * W2 + X) * CL)
|
||||
x[r] = buf[a:a + 2 * CL].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
def load_skip(P, bx, by, C):
|
||||
W, H = P.i32x2(0)
|
||||
sx, sy = P.i32x2(56)
|
||||
buf, o = P.ptr(16)
|
||||
x = np.zeros((64, C), F16)
|
||||
for r in range(64):
|
||||
Y = mirror(8 * by - sy + (r >> 3), H); X = mirror(8 * bx - sx + (r & 7), W)
|
||||
a = o + 2 * ((Y * W + X) * C)
|
||||
x[r] = buf[a:a + 2 * C].view(F16)
|
||||
return x
|
||||
|
||||
|
||||
def pwin_decoder(P, bx, by, H, C, CL, order=((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
xl = load_lowres(P, bx, by, CL)
|
||||
EXP = 2 * CL * 4 * C
|
||||
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
|
||||
skip = load_skip(P, bx, by, C)
|
||||
x0 = np.zeros((64, C), F16)
|
||||
for r in range(16):
|
||||
ly, lx = r >> 2, r & 3
|
||||
for q, (dy, dx) in enumerate(order):
|
||||
t = 8 * (2 * ly + dy) + 2 * lx + dx
|
||||
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
|
||||
L = Layout(H, C, C)
|
||||
return swin_core(Shifted(P, EXP + 8 * C), L, x0)
|
||||
|
||||
|
||||
def pwin_enc0(P, bx, by, H=2, C=32, COUT=64, CIN=16):
|
||||
"""first layer: token embedding x0 = in (CIN) W_e + b_e, then the core (position-only attention) and merge"""
|
||||
xin = load_window(P, bx, by, CIN)
|
||||
x0 = gemm(xin, weights(P, 2 * C, CIN, C, 512 * (CIN // 16)), np.tile(vec(P, 0, C), (64, 1)))
|
||||
x0 = np.maximum(x0, F16(0)) # ReLU after the embedding
|
||||
CB = 2 * C + 2 * CIN * C
|
||||
L = Layout(H, C, COUT)
|
||||
Pw = Shifted(P, CB)
|
||||
y = swin_core(Pw, L, x0, posattn=True)
|
||||
A = np.zeros((16, 4 * C), F16)
|
||||
for mt in range(16):
|
||||
my, mx = mt >> 2, mt & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
A[mt, C * q:C * q + C] = y[8 * (2 * my + dy) + 2 * mx + dx]
|
||||
Wp = weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C)
|
||||
bias = vec(Pw, L.PMB, L.NPA * H)
|
||||
cols = [gemm(A, Wp[:, L.NPA * w:L.NPA * (w + 1)], np.tile(bias[L.NPA * w:L.NPA * (w + 1)], (16, 1)))[:, :L.NPW] for w in range(H)]
|
||||
return y, np.concatenate(cols, axis=1)
|
||||
|
||||
|
||||
def pwin_dec0(P, bx, by, H=2, C=32, CL=64, NOUT=40, NOUTA=48):
|
||||
"""last layer: patch expand + skip, core (position-only attention), output head y W_h + b_h (NOUT of NOUTA)"""
|
||||
xl = load_lowres(P, bx, by, CL)
|
||||
EXP = 2 * CL * 4 * C
|
||||
E = gemm(xl, weights(P, 0, CL, 4 * C, 512 * (CL // 16)), np.tile(vec(P, EXP, 4 * C), (16, 1)))
|
||||
skip = load_skip(P, bx, by, C)
|
||||
x0 = np.zeros((64, C), F16)
|
||||
for r in range(16):
|
||||
ly, lx = r >> 2, r & 3
|
||||
for q, (dy, dx) in enumerate(((0, 0), (0, 1), (1, 0), (1, 1))):
|
||||
t = 8 * (2 * ly + dy) + 2 * lx + dx
|
||||
x0[t] = add16(E[r, q * C:(q + 1) * C], skip[t])
|
||||
L = Layout(H, C, C)
|
||||
Pw = Shifted(P, EXP + 8 * C)
|
||||
y = swin_core(Pw, L, x0, posattn=True)
|
||||
HB = L.PM # head bias, then the head weights
|
||||
Wh = weights(Pw, HB + 2 * NOUTA, C, NOUTA, 1024)
|
||||
out = gemm(y, Wh, np.tile(vec(Pw, HB, NOUTA), (64, 1)))
|
||||
return y, out[:, :NOUT]
|
||||
@@ -0,0 +1,113 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Compares native preset-M Swin-block outputs with the numpy model (swin_model.py) on sampled blocks of a dump.
|
||||
|
||||
usage: swin_check.py DUMP_DIR OUT_DIR [OUT_DIR...] [--blocks N]
|
||||
DUMP_DIR replay dump of one rrlite_{enc1,enc2,enc3_tube,dec1,dec2}_4x4 launch
|
||||
OUT_DIR dump_runner output of a native kernel on that dump
|
||||
The stored e4m3 outputs (block output, and the patch merge of encoders that have one) are compared code by code:
|
||||
the share of identical codes, of codes one e4m3 step apart, and the PSNR of the decoded values. The model is
|
||||
NVIDIA's arithmetic (e4m3 activations, f16 accumulation per k32 step); the production kernels keep activations
|
||||
f16 and accumulate in f32 (SWIN_NO_Q8, SWIN_F32ACC), so they are not expected to match it exactly.
|
||||
"""
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
import model_enc3 as me # noqa: E402
|
||||
import swin_model as sm # noqa: E402
|
||||
|
||||
# kernel -> (C, heads, NPM, CIN)
|
||||
LAYERS = {'enc1': (64, 2, 96, 0), 'enc2': (96, 4, 128, 0), 'enc3_tube': (128, 4, 0, 0), 'dec2': (96, 4, 0, 128),
|
||||
'dec1': (64, 2, 0, 96)}
|
||||
|
||||
# Preserve the existing Linux oracle by default; Windows strict replay enables
|
||||
# the audited 310.9 f16 patch-merge bound before final e4m3 quantisation.
|
||||
CLAMP_PATCH_MERGE = False
|
||||
|
||||
|
||||
def blocks(grid, n, seed=1):
|
||||
gx, gy = grid[0], grid[1]
|
||||
b = {(0, 0), (gx - 1, 0), (0, gy - 1), (gx - 1, gy - 1), (gx // 2, 0), (0, gy // 2), (gx - 1, gy // 2), (gx // 2, gy - 1)}
|
||||
r = random.Random(seed)
|
||||
while len(b) < min(n, gx * gy):
|
||||
b.add((r.randrange(gx), r.randrange(gy)))
|
||||
return sorted(b)
|
||||
|
||||
|
||||
def model_block(P, layer, bx, by):
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
pre = 4 * C * CIN + 8 * C if CIN else 0
|
||||
x0 = sm.decoder_input(P, bx, by, C, CIN, heads, P.p32) if CIN else sm.load_tokens(P, bx, by, C // 32)
|
||||
x2, out, L = sm.swin_block(P, x0, C, heads, pre)
|
||||
merged = None
|
||||
if NPM:
|
||||
A = sm.patch_merge_input(out, C)
|
||||
PM0 = L['end']
|
||||
cols = []
|
||||
for g in range(NPM // 32):
|
||||
Wg = me.E4M3[P.wbuf[sm.woff_table(PM0 + g * 128 * C, 512, 64 * C, 4 * C, 32)]]
|
||||
bias = P.f16vec(PM0 + 4 * C * NPM + 64 * g, 32)
|
||||
cols.append(me.mma_chain(A, Wg, C0=np.broadcast_to(bias, (16, 32)).copy()))
|
||||
merged = sm.pair_to_nat(np.concatenate(cols, axis=1))
|
||||
if CLAMP_PATCH_MERGE:
|
||||
limit = me.f16(2 * np.pi)
|
||||
merged = np.clip(merged, -limit, limit)
|
||||
merged = me.q8(merged)
|
||||
return out, merged
|
||||
|
||||
|
||||
def codes_at(buf, base, ptr, W, H, X, Y, C):
|
||||
"""e4m3 codes of tokens (X, Y) (in bounds) from planes of 32 channels at ptr"""
|
||||
c = np.arange(C)
|
||||
off = ptr - base + (((c[None, :] >> 5) * H + Y[:, None]) * W + X[:, None]) * 32 + (c[None, :] & 31)
|
||||
return buf[off]
|
||||
|
||||
|
||||
def main():
|
||||
args = [a for a in sys.argv[1:] if not a.startswith('--')]
|
||||
n = 8
|
||||
if '--blocks' in sys.argv:
|
||||
n = int(sys.argv[sys.argv.index('--blocks') + 1])
|
||||
args.remove(str(n))
|
||||
dump, outs = args[0], args[1:]
|
||||
P = me.Params(dump)
|
||||
layer = P.kernel.replace('rrlite_', '').replace('_4x4', '')
|
||||
C, heads, NPM, CIN = LAYERS[layer]
|
||||
bl = blocks(P.grid, n)
|
||||
models = {b: model_block(P, layer, *b) for b in bl}
|
||||
for out_dir in outs:
|
||||
got, want = [], []
|
||||
for (bx, by) in bl:
|
||||
base, _ = P.alloc_of(P.out)
|
||||
aid = next(i for i, (b0, buf) in P.allocs.items() if b0 == base)
|
||||
buf = np.fromfile(f'{out_dir}/alloc-{aid}.bin', np.uint8)
|
||||
X, Y = P.token_xy(bx, by)
|
||||
ok = (X >= 0) & (X < P.W) & (Y >= 0) & (Y < P.H)
|
||||
out, merged = models[(bx, by)]
|
||||
got.append(codes_at(buf, base, P.out, P.W, P.H, X[ok], Y[ok], C).ravel())
|
||||
want.append(out[ok].ravel())
|
||||
if NPM:
|
||||
W2, H2 = P.W // 2, P.H // 2
|
||||
r = np.arange(16)
|
||||
MX, MY = (8 * bx - P.sx) // 2 + (r & 3), (8 * by - P.sy) // 2 + (r >> 2)
|
||||
okm = (MX >= 0) & (MX < W2) & (MY >= 0) & (MY < H2)
|
||||
mbase, _ = P.alloc_of(P.p48)
|
||||
mid = next(i for i, (b0, _) in P.allocs.items() if b0 == mbase)
|
||||
mbuf = buf if mid == aid else np.fromfile(f'{out_dir}/alloc-{mid}.bin', np.uint8)
|
||||
got.append(codes_at(mbuf, mbase, P.p48, W2, H2, MX[okm], MY[okm], NPM).ravel())
|
||||
want.append(merged[okm].ravel())
|
||||
g = me.E4M3[np.concatenate(got)]
|
||||
w = np.concatenate(want).astype(np.float64)
|
||||
rank = lambda v: np.searchsorted(me._POS, np.abs(v)) * np.sign(v) # noqa: E731
|
||||
steps = np.abs(rank(g) - rank(w))
|
||||
psnr = 10 * np.log10(np.nanmax(np.abs(w)) ** 2 / max(np.nanmean((g - w) ** 2), 1e-30))
|
||||
print(f'{layer:9s} {os.path.basename(os.path.normpath(out_dir)):10s} {len(bl)} blocks {g.size} codes: '
|
||||
f'identical {100 * (g == w).mean():5.1f}% <=1 step {100 * (steps <= 1).mean():6.2f}% '
|
||||
f'max {int(steps.max())} steps psnr {psnr:.1f} dB')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,170 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Parametrised numpy model of the rrlite Swin blocks (enc3 spec generalised); checked stage by stage
|
||||
against ptxsim MMA captures."""
|
||||
import numpy as np
|
||||
from model_enc3 import (q8, f16, Params, E4M3, F16, exp_trick, row_sum16, C_SCALE, CLAMP, P1, P0,
|
||||
gelu_poly, mma_chain)
|
||||
|
||||
def gp(c):
|
||||
c = np.asarray(c)
|
||||
return 32 * (c // 32) + 2 * ((c % 32) // 4) + (c % 2) + 16 * ((c % 4) // 2)
|
||||
|
||||
def pair_to_nat(M):
|
||||
return M[:, gp(np.arange(M.shape[1]))]
|
||||
|
||||
def nat_to_pair(M):
|
||||
out = np.empty_like(M); out[:, gp(np.arange(M.shape[1]))] = M; return out
|
||||
|
||||
def woff_table(base, Ks, Ns, K, N):
|
||||
k = np.arange(K)[:, None]; n = np.arange(N)[None, :]; j = n // 8
|
||||
return base + Ks * (k // 32) + Ns * (j // 2) + 64 * (n % 8) + 16 * ((k % 16) // 4) + 8 * (j % 2) + 4 * ((k % 32) // 16) + k % 4
|
||||
|
||||
def load_tokens(P, bx, by, planes):
|
||||
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, P.W), P.mirror(Y, P.H)
|
||||
out = np.zeros((64, 32 * planes))
|
||||
for p in range(planes):
|
||||
off = P.inp - P.abase + ((p * P.H + Y) * P.W + X) * 32
|
||||
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
|
||||
return out
|
||||
|
||||
def sumsq_tokens(x):
|
||||
"""Kernel-exact f16 sum of squares of natural-order token rows (planes of 32; lane t holds words t, t+4)."""
|
||||
sq = f16(x * x); planes = x.shape[1] // 32
|
||||
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
|
||||
lanes = []
|
||||
for t in range(4):
|
||||
def pair(pp, word, k):
|
||||
c = 32 * pp + 4 * word + 2 * k
|
||||
return sq[:, c], sq[:, c + 1]
|
||||
def plane_sum(word):
|
||||
p = [add(pair(pp, word, 0), pair(pp, word, 1)) for pp in range(planes)]
|
||||
if planes == 4: return add(add(p[0], p[1]), add(p[2], p[3]))
|
||||
if planes == 3: return add(add(p[0], p[1]), p[2])
|
||||
if planes == 2: return add(p[0], p[1])
|
||||
raise ValueError(planes)
|
||||
lanes.append(add(plane_sum(t), plane_sum(t + 4)))
|
||||
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
|
||||
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
|
||||
return f16(b[0] + b[1])[:, None]
|
||||
|
||||
def rms_from_sum(x, s, gamma):
|
||||
r = f16(1.0 / np.sqrt(f16(s + f16(2.0 ** -13))))
|
||||
return q8(f16(x * f16(r * gamma[None, :])))
|
||||
|
||||
def rel_bias(P, off):
|
||||
raw = P.f16vec(off, 256)
|
||||
i, j = np.meshgrid(np.arange(16), np.arange(16), indexing='ij')
|
||||
lane = 4 * (i % 8) + (j % 8) // 2; word = i // 8 + 2 * (j // 8)
|
||||
return raw[8 * lane + 2 * word + (j % 2)]
|
||||
|
||||
def attention(Qq, V, bias):
|
||||
S = f16(Qq @ Qq.T + bias)
|
||||
t = np.clip(f16(S * C_SCALE), -CLAMP, CLAMP)
|
||||
e1 = f16(t * (-t) + P1); poly = f16(t * e1 + P0)
|
||||
w = exp_trick(poly); s = row_sum16(w)
|
||||
r = f16(np.float32(1.0) / s.astype(np.float32)); Pm = f16(w * r)
|
||||
return f16(Pm @ V)
|
||||
|
||||
def tree_norm2(x):
|
||||
"""Kernel-exact f16 sum of squares of pair-ordered C-fragment rows (NT8 = C/8 n8 tiles)."""
|
||||
sq = f16(x * x); nt8 = x.shape[1] // 8
|
||||
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
|
||||
lanes = []
|
||||
for t in range(4):
|
||||
T = [(sq[:, 8 * j + 2 * t], sq[:, 8 * j + 2 * t + 1]) for j in range(nt8)]
|
||||
def quad(js): return add(add(T[js[0]], T[js[1]]), add(T[js[2]], T[js[3]]))
|
||||
if nt8 == 16:
|
||||
even = add(quad([0, 2, 4, 6]), quad([8, 10, 12, 14])); odd = add(quad([1, 3, 5, 7]), quad([9, 11, 13, 15]))
|
||||
elif nt8 == 12:
|
||||
even = add(quad([0, 2, 4, 6]), add(T[8], T[10])); odd = add(quad([1, 3, 5, 7]), add(T[9], T[11]))
|
||||
elif nt8 == 8:
|
||||
even = quad([0, 2, 4, 6]); odd = quad([1, 3, 5, 7])
|
||||
else:
|
||||
raise ValueError(nt8)
|
||||
lanes.append(add(even, odd))
|
||||
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
|
||||
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
|
||||
return f16(b[0] + b[1])[:, None]
|
||||
|
||||
def layout(C, heads, pre=0):
|
||||
"""Weight offsets of a Swin block with C channels starting at byte `pre`."""
|
||||
HS = 96 * C + 512
|
||||
L = dict(g1=pre, head=[pre + 2 * C + HS * h for h in range(heads)])
|
||||
E = pre + 2 * C + heads * HS
|
||||
L.update(bo=E, g2=E + 2 * C)
|
||||
M0 = E + 4 * C; NC = C // 8
|
||||
L['b1'] = [M0 if c == 0 else M0 + 66 * C + 64 + (64 * C + 64) * (c - 1) for c in range(NC)]
|
||||
L['b2'] = M0 + 64 * C + 64
|
||||
L['end'] = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1)
|
||||
return L
|
||||
|
||||
def swin_block(P, x0, C, heads, pre=0, stages=None):
|
||||
"""x0: 64 x C natural (f16 values). Returns (x2 pair order, out natural fp8 values)."""
|
||||
fp8 = lambda T: E4M3[P.wbuf[T]]
|
||||
L = layout(C, heads, pre)
|
||||
h1 = rms_from_sum(x0, sumsq_tokens(x0), P.f16vec(L['g1'], C)[gp(np.arange(C))])
|
||||
O = np.zeros((64, 32 * heads))
|
||||
for h in range(heads):
|
||||
hb = L['head'][h]
|
||||
Q = mma_chain(h1, fp8(woff_table(hb, 1024, 512, C, 32)))
|
||||
V = mma_chain(h1, fp8(woff_table(hb + 32 * C, 1024, 512, C, 32)))
|
||||
B = rel_bias(P, hb + 64 * C)
|
||||
for m in range(4):
|
||||
O[16 * m:16 * m + 16, 32 * h:32 * h + 32] = attention(q8(Q[16 * m:16 * m + 16]), V[16 * m:16 * m + 16], B)
|
||||
Wo = np.zeros((32 * heads, C), np.int64)
|
||||
for s in range(heads):
|
||||
Wo[32 * s:32 * s + 32] = woff_table(L['head'][s] + 64 * C + 512, 0, 512, 32, C)
|
||||
x1 = mma_chain(pair_to_nat(q8(O)), fp8(Wo), C0=f16(nat_to_pair(x0) + P.f16vec(L['bo'], C)[None, :]))
|
||||
h2 = pair_to_nat(rms_from_sum(x1, tree_norm2(x1), P.f16vec(L['g2'], C)))
|
||||
if stages is not None:
|
||||
stages.update(h1=h1, x1=pair_to_nat(x1), h2=h2, O=pair_to_nat(q8(O)))
|
||||
x2 = f16(x1 + P.f16vec(L['b2'], C)[None, :])
|
||||
idx = gp(np.arange(32))
|
||||
for c, B1 in enumerate(L['b1']):
|
||||
z = mma_chain(h2, fp8(woff_table(B1, 512, 16 * C, C, 32)), C0=np.broadcast_to(P.f16vec(B1 + 32 * C, 32), (64, 32)).copy())
|
||||
g = q8(gelu_poly(z))[:, idx]
|
||||
x2 = f16(g @ fp8(woff_table(B1 + 32 * C + 64, 0, 512, 32, C)) + x2)
|
||||
return x2, q8(pair_to_nat(x2)), L
|
||||
|
||||
TOK_X = (np.arange(64) & 3) + 4 * ((np.arange(64) >> 4) & 1)
|
||||
TOK_Y = ((np.arange(64) >> 2) & 3) + 4 * (np.arange(64) >> 5)
|
||||
|
||||
def patch_merge_input(out, C):
|
||||
"""16 merged rows (r -> (r%4, r//4)) x 4C: column C*q + ch = out[token(2mx+q%2, 2my+q//2), ch]."""
|
||||
A = np.zeros((16, 4 * C))
|
||||
for r in range(16):
|
||||
mx, my = r % 4, r // 4
|
||||
for q in range(4):
|
||||
t = np.nonzero((TOK_X == 2 * mx + (q & 1)) & (TOK_Y == 2 * my + (q >> 1)))[0][0]
|
||||
A[r, C * q:C * q + C] = out[t]
|
||||
return A
|
||||
|
||||
def load_at(P, ptr, bx, by, planes, W=None, H=None, low=False):
|
||||
W = W or P.W; H = H or P.H
|
||||
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, W), P.mirror(Y, H)
|
||||
out = np.zeros((64, 32 * planes))
|
||||
for p in range(planes):
|
||||
off = ptr - P.abase + ((p * H + Y) * W + X) * 32
|
||||
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
|
||||
return out
|
||||
|
||||
def decoder_input(P, bx, by, C, CIN, NW, skip_ptr):
|
||||
"""x0 = f16(q8(patch_expand(low-res CIN)) + skip), expand output col G = (4C/NW) w + n -> q = G // C, ch = ginv(G % C)."""
|
||||
fp8 = lambda T: E4M3[P.wbuf[T]]
|
||||
W2, H2 = P.W // 2, P.H // 2
|
||||
A = np.zeros((16, CIN))
|
||||
for r in range(16):
|
||||
MX = P.mirror(np.array((8 * bx - P.sx) // 2 + r % 4), W2); MY = P.mirror(np.array((8 * by - P.sy) // 2 + r // 4), H2)
|
||||
for p in range(CIN // 32):
|
||||
off = P.inp - P.abase + ((p * H2 + MY) * W2 + MX) * 32
|
||||
A[r, 32 * p:32 * p + 32] = E4M3[P.arena[off:off + 32]]
|
||||
cols = 4 * C // NW
|
||||
E_ = np.concatenate([mma_chain(A, fp8(woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)),
|
||||
C0=np.broadcast_to(P.f16vec(4 * C * CIN + 2 * cols * w, cols), (16, cols)).copy()) for w in range(NW)], 1)
|
||||
skip = load_at(P, skip_ptr, bx, by, C // 32)
|
||||
x0 = np.zeros((64, C))
|
||||
for T in range(64):
|
||||
tx, ty = TOK_X[T], TOK_Y[T]; r = tx // 2 + 4 * (ty // 2); q = (tx & 1) + 2 * (ty & 1)
|
||||
G = np.arange(C) + q * C # global expand columns of sub-position q (pair order within C)
|
||||
x0[T] = q8(E_[r, G])[gp(np.arange(C))]
|
||||
return f16(x0 + skip)
|
||||
@@ -0,0 +1,141 @@
|
||||
// Standalone check of the texture-kernel GEMM tails (kernels/tex/enc0_tail.hip, dec0_head.hip) outside the
|
||||
// ZLUDA-compiled kernels they are linked into: random e4m3 weights and activations, several tile positions
|
||||
// (including partial tiles at the image border), all outputs written to one file. Builds of different WMMA
|
||||
// layouts (gfx11, the gfx12-layout shim, gfx12, gfx12 native FP8) are then compared byte for byte (same
|
||||
// arithmetic) or by value (gfx12 rounding). The tails touch no images, so this also runs in rocjitsu.
|
||||
//
|
||||
// build: clang++ -x hip --offload-arch=ARCH -O3 [-DD4R_WMMA_LAYOUT=12] [-DD4R_TEX_FP8] -I kernels/tex
|
||||
// kernels/tools/tail_check.hip -o tail_check usage: tail_check OUT.bin
|
||||
#include <hip/hip_runtime.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <vector>
|
||||
|
||||
// the surface-store helpers of tex_common.h reference ZLUDA's helper library, which the tails do not use
|
||||
extern "C" __device__ void __zluda_ptx_impl_surfobj_b_2d_v2_b16_zero(uint64_t, int __attribute__((ext_vector_type(2))),
|
||||
short __attribute__((ext_vector_type(2))))
|
||||
{
|
||||
}
|
||||
|
||||
// both tails define the same hook globals
|
||||
#define D4R_KERNEL_NAME tail_enc0
|
||||
#define d4r_prep_blocks d4r_prep_blocks_enc0
|
||||
#define d4r_prep_key_at d4r_prep_key_at_enc0
|
||||
#include "enc0_tail.hip"
|
||||
#undef d4r_prep_blocks
|
||||
#undef d4r_prep_key_at
|
||||
#define d4r_prep_blocks d4r_prep_blocks_dec0
|
||||
#define d4r_prep_key_at d4r_prep_key_at_dec0
|
||||
#include "dec0_head.hip"
|
||||
|
||||
extern "C" __global__ void __launch_bounds__(32) run_enc0(const uint32_t* lds_in, const uint8_t* w, uint32_t* out,
|
||||
uint32_t dimx, uint32_t dimy, int ox, int oy)
|
||||
{
|
||||
__shared__ uint32_t s[512 + 4];
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
s[i] = lds_in[i];
|
||||
__syncthreads();
|
||||
d4r_enc0_tail((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)out, dimx,
|
||||
dimy, ox, oy);
|
||||
}
|
||||
|
||||
extern "C" __global__ void __launch_bounds__(32) run_dec0(const uint8_t* w, const uint32_t* in, uint32_t* staged,
|
||||
uint32_t dimx, uint32_t dimy, int ox, int oy)
|
||||
{
|
||||
__shared__ uint32_t s[512 + 4];
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
s[i] = 0xdeadbeefu;
|
||||
__syncthreads();
|
||||
d4r_dec0_head((uint32_t)(uintptr_t)(__attribute__((address_space(3))) void*)s, (uint64_t)w, (uint64_t)in, dimx,
|
||||
dimy, ox, oy);
|
||||
__syncthreads();
|
||||
for (int i = threadIdx.x; i < 512; i += 32)
|
||||
staged[i] = s[i];
|
||||
}
|
||||
|
||||
#define CHECK(x) \
|
||||
do \
|
||||
{ \
|
||||
hipError_t e_ = (x); \
|
||||
if (e_ != hipSuccess) \
|
||||
{ \
|
||||
fprintf(stderr, "%s: %s\n", #x, hipGetErrorString(e_)); \
|
||||
return 1; \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
static uint8_t e4m3(uint32_t r)
|
||||
{
|
||||
// Keep the matrix values in a useful range so saturation cannot hide a wrong WMMA result.
|
||||
const uint8_t magnitude = (r & 31) == 0 ? 0 : (uint8_t)(0x18 + (r & 31));
|
||||
return (uint8_t)(magnitude | ((r >> 8) & 0x80));
|
||||
}
|
||||
static uint16_t f16_bits(float v)
|
||||
{
|
||||
return __builtin_bit_cast(uint16_t, (_Float16)v);
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc < 2)
|
||||
{
|
||||
fprintf(stderr, "usage: %s OUT.bin\n", argv[0]);
|
||||
return 2;
|
||||
}
|
||||
srand(12345);
|
||||
auto rnd = []() { return (uint32_t)rand() * 2654435761u + (uint32_t)rand(); };
|
||||
// weights: e4m3 fragments, f16 biases (enc0 at 12288, dec0 at 8192)
|
||||
std::vector<uint8_t> w(16384 + 512);
|
||||
for (auto& b : w)
|
||||
b = e4m3(rnd());
|
||||
for (int i = 0; i < 64; ++i)
|
||||
*(uint16_t*)&w[12288 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
|
||||
for (int i = 0; i < 128; ++i)
|
||||
*(uint16_t*)&w[8192 + 2 * i] = f16_bits((float)(int)(rnd() % 2001 - 1000) / 256.0f);
|
||||
const uint32_t dimx = 20, dimy = 14, W2 = dimx / 2, H2 = dimy / 2, rowstride = W2 * 8, plane = H2 * rowstride;
|
||||
std::vector<uint32_t> lds(512), in(2 * plane);
|
||||
for (auto& v : lds)
|
||||
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
|
||||
for (auto& v : in)
|
||||
v = (uint32_t)e4m3(rnd()) | (uint32_t)e4m3(rnd()) << 8 | (uint32_t)e4m3(rnd()) << 16 | (uint32_t)e4m3(rnd()) << 24;
|
||||
|
||||
uint8_t *dw, *dummy;
|
||||
uint32_t *dlds, *din, *dout, *dstaged;
|
||||
CHECK(hipMalloc(&dw, w.size()));
|
||||
CHECK(hipMalloc(&dlds, lds.size() * 4));
|
||||
CHECK(hipMalloc(&din, in.size() * 4));
|
||||
CHECK(hipMalloc(&dout, 2 * plane * 4));
|
||||
CHECK(hipMalloc(&dstaged, 512 * 4));
|
||||
CHECK(hipMalloc(&dummy, 512));
|
||||
CHECK(hipMemcpy(dw, w.data(), w.size(), hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(dlds, lds.data(), lds.size() * 4, hipMemcpyHostToDevice));
|
||||
CHECK(hipMemcpy(din, in.data(), in.size() * 4, hipMemcpyHostToDevice));
|
||||
// the prep kernels read the weight pointer at parameter offset 0
|
||||
Enc0Params pe{};
|
||||
pe.w = (uint64_t)dw;
|
||||
Dec0Params pd{};
|
||||
pd.w = (uint64_t)dw;
|
||||
hipLaunchKernelGGL(tail_enc0_prep, dim3(4), dim3(128), 0, 0, pe);
|
||||
hipLaunchKernelGGL(rrlite_dec0_4x4_folded_prep, dim3(4), dim3(128), 0, 0, pd);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
|
||||
FILE* f = fopen(argv[1], "wb");
|
||||
const int offsets[][2] = {{0, 0}, {8, 4}, {12, 6}, {16, 8}, {-4, -2}, {18, 12}};
|
||||
std::vector<uint32_t> host(2 * plane);
|
||||
for (const auto& o : offsets)
|
||||
{
|
||||
CHECK(hipMemset(dout, 0x5a, 2 * plane * 4));
|
||||
hipLaunchKernelGGL(run_enc0, dim3(1), dim3(32), 0, 0, dlds, dw, dout, dimx, dimy, o[0], o[1]);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(host.data(), dout, 2 * plane * 4, hipMemcpyDeviceToHost));
|
||||
fwrite(host.data(), 4, 2 * plane, f);
|
||||
hipLaunchKernelGGL(run_dec0, dim3(1), dim3(32), 0, 0, dw, din, dstaged, dimx, dimy, o[0], o[1]);
|
||||
CHECK(hipDeviceSynchronize());
|
||||
CHECK(hipMemcpy(host.data(), dstaged, 512 * 4, hipMemcpyDeviceToHost));
|
||||
fwrite(host.data(), 4, 512, f);
|
||||
}
|
||||
fclose(f);
|
||||
printf("wrote %s\n", argv[1]);
|
||||
return 0;
|
||||
}
|
||||
Reference in new issue
Block a user