Engine: native presets K and M in the D3D12 shim, wider game coverage
A standalone Vulkan inference engine (engine/) for DLSS presets K and M that runs from weights extracted from the user's own nvngx_dlss.dll, with d4r's own input, network and reconstruction shaders or the native HIP network layers, recorded into the game's command list by the shim. Opt-in ([Engine] Enabled / D4R_ENGINE=1); anything uncovered falls back to the CUDA/ZLUDA path with the reason logged. Byte-identical to NGX's CUDA backend under the phase-capable reference (1706x960 -> 2560x1440, native HIP network, eleven frames, RX 9070 XT): - K render- and display-resolution motion vectors, HDR and LDR - K game-supplied exposure and automatic exposure with any pre-exposure and exposure scale (one residual: effective exposure exactly 2 or 4) - K and M with the DLSS indicator's inverted-axis parameters - M with any exposure settings (NGX ignores them for M) K packages need the current compile_k.py (marker D4RO0003). Also: HIP shared buffers are pooled across feature recreation (ROCm leaks destroyed imports), hipnet.bin must match the preset, M-only games do not load the K model, install_d4r_runtime.sh stages k, k-ldr, m, m-ldr with hip/, the activation table has a generator (gen_activation_m.py), and package_release.sh can bundle the engine's shaders (never its weights). Signed-off-by: woutermans <[email protected]>
This commit is contained in:
woutermans committed
1 parent
16cfc1f063
commit
e0d1549a21
82 files changed
+13335
-96
No files matched your search
@@ -0,0 +1,47 @@
|
||||
# Vulkan kernels (experimental)
|
||||
|
||||
Tools for running DLSS texture kernels as Vulkan compute shaders instead of through CUDA/HIP. Nothing here is
|
||||
wired into the shim yet; it is a kernel-level proof with a byte-exact check against captured frames.
|
||||
|
||||
| File | Role |
|
||||
|---|---|
|
||||
| `ptx2glsl.py` | Translates one PTX kernel from your DLSS DLL into GLSL. Every instruction becomes the same operation `kernels/native/ptx2hip.py` emits, gotos become nested `if`/`else` and `do`/`while` (each branch rejoins at its immediate post-dominator), and registers that only hold half floats are typed `float16_t` / `f16vec2`. |
|
||||
| `kvk.cpp` | Headless runner: binds captured textures, surfaces, buffers and the parameter block, dispatches, writes the surfaces back and times the dispatch. |
|
||||
| `kcase.py` | Builds a `kvk` manifest from a capture and compares `kvk`'s output with the captured surfaces. |
|
||||
|
||||
## Reproducing the K output kernel check
|
||||
|
||||
```sh
|
||||
# 1. capture (bridge options that already exist): every launch's resources after it ran, plus the arguments
|
||||
D4R_CUDA_VERBOSE=1 D4R_CUDA_LAUNCH_STATS=1 D4R_CUDA_LAUNCH_DUMP_DIR=$CAP/dump \
|
||||
D4R_CUDA_REPLAY_DUMP_DIR=$CAP/replay D4R_CUDA_REPLAY_DUMP_FILTER=hiluma_engine_output D4R_CUDA_REPLAY_DUMP_LIMIT=8 \
|
||||
<run the NGX CUDA benchmark, preset K>
|
||||
# 2. translate and compile
|
||||
K=hiluma_engine_output_depthinv_mvlo_hdr_max_v2_rel
|
||||
python3 ptx2glsl.py MODULE.ptx $K > k.comp && glslangValidator -V --target-env vulkan1.3 k.comp -o k.spv
|
||||
g++ -O2 -std=c++17 -o kvk kvk.cpp -lvulkan
|
||||
# 3. run frame N and compare (the three linear-filtered textures are named by parameter offset)
|
||||
python3 kcase.py make $CAP $K N case 200:linear 224:linear 232:linear
|
||||
./kvk k.spv case/manifest 150 64 # 150 dispatches per submit, 64-wide subgroups
|
||||
python3 kcase.py cmp case
|
||||
```
|
||||
|
||||
Sampler settings come from the `cuTexObjectCreate` lines of the verbose log (all clamp, normalized coordinates;
|
||||
filter per object).
|
||||
|
||||
## Result (RX 9070 XT, RADV / Mesa 26.2.4, DLSS 310.7, 2026-10-08)
|
||||
|
||||
Byte-identical to the native HIP kernel's four output surfaces on six consecutive frames at 640x360 -> 1280x720 and
|
||||
one frame at 1920x1080 -> 3840x2160. At 4K: 1.18 ms per dispatch under sustained load (GPU clock 2.38 GHz),
|
||||
1.05 ms with 0.3 ms gaps (2.7 GHz). The native kernel measures 1.16 ms in the live benchmark at about 3.1 GHz. The
|
||||
two were not timed in the same run or clock state.
|
||||
|
||||
## Pitfalls found
|
||||
|
||||
- Shader-visible buffers must be allocated from a `DEVICE_LOCAL | HOST_VISIBLE` memory type. The first
|
||||
host-visible type on RADV is system memory; reading the network output from it made every variant take 2.9 ms.
|
||||
- A program-counter loop around the whole shader (`while (pc != END) { if (pc == k) ... }`) keeps every value
|
||||
live across the back edge: 35 ms. Guards without the loop: 2.7 ms (VRAM buffers). Real nesting: 1.69 ms.
|
||||
- Packing half floats into `uint` registers costs about 840 `v_mov_b16`; native half types: 1.35 -> 1.18 ms.
|
||||
- A 64K-entry table for `ex2.approx.f16x2` is exact but slower than `v_exp_f32`; dropping `precise` gains under 1%
|
||||
and changes about 0.005% of the output values.
|
||||
@@ -0,0 +1,50 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Condense a ptx2glsl.py translation for reading: drops declarations, prints float constants as decimals and
|
||||
substitutes registers that are assigned once and used once into their use. The result is not meant to compile.
|
||||
|
||||
usage: condense.py KERNEL.comp > KERNEL.txt
|
||||
"""
|
||||
import re, struct, sys
|
||||
src = open(sys.argv[1]).read().split('\n')
|
||||
start = next(i for i, l in enumerate(src) if l.startswith('void main'))
|
||||
body = [l for l in src[start:] if not re.match(r'\s*(precise )?(uint|float|bool|uint64_t|float16_t|f16vec2|int) [a-z0-9_]+ = [^;]*(, [a-z0-9_]+ = [^;,]*)*;\s*$', l) or '(' in l.split('=', 1)[1][:3]]
|
||||
def fl(m):
|
||||
v = struct.unpack('<f', struct.pack('<I', int(m.group(1), 16)))[0]
|
||||
return repr(round(v, 7)) if abs(v) < 1e7 and (abs(v) > 1e-5 or v == 0) else f'{v:.6g}'
|
||||
body = [re.sub(r'uintBitsToFloat\(0x([0-9A-Fa-f]{8})u\)', fl, l) for l in body]
|
||||
body = [re.sub(r'\b(\d+)u\b', r'\1', l) for l in body]
|
||||
body = [re.sub(r' \((\d+) & 31\)', r' \1', l) for l in body]
|
||||
tok = re.compile(r'\b(?:r|f|rs|rd|p)\d+\b')
|
||||
for _ in range(12):
|
||||
assigns, uses = {}, {}
|
||||
for i, l in enumerate(body):
|
||||
m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', l)
|
||||
rhs = l
|
||||
if m:
|
||||
assigns.setdefault(m.group(2), []).append(i); rhs = m.group(3)
|
||||
for t in tok.findall(rhs): uses.setdefault(t, []).append(i)
|
||||
changed = False
|
||||
for name, at in assigns.items():
|
||||
u = uses.get(name, [])
|
||||
if len(at) != 1 or len(u) != 1 or u[0] <= at[0]: continue
|
||||
a, b = at[0], u[0]
|
||||
if body[a] is None or body[b] is None: continue
|
||||
ind = lambda s: len(s) - len(s.lstrip())
|
||||
# only within straight-line code at the same nesting, and short enough to stay readable
|
||||
if any(body[k] is not None and ('{' in body[k] or '}' in body[k]) for k in range(a, b)): continue
|
||||
m = re.match(r'(\s*)((?:r|f|rs|rd|p)\d+) = (.*);\s*$', body[a])
|
||||
expr = m.group(3)
|
||||
if len(expr) > 90: continue
|
||||
# the expression's inputs must not be reassigned in between
|
||||
ins = set(tok.findall(expr))
|
||||
if any(body[k] is not None and re.match(r'\s*(%s) = ' % '|'.join(map(re.escape, ins)) if ins else r'$^', body[k]) for k in range(a + 1, b)): continue
|
||||
simple = re.fullmatch(r'[\w.]+|\w+\([^()]*\)', expr)
|
||||
body[b] = re.sub(r'\b%s\b' % name, expr if simple else f'({expr})', body[b], count=1)
|
||||
body[a] = None; changed = True
|
||||
body = [l for l in body if l is not None]
|
||||
if not changed: break
|
||||
for a, b in ((r'uint\(uint\(([^()]*)\)\)', r'\1'), (r'\buint\(((?:r|rs)\d+)\)', r'\1'), (r'uint64_t\((rd\d+)\)', r'\1')):
|
||||
body = [re.sub(a, b, l) for l in body]
|
||||
# scripts/check_redistributable.py refuses files carrying this line: the input was NVIDIA's code.
|
||||
print(f'// Generated by vulkan/condense.py from NVIDIA PTX ({sys.argv[1]}); do not commit or redistribute.')
|
||||
print('\n'.join(l.rstrip() for l in body if l.strip()))
|
||||
@@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Build a kvk manifest for one captured launch and compare kvk's outputs with the capture.
|
||||
|
||||
usage: kcase.py make CAPTURE_DIR KERNEL FRAME OUTDIR [--tex OFF:linear ...] (FRAME counts that kernel's launches from 0)
|
||||
kcase.py cmp OUTDIR
|
||||
CAPTURE_DIR holds dump/ (D4R_CUDA_LAUNCH_DUMP_DIR) and replay/ (D4R_CUDA_REPLAY_DUMP_DIR) of the same run.
|
||||
"""
|
||||
import glob, os, re, sys
|
||||
import numpy as np
|
||||
PING = set()
|
||||
|
||||
def launches(cap, kernel):
|
||||
return sorted({int(re.match(r'launch-(\d+)-', os.path.basename(f)).group(1)) for f in glob.glob(f'{cap}/dump/launch-*-{kernel}-arg*')})
|
||||
|
||||
def res(cap, kernel, seq):
|
||||
out = {}
|
||||
for f in glob.glob(f'{cap}/dump/launch-{seq:03d}-{kernel}-arg*'):
|
||||
m = re.search(r'-arg(\d+)-(texture|surface|buffer)-(.*)-fmt(-?\d+)\.bin$', f)
|
||||
out[int(m.group(1))] = (m.group(2), m.group(3), int(m.group(4)), f)
|
||||
return out
|
||||
|
||||
if sys.argv[1] == 'make':
|
||||
cap, kernel, frame, outdir = sys.argv[2], sys.argv[3], int(sys.argv[4]), sys.argv[5]
|
||||
linear = {int(a.split(':')[0]) for a in sys.argv[6:] if a.endswith(':linear')}
|
||||
os.makedirs(outdir, exist_ok=True)
|
||||
seqs = launches(cap, kernel); seq = seqs[frame]
|
||||
cur = res(cap, kernel, seq)
|
||||
man = open(f'{cap}/replay/replay-{seq:06d}-{kernel}/manifest.txt').read().split('\n')
|
||||
grid = next(l for l in man if l.startswith('launch')).split()
|
||||
lines = [f'param {cap}/replay/replay-{seq:06d}-{kernel}/args.bin', f'grid {grid[1]} {grid[2]}']
|
||||
ptr = {int(l.split()[1]): int(l.split()[3]) for l in man if l.startswith('pointer')}
|
||||
expect = []
|
||||
for off, (kind, shape, fmt, f) in sorted(cur.items()):
|
||||
if kind == 'buffer':
|
||||
lines.append(f'buf {off} {f} {ptr[off]}'); continue
|
||||
w, h, c = (int(x) for x in shape.split('x'))
|
||||
bits = {32: 32, 194: 8, 192: 8, 193: 8, 200: 9, 198: 9, 199: 9}.get(fmt, 16) # 19x: UNORM8, 198..200: SNORM8
|
||||
if kind == 'texture':
|
||||
lines.append(f'tex {off} {f} {w} {h} {c} {bits} {"linear" if off in linear else "nearest"}')
|
||||
else:
|
||||
# the surface's contents before this launch: its dump from the latest earlier launch that had the same data layout
|
||||
init = '-'
|
||||
for s in reversed([s for s in seqs if s < seq]):
|
||||
prev = res(cap, kernel, s)
|
||||
# ping-pong surfaces alternate; take the one whose size matches and which is not this frame's input
|
||||
if off in prev and prev[off][1] == shape and (s == seqs[frame - 2] if frame >= 2 and off in PING else True):
|
||||
init = prev[off][3]; break
|
||||
lines.append(f'surf {off} {w} {h} {c} {bits} {init} {outdir}/out-{off}.bin')
|
||||
expect.append((off, f, c, bits))
|
||||
open(f'{outdir}/manifest', 'w').write('\n'.join(lines) + '\n')
|
||||
open(f'{outdir}/expect', 'w').write('\n'.join(f'{o} {f} {c} {b}' for o, f, c, b in expect) + '\n')
|
||||
print('\n'.join(lines))
|
||||
else:
|
||||
outdir = sys.argv[2]
|
||||
ok = True
|
||||
for l in open(f'{outdir}/expect').read().split('\n'):
|
||||
if not l:
|
||||
continue
|
||||
off, f, c, bits = l.split()
|
||||
dt = np.uint16 if bits == '16' else np.uint32
|
||||
a = np.fromfile(f, dtype=dt); b = np.fromfile(f'{outdir}/out-{off}.bin', dtype=dt)
|
||||
diff = a != b
|
||||
fa = a.view(np.float16 if bits == '16' else np.float32).astype(np.float64); fb = b.view(np.float16 if bits == '16' else np.float32).astype(np.float64)
|
||||
with np.errstate(invalid='ignore'):
|
||||
d = np.abs(fa - fb); d = d[np.isfinite(d)]
|
||||
ok &= not diff.any()
|
||||
print(f'surface {off}: {a.size} values, {int(diff.sum())} differ ({diff.mean() * 100:.4f}%), max |diff| {d.max() if d.size else 0:.6g}, mean |ref| {np.nanmean(np.abs(fa[np.isfinite(fa)])):.4g}')
|
||||
print('BYTE-IDENTICAL' if ok else 'DIFFERENT')
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
// Runs one translated DLSS texture kernel as a Vulkan compute shader on captured inputs.
|
||||
// kvk SHADER.spv MANIFEST [repeats] [subgroup]
|
||||
// MANIFEST lines (bindings are assigned in the order ptx2glsl.py uses: parameter block, textures by
|
||||
// offset, surfaces by offset):
|
||||
// param FILE
|
||||
// grid X Y
|
||||
// tex OFF FILE W H CHANNELS BITS linear|nearest
|
||||
// surf OFF W H CHANNELS BITS INITFILE|- OUTFILE
|
||||
// buf OFF FILE DELTA (8-byte pointer at OFF = buffer address + DELTA)
|
||||
// lut FILE (storage buffer bound after the surfaces)
|
||||
// bufout OFF FILE (after the dispatch, write the pointer buffer of parameter offset OFF to FILE)
|
||||
// outbuf FILE BYTES (storage buffer bound after those, written to FILE after the dispatch)
|
||||
#include <vulkan/vulkan.h>
|
||||
#include <algorithm>
|
||||
#include <chrono>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <fstream>
|
||||
#include <map>
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#define CK(x) do { VkResult r_ = (x); if (r_ != VK_SUCCESS) { fprintf(stderr, "%s failed: %d\n", #x, r_); exit(1); } } while (0)
|
||||
static std::vector<char> slurp(const std::string& p)
|
||||
{
|
||||
std::ifstream f(p, std::ios::binary);
|
||||
if (!f) { fprintf(stderr, "cannot read %s\n", p.c_str()); exit(1); }
|
||||
return std::vector<char>((std::istreambuf_iterator<char>(f)), {});
|
||||
}
|
||||
static VkDevice dev; static VkPhysicalDeviceMemoryProperties mp;
|
||||
static VkDeviceMemory alloc(VkMemoryRequirements mr, VkMemoryPropertyFlags want, bool bda)
|
||||
{
|
||||
uint32_t mt = ~0u;
|
||||
for (uint32_t t = 0; t < mp.memoryTypeCount; ++t) if ((mr.memoryTypeBits >> t & 1) && (mp.memoryTypes[t].propertyFlags & want) == want) { mt = t; break; }
|
||||
VkMemoryAllocateFlagsInfo fi{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_FLAGS_INFO}; fi.flags = VK_MEMORY_ALLOCATE_DEVICE_ADDRESS_BIT;
|
||||
VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO, bda ? &fi : nullptr}; ai.allocationSize = mr.size; ai.memoryTypeIndex = mt;
|
||||
VkDeviceMemory m; CK(vkAllocateMemory(dev, &ai, nullptr, &m)); return m;
|
||||
}
|
||||
struct Buf { VkBuffer b; VkDeviceMemory m; void* map; VkDeviceSize size; uint64_t addr; };
|
||||
static Buf make_buf(VkDeviceSize size, VkBufferUsageFlags usage)
|
||||
{
|
||||
Buf r{}; r.size = size;
|
||||
VkBufferCreateInfo ci{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; ci.size = size; ci.usage = usage | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT;
|
||||
CK(vkCreateBuffer(dev, &ci, nullptr, &r.b));
|
||||
VkMemoryRequirements mr; vkGetBufferMemoryRequirements(dev, r.b, &mr);
|
||||
// shader-visible data must live in VRAM: the first host-visible type on RADV is system memory, which the GPU reads over PCIe
|
||||
const VkMemoryPropertyFlags vram = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT | VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;
|
||||
bool has_vram = false;
|
||||
for (uint32_t t = 0; t < mp.memoryTypeCount; ++t) has_vram |= (mr.memoryTypeBits >> t & 1) && (mp.memoryTypes[t].propertyFlags & vram) == vram;
|
||||
r.m = alloc(mr, has_vram ? vram : VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, true);
|
||||
CK(vkBindBufferMemory(dev, r.b, r.m, 0)); CK(vkMapMemory(dev, r.m, 0, size, 0, &r.map));
|
||||
VkBufferDeviceAddressInfo ai{VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO}; ai.buffer = r.b; r.addr = vkGetBufferDeviceAddress(dev, &ai);
|
||||
return r;
|
||||
}
|
||||
struct Img { VkImage i; VkImageView v; VkFormat f; uint32_t w, h, bpp; std::string init, out; bool surf, linear; int off; };
|
||||
static VkFormat fmt_of(int ch, int bits)
|
||||
{
|
||||
// 8: UNORM, 9: SNORM, one byte per channel (CUDA's UNORM_INT8Xn / SNORM_INT8Xn arrays)
|
||||
if (bits == 8) return ch == 1 ? VK_FORMAT_R8_UNORM : ch == 2 ? VK_FORMAT_R8G8_UNORM : VK_FORMAT_R8G8B8A8_UNORM;
|
||||
if (bits == 9) return ch == 1 ? VK_FORMAT_R8_SNORM : ch == 2 ? VK_FORMAT_R8G8_SNORM : VK_FORMAT_R8G8B8A8_SNORM;
|
||||
if (bits == 10 && ch == 4) return VK_FORMAT_A2B10G10R10_UNORM_PACK32;
|
||||
if (bits == 16) return ch == 1 ? VK_FORMAT_R16_SFLOAT : ch == 2 ? VK_FORMAT_R16G16_SFLOAT : VK_FORMAT_R16G16B16A16_SFLOAT;
|
||||
return ch == 1 ? VK_FORMAT_R32_SFLOAT : ch == 2 ? VK_FORMAT_R32G32_SFLOAT : VK_FORMAT_R32G32B32A32_SFLOAT;
|
||||
}
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
if (argc < 3) return 1;
|
||||
const int repeats = argc > 3 ? atoi(argv[3]) : 1; const uint32_t subgroup = argc > 4 ? atoi(argv[4]) : 32;
|
||||
VkApplicationInfo app{VK_STRUCTURE_TYPE_APPLICATION_INFO}; app.apiVersion = VK_API_VERSION_1_3;
|
||||
VkInstanceCreateInfo ici{VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO}; ici.pApplicationInfo = &app;
|
||||
VkInstance inst; CK(vkCreateInstance(&ici, nullptr, &inst));
|
||||
uint32_t np = 8; VkPhysicalDevice pds[8]; vkEnumeratePhysicalDevices(inst, &np, pds);
|
||||
VkPhysicalDevice pd = pds[0];
|
||||
for (uint32_t i = 0; i < np; ++i) { VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(pds[i], &p); if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) pd = pds[i]; }
|
||||
VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES};
|
||||
VkPhysicalDeviceVulkan12Features v12{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, &v13};
|
||||
VkPhysicalDeviceVulkan11Features v11{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, &v12};
|
||||
VkPhysicalDeviceFeatures2 f2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, &v11};
|
||||
vkGetPhysicalDeviceFeatures2(pd, &f2);
|
||||
vkGetPhysicalDeviceMemoryProperties(pd, &mp);
|
||||
uint32_t nq = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(pd, &nq, qf);
|
||||
uint32_t qi = 0; for (uint32_t i = 0; i < nq; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { qi = i; break; }
|
||||
float prio = 1; VkDeviceQueueCreateInfo qci{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qci.queueFamilyIndex = qi; qci.queueCount = 1; qci.pQueuePriorities = &prio;
|
||||
VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci;
|
||||
CK(vkCreateDevice(pd, &dci, nullptr, &dev));
|
||||
VkQueue q; vkGetDeviceQueue(dev, qi, 0, &q);
|
||||
VkCommandPoolCreateInfo cpci{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; cpci.queueFamilyIndex = qi; cpci.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
|
||||
VkCommandPool cp; CK(vkCreateCommandPool(dev, &cpci, nullptr, &cp));
|
||||
VkCommandBufferAllocateInfo cai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; cai.commandPool = cp; cai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; cai.commandBufferCount = 1;
|
||||
VkCommandBuffer cb; CK(vkAllocateCommandBuffers(dev, &cai, &cb));
|
||||
auto run = [&](auto&& record) {
|
||||
VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); record(); CK(vkEndCommandBuffer(cb));
|
||||
VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb;
|
||||
CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q));
|
||||
};
|
||||
// ---------------------------------------------------------------- manifest
|
||||
std::vector<char> param; uint32_t gx = 1, gy = 1;
|
||||
std::vector<Img> imgs; std::map<int, std::pair<Buf, int64_t>> bufs; std::vector<Buf> luts; std::vector<std::string> lutOut; std::map<int, std::string> bufOut;
|
||||
std::ifstream mf(argv[2]); std::string line;
|
||||
while (std::getline(mf, line))
|
||||
{
|
||||
std::istringstream ls(line); std::string k; ls >> k;
|
||||
if (k == "param") { std::string f; ls >> f; param = slurp(f); }
|
||||
else if (k == "grid") ls >> gx >> gy;
|
||||
else if (k == "tex" || k == "surf")
|
||||
{
|
||||
Img im{}; int ch, bits; std::string mode; im.surf = k == "surf";
|
||||
if (im.surf) ls >> im.off >> im.w >> im.h >> ch >> bits >> im.init >> im.out;
|
||||
else { ls >> im.off >> im.init >> im.w >> im.h >> ch >> bits >> mode; im.linear = mode == "linear"; }
|
||||
im.f = fmt_of(ch, bits); im.bpp = bits < 16 ? ch : ch * bits / 8;
|
||||
imgs.push_back(im);
|
||||
}
|
||||
else if (k == "lut")
|
||||
{
|
||||
std::string f; ls >> f; auto data = slurp(f);
|
||||
Buf b = make_buf(data.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(b.map, data.data(), data.size()); luts.push_back(b); lutOut.push_back("");
|
||||
}
|
||||
else if (k == "bufout") { int off; std::string f; ls >> off >> f; bufOut[off] = f; }
|
||||
else if (k == "outbuf")
|
||||
{
|
||||
std::string f; VkDeviceSize n; ls >> f >> n;
|
||||
Buf b = make_buf(n, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memset(b.map, 0, n); luts.push_back(b); lutOut.push_back(f);
|
||||
}
|
||||
else if (k == "buf")
|
||||
{
|
||||
int off; std::string f; int64_t delta; ls >> off >> f >> delta;
|
||||
auto data = slurp(f); Buf b = make_buf(data.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(b.map, data.data(), data.size());
|
||||
bufs[off] = {b, delta};
|
||||
}
|
||||
}
|
||||
for (auto& [off, b] : bufs) { uint64_t a = b.first.addr + b.second; memcpy(param.data() + off, &a, 8); }
|
||||
std::sort(imgs.begin(), imgs.end(), [](const Img& a, const Img& b) { return a.surf != b.surf ? !a.surf : a.off < b.off; });
|
||||
Buf pbuf = make_buf((param.size() + 3) & ~3ull, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT); memcpy(pbuf.map, param.data(), param.size());
|
||||
// ---------------------------------------------------------------- images
|
||||
VkSamplerCreateInfo sci{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO};
|
||||
sci.addressModeU = sci.addressModeV = sci.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; sci.maxLod = 0;
|
||||
VkSampler nearest, linear; CK(vkCreateSampler(dev, &sci, nullptr, &nearest));
|
||||
sci.magFilter = sci.minFilter = VK_FILTER_LINEAR; CK(vkCreateSampler(dev, &sci, nullptr, &linear));
|
||||
std::vector<Buf> staging;
|
||||
for (Img& im : imgs)
|
||||
{
|
||||
const bool r16u = im.surf && im.f == VK_FORMAT_R16_SFLOAT; // byte-addressed surface: bound as R16_UINT
|
||||
VkImageCreateInfo ci{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; ci.imageType = VK_IMAGE_TYPE_2D; ci.format = im.f; ci.extent = {im.w, im.h, 1};
|
||||
ci.mipLevels = 1; ci.arrayLayers = 1; ci.samples = VK_SAMPLE_COUNT_1_BIT;
|
||||
ci.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
|
||||
if (r16u) ci.flags = VK_IMAGE_CREATE_MUTABLE_FORMAT_BIT;
|
||||
CK(vkCreateImage(dev, &ci, nullptr, &im.i));
|
||||
VkMemoryRequirements mr; vkGetImageMemoryRequirements(dev, im.i, &mr);
|
||||
CK(vkBindImageMemory(dev, im.i, alloc(mr, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, false), 0));
|
||||
VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; vi.image = im.i; vi.viewType = VK_IMAGE_VIEW_TYPE_2D;
|
||||
vi.format = r16u ? VK_FORMAT_R16_UINT : im.f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
|
||||
CK(vkCreateImageView(dev, &vi, nullptr, &im.v));
|
||||
Buf st = make_buf((VkDeviceSize)im.w * im.h * im.bpp, VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
|
||||
memset(st.map, 0, st.size);
|
||||
if (im.init != "-") { auto d = slurp(im.init); memcpy(st.map, d.data(), std::min<size_t>(d.size(), st.size)); }
|
||||
staging.push_back(st);
|
||||
}
|
||||
auto barrier = [&](VkImage i, VkImageLayout a, VkImageLayout b) {
|
||||
VkImageMemoryBarrier mb{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; mb.oldLayout = a; mb.newLayout = b; mb.image = i;
|
||||
mb.srcAccessMask = mb.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT;
|
||||
mb.srcQueueFamilyIndex = mb.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; mb.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1};
|
||||
vkCmdPipelineBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, 0, 0, nullptr, 0, nullptr, 1, &mb);
|
||||
};
|
||||
auto upload = [&]() {
|
||||
for (size_t i = 0; i < imgs.size(); ++i)
|
||||
{
|
||||
barrier(imgs[i].i, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);
|
||||
VkBufferImageCopy r{}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {imgs[i].w, imgs[i].h, 1};
|
||||
vkCmdCopyBufferToImage(cb, staging[i].b, imgs[i].i, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &r);
|
||||
barrier(imgs[i].i, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL);
|
||||
}
|
||||
};
|
||||
run(upload);
|
||||
// ---------------------------------------------------------------- pipeline
|
||||
auto code = slurp(argv[1]);
|
||||
VkShaderModuleCreateInfo smi{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO}; smi.codeSize = code.size(); smi.pCode = (const uint32_t*)code.data();
|
||||
VkShaderModule sm; CK(vkCreateShaderModule(dev, &smi, nullptr, &sm));
|
||||
std::vector<VkDescriptorSetLayoutBinding> lb;
|
||||
lb.push_back({0, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr});
|
||||
for (size_t i = 0; i < imgs.size(); ++i)
|
||||
lb.push_back({(uint32_t)i + 1, imgs[i].surf ? VK_DESCRIPTOR_TYPE_STORAGE_IMAGE : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr});
|
||||
for (size_t i = 0; i < luts.size(); ++i)
|
||||
lb.push_back({(uint32_t)(imgs.size() + 1 + i), VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1, VK_SHADER_STAGE_COMPUTE_BIT, nullptr});
|
||||
VkDescriptorSetLayoutCreateInfo dli{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO}; dli.bindingCount = lb.size(); dli.pBindings = lb.data();
|
||||
VkDescriptorSetLayout dl; CK(vkCreateDescriptorSetLayout(dev, &dli, nullptr, &dl));
|
||||
VkPipelineLayoutCreateInfo pli{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO}; pli.setLayoutCount = 1; pli.pSetLayouts = &dl;
|
||||
VkPipelineLayout pl; CK(vkCreatePipelineLayout(dev, &pli, nullptr, &pl));
|
||||
VkPipelineShaderStageRequiredSubgroupSizeCreateInfo rs{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; rs.requiredSubgroupSize = subgroup;
|
||||
VkComputePipelineCreateInfo cpi{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO};
|
||||
cpi.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, &rs, VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT, VK_SHADER_STAGE_COMPUTE_BIT, sm, "main", nullptr};
|
||||
cpi.layout = pl;
|
||||
VkPipeline pipe;
|
||||
auto c0 = std::chrono::steady_clock::now();
|
||||
CK(vkCreateComputePipelines(dev, VK_NULL_HANDLE, 1, &cpi, nullptr, &pipe));
|
||||
fprintf(stderr, "pipeline compile: %.0f ms\n", std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - c0).count());
|
||||
VkDescriptorPoolSize dps[3] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 8}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 32}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 32}};
|
||||
VkDescriptorPoolCreateInfo dpi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; dpi.maxSets = 1; dpi.poolSizeCount = 3; dpi.pPoolSizes = dps;
|
||||
VkDescriptorPool dp; CK(vkCreateDescriptorPool(dev, &dpi, nullptr, &dp));
|
||||
VkDescriptorSetAllocateInfo dai{VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO}; dai.descriptorPool = dp; dai.descriptorSetCount = 1; dai.pSetLayouts = &dl;
|
||||
VkDescriptorSet ds; CK(vkAllocateDescriptorSets(dev, &dai, &ds));
|
||||
VkDescriptorBufferInfo pbi{pbuf.b, 0, VK_WHOLE_SIZE};
|
||||
std::vector<VkDescriptorImageInfo> ii(imgs.size()); std::vector<VkWriteDescriptorSet> wr;
|
||||
wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, 0, 0, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, nullptr, &pbi, nullptr});
|
||||
for (size_t i = 0; i < imgs.size(); ++i)
|
||||
{
|
||||
ii[i] = {imgs[i].linear ? linear : nearest, imgs[i].v, VK_IMAGE_LAYOUT_GENERAL};
|
||||
wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, (uint32_t)i + 1, 0, 1, imgs[i].surf ? VK_DESCRIPTOR_TYPE_STORAGE_IMAGE : VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, &ii[i], nullptr, nullptr});
|
||||
}
|
||||
std::vector<VkDescriptorBufferInfo> li(luts.size());
|
||||
for (size_t i = 0; i < luts.size(); ++i)
|
||||
{
|
||||
li[i] = {luts[i].b, 0, VK_WHOLE_SIZE};
|
||||
wr.push_back({VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, ds, (uint32_t)(imgs.size() + 1 + i), 0, 1, VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, nullptr, &li[i], nullptr});
|
||||
}
|
||||
vkUpdateDescriptorSets(dev, wr.size(), wr.data(), 0, nullptr);
|
||||
auto dispatch = [&](int n) {
|
||||
vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe);
|
||||
vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pl, 0, 1, &ds, 0, nullptr);
|
||||
for (int r = 0; r < n; ++r) vkCmdDispatch(cb, gx, gy, 1);
|
||||
};
|
||||
run([&] { dispatch(1); });
|
||||
// ---------------------------------------------------------------- read back
|
||||
run([&] {
|
||||
for (size_t i = 0; i < imgs.size(); ++i) if (imgs[i].surf)
|
||||
{
|
||||
VkBufferImageCopy r{}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {imgs[i].w, imgs[i].h, 1};
|
||||
vkCmdCopyImageToBuffer(cb, imgs[i].i, VK_IMAGE_LAYOUT_GENERAL, staging[i].b, 1, &r);
|
||||
}
|
||||
});
|
||||
for (size_t i = 0; i < imgs.size(); ++i) if (imgs[i].surf)
|
||||
std::ofstream(imgs[i].out, std::ios::binary).write((const char*)staging[i].map, staging[i].size);
|
||||
for (auto& [off, f] : bufOut)
|
||||
std::ofstream(f, std::ios::binary).write((const char*)bufs[off].first.map, bufs[off].first.size);
|
||||
for (size_t i = 0; i < luts.size(); ++i) if (!lutOut[i].empty())
|
||||
std::ofstream(lutOut[i], std::ios::binary).write((const char*)luts[i].map, luts[i].size);
|
||||
if (const char* gap = getenv("KVK_GAP_US"))
|
||||
{
|
||||
// one dispatch per submit with an idle gap, like one kernel inside a frame
|
||||
std::vector<double> t;
|
||||
for (int it = 0; it < 600; ++it)
|
||||
{
|
||||
VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); dispatch(1); CK(vkEndCommandBuffer(cb));
|
||||
VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb;
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q));
|
||||
if (it >= 100) t.push_back(std::chrono::duration<double, std::micro>(std::chrono::steady_clock::now() - t0).count());
|
||||
auto until = std::chrono::steady_clock::now() + std::chrono::microseconds(atoi(gap));
|
||||
while (std::chrono::steady_clock::now() < until) {}
|
||||
}
|
||||
std::sort(t.begin(), t.end());
|
||||
printf("%s: %.1f us per submitted dispatch with %s us gaps (median of %zu, min %.1f)\n", argv[1], t[t.size() / 2], gap, t.size(), t[0]);
|
||||
}
|
||||
else if (repeats > 1)
|
||||
{
|
||||
std::vector<double> t;
|
||||
for (int it = 0; it < 12; ++it)
|
||||
{
|
||||
VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; CK(vkBeginCommandBuffer(cb, &bi)); dispatch(repeats); CK(vkEndCommandBuffer(cb));
|
||||
VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb;
|
||||
auto t0 = std::chrono::steady_clock::now();
|
||||
CK(vkQueueSubmit(q, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(q));
|
||||
if (it >= 3) t.push_back(std::chrono::duration<double, std::micro>(std::chrono::steady_clock::now() - t0).count() / repeats);
|
||||
}
|
||||
std::sort(t.begin(), t.end());
|
||||
printf("%s: %.1f us per dispatch (median of %zu submits of %d, min %.1f)\n", argv[1], t[t.size() / 2], t.size(), repeats, t[0]);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,968 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Translate one NVIDIA PTX texture kernel into a Vulkan GLSL compute shader.
|
||||
|
||||
usage: ptx2glsl.py MODULE.ptx KERNEL [--tex OFF:linear|nearest ...] > KERNEL.comp
|
||||
|
||||
Every PTX instruction becomes the GLSL expression of the same operation that kernels/native/ptx2hip.py
|
||||
emits for HIP, so the shader computes what the translated kernel computes. Registers become locals and
|
||||
basic blocks become guarded sections of a program-counter loop (GLSL has no goto). The 8-byte handles in
|
||||
the parameter block become bindings: textures by their parameter offset (T<off>), surfaces (S<off>);
|
||||
pointers stay 64-bit addresses and are read through buffer references.
|
||||
"""
|
||||
import re
|
||||
import sys
|
||||
|
||||
REG_RE = re.compile(r'%([a-zA-Z]+)(\d+)$')
|
||||
GT = {'f': 'float', 'r': 'uint', 'rd': 'uint64_t', 'rs': 'uint', 'p': 'bool'}
|
||||
SREG = {'%tid.x': 'gl_LocalInvocationID.x', '%tid.y': 'gl_LocalInvocationID.y', '%tid.z': 'gl_LocalInvocationID.z',
|
||||
'%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z',
|
||||
'%ntid.x': 'gl_WorkGroupSize.x', '%ntid.y': 'gl_WorkGroupSize.y',
|
||||
'%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%laneid': 'gl_SubgroupInvocationID'}
|
||||
|
||||
|
||||
class Err(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def tokenize(body):
|
||||
out, i, n = [], 0, len(body)
|
||||
while i < n:
|
||||
c = body[i]
|
||||
if c.isspace():
|
||||
i += 1; continue
|
||||
if c == '{':
|
||||
out.append(('open', None)); i += 1; continue
|
||||
if c == '}':
|
||||
out.append(('close', None)); i += 1; continue
|
||||
m = re.match(r'(\$?[A-Za-z_][\w$]*):', body[i:])
|
||||
if m:
|
||||
out.append(('label', m.group(1))); i += m.end(); continue
|
||||
j = body.index(';', i)
|
||||
out.append(('stmt', ' '.join(body[i:j].split())))
|
||||
i = j + 1
|
||||
return out
|
||||
|
||||
|
||||
def split_operands(s):
|
||||
ops, depth, cur = [], 0, ''
|
||||
for ch in s:
|
||||
if ch in '{[':
|
||||
depth += 1
|
||||
elif ch in '}]':
|
||||
depth -= 1
|
||||
if ch == ',' and depth == 0:
|
||||
ops.append(cur.strip()); cur = ''
|
||||
else:
|
||||
cur += ch
|
||||
if cur.strip():
|
||||
ops.append(cur.strip())
|
||||
return ops
|
||||
|
||||
|
||||
class T:
|
||||
def __init__(self, ptx, kernel):
|
||||
text = re.sub(r'//[^\n]*', '', ptx)
|
||||
m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S)
|
||||
if not m:
|
||||
raise Err('kernel not found')
|
||||
pm = re.search(r'\.param\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', m.group(1))
|
||||
self.kernel, self.pname, self.psize = kernel, pm.group(2), int(pm.group(3))
|
||||
rest = text[m.end():]
|
||||
mt = re.search(r'\.maxntid\s+(\d+)\s*,\s*(\d+)\s*,\s*(\d+)', rest[:rest.index('{')])
|
||||
self.maxntid = [int(x) for x in mt.groups()] if mt else [256, 1, 1]
|
||||
start = rest.index('{'); depth, k = 0, start
|
||||
while True:
|
||||
if rest[k] == '{':
|
||||
depth += 1
|
||||
elif rest[k] == '}':
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
break
|
||||
k += 1
|
||||
self.toks = tokenize(rest[start + 1:k])
|
||||
self.regs = {}
|
||||
self.shared = {} # name -> byte base
|
||||
self.shared_bytes = 0
|
||||
self.scopes = [{}]
|
||||
self.scope_id = 0
|
||||
self.blocks = [[]] # lists of lines
|
||||
self.labels = {} # label -> block index
|
||||
self.textures, self.surfaces = set(), {}
|
||||
self.dyn_tex = False
|
||||
self.fast = False # --fast: no 'precise', the compiler may contract and reassociate
|
||||
self.collect = False # first pass: gather how every 16/32-bit register is used
|
||||
self.ctx, self.links, self.width, self.half = {}, [], {}, set()
|
||||
self.exp_lut = False
|
||||
self.prescan()
|
||||
|
||||
# ---------------------------------------------------------------- static facts about registers
|
||||
def prescan(self):
|
||||
defs = {}
|
||||
for kind, s in self.toks:
|
||||
if kind != 'stmt' or s.startswith('.'):
|
||||
continue
|
||||
s = re.sub(r'^@!?%p\d+\s+', '', s)
|
||||
parts = s.split(None, 1)
|
||||
if len(parts) < 2:
|
||||
continue
|
||||
a = split_operands(parts[1])
|
||||
for d in re.findall(r'%rd\d+', a[0]) if not a[0].startswith('[') else []:
|
||||
defs.setdefault(d, []).append((parts[0], a))
|
||||
self.pbase = {} # reg -> constant offset from the parameter block
|
||||
self.handle = {} # reg -> parameter offset its 8-byte value was loaded from
|
||||
changed = True
|
||||
while changed:
|
||||
changed = False
|
||||
for r, ds in defs.items():
|
||||
if len(ds) != 1:
|
||||
continue
|
||||
op, a = ds[0]
|
||||
if r not in self.pbase:
|
||||
v = None
|
||||
if op.startswith('mov') and a[1] == self.pname:
|
||||
v = 0
|
||||
elif op in ('add.s64', 'add.u64') and a[1] in self.pbase and re.match(r'-?\d+$', a[2]):
|
||||
v = self.pbase[a[1]] + int(a[2])
|
||||
if v is not None:
|
||||
self.pbase[r] = v; changed = True
|
||||
if r not in self.handle:
|
||||
v = None
|
||||
if op == 'ld.param.u64':
|
||||
v = self.param_off(a[1])
|
||||
elif op.startswith('mov') and a[1] in self.handle:
|
||||
v = self.handle[a[1]]
|
||||
if v is not None:
|
||||
self.handle[r] = v; changed = True
|
||||
|
||||
def param_off(self, tok):
|
||||
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
|
||||
base, off = m.group(1), int(m.group(2) or 0)
|
||||
if base == self.pname:
|
||||
return off
|
||||
if base in self.pbase:
|
||||
return self.pbase[base] + off
|
||||
return None
|
||||
|
||||
# ---------------------------------------------------------------- operands
|
||||
def reg(self, tok):
|
||||
tok = tok.strip()
|
||||
for sc in reversed(self.scopes):
|
||||
if tok in sc:
|
||||
return sc[tok]
|
||||
m = REG_RE.match(tok)
|
||||
if not m or m.group(1) not in GT:
|
||||
raise Err('not a register: ' + tok)
|
||||
name = m.group(1) + m.group(2)
|
||||
if m.group(1) in ('r', 'rs'):
|
||||
self.width[name] = 32 if m.group(1) == 'r' else 16
|
||||
self.regs[name] = self.htype(name) or GT[m.group(1)]
|
||||
return name
|
||||
|
||||
def htype(self, name):
|
||||
return ('f16vec2' if self.width[name] == 32 else 'float16_t') if name in self.half else None
|
||||
|
||||
def note(self, name, cls):
|
||||
if self.collect:
|
||||
self.ctx.setdefault(name, set()).add(cls)
|
||||
|
||||
def hk(self, name):
|
||||
return 'f16x2' if self.width.get(name) == 32 else 'f16'
|
||||
|
||||
def is_reg(self, tok):
|
||||
return tok.startswith('%') and tok not in SREG or any(tok in sc for sc in self.scopes)
|
||||
|
||||
def val(self, tok, t):
|
||||
tok = tok.strip()
|
||||
if tok in SREG:
|
||||
return f'uint({SREG[tok]})' if t != 's32' else f'int({SREG[tok]})'
|
||||
if tok in ('%SP', '%SPL'):
|
||||
return 'uint64_t(0ul)' # the local depot is its own array; addresses in it start at 0
|
||||
if self.is_reg(tok):
|
||||
n = self.reg(tok); g = self.regs[n]
|
||||
self.note(n, 'h' if t in ('f16', 'f16x2') else 'u')
|
||||
if g in ('float16_t', 'f16vec2'):
|
||||
if t in ('f16', 'f16x2'):
|
||||
return n
|
||||
bits = f'HU({n})' if g == 'float16_t' else f'H2U({n})'
|
||||
if t == 's16':
|
||||
return f's16({bits})'
|
||||
if t == 's32':
|
||||
return f'int({bits})'
|
||||
if t in ('u64', 'b64'):
|
||||
return f'uint64_t({bits})'
|
||||
return bits
|
||||
if t == 'f32':
|
||||
return n if g == 'float' else f'uintBitsToFloat({n})'
|
||||
if t in ('u32', 'b32'):
|
||||
return f'floatBitsToUint({n})' if g == 'float' else (f'uint({n})' if g == 'uint64_t' else n)
|
||||
if t == 's32':
|
||||
return f'floatBitsToInt({n})' if g == 'float' else f'int({n})'
|
||||
if t in ('u64', 'b64'):
|
||||
return n if g == 'uint64_t' else f'uint64_t({n})'
|
||||
if t == 's64':
|
||||
return f'int64_t({n})'
|
||||
if t in ('u16', 'b16'):
|
||||
return n
|
||||
if t == 's16':
|
||||
return f's16({n})'
|
||||
if t == 'f16':
|
||||
return f'H({n})'
|
||||
if t == 'f16x2':
|
||||
return f'H2({n})'
|
||||
if t == 'pred':
|
||||
return n
|
||||
raise Err('type ' + t)
|
||||
if tok in self.shared:
|
||||
return f'{self.shared[tok]}u'
|
||||
if re.match(r'0[fF][0-9A-Fa-f]{8}$', tok):
|
||||
bits = int(tok[2:], 16)
|
||||
return f'uintBitsToFloat(0x{bits:08X}u)' if t == 'f32' else f'0x{bits:08X}u'
|
||||
if re.match(r'-?\d+$', tok) or re.match(r'0x[0-9a-fA-F]+$', tok):
|
||||
v = int(tok, 0)
|
||||
if t == 'f32':
|
||||
raise Err('int immediate as float')
|
||||
if t in ('u64', 'b64'):
|
||||
return f'uint64_t({v & 0xffffffffffffffff}ul)'
|
||||
if t == 's64':
|
||||
return f'int64_t({v}l)'
|
||||
if t in ('u16', 'b16'):
|
||||
return f'{v & 0xffff}u'
|
||||
if t == 's16':
|
||||
v &= 0xffff
|
||||
return f'({v - 0x10000 if v >= 0x8000 else v})'
|
||||
if t == 'f16':
|
||||
return f'H({v & 0xffff}u)'
|
||||
if t == 'f16x2':
|
||||
return f'H2({v & 0xffffffff}u)'
|
||||
if t == 's32':
|
||||
v &= 0xffffffff
|
||||
return f'int(0x{v:08X}u)' if v >= 1 << 31 else f'{v}'
|
||||
return f'{v & 0xffffffff}u'
|
||||
raise Err('operand ' + tok)
|
||||
|
||||
def A(self, dst, expr, st):
|
||||
n = self.reg(dst); g = self.regs[n]
|
||||
self.note(n, 'h' if st in ('f16', 'f16x2') else 'u')
|
||||
if g == 'bool' and st in ('u32', 'u16', 'u64'):
|
||||
self.E(f'{n} = ({expr}) != 0u;'); return # mov.pred with an immediate
|
||||
if g in ('float16_t', 'f16vec2'):
|
||||
if st in ('f16', 'f16x2'):
|
||||
self.E(f'{n} = {expr};')
|
||||
else:
|
||||
self.E(f'{n} = {"H" if g == "float16_t" else "H2"}(uint({expr}));')
|
||||
return
|
||||
if st == 'float':
|
||||
e = expr if g == 'float' else f'floatBitsToUint({expr})'
|
||||
elif st == 'u32':
|
||||
e = f'uintBitsToFloat({expr})' if g == 'float' else (f'uint64_t({expr})' if g == 'uint64_t' else f'uint({expr})')
|
||||
elif st == 'u64':
|
||||
e = f'uint64_t({expr})'
|
||||
elif st == 'u16':
|
||||
e = f'(uint({expr}) & 0xffffu)'
|
||||
elif st == 'f16':
|
||||
e = f'HU({expr})'
|
||||
elif st == 'f16x2':
|
||||
e = f'H2U({expr})'
|
||||
elif st == 'bool':
|
||||
e = f'({expr})'
|
||||
else:
|
||||
raise Err('src type')
|
||||
self.E(f'{n} = {e};')
|
||||
|
||||
def E(self, s):
|
||||
self.blocks[-1].append(s)
|
||||
|
||||
# ---------------------------------------------------------------- driver
|
||||
def translate(self):
|
||||
for kind, text in self.toks:
|
||||
if kind == 'open':
|
||||
self.scopes.append({}); continue
|
||||
if kind == 'close':
|
||||
self.scopes.pop(); continue
|
||||
if kind == 'label':
|
||||
self.blocks.append([])
|
||||
self.labels[text] = len(self.blocks) - 1
|
||||
continue
|
||||
self.stmt(text)
|
||||
|
||||
def stmt(self, s):
|
||||
if s.startswith('.reg'):
|
||||
m = re.match(r'\.reg\s+\.(\w+)\s+(.*)$', s)
|
||||
ty, names = m.groups()
|
||||
if '<' in names:
|
||||
return
|
||||
self.scope_id += 1
|
||||
for nm in names.split(','):
|
||||
nm = nm.strip()
|
||||
c = f'l{self.scope_id}_' + re.sub(r'\W', '_', nm)
|
||||
self.scopes[-1][nm] = c
|
||||
if ty in ('f16', 'b16', 'f16x2', 'b32', 'u32', 's32'):
|
||||
self.width[c] = 16 if ty in ('f16', 'b16') else 32
|
||||
self.regs[c] = self.htype(c) if c in self.half else {'f16': 'uint', 'b16': 'uint', 'f32': 'float', 'b32': 'uint', 'pred': 'bool', 'f16x2': 'uint',
|
||||
'u32': 'uint', 's32': 'uint', 'b64': 'uint64_t', 'u64': 'uint64_t'}[ty]
|
||||
return
|
||||
if s.startswith('.local'):
|
||||
m = re.match(r'\.local\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
||||
self.local_bytes = max(getattr(self, 'local_bytes', 0), int(m.group(3)))
|
||||
return
|
||||
if re.match(r'mov\.u64 %SPL?, __local_depot', s):
|
||||
return
|
||||
if s.startswith('.shared'):
|
||||
m = re.match(r'\.shared\s+\.align\s+(\d+)\s+\.b8\s+(\w+)\[(\d+)\]', s)
|
||||
self.shared[m.group(2)] = self.shared_bytes
|
||||
self.shared_bytes += (int(m.group(3)) + 3) & ~3
|
||||
return
|
||||
if s.startswith('.'):
|
||||
return
|
||||
pred = None
|
||||
m = re.match(r'@(!?)(%p\d+)\s+(.*)$', s)
|
||||
if m:
|
||||
pred = ('!' if m.group(1) else '') + self.reg(m.group(2)); s = m.group(3)
|
||||
parts = s.split(None, 1)
|
||||
op = parts[0]
|
||||
a = split_operands(parts[1]) if len(parts) > 1 else []
|
||||
if op.startswith('bra'):
|
||||
tgt = ('L', a[0])
|
||||
if pred is not None:
|
||||
self.E(('BR', pred, tgt)) # conditional: falls through into a new block
|
||||
else:
|
||||
self.E(('BR', None, tgt))
|
||||
self.blocks.append([])
|
||||
return
|
||||
if op == 'ret':
|
||||
self.E(('BR', pred, ('END', None)))
|
||||
self.blocks.append([])
|
||||
return
|
||||
if pred is not None:
|
||||
self.E(f'if ({pred}) {{')
|
||||
self.op(op, a)
|
||||
self.E('}')
|
||||
return
|
||||
self.op(op, a)
|
||||
|
||||
def cmp(self, mods, isf):
|
||||
table = {'lt': '{0} < {1}', 'gt': '{0} > {1}', 'le': '{0} <= {1}', 'ge': '{0} >= {1}', 'eq': '{0} == {1}', 'ne': '{0} != {1}',
|
||||
'ltu': '!({0} >= {1})', 'gtu': '!({0} <= {1})', 'leu': '!({0} > {1})', 'geu': '!({0} < {1})',
|
||||
'equ': '!({0} != {1})', 'neu': '!({0} == {1})'}
|
||||
for k in table:
|
||||
if k in mods:
|
||||
if isf and k == 'ne':
|
||||
return '({0} < {1} || {0} > {1})'
|
||||
if isf and k == 'equ':
|
||||
return '!({0} < {1} || {0} > {1})'
|
||||
return table[k]
|
||||
raise Err('compare')
|
||||
|
||||
def texref(self, tok):
|
||||
m = re.match(r'\[\s*(%rd\d+)\s*,\s*\{([^}]*)\}\s*\]', tok)
|
||||
return m.group(1), [x.strip() for x in m.group(2).split(',')]
|
||||
|
||||
def op(self, op, a):
|
||||
o = op.split('.'); base = o[0]; mods = set(o[1:])
|
||||
E, A, V = self.E, self.A, self.val
|
||||
ft = next((t for t in ('f32', 'f16x2', 'f16') if t in mods), None)
|
||||
if base == 'bar':
|
||||
E('memoryBarrierShared(); barrier();'); return
|
||||
if base == 'mov':
|
||||
dst, src = a
|
||||
if dst.startswith('{'):
|
||||
n = [x.strip() for x in dst[1:-1].split(',')]
|
||||
sr = self.reg(src) if self.is_reg(src) else None
|
||||
if self.collect and sr:
|
||||
for x in [sr] + [self.reg(y) for y in n]:
|
||||
self.ctx.setdefault(x, set()).add('h')
|
||||
return
|
||||
if sr and (sr in self.half or any(self.reg(y) in self.half for y in n)):
|
||||
A(n[0], f'{V(src, "f16x2")}.x', 'f16'); A(n[1], f'{V(src, "f16x2")}.y', 'f16'); return
|
||||
A(n[0], f'{V(src, "u32")} & 0xffffu', 'u16'); A(n[1], f'{V(src, "u32")} >> 16', 'u16'); return
|
||||
if src.startswith('{'):
|
||||
n = [x.strip() for x in src[1:-1].split(',')]
|
||||
dr = self.reg(dst)
|
||||
if self.collect:
|
||||
for x in [dr] + [self.reg(y) for y in n if self.is_reg(y)]:
|
||||
self.ctx.setdefault(x, set()).add('h')
|
||||
return
|
||||
if dr in self.half:
|
||||
A(dst, f'f16vec2({V(n[0], "f16")}, {V(n[1], "f16")})', 'f16x2'); return
|
||||
A(dst, f'({V(n[0], "u16")} & 0xffffu) | ({V(n[1], "u16")} << 16)', 'u32'); return
|
||||
if self.is_reg(dst) and self.reg(dst) in self.width and not (mods & {'f32', 'u64', 'b64', 's64'}):
|
||||
dr = self.reg(dst)
|
||||
srcreg = self.reg(src) if self.is_reg(src) and src not in SREG else None
|
||||
isimm = bool(re.match(r'-?\d+$', src) or re.match(r'0x[0-9a-fA-F]+$', src))
|
||||
if self.collect and srcreg and srcreg in self.width:
|
||||
self.links.append((dr, srcreg)); return
|
||||
if self.collect and isimm:
|
||||
return
|
||||
if dr in self.half and (srcreg or isimm):
|
||||
A(dst, V(src, self.hk(dr)), self.hk(dr)); return
|
||||
if src == self.pname:
|
||||
A(dst, '0ul', 'u64'); return
|
||||
if 'f32' in mods:
|
||||
A(dst, V(src, 'f32'), 'float'); return
|
||||
if mods & {'u64', 'b64', 's64'}:
|
||||
A(dst, V(src, 'u64'), 'u64'); return
|
||||
if mods & {'u16', 'b16', 's16'}:
|
||||
A(dst, V(src, 'u16'), 'u16'); return
|
||||
A(dst, V(src, 'u32'), 'u32'); return
|
||||
# ------------------------------------------------ f32
|
||||
if ft == 'f32' and base in ('add', 'sub', 'mul'):
|
||||
if 'rz' in mods:
|
||||
A(a[0], f'round_half_away({V(a[1], "f32")})', 'float'); return # the roundf idiom: + copysign(0.5), then trunc
|
||||
A(a[0], f'{V(a[1], "f32")} {"+-*"[["add", "sub", "mul"].index(base)]} {V(a[2], "f32")}', 'float'); return
|
||||
if ft == 'f32' and base == 'fma':
|
||||
A(a[0], f'fma({V(a[1], "f32")}, {V(a[2], "f32")}, {V(a[3], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base in ('max', 'min'):
|
||||
A(a[0], f'{base}({V(a[1], "f32")}, {V(a[2], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base == 'neg':
|
||||
A(a[0], f'-{V(a[1], "f32")}', 'float'); return
|
||||
if ft == 'f32' and base == 'abs':
|
||||
A(a[0], f'abs({V(a[1], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base in ('rcp', 'ex2', 'lg2', 'rsqrt', 'sqrt'):
|
||||
fn = {'rcp': 'rcp_a', 'ex2': 'exp2', 'lg2': 'log2', 'rsqrt': 'inversesqrt', 'sqrt': 'sqrt'}[base]
|
||||
A(a[0], f'{fn}({V(a[1], "f32")})', 'float'); return
|
||||
if ft == 'f32' and base == 'div':
|
||||
A(a[0], f'{V(a[1], "f32")} * rcp_a({V(a[2], "f32")})', 'float'); return
|
||||
# ------------------------------------------------ f16
|
||||
if ft == 'f16x2' and base in ('add', 'sub', 'mul', 'fma'):
|
||||
x = [V(t, 'f16x2') for t in a[1:]]
|
||||
e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
|
||||
if 'sat' in mods:
|
||||
e = f'min(max({e}, f16vec2(0.0hf)), f16vec2(1.0hf))'
|
||||
A(a[0], e, 'f16x2'); return
|
||||
if ft == 'f16x2' and base == 'abs':
|
||||
A(a[0], f'abs({V(a[1], "f16x2")})', 'f16x2'); return
|
||||
if ft == 'f16x2' and base == 'ex2':
|
||||
if self.exp_lut:
|
||||
A(a[0], f'ex2_lut({V(a[1], "u32")})', 'u32'); return
|
||||
A(a[0], f'f16vec2(exp2(vec2({V(a[1], "f16x2")})))', 'f16x2'); return
|
||||
if ft in ('f16', 'f16x2') and base in ('min', 'max'):
|
||||
A(a[0], f'{base}({V(a[1], ft)}, {V(a[2], ft)})', ft); return
|
||||
if ft == 'f16' and base in ('add', 'sub', 'mul', 'fma'):
|
||||
x = [V(t, 'f16') for t in a[1:]]
|
||||
e = f'fma({x[0]}, {x[1]}, {x[2]})' if base == 'fma' else f'({x[0]} {"+-*"[["add", "sub", "mul"].index(base)]} {x[1]})'
|
||||
if 'sat' in mods:
|
||||
e = f'min(max({e}, 0.0hf), 1.0hf)'
|
||||
A(a[0], e, 'f16'); return
|
||||
if base == 'set' and 'f16x2' in mods and 'u32' in mods:
|
||||
c = self.cmp(mods, True); x, y = V(a[1], 'f16x2'), V(a[2], 'f16x2')
|
||||
A(a[0], f'(({c.format(x + ".x", y + ".x")}) ? 0xffffu : 0u) | (({c.format(x + ".y", y + ".y")}) ? 0xffff0000u : 0u)', 'u32'); return
|
||||
# ------------------------------------------------ compares / selects
|
||||
if base == 'setp':
|
||||
isf = bool(mods & {'f32', 'f16'})
|
||||
t = next(t for t in ('f32', 'f16', 's32', 'u32', 'b32', 's16', 'u16', 'b16', 's64', 'u64') if t in mods)
|
||||
if '|' in a[0]:
|
||||
raise Err('setp form')
|
||||
A(a[0], self.cmp(mods, isf).format(V(a[1], t), V(a[2], t)), 'bool'); return
|
||||
if base == 'selp' and not (mods & {'f32', 'u64', 'b64', 's64'}) and self.reg(a[0]) in self.width:
|
||||
dr = self.reg(a[0])
|
||||
srcs = [self.reg(x) for x in a[1:3] if self.is_reg(x) and x not in SREG]
|
||||
if self.collect:
|
||||
self.links += [(dr, x) for x in srcs if x in self.width]
|
||||
V(a[3], 'pred')
|
||||
if any(x not in self.width for x in srcs):
|
||||
self.note(dr, 'u')
|
||||
return
|
||||
if dr in self.half:
|
||||
k = self.hk(dr)
|
||||
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], k)} : {V(a[2], k)}', k); return
|
||||
if base == 'selp':
|
||||
t = 'f32' if 'f32' in mods else 'u16' if mods & {'u16', 'b16', 's16'} else 'u64' if mods & {'u64', 'b64', 's64'} else 'u32'
|
||||
A(a[0], f'{V(a[3], "pred")} ? {V(a[1], t)} : {V(a[2], t)}', {'f32': 'float', 'u16': 'u16', 'u64': 'u64', 'u32': 'u32'}[t]); return
|
||||
if 'pred' in mods and base in ('and', 'or', 'xor', 'not'):
|
||||
if base == 'not':
|
||||
A(a[0], f'!{V(a[1], "pred")}', 'bool'); return
|
||||
A(a[0], f'{V(a[1], "pred")} {{"and": "&&", "or": "||", "xor": "!="}}'.replace('{"and": "&&", "or": "||", "xor": "!="}', {'and': '&&', 'or': '||', 'xor': '!='}[base]) + f' {V(a[2], "pred")}', 'bool'); return
|
||||
# ------------------------------------------------ integer
|
||||
if base in ('and', 'or', 'xor', 'not', 'shl', 'shr', 'add', 'sub', 'mul', 'mad', 'min', 'max', 'div', 'rem', 'neg', 'abs') and ft is None:
|
||||
w = 64 if mods & {'u64', 's64', 'b64'} else 16 if mods & {'u16', 's16', 'b16'} else 32
|
||||
signed = bool(mods & {'s32', 's64', 's16'})
|
||||
ut, st = {64: ('u64', 's64'), 32: ('u32', 's32'), 16: ('u16', 's16')}[w]
|
||||
if base == 'mul' and 'wide' in mods:
|
||||
if 's32' in mods:
|
||||
A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")}))', 'u64'); return
|
||||
if 'u32' in mods:
|
||||
A(a[0], f'uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})', 'u64'); return
|
||||
if 'u16' in mods:
|
||||
A(a[0], f'{V(a[1], "u16")} * {V(a[2], "u16")}', 'u32'); return
|
||||
raise Err(op)
|
||||
if base == 'not':
|
||||
A(a[0], f'~{V(a[1], ut)}', ut); return
|
||||
if base == 'neg':
|
||||
A(a[0], f'0u - {V(a[1], ut)}', ut); return
|
||||
if base == 'shl':
|
||||
A(a[0], f'{V(a[1], ut)} << ({V(a[2], "u32")} & {w - 1}u)', ut); return
|
||||
if base == 'shr':
|
||||
x = V(a[1], st if signed else ut)
|
||||
A(a[0], f'{x} >> ({V(a[2], "u32")} & {w - 1}u)', ut); return
|
||||
if base in ('and', 'or', 'xor'):
|
||||
A(a[0], f'{V(a[1], ut)} {{"and": "&", "or": "|", "xor": "^"}}'.replace('{"and": "&", "or": "|", "xor": "^"}', {'and': '&', 'or': '|', 'xor': '^'}[base]) + f' {V(a[2], ut)}', ut); return
|
||||
if base in ('add', 'sub'):
|
||||
A(a[0], f'{V(a[1], ut)} {"+" if base == "add" else "-"} {V(a[2], ut)}', ut); return
|
||||
if base == 'mul' and 'hi' in mods:
|
||||
if signed:
|
||||
A(a[0], f'uint(int((int64_t({V(a[1], "s32")}) * int64_t({V(a[2], "s32")})) >> 32))', 'u32'); return
|
||||
A(a[0], f'uint((uint64_t({V(a[1], "u32")}) * uint64_t({V(a[2], "u32")})) >> 32)', 'u32'); return
|
||||
if base == 'mul':
|
||||
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)}', ut); return
|
||||
if base == 'mad':
|
||||
A(a[0], f'{V(a[1], ut)} * {V(a[2], ut)} + {V(a[3], ut)}', ut); return
|
||||
if base == 'abs':
|
||||
A(a[0], f'uint(abs({V(a[1], st)}))', ut); return
|
||||
if base == 'neg' and len(a) == 2:
|
||||
A(a[0], f'uint(-{V(a[1], st)})', ut); return
|
||||
x, y = V(a[1], st if signed else ut), V(a[2], st if signed else ut)
|
||||
if base in ('min', 'max'):
|
||||
A(a[0], f'{base}({x}, {y})', ut); return
|
||||
if base == 'div':
|
||||
A(a[0], f'{x} / {y}', ut); return
|
||||
if base == 'rem':
|
||||
A(a[0], f'{x} % {y}', ut); return
|
||||
if base == 'bfi':
|
||||
A(a[0], f'bfi32({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")}, {V(a[4], "u32")})', 'u32'); return
|
||||
if base == 'cvt':
|
||||
return self.cvt(op, o[1:], a)
|
||||
if base == 'cvta':
|
||||
A(a[0], V(a[1], 'u64'), 'u64'); return
|
||||
if base == 'ld':
|
||||
return self.ld(o[1:], a)
|
||||
if base == 'st':
|
||||
return self.st(o[1:], a)
|
||||
if base in ('tex', 'tld4'):
|
||||
dst = [x.strip() for x in a[0][1:-1].split(',')]
|
||||
h, c = self.texref(a[1])
|
||||
off = self.handle.get(h)
|
||||
if off is None:
|
||||
raise Err('texture handle of ' + h + ' is not a single parameter load')
|
||||
self.textures.add(off)
|
||||
if base == 'tld4':
|
||||
E(f'{{ vec4 t_ = textureGather(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), 0);')
|
||||
elif 'base' in mods:
|
||||
E(f'{{ vec4 t_ = texelFetch(T{off}, ivec2({V(c[0], "s32")}, {V(c[1], "s32")}), 0);')
|
||||
elif 'level' in mods:
|
||||
E(f'{{ vec4 t_ = textureLod(T{off}, vec2({V(c[0], "f32")}, {V(c[1], "f32")}), {V(a[2], "f32")});')
|
||||
else:
|
||||
raise Err(op)
|
||||
for i, d in enumerate(dst):
|
||||
if 'f16' in o:
|
||||
A(d, f'float16_t(t_[{i}])', 'f16') # the texel as a half float
|
||||
else:
|
||||
A(d, f't_[{i}]', 'float')
|
||||
E('}')
|
||||
return
|
||||
if base == 'sust':
|
||||
h, c = self.texref(a[0])
|
||||
off = self.handle.get(h)
|
||||
if off is None:
|
||||
raise Err('surface handle of ' + h)
|
||||
v = [x.strip() for x in a[1][1:-1].split(',')]
|
||||
if op == 'sust.p.2d.v4.b32.zero':
|
||||
self.surfaces.setdefault(off, 'p'); assert self.surfaces[off] == 'p'
|
||||
E(f'sust_p(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, vec4({", ".join(V(x, "f32") for x in v)}));'); return
|
||||
if op == 'sust.b.2d.b32.zero':
|
||||
self.surfaces.setdefault(off, 'b'); assert self.surfaces[off] == 'b'
|
||||
E(f'sust_b(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u32")});'); return
|
||||
if op == 'sust.b.2d.v2.b16.zero': # reading only: two 16-bit values at a byte-addressed position
|
||||
self.surfaces.setdefault(off, 'b')
|
||||
E(f'sust_b2x16(S{off}, {V(c[0], "s32")}, {V(c[1], "s32")}, {V(v[0], "u16")}, {V(v[1], "u16")});'); return
|
||||
raise Err(op)
|
||||
if op == 'shfl.sync.bfly.b32':
|
||||
d = a[0].split('|')
|
||||
E('{ bool ib_;')
|
||||
A(d[0], f'shfl_bfly({V(a[1], "u32")}, {V(a[2], "s32")}, {V(a[3], "u32")}, ib_)', 'u32')
|
||||
if len(d) > 1:
|
||||
A(d[1], 'ib_', 'bool')
|
||||
E('}')
|
||||
return
|
||||
if op == 'shfl.sync.idx.b32':
|
||||
d = a[0].split('|')
|
||||
A(d[0], f'subgroupShuffle({V(a[1], "u32")}, {V(a[2], "u32")} & 31u)', 'u32')
|
||||
if len(d) > 1:
|
||||
A(d[1], 'true', 'bool')
|
||||
return
|
||||
if op == 'tanh.approx.f16':
|
||||
A(a[0], f'tanh({V(a[1], "f16")})', 'f16'); return
|
||||
if op == 'prmt.b32':
|
||||
A(a[0], f'prmt({V(a[1], "u32")}, {V(a[2], "u32")}, {V(a[3], "u32")})', 'u32'); return
|
||||
if op.startswith('mma.sync'):
|
||||
# not executable here: the kernel is translated for reading only (tensor-core fragments are not emulated)
|
||||
self.E(f'// {op} ' + ' '.join(a))
|
||||
for d in a[0][1:-1].split(','):
|
||||
A(d.strip(), '0u', 'u32')
|
||||
return
|
||||
if op == 'vote.sync.ballot.b32':
|
||||
# the warp is the subgroup (32 lanes in thread order), as for the shuffles
|
||||
A(a[0], f'subgroupBallot({V(a[1], "pred")}).x', 'u32'); return
|
||||
raise Err('unhandled ' + op)
|
||||
|
||||
def cvt(self, op, o, a):
|
||||
A, V = self.A, self.val
|
||||
mods = set(o)
|
||||
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)]
|
||||
types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$|^e4m3x2$', t)]
|
||||
dt, st = types[0], types[1]
|
||||
# FP8 pairs (reading only: the helpers are not defined)
|
||||
if dt == 'f16x2' and st == 'e4m3x2':
|
||||
A(a[0], f'e4m3x2_to_f16x2({V(a[1], "u16")})', 'f16x2'); return
|
||||
if dt == 'e4m3x2' and st == 'f16x2':
|
||||
A(a[0], f'f16x2_to_e4m3x2({V(a[1], "f16x2")})', 'u16'); return
|
||||
if dt == 'f32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
if 'sat' in mods:
|
||||
A(a[0], f'min(max({x}, 0.0), 1.0)', 'float'); return
|
||||
fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven' if 'rni' in mods else 'ceil' if 'rpi' in mods else ''
|
||||
A(a[0], f'{fn}({x})', 'float'); return
|
||||
if dt == 's32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
fn = '' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven'
|
||||
A(a[0], f'uint(int({fn}({x})))', 'u32'); return
|
||||
if dt == 'u32' and st == 'f32':
|
||||
x = V(a[1], 'f32')
|
||||
fn = 'trunc' if 'rzi' in mods else 'floor' if 'rmi' in mods else 'roundEven'
|
||||
A(a[0], f'uint(clamp({fn}({x}), 0.0, 4294967040.0))', 'u32'); return # CUDA saturates
|
||||
if dt == 'f32' and st == 's32':
|
||||
A(a[0], f'float({V(a[1], "s32")})', 'float'); return
|
||||
if dt == 'f32' and st == 'u32':
|
||||
A(a[0], f'float({V(a[1], "u32")})', 'float'); return
|
||||
if dt == 'f32' and st == 'f16':
|
||||
A(a[0], f'float({V(a[1], "f16")})', 'float'); return
|
||||
if dt == 'f16' and st == 'f32':
|
||||
A(a[0], f'float16_t({V(a[1], "f32")})', 'f16'); return
|
||||
if dt == 'f16x2' and st == 'f32':
|
||||
A(a[0], f'f16vec2(float16_t({V(a[2], "f32")}), float16_t({V(a[1], "f32")}))', 'f16x2'); return
|
||||
if dt == 'f16' and st == 'f16':
|
||||
fn = 'roundEven' if 'rni' in mods else 'floor'
|
||||
A(a[0], f'{fn}({V(a[1], "f16")})', 'f16'); return
|
||||
if dt == 's32' and st == 'f16':
|
||||
fn = 'roundEven' if 'rni' in mods else 'floor'
|
||||
A(a[0], f'uint(int(float({fn}({V(a[1], "f16")}))))', 'u32'); return
|
||||
if dt == 'f16' and st == 's32':
|
||||
A(a[0], f'float16_t({V(a[1], "s32")})', 'f16'); return
|
||||
if dt == 's64' and st == 's32':
|
||||
A(a[0], f'uint64_t(int64_t({V(a[1], "s32")}))', 'u64'); return
|
||||
if dt == 'u64' and st == 'u32':
|
||||
A(a[0], f'uint64_t({V(a[1], "u32")})', 'u64'); return
|
||||
if dt in ('u16', 's16') and st in ('u32', 's32'):
|
||||
A(a[0], V(a[1], 'u32'), 'u16'); return
|
||||
if dt in ('u32', 's32') and st in ('u16',):
|
||||
A(a[0], V(a[1], 'u16'), 'u32'); return
|
||||
if dt == 's32' and st == 's16':
|
||||
A(a[0], f'uint({V(a[1], "s16")})', 'u32'); return
|
||||
raise Err('cvt ' + op)
|
||||
|
||||
def addr(self, tok):
|
||||
m = re.match(r'\[\s*([^\]+]+?)\s*(?:\+\s*(-?\d+))?\s*\]$', tok)
|
||||
return m.group(1), int(m.group(2) or 0)
|
||||
|
||||
def ld(self, o, a):
|
||||
A, V, E = self.A, self.val, self.E
|
||||
mods = set(o)
|
||||
ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t))
|
||||
vec = next((int(t[1]) for t in o if re.match(r'^v\d$', t)), 1)
|
||||
size = int(ty[1:]) // 8
|
||||
dsts = [x.strip() for x in a[0][1:-1].split(',')] if a[0].startswith('{') else [a[0]]
|
||||
base, off = self.addr(a[1])
|
||||
for i, d in enumerate(dsts):
|
||||
if 'param' in mods:
|
||||
po = self.param_off(a[1])
|
||||
if po is None:
|
||||
raise Err('param address ' + a[1])
|
||||
e = f'PARAM{size * 8}({po + i * size}u)'
|
||||
elif 'shared' in mods:
|
||||
e = f'SH{size * 8}({V(base, "u32")} + {off + i * size}u)'
|
||||
elif 'local' in mods:
|
||||
e = f'loc[(uint({V(base, "u64")}) + {off + i * size}u) >> 2]'
|
||||
else:
|
||||
e = f'G{size * 8}({V(base, "u64")} + uint64_t({off + i * size}l))'
|
||||
if size == 8:
|
||||
A(d, e, 'u64')
|
||||
elif size == 2:
|
||||
A(d, e, 'u16')
|
||||
elif size == 1:
|
||||
A(d, e, 'u16')
|
||||
else:
|
||||
A(d, e, 'u32')
|
||||
|
||||
def st(self, o, a):
|
||||
V, E = self.val, self.E
|
||||
mods = set(o)
|
||||
ty = next(t for t in o if re.match(r'^(f|s|u|b)\d+$', t))
|
||||
size = int(ty[1:]) // 8
|
||||
base, off = self.addr(a[0])
|
||||
if 'shared' in mods:
|
||||
# vector stores go word by word; 16-bit pairs share a word (the kernels keep them 4-byte aligned)
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
|
||||
if size == 4:
|
||||
for i, v in enumerate(vals):
|
||||
E(f'sh[({V(base, "u32")} + {off + 4 * i}u) >> 2] = {V(v, "u32")};')
|
||||
return
|
||||
if size == 2 and len(vals) % 2 == 0:
|
||||
for i in range(0, len(vals), 2):
|
||||
E(f'sh[({V(base, "u32")} + {off + 2 * i}u) >> 2] = (uint({V(vals[i], "u16")}) & 0xFFFFu) | (uint({V(vals[i + 1], "u16")}) << 16);')
|
||||
return
|
||||
if 'local' in mods and size == 4:
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')] if a[1].startswith('{') else [a[1]]
|
||||
for i, v in enumerate(vals):
|
||||
E(f'loc[(uint({V(base, "u64")}) + {off + 4 * i}u) >> 2] = {V(v, "u32")};')
|
||||
return
|
||||
if 'global' in mods and size == 2 and not a[1].startswith('{'):
|
||||
E(f'W16({V(base, "u64")} + uint64_t({off}l)).v[0] = uint16_t({V(a[1], "u16")});'); return
|
||||
if 'global' in mods and size == 2 and a[1].startswith('{'):
|
||||
vals = [x.strip() for x in a[1][1:-1].split(',')]
|
||||
E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); ' + ' '.join(f'w_.v[{i}] = uint16_t({V(v, "u16")});' for i, v in enumerate(vals)) + ' }'); return
|
||||
if 'global' in mods and size == 4:
|
||||
E(f'{{ W16 w_ = W16({V(base, "u64")} + uint64_t({off}l)); uint v_ = {V(a[1], "u32")}; w_.v[0] = uint16_t(v_ & 0xffffu); w_.v[1] = uint16_t(v_ >> 16); }}'); return
|
||||
raise Err('store ' + '.'.join(o))
|
||||
|
||||
# ---------------------------------------------------------------- output
|
||||
def emit(self, tex_filter):
|
||||
nb = len(self.blocks)
|
||||
out = []
|
||||
w = out.append
|
||||
w('#version 460')
|
||||
# scripts/check_redistributable.py refuses files carrying this line: they are NVIDIA's code.
|
||||
w(f'// Generated by vulkan/ptx2glsl.py from NVIDIA PTX ({self.kernel}); do not commit or redistribute.')
|
||||
for e in ('GL_EXT_shader_explicit_arithmetic_types', 'GL_EXT_buffer_reference', 'GL_EXT_buffer_reference2', 'GL_EXT_scalar_block_layout',
|
||||
'GL_KHR_shader_subgroup_basic', 'GL_KHR_shader_subgroup_shuffle', 'GL_KHR_shader_subgroup_ballot', 'GL_EXT_samplerless_texture_functions'):
|
||||
w(f'#extension {e} : require')
|
||||
bx = 16 if self.maxntid[0] == 256 else self.maxntid[0]
|
||||
by = self.maxntid[0] // bx
|
||||
w(f'layout(local_size_x = {bx}, local_size_y = {by}) in;')
|
||||
w(f'layout(set = 0, binding = 0, std430) readonly buffer ParamBlock {{ uint w[{(self.psize + 3) // 4}]; }} param;')
|
||||
b = 1
|
||||
for off in sorted(self.textures):
|
||||
w(f'layout(set = 0, binding = {b}) uniform sampler2D T{off}; // {tex_filter.get(off, "?")}')
|
||||
b += 1
|
||||
for off in sorted(self.surfaces):
|
||||
fmt = 'r16ui' if self.surfaces[off] == 'b' else 'rgba16f'
|
||||
typ = 'uimage2D' if self.surfaces[off] == 'b' else 'image2D'
|
||||
w(f'layout(set = 0, binding = {b}, {fmt}) uniform writeonly {typ} S{off};')
|
||||
b += 1
|
||||
if self.exp_lut:
|
||||
w(f'layout(set = 0, binding = {b}, scalar) readonly buffer ExpLut {{ uint16_t e[65536]; }} explut;')
|
||||
w('uint ex2_lut(uint u) { return uint(explut.e[u & 0xffffu]) | (uint(explut.e[u >> 16]) << 16); }')
|
||||
w(f'shared uint sh[{max(self.shared_bytes // 4, 1)}];')
|
||||
w(RUNTIME)
|
||||
w('void main()\n{')
|
||||
if getattr(self, 'local_bytes', 0):
|
||||
w(f' uint loc[{(self.local_bytes + 3) // 4}];')
|
||||
for g in ('float', 'float16_t', 'f16vec2', 'uint', 'uint64_t', 'bool'):
|
||||
names = sorted(n for n, t in self.regs.items() if t == g)
|
||||
for i in range(0, len(names), 24):
|
||||
init = {'float': '0.0', 'float16_t': '0.0hf', 'f16vec2': 'f16vec2(0.0hf)', 'uint': '0u', 'uint64_t': '0ul', 'bool': 'false'}[g]
|
||||
w(f' {"precise " if g in ("float", "float16_t", "f16vec2") and not self.fast else ""}{g} ' + ', '.join(f'{n} = {init}' for n in names[i:i + 24]) + ';')
|
||||
# control flow graph over the blocks; END is block nb
|
||||
succ = []
|
||||
for i, blk in enumerate(self.blocks):
|
||||
br = blk[-1] if blk and isinstance(blk[-1], tuple) else None
|
||||
if br is None:
|
||||
succ.append([min(i + 1, nb)])
|
||||
else:
|
||||
tgt = nb if br[2][0] == 'END' else self.labels[br[2][1]]
|
||||
succ.append([tgt] if br[1] is None else [tgt, i + 1])
|
||||
# ---- structured emission: if/else arms that rejoin at the branch's immediate post-dominator --------------
|
||||
END = nb
|
||||
# reachable nodes, back edges (to a node on the DFS stack)
|
||||
color, back = {}, set()
|
||||
st = [(0, iter(succ[0]))]; color[0] = 1
|
||||
while st:
|
||||
n, it = st[-1]
|
||||
for m in it:
|
||||
if m == END:
|
||||
continue
|
||||
if color.get(m) == 1:
|
||||
back.add((n, m))
|
||||
elif m not in color:
|
||||
color[m] = 1; st.append((m, iter(succ[m]))); break
|
||||
else:
|
||||
color[n] = 2; st.pop()
|
||||
nodes = sorted(color)
|
||||
# each loop header h gets a synthetic "continue" node: back edges go there, and it flows to the loop's exit
|
||||
headers = sorted({h for _, h in back})
|
||||
cont = {h: END + 1 + i for i, h in enumerate(headers)}
|
||||
g = {n: [(cont[m] if (n, m) in back else m) for m in succ[n]] for n in nodes}
|
||||
for h in headers:
|
||||
latches = [n for n, m in back if m == h]
|
||||
exits = {m for n in latches for m in succ[n] if (n, m) not in back}
|
||||
if len(exits) != 1:
|
||||
raise Err('loop without a single latch exit')
|
||||
g[cont[h]] = [exits.pop()]
|
||||
g[END] = []
|
||||
allnodes = list(g)
|
||||
# immediate post-dominators (iterative, on the acyclic graph g)
|
||||
rpo, seen = [], set()
|
||||
def dfs(n):
|
||||
stack = [(n, iter(g[n]))]; seen.add(n)
|
||||
while stack:
|
||||
x, it = stack[-1]
|
||||
for m in it:
|
||||
if m not in seen:
|
||||
seen.add(m); stack.append((m, iter(g[m]))); break
|
||||
else:
|
||||
rpo.append(x); stack.pop()
|
||||
dfs(0)
|
||||
topo = rpo[::-1] # entry first
|
||||
index = {n: i for i, n in enumerate(topo)}
|
||||
ipdom = {END: END}
|
||||
def inter(x, y):
|
||||
while x != y:
|
||||
while index[x] < index[y]:
|
||||
x = ipdom[x]
|
||||
while index[y] < index[x]:
|
||||
y = ipdom[y]
|
||||
return x
|
||||
for n in reversed(topo): # END first
|
||||
if n == END:
|
||||
continue
|
||||
d = None
|
||||
for m in g[n]:
|
||||
d = m if d is None else inter(d, m)
|
||||
ipdom[n] = d
|
||||
emitted = [0]
|
||||
def body(n, ind):
|
||||
if n > END:
|
||||
w(ind + f'again_{n - END - 1} = true;'); return
|
||||
for line in self.blocks[n]:
|
||||
if not isinstance(line, tuple):
|
||||
w(ind + line)
|
||||
emitted[0] += 1
|
||||
def region(n, stop, ind):
|
||||
while n != stop:
|
||||
if n in cont and (n, 'open') not in region.active:
|
||||
# loop: body once per iteration until no back edge is taken
|
||||
e = g[cont[n]][0]
|
||||
w(ind + f'bool again_{cont[n] - END - 1};'); w(ind + 'do\n' + ind + '{')
|
||||
w(ind + f' again_{cont[n] - END - 1} = false;')
|
||||
region.active.add((n, 'open'))
|
||||
region(n, e, ind + ' ')
|
||||
region.active.discard((n, 'open'))
|
||||
w(ind + f'}} while (again_{cont[n] - END - 1});')
|
||||
n = e
|
||||
continue
|
||||
body(n, ind)
|
||||
sc = g[n]
|
||||
if len(sc) == 1:
|
||||
n = sc[0]; continue
|
||||
br = self.blocks[n][-1]
|
||||
m = ipdom[n]
|
||||
t_empty, f_empty = sc[0] == m, sc[1] == m
|
||||
if t_empty and f_empty:
|
||||
pass
|
||||
elif f_empty:
|
||||
w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' '); w(ind + '}')
|
||||
elif t_empty:
|
||||
w(ind + f'if (!({br[1]}))\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}')
|
||||
else:
|
||||
w(ind + f'if ({br[1]})\n' + ind + '{'); region(sc[0], m, ind + ' ')
|
||||
w(ind + '}\n' + ind + 'else\n' + ind + '{'); region(sc[1], m, ind + ' '); w(ind + '}')
|
||||
n = m
|
||||
region.active = set()
|
||||
import sys as _s
|
||||
_s.setrecursionlimit(20000)
|
||||
region(0, END, ' ')
|
||||
w('}')
|
||||
sys.stderr.write(f'reachable blocks={len(nodes)} emitted block copies={emitted[0]} loops={len(headers)}\n')
|
||||
return '\n'.join(out) + '\n'
|
||||
|
||||
|
||||
RUNTIME = r'''
|
||||
layout(buffer_reference, scalar, buffer_reference_align = 2) readonly buffer R16 { uint16_t v[4]; };
|
||||
layout(buffer_reference, scalar, buffer_reference_align = 2) writeonly buffer W16 { uint16_t v[4]; };
|
||||
float16_t H(uint u) { return uint16BitsToFloat16(uint16_t(u)); }
|
||||
uint HU(float16_t h) { return uint(float16BitsToUint16(h)); }
|
||||
f16vec2 H2(uint u) { return uint16BitsToFloat16(u16vec2(uint16_t(u & 0xffffu), uint16_t(u >> 16))); }
|
||||
uint HU2_(u16vec2 u) { return uint(u.x) | (uint(u.y) << 16); }
|
||||
uint H2U(f16vec2 h) { return HU2_(float16BitsToUint16(h)); }
|
||||
int s16(uint u) { return (int(u) << 16) >> 16; }
|
||||
float rcp_a(float x) { return 1.0 / x; }
|
||||
// x + copysign(0.5, x) rounded toward zero and then truncated: round half away from zero, exactly
|
||||
float round_half_away(float x) { float t = trunc(x); return abs(x - t) >= 0.5 ? t + sign(x) : t; }
|
||||
uint PARAM32(uint o) { return param.w[o >> 2]; }
|
||||
uint64_t PARAM64(uint o) { return uint64_t(param.w[o >> 2]) | (uint64_t(param.w[(o >> 2) + 1u]) << 32); }
|
||||
uint PARAM16(uint o) { return (param.w[o >> 2] >> ((o & 2u) * 8u)) & 0xffffu; }
|
||||
uint PARAM8(uint o) { return (param.w[o >> 2] >> ((o & 3u) * 8u)) & 0xffu; }
|
||||
uint SH32(uint a) { return sh[a >> 2]; }
|
||||
uint SH16(uint a) { return (sh[a >> 2] >> ((a & 2u) * 8u)) & 0xffffu; }
|
||||
uint G16(uint64_t a) { return uint(R16(a).v[0]); }
|
||||
uint G32(uint64_t a) { R16 r = R16(a); return uint(r.v[0]) | (uint(r.v[1]) << 16); }
|
||||
uint bfi32(uint ins, uint base, uint pos, uint len)
|
||||
{
|
||||
pos &= 0xffu; len &= 0xffu;
|
||||
if (pos >= 32u) return base;
|
||||
uint mask = len >= 32u ? (0xffffffffu << pos) : (((1u << len) - 1u) << pos);
|
||||
return (~mask & base) | (mask & (ins << pos));
|
||||
}
|
||||
// prmt.b32 (default mode): byte i of the result is byte (selector >> 4i) & 7 of the eight bytes b:a
|
||||
uint prmt(uint a, uint b, uint sel)
|
||||
{
|
||||
uint r = 0u;
|
||||
for (uint i = 0u; i < 4u; ++i) { const uint k = (sel >> (4u * i)) & 7u; r |= (((k < 4u ? a : b) >> (8u * (k & 3u))) & 0xffu) << (8u * i); }
|
||||
return r;
|
||||
}
|
||||
// shfl.sync.bfly within a 32-thread warp (ZLUDA's SHFL_SYNC_IMPL)
|
||||
uint shfl_bfly(uint v, int delta, uint opts, out bool in_bounds)
|
||||
{
|
||||
int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u);
|
||||
int self = int(gl_SubgroupInvocationID & 31u);
|
||||
int sub_end = (section_mask & self) | (~section_mask & warp_end);
|
||||
int idx = self ^ (delta & 31);
|
||||
bool oob = idx > sub_end;
|
||||
if (oob) idx = self;
|
||||
in_bounds = !oob;
|
||||
return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));
|
||||
}
|
||||
void sust_p(writeonly image2D s, int x, int y, vec4 v)
|
||||
{
|
||||
ivec2 n = imageSize(s);
|
||||
if (uint(x) < uint(n.x) && uint(y) < uint(n.y)) imageStore(s, ivec2(x, y), v);
|
||||
}
|
||||
// byte-addressed 32-bit store into a 16-bit single-channel surface: two texels
|
||||
void sust_b(writeonly uimage2D s, int xb, int y, uint v)
|
||||
{
|
||||
ivec2 n = imageSize(s);
|
||||
int x = xb >> 1;
|
||||
if (uint(y) < uint(n.y))
|
||||
{
|
||||
if (uint(x) < uint(n.x)) imageStore(s, ivec2(x, y), uvec4(v & 0xffffu));
|
||||
if (uint(x + 1) < uint(n.x)) imageStore(s, ivec2(x + 1, y), uvec4(v >> 16));
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
if __name__ == '__main__':
|
||||
filt = {}
|
||||
args = sys.argv[1:]
|
||||
while '--tex' in args:
|
||||
i = args.index('--tex'); k, v = args[i + 1].split(':'); filt[int(k)] = v; del args[i:i + 2]
|
||||
lut = '--exp-lut' in args
|
||||
if lut:
|
||||
args.remove('--exp-lut')
|
||||
fast = '--fast' in args
|
||||
if fast:
|
||||
args.remove('--fast')
|
||||
src = open(args[0]).read()
|
||||
t = T(src, args[1])
|
||||
t.exp_lut = lut
|
||||
if '--no-half-types' not in args:
|
||||
t.collect = True
|
||||
t.translate()
|
||||
parent = {}
|
||||
def find(x):
|
||||
while parent.setdefault(x, x) != x:
|
||||
parent[x] = parent[parent[x]]; x = parent[x]
|
||||
return x
|
||||
for x, y in t.links:
|
||||
if t.width.get(x) == t.width.get(y):
|
||||
parent[find(x)] = find(y)
|
||||
else:
|
||||
t.ctx.setdefault(x, set()).add('u'); t.ctx.setdefault(y, set()).add('u')
|
||||
groups = {}
|
||||
for x in set(t.ctx) | set(parent):
|
||||
groups.setdefault(find(x), set()).update(t.ctx.get(x, set()))
|
||||
half = {x for x in set(t.ctx) | set(parent) if x in t.width and groups[find(x)] == {'h'}}
|
||||
sys.stderr.write(f'half-typed registers: {len(half)} of {len(t.width)}\n')
|
||||
lut2 = t.exp_lut
|
||||
t = T(src, args[1]); t.exp_lut = lut2; t.half = half
|
||||
t.fast = fast
|
||||
t.translate()
|
||||
sys.stdout.write(t.emit(filt))
|
||||
sys.stderr.write(f'blocks={len(t.blocks)} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n')
|
||||
Reference in new issue
Block a user