d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton
This commit is contained in:
ayo committed
commit
118173ecd7
128 files changed
+31333
No files matched your search
@@ -0,0 +1,167 @@
|
||||
// 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;
|
||||
const unsigned prepGrid = getenv("PREP_GRID") ? atoi(getenv("PREP_GRID")) : 88;
|
||||
|
||||
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 == "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,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,877 @@
|
||||
#!/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):
|
||||
src = open(path).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]
|
||||
for l in st.labels:
|
||||
if l in pending:
|
||||
active |= pending.pop(l)
|
||||
if not active.any():
|
||||
# jump to the nearest label someone waits at
|
||||
targets = sorted(self.labels[l] for l in 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
|
||||
raise NotImplementedError('divergent backward branch at ' + st.text)
|
||||
active = active & ~taking
|
||||
if taking.any():
|
||||
pending[st.args[0]] = pending.get(st.args[0], 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,408 @@
|
||||
#!/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
|
||||
|
||||
|
||||
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 = f16(np.sum(x0.astype(F32).astype(F64) ** 2, axis=1))
|
||||
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 = f16(np.sum(x0.astype(F32).astype(F64) ** 2, axis=1))
|
||||
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,168 @@
|
||||
#!/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):
|
||||
"""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)))
|
||||
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)
|
||||
Reference in new issue
Block a user