diff --git a/engine/build.sh b/engine/build.sh index 02068a4..912304e 100755 --- a/engine/build.sh +++ b/engine/build.sh @@ -12,9 +12,11 @@ if [[ -f "$ROCM/include/hip/hip_runtime_api.h" ]]; then HIP=(-DD4R_ENGINE_HIP "$ROOT/engine/hip_net.cpp" -I"$ROCM/include" -L"$ROCM/lib" -lamdhip64 "-Wl,-rpath,$ROCM/lib") fi "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers -fPIC -shared \ - "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" -lvulkan -o "$OUT/libd4r_engine.so" + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" -lvulkan -o "$OUT/libd4r_engine.so" "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ "$ROOT/engine/replay.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-k" "$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ "$ROOT/engine/replay_m.cpp" "${HIP[@]}" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-m" -printf 'Built %s and %s\n' "$OUT/d4r-k" "$OUT/d4r-m" +"$CXX" -std=c++17 -O3 -Wall -Wextra -Wno-missing-field-initializers \ + "$ROOT/engine/replay_rr.cpp" -L"$OUT" -ld4r_engine -lvulkan '-Wl,-rpath,$ORIGIN' -o "$OUT/d4r-rr" +printf 'Built %s, %s and %s\n' "$OUT/d4r-k" "$OUT/d4r-m" "$OUT/d4r-rr" diff --git a/engine/compile_rr.py b/engine/compile_rr.py new file mode 100755 index 0000000..6d34067 --- /dev/null +++ b/engine/compile_rr.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Build the engine's Ray Reconstruction model directory (preset D or E) from a capture of the denoiser. + +usage: compile_rr.py CAPTURE_DIR OUT_DIR [--ptx-dir DIR] [--frame N] + +CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py: replay/ with the launches' weight +allocations, dump/, run/run.log), taken with the preset the model is for. Writes: + weights-.bin each stage's weight image as NGX uploads it (argument 40) + enc0.args, hkpn.args the two argument-block templates whose constants the runtime keeps (engine/runtime_rr.inc) + .spv one shader per stage, exposure_rr0.spv / exposure_rr1.spv + rr.txt the manifest RModel reads +A stage's shader is d4r's own engine/rr/.comp (or _d.comp for preset D) when it exists; otherwise it is +generated from the kernel's PTX in the user's own nvngx_dlssd.dll by engine/rr/ptx2rr.py (--ptx-dir: the modules +kernels/tools/extract_dlss_ptx.py wrote, named NNNN-.ptx; default ~/.cache/d4r-rr-corpus). +The weights, the templates and any generated shader are derived from NVIDIA's DLL: they stay on the user's machine +and are never packaged (scripts/check_redistributable.py). +""" +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE / 'rr')) +from rr_capture import Capture # noqa: E402 +from test_layer import stage_params # noqa: E402 + +STAGES = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1', 'dec0', 'hkpn'] +KERNEL = {s: f'cuda_dldn_engine_swin_{s}_kernel' for s in STAGES[:-1]} +KERNEL['hkpn'] = 'cuda_dldn_engine_hkpn_output_kernel_transformer' +# preset D runs its own enc0, dec0 and output kernels +D_KERNEL = {'enc0': 'cuda_dldn_engine_swin_enc0_diamond_wallaby_kernel', + 'dec0': 'cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel', + 'hkpn': 'cuda_dldn_engine_hkpn_output_kernel_transformer_diamond_wallaby'} + + +# d4r's own stage shaders: their texture and surface bindings (argument offsets, ascending) +OWN_BINDINGS = {'hkpn': ([176, 192, 200, 208, 216, 280, 288, 296, 448], [(224, 'rgba16f'), (240, 'rgba16f')])} +# own shaders that implement the reset path only: engine/rr/hkpn.comp upsamples the current frame and has no history +# reprojection, so a game (Reset = 0) gets a blocky, jittering image from it (RE Requiem, 2026-10-12). The translated +# kernel is used unless --reset-only asks for the own shader (captures that reset every frame). +RESET_ONLY = {'hkpn'} +# stages compiled without driver optimizations (engine/runtime_rr.inc, "noopt"): measured per stage against capE3 on +# RADV/Mesa 26.2 gfx1201, only enc1 is better for it (its merged output 43 -> 64 dB); every other stage loses 2-6 dB +NOOPT = {'enc1'} +# formatted stores (sust.p) whose surface is not the store's own guess: dec0's recurrent feature is NGX's SNORM_INT8X4 +SURFACE_FORMATS = {'dec0': {144: 'rgba8_snorm'}} + + +def glsl(source, output, *defines): + subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines, str(source), '-o', str(output)], + check=True, stdout=subprocess.DEVNULL) + + +def find_module(ptx_dir, kernel): + for path in sorted(Path(ptx_dir).glob('*.ptx')): + if f'.entry {kernel}(' in path.read_text(errors='replace'): + return path + raise FileNotFoundError(f'no PTX module with the entry {kernel} in {ptx_dir}') + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('capture') + ap.add_argument('output', type=Path) + ap.add_argument('--ptx-dir', default=os.path.expanduser('~/.cache/d4r-rr-corpus')) + ap.add_argument('--frame', type=int, default=2, help='which captured frame supplies the templates (1-based)') + ap.add_argument('--reset-only', action='store_true', help='use the own output shader that has no history path') + a = ap.parse_args() + cap = Capture(a.capture) + preset_d = any('diamond_wallaby' in l.kernel for l in cap.launches) + kernels = dict(KERNEL, **(D_KERNEL if preset_d else {})) + a.output.mkdir(parents=True, exist_ok=True) + lines = ['D4RR0001', f'preset {4 if preset_d else 5}'] + info = {'preset': 'D' if preset_d else 'E', 'stages': {}} + for stage in STAGES: + launches = [l for l in cap.launches if l.kernel == kernels[stage]] + if len(launches) < a.frame: + raise SystemExit(f'the capture has {len(launches)} launches of {kernels[stage]}; frame {a.frame} needed') + launch = launches[a.frame - 1] + (a.output / f'weights-{stage}.bin').write_bytes(launch.weights) + if stage in ('enc0', 'hkpn'): + (a.output / f'{stage}.args').write_bytes(stage_params(launch)) + own = HERE / 'rr' / (f'{stage}_d.comp' if preset_d and stage in D_KERNEL else f'{stage}.comp') + spv = a.output / f'{stage}.spv' + if own.exists() and (stage not in RESET_ONLY or a.reset_only): + glsl(own, spv) + origin = str(own.relative_to(HERE.parent)) + textures, surfaces = OWN_BINDINGS[own.stem] + else: + generated = a.output / f'{stage}.generated.comp' + formats = [f'--surface={off}={fmt}' for off, fmt in SURFACE_FORMATS.get(stage, {}).items()] + subprocess.run([sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(a.ptx_dir, kernels[stage])), + kernels[stage], str(generated), *formats], check=True) + glsl(generated, spv) + origin = 'generated from the DLL' + binding = json.loads(generated.with_suffix('.json').read_text()) + textures, surfaces = binding['textures'], [tuple(x) for x in binding['surfaces']] + tex = ','.join(map(str, textures)) or '-' + surf = ','.join(f'{off}:{fmt}' for off, fmt in surfaces) or '-' + lines.append(f'stage {stage} {stage}.spv 32 weights-{stage}.bin {tex} {surf}' + (' noopt' if stage in NOOPT else '')) + info['stages'][stage] = {'kernel': kernels[stage], 'shader': origin, 'weights': len(launch.weights)} + print(f'{stage}: {kernels[stage]} ({origin}), {len(launch.weights)} weight bytes') + lines += ['template enc0 enc0.args', 'template hkpn hkpn.args'] + for p in (0, 1): + glsl(HERE / 'exposure_k.comp', a.output / f'exposure_rr{p}.spv', '-DMODE=2', f'-DPASS={p}') + (a.output / 'rr.txt').write_text('\n'.join(lines) + '\n') + (a.output / 'model.json').write_text(json.dumps(info, indent=1) + '\n') + print(f'Ray Reconstruction preset {info["preset"]} model in {a.output}') + + +if __name__ == '__main__': + main() diff --git a/engine/exposure_k.comp b/engine/exposure_k.comp index 2c6f636..f1efa02 100644 --- a/engine/exposure_k.comp +++ b/engine/exposure_k.comp @@ -3,6 +3,8 @@ // MODE 0 (preset K): exposure = key * exp(-mean of the log luma). // MODE 1 (preset M, the cuda_dldn_engine_* kernels): exposure = key / max(mean of the luma clamped to // [0, 10000], 0.005). +// MODE 2 (Ray Reconstruction, the same cuda_dldn_engine_* kernels): MODE 1's measurement, written to two 1x1 +// images instead of stage parameter blocks. // Both modes measure the same grid: the colour rectangle sampled at half the output size with an // endpoint-aligned nearest-texel lookup (the reference's cuda_luma_convert_kernel; the grid equals the // rectangle itself only when the render size is half the output size). @@ -87,10 +89,15 @@ layout(set = 0, binding = 2, std430) buffer State { float current, previous; } s layout(set = 0, binding = 3, std430) buffer InputParams { float w[]; } inp; // expo, expoIn, expoRatio at words 28..30 layout(set = 0, binding = 4, std430) buffer OutputParams { float w[]; } outp; // expo, expoRatio at words 37, 38 layout(set = 0, binding = 5) uniform sampler2D gameExposure; // the game's 1x1 exposure (flag 4) -#else +#elif MODE == 1 layout(set = 0, binding = 3, std430) buffer InputParams { float w[]; } inp; // enc0_m: expo, expoPrev at words 25, 26 layout(set = 0, binding = 4, std430) buffer PostParams { float w[]; } post; // post_m: words 14, 15 layout(set = 0, binding = 5, std430) buffer DownParams { float w[]; } down; // down_m: word 12 +#else +// MODE 2 (Ray Reconstruction): the denoiser's stages sample the exposure as two 1x1 textures, this frame's and +// the previous frame's (zero on the first frame and after a reset) +layout(set = 0, binding = 3, r16f) uniform writeonly image2D exposureOut; +layout(set = 0, binding = 4, r16f) uniform writeonly image2D previousOut; #endif void main() { @@ -143,10 +150,13 @@ void main() precise const float ratio = scaled / scaledPrevious; inp.w[28] = exposure; inp.w[29] = 1.0; inp.w[30] = ratio; outp.w[37] = exposure; outp.w[38] = ratio; -#else +#elif MODE == 1 inp.w[25] = current; inp.w[26] = previous; post.w[14] = current; post.w[15] = previous; down.w[12] = current; +#else + imageStore(exposureOut, ivec2(0), vec4(current)); + imageStore(previousOut, ivec2(0), vec4(previous)); #endif } #endif diff --git a/engine/game.cpp b/engine/game.cpp index e1a021b..bf678d9 100644 --- a/engine/game.cpp +++ b/engine/game.cpp @@ -10,6 +10,7 @@ bool gameColorFormat(VkFormat format) { switch (format) { case VK_FORMAT_R16G16B16A16_SFLOAT: case VK_FORMAT_R32G32B32A32_SFLOAT: case VK_FORMAT_B10G11R11_UFLOAT_PACK32: case VK_FORMAT_A2B10G10R10_UNORM_PACK32: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_B8G8R8A8_UNORM: + case VK_FORMAT_E5B9G9R9_UFLOAT_PACK32: // R9G9B9E5_SHAREDEXP (RE Requiem's colour and output): sampled and blitted return true; default: return false; } diff --git a/engine/game_rr.cpp b/engine/game_rr.cpp new file mode 100644 index 0000000..b53d08a --- /dev/null +++ b/engine/game_rr.cpp @@ -0,0 +1,145 @@ +#include "game_rr.h" +#include "game_internal.h" +#include +#include +#include +#include +#include + +namespace d4r { + +struct GameDenoiser::Impl { + RModel& model; + REngine engine; + uint32_t width, height, outputWidth, outputHeight; + game_detail::CanonicalInputs fallback; // colour, motion, depth + std::array guides; // the guides and optional inputs, when one cannot be sampled directly + bool started = false; + // debugging: D4R_ENGINE_RR_DUMP=DIR copies every engine image after each frame into a host-visible buffer and + // writes it at the next frame's recording (by then the game has waited for the frame), as DIR/F-NAME.bin + VkPhysicalDevice physical; + VkDevice device; + VkBuffer dumpBuffer = VK_NULL_HANDLE; + VkDeviceMemory dumpMemory = VK_NULL_HANDLE; + void* dumpMapped = nullptr; + std::vector> dumpPending; // name, bytes, in buffer order + uint32_t dumpFrame = 0; + Impl(RModel& m, VkPhysicalDevice p, VkDevice d, uint32_t ow, uint32_t oh, uint32_t rw, uint32_t rh) + : model(m), engine(m, rw, rh, ow, oh), width(rw), height(rh), outputWidth(ow), outputHeight(oh), + fallback(p, d, rw, rh, rw, rh), physical(p), device(d) {} + ~Impl() { + if (dumpBuffer) vkDestroyBuffer(device, dumpBuffer, nullptr); + if (dumpMemory) vkFreeMemory(device, dumpMemory, nullptr); + } + static size_t texel(VkFormat f) { + return f == VK_FORMAT_R16G16B16A16_SFLOAT ? 8 : f == VK_FORMAT_R16_SFLOAT ? 2 : 4; + } + void writeDump(const char* dir) { + size_t offset = 0; + for (const auto& [name, bytes] : dumpPending) { + std::ofstream(std::string(dir) + "/" + std::to_string(dumpFrame) + "-" + name + ".bin", std::ios::binary) + .write(static_cast(dumpMapped) + offset, static_cast(bytes)); + offset += (bytes + 255) & ~size_t(255); + } + dumpPending.clear(); + } + void recordDump(VkCommandBuffer cb, uint32_t frame) { + const auto images = engine.debugImages(); + size_t total = 0; + for (const auto& d : images) total += (size_t(d.width) * d.height * texel(d.format) + 255) & ~size_t(255); + if (!dumpBuffer) { + VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bi.size = total; bi.usage = VK_BUFFER_USAGE_TRANSFER_DST_BIT; + game_detail::check(vkCreateBuffer(device, &bi, nullptr, &dumpBuffer), "vkCreateBuffer"); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(device, dumpBuffer, &r); + VkPhysicalDeviceMemoryProperties mp; vkGetPhysicalDeviceMemoryProperties(physical, &mp); + const VkMemoryPropertyFlags want = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT; + uint32_t type = 0; + while (type < mp.memoryTypeCount && !((r.memoryTypeBits >> type & 1u) && (mp.memoryTypes[type].propertyFlags & want) == want)) ++type; + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; ai.memoryTypeIndex = type; + game_detail::check(vkAllocateMemory(device, &ai, nullptr, &dumpMemory), "vkAllocateMemory"); + game_detail::check(vkBindBufferMemory(device, dumpBuffer, dumpMemory, 0), "vkBindBufferMemory"); + game_detail::check(vkMapMemory(device, dumpMemory, 0, VK_WHOLE_SIZE, 0, &dumpMapped), "vkMapMemory"); + } + game_detail::barrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_MEMORY_WRITE_BIT, + VK_ACCESS_TRANSFER_READ_BIT); + size_t offset = 0; + for (const auto& d : images) { + VkBufferImageCopy region{}; + region.bufferOffset = offset; + region.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; + region.imageExtent = {d.width, d.height, 1}; + vkCmdCopyImageToBuffer(cb, d.image, VK_IMAGE_LAYOUT_GENERAL, dumpBuffer, 1, ®ion); + const size_t bytes = size_t(d.width) * d.height * texel(d.format); + dumpPending.emplace_back(d.name, bytes); + offset += (bytes + 255) & ~size_t(255); + } + game_detail::barrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_HOST_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_HOST_READ_BIT); + dumpFrame = frame; + } +}; +GameDenoiser::GameDenoiser(RModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, + uint32_t outputHeight, uint32_t renderWidth, uint32_t renderHeight) + : impl(std::make_unique(model, physical, device, outputWidth, outputHeight, renderWidth, renderHeight)) {} +GameDenoiser::~GameDenoiser() = default; + +void GameDenoiser::record(VkCommandBuffer cb, const GameFrameR& frame) { + auto& g = *impl; + const uint32_t w = g.width, h = g.height; + const char* dump = std::getenv("D4R_ENGINE_RR_DUMP"); + static uint32_t frameCount = 0; + ++frameCount; + if (dump && !g.dumpPending.empty()) g.writeDump(dump); + if (!gameColorFormat(frame.color.format) || !gameMotionFormat(frame.motion.format) || + !gameDepthFormat(frame.depth.format, frame.depth.aspect) || !gameOutputFormat(frame.output.format)) + throw std::runtime_error("unsupported game texture format"); + // The stages sample every input at its texel coordinates from the origin, so only views at origin zero are + // borrowed; anything else is copied first. Every input must be in GENERAL when borrowed (one layout for all). + auto direct = [](const GameImage& image) { + return image.view != VK_NULL_HANDLE && !image.x && !image.y && image.layout == VK_IMAGE_LAYOUT_GENERAL; + }; + const bool directColor = direct(frame.color), directMotion = direct(frame.motion), directDepth = direct(frame.depth); + // the four guides, then the two optional inputs (kept in their own one-channel formats when copied) + const GameImage* guides[6] = {&frame.normals, &frame.roughness, &frame.diffuseAlbedo, &frame.specularAlbedo, + &frame.specularHitDistance, &frame.disocclusionMask}; + const VkFormat guideFormats[6] = {VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16B16A16_SFLOAT, + VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R32_SFLOAT, VK_FORMAT_R16_SFLOAT}; + auto present = [&](int i) { return i < 4 || guides[i]->image != VK_NULL_HANDLE; }; + if (!g.started) { + if (!g.model.uploadRecorded()) g.model.recordUpload(cb); + g.engine.recordInitialize(cb); + g.started = true; + } + g.fallback.record(cb, frame.color, frame.motion, frame.depth, w, h, w, h, directColor, directMotion, directDepth); + for (int i = 0; i < 6; ++i) { + if (!present(i) || direct(*guides[i])) continue; + if (!g.guides[i].image) { + g.fallback.ensure(cb, g.guides[i], guideFormats[i], w, h); + } + game_detail::barrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, VK_ACCESS_TRANSFER_READ_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + game_detail::copyInput(cb, *guides[i], g.guides[i].image, w, h); + } + game_detail::barrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_ACCESS_MEMORY_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + RFrameImages images{}; + images.color = directColor ? frame.color.view : g.fallback.color.view; + images.motion = directMotion ? frame.motion.view : g.fallback.motion.view; + images.depth = directDepth ? frame.depth.view : g.fallback.depth.view; + VkImageView* guideViews[6] = {&images.normals, &images.roughness, &images.diffuseAlbedo, &images.specularAlbedo, + &images.specularHitDistance, &images.disocclusionMask}; + for (int i = 0; i < 6; ++i) + if (present(i)) *guideViews[i] = direct(*guides[i]) ? guides[i]->view : g.guides[i].view; + const bool directOutput = frame.output.view != VK_NULL_HANDLE && frame.output.format == VK_FORMAT_R16G16B16A16_SFLOAT && + frame.output.layout == VK_IMAGE_LAYOUT_GENERAL && !frame.output.x && !frame.output.y && + frame.output.image != frame.color.image && frame.output.image != frame.motion.image && frame.output.image != frame.depth.image; + g.engine.setFrame(images, frame.settings, directOutput ? FrameOutput{frame.output.view, frame.output.format} : FrameOutput{}); + g.engine.recordFrame(cb); + if (dump) g.recordDump(cb, frameCount); + if (directOutput) // the output stage wrote the game's texture: make it visible to whatever the game does next + game_detail::barrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); + else + game_detail::copyOutput(cb, g.engine.outputImage(), frame.output, g.outputWidth, g.outputHeight); +} +} // namespace d4r diff --git a/engine/game_rr.h b/engine/game_rr.h new file mode 100644 index 0000000..2acc6f8 --- /dev/null +++ b/engine/game_rr.h @@ -0,0 +1,27 @@ +// Adapter between a game's own textures and the Ray Reconstruction engine (REngine): input conversion where a +// texture cannot be sampled directly, the frame and the copy into the game's output texture, all recorded into the +// game's command buffer. +#pragma once +#include "game.h" + +namespace d4r { +struct GameFrameR { + GameImage color, depth, motion, normals, roughness, diffuseAlbedo, specularAlbedo, output; + // optional (image null: the game passes none): DLSSD.SpecularHitDistance, DLSS.DisocclusionMask + GameImage specularHitDistance, disocclusionMask; + RFrameSettings settings; +}; +class GameDenoiser { +public: + // The model must outlive this object. The render size is fixed. + GameDenoiser(RModel& model, VkPhysicalDevice physical, VkDevice device, uint32_t outputWidth, uint32_t outputHeight, + uint32_t renderWidth, uint32_t renderHeight); + ~GameDenoiser(); // only once every recorded frame has completed on the GPU + GameDenoiser(const GameDenoiser&) = delete; + GameDenoiser& operator=(const GameDenoiser&) = delete; + void record(VkCommandBuffer command, const GameFrameR& frame); +private: + struct Impl; + std::unique_ptr impl; +}; +} // namespace d4r diff --git a/engine/replay_rr.cpp b/engine/replay_rr.cpp new file mode 100644 index 0000000..f17fa68 --- /dev/null +++ b/engine/replay_rr.cpp @@ -0,0 +1,245 @@ +// Native Ray Reconstruction replay: runs REngine (engine/runtime_rr.inc) on a capture's game inputs, frame after frame, +// and writes every image the engine produces (REngine::debugImages) after each frame, to compare with the capture's +// per-stage surfaces. No CUDA, NGX or Wine. The inputs come from engine/rr/rr_inputs.py. +// +// usage: d4r-rr MODEL_DIR INPUT_DIR OUT_DIR +// INPUT_DIR: frames.txt ("size RW RH OW OH", then one line per frame: "frame F jx jy mvx mvy preExposure reset +// depthInverted" followed by 16 worldToView and 16 viewToClip floats) and F-color.bin (RGBA16F), F-motion.bin +// (RG16F), F-depth.bin (R32F), F-normals.bin and F-roughness.bin (RGBA16F), F-diffuse.bin and F-specular.bin +// (RGBA8 UNORM), all at the render size. +// OUT_DIR receives F-NAME.bin for every debug image of frame F. +#include "runtime.h" +#include +#include +#include +#include +#include +#include +#include + +namespace { +void check(VkResult r, const char* operation) { + if (r != VK_SUCCESS) throw std::runtime_error(std::string(operation) + ": VkResult " + std::to_string(r)); +} +#define CK(call) check((call), #call) +std::vector read(const std::string& file) { + std::ifstream stream(file, std::ios::binary | std::ios::ate); + if (!stream) throw std::runtime_error("cannot read " + file); + std::vector data(static_cast(stream.tellg())); + stream.seekg(0); + stream.read(data.data(), static_cast(data.size())); + return data; +} +uint32_t texelBytes(VkFormat f) { + switch (f) { + case VK_FORMAT_R16G16B16A16_SFLOAT: return 8; + case VK_FORMAT_R16G16_SFLOAT: case VK_FORMAT_R32_SFLOAT: case VK_FORMAT_R8G8B8A8_UNORM: case VK_FORMAT_R8G8B8A8_SNORM: return 4; + case VK_FORMAT_R16_SFLOAT: return 2; + default: throw std::runtime_error("texelBytes: format"); + } +} +struct Ctx { + VkPhysicalDevice physical; + VkDevice device; + VkQueue queue; + VkCommandPool pool; + VkPhysicalDeviceMemoryProperties memory; + uint32_t type(uint32_t bits, VkMemoryPropertyFlags flags) const { + for (uint32_t i = 0; i < memory.memoryTypeCount; ++i) + if ((bits & (1u << i)) && (memory.memoryTypes[i].propertyFlags & flags) == flags) return i; + throw std::runtime_error("memory type"); + } + template void submit(F&& record) { + VkCommandBufferAllocateInfo ai{VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO}; + ai.commandPool = pool; ai.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; ai.commandBufferCount = 1; + VkCommandBuffer cb; CK(vkAllocateCommandBuffers(device, &ai, &cb)); + VkCommandBufferBeginInfo bi{VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO}; + CK(vkBeginCommandBuffer(cb, &bi)); record(cb); CK(vkEndCommandBuffer(cb)); + VkSubmitInfo si{VK_STRUCTURE_TYPE_SUBMIT_INFO}; si.commandBufferCount = 1; si.pCommandBuffers = &cb; + CK(vkQueueSubmit(queue, 1, &si, VK_NULL_HANDLE)); CK(vkQueueWaitIdle(queue)); + vkFreeCommandBuffers(device, pool, 1, &cb); + } +}; +struct HostBuffer { + VkBuffer buffer; VkDeviceMemory memory; void* mapped; + HostBuffer(Ctx& c, VkDeviceSize size) { + VkBufferCreateInfo bi{VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO}; bi.size = size; + bi.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + CK(vkCreateBuffer(c.device, &bi, nullptr, &buffer)); + VkMemoryRequirements r; vkGetBufferMemoryRequirements(c.device, buffer, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; + ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT); + CK(vkAllocateMemory(c.device, &ai, nullptr, &memory)); CK(vkBindBufferMemory(c.device, buffer, memory, 0)); + CK(vkMapMemory(c.device, memory, 0, VK_WHOLE_SIZE, 0, &mapped)); + } +}; +void copyRegion(VkBufferImageCopy& r, uint32_t w, uint32_t h) { + r = {}; r.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; r.imageExtent = {w, h, 1}; +} +struct Input { + VkImage image; VkImageView view; VkFormat format; + Input(Ctx& c, VkFormat f, uint32_t w, uint32_t h) : format(f) { + VkImageCreateInfo ii{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; + ii.imageType = VK_IMAGE_TYPE_2D; ii.format = f; ii.extent = {w, h, 1}; ii.mipLevels = ii.arrayLayers = 1; + ii.samples = VK_SAMPLE_COUNT_1_BIT; ii.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; + CK(vkCreateImage(c.device, &ii, nullptr, &image)); + VkMemoryRequirements r; vkGetImageMemoryRequirements(c.device, image, &r); + VkMemoryAllocateInfo ai{VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO}; ai.allocationSize = r.size; + ai.memoryTypeIndex = c.type(r.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); + VkDeviceMemory m; CK(vkAllocateMemory(c.device, &ai, nullptr, &m)); CK(vkBindImageMemory(c.device, image, m, 0)); + VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; + vi.image = image; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = f; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + CK(vkCreateImageView(c.device, &vi, nullptr, &view)); + } +}; +void barrier(VkCommandBuffer cb, VkImage image, VkImageLayout from, VkImageLayout to) { + VkImageMemoryBarrier b{VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER}; + b.image = image; b.oldLayout = from; b.newLayout = to; + b.srcAccessMask = VK_ACCESS_MEMORY_WRITE_BIT; b.dstAccessMask = VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT; + b.srcQueueFamilyIndex = b.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED; b.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, &b); +} +} // namespace + +int main(int argc, char** argv) try { + if (argc != 4) { std::cerr << "usage: d4r-rr MODEL_DIR INPUT_DIR OUT_DIR\n"; return 2; } + const std::string modelDir = argv[1], inputDir = argv[2], outDir = argv[3]; + 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 instance; CK(vkCreateInstance(&ici, nullptr, &instance)); + uint32_t count = 8; VkPhysicalDevice devices[8]; vkEnumeratePhysicalDevices(instance, &count, devices); + Ctx c{}; + c.physical = devices[0]; + for (uint32_t i = 0; i < count; ++i) { + VkPhysicalDeviceProperties p; vkGetPhysicalDeviceProperties(devices[i], &p); + if (p.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) c.physical = devices[i]; + } + vkGetPhysicalDeviceMemoryProperties(c.physical, &c.memory); + VkPhysicalDeviceRobustness2FeaturesEXT rb2{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ROBUSTNESS_2_FEATURES_EXT}; + VkPhysicalDeviceShaderFloat8FeaturesEXT f8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT, &rb2}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR cm{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, &f8}; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &cm}; + 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(c.physical, &f2); + const char* extensions[] = {VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME, VK_EXT_SHADER_FLOAT8_EXTENSION_NAME, + VK_EXT_ROBUSTNESS_2_EXTENSION_NAME}; + uint32_t families = 8; VkQueueFamilyProperties qf[8]; vkGetPhysicalDeviceQueueFamilyProperties(c.physical, &families, qf); + uint32_t family = 0; + for (uint32_t i = 0; i < families; ++i) if (qf[i].queueFlags & VK_QUEUE_COMPUTE_BIT) { family = i; break; } + float priority = 1; + VkDeviceQueueCreateInfo qi{VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO}; qi.queueFamilyIndex = family; qi.queueCount = 1; qi.pQueuePriorities = &priority; + VkDeviceCreateInfo di{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; di.queueCreateInfoCount = 1; di.pQueueCreateInfos = &qi; + // D4R_RR_TEST_ROBUST2=1: robustBufferAccess2 / robustImageAccess2 as vkd3d-proton enables them (out-of-bounds + // reads return zero) + const bool robust2 = std::getenv("D4R_RR_TEST_ROBUST2") != nullptr; + if (!robust2) { rb2.robustBufferAccess2 = rb2.robustImageAccess2 = rb2.nullDescriptor = VK_FALSE; } + di.enabledExtensionCount = robust2 ? 3 : 2; di.ppEnabledExtensionNames = extensions; + // D4R_RR_TEST_NO_ROBUST=1: no robust buffer or image access (out-of-bounds reads are undefined) + if (std::getenv("D4R_RR_TEST_NO_ROBUST")) { f2.features.robustBufferAccess = VK_FALSE; v13.robustImageAccess = VK_FALSE; } + // D4R_RR_TEST_NO_COOPMAT=1: a device like vkd3d-proton's, without the cooperative-matrix and float8 features + if (std::getenv("D4R_RR_TEST_NO_COOPMAT")) { di.enabledExtensionCount = 0; v13.pNext = nullptr; } + CK(vkCreateDevice(c.physical, &di, nullptr, &c.device)); + vkGetDeviceQueue(c.device, family, 0, &c.queue); + VkCommandPoolCreateInfo pi{VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO}; pi.queueFamilyIndex = family; + pi.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + CK(vkCreateCommandPool(c.device, &pi, nullptr, &c.pool)); + + auto framesText = read(inputDir + "/frames.txt"); + std::istringstream frames(std::string(framesText.begin(), framesText.end())); + std::string key; + uint32_t rw, rh, ow, oh; + frames >> key >> rw >> rh >> ow >> oh; + d4r::RModel model(c.physical, c.device, modelDir); + d4r::REngine engine(model, rw, rh, ow, oh); + const char* names[] = {"color", "motion", "depth", "normals", "roughness", "diffuse", "specular"}; + const VkFormat formats[] = {VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16_SFLOAT, VK_FORMAT_R32_SFLOAT, + VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R16G16B16A16_SFLOAT, VK_FORMAT_R8G8B8A8_UNORM, + VK_FORMAT_R8G8B8A8_UNORM}; + std::vector inputs; + for (VkFormat f : formats) inputs.emplace_back(c, f, rw, rh); + const char* optionalNames[] = {"hitdistance", "disocclusion"}; + const VkFormat optionalFormats[] = {VK_FORMAT_R32_SFLOAT, VK_FORMAT_R16_SFLOAT}; + std::vector optional; + for (VkFormat f : optionalFormats) optional.emplace_back(c, f, rw, rh); + HostBuffer staging(c, VkDeviceSize(ow) * oh * 8); + // GPU time of every frame: REngine's four stamps (start, after exposure, after the network, after the output) + VkQueryPoolCreateInfo qp{VK_STRUCTURE_TYPE_QUERY_POOL_CREATE_INFO}; + qp.queryType = VK_QUERY_TYPE_TIMESTAMP; qp.queryCount = 4; + VkQueryPool stamps; CK(vkCreateQueryPool(c.device, &qp, nullptr, &stamps)); + VkPhysicalDeviceProperties props; vkGetPhysicalDeviceProperties(c.physical, &props); + const int repeat = std::getenv("D4R_RR_TIME_REPEAT") ? std::atoi(std::getenv("D4R_RR_TIME_REPEAT")) : 1; + bool first = true; + while (frames >> key) { + uint32_t f; + d4r::RFrameSettings s; + int reset, inverted; + frames >> f >> s.jitter[0] >> s.jitter[1] >> s.motionScale[0] >> s.motionScale[1] >> s.preExposure >> reset >> inverted; + for (float& v : s.worldToView) frames >> v; + for (float& v : s.viewToClip) frames >> v; + s.reset = reset != 0; s.depthInverted = inverted != 0; + for (size_t i = 0; i < inputs.size(); ++i) { + auto data = read(inputDir + "/" + std::to_string(f) + "-" + names[i] + ".bin"); + if (data.size() != size_t(rw) * rh * texelBytes(formats[i])) throw std::runtime_error(std::string("input size: ") + names[i]); + std::memcpy(staging.mapped, data.data(), data.size()); + c.submit([&](VkCommandBuffer cb) { + barrier(cb, inputs[i].image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL); + VkBufferImageCopy r; copyRegion(r, rw, rh); + vkCmdCopyBufferToImage(cb, staging.buffer, inputs[i].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &r); + barrier(cb, inputs[i].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL); + }); + } + d4r::RFrameImages images{}; + images.color = inputs[0].view; images.motion = inputs[1].view; images.depth = inputs[2].view; + images.normals = inputs[3].view; images.roughness = inputs[4].view; + images.diffuseAlbedo = inputs[5].view; images.specularAlbedo = inputs[6].view; + // optional inputs, present when the capture had them (rr_inputs.py) + for (int k = 0; k < 2; ++k) { + std::ifstream probe(inputDir + "/" + std::to_string(f) + "-" + optionalNames[k] + ".bin", std::ios::binary); + if (!probe) continue; + auto data = read(inputDir + "/" + std::to_string(f) + "-" + optionalNames[k] + ".bin"); + if (data.size() != size_t(rw) * rh * texelBytes(optionalFormats[k])) throw std::runtime_error(std::string("input size: ") + optionalNames[k]); + std::memcpy(staging.mapped, data.data(), data.size()); + c.submit([&](VkCommandBuffer cb) { + barrier(cb, optional[k].image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL); + VkBufferImageCopy region; copyRegion(region, rw, rh); + vkCmdCopyBufferToImage(cb, staging.buffer, optional[k].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, ®ion); + barrier(cb, optional[k].image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_GENERAL); + }); + (k == 0 ? images.specularHitDistance : images.disocclusionMask) = optional[k].view; + } + engine.setFrame(images, s); + if (first) c.submit([&](VkCommandBuffer cb) { model.recordUpload(cb); engine.recordInitialize(cb); first = false; }); + double best = 1e30, network = 0, output = 0, exposure = 0; + for (int r = 0; r < repeat; ++r) { + if (r) engine.setFrame(images, s); + c.submit([&](VkCommandBuffer cb) { + vkCmdResetQueryPool(cb, stamps, 0, 4); + engine.recordFrame(cb, stamps, 0); + }); + uint64_t t[4]; + CK(vkGetQueryPoolResults(c.device, stamps, 0, 4, sizeof(t), t, 8, VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT)); + const double ms = props.limits.timestampPeriod * 1e-6; + if ((t[3] - t[0]) * ms < best) { + best = (t[3] - t[0]) * ms; exposure = (t[1] - t[0]) * ms; network = (t[2] - t[1]) * ms; output = (t[3] - t[2]) * ms; + } + } + std::cout << "frame " << f << ": " << best << " ms (exposure " << exposure << ", network " << network + << ", output " << output << ")\n"; + for (const auto& d : engine.debugImages()) { + const size_t bytes = size_t(d.width) * d.height * texelBytes(d.format); + c.submit([&](VkCommandBuffer cb) { + VkBufferImageCopy r; copyRegion(r, d.width, d.height); + vkCmdCopyImageToBuffer(cb, d.image, VK_IMAGE_LAYOUT_GENERAL, staging.buffer, 1, &r); + }); + std::ofstream(outDir + "/" + std::to_string(f) + "-" + d.name + ".bin", std::ios::binary) + .write(static_cast(staging.mapped), static_cast(bytes)); + } + std::cout << "frame " << f << " done\n"; + } + return 0; +} catch (const std::exception& e) { + std::cerr << "d4r-rr: " << e.what() << "\n"; + return 1; +} diff --git a/engine/rr/README.md b/engine/rr/README.md new file mode 100644 index 0000000..0201652 --- /dev/null +++ b/engine/rr/README.md @@ -0,0 +1,173 @@ +# Ray Reconstruction stages (DLSS-D, NGX feature 13) + +d4r's own Vulkan compute shaders for the denoiser in `nvngx_dlssd.dll` 310.7 (presets D and E). One shader per +NVIDIA kernel of the frame; the engine records them in NVIDIA's order: + +| stage | NVIDIA kernel | shader | grid (640x360 render, 1280x720 output) | +|---|---|---|---| +| exposure | `cuda_dldn_engine_luma_convert_kernel`, `reduce_sum` x2-3, `auto_exposure_copy_kernel` | `engine/exposure_k.comp` (MODE 1) | | +| enc0 | `cuda_dldn_engine_swin_enc0_kernel` (D: `..._enc0_diamond_wallaby_kernel`) | `enc0.comp` | 81x49, 32 threads | +| enc1-enc5 | `cuda_dldn_engine_swin_enc1_kernel` ... `enc5` | `enc1.comp` ... `enc5.comp` | | +| dec4-dec1 | `cuda_dldn_engine_swin_dec4_kernel` ... `dec1` | `dec4.comp` ... `dec1.comp` | | +| dec0 | `cuda_dldn_engine_swin_dec0_kernel` (D: `..._dec0_diamond_wallaby_kernel`) | `dec0.comp` | 81x49, 32 threads | +| output | `cuda_dldn_engine_hkpn_output_kernel_transformer` (D: `..._diamond_wallaby`) | `hkpn.comp` | 80x45, 16x16 threads | + +The D variants of enc0, dec0 and the output kernel are different code; the other nine kernels are shared and only +their weights differ between presets. + +## Stage contract + +Every stage shader reads the same data NVIDIA's kernel reads, in the same layouts, so each one can be checked on +its own against a captured launch (`test_layer.py`) and the stages chain exactly like NGX's kernels: + +- binding 0: the parameter block, `std430` readonly. It is NVIDIA's argument block byte for byte, except that every + pointer into the plane arena holds its **byte offset in the arena** (a 64-bit value), the weight pointer is zero, + and a texture or surface handle is 1 where NGX bound one and 0 where it left the argument null (the kernels test + optional inputs' handles for null). Fields are read at NVIDIA's argument offsets. +- binding 1: the layer's weight image exactly as NGX uploads it (argument 40's allocation), `std430` readonly. +- binding 2: the plane arena, `std430` read/write. +- bindings 3, 4, ...: the textures in ascending argument offset (combined image samplers, clamp-to-edge, + normalized coordinates, linear or nearest as NVIDIA's texture object; 8-bit unsigned arrays are sampled as UNORM, + `SNORM_INT8X4` as SNORM), then the surfaces in ascending argument offset (storage images, `rgba16f` or + `rgba8_snorm`). +- The dispatch grid is NVIDIA's launch grid. The workgroup size is the shader's choice. + +## Checking a stage + +Captures (bridge launch and replay dumps of a harness run, see `rr_capture.py`) live outside the repository. + +```sh +glslangValidator -V --target-env vulkan1.3 engine/rr/enc1.comp -o build/rr/enc1.spv +python3 engine/rr/test_layer.py CAPTURE swin_enc1 2 build/rr/enc1.spv build/rr/work-enc1 [--time 50] +``` + +It starts from the arena as it was before NVIDIA's launch, with zeroed surfaces, and compares every surface and +the whole arena with what NVIDIA's launch left. `build/rr/kvk` is `vulkan/kvk.cpp` +(`g++ -O2 -std=c++17 -o build/rr/kvk vulkan/kvk.cpp -lvulkan`). + +## Native stages (preset E) + +`build_native.py CAPTURE MODEL_DIR` turns a `compile_rr.py` model directory into one where every Swin stage is +d4r's own cooperative-matrix shader instead of translated PTX. Each was derived phase by phase from ZLUDA replays +(`phase_dump.py`: every mma's operands of one window) into an exact numpy reference, then written against it: + +| stages | reference | shader | weights | +|---|---|---|---| +| enc1-enc5, dec4-dec1 | `swin_ref.py` (one generic model: C, heads, merge / expand) | `swin_m.comp` | `swin_pack.py` | +| dec0 | `dec0_ref.py` | `dec0_m.comp` | `dec0_pack.py` | +| enc0 | `enc0_ref.py` (after the per-pixel features) | `enc0_m.comp` from `gen_enc0_m.py` | `enc0_pack.py` | + +`enc1_m.comp` / `dec1_m.comp` are the hand-written C = 64 instances `swin_m.comp` generalises; `enc1_ref.py` holds +the shared numerics (inter, mm, the cubic softmax, the norm trees). enc0 keeps ptx2rr's translation of NVIDIA's feature +gather (waves 0-2 run its pixel loop) and is native from the embedding on. + +Binding 1 differs from the stage contract above: it holds the packed weights (enc0: NGX's image, then the packed +blob). Everything else is the contract, so `test_layer.py --weights FILE` checks a native stage the same way. Against +capE3 every native stage leaves the arena byte-identical to NVIDIA's; surfaces differ in a few values by one binary16 +ulp (`exp2` approximation, and decoder texture 88 where NVIDIA's bilinear unit rounds differently), and the whole +network's output is ~96 dB from NVIDIA's. + +The shaders run every product transposed (weights x activations) so results feed the next product in registers; +`enc1_m.comp`'s header explains the layouts and `rr_exact.glsl` the binary16 rounding RADV needs (it folds +`float(float16_t(x))` and does not fuse f16 `fma`). + +## Native models and preset D + +`engine/rr/build_native_model.sh CAPTURE OUT_DIR [NVNGX_DLSSD_DLL]` builds a model whose eleven network stages are +the native shaders (`swin_m.comp`, `dec0_m.comp`, the generated `enc0_m.comp`) for the preset of the capture, with the +translated output stage. + +The weights are read from the user's DLL itself (`dll_weights.py`), not from the capture. The DLL holds each model's +tensors as plain binary16 arrays in natural `[out][in]` order, stage after stage; nothing names them, so they are +found by structure: every attention head has a 100 x 100 relative position bias table whose entries depend only on +the offset between two window positions, and a model is eleven consecutive stages with 2 2 4 4 8 8 8 4 4 2 2 heads. +What NGX uploads is the same data rearranged: FP8 codes in the kernels' fragment order, the norm gains times +sqrt(C), the bias times 2 sqrt(2) for the kernels' 64 x 96 query / key order, and (preset E) the 4096-byte noise +table encoder 0 dithers with. `dll_weights.py DLL OUT --model N --check DIR` compares a reading with a phase-dump +extraction: for 310.7 every tensor of both presets is identical, and so are the packed weight files. The capture +only selects which of the DLL's models it ran (and supplies the two argument-block templates); a DLL whose network +kernels are not the ones the shaders were written against (`native_kernels.txt`, SHA-256 of the PTX modules) is +refused unless `D4R_RR_ANY_KERNELS=1`. + +`compile_rr.py` uses the translated output kernel for both presets: `engine/rr/hkpn.comp` has the reset path only +(`--reset-only` selects it for captures that reset every frame). + +Preset D (`..._diamond_wallaby` kernels) differs from E in three places: + +- **enc0**: another feature gather (kept as the translation, like E's) with 30 input features for E's 32, and no + noise table (its weight image is 4096 bytes smaller). +- **dec0**: the eight bytes per pixel for the output stage are channels 12, 13, 0, 18, 19, 20, 21, 0 + (`-DRR_PRESET_D`; E: 12, 13, 18-23). Everything else is E's kernel. +- **output stage**: its own kernel (closest-depth motion from texture 296, a 4x4 bicubic history fetch). It is the + translated stage; it has no subgroup operations and runs with 64-wide waves (`rr.txt`). + +The reset flag (enc0 byte 365, output stage byte 440) follows `Reset`, not the jitter. With `Reset = 0` the output +stage reprojects and blends its history: check stages against a temporal capture (the harness with +`D4R_HARNESS_JITTER_SCENE=1`, and `D4R_CUDA_REPLAY_DUMP_LIMIT=400` so every launch is saved), not only against +captures that reset every frame. `engine/rr/hkpn.comp` implements the reset path only. + +Preset D, RX 9070 XT, RADV 26.2.4, 2026-10-12, 1280x720 -> 3840x2160 (`d4r-rr`, `D4R_RR_TIME_REPEAT=20`): 5.0-5.2 ms +per frame (network 4.5, output 0.55-0.65), from 17.5 ms with every stage translated. Every network stage reproduces +the capture's arena byte for byte on reset and temporal frames. The presented image against the CUDA path over +temporal frames: 96-106 dB on the reset frame, then 41-52 dB (one-ulp differences in the filtered surfaces are +amplified by the recurrent feature; twelve frames stay at 41-45 dB). + +RE Requiem (2026-10-12, 1280x720 -> 3840x2160, R9G9B9E5 colour, roughness packed in the normals' w): with +`hkpn.comp` the picture was 3-pixel blocks that shook with the jitter (no history); with the translated output stage +presets D and E both give a stable, reconstructed image on the engine. Not compared with the CUDA path in that game. + +## Matching the reference bit for bit (2026-10-12) + +The reference is the CUDA path through ZLUDA on the same GPU, so its `rcp.approx`, `ex2.approx` and bilinear texture +fetches are the GPU's own instructions and texture unit, not correctly rounded arithmetic. Two places in the native +filters had been "corrected" to exact arithmetic and were one ulp off in ~0.01-0.2% of the filtered surfaces: + +- `rcp16` (the filters' softmax normalisation and gate) is now the plain hardware reciprocal; +- the decoders' texture 88 goes through the bilinear sampler at NVIDIA's coordinates (`rr_decfilter.glsl`, + `dec0_m.comp`) instead of a hand-computed 1/4, 3/4 blend. + +With both, every network stage (arena and surfaces 112 / 120 / 128 / 144 / 304) is identical to the capture on every +frame of capDT12, capD4K, capD1080 (preset D) and capE3 (preset E). What remains is the translated output stage: +~0.004% of surface 224 / 240 one ulp off (single channels, values just above 0.5; cause not found). The output is +107-112 dB on frames where the network's inputs are still exact. At 640x360 -> 1080p that holds for all twelve +frames of capDT12 (was 96 dB, then 41-45); at 1280x720 and 1920x1080 -> 4K the history's differing values reach +encoder 0 through texture 496 and the output falls to 52 / 45 / 43 dB (capD4K), 53 / 49 dB (capD1080) as before. + +Speed: the compiled stages spend about half their instructions on float conversions. `spv_exact_convert.py` marks +the conversions NoContraction so RADV keeps a binary16 rounding without the opaque OR (`RR_EXACT_CONVERT`), which is +bit-exact and ~8% faster: preset D 4.8 ms at 1280x720 -> 4K, 9.5 ms at 1920x1080 -> 4K (was 5.1 / 10.3). +`D4R_RR_RELAXED=1|2` builds drop NVIDIA's double rounding instead (8.2 / 7.6 ms at 1080p) but are not exact: about 1% +of a layer's FP8 codes change and the output is 39-44 dB from the reference on every frame. + +Profile of a full-resolution stage at 1080p (dec0, 3.0 ms): expand 0.31, key norm 0.16, attention 1.37 (per head: +projections 0.15, scores + softmax 0.26, P V 0.17, out-projection 0.10), MLP 0.69, output filter 0.50. The 136 +matrix products themselves are 0.45-0.5 ms; weight and table memory traffic is under 8%. enc0's translated feature +gather is 0.28 of its 3.0 ms. + +## Optional inputs: specular hit distance and disocclusion mask (2026-10-12) + +RE Requiem passes `DLSSD.SpecularHitDistance` (R32_FLOAT) and `DLSS.DisocclusionMask` (R16_FLOAT). Harness captures +with synthetic ones (`D4R_HARNESS_RR_EXTRA=hit,mask`; capDX, capDXhit, capDXmask) show the wiring: + +| input | encoder 0 texture | encoder 0 flag byte | output stage | +|---|---|---|---| +| specular hit distance | argument 248 | 536 | flag byte 160 | +| disocclusion mask | argument 296 | 539 | - | + +Encoder 0 also tests the handle at argument 248 for null before it takes the hit distance path (it writes a +reprojected specular motion into channels 2 and 3 of surface 304), which is why the parameter block now carries the +handles' presence. With both bound the network is identical to the capture on every frame (capDX). The output +stage then differs from the CUDA path on the image's border pixels only (63-67 dB on temporal frames): NVIDIA's +kernel fetches neighbours outside the image there, CUDA clamps such integer fetches to the edge (every denoiser +texture's address mode is clamp; the translation and rrsim do the same), and ZLUDA's image load returns zero. With +zero outside the image the stage matches the capture like before (99.997%); the engine keeps CUDA's clamp. + +Exact speed-ups since the first native build, all bit-identical per stage: conversions kept by NoContraction instead of +an opaque OR, the MLP activation and its FP8 conversion from a 64 KiB table (`enc1_pack.act_table`), FP8 conversions in +pairs, no saturation clamp on attention probabilities. Preset D: 4.4 ms at 1280x720 -> 4K, 8.6 ms at 1920x1080 -> 4K. + +Output stage, hypotheses tested for its remaining ~0.004% one-ulp values and rejected (2026-10-12): correctly rounded +exp2 / reciprocal (rrsim's exact arithmetic is further from the capture than the GPU), fusing the kernel's 34 +single-use multiply-subtract pairs into fma, merging the two constant multiplications after the log transform, +`precise` on the inlined exp2, NoContraction on its conversions. The differing values come in spatial clusters, one +colour channel at a time, always one ulp, mostly just above 0.5 in surface 224. diff --git a/engine/rr/build_native.py b/engine/rr/build_native.py new file mode 100644 index 0000000..38386fb --- /dev/null +++ b/engine/rr/build_native.py @@ -0,0 +1,110 @@ +#!/usr/bin/env python3 +"""Turn a preset-E Ray Reconstruction model directory (engine/compile_rr.py) into the native one: every Swin stage +from d4r's own shaders instead of translated PTX. + +usage: build_native.py CAPTURE MODEL_DIR [--frame N] [--work DIR] + + enc0 engine/rr/enc0_m.comp, generated (gen_enc0_m.py) from the model's translated enc0.generated.comp, + whose per-pixel feature gather it keeps; weights: NGX's image + enc0_pack.py's blob + enc1-enc5, engine/rr/swin_m.comp with swin_pack.py's per-stage header and weights + dec4-dec1 + dec0 engine/rr/dec0_m.comp with dec0_pack.py's weights +The weights are extracted from ZLUDA replays of the capture's launches (engine/rr/phase_dump.py: every mma's operands +of one window) and the stages' weight images; extraction checks itself against the dump. Like the model directory, +everything written is derived from the user's DLL and stays on the machine. Runs one replay or compile at a time +(phase_dump.py caps its replay at 24 GB). +""" +import argparse +import os +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +from rr_capture import Capture # noqa: E402 +import swin_ref # noqa: E402 +import dec0_ref # noqa: E402 +import enc0_ref # noqa: E402 + +SWIN = ['enc1', 'enc2', 'enc3', 'enc4', 'enc5', 'dec4', 'dec3', 'dec2', 'dec1'] + + +def run(*cmd, **kw): + subprocess.run([str(c) for c in cmd], check=True, **kw) + + +def glsl(source, output, *includes): + run('glslangValidator', '-V', '--target-env', 'vulkan1.3', *[f'-I{i}' for i in includes], source, '-o', output, + stdout=subprocess.DEVNULL) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('model') + ap.add_argument('--frame', type=int, default=2) + ap.add_argument('--work', help='dumps and build files (default MODEL/native-work)') + o = ap.parse_args() + model = Path(o.model) + work = Path(o.work) if o.work else model / 'native-work' + work.mkdir(parents=True, exist_ok=True) + cap = Capture(os.path.expanduser(o.capture)) + manifest = (model / 'rr.txt').read_text().split('\n') + + def dump(stage): + out = work / f'{stage}_all.npz' + launch = cap.find(f'swin_{stage}', o.frame) + block = '5,5' if launch.grid[0] > 5 and launch.grid[1] > 5 else '1,1' + if not out.exists(): + run(sys.executable, HERE / 'phase_dump.py', o.capture, stage, o.frame, block, out) + return out, block + + def stage_line(stage, shader, weights): + for i, line in enumerate(manifest): + f = line.split() + if len(f) > 4 and f[0] == 'stage' and f[1] == stage: + f[2], f[4] = shader, weights + manifest[i] = ' '.join(x for x in f if x != 'noopt') + + for name in ('rr_exact.glsl', 'rr_attn.glsl', 'rr_decfilter.glsl', 'enc0_m_decl.glsl', 'enc0_m_tail.glsl'): + shutil.copy(HERE / name, work / name) + for stage in SWIN: + d, block = dump(stage) + print(f'{stage}: extracting', flush=True) + l, ex, ey, sx, sy, inp, outs, tex = swin_ref.load(o.capture, stage, o.frame) + bx, by = (int(v) for v in block.split(',')) + swin_ref.extract(stage, d, l.weights, work / f'{stage}_w.npz', inp, 8 * bx - sx, 8 * by - sy, ex, ey) + sd = work / stage + sd.mkdir(exist_ok=True) + run(sys.executable, HERE / 'swin_pack.py', stage, work / f'{stage}_w.npz', model / f'weights-{stage}m.bin', + sd / 'swin_stage.h') + glsl(HERE / 'swin_m.comp', model / f'{stage}_m.spv', sd, work) + stage_line(stage, f'{stage}_m.spv', f'weights-{stage}m.bin') + # dec0 + d, _ = dump('dec0') + print('dec0: extracting', flush=True) + dec0_ref.extract(d, cap.find('swin_dec0', o.frame).weights, work / 'dec0_w.npz') + run(sys.executable, HERE / 'dec0_pack.py', work / 'dec0_w.npz', model / 'weights-dec0m.bin', work / 'dec0_m_w.h') + glsl(HERE / 'dec0_m.comp', model / 'dec0_m.spv', work) + stage_line('dec0', 'dec0_m.spv', 'weights-dec0m.bin') + # enc0: NGX's weight image (the translated feature gather reads it) followed by the native blob + d, _ = dump('enc0') + print('enc0: extracting', flush=True) + nv = cap.find('swin_enc0', o.frame).weights + enc0_ref.extract(d, nv, work / 'enc0_w.npz') + base = (len(nv) + 255) // 256 * 256 + run(sys.executable, HERE / 'enc0_pack.py', work / 'enc0_w.npz', work / 'enc0_native.bin', work / 'enc0_m_w.h', base) + (model / 'weights-enc0m.bin').write_bytes(nv + bytes(base - len(nv)) + (work / 'enc0_native.bin').read_bytes()) + run(sys.executable, HERE / 'gen_enc0_m.py', model / 'enc0.generated.comp', work / 'enc0_m.comp') + glsl(work / 'enc0_m.comp', model / 'enc0_m.spv', work) + stage_line('enc0', 'enc0_m.spv', 'weights-enc0m.bin') + (model / 'rr.txt').write_text('\n'.join(manifest)) + print(f'{model / "rr.txt"}: native stages written') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/build_native_model.sh b/engine/rr/build_native_model.sh new file mode 100755 index 0000000..ecd4da5 --- /dev/null +++ b/engine/rr/build_native_model.sh @@ -0,0 +1,98 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Builds a Ray Reconstruction model directory whose network stages are the native shaders (engine/rr/*_m.comp), +# for the preset the capture was taken with (D or E). +# +# usage: build_native_model.sh CAPTURE_DIR OUT_DIR [NVNGX_DLSSD_DLL] +# +# CAPTURE_DIR is a bridge capture of a few denoiser frames (engine/rr/rr_capture.py); the DLL is the one it was taken +# with (default: CAPTURE_DIR/run/app/nvngx_dlssd.dll, then $D4R_RR_DLL). The steps: +# 1. engine/compile_rr.py: the translated model (the argument-block templates, the exposure passes, the translated +# output stage and the translated encoder 0 whose feature gather the native one keeps) +# 2. engine/rr/dll_weights.py: the network's weights, read from the DLL itself. The capture only selects which of +# the DLL's models it ran, and a DLL whose kernels are not the ones the native shaders were written against is +# refused (D4R_RR_ANY_KERNELS=1 builds anyway) +# 3. per stage: the weights checked against the capture's frame (STAGE_ref.py), packed (STAGE_pack.py) and the +# native shader compiled against that packing's header +# 4. rr.txt rewritten to the native stages; the weight images NGX uploaded in the capture are removed +# Preset D differs from E in three kernels: encoder 0 (its feature gather and 30 input features for E's 32), decoder 0 +# (the eight bytes it hands the output stage: -DRR_PRESET_D) and the output stage, which is the translated kernel +# (it has no subgroup operations, so it runs with 64-wide waves). +# Everything written is derived from the user's own nvngx_dlssd.dll and stays on their machine. +# +# GPU jobs run one at a time, each in a memory-capped scope (a translated stage's pipeline compile can take tens of GB). +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +CAP="$(realpath "${1:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")" +OUT="$(realpath -m "${2:?usage: build_native_model.sh CAPTURE_DIR OUT_DIR}")" +RR="$ROOT/engine/rr" +WORK="$OUT/native-work" +DLL="${3:-${D4R_RR_DLL:-$CAP/run/app/nvngx_dlssd.dll}}" +[[ -f "$DLL" ]] || { echo "no nvngx_dlssd.dll at $DLL: pass the DLL the capture was taken with" >&2; exit 1; } +FRAME="${D4R_RR_MODEL_FRAME:-2}" +CAPPED=(systemd-run --user --scope -q -p MemoryMax=24G -p MemorySwapMax=0) +mkdir -p "$OUT" "$WORK" +cd "$ROOT" + +python3 engine/compile_rr.py "$CAP" "$OUT" > "$WORK/compile_rr.log" 2>&1 || { tail -5 "$WORK/compile_rr.log"; exit 1; } +python3 "$RR/dll_weights.py" "$DLL" "$WORK/dll" --capture "$CAP" ${D4R_RR_ANY_KERNELS:+--any-kernels} +PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")" +# RR_EXACT_CONVERT: the compiled module's float conversions are decorated NoContraction (spv_exact_convert.py), which +# keeps the driver from folding a binary16 rounding away without the opaque OR the shaders otherwise need +DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D) +# D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact +[[ -n "${D4R_RR_RELAXED:-}" ]] && DEFS+=("-DRR_RELAXED=$D4R_RR_RELAXED") +glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR... + local out="$1" src="$2"; shift 2 + local dir inc=(); for dir in "$@"; do inc+=("-I$dir"); done + "${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 || + { grep -E 'ERROR' "$WORK/glsl.log" | head -5; exit 1; } + python3 "$RR/spv_exact_convert.py" "$out" > /dev/null +} +native() { # native STAGE WEIGHT_FILE: point rr.txt's stage line at STAGE_m.spv + sed -i -E "s/^stage $1 $1\.spv 32 weights-$1\.bin (.*)/stage $1 $1_m.spv 32 $2 \1/; /^stage $1 /s/ noopt\$//" "$OUT/rr.txt" + grep -q "^stage $1 $1_m.spv " "$OUT/rr.txt" || { echo "rr.txt: no stage line for $1" >&2; exit 1; } +} + +# the Swin stages shared by both presets (engine/rr/swin_m.comp) +for s in enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1; do + b=1,1; case $s in enc1|enc2|dec1|dec2) b=5,5 ;; esac + d="$WORK/$s"; mkdir -p "$d" + cp "$WORK/dll/$s.npz" "$d/weights.npz" + python3 "$RR/swin_ref.py" "$CAP" "$s" "$FRAME" "$d/weights.npz" --blocks "$b" 0,0 | tail -2 | sed "s/^/$s: /" + python3 "$RR/swin_pack.py" "$s" "$d/weights.npz" "$OUT/weights-${s}m.bin" "$d/swin_stage.h" > /dev/null + glsl "$OUT/${s}_m.spv" "$RR/swin_m.comp" "$d" "$RR" + native "$s" "weights-${s}m.bin" +done + +# decoder 0 +d="$WORK/dec0"; mkdir -p "$d" +cp "$WORK/dll/dec0.npz" "$d/weights.npz" +python3 "$RR/dec0_ref.py" "$CAP" "$FRAME" "$d/weights.npz" --blocks 5,5 0,0 | tail -2 | sed 's/^/dec0: /' +python3 "$RR/dec0_pack.py" "$d/weights.npz" "$OUT/weights-dec0m.bin" "$d/dec0_m_w.h" > /dev/null +glsl "$OUT/dec0_m.spv" "$RR/dec0_m.comp" "$d" "$RR" +native dec0 weights-dec0m.bin + +# encoder 0: the translated feature gather + the native rest +d="$WORK/enc0"; mkdir -p "$d" +cp "$WORK/dll/enc0.npz" "$d/weights.npz" +# preset E's gather dithers with the noise table that starts NGX's weight image: the table, then the packing +BASE=0; [[ "$PRESET" == 4 ]] || BASE=4096 +python3 "$RR/enc0_pack.py" "$d/weights.npz" "$d/enc0_m.bin" "$d/enc0_m_w.h" "$BASE" > /dev/null +if [[ "$BASE" != 0 ]]; then + [[ "$(stat -c %s "$WORK/dll/noise.bin" 2>/dev/null)" == 4096 ]] || { echo "no noise table found in $DLL" >&2; exit 1; } + cat "$WORK/dll/noise.bin" "$d/enc0_m.bin" > "$OUT/weights-enc0m.bin" +else + cp "$d/enc0_m.bin" "$OUT/weights-enc0m.bin" +fi +python3 "$RR/gen_enc0_m.py" "$OUT/enc0.generated.comp" "$d/enc0_m.comp" > /dev/null +glsl "$OUT/enc0_m.spv" "$d/enc0_m.comp" "$d" "$RR" +native enc0 weights-enc0m.bin + +# the translated output stage (preset D; preset E's when engine/rr/hkpn.comp is not used) has no subgroup operations +if grep -q '"shader": "generated from the DLL"' <(python3 -c "import json,sys; print(json.dumps(json.load(open('$OUT/model.json'))['stages']['hkpn']))"); then + sed -i -E 's/^stage hkpn hkpn\.spv 32 /stage hkpn hkpn.spv 64 /' "$OUT/rr.txt" +fi +for s in enc0 enc1 enc2 enc3 enc4 enc5 dec4 dec3 dec2 dec1 dec0; do rm -f "$OUT/weights-$s.bin"; done +printf 'Native Ray Reconstruction model (preset %s) in %s\n' "$([[ "$PRESET" == 4 ]] && echo D || echo E)" "$OUT" +grep '^stage' "$OUT/rr.txt" | awk '{ printf " %s %s (%s-wide)\n", $2, $3, $4 }' diff --git a/engine/rr/dec0_m.comp b/engine/rr/dec0_m.comp new file mode 100644 index 0000000..85e0472 --- /dev/null +++ b/engine/rr/dec0_m.comp @@ -0,0 +1,348 @@ +#version 460 +// Native decoder 0 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec0_kernel computed +// like engine/rr/enc1_m.comp (whose header explains the transposed products and the slot orders). The arithmetic is +// engine/rr/dec0_ref.py's (bit-exact to NVIDIA except surface 112's bilinear taps, which NVIDIA's texture unit rounds +// differently in a few percent of the values by one ulp); binding 1 holds engine/rr/dec0_pack.py's packing. +// +// One workgroup per 8x8-token window at the render resolution, four waves: +// expand the 6 x 6 low-resolution tokens under the window and its border (dec1's 64 channels) -> 4 sub-positions +// x 32 channels; wave w computes sub-position w, FP8 into EX in slot order +// prologue one key token per thread: expanded + skip (enc0's tokens), RMS norm -> XT; window tokens -> XR (f16) +// block two 32-channel heads and four MLP branches of 32 -> 32 -> 32, as enc1_m with C = 32 +// outputs one token per thread: 8 FP8 channels for the output stage (argument 48, [y][x][8]), tanh of four +// channels to surface 144, and the kernel-prediction filter of texture 88 (bilinear from half resolution) +// blended with textures 96 / 104 to surface 112 +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +layout(local_size_x = 128) in; +#include "dec0_m_w.h" +#include "rr_exact.glsl" + +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +layout(set = 0, binding = 3) uniform sampler2D T88; +layout(set = 0, binding = 4) uniform sampler2D T96; +layout(set = 0, binding = 5) uniform sampler2D T104; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112; +layout(set = 0, binding = 7, rgba8_snorm) uniform writeonly image2D S144; + +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) Q8PAIRS(dst, src) +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; + +// shared memory, aliased by phase: expand and prologue [LOW | EX], then the attention [K | V^T] +shared uint POOL[48 * 16 + 48 * 32]; +#define LOW POOL // the low-resolution tokens (6 x 6, padded to 48): 64 FP8 channels, arena order +#define EX_AT 768u // their expansion: 4 sub-positions x 32 FP8 channels in slot order (32 words a row) +#define RR_KS POOL // the current head's K +#define RR_KS0 0u +#define RR_VS POOL // the current head's V^T +#define RR_VS0 768u +shared uint XT[96 * 8]; // the normalised key tokens: slot-major, 32 FP8 channels in slot order +shared float16_t XR[64 * 32]; // the window tokens before the norm (the residual), slot-major +shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions +#define OUTF XR // after the block: the final tokens, token-major, natural channels (f16) + +// the window's key tokens: the 8x8 window, then a border row above, below, a column left and right +ivec2 keyXY(uint t, int x0, int y0) +{ + if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u)); + const int i = int(t & 7u), side = int((t - 64u) >> 3u); + return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i)); +} +int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); } +float tanhNV(float t) +{ + const float m = abs(t), s = t * t; + const float big = m >= 9.0109138 ? 1.0 : fma(1.0 / (exp2(m * 2.885390043258667) + 1.0), -2.0, 1.0); + return m >= 0.6 ? (t < 0.0 ? -big : big) : + fma(fma(fma(fma(fma(0.01573968306183815, s, -0.05230396240949631), s, 0.13315297663211823), s, -0.33332768082618713), s, 0.0), t, t); +} + +void main() +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2; + const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1; + const uint lplane = uint(hx * hy) * 32u; + + // ---------------------------------------------------------------- expand + for (uint i = tid; i < 48u * 16u; i += 128u) { + const uint r = i >> 4u, wd = i & 15u; + uint u = 0u; + if (r < 36u) { + const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy); + u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)]; + } + LOW[i] = u; + } + barrier(); + { + A8 we[2][4]; + [[unroll]] for (uint mm = 0u; mm < 2u; ++mm) + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(we[mm][kt], W_WE, 64u, 2u * wv + mm, kt); + [[unroll]] for (uint n = 0u; n < 3u; ++n) { + B8 lb[4]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(lb[kt], LOW, n * 256u + 4u * kt, 16u, CM); + [[unroll]] for (uint mm = 0u; mm < 2u; ++mm) { + const uint m = 2u * wv + mm; + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); + K32(a, we[mm][0], lb[0], we[mm][1], lb[1]) + K32(a, we[mm][2], lb[2], we[mm][3], lb[3]) + B8 q; + Q8(q, a) + coopMatStore(q, POOL, EX_AT + n * 512u + 4u * m, 32u, CM); + } + } + } + barrier(); + + // ---------------------------------------------------------------- prologue: one key token per thread + if (tid < 96u) { + const ivec2 k = keyXY(inter(tid), x0, y0); + const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1)); + const int xx = mirror(k.x, ex), yy = mirror(k.y, ey); + const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u; + float16_t v[32]; + [[unroll]] for (uint wd = 0u; wd < 8u; ++wd) { + const uint e = POOL[EX_AT + r * 32u + 8u * s + wd], sk = arena.w[sb + wd]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + precise float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((e >> (8u * b)) & 0xffu))) + + float16_t(uintBitsToFloate4m3EXT(uint8_t((sk >> (8u * b)) & 0xffu))); + v[4u * wd + b] = x; + } + } + // dec0_ref.norm_plane: per word t the pair sums, words t and t + 4, then t ^ 1, t ^ 2, the halves + precise f16vec2 h[4]; + [[unroll]] for (uint t = 0u; t < 4u; ++t) { + precise f16vec2 p0 = f16vec2(v[4u * t], v[4u * t + 1u]) * f16vec2(v[4u * t], v[4u * t + 1u]) + + f16vec2(v[4u * t + 2u], v[4u * t + 3u]) * f16vec2(v[4u * t + 2u], v[4u * t + 3u]); + precise f16vec2 p1 = f16vec2(v[16u + 4u * t], v[17u + 4u * t]) * f16vec2(v[16u + 4u * t], v[17u + 4u * t]) + + f16vec2(v[18u + 4u * t], v[19u + 4u * t]) * f16vec2(v[18u + 4u * t], v[19u + 4u * t]); + h[t] = p0 + p1; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t sm = hs.x + hs.y; + precise float16_t se = sm + float16_t(0.0001220703125); + const float16_t rr = rsq16(se); + [[unroll]] for (uint wd = 0u; wd < 8u; ++wd) { + floate4m3_t c[4]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + const uint j = 4u * wd + b; + precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)]; + precise float16_t n = v[j] * g; + c[b] = q8h(n); + } + XT[tid * 8u + wd] = code(c[0]) | (code(c[1]) << 8u) | (code(c[2]) << 16u) | (code(c[3]) << 24u); + } + if (tid < 64u) [[unroll]] for (uint j = 0u; j < 32u; ++j) XR[tid * 32u + j] = v[j]; + } + barrier(); + + // ---------------------------------------------------------------- the residual: window tokens + table 0 + ACC R[2]; + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + HAC x; + coopMatLoad(x, XR, wv * 512u + 16u * m, 32u, CM); + [[unroll]] for (int i = 0; i < 8; ++i) { + precise float16_t c = float16_t(x[i]) + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)]; + R[m][i] = float(c); + } + } + + // ---------------------------------------------------------------- attention, two heads + [[unroll]] for (uint hd = 0u; hd < 2u; ++hd) { + if (hd == 1u) barrier(); + for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c) + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + A8 wk[2], xa[2]; + B8 wvb[2], xb[2]; + [[unroll]] for (uint kt = 0u; kt < 2u; ++kt) { + coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM); + coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM); + LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt); + LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); + } + ACC a = ACC(0.0); + K32(a, wk[0], xb[0], wk[1], xb[1]) + A8 ka; + Q8(ka, a) + coopMatStore(ka, POOL, n * 128u + 4u * c, 8u, RM); + ACC v = ACC(0.0); + K32(v, xa[0], wvb[0], xa[1], wvb[1]) + A8 va; + Q8(va, v) + coopMatStore(va, POOL, 768u + c * 384u + 4u * n, 24u, RM); + } + barrier(); +#include "rr_attn.glsl" + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(R[m], a0, Ob[0], a1, Ob[1]) + } + } + + // ---------------------------------------------------------------- MLP + B8 Nb[2]; + { + float16_t rv[2][8]; + [[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]); + // dec0_ref.norm_plane on slot order: lane half hl holds words 2hl, 2hl + 1 (tile 0) and 4 + 2hl, 5 + 2hl (tile 1) + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1]; + precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1]; + precise float16_t hv = p0 + p1; + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < 2u; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint k = 16u * m + 8u * hl + uint(i); + precise float16_t g = r * w16.w[W_GAIN1 / 2u + k]; + precise float16_t n = rv[m][i] * g; + Nb[m][i] = q8h(n); + precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k]; + R[m][i] = float(a); + } + } + for (uint br = 0u; br < 4u; ++br) { + const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + A8 a0, a1; + LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u); + K32(h, a0, Nb[0], a1, Nb[1]) + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); + } + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(R[m], a0, Hb[0], a1, Hb[1]) + } + } + + // ---------------------------------------------------------------- outputs, one token per thread + { + const uint t = inter(16u * wv + col); + [[unroll]] for (uint m = 0u; m < 2u; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 32u + inter(16u * m + 8u * hl + uint(i))] = float16_t(R[m][i]); + } + barrier(); + // texture 88 (one level down) at the window's 10 x 10 full-resolution tap positions (clamped like the filter's + // coordinates), through the bilinear sampler at NVIDIA's coordinates: the reference's texture unit does not round + // like f32 arithmetic (computing the 1/4, 3/4 blend by hand left ~0.2% of surface 112 one ulp off, capDT 2026-10-12) + if (tid < 100u) { + const int xi = clampTap(x0 - 1 + int(tid % 10u), ex), yi = clampTap(y0 - 1 + int(tid / 10u), ey); + const float iw = uintBitsToFloat(param.p[34]), ih = uintBitsToFloat(param.p[35]); + precise float cu = fma(iw, float(xi), iw * 0.5), cv = fma(ih, float(yi), ih * 0.5); + const vec4 bl = textureLod(T88, vec2(cu, cv), 0.0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) TAP[tid * 4u + c] = float16_t(bl[c]); + } + barrier(); + // two threads per token: half 0 filters channels 0 and 1, half 1 channel 2 and writes the output stage's codes + // and surface 144 + { + const uint t = tid >> 1u, hf = tid & 1u; + const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u); + if (x < 0 || y < 0 || x >= ex || y >= ey) return; + float16_t a[10]; + [[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 32u + c]; + if (hf == 1u) { + float16_t b[12]; + [[unroll]] for (uint c = 0u; c < 12u; ++c) b[c] = OUTF[t * 32u + 12u + c]; + const uint o = (outOff + uint(y * ex + x) * 8u) / 4u; +#ifdef RR_PRESET_D + // cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel: channels 12, 13, 0, 18, 19, 20, 21, 0 (its only + // difference from the preset E kernel; capDT, 2026-10-12) + arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 24u); + arena.w[o + 1u] = code(q8h(b[7])) | (code(q8h(b[8])) << 8u) | (code(q8h(b[9])) << 16u); +#else + arena.w[o] = code(q8h(b[0])) | (code(q8h(b[1])) << 8u) | (code(q8h(b[6])) << 16u) | (code(q8h(b[7])) << 24u); + arena.w[o + 1u] = code(q8h(b[8])) | (code(q8h(b[9])) << 8u) | (code(q8h(b[10])) << 16u) | (code(q8h(b[11])) << 24u); +#endif + imageStore(S144, ivec2(x, y), vec4(tanhNV(float(b[2])), tanhNV(float(b[3])), tanhNV(float(b[4])), tanhNV(float(b[5])))); + } + // surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9) + // blends it with texture 96; + texture 104 + const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]); + const float16_t mx = max(max(ma, mb), a[8]); + float e[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t rr = rcp16(s); + float16_t wg[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr)); + const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf; + float16_t f[2] = float16_t[2](0.0hf, 0.0hf); + bool mid = false; + [[unroll]] for (uint k = 0u; k < 9u; ++k) { + const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0; + [[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid); + } + if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare) + f[0] = 0.0hf; f[1] = 0.0hf; + for (uint k = 0u; k < 9u; ++k) { + const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0; + for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]); + } + } + const float16_t g = sig(a[9]); + precise float16_t om = float16_t(1.0) - g; + const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0); + float16_t o[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + precise float16_t fg = f[c] * g; + float16_t v = fma16(float16_t(t96[c0 + c]), om, fg); + const float16_t t4 = float16_t(t104[c0 + c]); + if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; } + o[c] = v; + } + const uint other = subgroupShuffleXor(bh(o[0]), 1u); + if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0)); + } +} diff --git a/engine/rr/dec0_pack.py b/engine/rr/dec0_pack.py new file mode 100644 index 0000000..a715e31 --- /dev/null +++ b/engine/rr/dec0_pack.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Pack decoder 0's extracted weights (engine/rr/dec0_ref.py --extract) for engine/rr/dec0_m.comp. + +usage: dec0_pack.py WEIGHTS.npz OUT.bin OUT.h + +The layout rules are enc1_pack.py's: transposed products, every weight's columns in the k order of the product that +consumes its result (C = 32: slot j holds natural channel inter(32)[j]). The expand's 128 columns are four +sub-positions of 32 channels, each permuted the same way, so a key token's expanded half is already in slot order. +""" +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from enc1_ref import inter # noqa: E402 +from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402 + +I32, I96 = inter(32), inter(96) +I128 = [32 * (j // 32) + I32[j % 32] for j in range(128)] + + +def main(): + src, out_bin, out_h = sys.argv[1:4] + w = dict(np.load(src)) + e = {'we': ('A', w['we'][:, I128]), 'be': ('H', w['be'][I128])} + for h in range(2): + e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32]) + e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32]) + e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32]) + e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T) + for b in range(4): + e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32]) + e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32]) + e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32]) + e['gain0'] = ('H', w['gain0']) + for k in ('table0', 'gain1', 'table1'): + e[k] = ('H', w[k][I32]) + e['cubic'] = ('U', cubic_table()) + e['act'] = ('Y', act_table()) + blob, defines = bytearray(), [] + for name, (kind, a) in e.items(): + while len(blob) % 256: + blob.append(0) + off = len(blob) + if kind == 'A': + blob += a_tiles(a).tobytes() + elif kind == 'B': + blob += b_tiles(a).tobytes() + elif kind == 'Y': + blob += np.ascontiguousarray(a, np.uint8).tobytes() + elif kind == 'U': + blob += np.ascontiguousarray(a, np.uint16).tobytes() + else: + blob += np.ascontiguousarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write('// generated by engine/rr/dec0_pack.py\n' + '\n'.join(defines) + '\n') + print(f'{len(blob)} bytes, {len(e)} entries') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/dec0_ref.py b/engine/rr/dec0_ref.py new file mode 100644 index 0000000..04b97b6 --- /dev/null +++ b/engine/rr/dec0_ref.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +"""Exact numpy reference of the denoiser's decoder 0 (cuda_dldn_engine_swin_dec0_kernel), one 8x8 window at a time, +from its weights and real inputs only: the specification the native shader implements. Derived like +engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py). + + expand the low-resolution tokens under the window and its border (6 x 6 of dec1's 64-channel output, arena order) + -> 128 = 4 sub-positions x 32 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1) + prologue key token = expanded + skip token (enc0's 32-channel output at argument 24), f16; RMS norm (enc1's per-word + tree on one plane) times the gain at W+8448; the residual is the token + the table at W+39232 + attention two heads of 32 channels (K = 32, one k32 step), as enc1 + MLP RMS norm, four branches of 32 -> 32 (+ bias, ramp activation) -> 32 onto the residual (+ table) + +usage: dec0_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies the phases of a dump, see --dump) + dec0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump) +""" +import argparse +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import e4m3, f16, q8 # noqa: E402 +from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, H, F # noqa: E402 + +I32 = inter(32) +UN32 = [I32.index(i) for i in range(32)] +I96 = inter(96) + + +def norm_plane(xn, gain): + """RMS norm of one 32-channel token (natural order) in NVIDIA's tree over the arena order; returns codes in k order""" + xa = xn[I32] + sq = (xa * xa).astype(H) + h = [] + for t in range(4): + def pair(w): + return add(sq[4 * w:4 * w + 2], sq[4 * w + 2:4 * w + 4]) + h.append(add(pair(t), pair(t + 4))) + h01 = [add(h[t], h[t ^ 1]) for t in range(4)] + h2 = add(h01[0], h01[2]) + s = H(H(h2[0] + h2[1]) + H(2.0 ** -13)) + r = H(F(1) / np.sqrt(F(s))) + return q8((xn * (r * gain).astype(H)).astype(H))[I32] + + +def keys(x0, y0): + border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)] + + [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)]) + return [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border + + +def window(w, lo, skip, x0, y0, ex, ey, dbg=None): + """lo: dec1's output [2][ey/2][ex/2][32] codes; skip: enc0's [ey][ex][32] codes. Returns the residual stream after + the block (64 x 32 f16, natural channels)""" + dbg = {} if dbg is None else dbg + hx, hy = ex // 2, ey // 2 + X0, Y0 = x0 // 2 - 1, y0 // 2 - 1 + low = np.zeros((48, 64), np.uint8) + for r in range(36): + X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy) + low[r] = np.concatenate([lo[0, Y, X], lo[1, Y, X]]) + Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0))) + dbg['p1_A'] = low + xs, ns = [], [] + for (x, y) in keys(x0, y0): + xx, yy = mirror(x, ex), mirror(y, ey) + r = (yy // 2 - Y0) * 6 + (xx // 2 - X0) if False else (y // 2 - Y0) * 6 + (x // 2 - X0) + s = 2 * (y & 1) + (x & 1) + xn = (e4m3(Ex[r, 32 * s:32 * s + 32]).astype(H) + e4m3(skip[yy, xx][UN32]).astype(H)).astype(H) + xs.append(xn) + ns.append(norm_plane(xn, w['gain0'])) + X = np.array(ns) + dbg['p2_A'] = X + R = None + for h in range(2): + qk = mm(X, w[f'wk{h}'], np.zeros((96, 32))) + v = mm(X, w[f'wv{h}'], np.zeros((96, 32))) + Q = q8(qk)[:, I32] + S = mm(Q[:64], Q.T, w[f'bias{h}']) + P = softmax_cubic(S) + O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32))) + C = add(np.array(xs[:64]), w['table0'][None, :]) if h == 0 else R + R = mm(q8(O)[:, I32], w[f'wo{h}'], C) + dbg[f'p{4 + 5 * h}_A'], dbg[f'p{5 + 5 * h}_A'], dbg[f'p{6 + 5 * h}_A'] = Q[:64], q8(P)[:, I96], q8(O)[:, I32] + dbg[f'p{6 + 5 * h}_C'], dbg[f'p{6 + 5 * h}_D'] = C, R + N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)]) + acc = add(R, w['table1'][None, :]) + for b in range(4): + hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0))) + dbg[f'p{12 + 2 * b}_A'] = N + dbg[f'p{13 + 2 * b}_A'], dbg[f'p{13 + 2 * b}_C'] = q8(hid)[:, I32], acc + acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc) + dbg[f'p{13 + 2 * b}_D'] = acc + return acc + + +OUT_CH = [12, 13, 18, 19, 20, 21, 22, 23] # the 8 FP8 channels of the output stage's input (argument 48) +OUT_CH_D = [12, 13, None, 18, 19, 20, 21, None] # preset D (..._dec0_diamond_wallaby_kernel): bytes 2 and 7 are zero +TANH_CH = [14, 15, 16, 17] # surface 144 (rgba8 snorm) + + +def fma16(a, b, c): + return H(np.float64(a) * np.float64(b) + np.float64(c)) + + +def tanh_nv(t): + """NVIDIA's tanh of the surface-144 channels: an odd polynomial below 0.6, 1 - 2 / (e^2x + 1) above, f32""" + t = F(t) + m = abs(t) + big = F(1) if m >= F(9.0109138) else F(np.float64(F(1) / F(np.exp2(F(m * F(2.8853900432586670))) + F(1))) * -2 + 1) + s = F(t * t) + p = F(0.01573968306183815) + for c in (-0.05230396240949631, 0.13315297663211823, -0.33332768082618713, 0.0): + p = F(np.float64(p) * np.float64(s) + np.float64(F(c))) + small = F(np.float64(p) * np.float64(t) + np.float64(t)) + return np.copysign(big, t) if m >= F(0.6) else small + + +def bilinear88(t88, xi, yi): + """texture 88 (half resolution, linear) at the full-resolution texel centre (xi, yi): weights 1/4 and 3/4, + clamped at the edge; f32 sum rounded to binary16 (NVIDIA's texture unit rounds a few of these differently)""" + def axis(i): + k = i // 2 + return ((k - 1, 0.25), (k, 0.75)) if i % 2 == 0 else ((k, 0.75), (k + 1, 0.25)) + h, w = t88.shape[:2] + v = np.zeros(4) + for (yy, wy) in axis(yi): + for (xx, wx) in axis(xi): + v += wx * wy * t88[min(max(yy, 0), h - 1), min(max(xx, 0), w - 1)].astype(np.float64) + return v.astype(H) + + +def outputs(acc, t88, t96, t104, x0, y0, ex, ey): + """per window token: the 8 output codes, surface 144 (4 floats before the snorm store), surface 112 (3 f16)""" + codes, s144, s112 = [], [], [] + for t in range(64): + x, y = x0 + t % 8, y0 + t // 8 + a = acc[t].astype(H) + codes.append(q8(a)[OUT_CH]) + s144.append([tanh_nv(a[c]) for c in TANH_CH]) + lg = a[:9] + m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8]) + m = max(max(m01[0], m01[1]), lg[8]) + e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)] + pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)] + sm = F(pair[0] + F(0)) + for k in (1, 2, 3): + sm = F(sm + pair[k]) + sm = F(sm + e[8]) + r = H(F(1) / sm) + wg = [H(H(e[i]) * r) for i in range(9)] + + def cl(v, n): + return 1 if v < 0 else (n - 2 if v >= n else v) + xs, ys = [cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)], [cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)] + f = np.zeros(3, H) + k = 0 + for yy in ys: + for xx in xs: + T = bilinear88(t88, xx, yy) + for c in range(3): + f[c] = fma16(T[c], wg[k], f[c]) + k += 1 + g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1))) + om = H(H(1) - g) + o = np.zeros(3, H) + for c in range(3): + if 0 <= x < ex and 0 <= y < ey: + o[c] = fma16(t96[y, x, c], om, H(f[c] * g)) + if not np.isnan(t104[y, x, c]): + o[c] = H(o[c] + t104[y, x, c]) + s112.append(o) + return np.array(codes), np.array(s144), np.array(s112) + + +def extract(dump, weights, out): + z = np.load(dump) + W = np.frombuffer(weights, np.uint8) + + def B(p): + return z[f'p{p}_w0_B'] + + def C(p): + return f16(z[f'p{p}_w0_C']).astype(H) + w = {'we': B(1), 'be': C(1)[0], 'gain0': W[8448:8448 + 64].view(H), 'table0': W[39232:39232 + 64].view(H), + 'gain1': W[39296:39296 + 64].view(H), 'table1': W[41472:41472 + 64].view(H)} + for h, base in ((0, 2), (1, 7)): + w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1) + w[f'bias{h}'] = C(base + 2) + w[f'wo{h}'] = B(base + 4) + for b in range(4): + w[f'w1_{b}'], w[f'b1_{b}'] = B(12 + 2 * b), C(12 + 2 * b)[0] + w[f'w2_{b}'] = B(13 + 2 * b) + np.savez_compressed(out, **w) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('frame', type=int) + ap.add_argument('weights') + ap.add_argument('--extract', metavar='DUMP') + ap.add_argument('--dump', help='phase dump of --blocks[0] to verify phase by phase') + ap.add_argument('--blocks', nargs='*', default=['5,5']) + o = ap.parse_args() + from rr_capture import Capture + c = Capture(os.path.expanduser(o.capture)) + l = c.find('swin_dec0', o.frame) + if o.extract: + extract(o.extract, l.weights, o.weights) + return + ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32)) + sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16)) + ptr = {p.offset: p.delta for p in l.pointers} + pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8) + lo = pre[ptr[8]:ptr[8] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32) + skip = pre[ptr[24]:ptr[24] + ex * ey * 32].reshape(ey, ex, 32) + w = dict(np.load(o.weights)) + img = {i.offset: i for i in l.images} + + def surf(off, dt=np.float16): + i = img[off] + return np.frombuffer(i.path.read_bytes(), dt).reshape(i.height, i.width, i.channels) + t88, t96, t104, r112, r144 = surf(88), surf(96), surf(104), surf(112), surf(144, np.int8) + post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8) + rout = post[ptr[48]:ptr[48] + ex * ey * 8].reshape(ey, ex, 8) + for blk in o.blocks: + bx, by = (int(v) for v in blk.split(',')) + d = {} + x0, y0 = 8 * bx - sx, 8 * by - sy + acc = window(w, lo, skip, x0, y0, ex, ey, d) + codes, s144, s112 = outputs(acc, t88, t96, t104, x0, y0, ex, ey) + if 'diamond' in l.kernel: + codes = np.array([[0 if ch is None else q8(acc[t].astype(H))[ch] for ch in OUT_CH_D] for t in range(64)]) + n = ok_c = ok_s = ok_t = 0 + for t in range(64): + x, y = x0 + t % 8, y0 + t // 8 + if not (0 <= x < ex and 0 <= y < r144.shape[0]): + continue + n += 1 + ok_c += np.array_equal(codes[t], rout[y, x]) + ok_t += np.array_equal(np.clip(np.round(np.array(s144[t], np.float64) * 127), -127, 127).astype(np.int8), r144[y, x]) + ok_s += np.array_equal(s112[t], r112[y, x, :3]) + print(f'block {bx},{by}: output codes {ok_c}/{n}, surface 144 {ok_t}/{n}, surface 112 {ok_s}/{n}') + if o.dump: + z = np.load(o.dump) + for k in sorted(d, key=lambda s: (int(s[1:].split('_')[0]), s)): + ref = z[f'{k.split("_")[0]}_w0_{k.split("_")[1]}'] + ours = np.asarray(d[k]) + ours = ours.view(np.uint16) if ours.dtype == H else ours + n = min(len(ref), len(ours)) + print(f'{k}: rows equal {int((ref[:n] == ours[:n]).all(1).sum())}/{n}') + o.dump = None + + +if __name__ == '__main__': + main() diff --git a/engine/rr/dec1_m.comp b/engine/rr/dec1_m.comp new file mode 100644 index 0000000..a8be455 --- /dev/null +++ b/engine/rr/dec1_m.comp @@ -0,0 +1,295 @@ +#version 460 +// Native decoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_dec1_kernel. +// engine/rr/dec1_ref.py's arithmetic (bit-exact to NVIDIA except surface 112's bilinear taps, see dec0_m.comp); +// binding 1 holds engine/rr/dec1_pack.py's packing. The expand and prologue are dec0_m.comp's (64 channels from dec2's +// 96), the block is enc1_m.comp's (whose header explains the layouts), the outputs are the tokens (argument 48) and +// dec0_m.comp's filter on the 64 -> 16 head. +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +layout(local_size_x = 128) in; +#include "dec1_m_w.h" +#include "rr_exact.glsl" + +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +layout(set = 0, binding = 3) uniform sampler2D T88; +layout(set = 0, binding = 4) uniform sampler2D T96; +layout(set = 0, binding = 5) uniform sampler2D T104; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112; + +#define SKIP 0 +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]); +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; + +// shared memory, aliased by phase: expand [LOW | EX], attention [KS | VS], outputs [TAP | OUTF] +shared uint POOL[48 * 24 + 48 * 64]; +#define LOW POOL // 48 low-resolution tokens (6 x 6, padded): 96 FP8 channels, arena order (24 words) +#define EX_AT (48u * 24u) // their expansion: 4 sub-positions x 64 FP8 channels in slot order (64 words) +#define KS POOL // the current head's K: token slot x 32 FP8 channels +#define VS_AT 768u // the current head's V^T: 32 channels x 96 token slots +#define RR_KS POOL +#define RR_KS0 0u +#define RR_VS POOL +#define RR_VS0 VS_AT +shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in slot order +shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions +shared float16_t OUTF[64 * 16]; // the head per window token (raster order) + +ivec2 keyXY(uint t, int x0, int y0) +{ + if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u)); + const int i = int(t & 7u), side = int((t - 64u) >> 3u); + return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i)); +} +int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); } +uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d) +{ + return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u); +} +float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); } + +void main() +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2; + const uint loOff = param.p[2], skOff = param.p[6], outOff = param.p[12], shift = param.p[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1; + const uint lplane = uint(hx * hy) * 32u, plane = uint(ex * ey) * 32u; + + // ---------------------------------------------------------------- expand: wave w computes sub-position w + for (uint i = tid; i < 48u * 24u; i += 128u) { + const uint r = i / 24u, wd = i % 24u; + uint u = 0u; + if (r < 36u) { + const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy); + u = arena.w[(loOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)]; + } + LOW[i] = u; + } + barrier(); + [[unroll]] for (uint n = 0u; n < 3u; ++n) { + B8 lb[6]; + [[unroll]] for (uint kt = 0u; kt < 6u; ++kt) coopMatLoad(lb[kt], POOL, n * 384u + 4u * kt, 24u, CM); + [[unroll]] for (uint mm = 0u; mm < 4u; ++mm) { + const uint m = 4u * wv + mm; + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < 6u; kt += 2u) { + A8 w0, w1; + LDA(w0, W_WE, 96u, m, kt); LDA(w1, W_WE, 96u, m, kt + 1u); + K32(a, w0, lb[kt], w1, lb[kt + 1u]) + } + B8 q; + Q8(q, a) + coopMatStore(q, POOL, EX_AT + n * 1024u + 4u * m, 64u, CM); + } + } + barrier(); + + // ---------------------------------------------------------------- prologue: one key token per thread + if (tid < 96u) { + const ivec2 k = keyXY(inter(tid), x0, y0); + const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1)); + const int xx = mirror(k.x, ex), yy = mirror(k.y, ey); + const uint sb = (skOff + uint(yy * ex + xx) * 32u) / 4u; + float16_t v[64]; + [[unroll]] for (uint wd = 0u; wd < 16u; ++wd) { + const uint e = POOL[EX_AT + r * 64u + 16u * s + wd], sk = arena.w[sb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(sk, b); v[4u * wd + b] = x; } + } + // enc1's input norm: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves + precise f16vec2 h[4]; + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + precise f16vec2 p[4]; + [[unroll]] for (uint kk = 0u; kk < 4u; ++kk) { + const uint o = 32u * (kk & 1u) + 4u * (q + 4u * (kk >> 1u)); + precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]); + precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]); + p[kk] = a + c; + } + precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3]; + h[q] = s0 + s1; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t sm = hs.x + hs.y; + precise float16_t se = sm + float16_t(0.0001220703125); + const float16_t rr = rsq16(se); + [[unroll]] for (uint wd = 0u; wd < 16u; ++wd) { + floate4m3_t c[4]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + const uint j = 4u * wd + b; + precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)]; + precise float16_t n = v[j] * g; + c[b] = q8h(n); + } + XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]); + } + } + barrier(); + + // ---------------------------------------------------------------- the residual: expanded + skip + table 7 + ACC R[4]; + { + const uint tq = inter(16u * wv + col); + const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u); + const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1)); + const uint sb = (skOff + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u; + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + const uint ea = EX_AT + r * 64u + 16u * s + 4u * m + 2u * hl; + const uint sa = sb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl; + const uvec2 e = uvec2(POOL[ea], POOL[ea + 1u]), sk = uvec2(arena.w[sa], arena.w[sa + 1u]); + [[unroll]] for (int i = 0; i < 8; ++i) { + precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(sk[i >> 2], uint(i) & 3u); + precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)]; + R[m][i] = float(c); + } + } + } + barrier(); // EX is read; the attention reuses its memory + + // ---------------------------------------------------------------- attention, two heads + [[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) { + // this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows + // to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots) + if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS + for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) { + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + A8 wk[4]; + B8 wvb[4], xb[4]; + A8 xa[4]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) { + coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM); + coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM); + LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt); + LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); + } + ACC a = ACC(0.0); + K32(a, wk[0], xb[0], wk[1], xb[1]) + K32(a, wk[2], xb[2], wk[3], xb[3]) + A8 ka; + Q8(ka, a) + coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM); + ACC v = ACC(0.0); + K32(v, xa[0], wvb[0], xa[1], wvb[1]) + K32(v, xa[2], wvb[2], xa[3], wvb[3]) + A8 va; + Q8(va, v) + coopMatStore(va, POOL, VS_AT + c * 384u + 4u * n, 24u, RM); + } + barrier(); +#include "rr_attn.glsl" + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(R[m], a0, Ob[0], a1, Ob[1]) + } + } + + // ---------------------------------------------------------------- MLP + B8 Nb[4]; + { + float16_t rv[4][8]; + [[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]); + // slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1), + // t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree) + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1]; + precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1]; + precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1]; + precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1]; + precise float16_t hv = (a + cc) + (b + d); + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < 4u; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint k = 16u * m + 8u * hl + uint(i); + precise float16_t g = r * w16.w[W_GAIN14 / 2u + k]; + precise float16_t n = rv[m][i] * g; + Nb[m][i] = q8h(n); + precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k]; + R[m][i] = float(a); + } + } + for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) { + const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + A8 a0, a1, a2, a3; + LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u); + K32(h, a0, Nb[0], a1, Nb[1]) + K32(h, a2, Nb[2], a3, Nb[3]) + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i]))); + } + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(R[m], a0, Hb[0], a1, Hb[1]) + } + } + + // ---------------------------------------------------------------- outputs + { + const uint t = inter(16u * wv + col); + const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + B8 Tb[4]; + const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u; + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + Q8(Tb[m], R[m]) + const uint at = 4u * (m & 1u) + 2u * hl; + if (inside) { + arena.w[o + (m >> 1u) * (plane >> 2u) + at] = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]); + arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]); + } + } + // the 64 -> 16 head, per token to OUTF + ACC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + A8 a0, a1, a2, a3; + LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u); + K32(hv, a0, Tb[0], a1, Tb[1]) + K32(hv, a2, Tb[2], a3, Tb[3]) + [[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]); + } + barrier(); +#include "rr_decfilter.glsl" +} diff --git a/engine/rr/dec1_pack.py b/engine/rr/dec1_pack.py new file mode 100644 index 0000000..5273f28 --- /dev/null +++ b/engine/rr/dec1_pack.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python3 +"""Pack decoder 1's extracted weights (engine/rr/dec1_ref.py --extract) for engine/rr/dec1_m.comp. + +usage: dec1_pack.py WEIGHTS.npz OUT.bin OUT.h + +enc1_pack.py's layout (transposed products, columns in the consumer's k order) plus the expand (dec0_pack.py): 256 +columns, four sub-positions of 64 channels each permuted to slot order. +""" +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from enc1_ref import inter # noqa: E402 +from enc1_pack import a_tiles, b_tiles, cubic_table # noqa: E402 + +I32, I64, I96 = inter(32), inter(64), inter(96) +I256 = [64 * (j // 64) + I64[j % 64] for j in range(256)] + + +def main(): + src, out_bin, out_h = sys.argv[1:4] + w = dict(np.load(src)) + e = {'we': ('A', w['we'][:, I256]), 'be': ('H', w['be'][I256])} + for h in range(2): + e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32]) + e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32]) + e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64]) + e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T) + for b in range(8): + e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32]) + e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32]) + e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64]) + e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh']) + e['gain0'] = ('H', w['gain0']) + for k in ('table7', 'gain14', 'table16'): + e[k] = ('H', w[k][I64]) + e['cubic'] = ('U', cubic_table()) + blob, defines = bytearray(), [] + for name, (kind, a) in e.items(): + while len(blob) % 256: + blob.append(0) + off = len(blob) + if kind == 'A': + blob += a_tiles(a).tobytes() + elif kind == 'B': + blob += b_tiles(a).tobytes() + elif kind == 'U': + blob += np.ascontiguousarray(a, np.uint16).tobytes() + else: + blob += np.ascontiguousarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write('// generated by engine/rr/dec1_pack.py\n' + '\n'.join(defines) + '\n') + print(f'{len(blob)} bytes, {len(e)} entries') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/dec1_ref.py b/engine/rr/dec1_ref.py new file mode 100644 index 0000000..e182c5c --- /dev/null +++ b/engine/rr/dec1_ref.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +"""Exact numpy reference of the denoiser's decoder 1 (cuda_dldn_engine_swin_dec1_kernel), one 8x8 window at a time. +Derived like engine/rr/enc1_ref.py, phase by phase against ZLUDA replays (engine/rr/phase_dump.py); NVIDIA runs it with +two warps, each one head and half of the expand's columns. + + expand the 6 x 6 level-2 tokens under the window and its border (dec2's 96 channels, arena order) -> 256 = + 4 sub-positions x 64 channels (+ bias), FP8; sub-position 2 (y & 1) + (x & 1) (dec0_ref) + prologue key token = expanded + skip token (enc1's 64-channel output at argument 24), f16; enc1's input norm (two + planes) times the gain at W+25088; the residual is the token + the table at W+62080 + block enc1's: two 32-channel heads, eight MLP branches 64 -> 32 -> 64 (gain W+62208, table W+66496) + outputs the tokens (argument 48); the 64 -> 16 head and dec0's filter (bilinear texture 88 from level 2, gate, + textures 96 / 104) to surface 112 + +usage: dec1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture) + dec1_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump) +""" +import argparse +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import e4m3, f16, q8 # noqa: E402 +from enc1_ref import inter, mm, softmax_cubic, act, mirror, add, input_norm, nv_rowsum_sq, H, F # noqa: E402 +from dec0_ref import keys, fma16, bilinear88 # noqa: E402 + +I32, I64, I96 = inter(32), inter(64), inter(96) +UN64 = [I64.index(i) for i in range(64)] +GAIN0, TABLE7, GAIN14, TABLE16 = 25088, 62080, 62208, 66496 + + +def window(w, lo, skip, x0, y0, ex, ey): + """lo: dec2's output [3][ey/2][ex/2][32] codes; skip: enc1's [2][ey][ex][32]. Returns (tokens 64 x 64 codes in arena + order, head 64 x 16 f16)""" + hx, hy = ex // 2, ey // 2 + X0, Y0 = x0 // 2 - 1, y0 // 2 - 1 + low = np.zeros((48, 96), np.uint8) + for r in range(36): + X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy) + low[r] = np.concatenate([lo[p, Y, X] for p in range(3)]) + Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0))) + xs, ns = [], [] + for (x, y) in keys(x0, y0): + xx, yy = mirror(x, ex), mirror(y, ey) + r = (y // 2 - Y0) * 6 + (x // 2 - X0) + s = 2 * (y & 1) + (x & 1) + sk = np.concatenate([skip[0, yy, xx], skip[1, yy, xx]])[UN64] + xn = (e4m3(Ex[r, 64 * s:64 * s + 64]).astype(H) + e4m3(sk).astype(H)).astype(H) + xs.append(xn) + xa = xn[I64] + n, _ = input_norm([xa[:32], xa[32:]], w['gain0']) + ns.append(q8(n)[I64]) + X = np.array(ns) + res = np.array(xs[:64]) + R = None + for h in range(2): + qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32)))) + v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32)))) + Q = q8(qk)[:, I32] + S = mm(Q[:64], Q.T, w[f'bias{h}']) + P = softmax_cubic(S) + O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32))) + C = add(res, w['table7'][None, :]) if h == 0 else R + R = mm(q8(O)[:, I32], w[f'wo{h}'], C) + N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, I64] + acc = add(R, w['table16'][None, :]) + for b in range(8): + hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0))) + acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc) + tokens = q8(acc)[:, I64] + head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0)) + return tokens, head + + +def filter112(head, t88, t96, t104, x0, y0, ex, ey): + """dec0_ref.outputs' surface 112 on the head's channels 0-8 (logits) and 9 (gate)""" + out = [] + for t in range(64): + x, y = x0 + t % 8, y0 + t // 8 + a = head[t].astype(H) + lg = a[:9] + m01 = np.maximum(np.maximum(np.maximum(lg[0:2], lg[2:4]), lg[4:6]), lg[6:8]) + m = max(max(m01[0], m01[1]), lg[8]) + e = [np.exp2(F(H(lg[i] - m)) * F(1.4426950216293335)) for i in range(9)] + pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)] + sm = F(pair[0] + F(0)) + for k in (1, 2, 3): + sm = F(sm + pair[k]) + sm = F(sm + e[8]) + r = H(F(1) / sm) + wg = [H(H(e[i]) * r) for i in range(9)] + + def cl(v, n): + return 1 if v < 0 else (n - 2 if v >= n else v) + f = np.zeros(3, H) + k = 0 + for yy in (cl(y - 1, ey), cl(y, ey), cl(y + 1, ey)): + for xx in (cl(x - 1, ex), cl(x, ex), cl(x + 1, ex)): + T = bilinear88(t88, xx, yy) + for c in range(3): + f[c] = fma16(T[c], wg[k], f[c]) + k += 1 + g = H(F(1) / (np.exp2(F(a[9]) * F(-1.4426950216293335)) + F(1))) + om = H(H(1) - g) + o = np.zeros(3, H) + if 0 <= x < ex and 0 <= y < ey: + for c in range(3): + o[c] = fma16(t96[y, x, c], om, H(f[c] * g)) + if not np.isnan(t104[y, x, c]): + o[c] = H(o[c] + t104[y, x, c]) + out.append(o) + return np.array(out) + + +def extract(dump, weights, out): + z = np.load(dump) + W = np.frombuffer(weights, np.uint8) + + def B(p, wv=0): + return z[f'p{p}_w{wv}_B'] + + def C(p, wv=0): + return f16(z[f'p{p}_w{wv}_C']).astype(H) + + def table(off, n=64): + return W[off:off + 2 * n].view(H) + w = {'we': np.hstack([B(1, 0), B(1, 1)]), 'be': np.concatenate([C(1, 0)[0], C(1, 1)[0]]), + 'gain0': table(GAIN0), 'table7': table(TABLE7), 'gain14': table(GAIN14), 'table16': table(TABLE16)} + for h in range(2): + w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(2, h), B(3, h), B(4, h), B(5, h) + w[f'bias{h}'] = C(6, h) + w[f'wo{h}'] = B(8 + h) + for b in range(8): + w[f'w1_{b}'], w[f'b1_{b}'] = B(10 + 2 * b), C(10 + 2 * b)[0] + w[f'w2_{b}'] = B(11 + 2 * b) + w['wh'], w['bh'] = B(26), C(26)[0] + np.savez_compressed(out, **w) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('frame', type=int) + ap.add_argument('weights') + ap.add_argument('--extract', metavar='DUMP') + ap.add_argument('--blocks', nargs='*', default=['5,5']) + o = ap.parse_args() + from rr_capture import Capture + c = Capture(os.path.expanduser(o.capture)) + l = c.find('swin_dec1', o.frame) + if o.extract: + extract(o.extract, l.weights, o.weights) + return + ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32)) + sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16)) + ptr = {p.offset: p.delta for p in l.pointers} + pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8) + post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8) + lo = pre[ptr[8]:ptr[8] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32) + skip = pre[ptr[24]:ptr[24] + 2 * ex * ey * 32].reshape(2, ey, ex, 32) + outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32) + img = {i.offset: i for i in l.images} + + def surf(off): + i = img[off] + return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) + t88, t96, t104, r112 = surf(88), surf(96), surf(104), surf(112) + w = dict(np.load(o.weights)) + for blk in o.blocks: + bx, by = (int(v) for v in blk.split(',')) + x0, y0 = 8 * bx - sx, 8 * by - sy + tokens, head = window(w, lo, skip, x0, y0, ex, ey) + s112 = filter112(head, t88, t96, t104, x0, y0, ex, ey) + inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey] + tk = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8], outp[1, y0 + t // 8, x0 + t % 8]])) + for t in inside) + sf = sum(np.array_equal(s112[t], r112[y0 + t // 8, x0 + t % 8, :3]) for t in inside) + print(f'block {bx},{by}: tokens {tk}/{len(inside)}, surface 112 {sf}/{len(inside)}') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/dll_weights.py b/engine/rr/dll_weights.py new file mode 100644 index 0000000..2382e5e --- /dev/null +++ b/engine/rr/dll_weights.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +"""Read the Ray Reconstruction network's weights straight from the user's nvngx_dlssd.dll. + +usage: dll_weights.py DLL --list + dll_weights.py DLL --kernels + dll_weights.py DLL OUT_DIR (--model N | --capture CAPTURE_DIR) [--check EXTRACTED_DIR] [--any-kernels] + +The DLL holds every model's tensors as plain binary16 arrays in natural ([out][in]) order, one stage after the other +(encoder 0 ... decoder 0), with nothing naming them. They are found by structure alone: every attention head has a +relative position bias table of 100 x 100 values (the 10 x 10 window with its border) whose entries depend only on the +offset between the two positions, and a model is eleven consecutive stages whose head counts are 2 2 4 4 8 8 8 4 4 2 2. +The other tensors of a stage sit at fixed distances from its tables (LAYOUT below). + +OUT_DIR/STAGE.npz gets the tensors in the form engine/rr/{swin,enc0,dec0}_ref.py --extract produced from a ZLUDA +phase dump, which is what the packers (swin_pack.py, enc0_pack.py, dec0_pack.py) take: + FP8 matrices e4m3 codes, [k][n] with k in NVIDIA's fragment order (natural input channel inter(K)[k]) + gains times sqrt(32), as NGX uploads them + bias [64 queries][96 keys] in the kernels' key order, times 2 sqrt(2), as NGX uploads it +--model picks the N-th model of the DLL (--list); --capture picks the one a bridge capture of the denoiser ran, by +looking up each stage's vectors in the weight image NGX uploaded (and fails if the capture is of another DLL). +--check compares every tensor with an extracted one (EXTRACTED_DIR/STAGE/weights.npz). + +The native shaders were written against the kernels of one DLL version. Before anything is written the DLL's network +kernels (the PTX modules, engine/rr/native_kernels.txt: their SHA-256) are compared with those; another version is +refused unless --any-kernels is given. --kernels prints the digests of a DLL. + +Everything read and written is NVIDIA's: it stays on the user's machine and is never packaged. +""" +import argparse +import os +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import e4m3, q8 # noqa: E402 +from enc1_ref import inter # noqa: E402 +from dec0_ref import keys # noqa: E402 + +H = np.float16 +# stage: (kind, C, C_LO (decoders: the channels one level down) or C_M (encoders: the merged channels), heads) +STAGES = { + 'enc0': ('enc0', 32, 64, 2), 'enc1': ('enc', 64, 96, 2), 'enc2': ('enc', 96, 128, 4), 'enc3': ('enc', 128, 160, 4), + 'enc4': ('enc', 160, 160, 8), 'enc5': ('enc', 160, 0, 8), 'dec4': ('dec', 160, 160, 8), 'dec3': ('dec', 128, 160, 4), + 'dec2': ('dec', 96, 128, 4), 'dec1': ('dec', 64, 96, 2), 'dec0': ('dec0', 32, 64, 2), +} +ORDER = list(STAGES) +HEADS = [STAGES[s][3] for s in ORDER] +TABLE = 100 * 100 # one head's relative position bias, in values +# the 96 keys of a window (its 64 tokens, then the border without the corners) in the 10 x 10 table +POS = [(y + 1) * 10 + x + 1 for (x, y) in keys(0, 0)] +BIAS_SCALE = np.float32(2.0 * np.sqrt(2.0)) + + +def bias_tables(u): + """offsets (in values) of every relative position bias table in the binary16 view u of the DLL""" + # entry (i, j) equals entry (i + 1, j + 1) and (i + 10, j + 10): one step right / down for both positions + e1 = u[:-1111] == u[101:-1010] + e10 = u[:-1111] == u[1010:-101] + e11 = u[:-1111] == u[1111:] + first = e1[:-18] & e1[1:-17] & e1[2:-16] & e10[:-18] & e10[1:-17] & e11[:-18] & e11[1:-17] & (u[:len(e1) - 18] != u[1:len(e1) - 17]) + found = [] + dy, dx = np.divmod(np.arange(100), 10) + rel = ((dy[None, :] - dy[:, None] + 9) * 19 + dx[None, :] - dx[:, None] + 9).reshape(-1) + for p in np.nonzero(first)[0]: + if p + TABLE > len(u) or (found and p < found[-1] + TABLE): + continue + t = u[p:p + TABLE] + r = np.zeros(361, np.uint16) + r[rel] = t + if np.array_equal(r[rel], t) and len(np.unique(t)) > 64: + found.append(int(p)) + return found + + +def models(dll): + """[[bias offset (bytes) of each of the eleven stages] per model], in file order""" + u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16) + tables = bias_tables(u) + runs = [] # (first table, count) of consecutive tables + for p in tables: + if runs and p == runs[-1][0] + TABLE * runs[-1][1]: + runs[-1][1] += 1 + else: + runs.append([p, 1]) + out, i = [], 0 + while i + len(HEADS) <= len(runs): + if [n for _, n in runs[i:i + len(HEADS)]] == HEADS: + out.append([2 * p for p, _ in runs[i:i + len(HEADS)]]) + i += len(HEADS) + else: + i += 1 + return out + + +def kernel_digests(dll): + """{(kernel, SHA-256 of its PTX module)} of the DLL's network kernels""" + import hashlib + import re + sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'kernels', 'tools')) + from extract_dlss_ptx import fatbins, ptx_entries + out = set() + for _, fatbin in fatbins(dll): + for text in ptx_entries(fatbin): + for name in re.findall(rb'\.entry (cuda_dldn_engine_swin_\w+)\(', text): + out.add((name.decode(), hashlib.sha256(text).hexdigest())) + return out + + +def check_kernels(dll, path): + """the kernels the native shaders were written against must all be in the DLL, and no other version of them""" + known = {tuple(l.split()) for l in open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'native_kernels.txt')) + if l.strip() and not l.startswith('#')} + have = kernel_digests(dll) + missing, other = sorted(known - have), sorted(have - known) + if missing or other: + raise SystemExit(f'{path}: its network kernels are not the ones the native shaders were written against ' + f'({len(missing)} missing, {len(other)} unknown; e.g. {(missing + other)[0][0]}). The native ' + 'stages are unverified for this DLL version: use the translated model (engine/compile_rr.py), or ' + '--any-kernels to build anyway.') + + +def noise_table(dll): + """preset E's encoder 0 dithers its features with 8 x 16 x 16 16-bit values, the first 4096 bytes of the weight + image NGX uploads. In the DLL they are the array of 2048 values whose low five bits are all clear or all set.""" + u = np.frombuffer(dll[:len(dll) // 2 * 2], np.uint16) + low = u & 31 + ok = np.concatenate([[False], (low == 0) | (low == 31), [False]]) + edge = np.nonzero(ok[1:] != ok[:-1])[0] + found = [] + for a, b in zip(edge[::2], edge[1::2]): + if b - a >= 2048: + t = u[b - 2048:b] + if len(np.unique(t)) > 1024 and (t & 31 == 31).sum() > 256: + found.append(t.tobytes()) + return found[0] if len(set(found)) == 1 else None # some DLLs hold the same table twice + + +class Reader: + def __init__(self, dll, at): + self.dll, self.at = dll, at + + def back(self, *shape): + self.at -= 2 * int(np.prod(shape)) + return np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape) + + def take(self, *shape): + a = np.frombuffer(self.dll, H, int(np.prod(shape)), self.at).reshape(shape) + self.at += 2 * int(np.prod(shape)) + return a + + +def fp8(nat): + """a natural [out][in] binary16 matrix -> NVIDIA's B operand: codes [k][n], k in fragment order""" + k = nat.shape[1] + return np.ascontiguousarray(q8(nat)[:, inter(k) if k % 32 == 0 else np.arange(k)].T) + + +def gain(g, c): + return (g.astype(np.float32) * np.float32(np.sqrt(np.float32(c)))).astype(H) + + +def read_stage(dll, stage, bias_at): + """the stage's tensors (see the module text); returns (dict, first byte, one past the last byte)""" + kind, C, CX, NH = STAGES[stage] + w = {} + r = Reader(dll, bias_at) + table_a = r.back(C) + wo = r.back(C, 32 * NH) + wv = r.back(32 * NH, C) + wk = r.back(32 * NH, C) + gain_a = r.back(C) + if kind == 'enc0': # the embedding of the 32 input features + w['b2'] = r.back(32).copy() + w['w2'] = fp8(r.back(32, 32)) + w['b1'] = r.back(32).copy() + # [32][inputs]: preset E has 32 input features, preset D 30 (NGX pads the image's matrix with zeros). The + # matrix is the widest one that starts with a weight and holds FP8 values only + end = r.at + for inputs in range(32, 15, -1): + m = np.frombuffer(dll, H, 32 * inputs, end - 64 * inputs).reshape(32, inputs) + if m[0, :4].any() and np.array_equal(e4m3(q8(m)).astype(H).view(np.uint16), m.view(np.uint16)): + break + else: + raise SystemExit('enc0: no input embedding before its bias') + r.at = end - 64 * inputs + w['w1'] = fp8(np.hstack([m, np.zeros((32, 32 - inputs), H)])) + first = r.at + r = Reader(dll, bias_at) + bias = r.take(NH, 100, 100) + gain_b = r.take(C) + w1 = r.take(4 * C, C) + b1 = r.take(4 * C) + w2 = r.take(C, 4 * C) + table_b = r.take(C) + for h in range(NH): + w[f'wk{h}'] = fp8(wk[32 * h:32 * h + 32]) + w[f'wv{h}'] = fp8(wv[32 * h:32 * h + 32]) + w[f'wo{h}'] = fp8(wo[:, 32 * h:32 * h + 32]) + w[f'bias{h}'] = (bias[h][np.ix_(POS[:64], POS)].astype(np.float32) * BIAS_SCALE).astype(H) + for b in range(C // 8): + w[f'w1_{b}'] = fp8(w1[32 * b:32 * b + 32]) + w[f'b1_{b}'] = b1[32 * b:32 * b + 32].copy() + w[f'w2_{b}'] = fp8(w2[:, 32 * b:32 * b + 32]) + names = ('gain0', 'table0', 'gain1', 'table1') if kind in ('enc0', 'dec0') else ('gain0', 'table7', 'gain14', 'table16') + for name, a in zip(names, (gain(gain_a, C), table_a.copy(), gain(gain_b, C), table_b.copy())): + w[name] = a + if kind in ('enc', 'enc0') and CX: # the 2x2 merge + w['wm'] = fp8(r.take(CX, 4 * C)) + w['bm'] = r.take(CX).copy() + if kind in ('dec', 'dec0'): # the expand of the level below + w['we'] = fp8(r.take(4 * C, CX)) + w['be'] = r.take(4 * C).copy() + if kind != 'dec0': # the 16-channel head of the stage's filter + w['wh'] = fp8(r.take(16, C)) + w['bh'] = r.take(16).copy() + return w, first, r.at + + +def read_model(dll, bias_offsets): + """{stage: tensors}; checks that the stages lie back to back in the DLL, which the layout implies""" + out, end = {}, None + for stage, at in zip(ORDER, bias_offsets): + out[stage], first, last = read_stage(dll, stage, at) + if end is not None and first != end: + raise SystemExit(f'{stage}: starts at {first}, the stage before ends at {end}: not the layout this reader knows') + end = last + return out + + +def in_capture(capture, model): + """whether the capture's launches uploaded this model: every stage's vectors must be in its weight image""" + from rr_capture import Capture + cap = Capture(os.path.expanduser(capture)) + for stage, w in model.items(): + image = next(l for l in cap.launches if f'swin_{stage}' in l.kernel).weights + for name, a in w.items(): + if a.dtype == H and a.ndim == 1 and bytes(image).find(a.tobytes()) < 0: + return False + return True + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument('dll') + ap.add_argument('output', nargs='?') + ap.add_argument('--list', action='store_true') + ap.add_argument('--model', type=int) + ap.add_argument('--capture') + ap.add_argument('--check', metavar='EXTRACTED_DIR') + ap.add_argument('--kernels', action='store_true') + ap.add_argument('--any-kernels', action='store_true') + o = ap.parse_args() + dll = Path(o.dll).read_bytes() + if o.kernels: + for name, digest in sorted(kernel_digests(dll)): + print(name, digest) + return + if o.output and not o.any_kernels: + check_kernels(dll, o.dll) + found = models(dll) + if not found: + raise SystemExit(f'{o.dll}: no Ray Reconstruction model with the layout this reader knows') + if o.list or not o.output: + for i, m in enumerate(found): + print(f'model {i}: bias tables at ' + ' '.join(map(str, m))) + return + if o.model is not None: + model = read_model(dll, found[o.model]) + index = o.model + else: + if not o.capture: + raise SystemExit('--model N or --capture CAPTURE_DIR') + index, model = next(((i, m) for i, m in ((i, read_model(dll, f)) for i, f in enumerate(found)) + if in_capture(o.capture, m)), (None, None)) + if model is None: + raise SystemExit(f'none of the {len(found)} models in {o.dll} is the one the capture ran: the capture was taken ' + 'with another nvngx_dlssd.dll') + out = Path(o.output) + out.mkdir(parents=True, exist_ok=True) + for stage, w in model.items(): + np.savez_compressed(out / f'{stage}.npz', **w) + noise = noise_table(dll) + if noise is not None: + if o.capture: # preset E's image starts with it; preset D's has none + from rr_capture import Capture + image = bytes(next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).weights) + if 'diamond' not in next(l for l in Capture(os.path.expanduser(o.capture)).launches if 'swin_enc0' in l.kernel).kernel \ + and image[:4096] != noise: + raise SystemExit('the noise table found in the DLL is not the one the capture uploaded') + (out / 'noise.bin').write_bytes(noise) + print(f'{o.dll}: model {index} of {len(found)}, {len(model)} stages in {out}') + if o.check: + bad = 0 + for stage, w in model.items(): + ref = dict(np.load(Path(o.check) / stage / 'weights.npz')) + for name in sorted(set(ref) | set(w)): + a, b = w.get(name), ref.get(name) + if a is None or b is None: + print(f'{stage}.{name}: only in the {"extraction" if a is None else "DLL reading"}') + bad += 1 + continue + if name in ('we', 'be'): # a multi-warp extraction repeats the columns + b = b[..., :a.shape[-1]] + same = a.shape == b.shape and np.array_equal(a.view(np.uint8) if a.dtype == np.uint8 else a.astype(H).view(np.uint16), + b.view(np.uint8) if b.dtype == np.uint8 else b.astype(H).view(np.uint16)) + if not same: + n = int((a != b).sum()) if a.shape == b.shape else -1 + print(f'{stage}.{name}: differs ({a.shape} vs {b.shape}, {n} values)') + bad += 1 + print('check: every tensor equals the extraction' if not bad else f'check: {bad} tensors differ') + sys.exit(1 if bad else 0) + + +if __name__ == '__main__': + main() diff --git a/engine/rr/enc0_m_decl.glsl b/engine/rr/enc0_m_decl.glsl new file mode 100644 index 0000000..b338440 --- /dev/null +++ b/engine/rr/enc0_m_decl.glsl @@ -0,0 +1,53 @@ +// Native part of engine/rr/enc0_m.comp (generated by engine/rr/gen_enc0_m.py): declarations, after the translated +// header. Binding 1 holds NGX's weight image followed by engine/rr/enc0_pack.py's blob (offsets in enc0_m_w.h already +// include its base); binding 2 is the arena. +#extension GL_EXT_shader_8bit_storage : require +#extension GL_KHR_shader_subgroup_vote : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +#include "enc0_m_w.h" +#include "rr_exact.glsl" +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) Q8PAIRS(dst, src) +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; +shared uint KVS[2 * 768]; // the current head's K, then V^T; before the attention the key pixels' features +#define XF KVS // the key pixels' 32 FP8 features (plane 0 | plane 1), slot-major +shared uint XT[96 * 8]; // the normalised key tokens +#define RR_KS KVS +#define RR_KS0 0u +#define RR_VS KVS +#define RR_VS0 768u +#define MG XT // the merge rows (16 x 128 FP8) reuse XT +shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums +ivec2 keyXY(uint t, int x0, int y0) +{ + if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u)); + const int i = int(t & 7u), side = int((t - 64u) >> 3u); + return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i)); +} +uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d) +{ + return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u); +} +uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); } +// the translated prologue's shared buffer sh: input_noise (10 x 10 pixels, 4 f16) at byte 0, input_tensor (two planes +// of 10 x 10 pixels x 16 FP8 features) at byte 800 +float16_t noiseAt(uint gx, uint gy, uint c) +{ + const uint b = 80u * gy + 8u * gx + 2u * c; + return hb((sh[b >> 2] >> (8u * (b & 2u))) & 0xffffu); +} diff --git a/engine/rr/enc0_m_tail.glsl b/engine/rr/enc0_m_tail.glsl new file mode 100644 index 0000000..043f9a3 --- /dev/null +++ b/engine/rr/enc0_m_tail.glsl @@ -0,0 +1,263 @@ + // ======================================================================== native part (engine/rr/enc0_ref.py) + memoryBarrierShared(); barrier(); // wave 0's translated prologue filled sh (features, noise) and wrote 304 / 1016 + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(param.w[0]), ey = int(param.w[1]); + const uint outOff = param.w[12], mrgOff = param.w[14], shift = param.w[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const uint lplane = uint((ex / 2) * (ey / 2)) * 32u; + // the key slots' features, from the 10 x 10 grid (origin x0 - 1, y0 - 1) + if (tid < 96u) { + const ivec2 k = keyXY(inter(tid), x0, y0); + const uint g = uint((k.y - y0 + 1) * 10 + (k.x - x0 + 1)); + [[unroll]] for (uint wd = 0u; wd < 4u; ++wd) { + XF[tid * 8u + wd] = sh[(800u + g * 16u) / 4u + wd]; + XF[tid * 8u + 4u + wd] = sh[(2400u + g * 16u) / 4u + wd]; + } + } + barrier(); + // the embedding x = W2 q8(relu(W1 f + b1)) + b2 for token tiles n = wv (this wave's queries: x is also their + // residual) and n = wv + 4 (the border keys, waves 0 and 1), then each key's norm in registers (dec0_ref.norm_plane + // on the accumulator layout, as dec0_m.comp's MLP norm) to XT + ACC R[2]; + for (uint n = wv; n < 6u; n += 4u) { + B8 fb0, fb1, hb_[2]; + coopMatLoad(fb0, XF, n * 128u, 8u, CM); + coopMatLoad(fb1, XF, n * 128u + 4u, 8u, CM); + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]); + A8 a0, a1; + LDA(a0, W_W1, 32u, m, 0u); LDA(a1, W_W1, 32u, m, 1u); + K32(a, a0, fb0, a1, fb1) + [[unroll]] for (int i = 0; i < 8; ++i) hb_[m][i] = q8s(max(a[i], 0.0)); + } + float16_t xv[2][8]; + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_B2 / 2u + 16u * m + 8u * hl + uint(i)]); + A8 a0, a1; + LDA(a0, W_W2, 32u, m, 0u); LDA(a1, W_W2, 32u, m, 1u); + K32(a, a0, hb_[0], a1, hb_[1]) + [[unroll]] for (int i = 0; i < 8; ++i) xv[m][i] = float16_t(a[i]); + if (n == wv) { // the residual: x + table 0 + [[unroll]] for (int i = 0; i < 8; ++i) { + precise float16_t c = xv[m][i] + w16.w[W_TABLE0 / 2u + 16u * m + 8u * hl + uint(i)]; + R[m][i] = float(c); + } + } + } + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t p0 = xv[0][i0] * xv[0][i0] + xv[0][i1] * xv[0][i1]; + precise float16_t p1 = xv[1][i0] * xv[1][i0] + xv[1][i1] * xv[1][i1]; + precise float16_t hv = p0 + p1; + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + B8 xq; + [[unroll]] for (int i = 0; i < 8; ++i) { + precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(16u * m + 8u * hl + uint(i))]; + precise float16_t nn = xv[m][i] * g; + xq[i] = q8h(nn); + } + coopMatStore(xq, XT, n * 128u + 4u * m, 8u, CM); + } + } + barrier(); // XT complete; XF is read (the attention reuses its memory) + // two heads + [[unroll]] for (uint hd = 0u; hd < 2u; ++hd) { + if (hd == 1u) barrier(); + for (uint j = 0u; j < 3u; ++j) { + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + A8 wk[2], xa[2]; + B8 wvb[2], xb[2]; + [[unroll]] for (uint kt = 0u; kt < 2u; ++kt) { + coopMatLoad(xb[kt], XT, n * 128u + 4u * kt, 8u, CM); + coopMatLoad(xa[kt], XT, n * 128u + 4u * kt, 8u, RM); + LDA(wk[kt], W_WK0 + HEAD * hd, 32u, c, kt); + LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); + } + ACC a = ACC(0.0); + K32(a, wk[0], xb[0], wk[1], xb[1]) + A8 ka; + Q8(ka, a) + coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM); + ACC vv = ACC(0.0); + K32(vv, xa[0], wvb[0], xa[1], wvb[1]) + A8 va; + Q8(va, vv) + coopMatStore(va, KVS, 768u + c * 384u + 4u * n, 24u, RM); + } + barrier(); +#include "rr_attn.glsl" + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(R[m], a0, Ob[0], a1, Ob[1]) + } + } + // MLP (dec0_m.comp's) + B8 Nb[2]; + { + float16_t rv[2][8]; + [[unroll]] for (uint m = 0u; m < 2u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]); + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t p0 = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1]; + precise float16_t p1 = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1]; + precise float16_t hv = p0 + p1; + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < 2u; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint k = 16u * m + 8u * hl + uint(i); + precise float16_t g = r * w16.w[W_GAIN1 / 2u + k]; + precise float16_t nn = rv[m][i] * g; + Nb[m][i] = q8h(nn); + precise float16_t a = rv[m][i] + w16.w[W_TABLE1 / 2u + k]; + R[m][i] = float(a); + } + } + for (uint br = 0u; br < 4u; ++br) { + const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + A8 a0, a1; + LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u); + K32(h, a0, Nb[0], a1, Nb[1]) + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); + } + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(R[m], a0, Hb[0], a1, Hb[1]) + } + } + // outputs: the tokens (one plane, dec0's skip) and the merge rows + const uint t = inter(16u * wv + col); + const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + B8 Tb[2]; + { + const uint ty2 = t >> 3u, tx2 = t & 7u; + const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u; + const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u); + const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u; + barrier(); // XT is read; the merge rows reuse it + [[unroll]] for (uint m = 0u; m < 2u; ++m) { + Q8(Tb[m], R[m]) + const uint q0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), q1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]); + const uint at = 4u * m + 2u * hl; + if (inside) { arena.w[o + at] = q0; arena.w[o + at + 1u] = q1; } + const uint gm = mr * 32u + blk * 8u + 4u * m + 2u * hl; + MG[gm] = q0; MG[gm + 1u] = q1; + } + } + { // the 32 -> 16 head and encoder 0's filter of the noise grid (two f16 fma chains, taps even / odd) + ACC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + A8 a0, a1; + LDA(a0, W_WH, 32u, 0u, 0u); LDA(a1, W_WH, 32u, 0u, 1u); + K32(hv, a0, Tb[0], a1, Tb[1]) + float16_t l[12]; + [[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]); + const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u); + l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u); + const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]); + const float16_t mx = max(max(ma, mb), l[8]); + float e[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t rr = rcp16(s); + float16_t wg[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr)); + const uint gx = t & 7u, gy = t >> 3u; // the token's 3 x 3 taps in the grid: (gx .. gx + 2, gy .. gy + 2) + float16_t o112[3], o120[3]; + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + float16_t lo = 0.0hf, hi = 0.0hf; + [[unroll]] for (uint k = 0u; k < 8u; k += 2u) { + lo = fma16(noiseAt(gx + k % 3u, gy + k / 3u, c), wg[k], lo); + hi = fma16(noiseAt(gx + (k + 1u) % 3u, gy + (k + 1u) / 3u, c), wg[k + 1u], hi); + } + precise float16_t sum = lo + hi; + const float16_t f = fma16(noiseAt(gx + 2u, gy + 2u, c), wg[8], sum); + precise float16_t a = sig(l[9u + c]) * f; + precise float16_t b = f - a; + o120[c] = a; o112[c] = b; + } + if (hl == 0u) { + if (inside) { + imageStore(S112, ivec2(tx, ty), vec4(o112[0], o112[1], o112[2], 0.0)); + imageStore(S120, ivec2(tx, ty), vec4(o120[0], o120[1], o120[2], 0.0)); + } + [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = o112[c]; + } + } + barrier(); + { // the 2x2 merge (128 -> 64): wave w computes channel tile w + const uvec2 mxy = mergeXY(col); + const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2; + const uint m = wv; + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint ks = 0u; ks < 4u; ++ks) { + B8 g0, g1; A8 w0, w1; + coopMatLoad(g0, MG, 8u * ks, 32u, CM); + coopMatLoad(g1, MG, 8u * ks + 4u, 32u, CM); + LDA(w0, W_WM, 128u, m, 2u * ks); LDA(w1, W_WM, 128u, m, 2u * ks + 1u); + K32(a, w0, g0, w1, g1) + } + B8 q; + Q8(q, a) + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + const uint at = (mrgOff + (m >> 1u) * lplane + uint(Y * hx + X) * 32u) / 4u + 4u * (m & 1u) + 2u * hl; + arena.w[at] = pack4(q[0], q[1], q[2], q[3]); + arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]); + } + } + if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards + const uint mx = lane & 3u, my = lane >> 2u; + const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my); + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint t00 = (2u * my) * 8u + 2u * mx; + vec4 out_ = vec4(0.0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25); + precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25); + precise float16_t b0 = a00 + a01, b1 = a10 + a11; + precise float16_t cc = b0 + b1; + out_[c] = float(cc); + } + imageStore(S128, ivec2(x >> 1, y >> 1), out_); + } + } diff --git a/engine/rr/enc0_pack.py b/engine/rr/enc0_pack.py new file mode 100644 index 0000000..398e3a4 --- /dev/null +++ b/engine/rr/enc0_pack.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Pack encoder 0's extracted weights (engine/rr/enc0_ref.py --extract) for engine/rr/enc0_m.comp's native part. + +usage: enc0_pack.py WEIGHTS.npz OUT.bin OUT.h [BASE] + +dec0_pack.py's layout (C = 32) plus the embedding (W1, W2: 32 -> 32 -> 32, columns in slot order), the 2x2 merge +(128 -> 64, columns in arena order) and the 32 -> 16 head. BASE (default 0) is added to every offset in OUT.h: the +shader binds NGX's weight image first (its translated feature prologue reads it) and this blob after it. +""" +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from enc1_ref import inter # noqa: E402 +from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402 + +I32, I64, I96 = inter(32), inter(64), inter(96) + + +def main(): + src, out_bin, out_h = sys.argv[1:4] + base = int(sys.argv[4]) if len(sys.argv) > 4 else 0 + w = dict(np.load(src)) + e = {'w1': ('A', w['w1'][:, I32]), 'b1': ('H', w['b1'][I32]), 'w2': ('A', w['w2'][:, I32]), 'b2': ('H', w['b2'][I32])} + for h in range(2): + e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32]) + e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32]) + e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I32]) + e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T) + for b in range(4): + e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32]) + e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32]) + e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I32]) + e['wm'], e['bm'] = ('A', w['wm'][:, I64]), ('H', w['bm'][I64]) + e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh']) + e['gain0'] = ('H', w['gain0']) + for k in ('table0', 'gain1', 'table1'): + e[k] = ('H', w[k][I32]) + e['cubic'] = ('U', cubic_table()) + e['act'] = ('Y', act_table()) + blob, defines = bytearray(), [] + for name, (kind, a) in e.items(): + while len(blob) % 256: + blob.append(0) + off = len(blob) + if kind == 'A': + blob += a_tiles(a).tobytes() + elif kind == 'B': + blob += b_tiles(a).tobytes() + elif kind == 'Y': + blob += np.ascontiguousarray(a, np.uint8).tobytes() + elif kind == 'U': + blob += np.ascontiguousarray(a, np.uint16).tobytes() + else: + blob += np.ascontiguousarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {base + off}u // {kind} {a.shape}') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write('// generated by engine/rr/enc0_pack.py\n' + '\n'.join(defines) + '\n') + print(f'{len(blob)} bytes, {len(e)} entries, base {base}') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/enc0_ref.py b/engine/rr/enc0_ref.py new file mode 100644 index 0000000..ffa4872 --- /dev/null +++ b/engine/rr/enc0_ref.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Exact numpy reference of the denoiser's encoder 0 (cuda_dldn_engine_swin_enc0_kernel) after its input features, +one 8x8 window at a time. The 32 FP8 features of the 96 key pixels (NVIDIA's per-pixel gather from the game's guides +and the temporal state, which engine/rr/enc0_m.comp keeps as ptx2rr.py's translation) are an input here. + + embedding F (96 x 32) -> q8(relu(F W1 + b1)) -> W2 + b2: the key tokens x (32 channels, f16) + block dec0_ref's: RMS norm (gain W+6272), residual x + table W+37056, two heads, four MLP branches + outputs the tokens (argument 48, dec0's skip), the 2x2 merge 128 -> 64 (argument 56, enc1's input) and the + 32 -> 16 head with enc1's filter to surfaces 112, 120 and 128 + +usage: enc0_ref.py CAPTURE FRAME WEIGHTS.npz --dump DUMP.npz [--block X,Y] (verifies the dump's block) + enc0_ref.py CAPTURE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump) +""" +import argparse +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import f16, q8 # noqa: E402 +from enc1_ref import inter, mm, softmax_cubic, act, add, filter_head, MERGE_ORDER, H # noqa: E402 +from dec0_ref import norm_plane # noqa: E402 + +I32, I64, I96 = inter(32), inter(64), inter(96) +GAIN0, TABLE0, GAIN1, TABLE1 = 6272, 37056, None, None + + +def block(w, feat): + """feat: the 96 key pixels' features (codes, NVIDIA's k order). Returns (tokens 64 x 32 codes in arena order, + merged 16 x 64 codes in arena order, head 64 x 16 f16, dbg)""" + dbg = {} + e1 = mm(feat, w['w1'], np.repeat(w['b1'][None, :], 96, 0)) + x = mm(q8(np.maximum(e1, H(0)))[:, I32], w['w2'], np.repeat(w['b2'][None, :], 96, 0)) + dbg['x'] = x + X = np.array([norm_plane(x[k], w['gain0']) for k in range(96)]) + R = None + for h in range(2): + qk = mm(X, w[f'wk{h}'], np.zeros((96, 32))) + v = mm(X, w[f'wv{h}'], np.zeros((96, 32))) + Q = q8(qk)[:, I32] + P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}'])) + O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32))) + C = add(x[:64], w['table0'][None, :]) if h == 0 else R + R = mm(q8(O)[:, I32], w[f'wo{h}'], C) + N = np.array([norm_plane(R[t], w['gain1']) for t in range(64)]) + acc = add(R, w['table1'][None, :]) + for b in range(4): + hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0))) + acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc) + tokens = q8(acc)[:, I32] + Mg = np.zeros((16, 128), np.uint8) + for r, (mx, my) in enumerate(MERGE_ORDER): + for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]): + Mg[r, 32 * blk:32 * blk + 32] = tokens[(2 * my + dy) * 8 + 2 * mx + dx] + merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, I64] + head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0)) + return tokens, merged, head, dbg + + +def extract(dump, weights, out): + z = np.load(dump) + W = np.frombuffer(weights, np.uint8) + + def B(p): + return z[f'p{p}_w0_B'] + + def C(p): + return f16(z[f'p{p}_w0_C']).astype(H) + w = {'w1': B(1), 'b1': C(1)[0], 'w2': B(2), 'b2': C(2)[0], 'gain0': W[GAIN0:GAIN0 + 64].view(H), + 'table0': W[TABLE0:TABLE0 + 64].view(H)} + # the input norm's gain and the residual table sit elsewhere in preset D's (smaller) weight image: when the preset + # E offsets do not reproduce the dump, find them by matching it, like the MLP's below + x = mm(q8(np.maximum(mm(z['p1_w0_A'], w['w1'], np.repeat(w['b1'][None, :], 96, 0)), H(0)))[:, I32], w['w2'], + np.repeat(w['b2'][None, :], 96, 0)) + cands = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)] + if not all(np.array_equal(norm_plane(x[k], w['gain0']), z['p3_w0_A'][k]) for k in (0, 9)): + w['gain0'] = next(t for t in cands if all(np.array_equal(norm_plane(x[k], t), z['p3_w0_A'][k]) for k in (0, 9))) + if not np.array_equal(add(x[:2], w['table0'][None]).view(np.uint16), C(7)[:2].view(np.uint16)): + w['table0'] = next(t for t in cands if np.array_equal(add(x[:2], t[None]).view(np.uint16), C(7)[:2].view(np.uint16))) + for h, base in ((0, 3), (1, 8)): + w[f'wk{h}'], w[f'wv{h}'] = B(base), B(base + 1) + w[f'bias{h}'] = C(base + 2) + w[f'wo{h}'] = B(base + 4) + for b in range(4): + w[f'w1_{b}'], w[f'b1_{b}'] = B(13 + 2 * b), C(13 + 2 * b)[0] + w[f'w2_{b}'] = B(14 + 2 * b) + w['wm'], w['bm'] = B(21), C(21)[0] + w['wh'], w['bh'] = B(22), C(22)[0] + # the MLP norm's gain and table, found by matching the dump + R = f16(z['p12_w0_D']).astype(H) + tables = [W[o:o + 64].view(H) for o in range(0, len(W) - 63, 2)] + w['gain1'] = next(t for t in tables if all(np.array_equal(norm_plane(R[k], t), z['p13_w0_A'][k]) for k in (0, 9))) + w['table1'] = next(t for t in tables if np.array_equal(add(R[:2], t[None]).view(np.uint16), C(14)[:2].view(np.uint16))) + np.savez_compressed(out, **w) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('frame', type=int) + ap.add_argument('weights') + ap.add_argument('--extract', metavar='DUMP') + ap.add_argument('--dump') + ap.add_argument('--block', default='5,5') + o = ap.parse_args() + from rr_capture import Capture + c = Capture(os.path.expanduser(o.capture)) + l = c.find('swin_enc0', o.frame) + if o.extract: + extract(o.extract, l.weights, o.weights) + return + z = np.load(o.dump) + w = dict(np.load(o.weights)) + ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32)) + sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16)) + bx, by = (int(v) for v in o.block.split(',')) + x0, y0 = 8 * bx - sx, 8 * by - sy + tokens, merged, head, dbg = block(w, z['p1_w0_A']) + for p, name in ((3, 'X'), ): + pass + ptr = {p.offset: p.delta for p in l.pointers} + post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8) + otok = post[ptr[48]:ptr[48] + ex * ey * 32].reshape(ey, ex, 32) + omrg = post[ptr[56]:ptr[56] + 2 * (ex // 2) * (ey // 2) * 32].reshape(2, ey // 2, ex // 2, 32) + img = {i.offset: i for i in l.images} + tex = {k: np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) for k, i in img.items() + if k in (112, 120, 160)} + inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey] + tk = sum(np.array_equal(tokens[t], otok[y0 + t // 8, x0 + t % 8]) for t in inside) + mg = sum(np.array_equal(merged[r], np.concatenate([omrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(2)])) + for r, (mx, my) in enumerate(MERGE_ORDER)) + print(f'block {bx},{by}: tokens {tk}/{len(inside)}, merged {mg}/16') + return head, x0, y0, ex, ey + + +if __name__ == '__main__': + main() diff --git a/engine/rr/enc1_m.comp b/engine/rr/enc1_m.comp new file mode 100644 index 0000000..abfbf54 --- /dev/null +++ b/engine/rr/enc1_m.comp @@ -0,0 +1,357 @@ +#version 460 +// Native encoder 1 of the DLSS-D (preset E) denoiser, register-resident: cuda_dldn_engine_swin_enc1_kernel +// computed like engine/layer_m.comp. The arithmetic is engine/rr/enc1_ref.py's (bit-exact to NVIDIA); the stage +// contract is the translated stage's, except binding 1, which holds engine/rr/enc1_pack.py's packing. +// +// Every product runs transposed, D^T = W^T X^T: weights are the A operand, activations the B operand, and a result's +// rows are channels, its columns tokens. On RADV gfx1201 an accumulator's element i in lane l is (row 8 * (l >> 4) +// + i, column l & 15), a B operand's the same and an A operand's the transpose (layer_m.comp): so a result is the +// next product's B operand element for element, and a result tile is also the A operand of its own transpose (the +// keys of the scores, the V^T of P V). enc1_pack.py orders every weight's columns so that the rows come out in the +// k order NVIDIA's next product contracts them in; keys and queries are token slots (slot s = token inter(s)). +// Accumulators stay f32 and are rounded to binary16 values in place (r16) after each logical k32 step. +// +// One workgroup per 8x8-token window, four waves; wave w owns query slots 16w..16w+15 and computes the window's +// keys and values itself (no barrier between the prologue and the outputs). Shared memory holds only the +// normalised tokens, the raw tokens and the merge rows. +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_KHR_shader_subgroup_vote : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +layout(local_size_x = 128) in; +#include "enc1_m_w.h" +#include "rr_exact.glsl" + +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +layout(set = 0, binding = 3) uniform sampler2D T88; +layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112; +layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128; + +#ifndef SKIP +#define SKIP 0 +#endif +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +// W^T tile (mt, kt) of a weight with K rows (enc1_pack.a_tiles); W tile (kt, nt) of one with N columns (b_tiles) +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +// binary16 rounding in place: a real conversion, kept by OR-ing in a zero the compiler cannot see (RADV folds a +// plain float(float16_t(x)) back to x); about 2 instructions per element against r16's 6 +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = q8s((src)[i_]); +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; + +shared uint XT[96 * 16]; // the normalised key tokens: slot-major, 64 FP8 channels in arena (= k) order +// the merge rows (16 x 256 FP8, enc1_ref.MERGE_ORDER, 4 tokens of 64 channels) reuse XT, dead after the projections +#define MG XT +shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums +shared uint KS[96 * 8]; // the current head's K, token slot x 32 FP8 channels (k order) +shared uint VS[32 * 24]; // the current head's V^T, 32 channels x 96 token slots + +uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); } +uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d) +{ + return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u); +} + +void main() +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(param.p[0]), ey = int(param.p[1]); + const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const uint plane = uint(ex * ey) * 32u; + + // ---------------------------------------------------------------- prologue: one key token per thread + if ((SKIP & 16) == 0 && tid < 96u) { + const uint t = inter(tid); + int x, y; + if (t < 64u) { + x = x0 + int(t & 7u); y = y0 + int(t >> 3u); + } else { + const int i = int(t & 7u), side = int((t - 64u) >> 3u); + x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8); + y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i); + } + x = mirror(x, ex); y = mirror(y, ey); + const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u; + uint raw[16]; + float16_t v[64]; + [[unroll]] for (uint wd = 0u; wd < 16u; ++wd) { + raw[wd] = arena.w[base + (wd >> 3u) * (plane >> 2u) + (wd & 7u)]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) + v[4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((raw[wd] >> (8u * b)) & 0xffu))); + } + // NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves + precise f16vec2 h[4]; + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + precise f16vec2 p[4]; + [[unroll]] for (uint k = 0u; k < 4u; ++k) { + const uint o = 32u * (k & 1u) + 4u * (q + 4u * (k >> 1u)); + precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]); + precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]); + p[k] = a + c; + } + precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3]; + h[q] = s0 + s1; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t s = hs.x + hs.y; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint wd = 0u; wd < 16u; ++wd) { + floate4m3_t c[4]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + const uint j = 4u * wd + b; + precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)]; + precise float16_t n = v[j] * g; + c[b] = q8h(n); + } + XT[tid * 16u + wd] = pack4(c[0], c[1], c[2], c[3]); + } + } + barrier(); + + // ---------------------------------------------------------------- the residual: raw tokens + table 7 + // (each lane's token, mirrored like the prologue's; its channels 16m + 8hl .. + 7 are 8 bytes of plane m >> 1) + ACC R[4]; + { + const uint tq = inter(16u * wv + col); + const int qx = mirror(x0 + int(tq & 7u), ex), qy = mirror(y0 + int(tq >> 3u), ey); + const uint qb = (inOff + uint(qy * ex + qx) * 32u) >> 2u; + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + const uint at = qb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl; + const uvec2 u = uvec2(arena.w[at], arena.w[at + 1u]); + [[unroll]] for (int i = 0; i < 8; ++i) { + const float16_t x = float16_t(uintBitsToFloate4m3EXT(uint8_t((u[i >> 2] >> (8u * (uint(i) & 3u))) & 0xffu))); + precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)]; + R[m][i] = float(c); + } + } + } + + // ---------------------------------------------------------------- attention, two heads + [[unroll]] for (uint hd = 0u; hd < ((SKIP & 64) != 0 ? 0u : 2u); ++hd) { + // this head's Q/K and V, shared: 12 jobs (token tile n, channel tile c), three per wave. K = (QK^T)^T rows + // to KS (token slot x 32 channels, also Q^T read column-major), V^T to VS (32 channels x 96 token slots) + if (hd == 1u) barrier(); // every wave is done with head 0's KS / VS + for (uint j = 0u; j < ((SKIP & 128) != 0 ? 0u : 3u); ++j) { + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + A8 wk[4]; + B8 wvb[4], xb[4]; + A8 xa[4]; + [[unroll]] for (uint kt = 0u; kt < 4u; ++kt) { + coopMatLoad(xb[kt], XT, n * 256u + 4u * kt, 16u, CM); + coopMatLoad(xa[kt], XT, n * 256u + 4u * kt, 16u, RM); + LDA(wk[kt], W_WK0 + HEAD * hd, 64u, c, kt); + LDB(wvb[kt], W_WV0 + HEAD * hd, 32u, kt, c); + } + ACC a = ACC(0.0); + K32(a, wk[0], xb[0], wk[1], xb[1]) + K32(a, wk[2], xb[2], wk[3], xb[3]) + A8 ka; + Q8(ka, a) + coopMatStore(ka, KS, n * 128u + 4u * c, 8u, RM); + ACC v = ACC(0.0); + K32(v, xa[0], wvb[0], xa[1], wvb[1]) + K32(v, xa[2], wvb[2], xa[3], wvb[3]) + A8 va; + Q8(va, v) + coopMatStore(va, VS, c * 384u + 4u * n, 24u, RM); + } + barrier(); +#include "rr_attn.glsl" + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(R[m], a0, Ob[0], a1, Ob[1]) + } + } + + // ---------------------------------------------------------------- MLP + B8 Nb[4]; + { + float16_t rv[4][8]; + [[unroll]] for (uint m = 0u; m < 4u; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]); + // slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + 2((i >> 1) & 1), + // t = 2hl + (i >> 2), c = i & 1 (enc1_ref.nv_rowsum_sq's tree) + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t a = rv[0][i0] * rv[0][i0] + rv[0][i1] * rv[0][i1]; + precise float16_t b = rv[1][i0] * rv[1][i0] + rv[1][i1] * rv[1][i1]; + precise float16_t cc = rv[2][i0] * rv[2][i0] + rv[2][i1] * rv[2][i1]; + precise float16_t d = rv[3][i0] * rv[3][i0] + rv[3][i1] * rv[3][i1]; + precise float16_t hv = (a + cc) + (b + d); + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < 4u; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint k = 16u * m + 8u * hl + uint(i); + precise float16_t g = r * w16.w[W_GAIN14 / 2u + k]; + precise float16_t n = rv[m][i] * g; + Nb[m][i] = q8h(n); + precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k]; + R[m][i] = float(a); + } + } + for (uint br = 0u; br < ((SKIP & 4) != 0 ? 1u : 8u); ++br) { + const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + A8 a0, a1, a2, a3; + LDA(a0, w1, 64u, c, 0u); LDA(a1, w1, 64u, c, 1u); LDA(a2, w1, 64u, c, 2u); LDA(a3, w1, 64u, c, 3u); + K32(h, a0, Nb[0], a1, Nb[1]) + K32(h, a2, Nb[2], a3, Nb[3]) + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = q8h(act(float16_t(h[i]))); + } + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(R[m], a0, Hb[0], a1, Hb[1]) + } + } + + // ---------------------------------------------------------------- outputs + const uint t = inter(16u * wv + col); + const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + B8 Tb[4]; + { + const uint ty2 = t >> 3u, tx2 = t & 7u; + const uint r = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u; + const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u); + const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u; + [[unroll]] for (uint m = 0u; m < 4u; ++m) { + Q8(Tb[m], R[m]) + const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]); + const uint at = 4u * (m & 1u) + 2u * hl; + if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; } + const uint g = r * 64u + 16u * blk + 4u * m + 2u * hl; + MG[g] = w0; MG[g + 1u] = w1; + } + } + if ((SKIP & 32) == 0) { // the 64 -> 16 head and its kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15) + ACC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + A8 a0, a1, a2, a3; + LDA(a0, W_WH, 64u, 0u, 0u); LDA(a1, W_WH, 64u, 0u, 1u); LDA(a2, W_WH, 64u, 0u, 2u); LDA(a3, W_WH, 64u, 0u, 3u); + K32(hv, a0, Tb[0], a1, Tb[1]) + K32(hv, a2, Tb[2], a3, Tb[3]) + float16_t l[12]; + [[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]); + const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u); + l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u); + const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]); + const float16_t mx = max(max(ma, mb), l[8]); + float e[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t rr = rcp16(s); + float wg[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr)); + const int x = tx, y = ty; + const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1; + const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1; + const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc), + ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp)); + float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0)); + [[unroll]] for (uint k = 0u; k < 8u; ++k) { + const vec4 tv = texelFetch(T88, pts[k], 0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); } + } + const vec4 t8 = texelFetch(T88, pts[8], 0); + float o112[3], o120[3]; + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + precise float f0 = acc[c][0] + acc[c][1]; + precise float f1 = t8[c] * wg[8] + r16(f0); + const float f = r16(f1); + precise float a = float(sig(l[9u + c])) * f; + const float ar = r16(a); + precise float b = f - ar; + o120[c] = ar; o112[c] = r16(b); + } + if (hl == 0u) { + if (inside) { + imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0)); + imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0)); + } + [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]); + } + } + barrier(); + if ((SKIP & 8) == 0 && wv < 3u) { // the 2x2 merge: merged^T (96 x 16 rows), channel tiles 2w and 2w + 1 = arena plane w + const uvec2 mxy = mergeXY(col); + const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y), hx = ex / 2, hy = ey / 2; + const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u; + [[unroll]] for (uint m2 = 0u; m2 < 2u; ++m2) { + const uint m = 2u * wv + m2; + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + for (uint ks = 0u; ks < 8u; ++ks) { + B8 g0, g1; A8 w0, w1; + coopMatLoad(g0, MG, 8u * ks, 64u, CM); + coopMatLoad(g1, MG, 8u * ks + 4u, 64u, CM); + LDA(w0, W_WM, 256u, m, 2u * ks); LDA(w1, W_WM, 256u, m, 2u * ks + 1u); + K32(a, w0, g0, w1, g1) + } + B8 q; + Q8(q, a) + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + arena.w[o + 4u * m2 + 2u * hl] = pack4(q[0], q[1], q[2], q[3]); + arena.w[o + 4u * m2 + 2u * hl + 1u] = pack4(q[4], q[5], q[6], q[7]); + } + } + } else if (lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards + const uint mx = lane & 3u, my = lane >> 2u; + const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my); + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint t00 = (2u * my) * 8u + 2u * mx; + vec4 out_ = vec4(0.0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25); + precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25); + precise float16_t b0 = a00 + a01, b1 = a10 + a11; + precise float16_t cc = b0 + b1; + out_[c] = float(cc); + } + imageStore(S128, ivec2(x >> 1, y >> 1), out_); + } + } +} diff --git a/engine/rr/enc1_native.comp b/engine/rr/enc1_native.comp new file mode 100644 index 0000000..cd1982a --- /dev/null +++ b/engine/rr/enc1_native.comp @@ -0,0 +1,559 @@ +#version 460 +// Native encoder 1 of the DLSS-D (preset E) denoiser: cuda_dldn_engine_swin_enc1_kernel rewritten for cooperative +// matrices, bit-exact to engine/rr/enc1_ref.py (which is bit-exact to NVIDIA's launch). Same stage contract as the +// translated stage (engine/rr/README.md) except binding 1, which holds engine/rr/native_pack.py's packing of the +// extracted weights (offsets in the generated enc1_w.h). +// +// One workgroup is one 8x8-token window, four waves. Products run tokens x weights (A = activations from shared +// memory, B = weight tiles), one 16x16 output tile at a time, with NVIDIA's arithmetic: per logical k32 step the +// running value enters the first of two 16x16x16 FP8 products as C, the sum stays f32 and is rounded to binary16 +// once, through shared memory (RADV folds a register-only f32 -> f16 -> f32). Every FP8 operand is stored in NVIDIA's +// k order: within each 32-block, slot j holds natural channel inter(j). The reductions (norms, softmax row sums) run +// in NVIDIA's exact order. +// +// 1 prologue the 96 key tokens (window, then a border row above, below, a column left and right; mirrored at the +// image edge), RMS norm (gain W_GAIN0) -> XQ; the residual + W_TABLE7 -> RES +// 2 QK and V per head, 96 x 64 -> 32 each: 24 tile jobs spread over the waves -> KV +// 3 attention per wave 16 query rows, both heads: scores + bias, cubic softmax, P V, out-projection onto RES +// 4 MLP per wave: RMS norm (W_GAIN14), RES += W_TABLE16, eight branches 64 -> 32 -> 64 onto RES +// 5 outputs the tokens (argument 48), the 2x2 merge to 96 channels (argument 56, waves 0-2), the 64 -> 16 head +// and its 3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128 +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_GOOGLE_include_directive : require +layout(local_size_x = 128) in; +#include "enc1_w.h" + +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +layout(set = 0, binding = 3) uniform sampler2D T88; +layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112; +layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128; + +#ifdef DBG +// debugging: intermediates of window (DBX, DBY) as raw words: XQ, KV, RES after each head and after the MLP +layout(set = 0, binding = 7, std430) buffer Dbg { uint d[]; } dbg; +#define DBGWIN (gl_WorkGroupID.x == DBX && gl_WorkGroupID.y == DBY) +#define DUMP8(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = code(arr[i_]); barrier(); } +#define DUMP16(arr, n, at) if (DBGWIN) { barrier(); for (uint i_ = tid; i_ < (n); i_ += 128u) dbg.d[(at) + i_] = bh(arr[i_]); barrier(); } +#else +#define DUMP8(arr, n, at) +#define DUMP16(arr, n, at) +#endif +#ifndef SKIP +#define SKIP 0 +#endif +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +#define SYNC { memoryBarrierShared(); subgroupBarrier(); } + +shared floate4m3_t XQ[96 * 64]; // the normalised key tokens; after the MLP the output tokens (64 rows) +shared float16_t RES[64 * 64]; // the residual stream, natural channel order +shared floate4m3_t KV[4 * 96 * 32]; // Q/K head 0, V head 0, Q/K head 1, V head 1; later the merge rows (16 x 256) +shared float16_t SW[4 * 1536]; // per wave: f16 results (scores 16 x 96), rounding scratch at +1024 +shared floate4m3_t PW[4 * 1536]; // per wave: FP8 operands (P 16 x 96; O 16 x 32; N 16 x 64 + hidden 16 x 32) + +// a weight tile: matrix W (K x N) packed by native_pack.py, k-tile kt, n-tile nt +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +// one binary16 rounding of an accumulator, through the wave's scratch +#define RND(acc) { \ + coopMatStore(HAC(acc), SW, scr, 16u, RM); SYNC; \ + HAC r_; coopMatLoad(r_, SW, scr, 16u, RM); acc = ACC(r_); SYNC; } + +// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the +// arena's channel order: arena slot uninter(k) holds natural channel k) +uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); } +uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); } +int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); } + +// cvt.rn.satfinite.e4m3 (NaN -> 0x7f) +floate4m3_t q8(float16_t h) +{ + return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0))); +} +uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); } +float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); } +uint bh(float16_t h) { return uint(float16BitsToUint16(h)); } + +// x rounded to binary16 precision (nearest even), staying f32. RADV folds float(float16_t(x)) back to x, so every +// f32 -> f16 rounding whose result is used as f32 again goes through this: adding and subtracting 1.5 * 2^(E + 13) +// (E = x's binary16 exponent) leaves x rounded to a multiple of its binary16 ulp; precise keeps the pair. +float r16(float x) +{ + const uint e = max(floatBitsToUint(x) & 0x7f800000u, 0x38800000u); + const float c = uintBitsToFloat(e + (13u << 23u)) * 1.5; + precise float t = x + c; + precise float r = t - c; + return r; +} +float16_t h16(float x) { return float16_t(r16(x)); } +// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1 +int cmp1(float m, float x) +{ + precise float t = m * x; + precise float e = fma(m, x, -t); + return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0))); +} +// binary16 rounding of 1 / x (x > 0) to nearest even, from the hardware approximation and two midpoint checks +float16_t rcp16(float x) +{ + float16_t q = h16(1.0 / x); + for (int i = 0; i < 2; ++i) { + const uint b = bh(q); + const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5; + const int cl = cmp1(lo, x), ch = cmp1(hi, x); + if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u); + else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u); + } + return q; +} +// binary16 rounding of 1 / sqrt(s) (s > 0) +float16_t rsq16(float16_t s16) +{ + const float s = float(s16); + float16_t q = h16(inversesqrt(s)); + for (int i = 0; i < 2; ++i) { + const uint b = bh(q); + const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5; + const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s); + if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u); + else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u); + } + return q; +} +// fma of binary16 operands with one binary16 rounding, exactly: the product is exact in f32, the sum's f32 rounding +// error comes from TwoSum and decides the f32 results that land on a binary16 midpoint +float16_t fma16(float16_t a, float16_t b, float16_t c) +{ + precise float p = float(a) * float(b); + precise float s = p + float(c); + precise float bb = s - p; + precise float err = (p - (s - bb)) + (float(c) - bb); + const uint u = floatBitsToUint(s); + if ((u & 0x1fffu) == 0x1000u && err != 0.0 && abs(s) >= 6.103515625e-05) + return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u)); + return h16(s); +} +float16_t act(float16_t x) +{ + precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0)); + precise float16_t a = float16_t(0.0810546875) * abs(y); + precise float16_t b = float16_t(0.412109375) - a; + precise float16_t c = y * b; + precise float16_t g = float16_t(0.5) + c; + precise float16_t r = x * g; + return r; +} +float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); } +uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); } + +void main() +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const int ex = int(param.p[0]), ey = int(param.p[1]); + const uint inOff = param.p[2], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const uint plane = uint(ex * ey) * 32u; + const uint sw = wv * 1536u, scr = sw + 1024u, r0 = wv * 16u; + + // ---------------------------------------------------------------- 1 prologue +#if !(SKIP & 1) + if (tid < 96u) { + int x, y; + if (tid < 64u) { + x = x0 + int(tid & 7u); y = y0 + int(tid >> 3u); + } else { + const int i = int(tid & 7u), side = int((tid - 64u) >> 3u); + x = side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8); + y = side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i); + } + x = mirror(x, ex); y = mirror(y, ey); + const uint base = (inOff + uint(y * ex + x) * 32u) >> 2u; + float16_t v[64]; + for (uint pl = 0u; pl < 2u; ++pl) + for (uint wd = 0u; wd < 8u; ++wd) { + const uint u = arena.w[base + pl * (plane >> 2u) + wd]; + for (uint b = 0u; b < 4u; ++b) + v[32u * pl + 4u * wd + b] = float16_t(uintBitsToFloate4m3EXT(uint8_t((u >> (8u * b)) & 0xffu))); + } + // NVIDIA's order: per word t (and t + 4) the f16x2 pair sums of both planes, then lanes ^1, ^2, the halves + precise f16vec2 h[4]; + for (uint t = 0u; t < 4u; ++t) { + precise f16vec2 p[4]; + for (uint k = 0u; k < 4u; ++k) { + const uint o = 32u * (k & 1u) + 4u * (t + 4u * (k >> 1u)); + precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]); + precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]); + p[k] = a + c; + } + precise f16vec2 s0 = p[0] + p[1], s1 = p[2] + p[3]; + h[t] = s0 + s1; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t s = hs.x + hs.y; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + for (uint j = 0u; j < 64u; ++j) { + precise float16_t g = r * w16.w[W_GAIN0 / 2u + inter(j)]; + precise float16_t n = v[j] * g; + XQ[tid * 64u + j] = q8(n); + } + if (tid < 64u) + for (uint k = 0u; k < 64u; ++k) { + precise float16_t c = v[uninter(k)] + w16.w[W_TABLE7 / 2u + k]; + RES[tid * 64u + k] = c; + } + } +#endif + barrier(); + DUMP8(XQ, 6144u, 0u) + DUMP16(RES, 4096u, 6144u) + + // ---------------------------------------------------------------- 2 QK and V projections +#if !(SKIP & 2) + for (uint job = wv; job < 24u; job += 4u) { + const uint mt = job % 6u, mat = job / 6u, hd = mat >> 1u; + const bool isV = (mat & 1u) != 0u; + const uint wa = (isV ? W_WV0A : W_WK0A) + 2048u * hd, wb = wa + 1024u; + A8 a0, a1, a2, a3; + coopMatLoad(a0, XQ, mt * 1024u, 64u, RM); + coopMatLoad(a1, XQ, mt * 1024u + 16u, 64u, RM); + coopMatLoad(a2, XQ, mt * 1024u + 32u, 64u, RM); + coopMatLoad(a3, XQ, mt * 1024u + 48u, 64u, RM); + for (uint nt = 0u; nt < 2u; ++nt) { + B8 b0, b1; + LDB(b0, wa, 32u, 0u, nt); LDB(b1, wa, 32u, 1u, nt); + ACC acc = coopMatMulAdd(a1, b1, coopMatMulAdd(a0, b0, ACC(0.0))); + RND(acc); + LDB(b0, wb, 32u, 0u, nt); LDB(b1, wb, 32u, 1u, nt); + acc = coopMatMulAdd(a3, b1, coopMatMulAdd(a2, b0, acc)); + coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM); + } + SYNC; + const uint row = lane >> 1u, cb = (lane & 1u) * 16u, kv = mat * 3072u; + if (!isV) { + for (uint j = cb; j < cb + 16u; ++j) KV[kv + (mt * 16u + row) * 32u + j] = q8(SW[sw + row * 32u + inter(j)]); + } else { + const uint dr = uninter(mt * 16u + row); + for (uint c = cb; c < cb + 16u; ++c) KV[kv + dr * 32u + c] = q8(SW[sw + row * 32u + c]); + } + SYNC; + } +#endif + barrier(); + DUMP8(KV, 12288u, 10240u) + + // ---------------------------------------------------------------- 3 attention, 16 query rows per wave +#if !(SKIP & 4) + for (uint hd = 0u; hd < 2u; ++hd) { + const uint qk = 6144u * hd, vv = qk + 3072u; + { + A8 q0, q1; + coopMatLoad(q0, KV, qk + r0 * 32u, 32u, RM); + coopMatLoad(q1, KV, qk + r0 * 32u + 16u, 32u, RM); + for (uint nt = 0u; nt < 6u; ++nt) { + HAC c; + coopMatLoad(c, w16.w, (W_BIAS0 + 12288u * hd) / 2u + r0 * 96u + nt * 16u, 96u, RM); + B8 k0, k1; + coopMatLoad(k0, KV, qk + nt * 512u, 32u, CM); + coopMatLoad(k1, KV, qk + nt * 512u + 16u, 32u, CM); + ACC acc = coopMatMulAdd(q1, k1, coopMatMulAdd(q0, k0, ACC(c))); + coopMatStore(HAC(acc), SW, sw + nt * 16u, 96u, RM); + } + } + SYNC; + { // the cubic-exponent softmax: lane = (row, column parity) + const uint row = lane >> 1u, par = lane & 1u; + float16_t e[48]; + for (uint j = 0u; j < 12u; ++j) + for (uint tt = 0u; tt < 4u; ++tt) { + const float16_t x = SW[sw + row * 96u + 8u * j + 2u * tt + par]; + precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375)); + const float16_t t = fma16(m, -m, float16_t(0.92724609375)); + const float16_t u = fma16(m, t, float16_t(1.375)); + e[4u * j + tt] = hb((bh(u) << 5u) & 0x7FE0u); + } + precise float16_t sm[4]; + for (uint tt = 0u; tt < 4u; ++tt) { + precise float16_t ev = ((e[tt] + e[8u + tt]) + (e[16u + tt] + e[24u + tt])) + (e[32u + tt] + e[40u + tt]); + precise float16_t od = ((e[4u + tt] + e[12u + tt]) + (e[20u + tt] + e[28u + tt])) + (e[36u + tt] + e[44u + tt]); + sm[tt] = ev + od; + } + precise float16_t out_ = (sm[0] + sm[1]) + (sm[2] + sm[3]); + const float16_t other = hb(subgroupShuffleXor(bh(out_), 1u)); + precise float16_t rs = par == 0u ? out_ + other : other + out_; + const float16_t r = rcp16(float(rs)); + for (uint j = 0u; j < 12u; ++j) + for (uint tt = 0u; tt < 4u; ++tt) { + precise float16_t pv = e[4u * j + tt] * r; + PW[sw + row * 96u + uninter(8u * j + 2u * tt + par)] = q8(pv); + } + } + SYNC; + if (hd == 0u) { DUMP16(SW, 6144u, 34816u) DUMP8(PW, 6144u, 40960u) } + for (uint nt = 0u; nt < 2u; ++nt) { // P V + ACC acc = ACC(0.0); + for (uint ks = 0u; ks < 3u; ++ks) { + A8 p0, p1; B8 v0, v1; + coopMatLoad(p0, PW, sw + ks * 32u, 96u, RM); + coopMatLoad(p1, PW, sw + ks * 32u + 16u, 96u, RM); + coopMatLoad(v0, KV, vv + ks * 1024u + nt * 16u, 32u, RM); + coopMatLoad(v1, KV, vv + ks * 1024u + 512u + nt * 16u, 32u, RM); + acc = coopMatMulAdd(p1, v1, coopMatMulAdd(p0, v0, acc)); + if (ks < 2u) RND(acc); + } + coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM); + } + SYNC; + { + const uint row = lane >> 1u, cb = (lane & 1u) * 16u; + for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 32u + j] = q8(SW[sw + row * 32u + inter(j)]); + } + SYNC; + { // out-projection onto the residual + A8 o0, o1; + coopMatLoad(o0, PW, sw, 32u, RM); + coopMatLoad(o1, PW, sw + 16u, 32u, RM); + for (uint nt = 0u; nt < 4u; ++nt) { + HAC c; + coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM); + B8 b0, b1; + LDB(b0, W_WO0 + 2048u * hd, 64u, 0u, nt); LDB(b1, W_WO0 + 2048u * hd, 64u, 1u, nt); + ACC acc = coopMatMulAdd(o1, b1, coopMatMulAdd(o0, b0, ACC(c))); + coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM); + } + } + SYNC; + DUMP16(RES, 4096u, 22528u + 4096u * hd) + } +#endif + + // ---------------------------------------------------------------- 4 MLP +#if !(SKIP & 8) + if (lane < 16u) { + const uint base = (r0 + lane) * 64u; + float16_t R[64]; + for (uint k = 0u; k < 64u; ++k) R[k] = RES[base + k]; + precise f16vec2 h[4]; + for (uint t = 0u; t < 4u; ++t) { + precise f16vec2 q[8]; + for (uint j = 0u; j < 8u; ++j) { + const f16vec2 x = f16vec2(R[8u * j + 2u * t], R[8u * j + 2u * t + 1u]); + q[j] = x * x; + } + precise f16vec2 a = q[0] + q[2], b = q[1] + q[3], c = q[4] + q[6], d = q[5] + q[7]; + precise f16vec2 ac = a + c, bd = b + d; + h[t] = ac + bd; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t s = hs.x + hs.y; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + for (uint j = 0u; j < 64u; ++j) { + const uint k = inter(j); + precise float16_t g = r * w16.w[W_GAIN14 / 2u + k]; + precise float16_t n = R[k] * g; + PW[sw + lane * 96u + j] = q8(n); + } + for (uint k = 0u; k < 64u; ++k) { + precise float16_t c = R[k] + w16.w[W_TABLE16 / 2u + k]; + RES[base + k] = c; + } + } + SYNC; + { + A8 n0, n1, n2, n3; + coopMatLoad(n0, PW, sw, 96u, RM); + coopMatLoad(n1, PW, sw + 16u, 96u, RM); + coopMatLoad(n2, PW, sw + 32u, 96u, RM); + coopMatLoad(n3, PW, sw + 48u, 96u, RM); + for (uint br = 0u; br < 8u; ++br) { + const uint w1 = W_W1_0 + 2048u * br, w2 = W_W2_0 + 2048u * br; + for (uint nt = 0u; nt < 2u; ++nt) { + HAC c; + coopMatLoad(c, w16.w, (W_B1_0 + 1024u * br) / 2u + nt * 16u, 32u, RM); + B8 b0, b1; + LDB(b0, w1, 32u, 0u, nt); LDB(b1, w1, 32u, 1u, nt); + ACC acc = coopMatMulAdd(n1, b1, coopMatMulAdd(n0, b0, ACC(c))); + RND(acc); + LDB(b0, w1, 32u, 2u, nt); LDB(b1, w1, 32u, 3u, nt); + acc = coopMatMulAdd(n3, b1, coopMatMulAdd(n2, b0, acc)); + coopMatStore(HAC(acc), SW, sw + nt * 16u, 32u, RM); + } + SYNC; + { + const uint row = lane >> 1u, cb = (lane & 1u) * 16u; + for (uint j = cb; j < cb + 16u; ++j) PW[sw + row * 96u + 64u + j] = q8(act(SW[sw + row * 32u + inter(j)])); + } + SYNC; + A8 g0, g1; + coopMatLoad(g0, PW, sw + 64u, 96u, RM); + coopMatLoad(g1, PW, sw + 80u, 96u, RM); + for (uint nt = 0u; nt < 4u; ++nt) { + HAC c; + coopMatLoad(c, RES, r0 * 64u + nt * 16u, 64u, RM); + B8 b0, b1; + LDB(b0, w2, 64u, 0u, nt); LDB(b1, w2, 64u, 1u, nt); + ACC acc = coopMatMulAdd(g1, b1, coopMatMulAdd(g0, b0, ACC(c))); + coopMatStore(HAC(acc), RES, r0 * 64u + nt * 16u, 64u, RM); + } + SYNC; + } + } + +#endif + DUMP16(RES, 4096u, 30720u) + + // ---------------------------------------------------------------- 5 outputs + { // the tokens, in arena order (= NVIDIA's k order) + const uint row = lane >> 1u, pl = lane & 1u, t = r0 + row; + const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u); + uint words[8]; + for (uint wd = 0u; wd < 8u; ++wd) { + uint u = 0u; + for (uint b = 0u; b < 4u; ++b) { + const uint j = 32u * pl + 4u * wd + b; + const floate4m3_t c = q8(RES[t * 64u + inter(j)]); + XQ[t * 64u + j] = c; + u |= code(c) << (8u * b); + } + words[wd] = u; + } + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint o = (outOff + pl * plane + uint(y * ex + x) * 32u) >> 2u; + for (uint wd = 0u; wd < 8u; ++wd) arena.w[o + wd] = words[wd]; + } + } + barrier(); +#if !(SKIP & 16) + for (uint i = 0u; i < 32u; ++i) { // the merge rows: (mx, my) in NVIDIA's order, 4 tokens of 64 channels + const uint idx = tid * 32u + i, r = idx >> 8u, k = idx & 255u, blk = k >> 6u; + const uvec2 m = mergeXY(r); + KV[idx] = XQ[((2u * m.y + (blk >> 1u)) * 8u + 2u * m.x + (blk & 1u)) * 64u + (k & 63u)]; + } + barrier(); + if (wv < 3u) { // the 2x2 merge, one 32-channel plane per wave + for (uint n2 = 0u; n2 < 2u; ++n2) { + const uint nt = 2u * wv + n2; + HAC c; + coopMatLoad(c, w16.w, W_BM / 2u + nt * 16u, 96u, RM); + ACC acc = ACC(c); + for (uint ks = 0u; ks < 8u; ++ks) { + A8 m0, m1; B8 b0, b1; + coopMatLoad(m0, KV, ks * 32u, 256u, RM); + coopMatLoad(m1, KV, ks * 32u + 16u, 256u, RM); + LDB(b0, W_WM, 96u, 2u * ks, nt); LDB(b1, W_WM, 96u, 2u * ks + 1u, nt); + acc = coopMatMulAdd(m1, b1, coopMatMulAdd(m0, b0, acc)); + if (ks < 7u) RND(acc); + } + coopMatStore(HAC(acc), SW, sw + n2 * 16u, 32u, RM); + } + SYNC; + const uint row = lane >> 1u, hf = lane & 1u; + const uvec2 m = mergeXY(row); + const int X = x0 / 2 + int(m.x), Y = y0 / 2 + int(m.y), hx = ex / 2, hy = ey / 2; + uint words[4]; + for (uint wd = 0u; wd < 4u; ++wd) { + uint u = 0u; + for (uint b = 0u; b < 4u; ++b) u |= code(q8(SW[sw + row * 32u + inter(16u * hf + 4u * wd + b)])) << (8u * b); + words[wd] = u; + } + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + const uint o = (mrgOff + wv * uint(hx * hy) * 32u + uint(Y * hx + X) * 32u) >> 2u; + for (uint wd = 0u; wd < 4u; ++wd) arena.w[o + 4u * hf + wd] = words[wd]; + } + SYNC; + } +#endif + { // the 64 -> 16 head + HAC c; + coopMatLoad(c, w16.w, W_BH / 2u, 16u, RM); + A8 t0, t1, t2, t3; B8 b0, b1; + coopMatLoad(t0, XQ, r0 * 64u, 64u, RM); + coopMatLoad(t1, XQ, r0 * 64u + 16u, 64u, RM); + coopMatLoad(t2, XQ, r0 * 64u + 32u, 64u, RM); + coopMatLoad(t3, XQ, r0 * 64u + 48u, 64u, RM); + LDB(b0, W_WH, 16u, 0u, 0u); LDB(b1, W_WH, 16u, 1u, 0u); + ACC acc = coopMatMulAdd(t1, b1, coopMatMulAdd(t0, b0, ACC(c))); + RND(acc); + LDB(b0, W_WH, 16u, 2u, 0u); LDB(b1, W_WH, 16u, 3u, 0u); + acc = coopMatMulAdd(t3, b1, coopMatMulAdd(t2, b0, acc)); + coopMatStore(HAC(acc), SW, sw, 16u, RM); + } + SYNC; + { // the kernel-prediction filter; lanes 16-31 repeat 0-15 so that the shuffles below stay uniform + const uint tl = lane & 15u, tk = r0 + tl; + const int x = x0 + int(tk & 7u), y = y0 + int(tk >> 3u); + float16_t l[12]; + for (uint i = 0u; i < 12u; ++i) l[i] = SW[sw + tl * 16u + i]; + const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]); + const float16_t mx = max(max(ma, mb), l[8]); + float e[9]; + for (uint i = 0u; i < 9u; ++i) { + precise float16_t d = l[i] - mx; + e[i] = exp2(float(d) * 1.4426950216293335); + } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t r = rcp16(s); + float wg[9]; + for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(r)); + const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1; + const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1; + const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc), + ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp)); + // f32 values that are binary16 numbers throughout (r16 at every NVIDIA f16 rounding) + float o112[3], o120[3]; + float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0)); + for (uint k = 0u; k < 8u; ++k) { + const vec4 tv = texelFetch(T88, pts[k], 0); + for (uint c = 0u; c < 3u; ++c) { + precise float f = tv[c] * wg[k] + acc[c][k & 1u]; + acc[c][k & 1u] = r16(f); + } + } + const vec4 t8 = texelFetch(T88, pts[8], 0); + for (uint c = 0u; c < 3u; ++c) { + precise float f0 = acc[c][0] + acc[c][1]; + precise float f1 = t8[c] * wg[8] + r16(f0); + const float f = r16(f1); + precise float a = float(sig(l[9u + c])) * f; + const float ar = r16(a); + precise float b = f - ar; + o120[c] = ar; o112[c] = r16(b); + } + // surface 128: the 2x2 sums of 0.25 * surface 112, rightwards first, then downwards + float q[3], cs[3]; + const uint nb = min(lane + 1u, 31u), dn = min(lane + 8u, 31u); + for (uint c = 0u; c < 3u; ++c) { + q[c] = o112[c] * 0.25; + precise float b = q[c] + subgroupShuffle(q[c], nb); + const float br = r16(b); + precise float d = br + subgroupShuffle(br, dn); + cs[c] = r16(d); + } + if (lane < 16u && x >= 0 && y >= 0 && x < ex && y < ey) { + imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0)); + imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0)); + if (((x | y) & 1) == 0) imageStore(S128, ivec2(x >> 1, y >> 1), vec4(cs[0], cs[1], cs[2], 0.0)); + } + } +} diff --git a/engine/rr/enc1_pack.py b/engine/rr/enc1_pack.py new file mode 100644 index 0000000..d2fa77b --- /dev/null +++ b/engine/rr/enc1_pack.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""Pack encoder 1's extracted weights (engine/rr/enc1_ref.py --extract) for engine/rr/enc1_m.comp. + +usage: enc1_pack.py WEIGHTS.npz OUT.bin OUT.h + +The shader runs every product transposed (D^T = W^T X^T: weights are the A operand, activations the B operand), so a +result's rows are output channels and its columns tokens, and the result is the next product's B operand in place. +The rows come out in the order of W's columns, so every weight's columns are permuted here into the k order NVIDIA's +next product contracts them in (engine/rr/enc1_ref.py: q8(x)[:, inter(n)]), and the f16 vectors that act on those +rows likewise. Keys and queries are token slots: slot s is token inter(96)[s] (the prologue's key list), so the +attention bias is stored transposed in slot order. + + A tiles W (K x N) -> W^T 16x16 tiles [mt][kt], row-major (m, k): 256 * (mt * K/16 + kt) bytes + B tiles the V projection runs X Wv (A = tokens): tiles [kt][nt], row-major (k, n) + f16 vectors as is; the attention bias as 96 x 64 [key slot][query slot] +""" +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from enc1_ref import inter # noqa: E402 + +I32, I64, I96 = inter(32), inter(64), inter(96) + + +def a_tiles(w): + k, n = w.shape + return np.ascontiguousarray(w.T).reshape(n // 16, 16, k // 16, 16).transpose(0, 2, 1, 3).reshape(-1) + + +def b_tiles(w): + k, n = w.shape + return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1) + + +CUBIC_MAX = 0x3873 # binary16 bits of 0.55615234375, the clamp of the scaled score + + +def act_table(): + """the MLP's hidden activation and its FP8 conversion, q8(act(h)), for every binary16 h, indexed by its bits: one + byte load in place of the six binary16 operations and the conversion (the shaders' ACTQ)""" + from enc1_ref import act + from rrmath import q8 + h = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16) + with np.errstate(all='ignore'): + return np.asarray(q8(act(h).astype(np.float64)), np.uint8) + + +def cubic_table(): + """the cubic softmax numerator's bits (enc1_ref.softmax_cubic) for every binary16 score, indexed by its bits (the + scaling and the clamp included; NaN scores take NVIDIA's min / max, which return the other operand)""" + x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16) + with np.errstate(invalid='ignore', over='ignore'): + m = (x * np.float16(0.017181396484375)).astype(np.float16) + lim = np.float16(0.55615234375) + m = np.where(np.isnan(m), lim, np.clip(m, -lim, lim)).astype(np.float16) + + def fma16(a, b, c): + return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(np.float16) + t = fma16(m, -m, np.full_like(m, np.float16(0.92724609375))) + u = fma16(m, t, np.full_like(m, np.float16(1.375))) + return ((u.view(np.uint16).astype(np.uint32) << 5) & 0x7FE0).astype(np.uint16) + + +def main(): + src, out_bin, out_h = sys.argv[1:4] + w = dict(np.load(src)) + e = {} + for h in range(2): + e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32]) + e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32]) + e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64]) + e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T) + for b in range(8): + e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32]) + e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32]) + e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64]) + e['wm'], e['bm'] = ('A', w['wm'][:, I96]), ('H', w['bm'][I96]) + e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh']) + e['gain0'] = ('H', w['gain0']) + for k in ('table7', 'gain14', 'table16'): + e[k] = ('H', w[k][I64]) + e['cubic'] = ('U', cubic_table()) + blob, defines = bytearray(), [] + for name, (kind, a) in e.items(): + while len(blob) % 256: + blob.append(0) + off = len(blob) + if kind == 'A': + blob += a_tiles(a).tobytes() + elif kind == 'B': + blob += b_tiles(a).tobytes() + elif kind == 'U': + blob += np.ascontiguousarray(a, np.uint16).tobytes() + else: + blob += np.ascontiguousarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write('// generated by engine/rr/enc1_pack.py\n' + '\n'.join(defines) + '\n') + print(f'{len(blob)} bytes, {len(e)} entries') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/enc1_ref.py b/engine/rr/enc1_ref.py new file mode 100644 index 0000000..d4aff3b --- /dev/null +++ b/engine/rr/enc1_ref.py @@ -0,0 +1,330 @@ +#!/usr/bin/env python3 +"""Exact numpy reference of the denoiser's encoder 1 (cuda_dldn_engine_swin_enc1_kernel), one 8x8 window at a time, +from its weights and real inputs only: the specification the native shader implements. + +Derived phase by phase against ZLUDA replays (engine/rr/phase_dump.py), every step bit-exact: + prologue the 96 key tokens (the window plus a one-token border, corners excluded; mirrored at the image edge) + from the arena's two 32-channel planes, f16 RMS norm (NVIDIA's summation order) times the gain at W+0 + attention two heads: QK and V projections (K = 64 in two k32 steps), scores = q8(QK[:64]) q8(QK)^T + bias, + the cubic-exponent softmax, P V, out-projection onto the residual (input tile + table at W+36992) + MLP RMS norm times the gain at W+37120, eight branches of 64 -> 32 (bias, ramp activation) -> 64, + accumulated onto the residual (+ table at W+41408) + outputs q8 tokens to param+48, the 2x2 merge (16 x 256 -> 96 + bias) to param+56, the 64 -> 16 head and its + 3x3 kernel-prediction filter of texture 88 to surfaces 112, 120 and (half resolution) 128 + +usage: enc1_ref.py CAPTURE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture) + WEIGHTS.npz: engine/rr/enc1_ref.py --extract DUMP.npz CAPTURE FRAME OUT.npz builds it from a phase dump +""" +import argparse +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import e4m3, f16, q8, mma_k32 # noqa: E402 + +H = np.float16 +F = np.float32 + + +def inter(n): + perm = [] + for b in range(n // 32): + for m in range(8): + perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 17 + 2 * m] + return perm + + +INTER64 = inter(64) +# arena channel order (NVIDIA's A-fragment order) -> natural: natural channel k is arena channel ARENA[k] +ARENA = [0, 1, 4, 5, 8, 9, 12, 13, 16, 17, 20, 21, 24, 25, 28, 29, 2, 3, 6, 7, 10, 11, 14, 15, 18, 19, 22, 23, 26, 27, + 30, 31] +ARENA = ARENA + [32 + c for c in ARENA] + + +def add(a, b): + return (a + b).astype(H) + + +def mm(a_codes, b_codes, c): + """a (M x K e4m3 codes, NVIDIA k order) . b (K x N codes) + c (M x N f16 values) with ZLUDA's k32 rounding""" + return f16(mma_k32(a_codes, b_codes, np.asarray(c, H).view(np.uint16))).astype(H) + + +def nv_rowsum_sq(x, eps=2.0 ** -13): + """f16(rsqrt(sum of squares + eps)) per row of a 64-column C/D tile, NVIDIA's order""" + sq = (x * x).astype(H) + + def tile(j, t): + return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2] + h = [] + for t in range(4): + a, b = add(tile(0, t), tile(2, t)), add(tile(1, t), tile(3, t)) + c, d = add(tile(4, t), tile(6, t)), add(tile(5, t), tile(7, t)) + h.append(add(add(a, c), add(b, d))) + h1 = [add(h[t], h[t ^ 1]) for t in range(4)] + h2 = add(h1[0], h1[2]) + s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps))) + return (F(1) / np.sqrt(s.astype(F))).astype(H) + + +def input_norm(tok_arena, gain): + """the prologue's per-token norm: tok_arena (2 planes x 32, arena order) -> f16 normalised token (natural)""" + sq = [(v * v).astype(H) for v in tok_arena] + h = [] + for t in range(4): + def pair(p, w): + s = sq[p][4 * w: 4 * w + 4] + return add(s[0:2], s[2:4]) + h.append(add(add(pair(0, t), pair(1, t)), add(pair(0, t + 4), pair(1, t + 4)))) + h01 = [add(h[t], h[t ^ 1]) for t in range(4)] + h2 = add(h01[0], h01[2]) + s = H(H(h2[0] + h2[1]) + H(2.0 ** -13)) + r = H(F(1) / np.sqrt(F(s))) + x = np.concatenate(tok_arena)[ARENA] + return (x * (r * gain).astype(H)).astype(H), x + + +def softmax_cubic(x): + """E5: the cubic-exponent softmax of a 64 x 96 score block (natural columns)""" + def fma16(a, b, c): + return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(H) + m = np.clip((x * H(0.017181396484375)).astype(H), H(-0.55615234375), H(0.55615234375)) + t = fma16(m, -m, np.full_like(m, H(0.92724609375))) + u = fma16(m, t, np.full_like(m, H(1.375))) + e = ((((u.view(np.uint16).astype(np.uint32)) << 5) & 0x7FE0).astype(np.uint16)).view(H) + + def tile(j, tt, half): + return e[:, 8 * j + 2 * tt + half] + s = {} + for tt in range(4): + for half in range(2): + ev = add(add(add(tile(0, tt, half), tile(2, tt, half)), add(tile(4, tt, half), tile(6, tt, half))), + add(tile(8, tt, half), tile(10, tt, half))) + od = add(add(add(tile(1, tt, half), tile(3, tt, half)), add(tile(5, tt, half), tile(7, tt, half))), + add(tile(9, tt, half), tile(11, tt, half))) + s[tt, half] = add(ev, od) + out = [add(add(s[0, h_], s[1, h_]), add(s[2, h_], s[3, h_])) for h_ in range(2)] + rs = add(out[0], out[1]) + r = (F(1) / rs.astype(F)).astype(H) + return (e * r[:, None]).astype(H) + + +def act(x): + y = np.clip(x, H(-2), H(2)).astype(H) + g = add(H(0.5), (y * (H(0.412109375) - (H(0.0810546875) * np.abs(y)).astype(H)).astype(H)).astype(H)) + return (x * g).astype(H) + + +MERGE_ORDER = [(0, 0), (0, 1), (1, 0), (1, 1), (2, 0), (2, 1), (3, 0), (3, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 2), + (2, 3), (3, 2), (3, 3)] + + +def mirror(v, n): + v = abs(v) + return min(v, 2 * n - 2 - v) + + +def window(w, planes, tex, x0, y0, ex, ey, dbg=None): + """one window: returns (tokens out [64 x 64 codes, arena order], merged [16 x 96 codes, A order], + surf112 / surf120 [64 x 3 f16], surf128 contributions)""" + border = ([(x, y0 - 1) for x in range(x0, x0 + 8)] + [(x, y0 + 8) for x in range(x0, x0 + 8)] + + [(x0 - 1, y) for y in range(y0, y0 + 8)] + [(x0 + 8, y) for y in range(y0, y0 + 8)]) + coords = [(x, y) for y in range(y0, y0 + 8) for x in range(x0, x0 + 8)] + border + norm, raw = [], [] + for (x, y) in coords: + xx, yy = mirror(x, ex), mirror(y, ey) + n_, r_ = input_norm([e4m3(planes[p, yy, xx]).astype(H) for p in (0, 1)], w['gain0']) + norm.append(n_) + raw.append(r_) + X = q8(np.array(norm))[:, INTER64] # 96 x 64 codes, NVIDIA k order + res = np.array(raw[:64]) + dbg = {} if dbg is None else dbg + dbg['XQ'], dbg['RES0'] = X, add(res, w['table7'][None, :]) + R = None + for h in range(2): + qk = mm(X[:, 32:], w[f'wk{h}b'], mm(X[:, :32], w[f'wk{h}a'], np.zeros((96, 32)))) + v = mm(X[:, 32:], w[f'wv{h}b'], mm(X[:, :32], w[f'wv{h}a'], np.zeros((96, 32)))) + Q = q8(qk)[:, inter(32)] + dbg[f'QK{h}'], dbg[f'V{h}'] = Q, q8(v)[inter(96)] + S = mm(Q[:64], Q.T, w[f'bias{h}']) + P = softmax_cubic(S) + O = mm(q8(P)[:, inter(96)], q8(v)[inter(96)], np.zeros((64, 32))) + dbg[f'S{h}'], dbg[f'P{h}'], dbg[f'O{h}'] = S, q8(P)[:, inter(96)], O + C = add(res, w['table7'][None, :]) if h == 0 else R + R = mm(q8(O)[:, inter(32)], w[f'wo{h}'], C) + dbg[f'R{h}'] = R + N = q8((R * (nv_rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, INTER64] + acc = add(R, w['table16'][None, :]) + for b in range(8): + hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0))) + acc = mm(q8(hid)[:, inter(32)], w[f'w2_{b}'], acc) + dbg['ACC'] = acc + tokens = q8(acc)[:, INTER64] + Mg = np.zeros((16, 256), np.uint8) + for r, (mx, my) in enumerate(MERGE_ORDER): + for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]): + Mg[r, 64 * blk: 64 * blk + 64] = tokens[(2 * my + dy) * 8 + 2 * mx + dx] + merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(96)] + head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0)) + s112, s120 = filter_head(head, tex, x0, y0, ex, ey) + return tokens, merged, s112, s120 + + +def filter_head(D, tex, x0, y0, ex, ey, ungated=False): + """surfaces 112 and 120 of the encoders (the gate's split of the filtered value); ungated (the bottleneck enc5): + the filtered value itself, for surface 112""" + def sig(v): + return H(F(1) / (np.exp2(F(v) * F(-1.4426950216293335)) + F(1))) + + def ex2(v): + return np.exp2(F(v) * F(1.4426950216293335)) + out112, out120 = np.zeros((64, 3), H), np.zeros((64, 3), H) + for tk in range(64): + x, y = x0 + tk % 8, y0 + tk // 8 + l_ = D[tk].astype(H) + m01 = np.maximum(np.maximum(np.maximum(l_[0:2], l_[2:4]), l_[4:6]), l_[6:8]) + m = max(max(m01[0], m01[1]), l_[8]) + e = [ex2(H(l_[i] - m)) for i in range(9)] + pair = [F(e[2 * k] + e[2 * k + 1]) for k in range(4)] + s = F(F(pair[0]) + F(0)) + for k in (1, 2, 3): + s = F(s + pair[k]) + s = F(s + e[8]) + r = H(F(1) / s) + wgt = [H(H(e[i]) * r) for i in range(9)] + + def cl(v, n): + return 1 if v < 0 else (n - 2 if v >= n else v) + xm, xc, xp = cl(x - 1, ex), cl(x, ex), cl(x + 1, ex) + ym, yc, yp = cl(y - 1, ey), cl(y, ey), cl(y + 1, ey) + + def T(xx, yy): + return tex[yy, xx].astype(H) + pts = [((xm, ym), (xc, ym)), ((xp, ym), (xm, yc)), ((xc, yc), (xp, yc)), ((xm, yp), (xc, yp))] + acc = np.zeros((3, 2), H) + for k, (pa, pb) in enumerate(pts): + for c in range(3): + for h_, pt in enumerate((pa, pb)): + acc[c, h_] = H(F(T(*pt)[c]) * F(wgt[2 * k + h_]) + F(acc[c, h_])) + f = [H(acc[c, 0] + acc[c, 1]) for c in range(3)] + f = [H(F(T(xp, yp)[c]) * F(wgt[8]) + F(f[c])) for c in range(3)] + if ungated: + out112[tk] = f + continue + g = [sig(l_[9]), sig(l_[10]), sig(l_[11])] + a = [H(g[c] * f[c]) for c in range(3)] + out120[tk] = a + out112[tk] = [H(f[c] - a[c]) for c in range(3)] + return out112, out120 + + +def extract(dump, cap, frame, out): + """the weights, from a phase dump of any block (they do not depend on the block) and the weight image""" + z = np.load(dump) + + def B(p): + return z[f'p{p}_w0_B'] + + def C(p): + return f16(z[f'p{p}_w0_C']).astype(H) + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from rr_capture import Capture + c = Capture(os.path.expanduser(cap)) + W = np.frombuffer(c.find('swin_enc1', frame).weights, np.uint8) + + def table(off, n=64): + return W[off:off + 2 * n].view(np.float16) + w = {'gain0': table(0), 'gain14': table(37120), 'table7': table(36992), 'table16': table(41408)} + for h, base in ((0, 1), (1, 8)): + w[f'wk{h}a'], w[f'wv{h}a'], w[f'wk{h}b'], w[f'wv{h}b'] = B(base), B(base + 1), B(base + 2), B(base + 3) + w[f'bias{h}'] = C(base + 4) + w[f'wo{h}'] = B(base + 6) + for b in range(8): + w[f'w1_{b}'], w[f'b1_{b}'] = B(15 + 2 * b), C(15 + 2 * b)[0] + w[f'w2_{b}'] = B(16 + 2 * b) + w['wm'], w['bm'] = B(31), C(31)[0] + w['wh'], w['bh'] = B(32), C(32)[0] + np.savez_compressed(out, **w) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('frame', type=int) + ap.add_argument('weights') + ap.add_argument('--extract', metavar='DUMP') + ap.add_argument('--blocks', nargs='*', default=['5,5']) + o = ap.parse_args() + if o.extract: + extract(o.extract, o.capture, o.frame, o.weights) + return + from rr_capture import Capture + c = Capture(os.path.expanduser(o.capture)) + l = c.find('swin_enc1', o.frame) + ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32)) + sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16)) + pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8) + post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8) + ptr = {p.offset: p.delta for p in l.pointers} + planes = pre[ptr[8]:ptr[8] + 2 * ex * ey * 32].reshape(2, ey, ex, 32) + outp = post[ptr[48]:ptr[48] + 2 * ex * ey * 32].reshape(2, ey, ex, 32) + mrg = post[ptr[56]:ptr[56] + 3 * (ex // 2) * (ey // 2) * 32].reshape(3, ey // 2, ex // 2, 32) + img = {i.offset: i for i in l.images} + + def surf(off): + i = img[off] + return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) + tex, s112, s120 = surf(88), surf(112), surf(120) + w = dict(np.load(o.weights)) + for blk in o.blocks: + bx, by = (int(v) for v in blk.split(',')) + x0, y0 = 8 * bx - sx, 8 * by - sy + tokens, merged, o112, o120 = window(w, planes, tex, x0, y0, ex, ey) + tk_ok = sum(np.array_equal(tokens[t], np.concatenate([outp[0, y0 + t // 8, x0 + t % 8], + outp[1, y0 + t // 8, x0 + t % 8]])) + for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey) + mg_ok = sum(np.array_equal(merged[r], np.concatenate([mrg[p, y0 // 2 + my, x0 // 2 + mx] for p in range(3)])) + for r, (mx, my) in enumerate(MERGE_ORDER) + if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2) + inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey] + s_ok = sum(np.array_equal(o112[t], s112[y0 + t // 8, x0 + t % 8, :3]) and + np.array_equal(o120[t], s120[y0 + t // 8, x0 + t % 8, :3]) for t in inside) + print(f'block {bx},{by}: tokens {tk_ok}/{len(inside)}, merged {mg_ok}, surfaces {s_ok}/{len(inside)}') + + +if __name__ == '__main__': + main() + + +def pairs_tree(v): + """NVIDIA's reduction over a list: (v0 + v1) + v2 for three, (v0 + v1) + (v2 + v3) for four, then the same + recursively on halves rounded up ((v0 + v1) + (v2 + v3)) + v4 for five""" + if len(v) == 1: + return v[0] + if len(v) == 2: + return add(v[0], v[1]) + half = 1 << ((len(v) - 1).bit_length() - 1) + return add(pairs_tree(v[:half]), pairs_tree(v[half:])) + + +def rowsum_sq(x, eps=2.0 ** -13): + """nv_rowsum_sq for any C = 8 n8-tiles: per lane word t the tiles of each parity (tile pairs 2 apart, tree), + the parities, then t ^ 1, t ^ 2 and the halves""" + sq = (x * x).astype(H) + nt = x.shape[1] // 8 + + def tile(j, t): + return sq[:, 8 * j + 2 * t: 8 * j + 2 * t + 2] + h = [] + for t in range(4): + ev = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(0, nt, 4)] if nt % 4 == 0 else + [tile(j, t) for j in range(0, nt, 2)]) + od = pairs_tree([add(tile(j, t), tile(j + 2, t)) for j in range(1, nt, 4)] if nt % 4 == 0 else + [tile(j, t) for j in range(1, nt, 2)]) + h.append(add(ev, od)) + h1 = [add(h[t], h[t ^ 1]) for t in range(4)] + h2 = add(h1[0], h1[2]) + s = add(add(h2[:, 0], h2[:, 1]), np.full(x.shape[0], H(eps))) + return (F(1) / np.sqrt(s.astype(F))).astype(H) diff --git a/engine/rr/f16_to_e4m3.npy b/engine/rr/f16_to_e4m3.npy new file mode 100644 index 0000000..37e6a96 Binary files /dev/null and b/engine/rr/f16_to_e4m3.npy differ diff --git a/engine/rr/gen_enc0_m.py b/engine/rr/gen_enc0_m.py new file mode 100644 index 0000000..87fdad8 --- /dev/null +++ b/engine/rr/gen_enc0_m.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 +"""Generate engine/rr/enc0_m.comp's source: encoder 0 with ptx2rr.py's translation of NVIDIA's per-pixel feature gather +(run by wave 0, as NVIDIA runs it in one warp) and the native rest (engine/rr/enc0_m_decl.glsl, enc0_m_tail.glsl). + +usage: gen_enc0_m.py ENC0.generated.comp OUT.comp + +The translation's main is cut at its first workgroup barrier: by then the 10 x 10 pixels' features and the noise grid +are in its shared buffer sh and the state surfaces 304 / 1016 are written. Model data derived from the user's DLL is in +the translated part; do not commit the output. +""" +import re +import sys + + +def main(): + src, out = sys.argv[1:3] + lines = open(src).read().split('\n') + m = next(i for i, l in enumerate(lines) if l.startswith('void main()')) + head = '\n'.join(lines[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;', + 'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;') + body = lines[m + 2:] + first = next(i for i, l in enumerate(body) if not re.match(r' (precise |bool |uint64_t |float |uint |int )', l)) + bar = next(i for i, l in enumerate(body) if l.strip() == 'memoryBarrierShared(); barrier();') + decls, prologue = body[:first], body[first:bar] + # NVIDIA's warp gathers the 96 key pixels in a loop of three 32-pixel steps, then the four corners: the steps go to + # waves 0-2 (the setup before the loop runs in every wave, the corners in wave 0) + do_ = next(i for i, l in enumerate(prologue) if l.strip() == 'do') + end = next(i for i, l in enumerate(prologue) if l.strip().startswith('} while (again_')) + loop = '\n'.join(prologue[do_ - 1:end + 1]) + inc = re.search(r'(r\d+) = uint\((r\d+) \+ 32u\);\n\s+(p\d+) = \(\1 < 96u\);', loop) + counter = inc.group(1) + loop = loop.replace(inc.group(0), inc.group(0).replace('+ 32u', '+ 96u')) + pre = [l.replace(f'{counter} = uint(0u);', f'{counter} = uint(32u * gl_SubgroupID);') for l in prologue[:do_ - 1]] + post = prologue[end + 1:] + head = head.replace('shared uint mms[384];', 'shared uint mms[384 * 4];') + text = (head + '\n#extension GL_GOOGLE_include_directive : require\n#include "enc0_m_decl.glsl"\n\nvoid main()\n{\n' + + '\n'.join(decls) + '\n // the translated feature gather (NVIDIA\'s one warp, its pixel loop split over the waves)\n' + + '\n'.join(pre) + '\n if (gl_SubgroupID < 3u) {\n' + loop + '\n }\n if (gl_SubgroupID == 0u) {\n' + + '\n'.join(post) + '\n }\n#include "enc0_m_tail.glsl"\n}\n') + open(out, 'w').write(text) + print(f'{out}: {len(decls)} declaration lines, {len(prologue)} prologue lines') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/hkpn.comp b/engine/rr/hkpn.comp new file mode 100644 index 0000000..91baaf3 --- /dev/null +++ b/engine/rr/hkpn.comp @@ -0,0 +1,298 @@ +#version 460 +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_EXT_shader_16bit_storage : require + +// Ray Reconstruction hkpn output stage, preset E (no diamond-wallaby branch). +// One workgroup per 16x16 output tile. Stage 1 stages the render-space neighbourhood of +// the tile into shared memory (colour log, per-texel weights, anisotropy); stage 2 gathers +// a 5x5 elliptical kernel per output pixel and writes the two output surfaces. +// Parameter block (binding 0) holds NVIDIA's kernel argument block; word index = byte / 4. + +layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; + +layout(std430, binding = 0) readonly buffer Params { uint P[]; }; +layout(std430, binding = 2) buffer Arena { uint A[]; }; +layout(binding = 3) uniform sampler2D colourTex; // arg 176, 640x360 rgba16f +layout(binding = 5) uniform sampler2D exposureTex; // arg 200, 1x1 rgba16f +layout(binding = 8) uniform sampler2D albedoA; // arg 280, 640x360 rgba8unorm +layout(binding = 9) uniform sampler2D albedoB; // arg 288, 640x360 rgba8unorm +layout(binding = 11) uniform sampler2D decTex; // arg 448, 640x384 rgba16f +layout(rgba16f, binding = 12) uniform writeonly image2D outLog; // arg 224 +layout(rgba16f, binding = 13) uniform writeonly image2D outLin; // arg 240 + +#define F_LOG2E uintBitsToFloat(0x3FB8AA3Bu) +#define F_2LOG2E uintBitsToFloat(0x4038AA3Bu) +#define F_LOG10_2 uintBitsToFloat(0x3E9A209Bu) +#define F_0_36 uintBitsToFloat(0x3EB851ECu) +#define F_100 uintBitsToFloat(0x42C80000u) +#define F_1E4 uintBitsToFloat(0x38D1B717u) +#define F_0_005 uintBitsToFloat(0x3BA3D70Au) +#define F_0_1 uintBitsToFloat(0x3DCCCCCDu) +#define F_1_7 uintBitsToFloat(0x3FD9999Au) +#define F_0_17 uintBitsToFloat(0x3E2E147Bu) +#define F_0_95 uintBitsToFloat(0x3F733333u) +#define F_0_6 uintBitsToFloat(0x3F19999Au) +#define F_9 uintBitsToFloat(0x41102CB4u) + +// shared tile: 21 entries per row, up to 21x21 entries (render scale <= 1) +#define TILE_STRIDE 21 +// binary16 numbers, stored as such: half the shared-memory traffic of the 5x5 gather +shared f16vec4 tLog[TILE_STRIDE * TILE_STRIDE]; // (L_R, L_G, L_B, alpha = 0) log-domain colour +shared f16vec4 tSig[TILE_STRIDE * TILE_STRIDE]; // (s0, s1, s2, s3) sigmoid of arena bytes 0..3 +shared vec4 tAn[TILE_STRIDE * TILE_STRIDE]; // (p0, p1, p2, p3) anisotropy parameters + +float pf(uint idx) { return uintBitsToFloat(P[idx]); } + +// e4m3 byte -> float (exact); e4m3fn NaN for 0x7F/0xFF. +float e4m3(uint b) { + uint e = (b >> 3u) & 15u; + uint m = b & 7u; + if (e == 15u && m == 7u) return uintBitsToFloat(0x7FC00000u); + float v = (e == 0u) ? float(m) * uintBitsToFloat(0x3B000000u) + : (1.0 + float(m) * 0.125) * exp2(float(e) - 7.0); + return ((b & 128u) != 0u) ? -v : v; +} + +// round to nearest even binary16, returned as f32: the hardware conversion (RNE), kept from being folded back to x by +// OR-ing an opaque zero into its bits (RADV folds a register-only f32 -> f16 -> f32) +// Beyond the binary16 range NVIDIA's arithmetic keeps a finite rounded value: that case takes the explicit rounding. +uint16_t zero16; +float f16rnSlow(float x) { + if (x == 0.0 || isnan(x) || isinf(x)) return x; + int e; + frexp(x, e); + int ee = max(e - 11, -24); + float q = uintBitsToFloat(uint(ee + 127) << 23); + return roundEven(x / q) * q; +} +float f16rn(float x) { + if (abs(x) >= 65504.0) return f16rnSlow(x); + return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | zero16)); +} + +// Sigmoid as evaluated in binary16 on this lowering: r = f16(2^(-x*log2e)), den = f16(1+r), +// s = f16(1/den), with one Newton step for small s and the observed fix-ups of r. +float sig16(float e) { + float x = f16rn(e); + float r = f16rn(exp2(-x * F_LOG2E)); + uint hb = uint(float16BitsToUint16(float16_t(r))); + if (hb == 0x1F79u || hb == 0x25CFu) { + r = f16rn(r - 0.0009765625); + } else if (hb == 0xC13Bu) { + r = f16rn(r + 0.00006103515625); + } else if (hb == 0xC1EFu) { + r = f16rn(r + 0.000030517578125); + } + float d32 = f16rn(1.0 + r); + float rc = 1.0 / d32; + float s16 = f16rn(rc); + uint sb = uint(float16BitsToUint16(float16_t(s16))) & 0x7FFFu; + if (sb != 0u && sb < 0x8Fu) { + float ee = fma(-d32, s16, 1.0); + s16 = f16rn(fma(rc, ee, s16)); + } + return s16; +} + +// tanh as evaluated: polynomial for small |x|, exponential form otherwise. +float tanhApprox(float x) { + float ax = abs(x); + float x2 = x * x; + float q = fma(uintBitsToFloat(0x3C80F082u), x2, uintBitsToFloat(0xBD563CAEu)); + q = fma(q, x2, uintBitsToFloat(0x3E085941u)); + q = fma(q, x2, uintBitsToFloat(0xBEAAA9EDu)); + q = fma(q, x2, 0.0); + float small = fma(q, x, x); + float ex = exp2(ax * F_2LOG2E); + float big = fma(1.0 / (ex + 1.0), -2.0, 1.0); + if (ax >= F_9) big = 1.0; + float bigSigned = uintBitsToFloat((floatBitsToUint(x) & 0x80000000u) | floatBitsToUint(big)); + return (ax >= F_0_6) ? bigSigned : small; +} + +// round half away from zero +int roundAway(float v) { + return int(trunc(v + ((v < 0.0) ? -0.5 : 0.5))); +} + +void main() { + zero16 = uint16_t(P[0] >> 24u); // the output width is below 2^24 + // ---- parameters (byte offsets / 4) ---- + uint outW = P[0]; + uint outH = P[1]; + float sx = pf(6u); // 24: render/out scale x + float sy = pf(7u); // 28 + float ox = pf(10u); // 40: jitter offset x + float oy = pf(11u); // 44 + float expScale = pf(37u); // 148 + uint decMaxX = P[14] - 1u; // 56: dec0 width + uint decMaxY = P[15] - 1u; // 60: dec0 height + uint colMaxX = P[35]; // 140: colour/albedo max x + uint colMaxY = P[36]; // 144: colour/albedo max y + uint arenaBase = P[64]; // 256: arena offset + uint arenaPitch = P[4]; // 16: render width (640) + + // ---- tile origin / size in render space ---- + uvec2 wg = gl_WorkGroupID.xy; + float tileX = ox + fma(float(wg.x * 16u) + 0.5, sx, -0.5); + float tileY = oy + fma(float(wg.y * 16u) + 0.5, sy, -0.5); + int c0 = int(floor(tileX)) - 2; + int r0 = int(floor(tileY)) - 2; + int tw = int(ceil(16.0 * sx)) + 5; + int th = int(ceil(16.0 * sy)) + 5; + + // ---- stage 1: tile load ---- + uint total = uint(tw * th); + for (uint idx = gl_LocalInvocationIndex; idx < total; idx += 256u) { + int col = int(idx) % tw; + int row = int(idx) / tw; + int ix = clamp(c0 + col, 0, int(decMaxX)); + int iy = clamp(r0 + row, 0, int(decMaxY)); + + vec3 d = texelFetch(decTex, ivec2(ix, iy), 0).rgb; + int ax = min(ix, int(colMaxX)); + int ay = min(iy, int(colMaxY)); + vec3 alb = texelFetch(albedoA, ivec2(ax, ay), 0).rgb + + texelFetch(albedoB, ivec2(ax, ay), 0).rgb; + float albSum = (alb.r + alb.g) + alb.b; + if (albSum != 0.0) { + d.r *= max(alb.r, F_0_005); + d.g *= max(alb.g, F_0_005); + d.b *= max(alb.b, F_0_005); + } + vec3 lg; + lg.r = f16rn(F_0_36 * (log2(max(fma(d.r, F_100, 1.0), F_1E4)) * F_LOG10_2)); + lg.g = f16rn(F_0_36 * (log2(max(fma(d.g, F_100, 1.0), F_1E4)) * F_LOG10_2)); + lg.b = f16rn(F_0_36 * (log2(max(fma(d.b, F_100, 1.0), F_1E4)) * F_LOG10_2)); + + uint wordIdx = (arenaBase >> 2u) + uint(arenaPitch * uint(iy) + uint(ix)) * 2u; + uint lo = A[wordIdx]; + uint hi = A[wordIdx + 1u]; + vec4 sg = vec4(sig16(e4m3(lo & 0xFFu)), + sig16(e4m3((lo >> 8u) & 0xFFu)), + sig16(e4m3((lo >> 16u) & 0xFFu)), + sig16(e4m3(lo >> 24u))); + vec4 an = vec4(e4m3(hi & 0xFFu), + e4m3((hi >> 8u) & 0xFFu), + e4m3((hi >> 16u) & 0xFFu), + e4m3(hi >> 24u)); + + uint se = uint(row * TILE_STRIDE + col); + tLog[se] = f16vec4(lg, 0.0); + tSig[se] = f16vec4(sg); + tAn[se] = an; + } + barrier(); + + // ---- stage 2: per output pixel ---- + uvec2 pix = gl_GlobalInvocationID.xy; + if (pix.x >= outW || pix.y >= outH) return; + + float expo = texelFetch(exposureTex, ivec2(0, 0), 0).r; + float e41 = ((expo == 0.0) ? 1.0 : expo) * expScale; + + float fx = (float(pix.x) + 0.5) * sx; + float fy = (float(pix.y) + 0.5) * sy; + float cx = fx + ox; + float cy = fy + oy; + + // colour fetch: log of exposure-scaled colour at the nearest render texel + int ux = clamp(roundAway((fx - 0.5) + ox), 0, int(colMaxX)); + int uy = clamp(roundAway((fy - 0.5) + oy), 0, int(colMaxY)); + vec3 col = texelFetch(colourTex, ivec2(ux, uy), 0).rgb; + vec3 gcol; + gcol.r = F_0_36 * (log2(max(fma(col.r * e41, F_100, 1.0), F_1E4)) * F_LOG10_2); + gcol.g = F_0_36 * (log2(max(fma(col.g * e41, F_100, 1.0), F_1E4)) * F_LOG10_2); + gcol.b = F_0_36 * (log2(max(fma(col.b * e41, F_100, 1.0), F_1E4)) * F_LOG10_2); + + // centre entry, anisotropy + float flx = floor(cx); + float fly = floor(cy); + int ix0 = int(flx) - c0; + int iy0 = int(fly) - r0; + float dx0 = (flx + 0.5) - cx; + float dy0 = (fly + 0.5) - cy; + float ks = (sy >= 0.5) ? 1.0 : (0.5 / sy); + + vec4 an = tAn[iy0 * TILE_STRIDE + ix0]; + float sa = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.x)), F_1_7, F_0_17); + float sb = fma(1.0 / (1.0 + exp2(-F_LOG2E * an.y)), F_1_7, F_0_17); + float rho = F_0_95 * tanhApprox(an.z); + + float sa2 = sa * sa; + float sb2 = sb * sb; + float denom = (sa2 * sb2) * (1.0 - rho * rho); + float invDenom = -(1.0 / denom); + float cross2 = rho * (2.0 * sqrt(sa2 * sb2)); + + // 5x5 gather + float sumw = 0.0; + vec3 accL = vec3(0.0); + vec4 accS = vec4(0.0); + for (int c = -2; c <= 2; c++) { + for (int r = -2; r <= 2; r++) { + float dx = ks * (dx0 + float(c)); + float dy = ks * (dy0 + float(r)); + uint se = uint((iy0 + r) * TILE_STRIDE + (ix0 + c)); + precise float t_c = (cross2 * dx) * dy; + precise float t_b = dx * (sb2 * dx); + precise float t_a = dy * (sa2 * dy); + precise float num = (t_c - t_b) - t_a; + precise float Q = invDenom * num; + float w = exp2((Q * -0.5) * F_LOG2E); + sumw += w; + vec4 lv = vec4(tLog[se]); + vec4 sv = vec4(tSig[se]); + accL = fma(vec3(w), lv.rgb, accL); + accS = fma(vec4(w), sv, accS); + } + } + + float inv = 1.0 / sumw; + vec3 Lm = accL * inv; + float S0 = accS.x * inv; + float S1 = accS.y * inv; + float S2 = accS.z * inv; + + // blend weight + float q = F_0_1 / (sx * sy); + float t1 = sumw * q; + float t2 = S0 * t1; + float t3 = (sx * sy) * t2; + float t4 = 4.0 * t3; + float t5 = (sy < 0.5) ? t4 : t2; + float wgt = (t5 < 0.0) ? 0.0 : ((t5 > 1.0) ? 1.0 : t5); + + float fA = fma(1.0 - S1, wgt, S1); + float fB = fma(1.0 - S2, wgt, S2); + + // surface 224 (E): fA * D + (1 - fA) * gcol; D = Lm (f16 tile log channels), + // gcol = log of the exposure-scaled colour texel (PTX linear texel fetch path); D.a = 0 so aA = 1 - fA + vec3 cA = fma(Lm, vec3(fA), gcol * (1.0 - fA)); + float aA = 1.0 - fA; + float sumA = cA.r + cA.g + cA.b + aA; + if (isnan(sumA) || isinf(sumA)) { + cA = Lm; + aA = 0.0; + } + imageStore(outLog, ivec2(pix), vec4(f16rn(cA.r), f16rn(cA.g), f16rn(cA.b), f16rn(aA))); + + // surface 240 (E, PTX BB0_83): {R,G,B,1} of lin = (10^(clamp(c)/0.36)-1)/100/e41, + // c = fma(Lm, fB, gcol*(1-fB)); c = Lm when the sum is not finite; stored RNE f16 + vec3 cB = fma(Lm, vec3(fB), gcol * (1.0 - fB)); + float sumB = cB.r + cB.g + cB.b; + if (isnan(sumB) || isinf(sumB)) cB = Lm; + cB = clamp(cB, vec3(F_1E4), vec3(3.0)); + // PTX divides by constants and takes lg2(10) with the hardware approximations at run time; keep the + // constants opaque so the compiler cannot fold them with different rounding. + uint zero = P[0] >> 24u; + float c036 = uintBitsToFloat(0x3EB851ECu | zero); + float c100 = uintBitsToFloat(0x42C80000u | zero); + float lg10 = log2(uintBitsToFloat(0x41200000u | zero)); + float rc036 = 1.0 / c036; + float rc100 = 1.0 / c100; + float rce41 = 1.0 / e41; + precise vec3 lin = (exp2((cB * rc036) * lg10) - 1.0) * rc100 * rce41; + imageStore(outLin, ivec2(pix), vec4(float(float16_t(lin.r)), float(float16_t(lin.g)), float(float16_t(lin.b)), 1.0)); +} diff --git a/engine/rr/native_kernels.txt b/engine/rr/native_kernels.txt new file mode 100644 index 0000000..50903a5 --- /dev/null +++ b/engine/rr/native_kernels.txt @@ -0,0 +1,15 @@ +# SHA-256 of the PTX modules of the network kernels engine/rr's native shaders (swin_m.comp, dec0_m.comp, enc0_m) +# were written against and verified with: nvngx_dlssd.dll 310.7. engine/rr/dll_weights.py DLL --kernels prints a DLL's. +cuda_dldn_engine_swin_dec0_diamond_wallaby_kernel f9420970b6e29c1cd100057ba9eb56a0465a95a4aaeaa21dd355605b75160e5b +cuda_dldn_engine_swin_dec0_kernel 5e533151fd1d469bb344548b113e67b5fc798db807cebaf627864ce877f01242 +cuda_dldn_engine_swin_dec1_kernel ac293a9826ea7c8be44168698e79f122d0bb2cb23dbf55999a094876e04ce32f +cuda_dldn_engine_swin_dec2_kernel 282d8f093bf343c7a9080c7379ff933443943d503c7dea647535444711494528 +cuda_dldn_engine_swin_dec3_kernel f2325cc78cac7c3f20d4b3b5d6604b8a350774dd83642da10ab525bd3fd1e5b0 +cuda_dldn_engine_swin_dec4_kernel 9bb7cfeba4380b023f2a1b2db5fa311edd40d4e4fdce62af2fe9227a83c52213 +cuda_dldn_engine_swin_enc0_diamond_wallaby_kernel 35a2b205de3feb52dacfd54ebdcb2eda4dc75fe78a181a28687876f9a1d2a17d +cuda_dldn_engine_swin_enc0_kernel e0927a196eeffc93996a0c53337fe6fecc8a869145d832a02c59d9e12741db9d +cuda_dldn_engine_swin_enc1_kernel c9df416b891e4dddbeb6311ce110b8362c1c7f1d9e9efb2595b9b4ccae155fce +cuda_dldn_engine_swin_enc2_kernel 58702ea68b478da919a31cb9052fdb82eba3d7d1ec0d94ef2bb2827686a3b36a +cuda_dldn_engine_swin_enc3_kernel 7f07b4cf84adcdc6bfb367c00767b1f124bbdb1956f704a5dbc72f9e1d3fb9ac +cuda_dldn_engine_swin_enc4_kernel d1466bfa00c702d1a67820d1ad12e351fc4185d90695b21286c1d9523e772d6e +cuda_dldn_engine_swin_enc5_kernel 720c531da5552ff9a0304c8628aea46d9eaeae3c60f9ee0cdb4f94bebfe1a6e0 diff --git a/engine/rr/native_pack.py b/engine/rr/native_pack.py new file mode 100644 index 0000000..08d175e --- /dev/null +++ b/engine/rr/native_pack.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +"""Pack a denoiser stage's weights (engine/rr/enc1_ref.py --extract) for the native stage shaders. + +usage: native_pack.py WEIGHTS.npz OUT.bin OUT.h + +Every FP8 matrix W (K x N, rows in NVIDIA's k order) becomes 16x16 tiles, tile (kt, nt) at +256 * (kt * N/16 + nt) bytes, row-major inside the tile (stride 16): a B operand of the cooperative matrix loads. +Every f16 matrix is stored as is (row-major); an f16 bias vector (a 1-D entry whose name starts with 'b') as 16 +equal rows, so that it loads as an accumulator tile. OUT.h lists the byte offset of every entry as a #define, +for the shader's include. +""" +import sys + +import numpy as np + + +def tiles(w): + k, n = w.shape + assert k % 16 == 0 and n % 16 == 0, w.shape + return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1) + + +def main(): + src, out_bin, out_h = sys.argv[1:4] + w = dict(np.load(src)) + blob = bytearray() + defines = [] + for name in sorted(w): + a = w[name] + while len(blob) % 256: + blob.append(0) + off = len(blob) + if a.dtype == np.uint8: + blob += tiles(a).tobytes() + elif a.ndim == 1 and name.startswith('b'): + blob += np.tile(np.asarray(a, np.float16), (16, 1)).tobytes() + else: + blob += np.asarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {off}u // {a.dtype} {a.shape}') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write('// generated by engine/rr/native_pack.py\n' + '\n'.join(defines) + '\n') + print(f'{len(blob)} bytes, {len(w)} entries') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/phase_dump.py b/engine/rr/phase_dump.py new file mode 100644 index 0000000..2eca0b6 --- /dev/null +++ b/engine/rr/phase_dump.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +"""Dump every mma of a denoiser layer as matrices, from a ZLUDA replay of a captured launch: the oracle for porting +a layer to native shaders. + +usage: phase_dump.py CAPTURE STAGE FRAME BX,BY OUT.npz [--phases A-B] + +For one thread block, debug stores added to the PTX record every mma's A, B, C and D registers (each mma's last +execution) into padding behind the plane arena; the replay runs through ZLUDA (engine/rr/zluda_replay.py). The +fragments are then put back together with the m16n8k32 layouts (A row-major 16x32 e4m3, B column-major 32x8 e4m3, +C/D 16x8 f16) into one matrix per phase: phase P's mma are grouped into m-tiles by A fragment and into n-tiles by +B fragment, k-steps follow the C <- D chains. OUT.npz holds, per phase p and wave w: p{p}_w{w}_A (M x K, e4m3 codes), +_B (K x N codes), _C (M x N f16 of the first k-step), _D (M x N f16 of the last). +""" +import argparse +import os +import re +import subprocess +import sys +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +ROOT = HERE.parents[1] +sys.path.insert(0, str(HERE)) +from rr_capture import Capture # noqa: E402 + +REG = r'%r\d+' + + +def kernel_module(kernel): + for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')): + if f'.entry {kernel}(' in p.read_text(errors='replace'): + return p + raise SystemExit(f'no corpus module for {kernel}') + + +def mma_statements(lines, start, end): + """(first line, last line, D, A, B, C) of every mma in [start, end)""" + out, i = [], start + while i < end: + if 'mma.sync.aligned.m16n8k32' in lines[i]: + j, text = i, '' + while True: + text += lines[j] + if lines[j].rstrip().endswith(';'): + break + j += 1 + groups = [re.findall(REG, g) for g in re.findall(r'\{([^}]*)\}', text)] + out.append((i + 1, j + 1, *groups[:4])) + i = j + 1 + else: + i += 1 + return out + + +def phases_of(mmas, lines): + """maximal runs of mma separated only by blank or comment lines""" + phases, cur = [], [] + for m in mmas: + if cur: + gap = [l.strip() for l in lines[cur[-1][1]:m[0] - 1]] + if any(g and not g.startswith('//') for g in gap): + phases.append(cur) + cur = [] + cur.append(m) + if cur: + phases.append(cur) + return phases + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('stage') + ap.add_argument('frame', type=int) + ap.add_argument('block') + ap.add_argument('out') + ap.add_argument('--phases', help='first-last (1-based), default all') + o = ap.parse_args() + # the capture decides the kernel: a preset D capture runs cuda_dldn_engine_swin_{enc0,dec0}_diamond_wallaby_kernel + cap = Capture(os.path.expanduser(o.capture)) + launch = [l for l in cap.launches if f'swin_{o.stage}' in l.kernel][o.frame - 1] + kernel = launch.kernel + module = kernel_module(kernel) + lines = module.read_text().replace('\r', '').split('\n') + k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l) + e = next((i for i in range(k + 1, len(lines)) if '.entry ' in lines[i]), len(lines)) + phases = phases_of(mma_statements(lines, k, e), lines) + first, last = (int(x) for x in o.phases.split('-')) if o.phases else (1, len(phases)) + T = launch.block[0] * launch.block[1] * launch.block[2] + arena_bytes = len(launch.pre(launch.arena_alloc)) + # any argument pointing into the arena locates its base (enc0 has none at 8) + aptr = next(p for p in launch.pointers if p.alloc == launch.arena_alloc) + delta8, off8 = aptr.delta, aptr.offset + bx, by = (int(v) for v in o.block.split(',')) + # one probe slot per register per mma: after the mma, all of A, B, C, D (C read after the mma is still C + # unless C and D share registers, which the k-step chains never do for the first step) + slots, inserts = [], {} + pname = f'{kernel}_param_0' + for p in range(first, last + 1): + for (l0, l1, D, A, B, C) in phases[p - 1]: + code = [] + for role, regs in (('D', D), ('A', A), ('B', B), ('C', C)): + for r in regs: + s = len(slots) + slots.append((p, l0, role, r)) + code.append(f'mov.b32 %dv, {r}; mad.wide.u32 %da, %dt, 4, %db; ' + f'add.s64 %da, %da, {4 * T * s}; @%dq2 st.global.u32 [%da], %dv;') + inserts[l1] = ('{ .reg .b64 %da, %db; .reg .b32 %dv, %dt, %dx, %dy; .reg .pred %dq, %dq2;\n' + f'mov.u32 %dx, %ctaid.x; mov.u32 %dy, %ctaid.y; setp.eq.u32 %dq, %dx, {bx}; ' + f'setp.eq.and.u32 %dq2, %dy, {by}, %dq;\n' + f'mov.u32 %dt, %tid.z; mov.u32 %dx, %tid.y; mov.u32 %dy, %tid.x; ' + f'mad.lo.u32 %dt, %dt, {launch.block[1]}, %dx; mad.lo.u32 %dt, %dt, {launch.block[0]}, %dy;\n' + f'ld.param.u64 %db, [{pname}+{off8}]; cvta.to.global.u64 %db, %db; ' + f'add.s64 %db, %db, {arena_bytes - delta8};\n' + '\n'.join(code) + '\n}') + out = list(lines) + for l1 in sorted(inserts, reverse=True): + out[l1:l1] = [inserts[l1]] + work = ROOT / 'build/rr' / f'phase-{o.stage}' + work.mkdir(parents=True, exist_ok=True) + (work / 'probe.ptx').write_text('\n'.join(out)) + pad = 4 * T * len(slots) + env = dict(os.environ, D4R_ZLUDA_WMMA='1', D4R_ZLUDA_WMMA_FP8='1', D4R_ZLUDA_WMMA_FP8_NATIVE='1', + D4R_ZLUDA_IMPLICIT_MAX_BLOCK='256') + cmd = ['systemd-run', '--user', '--scope', '-q', '-p', 'MemoryMax=24G', '-p', 'MemorySwapMax=0'] + cmd += [f'--setenv={k_}={v_}' for k_, v_ in env.items() if k_.startswith('D4R_ZLUDA')] + r = subprocess.run(cmd + [sys.executable, str(HERE / 'zluda_replay.py'), os.path.expanduser(o.capture), + f'swin_{o.stage}' + ('_diamond' if 'diamond' in kernel else ''), str(o.frame), str(work), + '--ptx', str(work / 'probe.ptx'), + '--pad', str(pad)], capture_output=True, text=True) + if 'failed' in r.stdout + r.stderr: + sys.exit(r.stdout[-2000:] + r.stderr[-3000:]) + raw = np.frombuffer((work / 'arena.out.bin').read_bytes()[arena_bytes:], np.uint32).reshape(len(slots), T) + val = {(p, l0, role, r): raw[s] for s, (p, l0, role, r) in enumerate(slots)} + result = {} + waves = T // 32 + for p in range(first, last + 1): + mm = phases[p - 1] + # k-step chains: C of a later mma = D of an earlier one + dset = {tuple(m[2]): n for n, m in enumerate(mm)} + prev = {n: dset.get(tuple(m[5])) for n, m in enumerate(mm)} + step = {} + for n in range(len(mm)): + s_, q = 0, n + while prev[q] is not None: + q = prev[q] + s_ += 1 + step[n] = s_ + ksteps = max(step.values()) + 1 + a_keys = list(dict.fromkeys(tuple(m[3]) for m in mm)) + b_keys = list(dict.fromkeys(tuple(m[4]) for m in mm)) + # m-tile of an A fragment: its index among the A fragments of the same k-step + a_tile = {} + for s_ in range(ksteps): + for t_, key in enumerate(dict.fromkeys(tuple(m[3]) for n, m in enumerate(mm) if step[n] == s_)): + a_tile[(s_, key)] = t_ + b_tile = {} + for s_ in range(ksteps): + for t_, key in enumerate(dict.fromkeys(tuple(m[4]) for n, m in enumerate(mm) if step[n] == s_)): + b_tile[(s_, key)] = t_ + mt = max(t for (s_, _), t in a_tile.items()) + 1 + nt = max(t for (s_, _), t in b_tile.items()) + 1 + for w in range(waves): + lanes = np.arange(32) + 32 * w + A = np.zeros((16 * mt, 32 * ksteps), np.uint8) + B = np.zeros((32 * ksteps, 8 * nt), np.uint8) + C = np.zeros((16 * mt, 8 * nt), np.uint16) + D = np.zeros((16 * mt, 8 * nt), np.uint16) + for n, (l0, l1, Dr, Ar, Br, Cr) in enumerate(mm): + s_ = step[n] + i, j = a_tile[(s_, tuple(Ar))], b_tile[(s_, tuple(Br))] + for lane in range(32): + g, t = lane >> 2, lane & 3 + av = [int(val[(p, l0, 'A', r)][lanes[lane]]) for r in Ar] + bv = [int(val[(p, l0, 'B', r)][lanes[lane]]) for r in Br] + for q in range(4): # A regs: (g, k 4t..), (g+8, 4t..), (g, 16+4t..), (g+8, 16+4t..) + row = 16 * i + g + 8 * (q & 1) + col = 32 * s_ + 4 * t + 16 * (q >> 1) + for byte in range(4): + A[row, col + byte] = (av[q] >> (8 * byte)) & 0xff + for q in range(2): # B regs: column g, k 4t.. and 16+4t.. + for byte in range(4): + B[32 * s_ + 4 * t + 16 * q + byte, 8 * j + g] = (bv[q] >> (8 * byte)) & 0xff + for q in range(2): # C/D regs: row g (+8), columns 2t, 2t+1 + row = 16 * i + g + 8 * q + if s_ == 0: + cv = int(val[(p, l0, 'C', Cr[q])][lanes[lane]]) + C[row, 8 * j + 2 * t] = cv & 0xffff + C[row, 8 * j + 2 * t + 1] = cv >> 16 + if s_ == ksteps - 1: + dv = int(val[(p, l0, 'D', Dr[q])][lanes[lane]]) + D[row, 8 * j + 2 * t] = dv & 0xffff + D[row, 8 * j + 2 * t + 1] = dv >> 16 + for name, mat in (('A', A), ('B', B), ('C', C), ('D', D)): + result[f'p{p}_w{w}_{name}'] = mat + print(f'phase {p}: {len(mm)} mma, {mt} m-tiles x {nt} n-tiles x {ksteps} k-steps -> ' + f'{16 * mt}x{8 * nt}x{32 * ksteps}, lines {mm[0][0]}-{mm[-1][1]}') + np.savez_compressed(o.out, **result) + + +if __name__ == '__main__': + main() diff --git a/engine/rr/ptx2rr.py b/engine/rr/ptx2rr.py new file mode 100755 index 0000000..836fa5d --- /dev/null +++ b/engine/rr/ptx2rr.py @@ -0,0 +1,729 @@ +#!/usr/bin/env python3 +"""Generate a GLSL compute shader for one Ray Reconstruction (DLSS-D) PTX kernel. + + python3 engine/rr/ptx2rr.py MODULE.ptx KERNEL OUT.comp [--block X,Y,Z] + +The translation core is vulkan/ptx2glsl.py. This module subclasses its translator and +implements the stage contract of engine/rr/README.md: binding 0 = parameter block, +binding 1 = weight buffer (argument 40), binding 2 = plane arena, then textures and +surfaces in ascending argument offset. One PTX block is one workgroup of X*Y*Z invocations. + +The generated shader is derived from the user's own DLL and is model data: write it to the +model/build directory only; do not commit or redistribute it. +""" +import argparse +import importlib.util +import json +import re +import struct +import sys +from pathlib import Path + +import numpy as np + +ROOT = Path(__file__).resolve().parents[2] +_spec = importlib.util.spec_from_file_location('ptx2glsl', ROOT / 'vulkan' / 'ptx2glsl.py') +base = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(base) + +# Workgroup shapes of the 15 Swin/HKPN kernels, from the captured launch manifests. +BLOCKS = { + 'enc0': (32, 1, 1), 'enc1': (32, 1, 1), 'enc2': (32, 1, 4), 'enc3': (32, 1, 4), 'enc4': (32, 1, 8), + 'enc5': (32, 1, 8), 'dec0': (32, 1, 1), 'dec1': (32, 1, 2), 'dec2': (32, 1, 4), 'dec3': (32, 1, 4), + 'dec4': (32, 1, 8), 'hkpn': (16, 16, 1), +} + + +def default_block(kernel): + m = re.search(r'swin_(enc|dec)(\d)', kernel) + key = f'{m.group(1)}{m.group(2)}' if m else ('hkpn' if 'hkpn_output' in kernel else None) + if key not in BLOCKS: + raise SystemExit(f'no default block shape for {kernel}; pass --block X,Y,Z') + return BLOCKS[key] + + +def braces(tok): + s = tok.strip() + if not (s.startswith('{') and s.endswith('}')): + raise base.Err(f'expected a {{...}} vector operand, got {tok}') + return [x.strip() for x in s[1:-1].split(',')] + + +def mem_shape(o): + """(bytes per element, element count) of ld/st.global modifiers o.""" + count = 2 if 'v2' in o else 4 if 'v4' in o else 1 + size = int(re.sub(r'\D', '', o[-1])) // 8 + if size not in (2, 4): + raise base.Err(f'unsupported global access width in {".".join(o)}') + return size, count + + +RR_RUNTIME = r''' +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 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)); +} + +// Global memory: weight buffer (binding 1, read-only) and plane arena (binding 2). Addresses are byte offsets. +#ifdef OOB_REPORT +// debugging (ptx2rr.py --oob-report): per buffer (0 weights read, 1 arena read, 2 arena write) a count and the lowest +// and highest out-of-range byte address, in dbg.d[3k .. 3k+2] +void OOB(uint k, uint a, uint size) { if (a + 2u > size) { atomicAdd(dbg.d[3u * k], 1u); atomicMin(dbg.d[3u * k + 1u], a); atomicMax(dbg.d[3u * k + 2u], a); } } +#define CHECKW(a) OOB(0u, a, uint(wbuf.wh.length()) * 2u) +#define CHECKR(a) OOB(1u, a, uint(abuf.ah.length()) * 2u) +#define CHECKS(a) OOB(2u, a, uint(abuf.ah.length()) * 2u) +#else +#define CHECKW(a) +#define CHECKR(a) +#define CHECKS(a) +#endif +uint GW16(uint a) { CHECKW(a); return uint(wbuf.wh[a >> 1]); } +uint GW32(uint a) { CHECKW(a + 2u); return uint(wbuf.wh[a >> 1]) | (uint(wbuf.wh[(a >> 1) + 1u]) << 16); } +uint GA16(uint a) { CHECKR(a); return uint(abuf.ah[a >> 1]); } +uint GA32(uint a) { CHECKR(a + 2u); return uint(abuf.ah[a >> 1]) | (uint(abuf.ah[(a >> 1) + 1u]) << 16); } +void SA16(uint a, uint v) { CHECKS(a); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); } +void SA32(uint a, uint v) { CHECKS(a + 2u); abuf.ah[a >> 1] = uint16_t(v & 0xffffu); abuf.ah[(a >> 1) + 1u] = uint16_t(v >> 16); } + +// Surface stores with .zero out-of-bounds semantics. +void sust_p4(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); +} +void sust_p1(writeonly image2D s, int x, int y, float v) { sust_p4(s, x, y, vec4(v, 0.0, 0.0, 0.0)); } +// Binary stores write the halves' bits. RADV folds float(float16_t(x)) back to x when the half came from a +// cvt.rn.f16.f32 in a register, and the image store then truncates x instead (the output stage's history surface +// was one binary16 ulp low in half its values, which the next frame's encoder 0 amplifies): OR-ing in a zero the +// compiler cannot see (the dispatch's z count is 1) keeps the conversion's rounding. +#define SUST_Z (gl_NumWorkGroups.z - 1u) +void sust_b2(writeonly image2D s, int x, int y, uint a, uint b) +{ + sust_p4(s, x >> 2, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), 0.0, 0.0)); +} +void sust_b4(writeonly image2D s, int x, int y, uint a, uint b, uint c, uint d) +{ + sust_p4(s, x >> 3, y, vec4(float(H(a | SUST_Z)), float(H(b | SUST_Z)), float(H(c | SUST_Z)), float(H(d | SUST_Z)))); +} + +// shfl.sync.up/down within a 32-lane warp (ZLUDA's SHFL_SYNC_IMPL). in_bounds is the predicate output. +uint shfl_dir(uint v, uint delta, uint opts, bool down, out bool in_bounds) +{ + int section_mask = int((opts >> 8) & 31u), warp_end = int(opts & 31u); + int self = int(gl_SubgroupInvocationID & 31u); + int d = int(delta & 31u); + int idx = down ? self + d : self - d; + int sub_end = (section_mask & self) | (~section_mask & (down ? warp_end : 0)); + bool oob = down ? idx > sub_end : idx < sub_end; + if (oob) idx = self; + in_bounds = !oob; + return subgroupShuffle(v, uint(idx) & 31u); +} + +// shfl.sync.idx (PTX ISA): j = minLane | (b & ~segmask), in bounds while j <= maxLane, otherwise the lane's own value +uint shfl_idx(uint v, uint b, uint opts, out bool in_bounds) +{ + uint mask = (opts >> 8) & 31u, lane = gl_SubgroupInvocationID & 31u; + uint maxLane = (lane & mask) | (opts & 31u & ~mask), j = (lane & mask) | (b & 31u & ~mask); + in_bounds = j <= maxLane; + return subgroupShuffle(v, in_bounds ? j : lane); +} +// movmatrix.sync.trans.aligned.m8n8.b16 on one b32 fragment register per lane: lane (g = lane>>2, t = lane&3) +// holds row g, columns 2t and 2t+1. The transpose gathers (row 2t, col g) and (row 2t+1, col g). +uint movm_trans(uint src) +{ + // through the subgroup's mma scratch rather than two shuffles: RADV miscompiled the shuffle form in the large + // stages (enc1's movmatrix results were wrong unless a store followed them; ZLUDA replay, 2026-10-11) + const uint base = gl_SubgroupID * 384u + 256u; + uint lane = gl_SubgroupInvocationID & 31u; + uint g = lane >> 2u, t = lane & 3u; + // order this operation's writes after every lane's reads of the previous one's results (the same scratch) + memoryBarrierShared(); subgroupBarrier(); + mms[base + lane] = src; + memoryBarrierShared(); subgroupBarrier(); + uint w1 = mms[base + 8u * t + (g >> 1u)]; + uint w2 = mms[base + 8u * t + 4u + (g >> 1u)]; + uint s = 16u * (g & 1u); + return ((w1 >> s) & 0xffffu) | (((w2 >> s) & 0xffffu) << 16); +} + +// PTX min/max: a NaN operand yields the other one (GLSL leaves NaN undefined, so the result depended on codegen) +float16_t pmin(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); } +float16_t pmax(float16_t a, float16_t b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); } +f16vec2 pmin2(f16vec2 a, f16vec2 b) { return f16vec2(pmin(a.x, b.x), pmin(a.y, b.y)); } +f16vec2 pmax2(f16vec2 a, f16vec2 b) { return f16vec2(pmax(a.x, b.x), pmax(a.y, b.y)); } +float pminf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : min(a, b); } +float pmaxf(float a, float b) { return isnan(a) ? b : isnan(b) ? a : max(a, b); } +// e4m3 <-> f16 (PTX cvt semantics; RNE, satfinite; NaN -> 0x7f e4m3 / 0x7e00 f16). +uint e4m3_to_h(uint b) +{ + uint e = (b >> 3) & 15u, m = b & 7u, s = (b & 0x80u) << 8; + if (e == 15u && m == 7u) return s | 0x7e00u; + float f = e == 0u ? ldexp(float(m), -9) : ldexp(float(8u + m), int(max(e, 1u)) - 10); + return s | HU(float16_t(f)); +} +#ifdef SOFTWARE_E4M3 +f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h(v & 0xffu) | (e4m3_to_h((v >> 8) & 0xffu) << 16)); } +#else +// the hardware conversion: exact for every non-NaN code (2026-10-12); NaN keeps e4m3_to_h's sign | 0x7e00 +uint e4m3_to_h_hw(uint b) +{ + if ((b & 0x7fu) == 0x7fu) return ((b & 0x80u) << 8) | 0x7e00u; + return uint(float16BitsToUint16(float16_t(uintBitsToFloate4m3EXT(uint8_t(b))))); +} +f16vec2 e4m3x2_to_f16x2(uint v) { return H2(e4m3_to_h_hw(v & 0xffu) | (e4m3_to_h_hw((v >> 8) & 0xffu) << 16)); } +#endif +uint h2e4m3(float f, bool relu) +{ + if (isnan(f)) return 0x7fu; + if (relu && !(f > 0.0)) return 0u; + uint s = f < 0.0 ? 0x80u : 0u; + float a = abs(f); + if (a >= 448.0) return s | 0x7eu; + if (a == 0.0) return s; + int ex; + frexp(a, ex); + int qe = max(ex - 4, -9); + uint n = uint(roundEven(ldexp(a, -qe))); + uint bits; + if (n < 8u) bits = n; + else { int p = findMSB(n); bits = (uint(qe + p + 7) << 3) | ((n >> uint(p - 3)) & 7u); } + return s | min(bits, 0x7eu); +} +#ifdef SOFTWARE_E4M3 +uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3(float(h.x), relu) | (h2e4m3(float(h.y), relu) << 8); } +#else +// cvt.rn.satfinite(.relu).e4m3x2.f16x2 on the hardware conversion: round to nearest even after clamping to +-448 +// (satfinite), NaN to 0x7f. Equal to the software h2e4m3 for all 65536 f16 inputs (2026-10-12) and about two thirds +// of a stage's time faster (dec0: 2.85 -> 1.03 ms). +uint h2e4m3_hw(float16_t h, bool relu) +{ + if (isnan(h)) return 0x7fu; + if (relu) h = h > float16_t(0.0) ? h : float16_t(0.0); + return uint(floate4m3BitsToUintEXT(floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0))))); +} +uint f16x2_to_e4m3x2(f16vec2 h, bool relu) { return h2e4m3_hw(h.x, relu) | (h2e4m3_hw(h.y, relu) << 8); } +#endif + +// mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 on the FP8 cooperative matrices, with ZLUDA's lowering of +// it (engine/layer_m.comp's TEXL): C enters the first of two 16x16x16 products, the sum stays f32 and is rounded to +// binary16 once. The operands pass through the subgroup's own scratch (MMS): NVIDIA's fragments are written in +// place (A row-major 16x32, B column-major 32x8 in a 16-column tile, C/D row-major f16 16x16), then loaded as +// whole matrices. Lane l = 4g + t holds A rows g and g+8 (k 4t.. and 16+4t..), B column g (same k), C/D row g +// and g+8, columns 2t and 2t+1. +#define MMA coopmat +#define MMB coopmat +#define MMF coopmat +#define MMH coopmat +void mma_e4m3_f16(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1, out uint d0, out uint d1) +{ + const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u; + const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u; + // order this operation's writes after every lane's reads of the previous one's results (the same scratch) + memoryBarrierShared(); subgroupBarrier(); + mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1; + mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3; + mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1; + mms[B + (g + 8u) * 8u + t] = 0u; mms[B + (g + 8u) * 8u + 4u + t] = 0u; + mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1; + mms[C + g * 8u + 4u + t] = 0u; mms[C + (g + 8u) * 8u + 4u + t] = 0u; + memoryBarrierShared(); subgroupBarrier(); + MMA x0, x1; MMB y0, y1; MMH c; + coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor); + MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c))); + memoryBarrierShared(); subgroupBarrier(); + coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor); + memoryBarrierShared(); subgroupBarrier(); + d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t]; +} +// Two m16n8k32 products with the same A fragment (NVIDIA issues every one of them in such adjacent pairs): one 16x16 +// product, the first's B, C and D in columns 0-7, the second's in columns 8-15. +void mma_e4m3_f16_pair(uint a0, uint a1, uint a2, uint a3, uint b0, uint b1, uint c0, uint c1, + uint e0, uint e1, uint f0, uint f1, out uint d0, out uint d1, out uint h0, out uint h1) +{ + const uint base = gl_SubgroupID * 384u, A = base, B = base + 128u, C = base + 256u; + const uint lane = gl_SubgroupInvocationID & 31u, g = lane >> 2u, t = lane & 3u; + memoryBarrierShared(); subgroupBarrier(); + mms[A + g * 8u + t] = a0; mms[A + (g + 8u) * 8u + t] = a1; + mms[A + g * 8u + 4u + t] = a2; mms[A + (g + 8u) * 8u + 4u + t] = a3; + mms[B + g * 8u + t] = b0; mms[B + g * 8u + 4u + t] = b1; + mms[B + (g + 8u) * 8u + t] = e0; mms[B + (g + 8u) * 8u + 4u + t] = e1; + mms[C + g * 8u + t] = c0; mms[C + (g + 8u) * 8u + t] = c1; + mms[C + g * 8u + 4u + t] = f0; mms[C + (g + 8u) * 8u + 4u + t] = f1; + memoryBarrierShared(); subgroupBarrier(); + MMA x0, x1; MMB y0, y1; MMH c; + coopMatLoad(x0, mms, A, 8u, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(x1, mms, A + 4u, 8u, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(y0, mms, B, 8u, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(y1, mms, B + 4u, 8u, gl_CooperativeMatrixLayoutColumnMajor); + coopMatLoad(c, mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor); + MMF acc = coopMatMulAdd(x1, y1, coopMatMulAdd(x0, y0, MMF(c))); + memoryBarrierShared(); subgroupBarrier(); + coopMatStore(MMH(acc), mms, C, 8u, gl_CooperativeMatrixLayoutRowMajor); + memoryBarrierShared(); subgroupBarrier(); + d0 = mms[C + g * 8u + t]; d1 = mms[C + (g + 8u) * 8u + t]; + h0 = mms[C + g * 8u + 4u + t]; h1 = mms[C + (g + 8u) * 8u + 4u + t]; +} +''' + + +def statement_lines(ptx, kernel): + """the source line of every statement of KERNEL, in the order vulkan/ptx2glsl.py's tokenizer yields them""" + text = re.sub(r'//[^\n]*', '', ptx.replace('\r', '')) + m = re.search(r'\.entry\s+' + re.escape(kernel) + r'\s*\((.*?)\)', text, re.S) + start = text.index('{', m.end()) + depth, k = 0, start + while True: + if text[k] == '{': + depth += 1 + elif text[k] == '}': + depth -= 1 + if depth == 0: + break + k += 1 + lines, i, line = [], start + 1, text.count('\n', 0, start + 1) + 1 + while i < k: + c = text[i] + if c.isspace(): + line += c == '\n' + i += 1 + continue + if c in '{}': + i += 1 + continue + lm = re.match(r'(\$?[A-Za-z_][\w$]*):', text[i:i + 200]) + if lm: + i += lm.end() + continue + j = text.index(';', i) + lines.append(line) + line += text.count('\n', i, j) + i = j + 1 + return lines + + +def base_helper(name): + """one helper function of vulkan/ptx2glsl.py's runtime, with its leading comment""" + lines = base.RUNTIME.split('\n') + start = next(i for i, l in enumerate(lines) if re.match(r'\w+ ' + name + r'\(', l)) + while lines[start - 1].startswith('//'): + start -= 1 + end = start + if not lines[start].rstrip().endswith('}'): # one-line helpers end on their own line + while lines[end] != '}': + end += 1 + return '\n'.join(lines[start:end + 1]) + '\n' + + +RR_RUNTIME += base_helper('prmt') + base_helper('shfl_bfly').replace( + 'return subgroupShuffle(v, (gl_SubgroupInvocationID & 32u) | uint(idx));', 'return subgroupShuffle(v, uint(idx) & 31u);') + + +class RRT(base.T): + def __init__(self, ptx, kernel, block): + bx, by, bz = block + if bx * by * bz % 32: + raise base.Err('block size must be a multiple of the 32-lane subgroup') + self.bx, self.by, self.bz = bx, by, bz + self.sreg = { + '%tid.x': f'(gl_LocalInvocationIndex % {bx}u)', + '%tid.y': f'((gl_LocalInvocationIndex / {bx}u) % {by}u)', + '%tid.z': f'(gl_LocalInvocationIndex / {bx * by}u)', + '%ntid.x': f'{bx}u', '%ntid.y': f'{by}u', '%ntid.z': f'{bz}u', + '%ctaid.x': 'gl_WorkGroupID.x', '%ctaid.y': 'gl_WorkGroupID.y', '%ctaid.z': 'gl_WorkGroupID.z', + '%nctaid.x': 'gl_NumWorkGroups.x', '%nctaid.y': 'gl_NumWorkGroups.y', '%nctaid.z': 'gl_NumWorkGroups.z', + '%laneid': 'gl_SubgroupInvocationID', + 'WARP_SZ': '32u', + } + self.origin = {} + self.f64 = {} # %fd registers: the kernels only load f64 immediates into them, to round them to f16 + self.probes = {} # debugging: PTX line -> registers stored to the debug buffer after it executes + self.probe_block = (0, 0) + self.probe_slots = 0 + self.oob = False + # every register precise, the integer ones too: they carry f16x2 and f32 bit patterns between operations, and + # without it the compiler may contract a product and a sum across them depending on scheduling (results then + # changed with unrelated code and with the device's robustness features) + self.precise_all = True + self.stmt_lines = statement_lines(ptx, kernel) + self.stmt_index = 0 + super().__init__(ptx, kernel) + self.pair_mma() + self.origins() + + def stmt(self, s): + line = self.stmt_lines[self.stmt_index] if self.stmt_index < len(self.stmt_lines) else None + self.stmt_index += 1 + super().stmt(s) + for reg in self.probes.get(line, ()): + # the debug buffer: 32-bit words [slot][local invocation], the last execution of the statement + n = self.reg(reg) + g = self.regs[n] + v = (f'floatBitsToUint({n})' if g == 'float' else f'uint({n})' if g in ('uint', 'uint64_t', 'bool') else + f'HU({n})' if g == 'float16_t' else f'H2U({n})') + bx, by = self.probe_block + self.E(f'if (gl_WorkGroupID.xy == uvec2({bx}u, {by}u)) dbg.d[{self.probe_slots}u * (gl_WorkGroupSize.x * gl_WorkGroupSize.y) + ' + f'gl_LocalInvocationIndex] = {v};') + self.probe_slots += 1 + + def pair_mma(self): + """merge adjacent mma statements with the same A fragment and independent operands into one paired product; + the second statement becomes a no-op so statement numbering (probes) is unchanged""" + stmts = [k for k, (kind, _) in enumerate(self.toks) if kind == 'stmt'] + n = 0 + while n + 1 < len(stmts): + i, j = stmts[n], stmts[n + 1] + x, y = self.toks[i][1], self.toks[j][1] + if x.startswith('mma.sync.aligned.m16n8k32') and y.startswith('mma.sync.aligned.m16n8k32'): + gx = [braces(t) for t in base.split_operands(x.split(None, 1)[1])] + gy = [braces(t) for t in base.split_operands(y.split(None, 1)[1])] + if gx[1] == gy[1] and not set(gx[0]) & set(gy[2] + gy[3] + gy[1]): + ops = ', '.join('{' + ', '.join(g) + '}' for g in (gx[0], gy[0], gx[1], gx[2], gx[3], gy[2], gy[3])) + self.toks[i] = ('stmt', 'mma.pair.e4m3 ' + ops) + self.toks[j] = ('stmt', 'mma.paired.noop') + n += 2 + continue + n += 1 + + # ------------------------------------------------ SREG remap onto the flat workgroup + def is_reg(self, tok): + return tok.strip() not in self.sreg and super().is_reg(tok) + + def val(self, tok, t): + tok = tok.strip() + if re.fullmatch(r'(0[xX][0-9a-fA-F]+|-?\d+)[uU]', tok): + tok = tok[:-1] + tok = re.sub(r'^0X', '0x', tok) + if tok in self.sreg: + e = self.sreg[tok] + return f'int({e})' if t == 's32' else f'uint64_t({e})' if t in ('u64', 'b64') else f'uint({e})' + return super().val(tok, t) + + # ------------------------------------------------ global memory origin: weight buffer or arena + def origins(self): + stmts = [] + 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) + stmts.append((parts[0], base.split_operands(parts[1]) if len(parts) > 1 else [])) + org = self.origin + changed = True + while changed: + changed = False + for op, a in stmts: + if not a or not re.fullmatch(r'%rd\d+', a[0]) or op.startswith(('ld.', 'st.', 'bra')): + if op != 'ld.param.u64': + continue + if op == 'ld.param.u64': + # the argument's true offset (the load may go through a register based inside the block) + new = {'w' if self.param_off(a[1]) == 40 else 'a'} + else: + new = set() + for x in a[1:]: + if re.fullmatch(r'%rd\d+', x): + new |= org.get(x, set()) + cur = org.setdefault(a[0], set()) + if new - cur: + cur |= new + changed = True + + def gorigin(self, reg): + s = self.origin.get(reg, set()) + if len(s) != 1: + raise base.Err(f'global pointer {reg} has origin {sorted(s)} (need exactly one of w/a)') + return next(iter(s)) + + def gaddr(self, tok): + m = re.fullmatch(r'\[\s*(%rd\d+)\s*(?:\+\s*(\d+))?\s*\]', tok.strip()) + if not m: + raise base.Err(f'unsupported global address {tok}') + return m.group(1), int(m.group(2) or 0) + + def gbyte(self, reg, off): + return f'uint({self.val(reg, "u64")}) + {off}u' + + def ld(self, o, a): + if 'global' not in o: + return super().ld(o, a) + size, count = mem_shape(o) + reg, off = self.gaddr(a[1]) + fn = 'GW' if self.gorigin(reg) == 'w' else 'GA' + dsts = braces(a[0]) if count > 1 else [a[0]] + for i, d in enumerate(dsts): + self.A(d, f'{fn}{size * 8}({self.gbyte(reg, off + i * size)})', 'u16' if size == 2 else 'u32') + + def st(self, o, a): + if 'global' not in o: + return super().st(o, a) + size, count = mem_shape(o) + reg, off = self.gaddr(a[0]) + if self.gorigin(reg) != 'a': + raise base.Err('store through the read-only weight pointer') + vals = braces(a[1]) if count > 1 else [a[1]] + for i, v in enumerate(vals): + if size == 2: + self.E(f'SA16({self.gbyte(reg, off + i * size)}, {self.val(v, "u16")});') + else: + self.E(f'SA32({self.gbyte(reg, off + i * size)}, {self.val(v, "u32")});') + + # ------------------------------------------------ textures and surfaces + def handle_off(self, h): + off = self.handle.get(h) + if off is None: + raise base.Err(f'{h} is not a texture or surface handle loaded from the parameter block') + return off + + def set_surface(self, off, fmt): + old = self.surfaces.get(off) + if old is not None and old != fmt: + raise base.Err(f'surface {off} used with formats {old} and {fmt}') + self.surfaces[off] = fmt + + def tex(self, op, a): + o = op.split('.')[1:] + dst = braces(a[0]) + h, c = self.texref(a[1]) + off = self.handle_off(h) + self.textures.add(off) + if 'base' in o and o[-1] == 's32': + # integer coordinates still go through the texture's address mode in CUDA (clamp for every denoiser + # texture); texelFetch does not clamp, and what it returns out of range depends on the device's robustness + fetch = (f'texelFetch(T{off}, clamp(ivec2({self.val(c[0], "s32")}, {self.val(c[1], "s32")}), ivec2(0), ' + f'textureSize(T{off}, 0) - 1), 0)') + elif 'base' in o and o[-1] == 'f32': + fetch = f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), 0.0)' + elif 'level' in o and o[-1] == 'f32': + fetch = (f'textureLod(T{off}, vec2({self.val(c[0], "f32")}, {self.val(c[1], "f32")}), ' + f'{self.val(a[2], "f32")})') + else: + raise base.Err(f'unsupported texture op {op}') + f16 = 'f16' in o + self.E(f'{{ vec4 t_ = {fetch};') + for i, d in enumerate(dst): + self.A(d, f'float16_t(t_[{i}])' if f16 else f't_[{i}]', 'f16' if f16 else 'float') + self.E('}') + + def sust(self, op, a): + h, c = self.texref(a[0]) + off = self.handle_off(h) + v = braces(a[1]) + x, y = self.val(c[0], 's32'), self.val(c[1], 's32') + S = f'S{off}' + if op == 'sust.p.2d.v4.b32.zero': + self.set_surface(off, 'rgba16f') + self.E(f'sust_p4({S}, {x}, {y}, vec4({", ".join(self.val(q, "f32") for q in v)}));') + elif op == 'sust.p.2d.b32.zero': + self.set_surface(off, 'r16f') + self.E(f'sust_p1({S}, {x}, {y}, {self.val(v[0], "f32")});') + elif op == 'sust.b.2d.v2.b16.zero': + self.set_surface(off, 'rg16f') + self.E(f'sust_b2({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});') + elif op == 'sust.b.2d.v4.b16.zero': + self.set_surface(off, 'rgba16f') + self.E(f'sust_b4({S}, {x}, {y}, {", ".join(self.val(q, "u16") for q in v)});') + else: + raise base.Err(f'unsupported surface op {op}') + + # ------------------------------------------------ warp-level ops + def mma(self, a): + d, x, b, c = (braces(t) for t in a) + if not (len(d) == 2 and len(x) == 4 and len(b) == 2 and len(c) == 2): + raise base.Err('mma operand shape') + args = ', '.join(self.val(r, 'u32') for r in x + b + c) + self.E('{ uint md0_, md1_;') + self.E(f'mma_e4m3_f16({args}, md0_, md1_);') + self.A(d[0], 'md0_', 'u32') + self.A(d[1], 'md1_', 'u32') + self.E('}') + + def shfl_dir(self, op, a): + down = op.split('.')[2] == 'down' + d = a[0].split('|') + self.E('{ bool ib_;') + self.A(d[0], f'shfl_dir({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, ' + f'{"true" if down else "false"}, ib_)', 'u32') + if len(d) > 1: + self.A(d[1], 'ib_', 'bool') + self.E('}') + + def cvt_to_e4m3(self, op, a): + relu = 'relu' in op.split('.') + self.A(a[0], f'f16x2_to_e4m3x2({self.val(a[1], "f16x2")}, {"true" if relu else "false"})', 'u16') + + def cvt(self, op, o, a): + types = [t for t in o if re.match(r'^(f|s|u|b)\d+(x2)?$', t)] + if types[:2] in (['u32', 'u64'], ['s32', 's64'], ['u32', 's64'], ['s32', 'u64']): + return self.A(a[0], f'uint({self.val(a[1], "u64")})', 'u32') + return super().cvt(op, o, a) + + def f64_op(self, op, a): + if op == 'mov.f64' and re.fullmatch(r'0[dD][0-9A-Fa-f]{16}', a[1].strip()): + self.f64[a[0].strip()] = struct.unpack('= 32u ? 0u : ({x}) << ({n}))', 'u32') + if parts[1] == 's32': + return self.A(a[0], f'uint(({x}) >> min({n}, 31u))', 'u32') + return self.A(a[0], f'(({n}) >= 32u ? 0u : ({x}) >> ({n}))', 'u32') + if op == 'shfl.sync.idx.b32': + d = a[0].split('|') + self.E('{ bool ib_;') + self.A(d[0], f'shfl_idx({self.val(a[1], "u32")}, {self.val(a[2], "u32")}, {self.val(a[3], "u32")}, ib_)', 'u32') + if len(d) > 1: + self.A(d[1], 'ib_', 'bool') + return self.E('}') + if op in ('shfl.sync.down.b32', 'shfl.sync.up.b32'): + return self.shfl_dir(op, a) + if parts[0] == 'sust': + return self.sust(op, a) + if parts[0] == 'tex': + return self.tex(op, a) + if parts[0] == 'cvt' and [t for t in parts[1:] if t in ('f16x2', 'e4m3x2')] == ['e4m3x2', 'f16x2']: + return self.cvt_to_e4m3(op, a) + if op == 'neg.f16': + return self.A(a[0], f'-{self.val(a[1], "f16")}', 'f16') + if parts[0] in ('min', 'max') and 'NaN' not in parts and any(t in parts for t in ('f16', 'f16x2', 'f32')): + t = 'f16x2' if 'f16x2' in parts else 'f16' if 'f16' in parts else 'f32' + fn = {'f16x2': 'p{}2', 'f16': 'p{}', 'f32': 'p{}f'}[t].format(parts[0]) + x, y = self.val(a[1], t), self.val(a[2], t) + return self.A(a[0], f'{fn}({x}, {y})', 'float' if t == 'f32' else t) + if op == 'neg.f16x2': + return self.A(a[0], f'-{self.val(a[1], "f16x2")}', 'f16x2') + if op == 'abs.f16x2': + return self.A(a[0], f'abs({self.val(a[1], "f16x2")})', 'f16x2') + if parts[0] == 'set' and parts[2] == 'f16': + c = ('(isnan({0}) || isnan({1}))' if parts[1] == 'nan' else + '!(isnan({0}) || isnan({1}))' if parts[1] == 'num' else self.cmp(set(parts), True)) + return self.A(a[0], f'(({c.format(self.val(a[1], "f16"), self.val(a[2], "f16"))}) ? 1.0hf : 0.0hf)', 'f16') + return super().op(op, a) + + # ------------------------------------------------ stage contract + def preamble(self, w, tex_filter): + w('#version 460') + w('#extension GL_EXT_shader_explicit_arithmetic_types : require') + w('#extension GL_EXT_shader_16bit_storage : require') + w('#extension GL_KHR_shader_subgroup_basic : require') + w('#extension GL_KHR_shader_subgroup_shuffle : require') + w('#extension GL_KHR_cooperative_matrix : require') + w('#extension GL_KHR_memory_scope_semantics : require') + w('#extension GL_EXT_float_e4m3 : require') + w(f'// Generated by engine/rr/ptx2rr.py from {self.kernel}. Model data derived from the user\'s DLL; do not commit.') + w(f'layout(local_size_x = {self.bx * self.by * self.bz}, local_size_y = 1, local_size_z = 1) in;') + w('layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint w[]; } param;') + w('layout(set = 0, binding = 1, std430) readonly buffer WeightBuf { uint16_t wh[]; } wbuf;') + w('layout(set = 0, binding = 2, std430) buffer ArenaBuf { uint16_t ah[]; } abuf;') + n = 3 + for off in sorted(self.textures): + w(f'layout(set = 0, binding = {n}) uniform sampler2D T{off};') + n += 1 + for off in sorted(self.surfaces): + w(f'layout(set = 0, binding = {n}, {self.surfaces[off]}) uniform writeonly image2D S{off};') + n += 1 + w(f'shared uint sh[{max(1, (self.shared_bytes + 3) // 4)}];') + if self.probe_slots or self.oob: + w(f'layout(set = 0, binding = {n}, std430) buffer Debug {{ uint d[]; }} dbg;') + if self.oob: + w('#define OOB_REPORT 1') + w(f'shared uint mms[{384 * (self.bx * self.by * self.bz // 32)}]; // mma_e4m3_f16 scratch, per subgroup') + + def runtime(self): + return RR_RUNTIME + + +def main(): + ap = argparse.ArgumentParser(description='PTX kernel -> GLSL compute shader (Ray Reconstruction stage contract)') + ap.add_argument('module') + ap.add_argument('kernel') + ap.add_argument('out') + ap.add_argument('--block', help='workgroup shape X,Y,Z (default: from the kernel table)') + ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT', + help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the " + "store's own guess, e.g. 144=rgba8_snorm (repeatable)") + ap.add_argument('--probe', action='append', default=[], metavar='LINE:%REG,...', + help='debugging: store the registers after the statement at that module line executes (last ' + 'execution) for workgroup --probe-block, into a debug buffer bound after the surfaces') + ap.add_argument('--probe-block', default='0,0') + ap.add_argument('--oob-report', action='store_true', help='debugging: count out-of-range buffer accesses') + ns = ap.parse_args() + block = tuple(int(x) for x in ns.block.split(',')) if ns.block else default_block(ns.kernel) + formats = {int(k): v for k, v in (x.split('=') for x in ns.surface)} + src = open(ns.module).read() + probes = {} + for spec in ns.probe: + line, regs = spec.split(':', 1) + probes[int(line)] = [r.strip() for r in regs.split(',')] + + def make(): + t = RRT(src, ns.kernel, block) + t.probes = probes + t.probe_block = tuple(int(x) for x in ns.probe_block.split(',')) + t.oob = ns.oob_report + return t + t = base.build(make, False, False, False) + for off, fmt in formats.items(): + if off in t.surfaces: + t.surfaces[off] = fmt + open(ns.out, 'w').write(t.emit({})) + # the binding order (stage contract): textures, then surfaces, in ascending argument offset + Path(ns.out).with_suffix('.json').write_text(json.dumps( + {'kernel': ns.kernel, 'block': block, 'textures': sorted(t.textures), + 'surfaces': [[off, t.surfaces[off]] for off in sorted(t.surfaces)], 'shared': t.shared_bytes, + 'probes': [[line, reg] for line in sorted(probes) for reg in probes[line]], 'oob': ns.oob_report}) + '\n') + sys.stderr.write(f'block={block} textures={sorted(t.textures)} surfaces={t.surfaces} shared={t.shared_bytes}\n') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/ptx_expr.py b/engine/rr/ptx_expr.py new file mode 100644 index 0000000..23d5493 --- /dev/null +++ b/engine/rr/ptx_expr.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Print the expression tree that defines a PTX register, followed back through arithmetic, pair packs/unpacks, +conversions and warp shuffles down to loads, constants and registers named as leaves. + +usage: ptx_expr.py STAGE LINE %REG [--depth N] [--stop %r1,%r2] + +LINE is a line at or after the register's definition (the search goes backwards from it). Brace-scoped temporaries +(`{ .reg .b16 hl, hu; ... }`) are resolved within their block. Shuffles print as `shfl.bfly^k(expr)`: the same +expression on the lane `laneid ^ k`. +""" +import argparse +import re +from pathlib import Path + + +def module_of(kernel): + for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')): + if f'.entry {kernel}(' in p.read_text(errors='replace'): + return p + + +class Tracer: + def __init__(self, lines, start): + self.lines = lines + self.start = start + self.consts = {} + self.cache = {} + + def find_def(self, reg, before, block_lo): + """the statement (line index) that last writes REG before line BEFORE""" + pat = re.compile(r'(^|[\s{,|])' + re.escape(reg) + r'([\s,|}]|$)') + lo = block_lo if not reg.startswith('%') else self.start + for i in range(before - 1, lo - 1, -1): + s = re.sub(r'//.*', '', self.lines[i]).strip().lstrip('{').strip() + s = re.sub(r'^@!?%p\d+\s+', '', s) + if not s or s.startswith(('.', '$', '}')): + continue + parts = s.split(None, 1) + if len(parts) < 2: + continue + rest = parts[1] + dsts = ([x.strip() for x in rest[1:rest.index('}')].split(',')] if rest.startswith('{') + else [x.strip() for x in rest.split(',')[0].split('|')]) + if reg in dsts and not parts[0].startswith(('st.', 'sust', 'bra')): + return i, parts[0], rest.rstrip(';').strip(), dsts + return None + + def block_start(self, line): + depth = 0 + for i in range(line, self.start, -1): + depth += self.lines[i].count('}') - self.lines[i].count('{') + if depth < 0: + return i + return self.start + + def expr(self, reg, before, depth, stop, half=None): + reg = reg.strip() + if re.fullmatch(r'-?\d+|0[fFdD][0-9A-Fa-f]+|0x[0-9a-fA-F]+', reg): + return reg + if reg in stop or depth <= 0: + return reg + (f'.{half}' if half is not None else '') + d = self.find_def(reg, before, self.block_start(before - 1)) + if d is None: + return reg + '?' + i, op, rest, dsts = d + args = [a.strip() for a in re.split(r',(?![^{]*\})', rest)] + srcs = args[1:] if not rest.startswith('{') else args[1:] + e = lambda r, h=None: self.expr(r, i, depth - 1, stop, h) + if op.startswith('mov.b32') and rest.startswith('{'): # unpack: {lo, hi} <- r + which = 'lo' if dsts.index(reg) == 0 else 'hi' + return e(srcs[0], which) + if op.startswith('mov.b32') and srcs and srcs[0].startswith('{'): # pack r <- {lo, hi} + lo, hi = [x.strip() for x in srcs[0][1:-1].split(',')] + if half == 'lo': + return e(lo) + if half == 'hi': + return e(hi) + return f'pack({e(lo)}, {e(hi)})' + if op.startswith(('mov.', 'cvt.')) and len(srcs) == 1: + inner = e(srcs[0], half) + return inner if op.startswith('mov.') else f'{op.split(".")[0]}.{op.split(".")[-2]}({inner})' + if op.startswith('shfl.sync.bfly'): + return f'shfl^{srcs[1]}({e(srcs[0], half)})' + if op.startswith('ld.'): + return f'load@L{i + 1}[{srcs[0]}]' + (f'.{half}' if half else '') + name = op.split('.')[0] + if op.endswith('x2') and half: + return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')' + return f'{name}(' + ', '.join(e(s, half) for s in srcs) + ')' + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('stage') + ap.add_argument('line', type=int) + ap.add_argument('reg') + ap.add_argument('--depth', type=int, default=30) + ap.add_argument('--stop', default='') + o = ap.parse_args() + kernel = f'cuda_dldn_engine_swin_{o.stage}_kernel' + lines = module_of(kernel).read_text().replace('\r', '').split('\n') + k = next(i for i, l in enumerate(lines) if f'.entry {kernel}(' in l) + t = Tracer(lines, k) + print(t.expr(o.reg, o.line, o.depth, set(filter(None, o.stop.split(','))))) + + +if __name__ == '__main__': + main() diff --git a/engine/rr/relations.py b/engine/rr/relations.py new file mode 100644 index 0000000..9dc82ca --- /dev/null +++ b/engine/rr/relations.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +"""Which earlier result each phase operand of a dumped layer (engine/rr/phase_dump.py) comes from. +usage: relations.py DUMP.npz [WAVE]""" +import sys +import numpy as np +sys.path.insert(0, __import__('os').path.dirname(__file__)) +from rrmath import f16, q8 + +z = np.load(sys.argv[1]); w = int(sys.argv[2]) if len(sys.argv) > 2 else 0 +phases = sorted({int(k.split('_')[0][1:]) for k in z.files}) +P = lambda p, n: z[f'p{p}_w{w}_{n}'] + + +def inter(n): + """the C/D -> A reinterpretation within each 32 block: column 4m+j <- 2m+j, 4m+2+j <- 16+2m+j""" + perm = [] + for b in range(n // 32): + for m in range(8): + perm += [32 * b + 2 * m, 32 * b + 2 * m + 1, 32 * b + 16 + 2 * m, 32 * b + 16 + 2 * m + 1] + return perm + + +def rel(op, M, p): + found = [] + if not M.any(): + return ['zero'] + for q in range(p - 1, 0, -1): + D = P(q, 'D') + if op == 'C': + for name, X in (('D', D),): + if X.shape == M.shape and np.array_equal(X, M): + found.append(f'{name}{q}') + continue + Q = q8(f16(D)) + cands = [] + if Q.shape[1] % 32 == 0: + cands += [('q8(D%d)[:, inter]' % q, Q[:, inter(Q.shape[1])])] + cands += [('q8(D%d)' % q, Q)] + for name, X in cands: + for tname, Y in ((name, X), (name + '.T', X.T)): + if Y.shape == M.shape and np.array_equal(Y, M): + found.append(tname) + elif Y.shape[1] == M.shape[1] and Y.shape[0] > M.shape[0] and np.array_equal(Y[:M.shape[0]], M): + found.append(tname + f'[:{M.shape[0]}]') + if op == 'B' and Y.shape[0] == M.shape[0] and Y.shape[1] > M.shape[1] and np.array_equal(Y[:, :M.shape[1]], M): + found.append(tname + f'[:, :{M.shape[1]}]') + if found: + break + return found or ['?'] + + +for p in phases: + A, B, C = P(p, 'A'), P(p, 'B'), P(p, 'C') + print(f'phase {p:2d} {A.shape[0]}x{B.shape[1]}x{A.shape[1]}: A <- {rel("A", A, p)} B <- {rel("B", B, p)} ' + f'C <- {rel("C", C, p)}') diff --git a/engine/rr/rr_attn.glsl b/engine/rr/rr_attn.glsl new file mode 100644 index 0000000..99797ec --- /dev/null +++ b/engine/rr/rr_attn.glsl @@ -0,0 +1,71 @@ +// The attention core shared by the 32-channel-head stages (engine/rr/enc1_m.comp, dec0_m.comp), included inside the +// head loop: this wave's 16 query slots against the window's 96 key slots. Expects KS (K, token slot x 32 FP8 +// channels in k order) and VS (V^T, 32 channels x 96 token slots) filled for head hd, the bias at W_BIAS0 + HEAD * hd +// (enc1_pack: 96 x 64, [key slot][query slot]) and the cubic table at W_CUBIC. Leaves Ob[2]: O^T (32 channels x +// this wave's queries) as the out-projection's B operand. RR_KS / RR_VS (+ word offsets RR_KS0 / RR_VS0) name +// the arrays when they are not KS and VS. +#ifndef RR_KS +#define RR_KS KS +#define RR_KS0 0u +#endif +#ifndef RR_VS +#define RR_VS VS +#define RR_VS0 0u +#endif + B8 Qb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(Qb[c], RR_KS, RR_KS0 + wv * 128u + 4u * c, 8u, CM); + // scores S^T (keys x this wave's queries) + bias, then the cubic softmax down each column + B8 Pb[6]; + { + float16_t e[6][8]; + [[unroll]] for (uint kt = 0u; kt < 6u; ++kt) { + HAC c; + coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM); + ACC s = ACC(c); + A8 k0, k1; + coopMatLoad(k0, RR_KS, RR_KS0 + kt * 128u, 8u, RM); + coopMatLoad(k1, RR_KS, RR_KS0 + kt * 128u + 4u, 8u, RM); + s = coopMatMulAdd(k1, Qb[1], coopMatMulAdd(k0, Qb[0], s)); // rounded once, by the conversion below + // the cubic numerator from enc1_pack.cubic_table, indexed by the score's binary16 bits + [[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(float16_t(s[i]))])); + } + // key slot 16kt + 8hl + i is NVIDIA's column 8j + 2tt + half with j = 4(kt >> 1) + (kt & 1) + 2((i >> 1) & 1), + // tt = 2hl + (i >> 2), half = i & 1 (enc1_ref.softmax_cubic's tree) + // packed: component h of a pair is column parity h (elements i = 4tq + h, i + 2), so the tree runs on f16x2 + f16vec2 ep[6][4]; + [[unroll]] for (uint kt = 0u; kt < 6u; ++kt) + [[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]); + precise f16vec2 st[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { // pairs q = 2tq (i0) and 2tq + 1 (i1 = i0 + 2) + const uint q0 = 2u * tq, q1 = q0 + 1u; + precise f16vec2 ev = ((ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1])) + (ep[4][q0] + ep[4][q1]); + precise f16vec2 od = ((ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1])) + (ep[5][q0] + ep[5][q1]); + st[tq] = ev + od; + } + precise f16vec2 loc = st[0] + st[1]; + const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u); + precise f16vec2 o = loc + unpackFloat2x16(other); + precise float16_t rs = o.x + o.y; + // f16(1 / rs) for a binary16 rs: x * m = 1 + d with |d| >= 2^-22 for every binary16 midpoint m, so the f32 + // reciprocal (within an ulp) already rounds correctly + const float16_t r = float16_t(1.0 / float(rs)); + const f16vec2 r2 = f16vec2(r, r); + [[unroll]] for (uint kt = 0u; kt < 6u; ++kt) + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + precise f16vec2 pv = ep[kt][q] * r2; + // a probability is in [0, 1]: the plain conversion, no saturation clamp, both in one instruction + { const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; } + } + } + // O^T = V^T P^T, then the out-projection onto the residual + B8 Ob[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC o = ACC(0.0); + [[unroll]] for (uint ks = 0u; ks < 3u; ++ks) { + A8 v0, v1; + coopMatLoad(v0, RR_VS, RR_VS0 + c * 384u + 8u * ks, 24u, RM); + coopMatLoad(v1, RR_VS, RR_VS0 + c * 384u + 8u * ks + 4u, 24u, RM); + K32(o, v0, Pb[2u * ks], v1, Pb[2u * ks + 1u]) + } + Q8(Ob[c], o) + } diff --git a/engine/rr/rr_capture.py b/engine/rr/rr_capture.py new file mode 100644 index 0000000..b0d827a --- /dev/null +++ b/engine/rr/rr_capture.py @@ -0,0 +1,156 @@ +"""Ray Reconstruction captures: the launches of one denoiser frame as the CUDA bridge recorded them. + +A capture directory holds what the bridge writes with + D4R_CUDA_VERBOSE=1 D4R_CUDA_REPLAY_DUMP_DIR=CAP/replay D4R_CUDA_REPLAY_DUMP_FILTER=cuda_dldn + D4R_CUDA_LAUNCH_STATS=1 D4R_CUDA_LAUNCH_STATS_FILTER=cuda_dldn D4R_CUDA_LAUNCH_DUMP_DIR=CAP/dump +plus the run's log as CAP/run/run.log (texture sampler settings and array shapes come from the verbose +cuTexObjectCreate / cuSurfObjectCreate / cuArrayCreate lines). + +replay/replay-NNNNNN-KERNEL/ args.bin, manifest.txt and every buffer the launch points into, *before* it ran +dump/launch-NNN-KERNEL-argOOO-{buffer,texture,surface}-... every resource *after* the launch ran + +Captures are reference data only: nothing at runtime reads them. +""" +from dataclasses import dataclass, field +from pathlib import Path +import re +import struct + +# CUarray_format codes the denoiser uses, and how a Vulkan image of the same texels is described +# (kvk's BITS column: 8 = UNORM bytes, 9 = SNORM bytes, 16 = binary16, 32 = binary32) +ARRAY_FORMATS = {16: 16, 32: 32, 1: 8, 200: 9, 194: 8} + + +@dataclass +class Image: + offset: int + kind: str # 'texture' or 'surface' + width: int + height: int + channels: int + format: int # CUarray_format + path: Path # post-launch contents + linear: bool = False # textures: bilinear filtering + + @property + def bits(self): + return ARRAY_FORMATS[self.format] + + @property + def texel_bytes(self): + return self.channels * (1 if self.bits in (8, 9) else self.bits // 8) + + +@dataclass +class Pointer: + offset: int + alloc: int + delta: int + + +@dataclass +class Launch: + sequence: int + kernel: str + directory: Path + grid: tuple + block: tuple + args: bytes + allocs: dict = field(default_factory=dict) # index -> (base, size) + pointers: list = field(default_factory=list) + images: list = field(default_factory=list) + + def pre(self, alloc): + """an allocation's contents before the launch""" + return (self.directory / f'alloc-{alloc}.bin').read_bytes() + + def pointer(self, offset): + return next(p for p in self.pointers if p.offset == offset) + + def image(self, offset): + return next(i for i in self.images if i.offset == offset) + + @property + def weights(self): + """the weight image (argument 40) as NGX uploaded it; 16 zero bytes for a stage without one (the output stage)""" + pointer = next((p for p in self.pointers if p.offset == 40), None) + return self.pre(pointer.alloc) if pointer is not None else bytes(16) + + @property + def arena_alloc(self): + """the plane arena: the allocation every arena pointer of the layer points into""" + candidates = [p.alloc for p in self.pointers if p.offset != 40] + return max(set(candidates), key=lambda a: self.allocs[a][1]) if candidates else None + + def post_buffer(self, capture, alloc): + base = self.allocs[alloc][0] + for path in (capture.root / 'dump').glob(f'launch-{self.sequence:03d}-{self.kernel}-arg*-buffer-*'): + m = re.search(r'-base0x([0-9a-f]+)\+', path.name) + if m and int(m.group(1), 16) == base: + return path.read_bytes() + raise FileNotFoundError(f'no post-launch dump of allocation {alloc} for launch {self.sequence}') + + def u32(self, offset): + return struct.unpack_from(' position: + break + if h == handle: + value = f + return value + + def _launch(self, directory): + manifest = (directory / 'manifest.txt').read_text().split('\n') + kernel = manifest[0].split()[1] + launch = [int(x) for x in manifest[1].split()[1:]] + sequence = int(re.search(r'replay-(\d+)-', directory.name).group(1)) + args = (directory / 'args.bin').read_bytes() + result = Launch(sequence, kernel, directory, tuple(launch[0:3]), tuple(launch[3:6]), args) + for line in manifest[2:]: + f = line.split() + if not f: + continue + if f[0] == 'alloc': + result.allocs[int(f[1])] = (int(f[2], 16), int(f[3])) + elif f[0] == 'pointer': + result.pointers.append(Pointer(int(f[1]), int(f[2]), int(f[3]))) + # the log position of this launch, so a reused texture handle resolves to its settings at the time + position = self.log.find(f'launch[{sequence}] {kernel}') + if position < 0: + position = len(self.log) + for path in sorted((self.root / 'dump').glob(f'launch-{sequence:03d}-{kernel}-arg*')): + m = re.search(r'-arg(\d+)-(texture|surface)-(\d+)x(\d+)x(\d+)-fmt(-?\d+)\.bin$', path.name) + if not m: + continue + offset = int(m.group(1)) + image = Image(offset, m.group(2), int(m.group(3)), int(m.group(4)), int(m.group(5)), int(m.group(6)), path) + if image.kind == 'texture': + handle = struct.unpack_from('> 1u, hf = tid & 1u; + const int x = x0 + int(t & 7u), y = y0 + int(t >> 3u); + if (x < 0 || y < 0 || x >= ex || y >= ey) return; + float16_t a[10]; + [[unroll]] for (uint c = 0u; c < 10u; ++c) a[c] = OUTF[t * 16u + c]; + // surface 112: softmax of channels 0-8 over the 3x3 taps, one f16 fma chain per channel; the gate (channel 9) + // blends it with texture 96; + texture 104 + const float16_t ma = max(max(max(a[0], a[2]), a[4]), a[6]), mb = max(max(max(a[1], a[3]), a[5]), a[7]); + const float16_t mx = max(max(ma, mb), a[8]); + float e[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = a[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t rr = rcp16(s); + float16_t wg[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = h16(r16(e[i]) * float(rr)); + const uint tx = uint(x - x0), ty = uint(y - y0), c0 = 2u * hf, nc = 2u - hf; + float16_t f[2] = float16_t[2](0.0hf, 0.0hf); + bool mid = false; + [[unroll]] for (uint k = 0u; k < 9u; ++k) { + const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0; + [[unroll]] for (uint c = 0u; c < 2u; ++c) f[c] = fma16q(TAP[at + c], wg[k], f[c], mid); + } + if (mid) { // an f32 result on a binary16 midpoint: redo the chain with the exact fma (rare) + f[0] = 0.0hf; f[1] = 0.0hf; + for (uint k = 0u; k < 9u; ++k) { + const uint at = ((ty + k / 3u) * 10u + tx + k % 3u) * 4u + c0; + for (uint c = 0u; c < 2u; ++c) f[c] = fma16(TAP[at + c], wg[k], f[c]); + } + } + const float16_t g = sig(a[9]); + precise float16_t om = float16_t(1.0) - g; + const vec4 t96 = texelFetch(T96, ivec2(x, y), 0), t104 = texelFetch(T104, ivec2(x, y), 0); + float16_t o[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + precise float16_t fg = f[c] * g; + float16_t v = fma16(float16_t(t96[c0 + c]), om, fg); + const float16_t t4 = float16_t(t104[c0 + c]); + if (!isnan(t4)) { precise float16_t vv = v + t4; v = vv; } + o[c] = v; + } + const uint other = subgroupShuffleXor(bh(o[0]), 1u); + if (hf == 0u) imageStore(S112, ivec2(x, y), vec4(o[0], o[1], hb(other), 0.0)); + } diff --git a/engine/rr/rr_exact.glsl b/engine/rr/rr_exact.glsl new file mode 100644 index 0000000..5aa1640 --- /dev/null +++ b/engine/rr/rr_exact.glsl @@ -0,0 +1,147 @@ +// Exact NVIDIA arithmetic helpers for the native Ray Reconstruction stages (engine/rr/*_m.comp, enc1_native.comp). +// Requires GL_EXT_shader_explicit_arithmetic_types, GL_EXT_float_e4m3. +// +// RADV folds a register-only float(float16_t(x)) back to x, so every binary16 rounding whose result is used again goes +// through a conversion with an opaque zero (RR_ZERO16, set in main from something the compiler cannot evaluate, e.g. +// gl_NumWorkGroups.z - 1) OR-ed into its bits: one instruction more than the plain conversion. +uint16_t RR_ZERO16; +#if defined(RR_RELAXED) +// Relaxed rounding (build_native_model.sh with D4R_RR_RELAXED=1): the accumulators stay f32 between 32-deep steps and +// are converted to FP8 in one rounding instead of NVIDIA's two (binary16, then FP8). About 1% of the FP8 codes of a +// layer then differ by one step, so a stage no longer reproduces the reference bit for bit; it saves the conversions +// that are a quarter of a stage's instructions. +float rnd16(float x) { return x; } +#elif defined(RR_EXACT_CONVERT) +// the module's conversions are decorated NoContraction (engine/rr/spv_exact_convert.py): the driver keeps them +float rnd16(float x) { return float(float16_t(x)); } +#else +float rnd16(float x) { return float(uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16)); } +#endif +// x rounded to binary16 (nearest even), as f32 / as f16, through a conversion the compiler keeps (see rnd16) +float r16(float x) { return rnd16(x); } +#if (defined(RR_RELAXED) && RR_RELAXED >= 2) || defined(RR_EXACT_CONVERT) +float16_t h16(float x) { return float16_t(x); } +#else +float16_t h16(float x) { return uint16BitsToFloat16(float16BitsToUint16(float16_t(x)) | RR_ZERO16); } +#endif +// NVIDIA's k order within a 32-block: slot j holds natural channel inter(j); uninter is the inverse (and also the +// arena's channel order: arena slot uninter(k) holds natural channel k) +uint inter(uint j) { return (j & ~31u) | (((j & 31u) >> 2u) << 1u) | (j & 1u) | (((j >> 1u) & 1u) << 4u); } +uint uninter(uint c) { return (c & ~31u) | (((c & 15u) >> 1u) << 2u) | (c & 1u) | (((c >> 4u) & 1u) << 1u); } +int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); } + +// cvt.rn.satfinite.e4m3 (NaN -> 0x7f) +floate4m3_t q8(float16_t h) +{ + return isnan(h) ? uintBitsToFloate4m3EXT(uint8_t(0x7fu)) : floate4m3_t(clamp(h, float16_t(-448.0), float16_t(448.0))); +} +// the same from f32 / f16 values in one saturating hardware conversion (satfinite; for an f32 value that is a binary16 +// number the result equals the f16 conversion's) +floate4m3_t q8s(float x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; } +floate4m3_t q8h(float16_t x) { floate4m3_t r; saturatedConvertEXT(r, x); return r; } + +uint code(floate4m3_t v) { return uint(floate4m3BitsToUintEXT(v)); } +float16_t hb(uint b) { return uint16BitsToFloat16(uint16_t(b)); } +uint bh(float16_t h) { return uint(float16BitsToUint16(h)); } + +// m * x compared with 1, exactly (the fma recovers the product's rounding error): -1, 0 or 1 +int cmp1(float m, float x) +{ + precise float t = m * x; + precise float e = fma(m, x, -t); + return t > 1.0 ? 1 : (t < 1.0 ? -1 : (e > 0.0 ? 1 : (e < 0.0 ? -1 : 0))); +} +// 1 / x of the filters' softmax and gate, rounded to binary16. The reference is the GPU's own reciprocal instruction +// (rcp.approx; ZLUDA lowers it to llvm.amdgcn.rcp.f32, which is what 1.0 / x compiles to here), not the correctly +// rounded quotient: correcting it left ~0.01% of the filtered surfaces one ulp off (capDT, 2026-10-12). +float16_t rcp16(float x) { return h16(1.0 / x); } +// binary16 rounding of 1 / sqrt(s) (s > 0) +float16_t rsq16(float16_t s16) +{ + const float s = float(s16); + float16_t q = h16(inversesqrt(s)); +#if defined(RR_RELAXED) && RR_RELAXED >= 2 + return q; +#endif + for (int i = 0; i < 2; ++i) { + const uint b = bh(q); + const float lo = (float(q) + float(hb(b - 1u))) * 0.5, hi = (float(q) + float(hb(b + 1u))) * 0.5; + const int cl = cmp1(lo * lo, s), ch = cmp1(hi * hi, s); + if (cl > 0 || (cl == 0 && (b & 1u) != 0u)) q = hb(b - 1u); + else if (ch < 0 || (ch == 0 && (b & 1u) != 0u)) q = hb(b + 1u); + } + return q; +} +// fma of binary16 operands with one binary16 rounding, exactly. The product is exact in f32, so an f32 fma rounds the +// exact result once; converting that to binary16 rounds twice, which differs only when the f32 value lands exactly +// on a binary16 midpoint: then the sum's rounding error (TwoSum, the rare branch) decides the direction. +float16_t fma16(float16_t a, float16_t b, float16_t c) +{ + const float fa = float(a), fb = float(b), fc = float(c); + precise float s = fma(fa, fb, fc); +#if defined(RR_RELAXED) && RR_RELAXED >= 2 + return float16_t(s); +#endif + const uint u = floatBitsToUint(s); + if ((u & 0x1fffu) == 0x1000u && abs(s) >= 6.103515625e-05) { + precise float p = fa * fb; + precise float bb = s - p; + precise float err = (p - (s - bb)) + (fc - bb); + if (err != 0.0) return h16(uintBitsToFloat(err * s > 0.0 ? u + 1u : u - 1u)); + } + return h16(s); +} +// the same in one f32 fma, rounded once more to binary16: mid is set when the f32 result lies on a binary16 midpoint, +// where that double rounding can differ (the caller then redoes the computation with fma16) +float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid) +{ + precise float s = fma(float(a), float(b), float(c)); +#if defined(RR_RELAXED) && RR_RELAXED >= 2 + return float16_t(s); +#endif + mid = mid || (floatBitsToUint(s) & 0x1fffu) == 0x1000u; +#ifdef RR_EXACT_CONVERT + return float16_t(s); +#else + return hb(bh(float16_t(s)) | uint(RR_ZERO16)); +#endif +} +// The saturating FP8 conversion of a whole accumulator, two elements per conversion instruction +#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, vec2((src)[i_], (src)[i_ + 1])); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; } +// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT) +#define ACTQ(h) uintBitsToFloate4m3EXT(w8u.w[W_ACT + bh(h)]) +float16_t act(float16_t x) +{ + precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0)); + precise float16_t a = float16_t(0.0810546875) * abs(y); + precise float16_t b = float16_t(0.412109375) - a; + precise float16_t c = y * b; + precise float16_t g = float16_t(0.5) + c; + precise float16_t r = x * g; + return r; +} +float16_t sig(float16_t v) { return rcp16(exp2(float(v) * -1.4426950216293335) + 1.0); } +// the cubic-exponent softmax numerator (enc1_ref.softmax_cubic): a cubic in the clamped scaled score whose binary16 +// bits, shifted into the exponent field, are the exponential +// the same in plain f32 fmas (double rounding); bad is set when an f32 result lies exactly on a binary16 midpoint, +// the only case where that differs: the caller then recomputes with cubic_e (subgroupAny, rare) +float16_t cubic_e_fast(float16_t x, inout bool bad) +{ + precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375)); + const float fm = float(m); + precise float t = fma(fm, -fm, 0.92724609375); + const float tr = rnd16(t); + precise float u = fma(fm, tr, 1.375); + // t - c and u - 1.375 are exact (Sterbenz), so one more fma gives the rounding error's sign: only an inexact + // midpoint is a double rounding + precise float et = fma(fm, fm, t - 0.92724609375), eu = fma(fm, -tr, u - 1.375); + bad = bad || ((floatBitsToUint(t) & 0x1fffu) == 0x1000u && et != 0.0) || ((floatBitsToUint(u) & 0x1fffu) == 0x1000u && eu != 0.0); + return hb((bh(float16_t(u)) << 5u) & 0x7FE0u); +} +float16_t cubic_e(float16_t x) +{ + precise float16_t m = clamp(x * float16_t(0.017181396484375), float16_t(-0.55615234375), float16_t(0.55615234375)); + const float16_t t = fma16(m, -m, float16_t(0.92724609375)); + const float16_t u = fma16(m, t, float16_t(1.375)); + return hb((bh(u) << 5u) & 0x7FE0u); +} diff --git a/engine/rr/rr_inputs.py b/engine/rr/rr_inputs.py new file mode 100644 index 0000000..d06b126 --- /dev/null +++ b/engine/rr/rr_inputs.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Write the game inputs of a capture's denoiser frames for engine/replay_rr.cpp (d4r-rr). + +usage: rr_inputs.py CAPTURE OUT_DIR [--frames N] [--no-reset] + +The inputs are enc0's textures (colour 160, motion 168, normals 184, roughness 192, the albedos 232 and 240, depth +512); the frame settings come from enc0's argument block: jitter (320), motion scale (344, in output pixels per +render pixel), depth inversion (528) and the camera, whose rotation-times-projection (596) and projection (660) give +back a world-to-view rotation and the view-to-clip matrix that reproduce the block. The harness passes reset on every +frame (scripts/run_rr_engine_compare.sh); --no-reset clears it after the first frame. +""" +import argparse +import os +import struct +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rr_capture import Capture # noqa: E402 + +OPTIONAL = {'hitdistance': 248, 'disocclusion': 296} +INPUTS = {'color': 160, 'motion': 168, 'normals': 184, 'roughness': 192, 'diffuse': 232, 'specular': 240, 'depth': 512} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('out', type=Path) + ap.add_argument('--frames', type=int, default=3) + ap.add_argument('--no-reset', action='store_true') + o = ap.parse_args() + cap = Capture(o.capture) + o.out.mkdir(parents=True, exist_ok=True) + lines = [] + for f in range(1, o.frames + 1): + enc0 = cap.find('swin_enc0', f) + a = enc0.args + u32 = lambda off: struct.unpack_from('> 7, (_c >> 3) & 15, _c & 7 + _v = np.nan if (_e == 15 and _m == 7) else (_m * 2.0 ** -9 if _e == 0 else (1 + _m / 8) * 2.0 ** (_e - 7)) + _E4M3[_c] = -_v if _s else _v + + +def e4m3(codes): + return _E4M3[np.asarray(codes, np.uint8)] + + +def f16(bits): + return np.asarray(bits, np.uint16).view(np.float16).astype(np.float64) + + +def f16bits(x): + return np.asarray(x, np.float64).astype(np.float16).view(np.uint16) + + +def mma_k32(a_codes, b_codes, c_f16bits, ksteps=None): + """D = A.B + C with ZLUDA's lowering: per 32-deep k-step, f32 sum of the exact products and C, one f16 rounding; + the next k-step takes that f16 value as C. (f64 stands in for f32: the products and short sums are exact.)""" + a, b = e4m3(a_codes), e4m3(b_codes) + acc = f16(c_f16bits) + for k in range(0, a.shape[1], 32): + acc = (a[:, k:k + 32] @ b[k:k + 32, :] + acc).astype(np.float32).astype(np.float16).astype(np.float64) + return f16bits(acc) + +# cvt.rn.satfinite.e4m3 of every f16 bit pattern (the hardware conversion, exhaustively equal to PTX's semantics; +# NaN -> 0x7f), see engine/rr/ptx2rr.py h2e4m3_hw +import os as _os +_Q8 = np.load(_os.path.join(_os.path.dirname(__file__), 'f16_to_e4m3.npy')) + + +def q8(x_f16): + """f16 values (float array) -> e4m3 codes""" + return _Q8[np.asarray(x_f16, np.float64).astype(np.float16).view(np.uint16)] diff --git a/engine/rr/spv_exact_convert.py b/engine/rr/spv_exact_convert.py new file mode 100644 index 0000000..a0795d5 --- /dev/null +++ b/engine/rr/spv_exact_convert.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +"""Decorate every float conversion of a SPIR-V module NoContraction, so the driver keeps it. + +usage: spv_exact_convert.py IN.spv [OUT.spv] + +RADV folds float(float16_t(x)) back to x unless the conversions are marked exact; glslang's `precise` does not reach +conversions. With the decoration a binary16 rounding is two instructions (engine/rr/rr_exact.glsl, RR_EXACT_CONVERT) +instead of three with an opaque zero OR-ed into the bits. +""" +import struct +import sys + +OP_DECORATE, OP_FCONVERT, NO_CONTRACTION = 71, 115, 42 + + +def main(): + src = sys.argv[1] + dst = sys.argv[2] if len(sys.argv) > 2 else src + data = open(src, 'rb').read() + words = list(struct.unpack(f'<{len(data) // 4}I', data)) + if words[0] != 0x07230203: + sys.exit(f'{src}: not SPIR-V') + insts, i = [], 5 + while i < len(words): + count = words[i] >> 16 + insts.append((words[i] & 0xffff, words[i:i + count])) + i += count + ids = [w[2] for op, w in insts if op == OP_FCONVERT] + first_type = next(k for k, (op, _) in enumerate(insts) if 19 <= op <= 39) # annotations end where types begin + out = list(words[:5]) + for k, (op, w) in enumerate(insts): + if k == first_type: + for result in ids: + out += [3 << 16 | OP_DECORATE, result, NO_CONTRACTION] + out += w + open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out)) + print(f'{dst}: {len(ids)} conversions marked exact') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/spv_float_controls.py b/engine/rr/spv_float_controls.py new file mode 100644 index 0000000..51d20dc --- /dev/null +++ b/engine/rr/spv_float_controls.py @@ -0,0 +1,71 @@ +#!/usr/bin/env python3 +"""Give a compute SPIR-V module PTX's floating-point environment: RoundingModeRTE for 16- and 32-bit results and +DenormPreserve for 16-bit values (SPV_KHR_float_controls, core in SPIR-V 1.4). + +usage: spv_float_controls.py IN.spv [OUT.spv] + +Without these execution modes the rounding of float conversions is the driver's choice, where every cvt.rn of +NVIDIA's kernels rounds to nearest even. glslang keeps only one execution mode per mode id, so the modes are inserted +into the compiled module instead. +""" +import struct +import sys + +OP_EXTENSION, OP_CAPABILITY, OP_ENTRY_POINT, OP_EXECUTION_MODE = 10, 17, 15, 16 +CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE = 4464, 4467 +MODES = ((4462, 16), (4462, 32), (4459, 16)) # RoundingModeRTE 16/32, DenormPreserve 16 +EXTENSION = 'SPV_KHR_float_controls' + + +def instructions(words): + i = 5 + while i < len(words): + count = words[i] >> 16 + yield words[i] & 0xffff, words[i:i + count] + i += count + + +def string_words(text): + data = text.encode() + b'\0' + data += b'\0' * (-len(data) % 4) + return list(struct.unpack(f'<{len(data) // 4}I', data)) + + +def patch(words): + insts = list(instructions(words)) + caps = {w[1] for op, w in insts if op == OP_CAPABILITY} + has_ext = any(op == OP_EXTENSION and EXTENSION.encode() in struct.pack(f'<{len(w) - 1}I', *w[1:]) for op, w in insts) + entry = next(w[2] for op, w in insts if op == OP_ENTRY_POINT) + modes = {(w[2], w[3]) for op, w in insts if op == OP_EXECUTION_MODE and w[1] == entry and len(w) > 3} + last_cap = max(i for i, (op, _) in enumerate(insts) if op == OP_CAPABILITY) + first_mode = min(i for i, (op, _) in enumerate(insts) if op == OP_EXECUTION_MODE) + out = list(words[:5]) + for i, (op, w) in enumerate(insts): + if i == first_mode: + for mode, width in MODES: + if (mode, width) not in modes: + out += [4 << 16 | OP_EXECUTION_MODE, entry, mode, width] + out += w + if i == last_cap: + for cap in (CAP_DENORM_PRESERVE, CAP_ROUNDING_RTE): + if cap not in caps: + out += [2 << 16 | OP_CAPABILITY, cap] + if not has_ext: + name = string_words(EXTENSION) + out += [(1 + len(name)) << 16 | OP_EXTENSION, *name] + return out + + +def main(): + src = sys.argv[1] + dst = sys.argv[2] if len(sys.argv) > 2 else src + data = open(src, 'rb').read() + words = list(struct.unpack(f'<{len(data) // 4}I', data)) + if words[0] != 0x07230203: + sys.exit(f'{src}: not SPIR-V') + out = patch(words) + open(dst, 'wb').write(struct.pack(f'<{len(out)}I', *out)) + + +if __name__ == '__main__': + main() diff --git a/engine/rr/swin_m.comp b/engine/rr/swin_m.comp new file mode 100644 index 0000000..11a9ec5 --- /dev/null +++ b/engine/rr/swin_m.comp @@ -0,0 +1,471 @@ +#version 460 +// Native Swin stage of the DLSS-D (preset E) denoiser, register-resident, for any stage of engine/rr/swin_ref.STAGES +// (enc2-enc5, dec2-dec4; enc1_m.comp / dec1_m.comp are the hand-written C = 64 instances it generalises). The stage's +// shape and weight offsets come from the include generated by engine/rr/swin_pack.py: +// KIND 0 encoder (tokens from the arena, 2x2 merge to CX channels, enc1's gated filter; CX = 0: the bottleneck, +// ungated filter, no merge), 1 decoder (expand of the CX channels one level down + skip, dec0's filter) +// C channels (KC = C / 32 planes, MT = C / 16 tiles), NH heads of 32 channels, C / 8 MLP branches +// The arithmetic is engine/rr/swin_ref.py's; enc1_m.comp's header explains the transposed products and slot orders. +// One workgroup per 8x8-token window, four waves, wave w owning query slots 16w..16w+15. +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_EXT_shader_explicit_arithmetic_types : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_float_e4m3 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require +#extension GL_GOOGLE_include_directive : require +layout(local_size_x = 128) in; +#include "swin_stage.h" +#include "rr_exact.glsl" + +#define KC (C / 32) +#define MT (C / 16) +#define NB (C / 8) + +layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param; +layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8; +layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16; +layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u; +layout(set = 0, binding = 1, std430) readonly buffer W8U { uint8_t w[]; } w8u; +layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena; +layout(set = 0, binding = 3) uniform sampler2D T88; +#if KIND == 1 +layout(set = 0, binding = 4) uniform sampler2D T96; +layout(set = 0, binding = 5) uniform sampler2D T104; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S112; +#else +layout(set = 0, binding = 4, rgba16f) uniform writeonly image2D S112; +#if CX +layout(set = 0, binding = 5, rgba16f) uniform writeonly image2D S120; +layout(set = 0, binding = 6, rgba16f) uniform writeonly image2D S128; +#endif +#endif + +#define ACC coopmat +#define HAC coopmat +#define A8 coopmat +#define B8 coopmat +#define RM gl_CooperativeMatrixLayoutRowMajor +#define CM gl_CooperativeMatrixLayoutColumnMajor +#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM) +#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM) +#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]); +#define Q8(dst, src) Q8PAIRS(dst, src) +#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); RND(acc) } +const uint HEAD = W_WK1 - W_WK0, BRANCH = W_W1_1 - W_W1_0; + +// NVIDIA's reduction over a short list (swin_ref.pairs_tree) +#if KC == 1 +#define TREE(v) (v[0]) +#elif KC == 2 +#define TREE(v) (v[0] + v[1]) +#elif KC == 3 +#define TREE(v) ((v[0] + v[1]) + v[2]) +#elif KC == 4 +#define TREE(v) ((v[0] + v[1]) + (v[2] + v[3])) +#else +#define TREE(v) (((v[0] + v[1]) + (v[2] + v[3])) + v[4]) +#endif + +#if KIND == 1 +// one pool, by phase: [EX | LOW] for the expand; the prologue reads EX and writes XT over LOW (dead) and beyond; +// the attention's K and V^T reuse EX (dead after the residual) +#define EX_W (48u * C) +#define POOL_W (EX_W + (12u * CX > 24u * C ? 12u * CX : 24u * C)) +shared uint POOL[POOL_W]; +#define EX POOL // the expansion: 48 rows x 4 sub-positions x C FP8 channels in slot order +#define LOW_AT EX_W // the low-resolution tokens (6 x 6, padded): CX FP8 channels, arena order +#define XT_AT EX_W // the normalised key tokens: slot-major, C FP8 channels in slot order +#define XTS POOL +#define RR_KS POOL +#define RR_KS0 0u +#define RR_VS POOL +#define RR_VS0 768u +shared float16_t TAP[100 * 4]; // texture 88 bilinear at the window's 10 x 10 filter tap positions +shared float16_t OUTF[64 * 16]; // the head per window token (raster order) +#else +shared uint XT[96 * C / 4]; // the normalised key tokens: slot-major, C FP8 channels in slot order +#define XT_AT 0u +#define XTS XT +shared uint KVS[2 * 768]; // the current head's K, then V^T +#define RR_KS KVS +#define RR_KS0 0u +#define RR_VS KVS +#define RR_VS0 768u +#define MG XT // the merge rows (16 x 4C FP8) reuse XT, dead after the projections +#if CX +shared float16_t O112[64 * 4]; // surface 112 per token, for surface 128's 2x2 sums +#endif +#endif + +ivec2 keyXY(uint t, int x0, int y0) +{ + if (t < 64u) return ivec2(x0 + int(t & 7u), y0 + int(t >> 3u)); + const int i = int(t & 7u), side = int((t - 64u) >> 3u); + return ivec2(side < 2 ? x0 + i : (side == 2 ? x0 - 1 : x0 + 8), side == 0 ? y0 - 1 : (side == 1 ? y0 + 8 : y0 + i)); +} +int clampTap(int v, int n) { return v < 0 ? 1 : (v >= n ? n - 2 : v); } +uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d) +{ + return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u); +} +float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); } +uvec2 mergeXY(uint r) { return r < 8u ? uvec2(r >> 1u, r & 1u) : uvec2((r - 8u) >> 1u, 2u + (r & 1u)); } + +void main() +{ + const uint tid = gl_LocalInvocationIndex, wv = gl_SubgroupID, lane = gl_SubgroupInvocationID; + const uint hl = lane >> 4u, col = lane & 15u; + RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u); + const int ex = int(param.p[0]), ey = int(param.p[1]), hx = ex / 2, hy = ey / 2; + const uint inOff = param.p[2], skOff = param.p[6], outOff = param.p[12], mrgOff = param.p[14], shift = param.p[20]; + const int x0 = 8 * int(gl_WorkGroupID.x) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u); + const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u; + +#if KIND == 1 + // ---------------------------------------------------------------- expand: wave w computes sub-position w + const int X0 = (x0 >> 1) - 1, Y0 = (y0 >> 1) - 1; + for (uint i = tid; i < 48u * (CX / 4u); i += 128u) { + const uint r = i / (CX / 4u), wd = i % (CX / 4u); + uint u = 0u; + if (r < 36u) { + const int X = mirror(X0 + int(r % 6u), hx), Y = mirror(Y0 + int(r / 6u), hy); + u = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)]; + } + POOL[LOW_AT + i] = u; + } + barrier(); + [[unroll]] for (uint n = 0u; n < 3u; ++n) { + B8 lb[CX / 16]; + [[unroll]] for (uint kt = 0u; kt < CX / 16; ++kt) coopMatLoad(lb[kt], POOL, LOW_AT + n * 4u * CX + 4u * kt, CX / 4u, CM); + for (uint mm = 0u; mm < MT; ++mm) { + const uint m = MT * wv + mm; + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < CX / 16; kt += 2u) { + A8 w0, w1; + LDA(w0, W_WE, CX, m, kt); LDA(w1, W_WE, CX, m, kt + 1u); + K32(a, w0, lb[kt], w1, lb[kt + 1u]) + } + B8 q; + Q8(q, a) + coopMatStore(q, EX, n * 16u * C + 4u * m, C, CM); + } + } + barrier(); +#endif + + // ---------------------------------------------------------------- prologue: one key token per thread + if (tid < 96u) { + const ivec2 k = keyXY(inter(tid), x0, y0); + const int xx = mirror(k.x, ex), yy = mirror(k.y, ey); + const uint tb = ((KIND == 1 ? skOff : inOff) + uint(yy * ex + xx) * 32u) / 4u; +#if KIND == 1 + const uint r = uint(((k.y >> 1) - Y0) * 6 + ((k.x >> 1) - X0)), s = uint(2 * (k.y & 1) + (k.x & 1)); +#endif + float16_t v[C]; + [[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) { + const uint t = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)]; +#if KIND == 1 + const uint e = EX[r * C + s * (C / 4u) + wd]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { precise float16_t x = e4(e, b) + e4(t, b); v[4u * wd + b] = x; } +#else + [[unroll]] for (uint b = 0u; b < 4u; ++b) v[4u * wd + b] = e4(t, b); +#endif + } + // swin_ref.input_norm: per lane word q (and q + 4) the pair sums of the planes (TREE), the halves, then + // q ^ 1, q ^ 2, the components + precise f16vec2 h[4]; + [[unroll]] for (uint q = 0u; q < 4u; ++q) { + precise f16vec2 hh[2]; + [[unroll]] for (uint half_ = 0u; half_ < 2u; ++half_) { + precise f16vec2 pr[KC]; + [[unroll]] for (uint p = 0u; p < KC; ++p) { + const uint o = 32u * p + 4u * (q + 4u * half_); + precise f16vec2 a = f16vec2(v[o], v[o + 1u]) * f16vec2(v[o], v[o + 1u]); + precise f16vec2 c = f16vec2(v[o + 2u], v[o + 3u]) * f16vec2(v[o + 2u], v[o + 3u]); + pr[p] = a + c; + } + precise f16vec2 tr = TREE(pr); + hh[half_] = tr; + } + h[q] = hh[0] + hh[1]; + } + precise f16vec2 h0 = h[0] + h[1], h2_ = h[2] + h[3]; + precise f16vec2 hs = h0 + h2_; + precise float16_t sm = hs.x + hs.y; + precise float16_t se = sm + float16_t(0.0001220703125); + const float16_t rr = rsq16(se); + [[unroll]] for (uint wd = 0u; wd < C / 4; ++wd) { + floate4m3_t c[4]; + [[unroll]] for (uint b = 0u; b < 4u; ++b) { + const uint j = 4u * wd + b; + precise float16_t g = rr * w16.w[W_GAIN0 / 2u + inter(j)]; + precise float16_t n = v[j] * g; + c[b] = q8h(n); + } + XTS[XT_AT + tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]); + } + } + barrier(); + + // ---------------------------------------------------------------- the residual: the window token + table 7 + ACC R[MT]; + { + const uint tq = inter(16u * wv + col); + const int qx = x0 + int(tq & 7u), qy = y0 + int(tq >> 3u); + const uint tb = ((KIND == 1 ? skOff : inOff) + uint(mirror(qy, ey) * ex + mirror(qx, ex)) * 32u) / 4u; +#if KIND == 1 + const uint r = uint(((qy >> 1) - Y0) * 6 + ((qx >> 1) - X0)), s = uint(2 * (qy & 1) + (qx & 1)); +#endif + [[unroll]] for (uint m = 0u; m < MT; ++m) { + const uint ta = tb + (m >> 1u) * (plane >> 2u) + 4u * (m & 1u) + 2u * hl; + const uvec2 t = uvec2(arena.w[ta], arena.w[ta + 1u]); +#if KIND == 1 + const uint ea = r * C + s * (C / 4u) + 4u * m + 2u * hl; + const uvec2 e = uvec2(EX[ea], EX[ea + 1u]); +#endif + [[unroll]] for (int i = 0; i < 8; ++i) { +#if KIND == 1 + precise float16_t x = e4(e[i >> 2], uint(i) & 3u) + e4(t[i >> 2], uint(i) & 3u); +#else + const float16_t x = e4(t[i >> 2], uint(i) & 3u); +#endif + precise float16_t c = x + w16.w[W_TABLE7 / 2u + 16u * m + 8u * hl + uint(i)]; + R[m][i] = float(c); + } + } + } + + // ---------------------------------------------------------------- attention, NH heads + [[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) { + barrier(); // every wave is done with the previous head's K / V (and, for the first, the expand's LOW) + for (uint j = 0u; j < 3u; ++j) { // 12 jobs (token tile n, channel tile c) + const uint job = wv + 4u * j, n = job >> 1u, c = job & 1u; + ACC a = ACC(0.0), v = ACC(0.0); + [[unroll]] for (uint kt = 0u; kt < 2u * KC; kt += 2u) { + B8 xb0, xb1, wv0, wv1; + A8 xa0, xa1, wk0, wk1; + coopMatLoad(xb0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, CM); + coopMatLoad(xb1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, CM); + coopMatLoad(xa0, XTS, XT_AT + n * 4u * C + 4u * kt, C / 4u, RM); + coopMatLoad(xa1, XTS, XT_AT + n * 4u * C + 4u * kt + 4u, C / 4u, RM); + LDA(wk0, W_WK0 + HEAD * hd, C, c, kt); LDA(wk1, W_WK0 + HEAD * hd, C, c, kt + 1u); + LDB(wv0, W_WV0 + HEAD * hd, 32u, kt, c); LDB(wv1, W_WV0 + HEAD * hd, 32u, kt + 1u, c); + K32(a, wk0, xb0, wk1, xb1) + K32(v, xa0, wv0, xa1, wv1) + } + A8 ka, va; + Q8(ka, a) + Q8(va, v) + coopMatStore(ka, RR_KS, RR_KS0 + n * 128u + 4u * c, 8u, RM); + coopMatStore(va, RR_VS, RR_VS0 + c * 384u + 4u * n, 24u, RM); + } + barrier(); +#include "rr_attn.glsl" + [[unroll]] for (uint m = 0u; m < MT; ++m) { + A8 a0, a1; + LDA(a0, W_WO0 + HEAD * hd, 32u, m, 0u); LDA(a1, W_WO0 + HEAD * hd, 32u, m, 1u); + K32(R[m], a0, Ob[0], a1, Ob[1]) + } + } + + // ---------------------------------------------------------------- MLP + B8 Nb[MT]; + { + float16_t rv[MT][8]; + [[unroll]] for (uint m = 0u; m < MT; ++m) [[unroll]] for (int i = 0; i < 8; ++i) rv[m][i] = float16_t(R[m][i]); + // swin_ref.rowsum_sq: slot 16m + 8hl + i is NVIDIA's channel 8j + 2t + c with j = 4(m >> 1) + (m & 1) + + // 2((i >> 1) & 1), t = 2hl + (i >> 2), c = i & 1: per parity (m & 1) the tile pairs (i, i + 2) of every + // 32-channel block, TREE over the blocks; the parities; t ^ 1 (local), t ^ 2 (the other lane half) + float16_t h1[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + float16_t ht[2]; + [[unroll]] for (uint tq = 0u; tq < 2u; ++tq) { + const uint i0 = 4u * tq + c, i1 = i0 + 2u; + precise float16_t par[2]; + [[unroll]] for (uint pa = 0u; pa < 2u; ++pa) { + precise float16_t pr[KC]; + [[unroll]] for (uint b = 0u; b < KC; ++b) + pr[b] = rv[2u * b + pa][i0] * rv[2u * b + pa][i0] + rv[2u * b + pa][i1] * rv[2u * b + pa][i1]; + precise float16_t tr = TREE(pr); + par[pa] = tr; + } + precise float16_t hv = par[0] + par[1]; + ht[tq] = hv; + } + precise float16_t hh = ht[0] + ht[1]; + h1[c] = hh; + } + const uint other = subgroupShuffleXor(bh(h1[0]) | (bh(h1[1]) << 16u), 16u); + precise float16_t g0 = h1[0] + hb(other & 0xffffu), g1 = h1[1] + hb(other >> 16u); + precise float16_t s = g0 + g1; + precise float16_t se = s + float16_t(0.0001220703125); + const float16_t r = rsq16(se); + [[unroll]] for (uint m = 0u; m < MT; ++m) + [[unroll]] for (int i = 0; i < 8; ++i) { + const uint k = 16u * m + 8u * hl + uint(i); + precise float16_t g = r * w16.w[W_GAIN14 / 2u + k]; + precise float16_t n = rv[m][i] * g; + Nb[m][i] = q8h(n); + precise float16_t a = rv[m][i] + w16.w[W_TABLE16 / 2u + k]; + R[m][i] = float(a); + } + } + [[dont_unroll]] for (uint br = 0u; br < NB; ++br) { + const uint w1 = W_W1_0 + BRANCH * br, b1 = W_B1_0 + BRANCH * br, w2 = W_W2_0 + BRANCH * br; + B8 Hb[2]; + [[unroll]] for (uint c = 0u; c < 2u; ++c) { + ACC h; + [[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { + A8 a0, a1; + LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u); + K32(h, a0, Nb[kt], a1, Nb[kt + 1u]) + } + [[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i])); + } + [[unroll]] for (uint m = 0u; m < MT; ++m) { + A8 a0, a1; + LDA(a0, w2, 32u, m, 0u); LDA(a1, w2, 32u, m, 1u); + K32(R[m], a0, Hb[0], a1, Hb[1]) + } + } + + // ---------------------------------------------------------------- outputs + const uint t = inter(16u * wv + col); + const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u); + const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey; + B8 Tb[MT]; + { + const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u; +#if KIND == 0 && CX + const uint ty2 = t >> 3u, tx2 = t & 7u; + const uint mr = (ty2 >> 1u) < 2u ? 2u * (tx2 >> 1u) + (ty2 >> 1u) : 8u + 2u * (tx2 >> 1u) + (ty2 >> 1u) - 2u; + const uint blk = 2u * (ty2 & 1u) + (tx2 & 1u); + barrier(); // every wave is done with XT (the merge rows reuse it) +#endif + [[unroll]] for (uint m = 0u; m < MT; ++m) { + Q8(Tb[m], R[m]) + const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]); + const uint at = 4u * (m & 1u) + 2u * hl; + if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; } +#if KIND == 0 && CX + const uint g = mr * C + blk * (C / 4u) + 4u * m + 2u * hl; + MG[g] = w0; MG[g + 1u] = w1; +#endif + } + } + // the C -> 16 head + ACC hv; + [[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]); + [[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) { + A8 a0, a1; + LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u); + K32(hv, a0, Tb[kt], a1, Tb[kt + 1u]) + } +#if KIND == 0 + { // enc1's kernel-prediction filter (lanes 0-15; 16-31 hold head channels 8-15) + float16_t l[12]; + [[unroll]] for (int i = 0; i < 8; ++i) l[i] = float16_t(hv[i]); + const uint u0 = subgroupShuffleXor(bh(l[0]) | (bh(l[1]) << 16u), 16u), u1 = subgroupShuffleXor(bh(l[2]) | (bh(l[3]) << 16u), 16u); + l[8] = hb(u0 & 0xffffu); l[9] = hb(u0 >> 16u); l[10] = hb(u1 & 0xffffu); l[11] = hb(u1 >> 16u); + const float16_t ma = max(max(max(l[0], l[2]), l[4]), l[6]), mb = max(max(max(l[1], l[3]), l[5]), l[7]); + const float16_t mx = max(max(ma, mb), l[8]); + float e[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) { precise float16_t d = l[i] - mx; e[i] = exp2(float(d) * 1.4426950216293335); } + precise float s = (e[0] + e[1]) + 0.0; + s = s + (e[2] + e[3]); + s = s + (e[4] + e[5]); + s = s + (e[6] + e[7]); + s = s + e[8]; + const float16_t rr = rcp16(s); + float wg[9]; + [[unroll]] for (uint i = 0u; i < 9u; ++i) wg[i] = r16(r16(e[i]) * float(rr)); + const int x = tx, y = ty; + const int xm = x - 1 < 0 ? 1 : x - 1, xc = x < 0 ? 1 : (x >= ex ? ex - 2 : x), xp = x + 1 >= ex ? ex - 2 : x + 1; + const int ym = y - 1 < 0 ? 1 : y - 1, yc = y < 0 ? 1 : (y >= ey ? ey - 2 : y), yp = y + 1 >= ey ? ey - 2 : y + 1; + const ivec2 pts[9] = ivec2[9](ivec2(xm, ym), ivec2(xc, ym), ivec2(xp, ym), ivec2(xm, yc), ivec2(xc, yc), + ivec2(xp, yc), ivec2(xm, yp), ivec2(xc, yp), ivec2(xp, yp)); + float acc[3][2] = float[3][2](float[2](0.0, 0.0), float[2](0.0, 0.0), float[2](0.0, 0.0)); + [[unroll]] for (uint k = 0u; k < 8u; ++k) { + const vec4 tv = texelFetch(T88, pts[k], 0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) { precise float f = tv[c] * wg[k] + acc[c][k & 1u]; acc[c][k & 1u] = r16(f); } + } + const vec4 t8 = texelFetch(T88, pts[8], 0); + float o112[3], o120[3]; + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + precise float f0 = acc[c][0] + acc[c][1]; + precise float f1 = t8[c] * wg[8] + r16(f0); + const float f = r16(f1); +#if CX + precise float a = float(sig(l[9u + c])) * f; + const float ar = r16(a); + precise float b = f - ar; + o120[c] = ar; o112[c] = r16(b); +#else + o112[c] = f; // the bottleneck: the filtered value itself +#endif + } + if (hl == 0u && inside) { + imageStore(S112, ivec2(x, y), vec4(o112[0], o112[1], o112[2], 0.0)); +#if CX + imageStore(S120, ivec2(x, y), vec4(o120[0], o120[1], o120[2], 0.0)); +#endif + } +#if CX + if (hl == 0u) [[unroll]] for (uint c = 0u; c < 3u; ++c) O112[t * 4u + c] = float16_t(o112[c]); +#endif + } +#if CX + barrier(); + { // the 2x2 merge: merged^T (CX x 16 rows), channel tiles m = wv, wv + 4, ... + const uvec2 mxy = mergeXY(col); + const int X = x0 / 2 + int(mxy.x), Y = y0 / 2 + int(mxy.y); + const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u; + for (uint m = wv; m < CX / 16u; m += 4u) { + ACC a; + [[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]); + for (uint ks = 0u; ks < C / 8u; ++ks) { + B8 g0, g1; A8 w0, w1; + coopMatLoad(g0, MG, 8u * ks, C, CM); + coopMatLoad(g1, MG, 8u * ks + 4u, C, CM); + LDA(w0, W_WM, 4u * C, m, 2u * ks); LDA(w1, W_WM, 4u * C, m, 2u * ks + 1u); + K32(a, w0, g0, w1, g1) + } + B8 q; + Q8(q, a) + if (X >= 0 && Y >= 0 && X < hx && Y < hy) { + const uint at = o + (m >> 1u) * (lplane >> 2u) + 4u * (m & 1u) + 2u * hl; + arena.w[at] = pack4(q[0], q[1], q[2], q[3]); + arena.w[at + 1u] = pack4(q[4], q[5], q[6], q[7]); + } + } + if (wv == 3u && lane < 16u) { // surface 128: 2x2 sums of 0.25 * surface 112, rightwards first, then downwards + const uint mx = lane & 3u, my = lane >> 2u; + const int x = x0 + 2 * int(mx), y = y0 + 2 * int(my); + if (x >= 0 && y >= 0 && x < ex && y < ey) { + const uint t00 = (2u * my) * 8u + 2u * mx; + vec4 out_ = vec4(0.0); + [[unroll]] for (uint c = 0u; c < 3u; ++c) { + precise float16_t a00 = O112[t00 * 4u + c] * float16_t(0.25), a01 = O112[(t00 + 1u) * 4u + c] * float16_t(0.25); + precise float16_t a10 = O112[(t00 + 8u) * 4u + c] * float16_t(0.25), a11 = O112[(t00 + 9u) * 4u + c] * float16_t(0.25); + precise float16_t b0 = a00 + a01, b1 = a10 + a11; + precise float16_t cc = b0 + b1; + out_[c] = float(cc); + } + imageStore(S128, ivec2(x >> 1, y >> 1), out_); + } + } + } +#endif +#else + [[unroll]] for (int i = 0; i < 8; ++i) OUTF[t * 16u + 8u * hl + uint(i)] = float16_t(hv[i]); + barrier(); +#define DEC_FILTER +#include "rr_decfilter.glsl" +#endif +} diff --git a/engine/rr/swin_pack.py b/engine/rr/swin_pack.py new file mode 100644 index 0000000..804f9a9 --- /dev/null +++ b/engine/rr/swin_pack.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +"""Pack a Swin stage's extracted weights (engine/rr/swin_ref.py --extract) for engine/rr/swin_m.comp. + +usage: swin_pack.py STAGE WEIGHTS.npz OUT.bin OUT.h + +enc1_pack.py's layout for any stage of engine/rr/swin_ref.STAGES: transposed products, every weight's columns in the +k order of the product that consumes its result (slot j holds natural channel inter(n)[j]); the expand's 4C columns +are four sub-positions of C channels, each permuted the same way. OUT.h also defines the stage's shape for the shader +(KIND, C, CX, NH). +""" +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from enc1_ref import inter # noqa: E402 +from enc1_pack import a_tiles, b_tiles, cubic_table, act_table # noqa: E402 +from swin_ref import STAGES # noqa: E402 + + +def main(): + stage, src, out_bin, out_h = sys.argv[1:5] + kind, C, CX, NH, _ = STAGES[stage] + w = dict(np.load(src)) + I32, I96, IC = inter(32), inter(96), inter(C) + e = {} + if kind == 'dec': + I4C = [C * (j // C) + IC[j % C] for j in range(4 * C)] + e['we'], e['be'] = ('A', w['we'][:, I4C]), ('H', w['be'][I4C]) + for h in range(NH): + e[f'wk{h}'] = ('A', w[f'wk{h}'][:, I32]) + e[f'wv{h}'] = ('B', w[f'wv{h}'][:, I32]) + e[f'wo{h}'] = ('A', w[f'wo{h}'][:, IC]) + e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(inter(64), I96)].T) + for b in range(C // 8): + e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32]) + e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32]) + e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, IC]) + if kind == 'enc' and CX: + e['wm'], e['bm'] = ('A', w['wm'][:, inter(CX)]), ('H', w['bm'][inter(CX)]) + e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh']) + e['gain0'] = ('H', w['gain0']) + for k in ('table7', 'gain14', 'table16'): + e[k] = ('H', w[k][IC]) + e['cubic'] = ('U', cubic_table()) + e['act'] = ('Y', act_table()) + blob, defines = bytearray(), [] + for name, (kd, a) in e.items(): + while len(blob) % 256: + blob.append(0) + off = len(blob) + if kd == 'A': + blob += a_tiles(a).tobytes() + elif kd == 'B': + blob += b_tiles(a).tobytes() + elif kd == 'Y': + blob += np.ascontiguousarray(a, np.uint8).tobytes() + elif kd == 'U': + blob += np.ascontiguousarray(a, np.uint16).tobytes() + else: + blob += np.ascontiguousarray(a, np.float16).tobytes() + defines.append(f'#define W_{name.upper()} {off}u // {kd} {a.shape}') + shape = [f'#define KIND {0 if kind == "enc" else 1}', f'#define C {C}', f'#define CX {CX}', f'#define NH {NH}'] + if NH < 2: + defines.append('#define W_WK1 (W_WK0 + 1u)') + if C // 8 < 2: + defines.append('#define W_W1_1 (W_W1_0 + 1u)') + open(out_bin, 'wb').write(bytes(blob)) + open(out_h, 'w').write(f'// generated by engine/rr/swin_pack.py for {stage}\n' + '\n'.join(shape + defines) + '\n') + print(f'{stage}: {len(blob)} bytes, {len(e)} entries') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/swin_ref.py b/engine/rr/swin_ref.py new file mode 100644 index 0000000..ab38ebf --- /dev/null +++ b/engine/rr/swin_ref.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 +"""Exact numpy reference of the denoiser's Swin stages enc1-enc5 and dec1-dec4 (cuda_dldn_engine_swin_*_kernel), one +8x8 window at a time, from weights and real inputs only. Generalises engine/rr/enc1_ref.py (encoders) and +engine/rr/dec1_ref.py (decoders); every stage was checked against its capture (main). + +A stage has C channels (32 * KC), NH heads of 32 channels, NB = C / 8 MLP branches of C -> 32 -> C: + prologue encoders: the 96 key tokens from the arena (C / 32 planes); decoders: the expand of the 6 x 6 tokens + one level down (C_LO channels -> 4 sub-positions x C) + the skip token. RMS norm (input_norm: per lane word + the planes summed pairwise (pairs_tree), the two word halves, then lanes ^1, ^2 and the halves) times gain0; the + residual is the token + table7 + heads Q/K and V projections over KC k32 steps, scores + bias, the cubic softmax, P V, out-projection onto the + residual, head after head + MLP RMS norm (rowsum_sq) times gain14, residual + table16, NB branches + outputs the tokens; encoders: the 2x2 merge (4C -> C_M) and the 64 -> 16 head with enc1's filter (surfaces 112, + 120, 128); decoders: the head with dec0's filter (surface 112) + +usage: swin_ref.py CAPTURE STAGE FRAME WEIGHTS.npz [--blocks X,Y ...] (verifies against the capture) + swin_ref.py CAPTURE STAGE FRAME OUT.npz --extract DUMP.npz (weights from a phase dump) +""" +import argparse +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rrmath import e4m3, f16, q8 # noqa: E402 +from enc1_ref import (inter, mm, softmax_cubic, act, mirror, add, rowsum_sq, pairs_tree, filter_head, # noqa: E402 + MERGE_ORDER, H, F) +from dec0_ref import keys # noqa: E402 +from dec1_ref import filter112 # noqa: E402 + +# kind, C, C_LO (decoders: channels one level down) or C_M (encoders: merged channels), heads, warps +STAGES = { + 'enc1': ('enc', 64, 96, 2, 1), 'enc2': ('enc', 96, 128, 4, 4), 'enc3': ('enc', 128, 160, 4, 4), + 'enc4': ('enc', 160, 160, 8, 8), 'enc5': ('enc', 160, 0, 8, 8), + 'dec4': ('dec', 160, 160, 8, 8), 'dec3': ('dec', 128, 160, 4, 4), 'dec2': ('dec', 96, 128, 4, 4), + 'dec1': ('dec', 64, 96, 2, 2), +} + + +def input_norm(xn, gain): + """xn: one token, natural channels (f16). NVIDIA's tree over the arena order; returns the normalised token + (natural, f16)""" + c = len(xn) + xa = xn[inter(c)] + sq = (xa * xa).astype(H) + + def pair(p, w): + return add(sq[32 * p + 4 * w:32 * p + 4 * w + 2], sq[32 * p + 4 * w + 2:32 * p + 4 * w + 4]) + + fold = pairs_tree # (p0 + p1) + p2, (p0 + p1) + (p2 + p3), ((p0 + p1) + (p2 + p3)) + p4 + h = [add(fold([pair(p, t) for p in range(c // 32)]), fold([pair(p, t + 4) for p in range(c // 32)])) + for t in range(4)] + h01 = [add(h[t], h[t ^ 1]) for t in range(4)] + h2 = add(h01[0], h01[2]) + s = H(H(h2[0] + h2[1]) + H(2.0 ** -13)) + r = H(F(1) / np.sqrt(F(s))) + return (xn * (r * gain).astype(H)).astype(H) + + +def block(stage, w, x0, y0, ex, ey, inp): + """the window: returns (tokens 64 x C codes in arena order, merged 16 x C_M codes or None, head 64 x 16 f16). + inp: 'tok' [C/32][ey][ex][32] (encoders) or 'lo' [C_LO/32][ey/2][ex/2][32] and 'skip' [C/32][ey][ex][32]""" + kind, C, CX, NH, _ = STAGES[stage] + IC = inter(C) + UN = [IC.index(i) for i in range(C)] + xs = [] + if kind == 'dec': + hx, hy = ex // 2, ey // 2 + X0, Y0 = x0 // 2 - 1, y0 // 2 - 1 + low = np.zeros((48, CX), np.uint8) + for r in range(36): + X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy) + low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)]) + Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0))) + for (x, y) in keys(x0, y0): + xx, yy = mirror(x, ex), mirror(y, ey) + r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1) + sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN] + xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H)) + else: + for (x, y) in keys(x0, y0): + xx, yy = mirror(x, ex), mirror(y, ey) + xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H)) + X = np.array([q8(input_norm(x, w['gain0']))[IC] for x in xs]) + R = add(np.array(xs[:64]), w['table7'][None, :]) + I32, I96 = inter(32), inter(96) + for h in range(NH): + qk, v = np.zeros((96, 32)), np.zeros((96, 32)) + for c in range(C // 32): + qk = mm(X[:, 32 * c:32 * c + 32], w[f'wk{h}'][32 * c:32 * c + 32], qk) + v = mm(X[:, 32 * c:32 * c + 32], w[f'wv{h}'][32 * c:32 * c + 32], v) + Q = q8(qk)[:, I32] + P = softmax_cubic(mm(Q[:64], Q.T, w[f'bias{h}'])) + O = mm(q8(P)[:, I96], q8(v)[I96], np.zeros((64, 32))) + R = mm(q8(O)[:, I32], w[f'wo{h}'], R) + N = q8((R * (rowsum_sq(R)[:, None] * w['gain14'][None, :]).astype(H)).astype(H))[:, IC] + acc = add(R, w['table16'][None, :]) + for b in range(C // 8): + hid = act(mm(N, w[f'w1_{b}'], np.repeat(w[f'b1_{b}'][None, :], 64, 0))) + acc = mm(q8(hid)[:, I32], w[f'w2_{b}'], acc) + tokens = q8(acc)[:, IC] + merged = None + if kind == 'enc' and CX: + Mg = np.zeros((16, 4 * C), np.uint8) + for r, (mx, my) in enumerate(MERGE_ORDER): + for blk, (dx, dy) in enumerate([(0, 0), (1, 0), (0, 1), (1, 1)]): + Mg[r, C * blk:C * blk + C] = tokens[(2 * my + dy) * 8 + 2 * mx + dx] + merged = q8(mm(Mg, w['wm'], np.repeat(w['bm'][None, :], 16, 0)))[:, inter(CX)] + head = mm(tokens, w['wh'], np.repeat(w['bh'][None, :], 64, 0)) + return tokens, merged, head + + +def phase_map(stage): + """dump phase numbers: projections [h][c] = (QK, V), bias[h] (scores phase, wave), out[h], mlp up / down, merge, + head, expand""" + kind, C, CX, NH, waves = STAGES[stage] + KC = C // 32 + m = {'proj': {}, 'bias': {}, 'out': {}} + if waves == 1: # heads in sequence: per head 2 KC projections, scores, P V, out-projection + p = 1 + for h in range(NH): + for c in range(KC): + m['proj'][h, c] = (p, p + 1, 0) + p += 2 + m['bias'][h] = (p, 0) + m['out'][h] = p + 2 + p += 3 + else: # one head per warp + p = 2 if kind == 'dec' else 1 + for h in range(NH): + for c in range(KC): + m['proj'][h, c] = (p + 2 * c, p + 2 * c + 1, h) + p += 2 * KC + for h in range(NH): + m['bias'][h] = (p, h) + m['out'][h] = p + 2 + h + p += 2 + NH + m['mlp'] = [(p + 2 * b, p + 2 * b + 1) for b in range(C // 8)] + p += 2 * (C // 8) + if kind == 'enc' and CX: + m['merge'] = p + p += 1 + m['head'] = p + return m + + +def read_b(W, base, K): + """an FP8 B matrix (K x 32) as NVIDIA stores it in the weight image: 4-byte k groups, offset 512 (kg >> 3) + + 16 (kg & 3) + 4 ((kg >> 2) & 1) + 64 (n & 7) + 8 ((n >> 3) & 1) + 16 K (n >> 4)""" + b = np.zeros((K, 32), np.uint8) + for kg in range(K // 4): + for n in range(32): + o = base + 512 * (kg >> 3) + 16 * (kg & 3) + 4 * ((kg >> 2) & 1) + 64 * (n & 7) + 8 * ((n >> 3) & 1) + 16 * K * (n >> 4) + b[4 * kg:4 * kg + 4, n] = W[o:o + 4] + return b + + +def merge_weights(z, W, p, waves, C, CX): + """the merge (4C -> CX) of a multi-warp encoder: warp w computes the 32-column chunks w, w + warps, ...; the dump + holds each warp's last chunk, so the chunks come from the weight image (located through warp 1's, chunk 1) and the + dumped ones check the layout""" + K, size = 4 * C, 4 * C * 32 + W32 = W[:len(W) // 4 * 4].view(np.uint32) + first = z[f'p{p}_w1_B'][0:4, 0].copy().view(np.uint32)[0] + pos1 = int(next(i for i in np.nonzero(W32 == first)[0] * 4 if np.array_equal(read_b(W, int(i), K), z[f'p{p}_w1_B']))) + base0 = pos1 - size + chunks = [read_b(W, base0 + size * j, K) for j in range(CX // 32)] + for v in range(waves): + last = max(j for j in range(CX // 32) if j % waves == v) if any(j % waves == v for j in range(CX // 32)) else None + if last is not None: + assert np.array_equal(chunks[last], z[f'p{p}_w{v}_B']), (v, last) + b1 = f16(z[f'p{p}_w1_C']).astype(H)[0].tobytes() + bpos1 = W.tobytes().find(b1) + bias = np.concatenate([W[bpos1 + 64 * (j - 1):bpos1 + 64 * j].view(H) for j in range(CX // 32)]) + for v in range(waves): + last = max((j for j in range(CX // 32) if j % waves == v), default=None) + if last is not None: + assert np.array_equal(bias[32 * last:32 * last + 32].view(np.uint16), z[f'p{p}_w{v}_C'][0]), ('bias', v) + return np.hstack(chunks), bias + + +def extract(stage, dump, weights, out, inp, x0, y0, ex, ey): + """weights from the phase dump; gains and tables are found in the weight image by matching the dump""" + kind, C, CX, NH, waves = STAGES[stage] + z = np.load(dump) + W = np.frombuffer(weights, np.uint8) + + def B(p, wv=0): + return z[f'p{p}_w{wv}_B'] + + def Cm(p, wv=0): + return f16(z[f'p{p}_w{wv}_C']).astype(H) + pm = phase_map(stage) + w = {} + if kind == 'dec': + w['we'] = np.hstack([B(1, v) for v in range(waves)]) + w['be'] = np.concatenate([Cm(1, v)[0] for v in range(waves)]) + for h in range(NH): + w[f'wk{h}'] = np.vstack([B(pm['proj'][h, c][0], pm['proj'][h, c][2]) for c in range(C // 32)]) + w[f'wv{h}'] = np.vstack([B(pm['proj'][h, c][1], pm['proj'][h, c][2]) for c in range(C // 32)]) + w[f'bias{h}'] = Cm(*pm['bias'][h]) + w[f'wo{h}'] = B(pm['out'][h]) + for b, (up, dn) in enumerate(pm['mlp']): + w[f'w1_{b}'], w[f'b1_{b}'], w[f'w2_{b}'] = B(up), Cm(up)[0], B(dn) + if 'merge' in pm: + mw = min(waves, 4) # the merge runs on (at most) four warps + if waves > 1 and CX // 32 > mw: # a warp computes several chunks: read them from the weight image + w['wm'], w['bm'] = merge_weights(z, W, pm['merge'], mw, C, CX) + elif waves > 1: + w['wm'] = np.hstack([B(pm['merge'], v) for v in range(CX // 32)]) + w['bm'] = np.concatenate([Cm(pm['merge'], v)[0] for v in range(CX // 32)]) + else: + w['wm'], w['bm'] = B(pm['merge']), Cm(pm['merge'])[0] + w['wh'], w['bh'] = B(pm['head']), Cm(pm['head'])[0] + # the four f16 vectors: candidates at every even offset, checked against the first token rows of the dump + tables = [W[o:o + 2 * C].view(H) for o in range(0, len(W) - 2 * C + 1, 2)] + w0 = {k: v for k, v in w.items()} + # gain0: the key tokens' norm against the first projection's A (window row 0..) + xs = key_tokens(stage, w0, x0, y0, ex, ey, inp) + IC = inter(C) + A1 = np.hstack([z[f'p{pm["proj"][0, c][0]}_w{pm["proj"][0, c][2]}_A'] for c in range(C // 32)]) + w['gain0'] = next(t for t in tables if np.array_equal(q8(input_norm(xs[0], t))[IC], A1[0]) and + np.array_equal(q8(input_norm(xs[5], t))[IC], A1[5])) + p_out = pm['out'][0] + C0 = Cm(p_out, 0) + rows = np.array(xs[:64]) + w['table7'] = next(t for t in tables if np.array_equal(add(rows[:len(C0)][:2], t[None]).view(np.uint16), C0[:2].view(np.uint16))) + up, dn = pm['mlp'][0] + Rt = f16(z[f'p{pm["out"][NH - 1]}_w0_D']).astype(H) + Aup, Cdn = z[f'p{up}_w0_A'], Cm(dn) + w['gain14'] = next(t for t in tables if np.array_equal( + q8((Rt[:2] * (rowsum_sq(Rt[:2])[:, None] * t[None]).astype(H)).astype(H))[:, IC], Aup[:2])) + w['table16'] = next(t for t in tables if np.array_equal(add(Rt[:2], t[None]).view(np.uint16), Cdn[:2].view(np.uint16))) + np.savez_compressed(out, **w) + + +def key_tokens(stage, w, x0, y0, ex, ey, inp): + kind, C, CX, NH, _ = STAGES[stage] + IC = inter(C) + UN = [IC.index(i) for i in range(C)] + xs = [] + if kind == 'dec': + hx, hy = ex // 2, ey // 2 + X0, Y0 = x0 // 2 - 1, y0 // 2 - 1 + low = np.zeros((48, CX), np.uint8) + for r in range(36): + X, Y = mirror(X0 + r % 6, hx), mirror(Y0 + r // 6, hy) + low[r] = np.concatenate([inp['lo'][p, Y, X] for p in range(CX // 32)]) + Ex = q8(mm(low, w['we'], np.repeat(w['be'][None, :], 48, 0))) + for (x, y) in keys(x0, y0): + xx, yy = mirror(x, ex), mirror(y, ey) + if kind == 'dec': + r, s = (y // 2 - Y0) * 6 + (x // 2 - X0), 2 * (y & 1) + (x & 1) + sk = np.concatenate([inp['skip'][p, yy, xx] for p in range(C // 32)])[UN] + xs.append((e4m3(Ex[r, C * s:C * s + C]).astype(H) + e4m3(sk).astype(H)).astype(H)) + else: + xs.append(e4m3(np.concatenate([inp['tok'][p, yy, xx] for p in range(C // 32)])[UN]).astype(H)) + return xs + + +def load(cap, stage, frame): + from rr_capture import Capture + c = Capture(os.path.expanduser(cap)) + l = c.find(f'swin_{stage}', frame) + kind, C, CX, NH, _ = STAGES[stage] + ex, ey = (int(v) for v in np.frombuffer(l.args[:8], np.uint32)) + sx, sy = (int(v) for v in np.frombuffer(l.args[80:84], np.uint16)) + ptr = {p.offset: p.delta for p in l.pointers} + pre = np.frombuffer(l.pre(l.arena_alloc), np.uint8) + post = np.frombuffer(l.post_buffer(c, l.arena_alloc), np.uint8) + inp = {} + if kind == 'dec': + inp['lo'] = pre[ptr[8]:ptr[8] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32) + inp['skip'] = pre[ptr[24]:ptr[24] + C * ex * ey].reshape(C // 32, ey, ex, 32) + else: + inp['tok'] = pre[ptr[8]:ptr[8] + C * ex * ey].reshape(C // 32, ey, ex, 32) + outs = {'tok': post[ptr[48]:ptr[48] + C * ex * ey].reshape(C // 32, ey, ex, 32)} + if kind == 'enc' and CX: + outs['mrg'] = post[ptr[56]:ptr[56] + CX * (ex // 2) * (ey // 2)].reshape(CX // 32, ey // 2, ex // 2, 32) + img = {i.offset: i for i in l.images} + + def surf(off): + i = img[off] + return np.frombuffer(i.path.read_bytes(), np.float16).reshape(i.height, i.width, 4) + tex = {k: surf(k) for k in img} + return l, ex, ey, sx, sy, inp, outs, tex + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('stage') + ap.add_argument('frame', type=int) + ap.add_argument('weights') + ap.add_argument('--extract', metavar='DUMP') + ap.add_argument('--dump-block', default='1,1', help='the block the dump was taken of (for --extract)') + ap.add_argument('--blocks', nargs='*', default=['1,1']) + o = ap.parse_args() + l, ex, ey, sx, sy, inp, outs, tex = load(o.capture, o.stage, o.frame) + kind, C, CX, NH, _ = STAGES[o.stage] + if o.extract: + bx, by = (int(v) for v in o.dump_block.split(',')) + extract(o.stage, o.extract, l.weights, o.weights, inp, 8 * bx - sx, 8 * by - sy, ex, ey) + return + w = dict(np.load(o.weights)) + for blk in o.blocks: + bx, by = (int(v) for v in blk.split(',')) + x0, y0 = 8 * bx - sx, 8 * by - sy + tokens, merged, head = block(o.stage, w, x0, y0, ex, ey, inp) + inside = [t for t in range(64) if 0 <= x0 + t % 8 < ex and 0 <= y0 + t // 8 < ey] + tk = sum(np.array_equal(tokens[t], np.concatenate([outs['tok'][p, y0 + t // 8, x0 + t % 8] for p in range(C // 32)])) + for t in inside) + msg = f'block {bx},{by}: tokens {tk}/{len(inside)}' + if merged is not None: + mg = sum(np.array_equal(merged[r], np.concatenate([outs['mrg'][p, y0 // 2 + my, x0 // 2 + mx] for p in range(CX // 32)])) + for r, (mx, my) in enumerate(MERGE_ORDER) if 0 <= x0 // 2 + mx < ex // 2 and 0 <= y0 // 2 + my < ey // 2) + msg += f', merged {mg}' + if kind == 'enc' and not CX: # the bottleneck: the ungated filter to surface 112 only + o112, _ = filter_head(head, tex[88], x0, y0, ex, ey, ungated=True) + ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside) + elif kind == 'enc': + o112, o120 = filter_head(head, tex[88], x0, y0, ex, ey) + ok = sum(np.array_equal(o112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) and + np.array_equal(o120[t], tex[120][y0 + t // 8, x0 + t % 8, :3]) for t in inside) + else: + s112 = filter112(head, tex[88], tex[96], tex[104], x0, y0, ex, ey) + ok = sum(np.array_equal(s112[t], tex[112][y0 + t // 8, x0 + t % 8, :3]) for t in inside) + print(msg + f', surfaces {ok}/{len(inside)}') + + +if __name__ == '__main__': + main() diff --git a/engine/rr/test_layer.py b/engine/rr/test_layer.py new file mode 100755 index 0000000..197200d --- /dev/null +++ b/engine/rr/test_layer.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +"""Run one Ray Reconstruction stage shader on a captured launch and compare it with NVIDIA's result. + +usage: test_layer.py CAPTURE KERNEL FRAME SHADER.spv WORK [--kvk build/rr/kvk] [--time N] [--subgroup 32|64] + +KERNEL is a substring of the denoiser kernel's name (e.g. swin_enc1, hkpn_output_kernel_transformer), FRAME the +1-based occurrence in the capture. The shader follows the RR stage contract (engine/rr/README.md): + + binding 0 the parameter block: NVIDIA's argument block with every pointer into the plane arena + replaced by its byte offset in the arena, the weight pointer and every texture or surface + handle by zero (stage_params below; the runtime builds the same block) + binding 1 the weight image exactly as NGX uploads it (argument 40), read-only + binding 2 the plane arena, read/write + bindings 3, 4, ... the textures in ascending argument offset, then the surfaces in ascending argument offset + +Every surface starts zeroed and the arena starts as it was before NVIDIA's launch; afterwards each surface and the +whole arena are compared with what NVIDIA's launch left. Exit status 0 only when everything is byte-identical. +""" +import argparse +import json +import os +import struct +import subprocess +import sys +from pathlib import Path + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rr_capture import Capture # noqa: E402 + + +# GLSL image formats of the stage shaders' surfaces -> (channels, kvk bits) +SURFACE_FORMATS = {'rgba16f': (4, 16), 'rg16f': (2, 16), 'r16f': (1, 16), 'rgba8_snorm': (4, 9), 'rgba8': (4, 8)} + + +def stage_params(launch): + args = bytearray(launch.args) + arena = launch.arena_alloc + for p in launch.pointers: + value = p.delta if p.alloc == arena and p.offset != 40 else 0 + struct.pack_into('> 7, (c >> 3) & 15, c & 7 + v = np.nan if (e == 15 and m == 7) else (m * 2.0 ** -9 if e == 0 else (1 + m / 8) * 2.0 ** (e - 7)) + t[c] = -v if s else v + return t + + +def psnr(a, b): + finite = np.isfinite(a) & np.isfinite(b) + if not finite.any(): + return float('nan') + err = np.mean((a[finite] - b[finite]) ** 2) + peak = max(np.max(np.abs(b[finite])), 1e-6) + return float('inf') if err == 0 else 10 * np.log10(peak * peak / err) + + +def compare_image(name, image, got, ref, ignore_channels=()): + """IGNORE_CHANNELS: channels NVIDIA's kernel fills from an undefined register (left out of the comparison)""" + if image.bits == 16: + a, b = np.frombuffer(got, np.float16).astype(np.float64), np.frombuffer(ref, np.float16).astype(np.float64) + elif image.bits == 32: + a, b = np.frombuffer(got, np.float32).astype(np.float64), np.frombuffer(ref, np.float32).astype(np.float64) + else: + dt = np.int8 if image.bits == 9 else np.uint8 + a, b = np.frombuffer(got, dt).astype(np.float64), np.frombuffer(ref, dt).astype(np.float64) + keep = np.ones(len(a), bool) + for c in ignore_channels: + keep[c::image.channels] = False + a, b = a[keep], b[keep] + element_bytes = len(got) // len(keep) + same = (np.frombuffer(got, np.uint8).reshape(-1, element_bytes) == np.frombuffer(ref, np.uint8).reshape(-1, element_bytes))[keep] + elements = len(a) + equal = np.sum(a == b) + np.sum(np.isnan(a) & np.isnan(b)) + worst = np.nanmax(np.abs(a - b)) if elements else 0 + print(f' {name}: {image.width}x{image.height}x{image.channels} fmt{image.format}' + f'{" (channels " + ",".join(map(str, ignore_channels)) + " ignored)" if ignore_channels else ""}: ' + f'{100.0 * equal / elements:.4f}% identical values, max |diff| {worst:.6g}, PSNR {psnr(a, b):.2f} dB, ' + f'nonfinite ours/ref {int(np.sum(~np.isfinite(a)))}/{int(np.sum(~np.isfinite(b)))}') + return bool(same.all()) + + +def compare_arena(got, before, ref): + g, p, r = (np.frombuffer(x, np.uint8) for x in (got, before, ref)) + written = (r != p) | (g != p) + differ = g != r + print(f' arena: {int(written.sum())} bytes written by either side, {int(differ.sum())} differ', end='') + if differ.any(): + idx = np.nonzero(differ)[0] + t = e4m3_table() + dv = np.abs(t[g[idx]] - t[r[idx]]) + print(f' (first at byte {idx[0]:#x}, last at {idx[-1]:#x}; as e4m3 values: max |diff| {np.nanmax(dv):.6g}, ' + f'{int(np.sum(dv <= 0.0625 * np.abs(t[r[idx]]) + 1e-12))} within one e4m3 step)') + else: + print() + return not differ.any() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('kernel') + ap.add_argument('frame', type=int) + ap.add_argument('shader') + ap.add_argument('work') + ap.add_argument('--kvk', default=str(Path(__file__).resolve().parents[2] / 'build/rr/kvk')) + ap.add_argument('--time', type=int, default=1) + ap.add_argument('--subgroup', type=int, default=32) + ap.add_argument('--weights', help='bind this file at binding 1 instead of NGX\'s weight image (native stages)') + ap.add_argument('--grid', help='X,Y workgroups (default: the launch grid)') + ap.add_argument('--ignore', action='append', default=[], metavar='SURFACE:CHANNEL', + help='leave a channel out of the comparison, e.g. 112:3 for a value NVIDIA stores from an ' + 'undefined register (repeatable)') + o = ap.parse_args() + + cap = Capture(o.capture) + launch = [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or + ('diamond' in o.kernel and o.kernel in l.kernel)][o.frame - 1] + work = Path(o.work) + work.mkdir(parents=True, exist_ok=True) + print(f'{launch.kernel} (launch {launch.sequence}, grid {launch.grid}, block {launch.block})') + + (work / 'params.bin').write_bytes(stage_params(launch)) + (work / 'weights.bin').write_bytes(Path(o.weights).read_bytes() if o.weights else launch.weights) + arena = launch.arena_alloc + arena_before = launch.pre(arena) + (work / 'arena.bin').write_bytes(arena_before) + gx, gy = (int(v) for v in o.grid.split(',')) if o.grid else launch.grid[:2] + lines = [f'param {work}/params.bin', f'grid {gx} {gy}'] + # Bind in the shader's order (its ptx2rr.py sidecar, SHADER.json: textures then surfaces by argument offset). + # A handle the shader reads but NGX left null gets a zeroed 1x1 image. + sidecar = Path(o.shader).with_suffix('.json') + if sidecar.exists(): + info = json.loads(sidecar.read_text()) + order = [('texture', off, None) for off in info['textures']] + [('surface', off, f) for off, f in info['surfaces']] + else: + order = [(i.kind, i.offset, None) for i in launch.images] + by_key = {(i.kind, i.offset): i for i in launch.images} + (work / 'zero.bin').write_bytes(bytes(64)) + for kind, offset, fmt in order: + image = by_key.get((kind, offset)) + if image is None: + ch, bits = SURFACE_FORMATS.get(fmt, (4, 16)) + lines.append(f'tex {offset} {work}/zero.bin 1 1 4 16 nearest' if kind == 'texture' else + f'surf {offset} 1 1 {ch} {bits} - {work}/dummy{offset}.bin') + elif kind == 'texture': + lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} {image.bits} ' + f'{"linear" if image.linear else "nearest"}') + else: + lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {image.bits} - ' + f'{work}/surf{image.offset}.bin') + n_images = len(order) + lines.append(f'lut {work}/weights.bin') + lines.append(f'inout {work}/arena.bin {work}/arena.out.bin') + probes = info.get('probes', []) if sidecar.exists() else [] + oob = info.get('oob', False) if sidecar.exists() else False + if probes: # ptx2rr.py --probe: the debug buffer after the surfaces + threads = launch.block[0] * launch.block[1] * launch.block[2] + lines.append(f'outbuf {work}/debug.bin {4 * threads * len(probes)}') + elif oob: # ptx2rr.py --oob-report: count, lowest and highest address per buffer + (work / 'oob.init').write_bytes(struct.pack('<9I', *([0, 0xffffffff, 0] * 3))) + lines.append(f'inout {work}/oob.init {work}/oob.bin') + lines.append('bindings 0 ' + ' '.join(str(3 + i) for i in range(n_images)) + ' 1 2' + + (f' {3 + n_images}' if probes or oob else '')) + (work / 'manifest').write_text('\n'.join(lines) + '\n') + r = subprocess.run([o.kvk, o.shader, str(work / 'manifest'), str(o.time), str(o.subgroup)], + capture_output=True, text=True) + sys.stdout.write(r.stdout) + if r.returncode != 0: + sys.stderr.write(r.stderr) + sys.exit(f'kvk failed ({r.returncode})') + + identical = True + bound = {(kind, offset) for kind, offset, _ in order} + for image in launch.images: + if image.kind != 'surface': + continue + if ('surface', image.offset) not in bound: + print(f' surface {image.offset}: not written by this kernel') + continue + got = (work / f'surf{image.offset}.bin').read_bytes() + waived = [int(s.split(':')[1]) for s in o.ignore if int(s.split(':')[0]) == image.offset] + identical &= compare_image(f'surface {image.offset}', image, got, image.path.read_bytes(), waived) + if arena is not None: + identical &= compare_arena((work / 'arena.out.bin').read_bytes(), arena_before, + launch.post_buffer(cap, arena)) + if oob: + r = struct.unpack('<9I', (work / 'oob.bin').read_bytes()[:36]) + for k, name in enumerate(('weight reads', 'arena reads', 'arena writes')): + print(f' out of range {name}: {r[3 * k]}' + (f' (bytes {r[3 * k + 1]:#x}..{r[3 * k + 2]:#x}, size ' + f'{len(launch.weights) if k == 0 else len(arena_before):#x})' if r[3 * k] else '')) + print('IDENTICAL' if identical else 'DIFFERENT') + sys.exit(0 if identical else 1) + + +if __name__ == '__main__': + main() diff --git a/engine/rr/zluda_replay.py b/engine/rr/zluda_replay.py new file mode 100644 index 0000000..020e21d --- /dev/null +++ b/engine/rr/zluda_replay.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Replay one captured denoiser launch through ZLUDA on Linux (tools/replay_rr_launch.cpp) and compare it with the +capture, like test_layer.py does for a stage shader. + +usage: zluda_replay.py CAPTURE KERNEL FRAME WORK [--ptx MODULE] [--zluda DIR] + +KERNEL and FRAME select the launch as in test_layer.py. MODULE defaults to the corpus module holding the kernel +(~/.cache/d4r-rr-corpus); pass an edited copy to add debug stores. ZLUDA's own settings (D4R_ZLUDA_WMMA, +D4R_ZLUDA_WMMA_FP8, D4R_ZLUDA_WMMA_FP8_NATIVE, ...) are taken from the environment. Outputs land in WORK +(surfOFFSET.bin, arena.out.bin) so the two replays can be compared with each other too. +""" +import argparse +import os +import subprocess +import sys +from pathlib import Path + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from rr_capture import Capture # noqa: E402 +from test_layer import compare_arena, compare_image # noqa: E402 + +ROOT = Path(__file__).resolve().parents[2] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('capture') + ap.add_argument('kernel') + ap.add_argument('frame', type=int) + ap.add_argument('work') + ap.add_argument('--ptx') + ap.add_argument('--zluda', default=os.path.expanduser('~/.cache/d4r-zluda-rr')) + ap.add_argument('--rocm', default='/opt/rocm/lib') + ap.add_argument('--pad', type=int, default=0, help='bytes of zeros appended to the arena (debug stores)') + o = ap.parse_args() + cap = Capture(o.capture) + launch = [l for l in cap.launches if o.kernel in l.kernel and 'diamond' not in l.kernel or + ('diamond' in o.kernel and o.kernel in l.kernel)][o.frame - 1] + work = Path(o.work) + work.mkdir(parents=True, exist_ok=True) + module = o.ptx or next(str(p) for p in sorted(Path('~/.cache/d4r-rr-corpus').expanduser().glob('*.ptx')) + if f'.entry {launch.kernel}(' in p.read_text(errors='replace')) + (work / 'args.bin').write_bytes(launch.args) + arena = launch.arena_alloc + lines = [f'args {work}/args.bin', 'grid ' + ' '.join(map(str, launch.grid)), + 'block ' + ' '.join(map(str, launch.block))] + if o.pad: + (work / 'arena-pad.bin').write_bytes(launch.pre(arena) + bytes(o.pad)) + for p in launch.pointers: + out = f' {work}/arena.out.bin' if p.alloc == arena else '' + source = f'{work}/arena-pad.bin' if o.pad and p.alloc == arena else f'{launch.directory}/alloc-{p.alloc}.bin' + lines.append(f'buf {p.offset} {source} {p.delta}{out}') + for image in launch.images: + if image.kind == 'texture': + # SNORM_INT8X4 (200): a SIGNED_INT8 array read as normalized floats gives the same values + fmt = 8 if image.format == 200 else image.format + lines.append(f'tex {image.offset} {image.path} {image.width} {image.height} {image.channels} ' + f'{fmt} {"linear" if image.linear else "nearest"}') + else: + # ZLUDA has no SNORM_INT8X4 arrays (format 200); surface stores write the same bytes to a SIGNED_INT8 one + fmt = 8 if image.format == 200 else image.format + lines.append(f'surf {image.offset} {image.width} {image.height} {image.channels} {fmt} ' + f'{work}/surf{image.offset}.bin') + (work / 'manifest').write_text('\n'.join(lines) + '\n') + env = dict(os.environ, LD_LIBRARY_PATH=f'{o.zluda}:{o.rocm}:' + os.environ.get('LD_LIBRARY_PATH', '')) + print(f'{launch.kernel} (launch {launch.sequence}) through ZLUDA {o.zluda}') + r = subprocess.run([str(ROOT / 'build/replay_rr_launch'), module, launch.kernel, str(work / 'manifest')], + env=env, capture_output=True, text=True) + sys.stdout.write(r.stdout) + if r.returncode: + sys.stderr.write(r.stderr[-4000:]) + sys.exit('replay failed') + same = True + for image in launch.images: + if image.kind == 'surface': + same &= compare_image(f'surface {image.offset}', image, (work / f'surf{image.offset}.bin').read_bytes(), + image.path.read_bytes()) + if arena is not None and not o.pad: + same &= compare_arena((work / 'arena.out.bin').read_bytes(), launch.pre(arena), launch.post_buffer(cap, arena)) + print('IDENTICAL' if same else 'DIFFERENT') + + +if __name__ == '__main__': + main() diff --git a/engine/runtime.cpp b/engine/runtime.cpp index 1377736..1b304f8 100644 --- a/engine/runtime.cpp +++ b/engine/runtime.cpp @@ -2,8 +2,11 @@ #include "vk_dispatch.h" #include #include +#include #include #include +#include +#include #include #include @@ -139,7 +142,7 @@ struct Pipeline { if (descriptors) vkDestroyDescriptorSetLayout(device, descriptors, nullptr); } void create(VkDevice dev, const std::string& file, const std::vector& types, - uint32_t subgroup, VkPipelineCache cache) { + uint32_t subgroup, VkPipelineCache cache, VkPipelineCreateFlags flags = 0) { device = dev; std::vector bindings; for (uint32_t i = 0; i < types.size(); ++i) @@ -161,7 +164,7 @@ struct Pipeline { VkPipelineShaderStageRequiredSubgroupSizeCreateInfo required{VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_REQUIRED_SUBGROUP_SIZE_CREATE_INFO}; required.requiredSubgroupSize = subgroup; VkComputePipelineCreateInfo ci{VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO}; - ci.layout = layout; + ci.layout = layout; ci.flags = flags; ci.stage = {VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, &required, VK_PIPELINE_SHADER_STAGE_CREATE_REQUIRE_FULL_SUBGROUPS_BIT, VK_SHADER_STAGE_COMPUTE_BIT, shader, "main", nullptr}; VkResult result = vkCreateComputePipelines(dev, cache, 1, &ci, nullptr, &handle); @@ -779,4 +782,5 @@ VkDeviceSize Engine::networkHeadBytes(uint32_t outputWidth, uint32_t outputHeigh } VkDeviceSize Engine::activationBytes() const { return impl->activations; } #include "runtime_m.inc" +#include "runtime_rr.inc" } // namespace d4r diff --git a/engine/runtime.h b/engine/runtime.h index a29ae66..3956a05 100644 --- a/engine/runtime.h +++ b/engine/runtime.h @@ -244,4 +244,64 @@ private: struct Impl; std::unique_ptr impl; }; + +// ---------------------------------------------------------------- Ray Reconstruction (DLSS-D, presets D and E) +// NVIDIA's denoiser frame as d4r's own shaders (engine/rr): exposure, six encoders, five decoders and the output +// stage, each reading and writing NGX's layouts so it can be checked against a captured launch. The network runs at +// the render size, the output stage at the output size. Same conventions as the classes above. +struct RFrameImages { + // Render-resolution inputs, sampled: colour, motion (two channels, before MV.Scale), depth (one channel), and the + // guides the denoiser reads (GBuffer.Normals, GBuffer.Roughness, the diffuse and specular albedo). + VkImageView color, motion, depth, normals, roughness, diffuseAlbedo, specularAlbedo; + // Optional inputs (null: the game passes none): DLSSD.SpecularHitDistance and DLSS.DisocclusionMask, one channel each + VkImageView specularHitDistance = VK_NULL_HANDLE, disocclusionMask = VK_NULL_HANDLE; + VkImageLayout layout = VK_IMAGE_LAYOUT_GENERAL; // of every view +}; +struct RFrameSettings { + float jitter[2] = {0, 0}; // Jitter.Offset, render pixels + float motionScale[2] = {1, 1}; // MV.Scale + float preExposure = 1; // DLSS.Pre.Exposure + float frameTimeMs = 16.666667f; // FrameTimeDeltaInMsec + bool reset = false; + bool depthInverted = false; // creation flag NVSDK_NGX_DLSS_Feature_Flags_DepthInverted + bool roughnessPacked = false; // DLSS.Roughness.Mode 1: the roughness is the normals' w (no roughness texture) + float worldToView[16] = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; // row-major, as the SDK passes them + float viewToClip[16] = {1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; +}; +class RModel { +public: + RModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache = VK_NULL_HANDLE); + ~RModel(); + RModel(const RModel&) = delete; + RModel& operator=(const RModel&) = delete; + void recordUpload(VkCommandBuffer command); + bool uploadRecorded() const; + uint32_t preset() const; // NVSDK_NGX_RayReconstruction_Hint_Render_Preset: 4 = D, 5 = E +private: + struct Impl; + std::unique_ptr impl; + friend class REngine; +}; +class REngine { +public: + REngine(RModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight); + ~REngine(); + REngine(const REngine&) = delete; + REngine& operator=(const REngine&) = delete; + void recordInitialize(VkCommandBuffer command); + // A borrowed RGBA16F finalOutput (GENERAL, origin zero) receives the presented image as a copy after the frame; + // otherwise read outputImage(). Views must outlive this engine. + void setFrame(const RFrameImages& images, const RFrameSettings& settings, FrameOutput finalOutput = {}); + // With a timestamp query pool, four stamps are written from query `first`: start, after the exposure, after the + // network (enc0 to dec0) and after the output stage. + void recordFrame(VkCommandBuffer command, VkQueryPool stamps = VK_NULL_HANDLE, uint32_t first = 0); + VkImage outputImage() const; // RGBA16F, GENERAL: the presented image of the last recorded frame + // Debugging (engine/replay_rr.cpp): every image the frame writes, named after the stage and argument offset that + // writes it ("enc1-112", "dec0-144", "q0", ...), GENERAL layout + struct DebugImage { std::string name; VkImage image; VkFormat format; uint32_t width, height; }; + std::vector debugImages() const; +private: + struct Impl; + std::unique_ptr impl; +}; } // namespace d4r diff --git a/engine/runtime_rr.inc b/engine/runtime_rr.inc new file mode 100644 index 0000000..ed83500 --- /dev/null +++ b/engine/runtime_rr.inc @@ -0,0 +1,551 @@ +// Ray Reconstruction (DLSS-D presets D and E): the runtime around engine/rr's stage shaders. Included by runtime.cpp +// inside namespace d4r, after runtime_m.inc (it uses that file's Tex, bindRange and exposureGrid). +// +// Every stage reads and writes NGX's own layouts (engine/rr/README.md): its parameter block is NVIDIA's argument +// block with arena byte offsets in place of pointers, the network's activations live in one plane arena at NGX's +// offsets, and the per-level surfaces have NGX's sizes. So a model is NGX's weight images plus two argument-block +// templates (enc0's and the output stage's constants), and this file fills every field that depends on the frame or +// on the geometry. +namespace { +constexpr unsigned RStages = 12; +enum : unsigned { REnc0, REnc1, REnc2, REnc3, REnc4, REnc5, RDec4, RDec3, RDec2, RDec1, RDec0, RHkpn }; +constexpr const char* RNames[RStages] = {"enc0", "enc1", "enc2", "enc3", "enc4", "enc5", + "dec4", "dec3", "dec2", "dec1", "dec0", "hkpn"}; +constexpr uint32_t RParamBytes[RStages] = {1032, 144, 144, 144, 144, 144, 144, 144, 144, 144, 152, 480}; +constexpr int RLevel[RStages] = {0, 1, 2, 3, 4, 5, 4, 3, 2, 1, 0, -1}; +// The plane arena: 70 * P0 bytes, P0 = the render size rounded up to 32 on both axes (W0 x H0). Every region starts +// at a multiple of P0 / 32 (NGX's layout, the same at 640x360 and 800x450): per stage, the arena offset of its +// input (argument 8), skip (24), output (48) and merged output (56) in units of P0 / 32, -1 where unused. +constexpr uint32_t RArenaUnits = 70 * 32; +constexpr int RArena[RStages][4] = { + {-1, -1, 0, 1024}, // enc0: 32-channel tokens (dec0's skip), the 2x2 merge for enc1 + {1024, -1, 1536, 2048}, // enc1 + {2048, -1, 1024, 1216}, // enc2 + {1216, -1, 1280, 1344}, // enc3 + {1344, -1, 1216, 1236}, // enc4 + {1236, -1, 1241, -1}, // enc5, the bottleneck + {1241, 1216, 1344, -1}, // dec4 + {1344, 1280, 1216, -1}, // dec3 + {1216, 1024, 2048, -1}, // dec2 + {2048, 1536, 1024, -1}, // dec1 + {1024, 0, 1536, -1}, // dec0: its output is the output stage's argument 256 + {-1, -1, -1, -1}}; +constexpr uint32_t RHkpnArena = 1536; +// NGX's window shift of a stage, packed (x | y << 16) at argument 80: it alternates every frame, the decoders opposite +// to the encoders; the bottleneck has none +uint32_t rShift(unsigned stage, uint32_t frame) { + if (stage == REnc5 || stage == RHkpn) return 0; + const bool even = ((frame & 1u) == 0) == (stage <= REnc5); + return even ? 0x00060002u : 0x00020006u; +} +void put32(std::vector& p, uint32_t offset, uint32_t value) { std::memcpy(p.data() + offset, &value, 4); } +void putF(std::vector& p, uint32_t offset, float value) { std::memcpy(p.data() + offset, &value, 4); } +void put64(std::vector& p, uint32_t offset, uint64_t value) { std::memcpy(p.data() + offset, &value, 8); } +using Mat4 = std::array; +Mat4 mul(const Mat4& a, const Mat4& b) { + Mat4 r{}; + for (int i = 0; i < 4; ++i) + for (int j = 0; j < 4; ++j) + for (int k = 0; k < 4; ++k) r[i * 4 + j] += a[i * 4 + k] * b[k * 4 + j]; + return r; +} +Mat4 inverse(const Mat4& m) { + // Gauss-Jordan with partial pivoting; a singular camera falls back to the identity + Mat4 a = m, r{1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + for (int c = 0; c < 4; ++c) { + int pivot = c; + for (int i = c + 1; i < 4; ++i) if (std::fabs(a[i * 4 + c]) > std::fabs(a[pivot * 4 + c])) pivot = i; + if (std::fabs(a[pivot * 4 + c]) < 1e-30) return Mat4{1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + for (int j = 0; j < 4; ++j) { std::swap(a[c * 4 + j], a[pivot * 4 + j]); std::swap(r[c * 4 + j], r[pivot * 4 + j]); } + const double d = a[c * 4 + c]; + for (int j = 0; j < 4; ++j) { a[c * 4 + j] /= d; r[c * 4 + j] /= d; } + for (int i = 0; i < 4; ++i) { + if (i == c) continue; + const double f = a[i * 4 + c]; + for (int j = 0; j < 4; ++j) { a[i * 4 + j] -= f * a[c * 4 + j]; r[i * 4 + j] -= f * r[c * 4 + j]; } + } + } + return r; +} +void putMat(std::vector& p, uint32_t offset, const Mat4& m) { + for (int i = 0; i < 16; ++i) putF(p, offset + 4 * i, static_cast(m[i])); +} +} // namespace + +struct RModel::Impl { + Device device; + uint32_t preset = 0; + // per stage, the shader's bindings after the three buffers (engine/rr/README.md): texture argument offsets, then + // surface argument offsets with their GLSL image format, both ascending + std::array, RStages> textures; + std::array>, RStages> surfaces; + std::array weights, uploads; + std::array stages; + std::array exposure; + std::vector enc0Template, hkpnTemplate; + bool uploaded = false; + Impl(VkPhysicalDevice physical, VkDevice dev) : device(physical, dev) {} +}; +RModel::RModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache) + : impl(std::make_unique(physical, device)) { + auto& m = *impl; + // rr.txt: "D4RR0001", "preset N", one "stage NAME SHADER.spv SUBGROUP WEIGHTS TEXTURES SURFACES" line per stage in + // RNames order (TEXTURES: comma-separated argument offsets or "-"; SURFACES: OFFSET:FORMAT,... or "-") and + // "template enc0|hkpn FILE" (the argument blocks whose constants the runtime keeps) + auto text = read(directory + "/rr.txt"); + std::istringstream lines(std::string(text.begin(), text.end())); + std::string line; + if (!std::getline(lines, line) || line.rfind("D4RR0001", 0) != 0) throw std::runtime_error("not a Ray Reconstruction model (rr.txt D4RR0001)"); + unsigned stage = 0; + const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; + auto split = [](const std::string& list) { + std::vector items; + if (list == "-") return items; + std::istringstream in(list); + for (std::string item; std::getline(in, item, ',');) items.push_back(item); + return items; + }; + while (std::getline(lines, line)) { + std::istringstream f(line); + std::string key; f >> key; + if (key == "preset") f >> m.preset; + else if (key == "stage") { + std::string name, shader, weightFile, textureList, surfaceList, option; uint32_t subgroup = 32; + f >> name >> shader >> subgroup >> weightFile >> textureList >> surfaceList >> option; + // "noopt": the stage is compiled without the driver's optimizations. RADV's ACO scheduler (Mesa 26.2, + // gfx1201) reorders the translated stages' cross-lane code so that enc1 lands at 20-43 dB against the + // reference; VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT gives 42-64 dB at about twice the time. + const VkPipelineCreateFlags flags = option == "noopt" ? VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT : 0; + if (stage >= RStages || name != RNames[stage] || (subgroup != 32 && subgroup != 64) || surfaceList.empty()) + throw std::runtime_error("rr.txt: stages out of order or invalid: " + line); + for (const auto& item : split(textureList)) m.textures[stage].push_back(static_cast(std::stoul(item))); + for (const auto& item : split(surfaceList)) { + const auto colon = item.find(':'); + if (colon == std::string::npos) throw std::runtime_error("rr.txt: surface without a format: " + line); + m.surfaces[stage].emplace_back(static_cast(std::stoul(item.substr(0, colon))), item.substr(colon + 1)); + } + std::vector types{b, b, b}; + types.insert(types.end(), m.textures[stage].size(), t); + types.insert(types.end(), m.surfaces[stage].size(), s); + m.stages[stage].create(device, directory + "/" + shader, types, subgroup, cache, flags); + auto blob = read(directory + "/" + weightFile); + const VkBufferUsageFlags usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + m.weights[stage].create(m.device, blob.size(), usage); + m.uploads[stage].create(m.device, blob.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT, true); + std::memcpy(m.uploads[stage].mapped, blob.data(), blob.size()); + ++stage; + } else if (key == "template") { + std::string which, file; f >> which >> file; + auto bytes = read(directory + "/" + file); + auto& target = which == "enc0" ? m.enc0Template : m.hkpnTemplate; + if (bytes.size() != RParamBytes[which == "enc0" ? REnc0 : RHkpn]) throw std::runtime_error("rr.txt: bad template " + file); + target.assign(bytes.begin(), bytes.end()); + } + } + if (stage != RStages || m.enc0Template.empty() || m.hkpnTemplate.empty() || (m.preset != 4 && m.preset != 5)) + throw std::runtime_error("rr.txt: incomplete Ray Reconstruction model"); + m.exposure[0].create(device, directory + "/exposure_rr0.spv", {b, t, b}, 64, cache); + m.exposure[1].create(device, directory + "/exposure_rr1.spv", {b, b, b, s, s}, 64, cache); +} +RModel::~RModel() = default; +bool RModel::uploadRecorded() const { return impl->uploaded; } +uint32_t RModel::preset() const { return impl->preset; } +void RModel::recordUpload(VkCommandBuffer cb) { + auto& m = *impl; + if (m.uploaded) throw std::runtime_error("model upload already recorded"); + for (unsigned i = 0; i < RStages; ++i) { + VkBufferCopy copy{0, 0, m.uploads[i].size}; + vkCmdCopyBuffer(cb, m.uploads[i].handle, m.weights[i].handle, 1, ©); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT); + m.uploaded = true; +} + +struct REngine::Impl { + RModel::Impl& model; + uint32_t rw, rh, ow, oh, w0, h0; + Buffer arena, expP, expSums, expState; + std::array params; + std::array, RStages> blocks; // this frame's parameter blocks + std::array s112, s120; // per level: the encoders' surfaces (decoders' skip textures 96, 104) + std::array s128; // the encoders' half-resolution surface, the next encoder's texture 88 + std::array d112, d120; // the decoders' surfaces (texture 88 of the next decoder; dec0's 112: output stage) + std::array q, r, h, e; // ping-pong: enc0's 304 / 224, dec0's 144 / enc0's 488, the output-size history, + // and e: this frame's and the previous frame's exposure (1x1) + Tex out; // the presented image (the output stage's surface 240) + Tex zero; // a zeroed 1x1 texture for the handles NGX leaves null (they read zero) + std::map spare; // per format, a 1x1 surface for stores NGX leaves unbound + VkDescriptorPool pool = VK_NULL_HANDLE; + VkDescriptorSet exp1Set = VK_NULL_HANDLE; + VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; + ExposureParams exposure{}; + struct Bindings { + RFrameImages images{}; + FrameOutput finalOutput{}; + VkDescriptorPool pool = VK_NULL_HANDLE; + std::array, RStages> stages{}; // per stage and frame parity + VkDescriptorSet exposure = VK_NULL_HANDLE; + }; + // what a stage's texture or surface at an argument offset is (the capture's wiring, engine/rr/README.md), for + // frame parity `par`; null view: NGX leaves it unbound + struct Resource { VkImageView view = VK_NULL_HANDLE; bool linear = false; bool game = false; }; + Resource resolve(unsigned stage, bool surface, uint32_t offset, uint32_t par, const RFrameImages& in, VkImageView output) const { + const int l = RLevel[stage]; + auto T = [](const Tex& t, bool lin = false) { return Resource{t.view, lin, false}; }; + auto G = [](VkImageView v, bool lin = false) { return Resource{v, lin, true}; }; + if (stage == REnc0) { + if (!surface) switch (offset) { + case 160: return G(in.color, true); case 168: return G(in.motion); case 184: return G(in.normals); + case 192: return G(in.roughness); case 200: return T(e[0]); case 208: return T(e[1]); + case 224: return T(q[par ^ 1], true); case 232: return G(in.diffuseAlbedo); case 240: return G(in.specularAlbedo); + // the optional inputs (capDX, 2026-10-12): unbound without them, like NGX + case 248: return in.specularHitDistance ? G(in.specularHitDistance) : Resource{}; + case 296: return in.disocclusionMask ? G(in.disocclusionMask) : Resource{}; + case 488: return T(r[par ^ 1], true); case 496: return T(h[par ^ 1], true); case 512: return G(in.depth); + } else switch (offset) { + case 112: return T(s112[0]); case 120: return T(s120[0]); case 128: return T(s128[0]); case 304: return T(q[par]); + } + } else if (stage <= REnc5) { + if (!surface && offset == 88) return T(s128[l - 1], true); + if (surface && offset == 112) return T(s112[l]); + if (surface && offset == 120) return T(s120[l]); + if (surface && offset == 128 && l < 5) return T(s128[l]); + } else if (stage <= RDec0) { + if (!surface && offset == 88) return T(stage == RDec4 ? s112[5] : d112[l + 1], true); + if (!surface && offset == 96) return T(s112[l]); + if (!surface && offset == 104) return T(s120[l]); + if (surface && offset == 112) return T(d112[l]); + if (surface && offset == 120) return T(d120[l]); + if (surface && offset == 144 && stage == RDec0) return T(r[par]); + } else { + if (!surface) switch (offset) { + case 176: return G(in.color, true); case 192: return G(in.depth); case 200: return T(e[0]); case 208: return T(e[1]); + case 216: return T(h[par ^ 1], true); case 280: return G(in.diffuseAlbedo); case 288: return G(in.specularAlbedo); + case 296: return T(q[par]); case 448: return T(d112[0]); + } else switch (offset) { + case 224: return T(h[par]); case 240: return Resource{output, false, false}; + } + } + return {}; + } + std::vector bindings; + const Bindings* bound = nullptr; + float prevJitter[2] = {0, 0}; + uint32_t frame = 0; + bool reset = true; + explicit Impl(RModel::Impl& m) : model(m) {} + VkFormat surfaceFormat(const std::string& f) const { + if (f == "rgba16f") return VK_FORMAT_R16G16B16A16_SFLOAT; + if (f == "rg16f") return VK_FORMAT_R16G16_SFLOAT; + if (f == "r16f") return VK_FORMAT_R16_SFLOAT; + if (f == "rgba8_snorm") return VK_FORMAT_R8G8B8A8_SNORM; + throw std::runtime_error("Ray Reconstruction stage surface format " + f + " is not implemented"); + } + ~Impl() { + VkDevice d = model.device.handle; + for (auto& b : bindings) vkDestroyDescriptorPool(d, b.pool, nullptr); + if (pool) vkDestroyDescriptorPool(d, pool, nullptr); + if (nearest) vkDestroySampler(d, nearest, nullptr); + if (linear) vkDestroySampler(d, linear, nullptr); + } + uint32_t lw(int level) const { return w0 >> level; } + uint32_t lh(int level) const { return h0 >> level; } + uint64_t arenaOffset(int units) const { return static_cast(units) * (static_cast(w0) * h0 / 32); } +}; + +REngine::REngine(RModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight) + : impl(std::make_unique(*model.impl)) { + auto& e = *impl; + const Device& d = e.model.device; + if (renderWidth < 64 || renderHeight < 64 || renderWidth > outputWidth || renderHeight > outputHeight || outputWidth > 16384 || outputHeight > 16384) + throw std::runtime_error("unsupported Ray Reconstruction dimensions"); + e.rw = renderWidth; e.rh = renderHeight; e.ow = outputWidth; e.oh = outputHeight; + e.w0 = (renderWidth + 31) & ~31u; e.h0 = (renderHeight + 31) & ~31u; + const VkBufferUsageFlags storage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; + e.arena.create(d, e.arenaOffset(RArenaUnits), storage); + for (unsigned i = 0; i < RStages; ++i) { + e.params[i].create(d, RParamBytes[i], storage); + e.blocks[i].assign(RParamBytes[i], 0); + } + e.expP.create(d, sizeof(ExposureParams), storage); + const uint32_t gx = exposureGrid(e.rw, e.ow), gy = exposureGrid(e.rh, e.oh); + e.exposure.blocksX = (gx + 31) / 32; + e.exposure.blocks = e.exposure.blocksX * ((gy + 31) / 32); + e.expSums.create(d, e.exposure.blocks * 4, storage); + e.expState.create(d, 8, storage); + const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT; + for (int l = 0; l < 6; ++l) { e.s112[l].create(d, e.lw(l), e.lh(l), half4); e.s120[l].create(d, e.lw(l), e.lh(l), half4); } + for (int l = 0; l < 5; ++l) { + e.s128[l].create(d, e.lw(l + 1), e.lh(l + 1), half4); + e.d112[l].create(d, e.lw(l), e.lh(l), half4); e.d120[l].create(d, e.lw(l), e.lh(l), half4); + } + for (int i = 0; i < 2; ++i) { + e.q[i].create(d, e.rw, e.rh, half4); + e.r[i].create(d, e.rw, e.rh, VK_FORMAT_R8G8B8A8_SNORM); + e.h[i].create(d, e.ow, e.oh, half4); + e.e[i].create(d, 1, 1, VK_FORMAT_R16_SFLOAT); + } + e.out.create(d, e.ow, e.oh, half4); + VkSamplerCreateInfo si{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; + si.addressModeU = si.addressModeV = si.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; si.maxLod = 0; + VKCHECK(vkCreateSampler(d.handle, &si, nullptr, &e.nearest)); + si.magFilter = si.minFilter = VK_FILTER_LINEAR; + VKCHECK(vkCreateSampler(d.handle, &si, nullptr, &e.linear)); + e.zero.create(d, 1, 1, half4); + for (unsigned stage = 0; stage < RStages; ++stage) + for (const auto& [offset, format] : e.model.surfaces[stage]) + if (!e.spare.count(format)) e.spare[format].create(d, 1, 1, e.surfaceFormat(format)); + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 4}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 4}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = 1; pi.poolSizeCount = 2; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(d.handle, &pi, nullptr, &e.pool)); + e.exp1Set = descriptorSet(d.handle, e.pool, e.model.exposure[1].descriptors); + bindBuffer(d.handle, e.exp1Set, 0, e.expP); + bindBuffer(d.handle, e.exp1Set, 1, e.expSums); + bindBuffer(d.handle, e.exp1Set, 2, e.expState); + bindImage(d.handle, e.exp1Set, 3, e.e[0].view); + bindImage(d.handle, e.exp1Set, 4, e.e[1].view); + // the parameter blocks' fixed fields + for (unsigned stage = REnc1; stage <= RDec0; ++stage) { + auto& p = e.blocks[stage]; + const int l = RLevel[stage]; + put32(p, 0, e.lw(l)); put32(p, 4, e.lh(l)); + const int offsets[4] = {8, 24, 48, 56}; + for (int k = 0; k < 4; ++k) + if (RArena[stage][k] >= 0) put64(p, offsets[k], e.arenaOffset(RArena[stage][k])); + putF(p, 136, 1.0f / static_cast(e.lw(l))); putF(p, 140, 1.0f / static_cast(e.lh(l))); + } +} +REngine::~REngine() = default; +VkImage REngine::outputImage() const { return impl->out.handle; } +std::vector REngine::debugImages() const { + const auto& e = *impl; + std::vector list; + const VkFormat half4 = VK_FORMAT_R16G16B16A16_SFLOAT; + for (int l = 0; l < 6; ++l) { + const std::string enc = "enc" + std::to_string(l); + list.push_back({enc + "-112", e.s112[l].handle, half4, e.lw(l), e.lh(l)}); + list.push_back({enc + "-120", e.s120[l].handle, half4, e.lw(l), e.lh(l)}); + if (l < 5) list.push_back({enc + "-128", e.s128[l].handle, half4, e.lw(l + 1), e.lh(l + 1)}); + if (l < 5) list.push_back({"dec" + std::to_string(l) + "-112", e.d112[l].handle, half4, e.lw(l), e.lh(l)}); + } + for (int i = 0; i < 2; ++i) { + list.push_back({"q" + std::to_string(i), e.q[i].handle, half4, e.rw, e.rh}); + list.push_back({"r" + std::to_string(i), e.r[i].handle, VK_FORMAT_R8G8B8A8_SNORM, e.rw, e.rh}); + list.push_back({"h" + std::to_string(i), e.h[i].handle, half4, e.ow, e.oh}); + list.push_back({"e" + std::to_string(i), e.e[i].handle, VK_FORMAT_R16_SFLOAT, 1, 1}); + } + list.push_back({"out", e.out.handle, half4, e.ow, e.oh}); + return list; +} + +void REngine::recordInitialize(VkCommandBuffer cb) { + auto& e = *impl; + std::vector images; + for (auto* group : {&e.s112, &e.s120}) for (auto& t : *group) images.push_back(t.handle); + for (auto* group : {&e.s128, &e.d112, &e.d120}) for (auto& t : *group) images.push_back(t.handle); + for (auto* group : {&e.q, &e.r, &e.h, &e.e}) for (auto& t : *group) images.push_back(t.handle); + images.push_back(e.out.handle); + images.push_back(e.zero.handle); + for (auto& [format, t] : e.spare) images.push_back(t.handle); + VkClearColorValue zero{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (VkImage image : images) { + imageBarrier(cb, image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, + VK_PIPELINE_STAGE_TRANSFER_BIT, 0, VK_ACCESS_TRANSFER_WRITE_BIT); + vkCmdClearColorImage(cb, image, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + } + vkCmdFillBuffer(cb, e.arena.handle, 0, VK_WHOLE_SIZE, 0); + vkCmdFillBuffer(cb, e.expState.handle, 0, VK_WHOLE_SIZE, 0); + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); + e.frame = 0; e.reset = true; +} + +void REngine::setFrame(const RFrameImages& images, const RFrameSettings& settings, FrameOutput finalOutput) { + auto& e = *impl; + const Device& d = e.model.device; + if (!images.color || !images.motion || !images.depth || !images.normals || !images.roughness || !images.diffuseAlbedo || !images.specularAlbedo) + throw std::runtime_error("Ray Reconstruction needs colour, motion, depth and the four guides"); + if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT) + throw std::runtime_error("Ray Reconstruction writes RGBA16F output only"); + const bool reset = settings.reset || e.frame == 0; + const uint32_t f = e.frame; + const float rw = static_cast(e.rw), rh = static_cast(e.rh), ow = static_cast(e.ow), oh = static_cast(e.oh); + const float prevJitter[2] = {reset ? 0.0f : e.prevJitter[0], reset ? 0.0f : e.prevJitter[1]}; + // the network layers: only the window shift changes + for (unsigned stage = REnc1; stage <= RDec0; ++stage) put32(e.blocks[stage], 80, rShift(stage, f)); + // enc0 (the argument block's constants from the model; offsets as observed on harness captures, 2026-10-11) + auto& p = e.blocks[REnc0]; + p = e.model.enc0Template; + put32(p, 0, e.w0); put32(p, 4, e.h0); + put64(p, 48, e.arenaOffset(RArena[REnc0][2])); put64(p, 56, e.arenaOffset(RArena[REnc0][3])); + put32(p, 80, rShift(REnc0, f)); + putF(p, 136, 1.0f / static_cast(e.w0)); putF(p, 140, 1.0f / static_cast(e.h0)); + put32(p, 144, e.rw); put32(p, 148, e.rh); put32(p, 152, e.w0); put32(p, 156, e.h0); + putF(p, 312, prevJitter[0]); putF(p, 316, prevJitter[1]); putF(p, 320, settings.jitter[0]); putF(p, 324, settings.jitter[1]); + putF(p, 328, rw / ow); putF(p, 332, rh / oh); putF(p, 336, ow / rw); putF(p, 340, oh / rh); + putF(p, 344, settings.motionScale[0] * (ow / rw)); putF(p, 348, settings.motionScale[1] * (oh / rh)); + // byte 365 (and the output stage's byte 440): the reset flag. Temporal harness captures (capDT, capDfr, capDst, + // 2026-10-12) show it follows Reset alone, not the jitter: set with reset and jitter, clear without both. + put32(p, 364, reset ? 256u : 0u); + // byte 366: DLSS.Roughness.Mode packed. enc0 then takes the roughness from the normals' w and never reads the + // roughness texture (argument 192, bound to the normals as well) + p[366] = settings.roughnessPacked ? 1u : 0u; + // bytes 536 / 539: DLSSD.SpecularHitDistance (texture 248) / DLSS.DisocclusionMask (texture 296) are bound. The hit + // distance also sets the output stage's byte 160 (harness captures capDXhit, capDXmask, capDX, 2026-10-12) + p[536] = images.specularHitDistance ? 1u : 0u; + p[539] = images.disocclusionMask ? 1u : 0u; + put32(p, 432, e.rw - 1); put32(p, 436, e.rh - 1); put32(p, 472, e.ow); put32(p, 476, e.oh); + put32(p, 528, settings.depthInverted ? 1u : 0u); + // the camera: the view's rotation (its translation row dropped) times the projection, its inverse, and both + // composed with the render viewport (NDC <-> render pixels, y down) + Mat4 view, clip; + for (int i = 0; i < 16; ++i) { view[i] = settings.worldToView[i]; clip[i] = settings.viewToClip[i]; } + view[12] = view[13] = view[14] = 0.0; + const Mat4 rc = mul(view, clip), inv = inverse(rc); + const Mat4 toNdc{2.0 / e.rw, 0, 0, 0, 0, -2.0 / e.rh, 0, 0, 0, 0, 1, 0, -1, 1, 0, 1}; + const Mat4 toPixels{e.rw / 2.0, 0, 0, 0, 0, -(e.rh / 2.0), 0, 0, 0, 0, 1, 0, e.rw / 2.0, e.rh / 2.0, 0, 1}; + putMat(p, 596, rc); putMat(p, 660, clip); putMat(p, 724, clip); putMat(p, 788, inv); + putMat(p, 852, mul(toNdc, inv)); putMat(p, 916, mul(rc, toPixels)); + putF(p, 756, settings.jitter[0] * (2.0f / rw)); putF(p, 760, -settings.jitter[1] * (2.0f / rh)); + put32(p, 984, f); + put32(p, 1024, e.rw); put32(p, 1028, e.rh); + // the output stage + auto& o = e.blocks[RHkpn]; + o = e.model.hkpnTemplate; + put32(o, 0, e.ow); put32(o, 4, e.oh); put32(o, 8, e.ow); put32(o, 12, e.oh); put32(o, 16, e.rw); put32(o, 20, e.rh); + putF(o, 24, rw / ow); putF(o, 28, rh / oh); putF(o, 32, ow / rw); putF(o, 36, oh / rh); + putF(o, 40, settings.jitter[0]); putF(o, 44, settings.jitter[1]); + put32(o, 56, e.w0); put32(o, 60, e.h0); put32(o, 140, e.rw - 1); put32(o, 144, e.rh - 1); + putF(o, 152, settings.frameTimeMs); + o[160] = images.specularHitDistance ? 1u : 0u; + put64(o, 256, e.arenaOffset(RHkpnArena)); + putF(o, 336, settings.motionScale[0] * (ow / rw)); putF(o, 340, settings.motionScale[1] * (oh / rh)); + put32(o, 440, (reset ? 1u : 0u) | (settings.depthInverted ? 0x10000u : 0u)); + // the exposure: NGX's MODE 1 measurement (engine/exposure_k.comp) on 1 / pre-exposure scaled colour + const uint32_t gx = exposureGrid(e.rw, e.ow), gy = exposureGrid(e.rh, e.oh); + e.exposure.origin[0] = e.exposure.origin[1] = 0; + e.exposure.size[0] = static_cast(e.rw); e.exposure.size[1] = static_cast(e.rh); + e.exposure.sampleSize[0] = static_cast(gx); e.exposure.sampleSize[1] = static_cast(gy); + e.exposure.invCount = 1.0f / static_cast(gx * gy); + e.exposure.scale = 1.0f / (settings.preExposure > 0.0f ? settings.preExposure : 1.0f); + // cuda_dldn_engine_auto_exposure_copy_kernel divides a constant of its own (0d3FCE23E8E0000000), not its argument 0 + // (0.2195...), by max(mean luma, 0.005): the capture's 0.919 = 0.23547 / 0.2562 + e.exposure.key = 0.23547087609767914f; + // NGX's reset keeps the exposure state, the recurrent feature and the output history and only clears the previous + // frame's enc0 feature q (a harness run resetting every frame: enc0's textures 208, 488 and 496 carry the previous + // frame from frame 2 on, 224 stays zero; capture capE3, 2026-10-11). The first frame starts from zero entirely. + e.exposure.flags = 1u | (e.frame == 0 ? 2u : 0u); + e.reset = reset; + // descriptors for the caller's images, cached per combination: every stage, both frame parities, bound by the + // argument offsets its shader declares (rr.txt) + // texture and surface handles of the two stages with optional inputs: 1 where something is bound, 0 where NGX + // leaves the argument null. The kernels branch on that (enc0 takes its specular hit distance path only when + // argument 248's handle is not null; capDX, 2026-10-12). + for (unsigned stage : {static_cast(REnc0), static_cast(RHkpn)}) { + for (uint32_t offset : e.model.textures[stage]) + put64(e.blocks[stage], offset, e.resolve(stage, false, offset, 0, images, e.out.view).view ? 1u : 0u); + for (const auto& [offset, format] : e.model.surfaces[stage]) + put64(e.blocks[stage], offset, e.resolve(stage, true, offset, 0, images, e.out.view).view ? 1u : 0u); + } + auto same = [&](const Impl::Bindings& b) { + return !std::memcmp(&b.images, &images, sizeof(images)) && b.finalOutput.view == finalOutput.view; + }; + auto found = std::find_if(e.bindings.begin(), e.bindings.end(), same); + if (found == e.bindings.end()) { + if (e.bindings.size() >= 16) throw std::runtime_error("too many distinct Ray Reconstruction input view sets"); + Impl::Bindings b; + b.images = images; b.finalOutput = finalOutput; + uint32_t textures = 0, surfaces = 0; + for (unsigned stage = 0; stage < RStages; ++stage) { + textures += static_cast(e.model.textures[stage].size()); + surfaces += static_cast(e.model.surfaces[stage].size()); + } + VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 2 * 3 * RStages + 3}, + {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 2 * textures + 1}, + {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 2 * surfaces}}; + VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; + pi.maxSets = 2 * RStages + 1; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; + VKCHECK(vkCreateDescriptorPool(d.handle, &pi, nullptr, &b.pool)); + const VkImageView output = finalOutput.view ? finalOutput.view : e.out.view; + for (unsigned stage = 0; stage < RStages; ++stage) { + for (uint32_t par = 0; par < 2; ++par) { + VkDescriptorSet set = b.stages[stage][par] = descriptorSet(d.handle, b.pool, e.model.stages[stage].descriptors); + bindBuffer(d.handle, set, 0, e.params[stage]); + bindBuffer(d.handle, set, 1, e.model.weights[stage]); + bindBuffer(d.handle, set, 2, e.arena); + uint32_t binding = 3; + for (uint32_t offset : e.model.textures[stage]) { + const auto res = e.resolve(stage, false, offset, par, images, output); + bindImage(d.handle, set, binding++, res.view ? res.view : e.zero.view, res.linear ? e.linear : e.nearest, + res.game ? images.layout : VK_IMAGE_LAYOUT_GENERAL); + } + for (const auto& [offset, format] : e.model.surfaces[stage]) { + const auto res = e.resolve(stage, true, offset, par, images, output); + bindImage(d.handle, set, binding++, res.view ? res.view : e.spare.at(format).view); + } + } + } + b.exposure = descriptorSet(d.handle, b.pool, e.model.exposure[0].descriptors); + bindBuffer(d.handle, b.exposure, 0, e.expP); + bindImage(d.handle, b.exposure, 1, images.color, e.nearest, images.layout); + bindBuffer(d.handle, b.exposure, 2, e.expSums); + e.bindings.push_back(b); + found = e.bindings.end() - 1; + } + e.bound = &*found; + // debugging: D4R_ENGINE_RR_DUMP=DIR writes every frame's parameter blocks as DIR/F-STAGE.bin (F counts from 1) + if (const char* dump = std::getenv("D4R_ENGINE_RR_DUMP")) { + for (unsigned stage = 0; stage < RStages; ++stage) + std::ofstream(std::string(dump) + "/" + std::to_string(f + 1) + "-" + RNames[stage] + ".bin", std::ios::binary) + .write(reinterpret_cast(e.blocks[stage].data()), static_cast(e.blocks[stage].size())); + std::ofstream(std::string(dump) + "/" + std::to_string(f + 1) + "-exposure.bin", std::ios::binary) + .write(reinterpret_cast(&e.exposure), sizeof(e.exposure)); + } + e.prevJitter[0] = settings.jitter[0]; e.prevJitter[1] = settings.jitter[1]; +} + +void REngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { + auto& e = *impl; + if (!e.bound) throw std::runtime_error("REngine::setFrame was not called"); + const auto& b = *e.bound; + const uint32_t f = e.frame, cur = f & 1u; + auto stamp = [&](uint32_t n) { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, stamps, first + n); }; + auto step = [&] { + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + }; + stamp(0); + // the frame's parameter blocks, then (on a reset) a clean temporal state + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_TRANSFER_WRITE_BIT); + for (unsigned i = 0; i < RStages; ++i) vkCmdUpdateBuffer(cb, e.params[i].handle, 0, e.blocks[i].size(), e.blocks[i].data()); + vkCmdUpdateBuffer(cb, e.expP.handle, 0, sizeof(ExposureParams), &e.exposure); + if (e.reset) { + VkClearColorValue zero{}; + VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; + for (auto& t : e.q) vkCmdClearColorImage(cb, t.handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); + } + memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_ACCESS_TRANSFER_WRITE_BIT, + VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); + // exposure + dispatch(cb, e.model.exposure[0], b.exposure, e.exposure.blocksX, e.exposure.blocks / e.exposure.blocksX); + step(); + dispatch(cb, e.model.exposure[1], e.exp1Set, 1, 1); + step(); + stamp(1); + // the network, in NGX's order; each stage's grid covers its level plus that frame's window shift + for (unsigned stage = REnc0; stage <= RDec0; ++stage) { + const int l = RLevel[stage]; + const uint32_t shift = rShift(stage, f), sx = shift & 0xffffu, sy = shift >> 16; + const VkDescriptorSet set = b.stages[stage][cur]; + dispatch(cb, e.model.stages[stage], set, (e.lw(l) + sx + 7) / 8, (e.lh(l) + sy + 7) / 8); + step(); + } + stamp(2); + dispatch(cb, e.model.stages[RHkpn], b.stages[RHkpn][cur], (e.ow + 15) / 16, (e.oh + 15) / 16); + memoryBarrier(cb, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT, + VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_READ_BIT); + stamp(3); + ++e.frame; +} diff --git a/engine/vk_dispatch.cpp b/engine/vk_dispatch.cpp index 3fbd116..b404299 100644 --- a/engine/vk_dispatch.cpp +++ b/engine/vk_dispatch.cpp @@ -22,6 +22,7 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD ok &= (vkTable.CmdCopyBufferToImage = reinterpret_cast(deviceProc(device, "vkCmdCopyBufferToImage"))) != nullptr; ok &= (vkTable.CmdCopyImageToBuffer = reinterpret_cast(deviceProc(device, "vkCmdCopyImageToBuffer"))) != nullptr; ok &= (vkTable.CmdDispatch = reinterpret_cast(deviceProc(device, "vkCmdDispatch"))) != nullptr; + ok &= (vkTable.CmdFillBuffer = reinterpret_cast(deviceProc(device, "vkCmdFillBuffer"))) != nullptr; ok &= (vkTable.CmdPipelineBarrier = reinterpret_cast(deviceProc(device, "vkCmdPipelineBarrier"))) != nullptr; ok &= (vkTable.CmdUpdateBuffer = reinterpret_cast(deviceProc(device, "vkCmdUpdateBuffer"))) != nullptr; ok &= (vkTable.CmdWriteTimestamp = reinterpret_cast(deviceProc(device, "vkCmdWriteTimestamp"))) != nullptr; diff --git a/engine/vk_dispatch.h b/engine/vk_dispatch.h index baef618..bf793c6 100644 --- a/engine/vk_dispatch.h +++ b/engine/vk_dispatch.h @@ -20,6 +20,7 @@ struct VkTable { PFN_vkCmdCopyBufferToImage CmdCopyBufferToImage = nullptr; PFN_vkCmdCopyImageToBuffer CmdCopyImageToBuffer = nullptr; PFN_vkCmdDispatch CmdDispatch = nullptr; + PFN_vkCmdFillBuffer CmdFillBuffer = nullptr; PFN_vkCmdPipelineBarrier CmdPipelineBarrier = nullptr; PFN_vkCmdUpdateBuffer CmdUpdateBuffer = nullptr; PFN_vkCmdWriteTimestamp CmdWriteTimestamp = nullptr; @@ -67,6 +68,7 @@ bool loadVulkan(PFN_vkGetInstanceProcAddr instanceProc, VkInstance instance, VkD #define vkCmdCopyBufferToImage ::d4r::vkTable.CmdCopyBufferToImage #define vkCmdCopyImageToBuffer ::d4r::vkTable.CmdCopyImageToBuffer #define vkCmdDispatch ::d4r::vkTable.CmdDispatch +#define vkCmdFillBuffer ::d4r::vkTable.CmdFillBuffer #define vkCmdPipelineBarrier ::d4r::vkTable.CmdPipelineBarrier #define vkCmdUpdateBuffer ::d4r::vkTable.CmdUpdateBuffer #define vkCmdWriteTimestamp ::d4r::vkTable.CmdWriteTimestamp diff --git a/kernels/rr/rrswin_enc0.hip b/kernels/rr/rrswin_enc0.hip index 4a34e82..7683221 100644 --- a/kernels/rr/rrswin_enc0.hip +++ b/kernels/rr/rrswin_enc0.hip @@ -39,9 +39,10 @@ // helpers and ZLUDA's own device functions are available. Compiled with -nogpuinc, // so only plain clang builtins and attributes are used. #include "../tex/tex_common.h" -#ifdef D4R_RRSWIN_NATIVE_GEMM +// The staged-image engine: the body's sixteen plain GEMM phases read staged images and prepared +// slots through it, so the header is needed unconditionally (the phase-8 pilot below keeps its +// own D4R_RRSWIN_NATIVE_GEMM guard). #include "rrswin_engine.h" -#endif // --------------------------------------------------------------------------- // Parameter block. Fields are addressed by their byte offset in the PTX, so the @@ -65,7 +66,7 @@ enum : int #define DEVF __attribute__((device)) static inline typedef float lds_f32 __attribute__((address_space(3))); -typedef uint32_t lds_u32 __attribute__((address_space(3))); +using rrswin::lds_u32; // the engine header's AS3 word (the body defines none of its own) typedef uint8_t lds_u8 __attribute__((address_space(3))); #ifdef D4R_RRSWIN_NATIVE_GEMM @@ -75,8 +76,8 @@ DEVF void gemm_native_unbiased(const uint32_t* as, const uint32_t* bs, int nmi, uint32_t scratch, lds_u32* out) { lds_u8* const shared = (lds_u8*)(uintptr_t)scratch; - uint8_t* const image = (uint8_t*)shared; - uint8_t* const weights = (uint8_t*)(shared + 512); + lds_u8* const image = shared; + lds_u8* const weights = shared + 512; __attribute__((address_space(3))) half_t* const result = (__attribute__((address_space(3))) half_t*)(shared + 1536); const uint32_t lane = lane_id(); @@ -102,7 +103,7 @@ DEVF void gemm_native_unbiased(const uint32_t* as, const uint32_t* bs, for (int mt = 0; mt < 2; ++mt) { f8v c = {0, 0, 0, 0, 0, 0, 0, 0}; - c = rrswin::swin_gemm_tile(c, (const kslot*)weights, 16 * mt, 32, image, 32, 0); + c = rrswin::swin_gemm_tile(c, weights, 16 * mt, 32, image, 32, 0); #pragma unroll for (int i = 0; i < 8; ++i) result[32 * (lane & 15) + 16 * mt + i + 8 * (lane >> 4)] = (half_t)c[i]; @@ -220,17 +221,8 @@ DEVF uint32_t pair_of(half_t a, half_t b) | ((uint32_t)__builtin_bit_cast(uint16_t, b) << 16); } -// ZLUDA's relu_keep_nan_halves: a negative (or -0) half becomes +0, NaN is kept. -DEVF uint32_t relu_pair(uint32_t p) -{ - if ((p & 0x8000u) != 0 && (p & 0x7fffu) <= 0x7c00u) { - p &= 0xffff0000u; - } - if ((p & 0x80000000u) != 0 && ((p >> 16) & 0x7fffu) <= 0x7c00u) { - p &= 0x0000ffffu; - } - return p; -} +// ZLUDA's relu_keep_nan_halves (the phase-2 staging's mode 1) lives in the engine header +// as rrswin::swin_relu_pair; this body has no local copy. // --------------------------------------------------------------------------- // One GEMM: 96 rows (six m-tiles of 16) x 32 columns, K=32, f16 output. @@ -783,6 +775,225 @@ DEVF void cres_regs(const uint32_t d[4][2][4], uint32_t c[64]) } } +// --------------------------------------------------------------------------- +// The staged-image engine's plumbing: the four contract-§4 primitives live in +// rrswin_engine.h (rrswin::lds_u32, swin_slots_from_wimg, swin_dcell_exchange_write/read, +// swin_stage_operand_p with modes 0/1/2); what remains here is the body's own drivers. + +// The engine tile's f16 pair compaction -- the one f16 rounding of a phase -- and its C +// seeds. Per-column f16 tables widen to f32; a residual enters as the producer's own tile +// pairs (contract §2's identity: same tile, same lane, pair i>>1 half i&1), with a +// per-column table added in f16 first where one exists (phases 7 and 14). +DEVF void enc0_pack(uint32_t pairs[4], const f8v& c) +{ +#pragma unroll + for (int p = 0; p < 4; ++p) + pairs[p] = pair_of((half_t)c[2 * p], (half_t)c[2 * p + 1]); +} + +// The lane's 16 f16 of a per-column blob = 8 u32: features 16*mt + 8h + 2j and +1 sit at +// bytes 32*mt + 16h + 4j (column c at byte 2c, contract §2). +DEVF void enc0_load_cols(uint32_t wb[2][4], const uint8_t* w, int base) +{ + const uint32_t h = lane_id() >> 4; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int j = 0; j < 4; ++j) + wb[mt][j] = *(const uint32_t*)(w + base + 32 * mt + 16 * h + 4 * j); +} + +DEVF void enc0_seed_bias(f8v& c, int mt, const uint32_t wb[2][4]) +{ +#pragma unroll + for (int i = 0; i < 8; ++i) + c[i] = h2f((i & 1) ? (wb[mt][i >> 1] >> 16) : wb[mt][i >> 1]); +} + +DEVF void enc0_seed_res(f8v& c, int mt, int nt, const uint32_t* pp, int n_nt) +{ +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t r = pp[(mt * n_nt + nt) * 4 + (i >> 1)]; + c[i] = (float)((i & 1) ? half_hi(r) : half_lo(r)); + } +} + +DEVF void enc0_seed_res_tab(f8v& c, int mt, int nt, const uint32_t* pp, int n_nt, + const uint32_t* wt) +{ +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t r = pp[(mt * n_nt + nt) * 4 + (i >> 1)]; + const uint32_t tb = wt[(mt << 2) + (i >> 1)]; + c[i] = (float)((i & 1) ? half_hi(r) + half_hi(tb) : half_lo(r) + half_lo(tb)); + } +} + +// One phase's dcell exchange: its 2*n_nt engine tiles do not fit the 2048-byte dead image +// (128 u32 per tile block), so the exchange runs in 2x2-tile chunks of 2048 bytes, chunk c +// refilling d[2c] and d[2c+1]. pairs is the lane's [2][n_nt][4] pair array, d its +// [n_nt][2][4] accumulator view. +DEVF void enc0_dcell_exchange(const uint32_t* pairs, int n_nt, lds_u32* ex, uint32_t* d) +{ + for (int c = 0; c < n_nt / 2; ++c) { + uint32_t xp[16]; +#pragma unroll + for (int mt = 0; mt < 2; ++mt) +#pragma unroll + for (int l = 0; l < 2; ++l) +#pragma unroll + for (int p = 0; p < 4; ++p) + xp[(mt * 2 + l) * 4 + p] = pairs[(mt * n_nt + 2 * c + l) * 4 + p]; + rrswin::swin_dcell_exchange_write(ex, xp, 2, 2); + __builtin_amdgcn_wave_barrier(); + rrswin::swin_dcell_exchange_read(ex, d + 16 * c, 2, 4, 2); + __builtin_amdgcn_wave_barrier(); + } +} + +// Phase 1: the staged block's fragments stage into the image per 48-token half (the six +// m-tiles restaged at local rows 0..47), the C seed is the per-column bias at w+5120, and +// the result stays in registers as pairs. +DEVF void enc0_phase1(const uint32_t* as1, const uint8_t* w, lds_u8* slots, lds_u8* img, + uint32_t pairs1[2][6][4]) +{ + uint32_t wb[2][4]; + rrswin::swin_slots_from_wimg(slots, w, 4096, 32); + rrswin::swin_stage_fragment(img, 32, as1, 3); + __builtin_amdgcn_wave_barrier(); + enc0_load_cols(wb, w, 5120); + for (int half = 0; half < 2; ++half) { + if (half != 0) { + rrswin::swin_stage_fragment(img, 32, as1 + 12, 3); + __builtin_amdgcn_wave_barrier(); + } +#pragma unroll + for (int nt = 3 * half; nt < 3 * half + 3; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + f8v c; + enc0_seed_bias(c, mt, wb); + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, + 16 * (nt - 3 * half)); + enc0_pack(pairs1[mt][nt], c); + } + __builtin_amdgcn_wave_barrier(); + } +} + +// Phase 2: phase 1's pairs staged with the relu (mode 1) as the A image, the per-column +// bias at w+6208 as C; the caller keeps pairs2 for phase 7's C and exchanges the dcell view. +DEVF void enc0_phase2(const uint32_t pairs1[2][6][4], const uint8_t* w, lds_u8* slots, + lds_u8* img, uint32_t pairs2[2][6][4]) +{ + uint32_t wb[2][4]; + rrswin::swin_slots_from_wimg(slots, w, 5184, 32); + __builtin_amdgcn_wave_barrier(); + enc0_load_cols(wb, w, 6208); + for (int half = 0; half < 2; ++half) { +#pragma unroll + for (int nt = 3 * half; nt < 3 * half + 3; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + rrswin::swin_stage_operand_p(img, 32, 16 * (nt - 3 * half), mt, pairs1[mt][nt], 1); + __builtin_amdgcn_wave_barrier(); +#pragma unroll + for (int nt = 3 * half; nt < 3 * half + 3; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + f8v c; + enc0_seed_bias(c, mt, wb); + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, + 16 * (nt - 3 * half)); + enc0_pack(pairs2[mt][nt], c); + } + __builtin_amdgcn_wave_barrier(); + } +} + +// Phases 3, 4, 8 and 9: the fused projections of the normed tile -- a 96-token phase whose +// A image is the as3 fragment pack staged per half and whose D is consumed only as the +// NVIDIA d[mi][row][nt] view (the score and P·V GEMMs walk it scalar). +DEVF void enc0_proj96(const uint32_t* as3, const uint8_t* w, int base, lds_u8* slots, + lds_u8* img, lds_u32* ex, uint32_t* d) +{ + uint32_t pr[2][6][4]; + rrswin::swin_slots_from_wimg(slots, w, base, 32); + rrswin::swin_stage_fragment(img, 32, as3, 3); + __builtin_amdgcn_wave_barrier(); + for (int half = 0; half < 2; ++half) { + if (half != 0) { + rrswin::swin_stage_fragment(img, 32, as3 + 12, 3); + __builtin_amdgcn_wave_barrier(); + } +#pragma unroll + for (int nt = 3 * half; nt < 3 * half + 3; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + f8v c = {0, 0, 0, 0, 0, 0, 0, 0}; + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, + 16 * (nt - 3 * half)); + enc0_pack(pr[mt][nt], c); + } + __builtin_amdgcn_wave_barrier(); + } + enc0_dcell_exchange(&pr[0][0][0], 6, ex, d); +} + +// The four MLP branches' first GEMMs (13, 15, 17, 19): the normed tile as the A image, the +// per-column bias as C, the result kept as pairs for the mode-2 staging and the residual. +DEVF void enc0_branch(const uint32_t* asN, const uint8_t* w, int base, int bias, + lds_u8* slots, lds_u8* img, uint32_t pairs[2][4][4]) +{ + uint32_t wb[2][4]; + rrswin::swin_slots_from_wimg(slots, w, base, 32); + rrswin::swin_stage_fragment(img, 32, asN, 4); + __builtin_amdgcn_wave_barrier(); + enc0_load_cols(wb, w, bias); +#pragma unroll + for (int nt = 0; nt < 4; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + f8v c; + enc0_seed_bias(c, mt, wb); + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, 16 * nt); + enc0_pack(pairs[mt][nt], c); + } + __builtin_amdgcn_wave_barrier(); +} + +// The four branches' second GEMMs (14, 16, 18, 20): the previous branch's pairs staged with +// the clamped cubic transform (mode 2) as the A image (pp), the running residual as C (pc) +// -- phase 12's pairs plus the per-column table for phase 14 (wt), the producer pairs +// without one for 16, 18 and 20 (wt null). +DEVF void enc0_branch_act(const uint32_t pp[2][4][4], const uint32_t* pc, const uint32_t* wt, + const uint8_t* w, int base, lds_u8* slots, lds_u8* img, + uint32_t pairs[2][4][4]) +{ + rrswin::swin_slots_from_wimg(slots, w, base, 32); +#pragma unroll + for (int nt = 0; nt < 4; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) + rrswin::swin_stage_operand_p(img, 32, 16 * nt, mt, pp[mt][nt], 2); + __builtin_amdgcn_wave_barrier(); +#pragma unroll + for (int nt = 0; nt < 4; ++nt) +#pragma unroll + for (int mt = 0; mt < 2; ++mt) { + f8v c; + if (wt != nullptr) { + enc0_seed_res_tab(c, mt, nt, pc, 4, wt); + } else { + enc0_seed_res(c, mt, nt, pc, 4); + } + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, 16 * nt); + enc0_pack(pairs[mt][nt], c); + } + __builtin_amdgcn_wave_barrier(); +} + // --------------------------------------------------------------------------- // Ground truth for a stage comes from kernels/tex/make_ptx.py's debug recipe: // with D4R_RRSWIN_DEBUG=p1 (p2, p3, p4, or a) the same splice stores that stage's @@ -1003,6 +1214,10 @@ DEVF void band21(const uint32_t a[16], const uint8_t* w, const uint8_t* cbias, i // twelve as six u32. DEVF void surface_scatter21(const uint8_t* prm, const uint32_t d[8][2], int lane) { +#if defined(D4R_RRSWIN_SC1) && D4R_RRSWIN_SC1 == 1 + (void)prm; (void)d; (void)lane; + return; // regression probe: empty call +#endif const int r1 = 8 * (int)__builtin_amdgcn_workgroup_id_x() - (int)*(const uint16_t*)(prm + 80); const int r6 = 8 * (int)(__builtin_amdgcn_workgroup_id_y() @@ -1017,35 +1232,105 @@ DEVF void surface_scatter21(const uint8_t* prm, const uint32_t d[8][2], int lane const int y2 = y1 + 1; const int x1 = (r1 >> 1) + ((lane >> 3) & 3); const int x2 = (r1 >> 1) + ((g >> 1) & 3); - uint8_t* const dst = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + 56); - uint16_t q[16]; + uint8_t* const dst_raw = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + 56); + // quadrant bisect: 1 = the four prm reads and the target load, 2 = the q fill, + // 3 = the index arithmetic and its two gates, 4 = the store groups +#ifndef D4R_RRSWIN_Q +#define D4R_RRSWIN_Q 4 +#endif +#if D4R_RRSWIN_Q == 1 + (void)d; (void)lane; + return; +#endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 19 + uint8_t* const dst = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + 48); // the dumped arena +#else + uint8_t* const dst = dst_raw; +#endif +#if D4R_RRSWIN_DEBUG_STAGE == 19 + { + // pointer capture: prm's own words, the +56 target and the geometry, in the + // debug-store convention (stage_regs4 layout filler is irrelevant here) + uint32_t cap[36]; + for (int i = 0; i < 32; ++i) { + cap[i] = *(const uint32_t*)(prm + 4 * i); + } + cap[32] = (uint32_t)(uintptr_t)(prm + 56); + *(uint64_t*)(cap + 32) = *(const uint64_t*)(prm + 56); + debug_store(prm, cap, 34); + } +#endif + // The sixteen payloads as the eight u32 words the four store groups write. + // pair16(q[2n], q[2n+1]) with q[k] = (uint16_t)e4m3x2(d[n][row]) is exactly + // e4m3x2(d[n][0]) | (e4m3x2(d[n][1]) << 16): + // e4m3x2 returns two e4m3 bytes, each in [0,255] with a zero upper half, so the + // uint16_t truncation in q[] is a no-op and the packing is bit for bit the same. + // Forming the words directly (rather than through a uint16_t q[16] that -O2/-O3 + // widens into a sixteen-lane <16 x i16> built by a shuffle whose mask indexes + // past its operands -- the last two lanes come back `poison`) keeps each payload in + // one u32 register and restores the -O0/-O1 lowering of the stores. + uint32_t w[8]; #pragma unroll for (int nt = 0; nt < 8; ++nt) { - for (int row = 0; row < 2; ++row) { - q[2 * nt + row] = (uint16_t)e4m3x2(d[nt][row]); - } + w[nt] = e4m3x2(d[nt][0]) | (e4m3x2(d[nt][1]) << 16); } +#if D4R_RRSWIN_Q == 2 + return; +#endif const int b1 = sx * y1 + 8 * x1; const int b2 = sx * y2 + 8 * x2; const int st = sx * ey; const bool ok1 = (unsigned)y1 < (unsigned)ey && (unsigned)x1 < (unsigned)ex; const bool ok2 = (unsigned)y2 < (unsigned)ey && (unsigned)x2 < (unsigned)ex; +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 18 + return; // bisect: the scatter's prm/geometry reads only, no stores +#endif +#if D4R_RRSWIN_Q == 3 + return; +#endif +#if D4R_RRSWIN_Q == 9 + // same as Q5, but the target pointer is re-read from prm inside the guarded block if (ok1) { - *(uint32_t*)(dst + 4 * (uint32_t)(b1 + t)) = pair16(q[0], q[1]); - *(uint32_t*)(dst + 4 * (uint32_t)(b1 + (t | 4))) = pair16(q[2], q[3]); - } - if (ok2) { - *(uint32_t*)(dst + 4 * (uint32_t)(b2 + t)) = pair16(q[4], q[5]); - *(uint32_t*)(dst + 4 * (uint32_t)(b2 + (t | 4))) = pair16(q[6], q[7]); + uint8_t* const d2 = (uint8_t*)(uintptr_t)*(const uint64_t*)(prm + 56); + *(uint32_t*)(d2 + 4 * (uint32_t)(b1 + t)) = w[0]; + *(uint32_t*)(d2 + 4 * (uint32_t)(b1 + (t | 4))) = w[1]; } + return; +#endif + // Q5..Q8 keep exactly one store group; Q4 keeps all four +#if D4R_RRSWIN_Q >= 4 if (ok1) { - *(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + t)) = pair16(q[8], q[9]); - *(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + (t | 4))) = pair16(q[10], q[11]); + *(uint32_t*)(dst + 4 * (uint32_t)(b1 + t)) = w[0]; + *(uint32_t*)(dst + 4 * (uint32_t)(b1 + (t | 4))) = w[1]; +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 15 + return; +#endif } +#endif +#if D4R_RRSWIN_Q == 4 || D4R_RRSWIN_Q >= 6 if (ok2) { - *(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + t)) = pair16(q[12], q[13]); - *(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + (t | 4))) = pair16(q[14], q[15]); + *(uint32_t*)(dst + 4 * (uint32_t)(b2 + t)) = w[2]; + *(uint32_t*)(dst + 4 * (uint32_t)(b2 + (t | 4))) = w[3]; +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 16 + return; +#endif } +#endif +#if D4R_RRSWIN_Q == 4 || D4R_RRSWIN_Q == 7 + if (ok1) { + *(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + t)) = w[4]; + *(uint32_t*)(dst + 4 * (uint32_t)(b1 + st + (t | 4))) = w[5]; +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 17 + return; +#endif + } +#endif +#if D4R_RRSWIN_Q == 4 || D4R_RRSWIN_Q == 8 + if (ok2) { + *(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + t)) = w[6]; + *(uint32_t*)(dst + 4 * (uint32_t)(b2 + st + (t | 4))) = w[7]; + } +#endif } @@ -1065,6 +1350,15 @@ DEVF void surface_scatter21(const uint8_t* prm, const uint32_t d[8][2], int lane // the callee reads 0x0 in all eighteen slots, and an 18 x i32 C++ struct return is an sret pointer, // a different LLVM signature that will not link. `scratch` is the base of the entry's `input_tensor` // (.shared 3200 bytes, dead after the cut), 72 bytes apart per lane. +#ifndef D4R_RRSWIN_PAD +#define D4R_RRSWIN_PAD 0 +#endif +#ifndef D4R_RRSWIN_VGPRC +#define D4R_RRSWIN_VGPRC 0 +#endif +#if D4R_RRSWIN_VGPRC +__attribute__((amdgpu_max_vgpr_count(D4R_RRSWIN_VGPRC))) +#endif extern "C" __attribute__((device)) void d4r_swin_enc0( uint32_t a0, uint32_t a1, uint32_t a2, uint32_t a3, uint32_t a4, uint32_t a5, uint32_t a6, uint32_t a7, uint32_t a8, uint32_t a9, uint32_t a10, uint32_t a11, uint32_t a12, uint32_t a13, @@ -1074,7 +1368,8 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( { const uint32_t as1[24] = {a0, a1, a2, a3, a4, a5, a6, a7, a8, a9, a10, a11, a12, a13, a14, a15, a16, a17, a18, a19, a20, a21, a22, a23}; - const uint32_t seeds[4] = {b0, b1, b2, b3}; + // b0..b3 (phase 1's bias C seeds, this lane's t-slice of the per-column blob) are unused + // by the engine body: its column-major C seed reads the whole blob at w+5120 directly. // The epilogue's eighteen registers: %r17597, %r18203 (%tid.x) and phase 22's eight mma D pairs. // Address space 3 (the shared scratch) explicitly: a generic store lowers to flat, which faults. __attribute__((address_space(3))) uint32_t* out = @@ -1083,33 +1378,49 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( const int g = (int)(lane_id() >> 2); // ... and of its sixteen rows const uint32_t lane = lane_id(); const half_t eps = (half_t)0.0009765625f; // the tree's sum-of-squares epsilon +#if D4R_RRSWIN_PAD + volatile uint32_t pad[D4R_RRSWIN_PAD]; +#pragma unroll + for (int i = 0; i < D4R_RRSWIN_PAD; ++i) { + pad[i] = (uint32_t)i; + } +#endif + + // The staged-image engine's shared regions: the prepared slots (1024 B) at +0, the + // phase's A image (2048 B) at +1024, and the dcell exchange overlaying the dead image + // (3072 of the 3200 bytes; the score_gemm/pv_gemm `slot` = scratch use comes later, + // by then both are dead). AS3 pointers throughout: generic pointers built from the + // raw shared-offset value lower to flat accesses at VA ~0 and fault. + lds_u8* const slots = (lds_u8*)(uintptr_t)scratch; + lds_u8* const img = slots + 1024; + lds_u32* const ex = (lds_u32*)(uintptr_t)(scratch + 1024); // ---- phase 1 ----------------------------------------------------------- - uint32_t d1[6][2][4]; // [m-tile][row half][n-tile]: f16 pair - gemm(as1, w, 4096, nullptr, seeds, nullptr, 6, t, d1); + // The staged block's fragments stage into the image per 48-token half (the six m-tiles + // restaged at local rows 0..47); the C seed is the per-column bias at w+5120. The + // result stays in registers as pairs; E1's relu requantisation happens inside phase 2's + // mode-1 staging. + uint32_t pairs1[2][6][4]; + enc0_phase1(as1, w, slots, img, pairs1); #if D4R_RRSWIN_DEBUG_STAGE == 1 { + uint32_t d1o[6][2][4]; + enc0_dcell_exchange(&pairs1[0][0][0], 6, ex, &d1o[0][0][0]); uint32_t regs[48]; - stage_regs(d1, regs); + stage_regs(d1o, regs); debug_store(prm, regs, 48); } #endif // ---- phase 2, on phase 1's relu requantisation -------------------------- - // cvt.rn.satfinite.relu.e4m3x2.f16x2, packed by the epilogue as {n-tile 0, - // n-tile 2} into the k-group 0 register and {n-tile 1, n-tile 3} into k-group 1. - uint32_t as2[24]; - for (int mi = 0; mi < 6; ++mi) { - for (int row = 0; row < 2; ++row) { - const uint32_t* d = d1[mi][row]; - const uint32_t c0 = e4m3x2(relu_pair(d[0])) | (e4m3x2(relu_pair(d[2])) << 16); - const uint32_t c1 = e4m3x2(relu_pair(d[1])) | (e4m3x2(relu_pair(d[3])) << 16); - as2[4 * mi + row] = c0; - as2[4 * mi + 2 + row] = c1; - } - } + // E1's relu (cvt.rn.satfinite.relu.e4m3x2.f16x2) is a per-pair transform, so staging + // phase 1's pairs with swin_stage_operand_p's mode 1 replaces the old as2 pack + // bit-identically; the C seed is the per-column bias at w+6208. The dcell exchange + // feeds norm_as (and stage 2); the pairs themselves stay for phase 7's C. + uint32_t pairs2[2][6][4]; + enc0_phase2(pairs1, w, slots, img, pairs2); uint32_t d2[6][2][4]; - gemm(as2, w, 5184, w + 6208, nullptr, nullptr, 6, t, d2); + enc0_dcell_exchange(&pairs2[0][0][0], 6, ex, &d2[0][0][0]); #if D4R_RRSWIN_DEBUG_STAGE == 2 { uint32_t regs[48]; @@ -1128,11 +1439,30 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( uint32_t as3[24]; norm_as(d2, w, 6272, eps, 6, as3); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 1 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 1) { + return; + } +#endif + // ---- phases 3 and 4, the fused Q/V projection, on that A operand ------- + // Both phases read the same normed-tile image halves (as3's fragment groups restaged + // per half); their D is consumed only as the dcell view (phases 5/6 walk it scalar). uint32_t d3[6][2][4]; uint32_t d4[6][2][4]; - gemm(as3, w, 6336, nullptr, nullptr, nullptr, 6, t, d3); - gemm(as3, w, 7360, nullptr, nullptr, nullptr, 6, t, d4); + enc0_proj96(as3, w, 6336, slots, img, ex, &d3[0][0][0]); + enc0_proj96(as3, w, 7360, slots, img, ex, &d4[0][0][0]); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 2 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 2) { + return; + } +#endif #if D4R_RRSWIN_DEBUG_STAGE == 3 { uint32_t regs[48]; @@ -1207,12 +1537,11 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( #endif // ---- E6 and phase 7, the first window's output projection ---------------- - // A is phase 6's D packed as before; C is phase 2's D plus the per-column table at - // weight bytes 37056..37120 (four words of eight f16, indexed by the 64-byte record - // each group of four lanes shares), taken from the d2 m-tile and n-tile the register - // order names: the mma's C pair for (m-tile mi, n-tile nt) is d2[mi][row][nt] -- the - // identity: phase 7's first four m-tiles are phase 2's rows 0..63 (VERIFIED, stage e6). - // Its output is also phase 12's C operand, so it has to exist before the second window. + // A is phase 6's D packed as before (the as7 fragment pack stages as the image); C is + // phase 2's pairs plus the per-column table at weight bytes 37056..37120, seeded per + // tile from the producer's own pairs (contract §2's residual identity, VERIFIED, stage + // e6). The pairs stay for phase 12's C; the dcell view only the stage 7 and stage 112 + // c12 rebuilds need, so its exchange is debug-only. uint32_t as7[16]; for (int mi = 0; mi < 4; ++mi) { for (int row = 0; row < 2; ++row) { @@ -1220,21 +1549,40 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( as7[4 * mi + 2 + row] = e4m3x2(d6[mi][row][1]) | (e4m3x2(d6[mi][row][3]) << 16); } } - uint32_t c7[64]; // [(4*m-tile + row)][n-tile] - for (int mi = 0; mi < 4; ++mi) { - for (int nt = 0; nt < 4; ++nt) { - const uint32_t tab = *(const uint32_t*)(w + 37056 + 16 * nt + ((lane << 2) & 12)); - for (int row = 0; row < 2; ++row) { - const uint32_t r = d2[mi][row][nt]; - c7[(4 * mi + row) * 4 + nt] = pair_of( - (half_t)(half_lo(tab) + half_lo(r)), (half_t)(half_hi(tab) + half_hi(r))); - } + uint32_t wt7[2][4]; + enc0_load_cols(wt7, w, 37056); + uint32_t pairs7[2][4][4]; + rrswin::swin_slots_from_wimg(slots, w, 20672, 32); + rrswin::swin_stage_fragment(img, 32, as7, 4); + __builtin_amdgcn_wave_barrier(); + for (int nt = 0; nt < 4; ++nt) { + for (int mt = 0; mt < 2; ++mt) { + f8v c; + enc0_seed_res_tab(c, mt, nt, &pairs2[0][0][0], 6, &wt7[0][0]); + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, 16 * nt); + enc0_pack(pairs7[mt][nt], c); } } - uint32_t d7[4][2][4]; - gemm(as7, w, 20672, nullptr, nullptr, c7, 4, t, d7); + __builtin_amdgcn_wave_barrier(); +#if D4R_RRSWIN_DEBUG_STAGE == 7 || D4R_RRSWIN_DEBUG_STAGE == 112 + uint32_t d7o[4][2][4]; + enc0_dcell_exchange(&pairs7[0][0][0], 4, ex, &d7o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 202 { + // the oracle's C array, rebuilt from phase 2's dcell view (which norm_as read in + // production) plus the per-column table + uint32_t c7[64]; + for (int mi = 0; mi < 4; ++mi) { + for (int nt = 0; nt < 4; ++nt) { + const uint32_t tab = *(const uint32_t*)(w + 37056 + 16 * nt + ((lane << 2) & 12)); + for (int row = 0; row < 2; ++row) { + const uint32_t r = d2[mi][row][nt]; + c7[(4 * mi + row) * 4 + nt] = pair_of( + (half_t)(half_lo(tab) + half_lo(r)), (half_t)(half_hi(tab) + half_hi(r))); + } + } + } uint32_t regs[48]; int i = 0; for (int mi = 0; mi < 4; ++mi) { @@ -1254,16 +1602,18 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( #if D4R_RRSWIN_DEBUG_STAGE == 7 { uint32_t regs[32]; - int i = 0; - for (int mi = 0; mi < 4; ++mi) { - for (int nt = 0; nt < 4; ++nt) { - regs[i++] = d7[mi][0][nt]; - regs[i++] = d7[mi][1][nt]; - } - } + stage_regs4(d7o, regs); debug_store(prm, regs, 32); } #endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 3 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 3) { + return; + } +#endif // ---- phases 8 and 9: the second window's two projections ----------------- // The same LayerNorm tile as the A operand of phases 3 and 4 (all four mma read the @@ -1272,8 +1622,8 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( // its V from D9. E7 and E8 between them only load the two weight tiles this reads. uint32_t d8[6][2][4]; uint32_t d9[6][2][4]; - gemm(as3, w, 21696, nullptr, nullptr, nullptr, 6, t, d8); - gemm(as3, w, 22720, nullptr, nullptr, nullptr, 6, t, d9); + enc0_proj96(as3, w, 21696, slots, img, ex, &d8[0][0][0]); + enc0_proj96(as3, w, 22720, slots, img, ex, &d9[0][0][0]); #if D4R_RRSWIN_DEBUG_STAGE == 8 { uint32_t regs[48]; @@ -1340,11 +1690,19 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( } #endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 4 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 4) { + return; + } +#endif // ---- E11 and phase 12: the second window's output projection -------------- - // Unlike E6 there is no table: phase 11's output is requantised only, because the - // residual of the first window is phase 12's C operand (the mma's C pair for m-tile - // mi and n-tile nt is phase 7's D register pair %r6334+10*(4*mi+nt) -- VERIFIED - // against the oracle's own C registers, debug stage p12o). + // Unlike E6 there is no table: the residual of the first window is phase 12's C + // operand, seeded per tile from phase 7's own pairs (contract §2's identity, VERIFIED + // against the oracle's own C registers, debug stage p12o). The dcell view feeds E12's + // RMS norm (and stage 12); the pairs stay for phase 14's C. uint32_t as12[16]; for (int mi = 0; mi < 4; ++mi) { for (int row = 0; row < 2; ++row) { @@ -1352,16 +1710,21 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( as12[4 * mi + 2 + row] = e4m3x2(d11[mi][row][1]) | (e4m3x2(d11[mi][row][3]) << 16); } } - uint32_t c12[64]; // [(4*m-tile + row)][n-tile]: phase 7's accumulators - for (int mi = 0; mi < 4; ++mi) { - for (int nt = 0; nt < 4; ++nt) { - for (int row = 0; row < 2; ++row) { - c12[(4 * mi + row) * 4 + nt] = d7[mi][row][nt]; - } + uint32_t pairs12[2][4][4]; + rrswin::swin_slots_from_wimg(slots, w, 36032, 32); + rrswin::swin_stage_fragment(img, 32, as12, 4); + __builtin_amdgcn_wave_barrier(); + for (int nt = 0; nt < 4; ++nt) { + for (int mt = 0; mt < 2; ++mt) { + f8v c; + enc0_seed_res(c, mt, nt, &pairs7[0][0][0], 4); + c = rrswin::swin_gemm_tile(c, slots, 16 * mt, 32, img, 32, 16 * nt); + enc0_pack(pairs12[mt][nt], c); } } + __builtin_amdgcn_wave_barrier(); uint32_t d12[4][2][4]; - gemm(as12, w, 36032, nullptr, nullptr, c12, 4, t, d12); + enc0_dcell_exchange(&pairs12[0][0][0], 4, ex, &d12[0][0][0]); #if D4R_RRSWIN_DEBUG_STAGE == 12 { uint32_t regs[32]; @@ -1378,7 +1741,10 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( #if D4R_RRSWIN_DEBUG_STAGE == 112 { // phase 12's mma operands: the four A fragments, the four B fragment pairs and - // the 32 C seeds in the oracle's (m-tile, n-tile) pair order + // the 32 C seeds in the oracle's (m-tile, n-tile) pair order -- the C array + // rebuilt from phase 7's debug-only dcell view + uint32_t c12[64]; + cres_regs(d7o, c12); uint32_t regs[56]; for (int o = 0; o < 16; ++o) { regs[o] = as12[o]; @@ -1389,6 +1755,14 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( } #endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 5 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 5) { + return; + } +#endif // ---- E12 and phases 13 to 20: the MLP ------------------------------------ // E12's RMS norm is the same sum-of-squares tree as the LayerNorm above, on phase 12's // accumulators and with the gain at weight bytes 37120..37184; its epsilon is E12's own @@ -1399,15 +1773,19 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( norm_as(d12, w, 37120, (half_t)0.0001220703125f, 4, asN); // ---- phase 13 and E13, the first branch ----------------------------------- - // A = the normed tile, B = the weights at 37184, C = the bias at 38208 (one word per - // n-tile, replicated over the 16 rows of its tile). E13 activates the result and adds - // the per-column table at 39296 to the activation's *input*, which is phase 14's C. - uint32_t d13[4][2][4]; - gemm(asN, w, 37184, w + 38208, nullptr, nullptr, 4, t, d13); + // A = the normed tile image, B = the weights at 37184, C = the per-column bias at + // 38208. E13's activation stages into phase 14's image as mode 2; the d13 view only + // stages 13 and 114 need, so its exchange is debug-only. + uint32_t pairs13[2][4][4]; + enc0_branch(asN, w, 37184, 38208, slots, img, pairs13); +#if D4R_RRSWIN_DEBUG_STAGE == 13 || D4R_RRSWIN_DEBUG_STAGE == 114 + uint32_t d13o[4][2][4]; + enc0_dcell_exchange(&pairs13[0][0][0], 4, ex, &d13o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 13 { uint32_t regs[32]; - stage_regs4(d13, regs); + stage_regs4(d13o, regs); debug_store(prm, regs, 32); } #endif @@ -1422,21 +1800,31 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( debug_store(prm, regs, 28); } #endif - uint32_t as14[16]; - act_pack(d13, as14); - uint32_t c14[64]; - tab_seed(d12, w + 39296, (int)lane, c14); - uint32_t d14[4][2][4]; - gemm(as14, w, 38272, nullptr, nullptr, c14, 4, t, d14); + // E13's activation is phase 14's A image (mode 2 = the clamped cubic per pair, + // bit-identical to the old act_pair + e4m3x2 pack); C is phase 12's pairs plus the + // per-column table at 39296, added in f16 before the widen (VERIFIED, stage p14o). + uint32_t wt14[2][4]; + enc0_load_cols(wt14, w, 39296); + uint32_t pairs14[2][4][4]; + enc0_branch_act(pairs13, &pairs12[0][0][0], &wt14[0][0], w, 38272, slots, img, pairs14); +#if D4R_RRSWIN_DEBUG_STAGE == 14 || D4R_RRSWIN_DEBUG_STAGE == 116 + uint32_t d14o[4][2][4]; + enc0_dcell_exchange(&pairs14[0][0][0], 4, ex, &d14o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 14 { uint32_t regs[32]; - stage_regs4(d14, regs); + stage_regs4(d14o, regs); debug_store(prm, regs, 32); } #endif #if D4R_RRSWIN_DEBUG_STAGE == 114 { + // the oracle's as14 pack and C array, rebuilt from the d13/d12 dcell views + uint32_t as14[16]; + act_pack(d13o, as14); + uint32_t c14[64]; + tab_seed(d12, w + 39296, (int)lane, c14); uint32_t regs[56]; for (int o = 0; o < 16; ++o) { regs[o] = as14[o]; @@ -1450,12 +1838,16 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( // ---- phases 15 and 16, the second branch --------------------------------- // The same pair with the weights at 39360 and 40448 and the bias at 40384: no table // add, and phase 16's C is branch one's accumulators. - uint32_t d15[4][2][4]; - gemm(asN, w, 39360, w + 40384, nullptr, nullptr, 4, t, d15); + uint32_t pairs15[2][4][4]; + enc0_branch(asN, w, 39360, 40384, slots, img, pairs15); +#if D4R_RRSWIN_DEBUG_STAGE == 15 || D4R_RRSWIN_DEBUG_STAGE == 116 + uint32_t d15o[4][2][4]; + enc0_dcell_exchange(&pairs15[0][0][0], 4, ex, &d15o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 15 { uint32_t regs[32]; - stage_regs4(d15, regs); + stage_regs4(d15o, regs); debug_store(prm, regs, 32); } #endif @@ -1470,21 +1862,26 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( debug_store(prm, regs, 28); } #endif - uint32_t as16[16]; - act_pack(d15, as16); - uint32_t c16[64]; - cres_regs(d14, c16); - uint32_t d16[4][2][4]; - gemm(as16, w, 40448, nullptr, nullptr, c16, 4, t, d16); + uint32_t pairs16[2][4][4]; + enc0_branch_act(pairs15, &pairs14[0][0][0], nullptr, w, 40448, slots, img, pairs16); +#if D4R_RRSWIN_DEBUG_STAGE == 16 || D4R_RRSWIN_DEBUG_STAGE == 118 + uint32_t d16o[4][2][4]; + enc0_dcell_exchange(&pairs16[0][0][0], 4, ex, &d16o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 16 { uint32_t regs[32]; - stage_regs4(d16, regs); + stage_regs4(d16o, regs); debug_store(prm, regs, 32); } #endif #if D4R_RRSWIN_DEBUG_STAGE == 116 { + // the oracle's as16 pack and C array, rebuilt from the d15/d14 dcell views + uint32_t as16[16]; + act_pack(d15o, as16); + uint32_t c16[64]; + cres_regs(d14o, c16); uint32_t regs[56]; for (int o = 0; o < 16; ++o) { regs[o] = as16[o]; @@ -1495,13 +1892,25 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( } #endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 6 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 6) { + return; + } +#endif // ---- phases 17 and 18, the third branch ---------------------------------- - uint32_t d17[4][2][4]; - gemm(asN, w, 41472, w + 42496, nullptr, nullptr, 4, t, d17); + uint32_t pairs17[2][4][4]; + enc0_branch(asN, w, 41472, 42496, slots, img, pairs17); +#if D4R_RRSWIN_DEBUG_STAGE == 17 || D4R_RRSWIN_DEBUG_STAGE == 118 + uint32_t d17o[4][2][4]; + enc0_dcell_exchange(&pairs17[0][0][0], 4, ex, &d17o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 17 { uint32_t regs[32]; - stage_regs4(d17, regs); + stage_regs4(d17o, regs); debug_store(prm, regs, 32); } #endif @@ -1516,21 +1925,26 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( debug_store(prm, regs, 28); } #endif - uint32_t as18[16]; - act_pack(d17, as18); - uint32_t c18[64]; - cres_regs(d16, c18); - uint32_t d18[4][2][4]; - gemm(as18, w, 42560, nullptr, nullptr, c18, 4, t, d18); + uint32_t pairs18[2][4][4]; + enc0_branch_act(pairs17, &pairs16[0][0][0], nullptr, w, 42560, slots, img, pairs18); +#if D4R_RRSWIN_DEBUG_STAGE == 18 + uint32_t d18o[4][2][4]; + enc0_dcell_exchange(&pairs18[0][0][0], 4, ex, &d18o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 18 { uint32_t regs[32]; - stage_regs4(d18, regs); + stage_regs4(d18o, regs); debug_store(prm, regs, 32); } #endif #if D4R_RRSWIN_DEBUG_STAGE == 118 { + // the oracle's as18 pack and C array, rebuilt from the d17/d16 dcell views + uint32_t as18[16]; + act_pack(d17o, as18); + uint32_t c18[64]; + cres_regs(d16o, c18); uint32_t regs[56]; for (int o = 0; o < 16; ++o) { regs[o] = as18[o]; @@ -1542,12 +1956,16 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( #endif // ---- phases 19 and 20, the fourth branch --------------------------------- - uint32_t d19[4][2][4]; - gemm(asN, w, 43584, w + 44608, nullptr, nullptr, 4, t, d19); + uint32_t pairs19[2][4][4]; + enc0_branch(asN, w, 43584, 44608, slots, img, pairs19); +#if D4R_RRSWIN_DEBUG_STAGE == 19 + uint32_t d19o[4][2][4]; + enc0_dcell_exchange(&pairs19[0][0][0], 4, ex, &d19o[0][0][0]); +#endif #if D4R_RRSWIN_DEBUG_STAGE == 19 { uint32_t regs[32]; - stage_regs4(d19, regs); + stage_regs4(d19o, regs); debug_store(prm, regs, 32); } #endif @@ -1562,13 +1980,28 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( debug_store(prm, regs, 28); } #endif - uint32_t as20[16]; - act_pack(d19, as20); - uint32_t c20[64]; - cres_regs(d18, c20); + uint32_t pairs20[2][4][4]; + enc0_branch_act(pairs19, &pairs18[0][0][0], nullptr, w, 44672, slots, img, pairs20); + // E20's r20 build walks the NVIDIA d[mi][row][nt] view uint32_t d20[4][2][4]; - gemm(as20, w, 44672, nullptr, nullptr, c20, 4, t, d20); + enc0_dcell_exchange(&pairs20[0][0][0], 4, ex, &d20[0][0][0]); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 8 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 8) { + return; + } +#endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 7 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 7) { + return; + } +#endif // ---- E20 and phases 21 and 22: the merge --------------------------------- // E20 requantises phase 20 into sixteen packed words (r20, the four A fragments of // phase 22), scatters them into the plane arena, and merges the lanes of each group of @@ -1582,19 +2015,77 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( r20[4 * mi + 2 + row] = a1; } } +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 9 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 9) { + return; + } +#endif arena_scatter(prm, r20, (int)lane); uint32_t m21[16]; merge_frag(r20, scratch, (int)lane, m21); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 10 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 10) { + return; + } +#endif uint32_t d21[8][2]; band21(m21, w, w + 53888, (int)lane, d21); +#if D4R_RRSWIN_KEEP21 == 1 + // band21 live, but no E21 store: d21's only real consumer is the scatter, so without + // this the whole phase-21 band gemm is dead-code-eliminated + if (lane == 0x7ffffffu) { + out[0] = d21[0][0]; + out[1] = d21[7][1]; + out[2] = d21[3][0]; + return; + } +#endif // ---- E21 and phase 22 --------------------------------------------------- // E21 requantises phase 21 and scatters it into the buffer at parameter block +56; // phase 22 takes the sixteen words E20 packed as its A operand (four m-tiles, K=32) // with the weights at 54016 and the two bias words at 54528. +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT == 20 + // capture: the param block's words (arena/extents/buffer pointers/geometry) handed back + // through the 18-word cut ABI, so the launch dump's arg112 surface bytes carry them + for (int i = 0; i < 18; ++i) { + out[i] = *(const uint32_t*)(prm + 4 * i); + } + return; +#endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 12 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 12) { + return; + } +#endif surface_scatter21(prm, d21, (int)lane); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 14 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 14) { + return; + } +#endif uint32_t d22[4][2][4]; gemm(r20, w, 54016, w + 54528, nullptr, nullptr, 4, t, d22, 2); +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 11 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 11) { + return; + } +#endif #if D4R_RRSWIN_DEBUG_STAGE == 20 { uint32_t regs[32]; @@ -1626,6 +2117,14 @@ extern "C" __attribute__((device)) void d4r_swin_enc0( } #endif +#if defined(D4R_RRSWIN_BISECT) && D4R_RRSWIN_BISECT >= 13 + out[0] = *(const uint32_t*)prm; + out[1] = (uint32_t)__builtin_amdgcn_workitem_id_x(); + for (int i = 2; i < 18; ++i) out[i] = 0; + if (D4R_RRSWIN_BISECT == 13) { + return; + } +#endif // ---- handing the accumulators back ------------------------------------ // %r17597 is param_0[0] (ld.param.v2.u32 at line 44316) and %r18203 is %tid.x // (line 44456); both sit inside the cut, so the epilogue cannot recompute them. diff --git a/kernels/rr/rrswin_enc0_engine.md b/kernels/rr/rrswin_enc0_engine.md new file mode 100644 index 0000000..13001b6 --- /dev/null +++ b/kernels/rr/rrswin_enc0_engine.md @@ -0,0 +1,163 @@ +# `cuda_dldn_engine_swin_enc0_kernel` on the staged-image engine + +This file is the implementation record for enc0's port from the mirror-cut scalar body to the +staged-image engine (`kernels/rr/rrswin_engine.h`). §1–§5 are the design contract the code +implements; §6 is the evidence table that accrues as each offset and layout claim is +validated (per-phase oracle captures, engine unit tests, output-surface compares). + +Status: **design fixed; primitives and body port dispatched.** Gates and their results are +recorded at the bottom as they are measured. + +## 1. What the port changes + +The mirror-cut body `d4r_swin_enc0` (`kernels/rr/rrswin_enc0.hip`) keeps + +* its cut ABI (called from NVIDIA's PTX between the prologue and E22; arguments + `a0..a23` = the staged block's 24 A-fragment words, `b0..b3` = phase 1's C seeds, + `w` = arg040 (the 54560-byte prepared weight image), `prm` = the 1032-byte parameter + block, `scratch` = the cut's `input_tensor`, 3200 bytes of shared memory), +* every epilogue (E5/E10 softmax, E6/E12 norms and table adds, E13/15/17/19 activations, + E20 arena scatter + lane merge, E21 surface scatter, E22 handover), +* the score/P·V phases 5, 6, 10, 11 (`score_gemm`, `pv_gemm`), phase 21 (`band21`) and + phase 22's GEMM as they are (scalar f32 dots over exact e4m3 decodes, 136 of 536 mma), + +and replaces the sixteen plain GEMM phases **1, 2, 3, 4, 7, 8, 9, 12, 13, 14, 15, 16, 17, +18, 19, 20** (400 of 536 mma) with engine GEMMs: operands read from staged e4m3 images and +prepared slots, one `k32_e4m3` per (m, n) tile, accumulators rounded to f16 after each k32 +step — the conservative chain that is bit-identical to the translated kernel's +`mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16` (see §2 of docs/rr-engine.md and +`rrswin_engine_k16_test.cpp` step 2). + +Replaced with this port: the scalar `gemm()` calls for those sixteen phases, whose per-mma +cost is 64 byte loads from the weight image, 64 e4m3 decodes, 128 f32 FMAs and twelve DPP +permutes per lane. + +## 2. Orientation and the two identities that make the port mechanical + +Measured facts (all previously validated, not re-derived here): + +* **D layout** (`kernels/tools/rrswin_operand_fit.cpp`): engine tile D — lane `l`, VGPR `i` + holds row `i + 8h`, column `l & 15`, `h = l >> 4`. With the engine's operand roles + (`swin_gemm_tile`: weights = A via `kop_image`, activations = B via `swin_operand`) the + tile's D is `[weight column wr][token tk]` — NVIDIA's D transposed: + `D[wr = i + 8h + 16*mt][tk = (l&15) + 16*nt]` for tile `(mt, nt)`. +* **Staging D into the next phase's A image** is exactly the existing + `swin_stage_operand(img, stride, 16*nt, mt, acc)`: it writes + `img[(16*nt + l&15)*stride + 16*mt + 8h + i] = e4m3(acc[i])`, i.e. image row = token, image + k = weight column — the A-operand form the next phase's `swin_operand` reads. +* **The residual identity**: the next phase's C operand is a cell-for-cell copy of the + producer's D. NVIDIA's mma computes `D = A·B + C` where `C[m][n]` sits at `(token, weight + column)`; the engine tile's D cell `(wr, tk)` is `C[tk][wr]`. So phase P+1's engine tile + seed is `c[i] = f16(P's own tile (mt, nt) at VGPR i)` — the producer's accumulator compacted + to f16 pairs, no exchange, no shuffle. This is why phases 12, 16, 18, 20 (residual = + previous phase's D) and 7 and 14 (residual + per-column table) need no C plumbing. +* **Per-column biases and tables are column-major f16**: every read of the form + `*(u32*)(w + base + 16*nt + ((laneid<<2)&12))` (E9's seeds, `wbias_regs`, E6/E13's tables) + places column `c` at byte `2*c` of the blob. Evidence: the word for `(nt, lane)` covers + columns `8nt+2t, 8nt+2t+1` at bytes `16nt+4t .. +3 = 2*(8nt+2t) .. +3`. The engine tile + seed for a per-column C is therefore `c[i] = f32(f16 at w + 2*(i + 8h + 16*mt))`. + +## 3. The dcell exchange (the one cross-lane primitive the port adds) + +Epilogues that consume the NVIDIA `d[mi][rh][nt]` view (`norm_as` for phases 2 and 12; the +scalar `score_gemm`/`pv_gemm` for phases 3, 4, 8, 9; E20's `r20` build for phase 20) get it +from the engine tiles through one shared-memory exchange. + +* Writer (per tile `(mt, nt)`, lane `l`, its own four f16 pairs + `pairs[p] = f16(D[2p]) | f16(D[2p+1]) << 16`, `p = 0..3`): + `lds[(mt*n_nt + nt)*128 + lane*4 + p] = pairs[(mt*n_nt + nt)*4 + p]` — all of it in + **u32 units: 128 words (= 512 bytes) per tile**, 32 lanes × 4 words. +* Reader: `d[mi][rh][nt] = lds[((nt>>1)*n_nt + mi)*128 + (8rh + g + 16*(nt&1))*4 + t]` + with `g = laneid>>2`, `t = laneid&3` — same units. + +Derivation (checkable against the two measured facts above): the wanted cell is +`D[wr = 8nt + 2t + cb][tk = 16mi + 8rh + g]`; its engine location is tile +`(mt_e = nt>>1, nt_e = mi)`, lane `(tk%16) + 16*((wr%16)>>3) = 8rh + g + 16*(nt&1)`, +VGPR `wr%8 = 2t + cb` — i.e. the reader's own pair index `t`, both halves. A full +twelve-tile exchange is 12 × 512 B = 6144 B and an eight-tile one 4096 B — both exceed +the scratch at exchange time, so **every exchange runs in 2048-byte chunks** (four tiles +at a time) overlaying the just-dead image region, with the chunk loop's barriers owned by +the phase code (§5). + +## 4. New primitives (kernels/rr/rrswin_engine.h) + +All are thin, checked by `kernels/tools/rrswin_enc0_engine_test.cpp` against independent +host models before the body relies on them. + +| primitive | signature | semantics | +|---|---|---| +| `swin_slots_from_wimg` | `(uint8_t* slots, const uint8_t* w, int base, int n_rows)` | One phase's B fragments (8 u32 per lane at `w + base + 512*t + 16*lane + 4*word`, `t = 0..n_rows/16-1`, `word = 0..3`) scattered into slot layout: slot `(n_rows*(k>>4) + n)*16 + (k&15)`, with `(n, k)` from the verified map `n = 16t + (lane>>2) + 8*(word>>1)`, `k = 2t' + (j&1) + 16*(j>>1) + 8*(word&1)`. Matches `swin_gemm_tile`'s read with `slot_base = 16*mt`, `slot_step = n_rows`. | +| `swin_dcell_exchange` | writer `(lds_u32*, const uint32_t* pairs, int n_mt, int n_nt)` / reader `(const lds_u32*, uint32_t* d, int n_mi, int n_nt4, ...)` | §3. Writer and reader are separate calls; the wave barrier between them belongs to the caller. | +| `swin_stage_operand_p` | `(uint8_t* img, int stride, int row_base, int chunk, const uint32_t pairs[4], int mode)` | `swin_stage_operand` from a tile's four f16 pairs with a per-pair transform before `e4m3x2`: mode 0 plain, 1 `relu_pair` (ZLUDA's `relu_keep_nan_halves`), 2 `act_pair` (the clamped cubic). Byte `i` of the image row holds `e4m3(transform(pair[i/2]) half i%2)`. | + +`swin_stage_fragment` needs no row-offset variant: the second half of a split 96-row GEMM +restages its fragment tile group at *local* rows 0–47 (`swin_stage_fragment(img, 32, +frag + 12, 3)`) and runs its GEMM tiles against local `row_base = 0`, emitting results to +the tiles' *global* indices (nt 3–5). The D output of every split phase is registers +(pairs), so no image write needs the global row position. + +The relu and act transforms are per-element, so `e4m3x2(transform(pair))` equals the +translated kernel's `transform`-then-`cvt.rn.satfinite.e4m3x2.f16x2` element for element. + +## 5. Shared-memory map (3200 bytes, the cut's `input_tensor`) + +`input_noise` (800 B) is a separate `.shared` variable whose offset relative to +`input_tensor` the linked body cannot know, so the budget is the 3200-byte scratch only. + +| region | size | live during | +|---|---|---| +| image buffer | 1536 (48 rows × 32 B) or 2048 (64 rows × 32 B) | one phase's A image, split into halves for the 96-row phases | +| slots | 1024 (64 slots × 16 B; 512 for phase 22, unused — phase 22 stays scalar) | one phase's GEMM | +| dcell exchange | 128 × (tile count) | between a phase's GEMM and its dcell-consuming epilogue | + +The 96-row phases (1, 2, 3, 4, 8, 9) cannot hold a 3072-byte image and 1024 bytes of slots +at once, so their GEMM runs in two token-tile halves (nt 0–2, nt 3–5): stage rows 0–47, +run, restage rows 48–95, run. The halves' staging sources are registers (the cut's `as1`, +the normed tile `as3`), so restaging is cheap. Phases 7 and 12–20 have 64-row images +(2048 B) and fit without splitting. Every exchange overlays the just-dead image or slot +region of the same phase. + +## 6. Per-phase implementation table + +| phase | mma | A (engine B operand) | B (engine A operand) | C seed | D consumed as | +|---|---|---|---|---|---| +| 1 | 24 | `as1` fragments → image halves | slots from `w+4096` | per-column, `w+5120` | relu-staged into phase 2's image (mode 1), no dcell | +| 2 | 24 | relu-staged image of D1 | slots from `w+5184` | per-column, `w+6208` | dcell exchange → `norm_as` (gain `w+6272`, eps 2⁻¹³·2⁻⁴? see body) → `as3`; pairs kept for phase 7's C | +| 3 | 24 | `as3` halves | slots from `w+6336` | — | dcell exchange → `d3` | +| 4 | 24 | `as3` halves | slots from `w+7360` | — | dcell exchange → `d4` | +| 5–6 | 48+48 | scalar `score_gemm`/`pv_gemm` unchanged (`d3`, `d4`, bias `w+8384`) | | | | +| 7 | 16 | `as7` (from `d6`) → image | slots from `w+20672` | `d2` pairs + table `w+37056` | pairs kept (phase 12's C) | +| 8 | 24 | `as3` halves | slots from `w+21696` | — | dcell exchange → `d8` | +| 9 | 24 | `as3` halves | slots from `w+22720` | — | dcell exchange → `d9` | +| 10–11 | 48+48 | scalar `score_gemm`/`pv_gemm` unchanged (`d8`, `d9`, bias `w+23744`) | | | | +| 12 | 16 | `as12` (from `d11`) → image | slots from `w+36032` | `d7` pairs | pairs kept (phase 14's C); dcell exchange → `norm_as` (gain `w+37120`, eps 2⁻¹³) → `asN` | +| 13 | 16 | `asN` → image | slots from `w+37184` | per-column, `w+38208` | act-staged into phase 14's image (mode 2) | +| 14 | 16 | act-staged image of D13 | slots from `w+38272` | `d12` pairs + table `w+39296` | pairs kept (phase 16's C) | +| 15 | 16 | `asN` → image | slots from `w+39360` | per-column, `w+40384` | act-staged into phase 16's image | +| 16 | 16 | act-staged image of D15 | slots from `w+40448` | `d14` pairs | pairs kept (phase 18's C) | +| 17 | 16 | `asN` → image | slots from `w+41472` | per-column, `w+42496` | act-staged into phase 18's image | +| 18 | 16 | act-staged image of D17 | slots from `w+42560` | `d16` pairs | pairs kept (phase 20's C) | +| 19 | 16 | `asN` → image | slots from `w+43584` | per-column, `w+44608` | act-staged into phase 20's image | +| 20 | 16 | act-staged image of D19 | slots from `w+44672` | `d18` pairs | dcell exchange → E20's `r20` build (unchanged) | +| 21 | 32 | scalar `band21` unchanged (merged `m21`, weights `w+45696`, bias `w+53888`) | | | | +| 22 | 8 | scalar `gemm(..., nnt=2)` unchanged (`r20`, `w+54016`, bias `w+54528`) | | | | + +Weight-image byte ranges are the body's own verified table (`rrswin_enc0.hip`'s +"Weight image layout of the 22 phases" comment); `waddr`'s formula is verified on all 1024 +`(n, k)` pairs of a tile. + +## 7. Validation gates and results + +*(filled as measured)* + +| gate | requirement | result | +|---|---|---| +| translated reference set | 22 oracle stages `p1..p22` captured (`/home/wouter/rr/gate-enc0/oracle-p/dumps`), 6 launches each | **captured** | +| harness determinism | `compare_launch_dump.py` oracle-p1 vs oracle-p3: 108 args, 6 differ — only arg048, first differing byte exactly at 46,000,000 (the debug payload); all surfaces/textures/weights byte-identical | **proven** | +| performance anchor (translated) | `D4R_HARNESS_FRAME_WAIT_MS=0` (continuous, matching the 16.2 ms config), 30 frames, frames 13–30: **enc0 = 3.879 ms/frame**, total 16.79 — the documented 3.79 ms baseline reproduced within ~2 % | **measured** | +| full-frame reference | `control-80`: 80-frame run of the installed runtime with `D4R_HARNESS_SAVE_FRAMES`, 80 saved frames (`/home/wouter/rr/gate-enc0/control-80/out.raw.frameN`) | **captured** | +| primitives | `rrswin_enc0_engine_test.cpp` steps vs independent host models | pending | +| per-phase oracle | engine `D4R_RRSWIN_DEBUG_STAGE=` vs the translated stage, payload-window compare | pending | +| layer surfaces | `compare_launch_dump.py` translated vs engine enc0 output surfaces | pending | +| full frame | 80-frame engine run, frames 21–80 byte-identical to `control-80` | pending | +| performance | enc0 < 3.879 ms/frame in the anchored continuous configuration | pending | diff --git a/kernels/rr/rrswin_engine.h b/kernels/rr/rrswin_engine.h index 6a7c7d7..dbeac63 100644 --- a/kernels/rr/rrswin_engine.h +++ b/kernels/rr/rrswin_engine.h @@ -47,17 +47,28 @@ namespace rrswin { +// The engine's LDS-typed register vectors: the images, the prepared slots and the exchange +// live in the caller's shared arena, and a pointer here that is typed with no address space +// is a generic pointer whose accesses lower to flat -- in the cut context, where this body +// is linked, that faults on the raw shared-offset value (VA ~0). Everything the engine +// reads or writes in shared memory is therefore typed with address_space(3) end to end +// (lds_u8 bytes, lds_v2 the 8-byte k16 operand read, lds_v4/u8 the widened variants of the +// non-native-FP8 layout); global pointers (the weight image, the _prep writes) stay generic. +typedef __attribute__((address_space(3))) uint2_t lds_v2; +typedef __attribute__((address_space(3))) u4v lds_v4; +typedef __attribute__((address_space(3))) u8v lds_v8; + // This lane's operand for k16 chunk `chunk` of image row `row_base + (lane & 15)`; rows are // `stride` bytes apart. The row comes from the lane, as in enc0_tail (`m = l & 15`), so a // caller cannot pass a row that does not match its lane. A row holds consecutive k, 16 // bytes per k16 chunk, so a k32 image needs stride >= 32; a smaller stride silently reads // the next row. -__attribute__((device)) static inline kop swin_operand(const uint8_t* img, int stride, int row_base, int chunk) +__attribute__((device)) static inline kop swin_operand(const lds_u8* img, int stride, int row_base, int chunk) { const uint32_t l = lane_id(); const uint32_t half_bytes = (l >= 16) ? 8u : 0u; const uint32_t row = (uint32_t)row_base + (l & 15u); - const uint2_t* p = (const uint2_t*)(img + (size_t)row * (size_t)stride + 16u * (uint32_t)chunk + half_bytes); + const lds_v2* p = (const lds_v2*)(img + (size_t)row * (size_t)stride + 16u * (uint32_t)chunk + half_bytes); return kop_from8(*p); } @@ -74,14 +85,28 @@ __attribute__((device)) static inline kop swin_operand(const uint8_t* img, int s // with slot_step fixed at 16 a 32-row phase would read tile 1's rows as tile 0's second k16. // `k32_index` selects which 32-wide slice of the activations this step reads: a phase with more // than one k32 step must advance the image's chunks too, or every step re-reads k 0..31. -__attribute__((device)) static inline f8v swin_gemm_tile(f8v c, const kslot* w, int slot_base, - int slot_step, const uint8_t* img, int stride, +// The engine's own slot reader: the prepared slots live in the caller's shared arena, so the +// read must be AS3-typed (tex_common.h's generic kop_image is for global slot banks and stays +// for those callers). Byte offsets are the same as the generic one's: 8 bytes per kslot under +// the native-FP8 layout, 16 under the widened one. +__attribute__((device)) static inline kop kop_image_lds(const lds_u8* img, int idx) +{ + const int i = 2 * idx + (int)(lane_id() >> 4); +#if defined(D4R_TEX_FP8) + return *(const lds_v2*)(img + 8 * (unsigned)i); +#else + return *(const lds_v4*)(img + 16 * (unsigned)i); +#endif +} + +__attribute__((device)) static inline f8v swin_gemm_tile(f8v c, const lds_u8* w, int slot_base, + int slot_step, const lds_u8* img, int stride, int row_base, int k32_index = 0) { const int slot = slot_base + (int)(lane_id() & 15u); - const kop a0 = kop_image(w, slot); + const kop a0 = kop_image_lds(w, slot); const kop b0 = swin_operand(img, stride, row_base, 2 * k32_index); - const kop a1 = kop_image(w, slot + slot_step); + const kop a1 = kop_image_lds(w, slot + slot_step); const kop b1 = swin_operand(img, stride, row_base, 2 * k32_index + 1); return k32_e4m3(c, a0, b0, a1, b1); } @@ -94,12 +119,12 @@ __attribute__((device)) static inline f8v swin_gemm_tile(f8v c, const kslot* w, // `row_base` is the first token column this tile covers and `chunk` the k16 chunk the tile's // rows land in (a 16-row m-tile is one k16 chunk), so a phase with several tiles stages each // into its own place. -__attribute__((device)) static inline void swin_stage_operand(uint8_t* img, int stride, int row_base, +__attribute__((device)) static inline void swin_stage_operand(lds_u8* img, int stride, int row_base, int chunk, const f8v acc) { const uint32_t l = lane_id(); const uint32_t h = l >> 4, col = (uint32_t)row_base + (l & 15u); - uint8_t* p = img + (size_t)col * (size_t)stride + 16u * (uint32_t)chunk + 8u * h; + lds_u8* p = img + (size_t)col * (size_t)stride + 16u * (uint32_t)chunk + 8u * h; #pragma unroll for (int i = 0; i < 8; i += 2) { @@ -178,16 +203,16 @@ __attribute__((device)) static inline uint32_t swin_act_pair(uint32_t x) // phase's 4*m_tiles registers cover 16*m_tiles rows. Each lane scatters its own 16 bytes, which // is a shared-memory store per byte -- the whole point of the engine is that this happens once // per phase instead of once per mma. -__attribute__((device)) static inline void swin_stage_fragment(uint8_t* img, int stride, - const uint32_t* frag, int m_tiles) +__attribute__((device)) static inline void swin_stage_fragment(lds_u8* img, int stride, + const uint32_t* frag, int m_tiles) { const uint32_t l = lane_id(); const uint32_t g = l >> 2, t = l & 3u; for (int mi = 0; mi < m_tiles; ++mi) { const uint32_t* f = frag + 4 * mi; - uint8_t* r0 = img + (size_t)(16 * mi + g) * (size_t)stride; - uint8_t* r1 = img + (size_t)(16 * mi + g + 8) * (size_t)stride; + lds_u8* r0 = img + (size_t)(16 * mi + g) * (size_t)stride; + lds_u8* r1 = img + (size_t)(16 * mi + g + 8) * (size_t)stride; #pragma unroll for (int j = 0; j < 4; ++j) { @@ -217,9 +242,9 @@ struct Phase // Run one phase from `src`, staging requantized accumulators into `dst` or retaining them. // `weight_rows` is the output row count used to prepare each pair of k16 weight slot banks. // When retaining results, `acc` must hold m_tiles * n_tiles entries in row-major tile order. -__attribute__((device)) static inline void swin_phase(const Phase& ph, const kslot* w, int weight_rows, - const uint8_t* src, int src_stride, - uint8_t* dst, int dst_stride, f8v* acc) +__attribute__((device)) static inline void swin_phase(const Phase& ph, const lds_u8* w, int weight_rows, + const lds_u8* src, int src_stride, + lds_u8* dst, int dst_stride, f8v* acc) { for (int mt = 0; mt < ph.m_tiles; ++mt) for (int nt = 0; nt < ph.n_tiles; ++nt) @@ -242,4 +267,116 @@ __attribute__((device)) static inline int swin_waddr_enc0(int k, int n) + 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1); } +// ------------------------------------------------------------------ the enc0 port +// Address space 3 (the shared arena) explicitly: a generic access lowers to flat, which +// faults (rrswin_enc0.hip's typedef comment). The cut body consumes rrswin::lds_u32 and +// must not define its own. +typedef unsigned lds_u32 __attribute__((address_space(3))); + +// ZLUDA's relu_keep_nan_halves on a packed pair, byte for byte the body's relu_pair +// (rrswin_enc0.hip): a negative or -0 half becomes +0, NaN kept, positive unchanged. +__attribute__((device)) static inline uint32_t swin_relu_pair(uint32_t p) +{ + if ((p & 0x8000u) != 0 && (p & 0x7fffu) <= 0x7c00u) + { + p &= 0xffff0000u; + } + if ((p & 0x80000000u) != 0 && ((p >> 16) & 0x7fffu) <= 0x7c00u) + { + p &= 0x0000ffffu; + } + return p; +} + +// The dcell exchange (contract §3), in u32 units: tile block (mt*n_nt + nt) is 128 words, +// the writer's own lane slot l*4 + p inside it, the reader walking the same words at +// (8rh + g + 16*(nt&1))*4 + t. Writer and reader are separate calls; the wave barriers +// between them belong to the caller. +__attribute__((device)) static inline void swin_dcell_exchange_write(lds_u32* lds, + const uint32_t* pairs, + int n_mt, int n_nt) +{ + const uint32_t l = lane_id() & 31u; + for (int mt = 0; mt < n_mt; ++mt) + for (int nt = 0; nt < n_nt; ++nt) +#pragma unroll + for (int p = 0; p < 4; ++p) + lds[(mt * n_nt + nt) * 128 + l * 4 + p] = pairs[(mt * n_nt + nt) * 4 + p]; +} + +// The matching reader: the epilogue's d[mi][rh][nt] from the same lds words. +__attribute__((device)) static inline void swin_dcell_exchange_read(const lds_u32* lds, uint32_t* d, + int n_mi, int n_nt4, int n_nt) +{ + const uint32_t l = lane_id(); + const uint32_t g = (l >> 2) & 7u, t = l & 3u; + for (int mi = 0; mi < n_mi; ++mi) + for (int rh = 0; rh < 2; ++rh) +#pragma unroll + for (int nt = 0; nt < n_nt4; ++nt) + d[(mi * 2 + rh) * n_nt4 + nt] = + lds[((nt >> 1) * n_nt + mi) * 128 + (8 * rh + g + 16 * (nt & 1)) * 4 + t]; +} + +// ------------------------------------------------------------------ slots from wimg +// One phase's B fragments (the WMMA slot order, rrswin_enc0.hip's documented layout) +// scattered into the slots swin_gemm_tile reads, without touching the image layout: +// one tile's 8 u32 per lane wait at w + base + 512*t + 16*lane + 4*word for tile t and +// fragment word, and their bytes are (n, k) at +// (t, lane, word, j) -> n = 16*t + (lane>>2) + 8*(word>>1), k = 2*(lane&3) + (j&1) +// + 16*(j>>1) + 8*(word&1) +// -- gemm_native_unbiased, enc0 phases 1-20 being a +// k32-steps variant of the same D4R_ZLUDA_WMMA_FP8 + bfly +// shape (D layout, wmma_layout.h) +// The destination layout is the slot one: slot (n_rows*(k>>4) + n)*16 + (k&15) of the +// phase's own bank, with the bytes in k order, exactly swin_gemm_tile's read with +// slot_base = 16*mt and slot_step = n_rows. Like the body's loops that make the same +// scatter, it is wave-wide (no barrier inside); the caller orders it against its GEMM. +__attribute__((device)) static inline void swin_slots_from_wimg(lds_u8* slots, const uint8_t* w, + int base, int n_rows) +{ + const uint32_t l = lane_id(); + for (int t = 0; t < n_rows / 16; ++t) + for (int word = 0; word < 4; ++word) + { + const uint32_t n = 16u * (uint32_t)t + (l >> 2) + 8u * ((uint32_t)word >> 1); + const uint32_t src = *(const uint32_t*)(w + base + 512 * t + 16 * l + 4 * word); +#pragma unroll + for (int j = 0; j < 4; ++j) + { + const uint32_t k = 2u * (l & 3u) + (j & 1) + 16u * (j >> 1) + 8u * (word & 1); + slots[(size_t)(n_rows * (k >> 4) + (int)n) * 16 + (k & 15)] = + (uint8_t)(src >> (8 * j)); + } + } +} + +// ---------------------------------------------- stage a tile's pairs with a transform +// swin_stage_operand from a tile's four f16 pairs (p = f16(D[2p]) | f16(D[2p+1]) << 16) +// with the epilogue's per-pair transform before e4m3x2: mode 0 plain, 1 the relu +// (swin_relu_pair), 2 the clamped cubic (swin_act_pair). Both transforms are +// per-element, so each is bit-identical to the body's packed-register path. This +// mirrors swin_stage_operand byte for byte, with the caller's pair array in place of +// the accumulator. +__attribute__((device)) static inline void swin_stage_operand_p(lds_u8* img, int stride, + int row_base, int chunk, + const uint32_t pairs[4], int mode) +{ + const uint32_t l = lane_id(); + const uint32_t h = l >> 4, col = (uint32_t)row_base + (l & 15u); + lds_u8* p = img + (size_t)col * (size_t)stride + 16u * (uint32_t)chunk + 8u * h; +#pragma unroll + for (int p_ = 0; p_ < 4; ++p_) + { + uint32_t pair = pairs[p_]; + if (mode == 1) + pair = swin_relu_pair(pair); + else if (mode == 2) + pair = swin_act_pair(pair); + const uint32_t code = e4m3x2(pair); + p[2 * p_] = (uint8_t)(code & 0xffu); + p[2 * p_ + 1] = (uint8_t)(code >> 8); + } +} + } // namespace rrswin diff --git a/kernels/tools/ptxsim.py b/kernels/tools/ptxsim.py old mode 100644 new mode 100755 index 478a95b..e3a991b --- a/kernels/tools/ptxsim.py +++ b/kernels/tools/ptxsim.py @@ -1,17 +1,32 @@ #!/usr/bin/env python3 -"""Minimal vectorised PTX interpreter for DLSS rrlite kernels. +"""Vectorised PTX interpreter for one thread block (all threads as numpy lanes). -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. +Runs one thread block of a kernel on the CPU against memory restored from a +capture (the replay directory: allocations before the launch). Every statement +can be observed (hooks, Stmt.line gives the .ptx line), which is how the +intermediates of a kernel are read. Textures and surfaces are bound per launch +through Block.resources (handle value -> Texture / Surface). -Only the instruction forms the rrlite kernels use are implemented; anything -else raises NotImplementedError naming the statement. +Only the instruction forms the RR and rrlite kernels use are implemented; +anything else raises NotImplementedError naming the statement. + +Fidelity against the GPU reference (ZLUDA on gfx1201), where the last bit can differ: +- add/sub/mul/div/sqrt of two f16 or f32 operands: exact (f64 evaluation and one rounding; + double rounding cannot occur for a single operation of binary16/binary32 operands). +- fma/mad (f16, f32): x*y+z is evaluated in f64 and rounded once; the product is exact but the + sum is rounded in f64 first, so a rare last-bit difference from a true fused multiply-add is possible. +- .approx transcendentals (ex2, lg2, rcp, rsqrt, sqrt.approx, div.approx, tanh.approx): correctly + rounded here. The AMD instructions behind the reference (v_exp_f32, v_rcp_f32, v_rsq_f32, v_log_f32) + are approximate, about 1 ulp f32, so these can differ in the last bit. +- .ftz: denormal inputs and results of f32 operations are flushed to signed zero. +- mma.sync m16n8k32 e4m3 and m16n8k16 f16: products are exact; the k-sum is evaluated in f64 and rounded + once to the accumulator type (MMA_MODE='exact'). MMA_MODE='split16' models the ZLUDA lowering of a k32 + step as two k16 WMMAs with C added in the first (each rounded to the accumulator type). +- tex (linear, normalised): texel values are exact; the bilinear weights are exact in f64 unless + TEX_FRAC_BITS quantises them (hardware filters often keep fewer fraction bits); one rounding to the result type. +- shfl, movmatrix, prmt, cvt (rounding modes), integer ops: exact. """ import re -import struct import sys import numpy as np @@ -20,6 +35,11 @@ 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} +F32_TINY = 2.0 ** -126 +F16_TINY = 2.0 ** -14 + +MMA_MODE = 'exact' # 'exact' | 'split16' (see module docstring) +TEX_FRAC_BITS = None # None: exact bilinear weights; N: weights quantised to 1/2**N def _e4m3_decode_table(): @@ -74,8 +94,13 @@ def f32_to_bits(v): return np.asarray(v, F32).view(np.uint32).astype(U64) +def flush(x, tiny): + """flush denormal values (held in float64) to signed zero""" + return np.where(np.abs(x) < tiny, np.copysign(0.0, x), x) + + class Stmt: - __slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index') + __slots__ = ('labels', 'pred', 'neg', 'op', 'args', 'text', 'index', 'line') def split_args(s): @@ -96,20 +121,35 @@ def split_args(s): def parse_kernel(path, name): + """(stmts, types, shared) of entry NAME in a .ptx file; Stmt.line is the source line.""" 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:] + k = src.index(')', i) + 1 # end of the parameter list + base = src.count('\n', 0, k) + 1 # source line of body position 0 # 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) + body = src[k:j] + + def _blank(m): # blank a match but keep its newlines: line numbers stay exact + return re.sub(r'[^\n]', ' ', m.group(0)) + body = re.sub(r'//[^\n]*', _blank, body) + body = re.sub(r'^[ \t]*\.(loc|file|pragma)[^\n]*', _blank, body, flags=re.M) + body = re.sub(r'\.(maxntid|reqntid|minnctapersm|maxnreg)[^\n{]*', _blank, body) + nl = [mm.start() for mm in re.finditer(r'\n', body)] + nli = 0 stmts, types, shared = [], {}, {} pending_labels = [] - for raw in re.sub(r'\s*\n\s*', ' ', body).split(';'): - s = raw.strip() + pos = 0 + for m in re.finditer(r';|\Z', body): + raw = body[pos:m.start()] + p0 = pos + len(raw) - len(raw.lstrip()) # first character of the statement + while nli < len(nl) and nl[nli] < p0: + nli += 1 + seg_line = base + nli + pos = m.end() + s = re.sub(r'\s*\n\s*', ' ', raw).strip() while True: s = s.strip() if s.startswith('{') and not s.startswith('{%'): @@ -118,18 +158,18 @@ def parse_kernel(path, name): 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():] + mm = re.match(r'(\$[\w]+):\s*', s) + if mm: + pending_labels.append(mm.group(1)) + s = s[mm.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(','): + mr = re.match(r'\.reg\s*\.(\w+)\s+(.*)', s) + ty = mr.group(1) + for decl in mr.group(2).split(','): decl = decl.strip() m2 = re.match(r'(%[\w]+)<(\d+)>', decl) if m2: @@ -138,30 +178,32 @@ def parse_kernel(path, name): types[decl] = ty continue if s.startswith('.shared'): - m = re.search(r'(\w+)\[(\d+)\]', s) - shared[m.group(1)] = int(m.group(2)) + ms = re.search(r'(\w+)\[(\d+)\]', s) + shared[ms.group(1)] = int(ms.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():] + mp = re.match(r'@(!?)(%\w+)\s+', s) + if mp: + st.neg, st.pred = mp.group(1) == '!', mp.group(2) + s = s[mp.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) + st.line = seg_line 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) + st.line = base + body.count('\n') stmts.append(st) - for k, v in module_shared.items(): - shared.setdefault(k, v) + for kk, v in module_shared.items(): + shared.setdefault(kk, v) return stmts, types, shared @@ -175,7 +217,7 @@ def reg_width(types, name): class Memory: - """Global memory: the replay dump's allocations at their captured addresses.""" + """Global memory: the allocations of a replay directory (manifest.txt, alloc-N.bin, args.bin).""" def __init__(self, dump_dir): manifest = open(dump_dir + '/manifest.txt').read().split('\n') @@ -195,7 +237,8 @@ class Memory: 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 + # (alloc id, element offsets within the allocation, raw bytes per lane, element type or None) + self.stores = [] def locate(self, addr): idx = np.searchsorted(self.base_arr, addr, side='right') - 1 @@ -214,14 +257,90 @@ class Memory: out[sel] = self.data[a][off[:, None] + np.arange(nbytes)] return out - def store(self, addr, raw, mask): + def store(self, addr, raw, mask, ty=None): 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())) + self.stores.append((self.ids[a], off.copy(), raw[sel].copy(), ty)) + + +class Texture: + """A texture object over captured texels. Texels are decoded to float64 RGBA; missing channels + are filled as CUDA does (0, 0, 0, 1). bits: 16/32 float, 8 UNORM (v/255), 9 SNORM (max(v/127, -1)).""" + + def __init__(self, path, width, height, channels, bits, linear): + raw = np.fromfile(str(path), np.uint8) + per = bits // 8 if bits >= 16 else 1 + if raw.size != width * height * channels * per: + raise ValueError(f'texture {path}: {raw.size} bytes, expected {width * height * channels * per}') + if bits == 16: + x = raw.view(F16).astype(F64) + elif bits == 32: + x = raw.view(F32).astype(F64) + elif bits == 8: + x = raw.astype(F64) / 255.0 + elif bits == 9: + x = np.maximum(raw.view(np.int8).astype(F64) / 127.0, -1.0) + else: + raise NotImplementedError(f'texture bits {bits}') + x = x.reshape(height, width, channels) + full = np.zeros((height, width, 4), F64) + full[..., :channels] = x + for c in range(channels, 4): + full[..., c] = 1.0 if c == 3 else 0.0 + self.texels = full + self.H, self.W = height, width + self.linear = linear + + def fetch(self, x, y): + """point fetch of texel (x, y) (int64 arrays), clamped to the texture: (N, 4) float64""" + return self.texels[np.clip(y, 0, self.H - 1), np.clip(x, 0, self.W - 1)] + + def sample(self, u, v): + """normalised coordinates, LOD 0: bilinear (linear) or nearest. (N, 4) float64""" + x = u * self.W + y = v * self.H + if not self.linear: + return self.fetch(np.floor(x).astype(np.int64), np.floor(y).astype(np.int64)) + x0, y0 = np.floor(x - 0.5), np.floor(y - 0.5) + fx, fy = (x - 0.5) - x0, (y - 0.5) - y0 + if TEX_FRAC_BITS is not None: + q = float(2 ** TEX_FRAC_BITS) + fx, fy = np.floor(fx * q) / q, np.floor(fy * q) / q + x0, y0 = x0.astype(np.int64), y0.astype(np.int64) + fx, fy = fx[:, None], fy[:, None] + top = (1 - fx) * self.fetch(x0, y0) + fx * self.fetch(x0 + 1, y0) + bot = (1 - fx) * self.fetch(x0, y0 + 1) + fx * self.fetch(x0 + 1, y0 + 1) + return (1 - fy) * top + fy * bot + + +class Surface: + """A surface: byte-addressed texels (channels x bits each), zero-initialised. `written` marks stored bytes; + stores outside the surface are dropped (zero boundary mode).""" + + def __init__(self, width, height, channels, bits): + self.W, self.H, self.C, self.bits = width, height, channels, bits + self.texel_bytes = channels * bits // 8 + self.data = np.zeros(width * height * self.texel_bytes, np.uint8) + self.written = np.zeros(self.data.size, bool) + + def write(self, x, y, vals, bits): + if bits != self.bits: + raise ValueError(f'surface of {self.bits}-bit channels written as b{bits}') + ok = (x >= 0) & (x < self.W) & (y >= 0) & (y < self.H) + base = (y[ok] * self.W + x[ok]) * self.texel_bytes + nb = bits // 8 + for c, v in enumerate(vals): + if c >= self.C: + raise ValueError(f'sust writes channel {c} of a {self.C}-channel surface') + v = v[ok] + for b in range(nb): + idx = base + c * nb + b + self.data[idx] = ((v >> U64(8 * b)) & U64(0xff)).astype(np.uint8) + self.written[idx] = True class Block: @@ -247,11 +366,14 @@ class Block: 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 + for lab in st.labels: + self.labels[lab] = st.index self.hooks = [] self.retired = np.zeros(self.T, bool) self.prov = {} + self.resources = {} # texture/surface handle value -> Texture / Surface + self.mma_hook = None + self.steps = 0 # --- operands ----------------------------------------------------------- def width(self, name): @@ -282,7 +404,7 @@ class Block: 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) + v = int(s.rstrip('Uu'), 0) # inline-asm literals carry a C suffix (0x3fb8aa3bU) return np.full(self.T, v & MASK[64], U64) def get(self, s, width=None): @@ -313,9 +435,9 @@ class Block: w = self.width(name) old = self.get(name) if w == 1: - self.regs[name] = np.where(mask, value.astype(bool), old) + self.regs[name] = np.where(mask, np.asarray(value).astype(bool), old) else: - self.regs[name] = np.where(mask, value.astype(U64) & np.uint64(MASK[w]), old) + self.regs[name] = np.where(mask, np.asarray(value).astype(U64) & np.uint64(MASK[w]), old) def vec(self, s): return [x.strip() for x in s.strip()[1:-1].split(',')] @@ -331,6 +453,12 @@ class Block: 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]) + def resource(self, h, cls): + obj = self.resources.get(int(h)) + if not isinstance(obj, cls): + raise NotImplementedError(f'no {cls.__name__} bound to handle {int(h):#x}') + return obj + # --- execution ------------------------------------------------------------ def run(self, limit=None): stmts = self.stmts @@ -340,12 +468,12 @@ class Block: steps = 0 while pc < len(stmts): st = stmts[pc] - for l in st.labels: - if l in pending: - active |= pending.pop(l) + for lab in st.labels: + if lab in pending: + active |= pending.pop(lab) if not active.any(): # jump to the nearest label someone waits at - targets = sorted(self.labels[l] for l in pending) + targets = sorted(self.labels[lab] for lab in pending) if not targets: break pc = targets[0] @@ -365,7 +493,7 @@ class Block: if not taking.any(): pc += 1 continue - raise NotImplementedError('divergent backward branch at ' + st.text) + raise NotImplementedError(f'divergent backward branch at line {st.line}: {st.text}') active = active & ~taking if taking.any(): pending[st.args[0]] = pending.get(st.args[0], np.zeros(self.T, bool)) | taking @@ -379,6 +507,7 @@ class Block: steps += 1 if limit and steps >= limit: break + self.steps = steps return steps def execute(self, st, mask): @@ -387,11 +516,13 @@ class Block: base = parts[0] fn = getattr(self, 'op_' + base, None) if fn is None: - raise NotImplementedError(st.text) + raise NotImplementedError(f'line {st.line}: {st.text}') try: fn(st, parts, a, mask) + except NotImplementedError as e: + raise NotImplementedError(f'{e} (line {st.line}: {st.text})') from e except Exception as e: - raise RuntimeError(f'{type(e).__name__} {e} in: {st.text}') from e + raise RuntimeError(f'{type(e).__name__} {e} at line {st.line}: {st.text}') from e # --- helpers -------------------------------------------------------------- @staticmethod @@ -401,8 +532,8 @@ class Block: 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 + def typed(self, s, ty, ftz=False): + w = 1 if ty == 'pred' else int(re.sub(r'\D', '', ty) or 32) v = self.get(s, w) if ty.startswith('s'): return self.sext(v, w) @@ -410,38 +541,51 @@ class Block: if ty == 'f16': return bits_to_f16(v).astype(F64) if ty == 'f32': - return bits_to_f32(v).astype(F64) + x = bits_to_f32(v).astype(F64) + return flush(x, F32_TINY) if ftz else x return v.view(F64) return (v & np.uint64(MASK[w])).astype(np.int64) if w < 64 else v - def result(self, ty, value): + def result(self, ty, value, ftz=False): if ty == 'f16': - return f16_to_bits(np.asarray(value).astype(F16)) + return f16_to_bits(np.asarray(value, F64).astype(F16)) if ty == 'f32': - return f32_to_bits(np.asarray(value).astype(F32)) + v = np.asarray(value, F64).astype(F32) + if ftz: + v = flush(v.astype(F64), F32_TINY).astype(F32) + return f32_to_bits(v) 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]) + @staticmethod + def relu(r, parts): + return np.where(r < 0, 0.0, r) if 'relu' in parts else r + # --- instructions ----------------------------------------------------------- def op_nop(self, st, parts, a, mask): pass - def op_bar(self, st, parts, a, mask): - pass + op_bar = op_nop + op_membar = op_nop + op_fence = op_nop + op_prefetch = op_nop + op_ldu = None # set below def op_ret(self, st, parts, a, mask): self.retired |= mask + op_exit = op_ret + 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) + self.put(a[0], self.get(a[1]) ^ U64(0x80008000), mask) elif ty == 'f16': - self.put(a[0], self.get(a[1]) ^ np.uint64(0x8000), mask) + self.put(a[0], self.get(a[1]) ^ U64(0x8000), mask) elif ty == 'f32': - self.put(a[0], self.get(a[1]) ^ np.uint64(0x80000000), mask) + self.put(a[0], self.get(a[1]) ^ U64(0x80000000), mask) else: self.put(a[0], self.result('u' + ty[1:], -self.typed(a[1], ty)), mask) @@ -455,14 +599,14 @@ class Block: 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) + self.put(n, (v >> U64(i * w)) & U64(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) + v |= (self.get(n) & U64(MASK[w])) << U64(i * w) self.put(d, v, mask) return self.put(d, self.get(s), mask) @@ -484,10 +628,10 @@ class Block: elif space == 'param': # ld.param through a computed address (mov.b64 %rd, ; 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) + off = np.where(mask, addr - U64(PARAM_BASE), 0).astype(np.int64) data = raw[off[:, None] + np.arange(total)] elif kind == 'shared' or space == 'shared': - off = addr.astype(np.int64) + off = np.where(mask, addr, 0).astype(np.int64) data = self.shared_mem[off[:, None] + np.arange(total)] else: data = self.mem.load(addr, total, mask) @@ -496,15 +640,16 @@ class Block: 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) + v |= chunk[:, b].astype(U64) << U64(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) + self.prov[n] = (st.index, kind if kind != 'reg' else space, addr + U64(i * nbytes), nbytes) + + op_ldu = op_ld 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]) @@ -513,88 +658,123 @@ class Block: 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) + raw[:, i * nbytes + b] = ((v >> U64(8 * b)) & U64(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) + self.mem.store(addr, raw, mask, ty) 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 h2(self, sa, sb, fn, sc=None, relu=False): + va, vb = self.get(sa), self.get(sb) + vc = self.get(sc) if sc is not None else None + out = np.zeros(self.T, U64) + for i in range(2): + x = bits_to_f16(va >> U64(16 * i)).astype(F64) + y = bits_to_f16(vb >> U64(16 * i)).astype(F64) + if vc is not None: + z = bits_to_f16(vc >> U64(16 * i)).astype(F64) + r = fn(x, y, z) + else: + r = fn(x, y) + if relu: + r = np.where(r < 0, 0.0, r) + out |= f16_to_bits(np.asarray(r).astype(F16)) << U64(16 * i) + return out + + def h2u(self, sa, fn): + va = self.get(sa) + out = np.zeros(self.T, U64) + for i in range(2): + x = bits_to_f16(va >> U64(16 * i)).astype(F64) + with np.errstate(all='ignore'): + out |= f16_to_bits(np.asarray(fn(x)).astype(F16)) << U64(16 * i) + return out + def op_add(self, st, parts, a, mask): - ty = parts[-1] + ty, ftz = parts[-1], 'ftz' in parts if ty in ('f16', 'f32'): - self.put(a[0], self.result(ty, self.typed(a[1], ty) + self.typed(a[2], ty)), mask) + r = self.typed(a[1], ty, ftz) + self.typed(a[2], ty, ftz) + self.put(a[0], self.result(ty, self.relu(r, parts), ftz), mask) elif ty == 'f16x2': - self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y), mask) + self.put(a[0], self.h2(a[1], a[2], lambda x, y: x + y, relu='relu' in parts), mask) else: self.int_binop(parts, a, mask, lambda x, y: x + y) def op_sub(self, st, parts, a, mask): - ty = parts[-1] + ty, ftz = parts[-1], 'ftz' in parts if ty in ('f16', 'f32'): - self.put(a[0], self.result(ty, self.typed(a[1], ty) - self.typed(a[2], ty)), mask) + self.put(a[0], self.result(ty, self.typed(a[1], ty, ftz) - self.typed(a[2], ty, ftz), ftz), 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): + def int_mul(self, parts, a, mask, z=None): + """mul/mad .wide and .hi (integers): products through 64-bit modular arithmetic (exact)""" 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) + w = int(re.sub(r'\D', '', ty)) + x = self.typed(a[1], ty).astype(U64) + y = self.typed(a[2], ty).astype(U64) + p = x * y + if 'hi' in parts: + out = p >> U64(w) + if z is not None: + out = out + self.get(z) + self.put(a[0], out & U64(MASK[w]), mask) + else: # wide + out = p + (self.get(z) if z is not None else U64(0)) + self.put(a[0], out, mask) + + def op_mul(self, st, parts, a, mask): + ty, ftz = parts[-1], 'ftz' in parts + if 'wide' in parts or 'hi' in parts: + self.int_mul(parts, a, mask) elif ty in ('f16', 'f32'): - self.put(a[0], self.result(ty, self.typed(a[1], ty) * self.typed(a[2], ty)), mask) + r = self.typed(a[1], ty, ftz) * self.typed(a[2], ty, ftz) + self.put(a[0], self.result(ty, self.relu(r, parts), ftz), mask) elif ty == 'f16x2': - self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y), mask) + self.put(a[0], self.h2(a[1], a[2], lambda x, y: x * y, relu='relu' in parts), 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 + ty, ftz = parts[-1], 'ftz' in parts + if ty in ('f16', 'f32'): + r = self.typed(a[1], ty, ftz) * self.typed(a[2], ty, ftz) + self.typed(a[3], ty, ftz) + self.put(a[0], self.result(ty, self.relu(r, parts), ftz), mask) + elif 'wide' in parts or 'hi' in parts: + self.int_mul(parts, a, mask, z=a[3]) + else: # mad.lo + w = int(re.sub(r'\D', '', ty)) + x = self.typed(a[1], ty).astype(U64) + y = self.typed(a[2], ty).astype(U64) + self.put(a[0], (x * y + self.get(a[3])) & U64(MASK[w]), mask) def op_fma(self, st, parts, a, mask): - ty = parts[-1] + ty, ftz = parts[-1], 'ftz' in parts if ty == 'f16x2': - self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3]), mask) + self.put(a[0], self.h2(a[1], a[2], lambda x, y, z: x * y + z, a[3], relu='relu' in parts), 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) + r = self.typed(a[1], ty, ftz) * self.typed(a[2], ty, ftz) + self.typed(a[3], ty, ftz) + self.put(a[0], self.result(ty, self.relu(r, parts), ftz), mask) else: raise NotImplementedError(st.text) def minmax(self, parts, a, mask, is_max): - ty = parts[-1] + ty, ftz = parts[-1], 'ftz' in parts + f = np.fmax if is_max else np.fmin 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) + self.put(a[0], self.result(ty, f(self.typed(a[1], ty, ftz), self.typed(a[2], ty, ftz)), ftz), mask) else: - f = np.maximum if is_max else np.minimum - self.int_binop(parts, a, mask, f) + self.int_binop(parts, a, mask, np.maximum if is_max else np.minimum) def op_min(self, st, parts, a, mask): self.minmax(parts, a, mask, False) @@ -606,15 +786,12 @@ class Block: 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) + self.put(a[0], self.get(a[1]) & U64(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) + 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) @@ -630,20 +807,20 @@ class Block: 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)))) + n = self.get(a[2]) & U64(0xff) + v = np.where(n >= w, 0, (self.get(a[1]) << np.minimum(n, U64(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) + n = (self.get(a[2]) & U64(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]) + v = self.get(a[1]) & U64(MASK[w]) r = np.where(n >= w, 0, v >> np.minimum(n, 63).astype(U64)) self.put(a[0], r.astype(U64), mask) @@ -662,107 +839,204 @@ class Block: out[i] = (int(b[i]) & ~m & MASK[w]) | ((int(f[i]) << p) & m) self.put(a[0], out, mask) + def op_bfe(self, st, parts, a, mask): + """bfe.{u,s}N d, a, pos, len""" + ty = parts[-1] + w = int(re.sub(r'\D', '', ty)) + x = self.get(a[1]) & U64(MASK[w]) + pos = (self.get(a[2]) & U64(0xff)).astype(np.int64) + ln = (self.get(a[3]) & U64(0xff)).astype(np.int64) + pos = np.minimum(pos, w) + ln = np.minimum(ln, w - pos) + field = (x >> pos.astype(U64)) & ((U64(1) << ln.astype(U64)) - U64(1)) + if ty.startswith('s'): + sgn_bit = np.where(ln > 0, (field >> (ln - 1).astype(U64)) & U64(1), 0).astype(bool) + field = np.where(sgn_bit & (ln > 0), field | (U64(MASK[64]) << ln.astype(U64)), field) + self.put(a[0], field & U64(MASK[w]) if w < 64 else field, mask) + + def op_prmt(self, st, parts, a, mask): + """prmt.b32 d, a, b, c (default mode): bytes 0-3 are a, 4-7 are b""" + if len(parts) > 2: + raise NotImplementedError(st.text) + comb = self.get(a[1]) & U64(0xffffffff) + comb = comb | ((self.get(a[2]) & U64(0xffffffff)) << U64(32)) + sel = self.get(a[3]) + out = np.zeros(self.T, U64) + for i in range(4): + s = (sel >> U64(4 * i)) & U64(0xf) + idx = s & U64(7) + byte = (comb >> (idx * U64(8))) & U64(0xff) + sign = (byte >> U64(7)) & U64(1) + byte = np.where((s & U64(8)) != 0, np.where(sign == 1, U64(0xff), U64(0)), byte) + out |= byte << U64(8 * i) + 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] + cmp, ty, ftz = parts[1], parts[-1], 'ftz' in parts + x, y = self.typed(a[1], ty, ftz), self.typed(a[2], ty, ftz) + un = np.isnan(x) | np.isnan(y) if ty in ('f16', 'f32', 'f64') else np.zeros(self.T, bool) + base = {'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, + 'equ': (x == y) | un, 'neu': (x != y) | un, 'ltu': (x < y) | un, 'leu': (x <= y) | un, + 'gtu': (x > y) | un, 'geu': (x >= y) | un, 'num': ~un, 'nan': un}[cmp] + r = np.asarray(base, bool) dests = a[0].split('|') + if len(a) == 4: # setp.CMP.BOOLOP p[|q], a, b, c + q = self.get(a[3].lstrip('!')) + if a[3].startswith('!'): + q = ~q + op = parts[2] + r = {'and': r & q, 'or': r | q, 'xor': r ^ q}[op] self.put(dests[0], r, mask) + if len(dests) > 1: + self.put(dests[1], ~r, mask) + + def op_set(self, st, parts, a, mask): + """set.CMP.DTYPE.STYPE d, a, b: d = 1.0 (a float DTYPE) or all ones (an integer DTYPE) where the comparison + holds, else 0""" + cmp, dty, ty = parts[1], parts[-2], parts[-1] + x, y = self.typed(a[1], ty, 'ftz' in parts), self.typed(a[2], ty, 'ftz' in parts) + un = np.isnan(x) | np.isnan(y) if ty in ('f16', 'f32', 'f64') else np.zeros(self.T, bool) + r = np.asarray({'lt': x < y, 'le': x <= y, 'gt': x > y, 'ge': x >= y, 'eq': x == y, 'ne': x != y, + 'equ': (x == y) | un, 'neu': (x != y) | un, 'ltu': (x < y) | un, 'leu': (x <= y) | un, + 'gtu': (x > y) | un, 'geu': (x >= y) | un, 'num': ~un, 'nan': un}[cmp], bool) + if dty.startswith('f'): + self.put(a[0], self.result(dty, np.where(r, 1.0, 0.0)), mask) + else: + w = int(re.sub(r'\D', '', dty)) + self.put(a[0], np.where(r, np.uint64(MASK[w]), np.uint64(0)).astype(U64), 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 unary(self, parts, a, mask, fn): + """unary float op on f16, f16x2 or f32 (ftz honoured for f32)""" + ty, ftz = parts[-1], 'ftz' in parts + with np.errstate(all='ignore'): + if ty == 'f16x2': + self.put(a[0], self.h2u(a[1], fn), mask) + elif ty == 'f16': + self.put(a[0], f16_to_bits(np.asarray(fn(self.typed(a[1], 'f16'))).astype(F16)), mask) + elif ty in ('f32', 'f64'): + self.put(a[0], self.result(ty, fn(self.typed(a[1], ty, ftz)), ftz), mask) + else: + raise NotImplementedError(st.text if False else f'unary {ty}') + 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) + self.unary(parts, a, mask, np.exp2) + + def op_lg2(self, st, parts, a, mask): + self.unary(parts, a, mask, np.log2) + + def op_tanh(self, st, parts, a, mask): + self.unary(parts, a, mask, np.tanh) + + def op_sqrt(self, st, parts, a, mask): + self.unary(parts, a, mask, np.sqrt) 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) + self.unary(parts, a, mask, lambda x: 1.0 / np.sqrt(x)) 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) + self.unary(parts, a, mask, lambda x: 1.0 / x) 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) + ty, ftz = parts[-1], 'ftz' in parts + if ty in ('f32', 'f64', 'f16'): + with np.errstate(all='ignore'): + x, y = self.typed(a[1], ty, ftz), self.typed(a[2], ty, ftz) + self.put(a[0], self.result(ty, x / y, ftz), mask) + return + w = int(re.sub(r'\D', '', ty)) + x = self.typed(a[1], ty) + y = self.typed(a[2], ty) + safe = np.where(y == 0, 1, y) + q = np.sign(x) * np.sign(safe) * (np.abs(x) // np.abs(safe)) + self.put(a[0], np.where(y == 0, 0, q).astype(np.int64).astype(U64) & U64(MASK[w]), mask) + + CVT_MODS = {'rn', 'rz', 'rm', 'rp', 'rni', 'rzi', 'rmi', 'rpi', 'sat', 'satfinite', 'relu', 'ftz', 'approx'} 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')] + types = [p for p in parts[1:] if p not in self.CVT_MODS] dt, stp = types[0], types[1] - relu = 'relu' in parts + ftz = 'ftz' 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: + h = (v >> U64(16 * i)) & U64(0xffff) + if 'relu' in parts: 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) + out |= F16_TO_E4M3[h.astype(np.int64)].astype(U64) << U64(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) + c = ((v >> U64(8 * i)) & U64(0xff)).astype(np.int64) + out |= f16_to_bits(E4M3[c].astype(F16)) << U64(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) + elif dt[0] in 'us' and stp.startswith('f'): # float -> integer, saturating, NaN -> 0 + x = self.typed(a[1], stp, ftz) + if any(p in parts for p in ('rzi', 'rz')) or (not any(p in parts for p in ('rni', 'rn', 'rmi', 'rm', 'rpi', 'rp'))): + r = np.trunc(x) + elif any(p in parts for p in ('rni', 'rn')): + r = np.rint(x) + elif any(p in parts for p in ('rmi', 'rm')): + r = np.floor(x) else: - self.put(a[0], f32_to_bits(iv.astype(F64).astype(F32)), mask) + r = np.ceil(x) + ws = int(re.sub(r'\D', '', dt)) + lo, hi = (-2.0 ** (ws - 1), 2.0 ** (ws - 1) - 1) if dt[0] == 's' else (0.0, 2.0 ** ws - 1) + r = np.where(np.isnan(r), 0.0, np.clip(r, lo, hi)) + self.put(a[0], r.astype(np.int64).astype(U64) & U64(MASK[ws]), mask) + elif dt.startswith('f') and (stp[0] in 'us'): # integer -> float + iv = self.typed(a[1], stp) + self.put(a[0], self.result(dt, np.asarray(iv).astype(F64), ftz), mask) + elif dt.startswith('f') and stp.startswith('f'): # float -> float (rzi / rni / rmi / rpi: to an integral value) + x = self.typed(a[1], stp, ftz) + for mod, fn in (('rzi', np.trunc), ('rni', np.rint), ('rmi', np.floor), ('rpi', np.ceil)): + if mod in parts: + x = fn(x) + self.put(a[0], self.result(dt, x, ftz), 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) + val = self.sext(v, ws) if stp[0] == 's' else (v & U64(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): + """shfl.sync.MODE.b32 d[|p], a, b, c, membermask (warp-collective; membermask is not checked)""" mode = parts[2] - v = self.get(a[1]) + dst, pred = (a[0].split('|') + [None])[:2] + v = self.get(a[1]) & U64(0xffffffff) 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) + c = self.get(a[3]).astype(np.int64) if len(a) > 3 else np.full(self.T, 31, np.int64) + clamp, seg = c & 0x1f, (c >> 8) & 0x1f + lane = self.lane.astype(np.int64) + minl = lane & seg + maxl = minl | (clamp & ~seg & 0x1f) + if mode == 'up': + j = lane - b + ok = j >= minl + elif mode == 'down': + j = lane + b + ok = j <= maxl + elif mode == 'bfly': + j = lane ^ (b & 0x1f) + ok = (j >= minl) & (j <= maxl) + elif mode == 'idx': + j = minl | (b & ~seg & 0x1f) + ok = j <= maxl + else: + raise NotImplementedError(st.text) + src = np.where(ok, j, lane) & 31 + gidx = self.warp.astype(np.int64) * 32 + src + self.put(dst, v[gidx], mask) + if pred is not None: + self.put(pred, ok, mask) def op_movmatrix(self, st, parts, a, mask): v = self.get(a[1]) @@ -779,57 +1053,91 @@ class Block: 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) + @staticmethod + def _mma_maps(k): + """per-lane (row, col) of every fragment element: A (16 x k), B (k x 8), C/D (16 x 8)""" + lane = np.arange(32) + g, t = lane >> 2, lane & 3 + if k == 32: + i = np.arange(16) + ar = np.where((i < 4) | ((i >= 8) & (i < 12)), g[:, None], g[:, None] + 8) + ac = t[:, None] * 4 + (i & 3) + np.where(i >= 8, 16, 0) + j = np.arange(8) + br = t[:, None] * 4 + (j & 3) + np.where(j >= 4, 16, 0) + else: + i = np.arange(8) + ar = np.where((i < 2) | ((i >= 4) & (i < 6)), g[:, None], g[:, None] + 8) + ac = t[:, None] * 2 + (i & 1) + np.where(i >= 4, 8, 0) + j = np.arange(4) + br = t[:, None] * 2 + (j & 1) + np.where(j >= 2, 8, 0) + bc = np.broadcast_to(g[:, None], br.shape) + i = np.arange(4) + cr = g[:, None] + np.where(i >= 2, 8, 0)[None, :] + cc = t[:, None] * 2 + (i & 1)[None, :] + return ar, ac, br, bc, cr, cc + + def _frag(self, names, ty): + """register fragment -> (W, 32, n) float64 elements in PTX fragment order""" + W = self.T // 32 + cols = [] + for n in names: + v = self.get(n).reshape(W, 32) + if ty == 'e4m3': + for k in range(4): + cols.append(E4M3[((v >> U64(8 * k)) & U64(0xff)).astype(np.int64)]) + elif ty == 'f16': + for k in range(2): + cols.append(bits_to_f16((v >> U64(16 * k)) & U64(0xffff)).astype(F64)) + elif ty == 'f32': + cols.append(bits_to_f32(v).astype(F64)) 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) + raise NotImplementedError(f'mma fragment type {ty}') + return np.stack(cols, axis=-1) + + def op_mma(self, st, parts, a, mask): + """mma.sync.aligned.SHAPE.row.col.DTYPE.ATYPE.BTYPE.CTYPE (m16n8k32 e4m3; m16n8k16 f16)""" + shape = parts[3] + dtype, atype, btype, ctype = parts[6], parts[7], parts[8], parts[9] + if shape == 'm16n8k32' and atype == 'e4m3' and btype == 'e4m3': + k = 32 + elif shape == 'm16n8k16' and atype == 'f16' and btype == 'f16': + k = 16 + else: + raise NotImplementedError(st.text) + d_names, a_names, b_names, c_names = (self.vec(x) for x in a) + W = self.T // 32 + ar, ac, br, bc, cr, cc = self._mma_maps(k) + Av = self._frag(a_names, atype) + Bv = self._frag(b_names, btype) + Cv = self._frag(c_names, ctype) + Am = np.zeros((W, 16, k)) + Bm = np.zeros((W, k, 8)) + Cm = np.zeros((W, 16, 8)) + Am[:, ar, ac] = Av + Bm[:, br, bc] = Bv + Cm[:, cr, cc] = Cv + + def rnd(x): + return x.astype(F16 if dtype == 'f16' else F32).astype(F64) + if MMA_MODE == 'split16' and k == 32: + d1 = rnd(np.matmul(Am[:, :, :16], Bm[:, :16]) + Cm) + Dm = rnd(np.matmul(Am[:, :, 16:], Bm[:, 16:]) + d1) + else: + Dm = rnd(np.matmul(Am, Bm) + Cm) + if self.mma_hook: + pa = [self.prov.get(n) for n in a_names] + pb = [self.prov.get(n) for n in b_names] + for w in range(W): + self.mma_hook(self, st, w, Am[w], Bm[w], Cm[w], Dm[w], pa, pb) + Dv = Dm[:, cr, cc] # (W, 32, 4) + if dtype == 'f16': + bits = Dv.astype(F16).view(np.uint16).astype(U64) + outs = [bits[..., 2 * r] | (bits[..., 2 * r + 1] << U64(16)) for r in range(2)] + else: + bits = Dv.astype(F32).view(np.uint32).astype(U64) + outs = [bits[..., i] for i in range(4)] + for n, vals in zip(d_names, outs): + self.put(n, vals.reshape(self.T), mask) mma_hook = None @@ -846,6 +1154,77 @@ class Block: '{%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) + # --- textures and surfaces ---------------------------------------------------------- + def op_tex(self, st, parts, a, mask): + """tex.{base,level}.2d.v4.DTYPE.CTYPE {d0..d3}, [h, {x, y}] (, lod) + CTYPE s32: unnormalised integer texel coordinates, point fetch. CTYPE f32: normalised, bilinear + when the texture is linear (level 0 only).""" + dtype, ctype = parts[4], parts[5] + names = self.vec(a[0]) + hname, coord = a[1].strip()[1:-1].split(',', 1) + h = self.get(hname.strip()) + xs, ys = self.vec(coord.strip()) + if ctype == 's32': + x = self.typed(xs, 's32').astype(F64) + y = self.typed(ys, 's32').astype(F64) + else: + x = bits_to_f32(self.get(xs)).astype(F64) + y = bits_to_f32(self.get(ys)).astype(F64) + if parts[1] == 'level' and len(a) > 2: + lod = bits_to_f32(self.get(a[2])).astype(F64) + if np.any(lod[mask] != 0.0): + raise NotImplementedError('tex with LOD != 0 (captured textures are single-level)') + out = np.zeros((self.T, 4), F64) + for hv in np.unique(h[mask]): + tex = self.resource(hv, Texture) + sel = mask & (h == hv) + if ctype == 's32': + out[sel] = tex.fetch(x[sel].astype(np.int64), y[sel].astype(np.int64)) + else: + out[sel] = tex.sample(x[sel], y[sel]) + for i, n in enumerate(names): + bits = f32_to_bits(out[:, i]) if dtype == 'f32' else f16_to_bits(out[:, i]) + self.put(n, bits, mask) + + def op_sust(self, st, parts, a, mask): + """sust.{b,p}.2d.vN.bBITS.zero [h, {x, y}], {v0..}""" + nch = int(next(p[1:] for p in parts if re.fullmatch(r'v\d', p))) + bits = int(re.sub(r'\D', '', parts[-2])) + hname, coord = a[0].strip()[1:-1].split(',', 1) + h = self.get(hname.strip()) + xs, ys = self.vec(coord.strip()) + x = self.typed(xs, 's32') + y = self.typed(ys, 's32') + data = self.vec(a[1]) + if len(data) != nch: + raise ValueError(f'sust v{nch} with {len(data)} values') + vals = [self.get(n) & U64(MASK[bits]) for n in data] + for hv in np.unique(h[mask]): + surf = self.resource(hv, Surface) + sel = mask & (h == hv) + if parts[1] == 'p': + # formatted store: x counts texels and the f32 values are converted to the surface's format + if bits != 32 or surf.bits not in (9, 16): + raise NotImplementedError(f'sust.p of b{bits} values to a {surf.bits}-bit surface') + f = [bits_to_f32(v[sel]).astype(F64) for v in vals] + if surf.bits == 16: + # ZLUDA's lowering of the formatted store truncates (round toward zero), as the capture shows + conv = [] + for c in f: + h = c.astype(F16) + over = np.abs(h.astype(F64)) > np.abs(c) + hb = f16_to_bits(h) + conv.append(np.where(over, hb - U64(1), hb)) + else: # SNORM8: round(clamp(v, -1, 1) * 127) + conv = [np.rint(np.clip(np.nan_to_num(c), -1.0, 1.0) * 127.0).astype(np.int64).astype(U64) & U64(0xff) for c in f] + surf.write(x[sel], y[sel], conv, surf.bits) + else: + # binary store: x counts bytes + xb = x[sel] + if np.any(xb % surf.texel_bytes): + raise ValueError('sust.b at an x that is not a texel boundary') + surf.write(xb // surf.texel_bytes, y[sel], [v[sel] for v in vals], bits) + class Stmt_like: def __init__(self, st, op, args): @@ -855,7 +1234,8 @@ class Stmt_like: setattr(self, k, getattr(st, k, None)) -def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None): +def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None, resources=None): + """Run one block of KERNEL from PTX against a replay directory. Returns (block, memory, steps).""" 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() @@ -864,14 +1244,15 @@ def run_block(ptx, kernel, dump, ctaid, hooks=(), mma_hook=None, limit=None): b = Block(stmts, types, shared, mem, kernel + '_param_0', block, grid, ctaid) b.hooks = list(hooks) b.mma_hook = mma_hook + b.resources = dict(resources or {}) steps = b.run(limit) return b, mem, steps if __name__ == '__main__': + import time 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') diff --git a/kernels/tools/rrsim.py b/kernels/tools/rrsim.py new file mode 100755 index 0000000..0ccd5c8 --- /dev/null +++ b/kernels/tools/rrsim.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""rrsim: run selected thread blocks of a captured Ray Reconstruction CUDA launch on the CPU and check +every global and surface store against the GPU reference (the capture's post-launch arena and surfaces). + +USAGE + python3 kernels/tools/rrsim.py CAPTURE KERNEL_SUBSTRING FRAME --blocks X,Y[,X,Y...] + [--watch 'LINE:%reg,...' ...] [--mma exact|split16] [--watch-max N] + + CAPTURE capture directory, e.g. ~/.cache/d4r-rr-work/capE + KERNEL_SUBSTRING substring of the launch's kernel name, e.g. swin_enc1, swin_dec2, swin_enc0, swin_dec0, hkpn_output + FRAME 1-based index among the launches that match KERNEL_SUBSTRING (enc1 frame 2 = 2) + --blocks ctaid.x,ctaid.y pairs, e.g. 0,0,20,12 + --watch print the named registers each time the statement at .ptx line LINE executes + (LINE is the line in the corpus module, as in grep -n). Repeatable. + --mma MMA accumulation model (see ptxsim.py module docstring): exact (default) or split16 + --watch-max records printed per --watch (default 16) + + For each block: runtime, statements executed, then one line per compared store group: + identical / differing counts, the ulp histogram of differing values, and max ulp. + +LIBRARY + import sys; sys.path.insert(0, 'kernels/tools') + from rrsim import run_block, compare + r = run_block('~/.cache/d4r-rr-work/capE', 'swin_enc1', 2, (0, 0), + watch={1234: ['%f1']}) # watch: {ptx line: [registers]}, optional + r['regs'] final register file: name -> uint64 array (T lanes), bool for predicates + r['stores'] store log: list of (alloc_id, byte offsets per lane, raw bytes (lanes, bytes), ptx type) + r['surfaces'] [(Image, Surface)] surface writes of the block (Surface.data, Surface.written) + r['watch'] list of {line, stmt, reg, width, values, mask} records, in execution order + r['seconds'], r['steps'], r['block'] (ptxsim.Block), r['launch'], r['capture'], r['ptx'] + compare(r) -> list of per-group dicts (n, identical, differing, max ulp, ulp histogram) + +FIDELITY (details in ptxsim.py) + Exact: integer and bit ops, shfl/movmatrix/prmt, sust/tex.base point fetch, cvt, f16/f32 add/mul/sub + (one rounding per op). The MMA accumulation model and transcendental approximations may differ in the last + bit from the GPU reference; compare() reports the differences as ulp counts so they can be attributed. + Surfaces start at zero and textures are the post-launch dump (the capture has no pre-launch surface or + texture bytes). Only stores of the simulated blocks are compared; the arena is the replay's pre-launch + copy with only those blocks executed. +""" +import argparse +import os +import struct +import sys +import time +from pathlib import Path + +import numpy as np + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +sys.path.insert(0, str(HERE.parents[1] / 'engine' / 'rr')) +import ptxsim # noqa: E402 +from rr_capture import Capture # noqa: E402 + +U64 = np.uint64 +CORPUS = Path.home() / '.cache' / 'd4r-rr-corpus' +NB = {'b8': 1, 'u8': 1, 's8': 1, 'b16': 2, 'u16': 2, 's16': 2, 'f16': 2, + 'b32': 4, 'u32': 4, 's32': 4, 'f32': 4, 'b64': 8, 'u64': 8, 's64': 8, 'f64': 8} +_PTX = {} + + +def find_ptx(kernel): + """the corpus module that defines the entry KERNEL (the PREPASS entry is never selected)""" + if kernel not in _PTX: + hits = [p for p in sorted(CORPUS.glob('*.ptx')) if f'.entry {kernel}(' in p.read_text(errors='replace')] + if len(hits) != 1: + raise FileNotFoundError(f'{len(hits)} corpus modules define entry {kernel}: {hits}') + _PTX[kernel] = hits[0] + return _PTX[kernel] + + +def bind(launch): + """resource handles of a launch: handle value (read from the argument block) -> Texture / fresh Surface""" + handles, surfaces, textures = {}, [], {} + for img in launch.images: + h = struct.unpack_from(' uint64 bit patterns""" + out = np.zeros(raw.shape[0], U64) + for b in range(raw.shape[1]): + out |= raw[:, b].astype(U64) << U64(8 * b) + return out + + +def _decode(bits, ty): + if ty == 'f16': + return (bits & U64(0xffff)).astype(np.uint16).view(np.float16).astype(np.float64) + return (bits & U64(0xffffffff)).astype(np.uint32).view(np.float32).astype(np.float64) + + +def _ordered(bits, n): + """monotone integer key of an IEEE bit pattern: adjacent floats differ by 1""" + sign = (bits >> U64(n - 1)) & U64(1) + mag = (bits & U64((1 << (n - 1)) - 1)).astype(np.int64) + return np.where(sign == 1, -mag, mag) + + +def compare_elements(sim, ref, ty): + """sim/ref: uint64 bit patterns of the same elements (type ty). Returns counts and ulp statistics.""" + same = sim == ref + out = {'n': int(sim.size), 'identical': int(same.sum()), 'differing': int((~same).sum()), + 'nan_mismatch': 0, 'max_ulp': 0, 'ulp_hist': {}} + if out['differing'] == 0: + return out + diff = np.nonzero(~same)[0] + if ty in ('f16', 'f32'): + n = 16 if ty == 'f16' else 32 + with np.errstate(all='ignore'): + nan = np.isnan(_decode(sim[diff], ty)) | np.isnan(_decode(ref[diff], ty)) + out['nan_mismatch'] = int(nan.sum()) + ulp = np.abs(_ordered(sim[diff], n) - _ordered(ref[diff], n))[~nan] + else: + ulp = np.abs(sim[diff].astype(np.int64) - ref[diff].astype(np.int64)) + if ulp.size: + out['max_ulp'] = int(ulp.max()) + vals, counts = np.unique(np.minimum(ulp, 9), return_counts=True) + out['ulp_hist'] = {(f'>{8}' if v == 9 else int(v)): int(c) for v, c in zip(vals, counts)} + return out + + +def compare(result): + """Compare the stores and surface writes of a run_block result with the capture's post-launch state.""" + cap, launch, mem = result['capture'], result['launch'], result['mem'] + groups, refs, report = {}, {}, [] + for aid, off, raw, ty in mem.stores: + nb = NB[ty] + m = raw.shape[1] // nb + groups.setdefault((aid, ty), []).append((off[:, None] + np.arange(m) * nb).reshape(-1)) + for (aid, ty), plist in sorted(groups.items()): + pos = np.unique(np.concatenate(plist)) + nb = NB[ty] + idx = pos[:, None] + np.arange(nb) + if aid not in refs: + refs[aid] = np.frombuffer(launch.post_buffer(cap, aid), np.uint8) + k = mem.ids.index(aid) + st = compare_elements(le_bits(mem.data[k][idx]), le_bits(refs[aid][idx]), ty) + st['name'] = f'global alloc {aid} ({ty}, {"arena" if aid == launch.arena_alloc else "other"})' + report.append(st) + for img, s in result['surfaces']: + ref = np.fromfile(img.path, np.uint8) + if ref.size != s.data.size: + raise ValueError(f'surface arg{img.offset}: dump {ref.size} bytes, surface {s.data.size}') + ntex, nb = s.W * s.H, s.bits // 8 + written = s.written.reshape(ntex, s.C, nb).all(axis=2) + ty = {16: 'f16', 32: 'f32'}.get(s.bits, 'u8') + st = compare_elements(le_bits(s.data.reshape(ntex, s.C, nb)[written]), + le_bits(ref.reshape(ntex, s.C, nb)[written]), ty) + st['name'] = f'surface arg{img.offset} {img.width}x{img.height}x{img.channels} {ty}' + report.append(st) + return report + + +def format_group(st): + n = st['n'] + pct = 100.0 * st['identical'] / n if n else 100.0 + line = f" {st['name']}: elements {n}, identical {st['identical']} ({pct:.3f}%), differing {st['differing']}" + if st['differing']: + line += f", max ulp {st['max_ulp']}, ulp histogram {st['ulp_hist']}" + if st['nan_mismatch']: + line += f", NaN mismatches {st['nan_mismatch']}" + return line + + +def print_watch(records, limit): + for rec in records[:limit]: + lanes = np.nonzero(rec['mask'])[0][:8] + parts = [] + for ln in lanes: + v = int(rec['values'][ln]) + if rec['width'] == 16: + dec = f'{float(_decode(np.array([v], U64), "f16")[0]):g}' + elif rec['width'] == 32: + dec = f'{float(_decode(np.array([v], U64), "f32")[0]):g}' + elif rec['width'] == 1: + dec = str(bool(v)) + else: + dec = str(v) + parts.append(f'lane{ln}={v:#x}({dec})') + print(f" watch line {rec['line']} {rec['reg']} [{rec['stmt'][:60]}] active={int(rec['mask'].sum())}: " + ' '.join(parts)) + if len(records) > limit: + print(f' ... {len(records) - limit} more watch records') + + +def main(argv=None): + ap = argparse.ArgumentParser(description='run captured RR CUDA blocks on the CPU and compare with the capture') + ap.add_argument('capture') + ap.add_argument('kernel') + ap.add_argument('frame', type=int) + ap.add_argument('--blocks', required=True) + ap.add_argument('--watch', action='append', default=[]) + ap.add_argument('--mma', choices=['exact', 'split16'], default='exact') + ap.add_argument('--watch-max', type=int, default=16) + args = ap.parse_args(argv) + nums = [int(v) for v in args.blocks.split(',')] + if len(nums) % 2: + ap.error('--blocks needs X,Y pairs') + blocks = list(zip(nums[0::2], nums[1::2])) + watch = {} + for spec in args.watch: + line, names = spec.split(':', 1) + watch.setdefault(int(line), []).extend(n.strip() for n in names.split(',') if n.strip()) + ptxsim.MMA_MODE = args.mma + + cap, launch = resolve(args.capture, args.kernel, args.frame) + print(f'launch {launch.sequence} {launch.kernel} grid {launch.grid} block {launch.block} ' + f'ptx {find_ptx(launch.kernel).name} mma {args.mma}') + for bx, by in blocks: + r = run_block(cap, args.kernel, args.frame, (bx, by), watch=watch or None) + print(f'block ({bx},{by}): {r["steps"]} statements in {r["seconds"]:.2f}s') + if watch: + print_watch(r['watch'], args.watch_max) + for st in compare(r): + print(format_group(st)) + + +if __name__ == '__main__': + main() diff --git a/kernels/tools/rrswin_enc0_engine_test.cpp b/kernels/tools/rrswin_enc0_engine_test.cpp new file mode 100644 index 0000000..fcfcb55 --- /dev/null +++ b/kernels/tools/rrswin_enc0_engine_test.cpp @@ -0,0 +1,1013 @@ +// Steps 1-5 of the staged-image engine's enc0-port validation: each new primitive in +// kernels/rr/rrswin_engine.h run on the GPU with checked HIP calls and compared against +// an independently written host model. +// +// step 1 swin_slots_from_wimg -- the slot bytes directly, then through a m16n16 +// swin_gemm_tile against a random staged image (the "from wimg" part only +// reshuffles layout, the GEMM proves the (row, k) map as the k16 test's step 4 +// does); 32 rows, two m-tiles, so slot_step = n_rows is exercised +// step 2 the dcell exchange (writer, wave barrier, reader): bit-exact against the D +// cells the host model derives from the engine-tile values; a twelve-tile +// phase (n_mt = 2, n_nt = 6), also driven in the body's (2,2)-tile chunks, and +// an eight-tile phase (n_mt = 2, n_nt = 4) +// step 3 swin_stage_operand_p modes 0/1/2: the staged image bytes compared bit for +// byte against the host's own requantisation (an exhaustive u16->u8 e4m3 table +// built at startup from the documented satfinite semantics), the activation +// function at a one-ulp f16 bar per half (the k16 test's step 6 style) +// step 4 an integrated engine phase: per-column f16 C seeds for a 32-row phase in +// the column-major f16 blob form (§2), two m-tiles x six n-tiles, k32, pairs +// rounded once to f16 like the engine loop does +// step 5 the phase1->2 chain: D1 (per-column C), relu-staged (mode 1; the staged +// image is dumped and compared bit for byte), phase 2 on the staged image with +// its own per-column C, pairs rounded once to f16 +// +// The host models are written from the contract (kernels/rr/rrswin_enc0_engine.md) and +// the translated body's documented layouts, not from the device code: the enc0 weight +// waddr, the fragment->(row, k) map, the dcell exchange's lds addresses and the +// satfinite-e4m3 rounding. They share no code with the device side. +// +// Build and run: +// hipcc --offload-arch=gfx1201 -DD4R_TEX_FP8 -DD4R_WMMA_LAYOUT=12 -O2 \ +// -o /tmp/enc0_engine_test kernels/tools/rrswin_enc0_engine_test.cpp \ +// && /tmp/enc0_engine_test +#include +#include +#include +#include +#include +#include +#include + +#include "../rr/rrswin_engine.h" + +#define CHECK(expr) \ + do \ + { \ + hipError_t e_ = (expr); \ + if (e_ != hipSuccess) \ + { \ + std::printf("HIP error at %s:%d: %s -> %s\n", __FILE__, __LINE__, #expr, \ + hipGetErrorString(e_)); \ + std::exit(2); \ + } \ + } while (0) + +// --------------------------------------------------------------------------- +// independent host models ----------------------------------------------------------- + +// the enc0 weight-image layout, from the contract (rrswin_enc0.hip's "byte j of a +// lane's 16-byte weight tile holds (k, n) at addr(n, k) = ...") +static int waddr_host(int k, int n) +{ + return 512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1) + + 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1); +} + +// exact e4m3 decode: full f32 (subnormals and the NaN included), from the format's +// definition (3 mantissa bits, bias 7, subnormals 2^-9 apart) +static double dec(uint8_t c) +{ + const uint32_t sign = c & 0x80u, e = (c >> 3) & 0xfu, m = c & 7u; + const double mag = e == 0 ? (m / 8.0) * std::pow(2.0, -6.0) + : (1.0 + m / 8.0) * std::pow(2.0, (int)e - 7.0); + return sign ? -mag : mag; +} + +// f16 bits <-> value, host-side +__attribute__((device, host)) static inline _Float16 h_of(uint16_t b) +{ + _Float16 h; + std::memcpy(&h, &b, sizeof(h)); + return h; +} +__attribute__((device, host)) static inline uint16_t b_of(_Float16 h) +{ + uint16_t b; + std::memcpy(&b, &h, sizeof(b)); + return b; +} + +// --------------------------------------------------------------------------- +// the RNE-satfinite e4m3 rounding table --------------------------------------------- + +// For f16 bits: the exact f64 value, then, at every binary step up 2^-24..2^25, bits +// 0..7: bit 7 into code (the merged exponent covers normal and subnormal), the bit +// below the cut -> a half, anything below -> sticky. At the top of the walk the code +// saturates at 0x7e; at the bottom subnormals emit 2^-9-spaced codes. Overflow and +// NaN are the doc's two cases (0x7e and the whole-tab 0x7f). +static uint8_t tab[0x10000]; +static void tab_build() +{ + for (uint32_t b = 0; b < 0x10000u; ++b) + { + const uint16_t bits = (uint16_t)b; + const uint32_t mag = bits & 0x7fffu, sign = (bits >> 8) & 0x80u; + if (mag > 0x7c00u) + { + tab[b] = (uint8_t)(sign | 0x7fu); // NaN keeps its sign + continue; + } + if (mag >= 0x7c00u) + { + tab[b] = (uint8_t)(sign | 0x7eu); // inf: out of range -> max finite + continue; + } + if (mag == 0) + { + tab[b] = (uint8_t)sign; // +/-0 preserved + continue; + } + const int e = (int)(bits >> 10) & 0x1f, m = (int)(bits & 0x3ffu); + // the exact f64 of the f16 magnitude + double v; + if (e > 0) + v = (1.0 + m / 1024.0) * std::pow(2.0, e - 15.0); + else + v = (m / 1024.0) * std::pow(2.0, -14.0); + // the walk: the leading power of two (2^9 .. 2^-6) fixes the code's merged + // exponent, the three steps below it pick the mantissa bits, and the first + // dropped power of two is the rounding bit, the rest sticky. Under the whole + // subnormal grid (v < 2^-9) the same walk on the 2^-9-spaced subnormal codes + // decides, rounding on 2^-10. + uint32_t code = 0; + uint32_t round = 0, sticky = 0; + int base = 0; + int hit = 0; + for (int q = 9; q >= -6; --q) // the merged exponent, q >= -6 + { + const double step = std::pow(2.0, q); + if (v >= step) + { + code = (uint32_t)((q + 7) << 3); + base = q; + v -= step; + hit = 1; + break; + } + } + if (hit) + { + for (int q = base - 1; q >= base - 3; --q) + { + const double step = std::pow(2.0, q); + if (v >= step) + { + v -= step; + code |= 1u << (q - (base - 3)); + } + } + // the first dropped power of two is the rounding bit, the rest sticky + const double halfstep = std::pow(2.0, (double)(base - 4)); + if (v >= halfstep) + round = 1; + if (v > halfstep) + sticky = 1; + } + else + { + // subnormal: codes m * 2^-9; walk the three steps 2^-7, 2^-8, 2^-9 (m bits + // 2,1,0), then the half (2^-10) and sticky + for (int q = -7; q >= -9; --q) + { + const double step = std::pow(2.0, q); + if (v >= step) + { + v -= step; + code |= 1u << (q + 9); + } + } + const double halfstep = std::pow(2.0, -10.0); + if (v >= halfstep) + round = 1; + if (v > halfstep) + sticky = 1; + } + if (round && (sticky || (code & 1u))) + ++code; // RNE: tie -> even + if (code > 0x7eu) + code = 0x7eu; + tab[b] = (uint8_t)(sign | code); + } +} + +// RNE-satfinite e4m3 from f16 bits, through the table (the device's +// cvt.rn.satfinite.e4m3x2.f16x2, low half -> low byte) +static uint8_t requant_code(uint16_t bits) { return tab[bits]; } +static uint8_t requant_code_bytes(double v) { return requant_code(b_of((_Float16)v)); } + +// The requantisation the device performs (dec->f16->satfinite->dec, the k16 test's +// requant with the table behind it) +static double requant(double v) { return dec(requant_code_bytes(v)); } + +// relu_keep_nan_halves per half, host model: negative or -0 becomes +0, NaN kept +static uint16_t relu_half(uint16_t bits) +{ + return (bits & 0x8000u) && (bits & 0x7fffu) <= 0x7c00u ? 0u : bits; +} + +// The clamped cubic activation per half, host model from the spec formula in +// rrswin_enc0.hip's comment, f16 throughout (each intermediate cast back), one-ulp bar +// for the packed-f16x2 device ops +static uint16_t act_model(uint16_t bits) +{ + const _Float16 x = h_of(bits); + const _Float16 two = h_of(0x4000), slope = h_of(0x2d30), mid = h_of(0x3698), half0 = h_of(0x3800); + _Float16 y = x > -two ? x : -two; + y = y < two ? y : two; + y = y > -y ? y : -y; + const _Float16 g = (_Float16)(half0 + y * (mid - slope * y)); + return b_of((_Float16)(x * g)); +} + +// --------------------------------------------------------------------------- +// step 1 kernels: the slots (dumped for the byte check), and one m16n16 GEMM per +// m-tile against a random staged image (tile (mt, nt) of a 32-row x 32-token block) +__global__ void step1_kernel(const uint8_t* wimg, uint8_t* slots_dump, const uint8_t* ximg, + float* out) +{ + __shared__ uint8_t slots[2048]; + __shared__ uint8_t img[32 * 32]; + for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x) + img[i] = ximg[i]; + __syncthreads(); + rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 0, 32); + __syncthreads(); + for (int i = threadIdx.x; i < 2048; i += blockDim.x) + slots_dump[i] = slots[i]; + for (int mt = 0; mt < 2; ++mt) + { + f8v c = {0, 0, 0, 0, 0, 0, 0, 0}; + c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 0); + for (int i = 0; i < 8; ++i) + out[256 * mt + i * 32 + threadIdx.x] = c[i]; + } +} + +// --------------------------------------------------------------------------- +// step 2 kernel: the dcell exchange. The value source is random f32 D values, shaped +// as the lane's own [tiles][8] VGPRs (pair j covers VGPRs 2j, 2j + 1). The kernel packs +// pairs per tile, writes, barriers, reads, and dumps d flat as u32 words per lane. +// `chunked`: 0 = the whole twelve-tile buffer in one call, 1 = the body's (2,2)-tile +// chunk driver over the same [12][8] source (d's [6][2][4] accumulated through three +// chunks), 2 = the eight-tile phase (its own [8][8] source buffer). +__global__ void step2_kernel(const float* vin, uint32_t* out, int chunked) +{ + __shared__ __attribute__((aligned(1024))) unsigned char sraw[12 * 128 * 4]; + rrswin::lds_u32* const lds = (rrswin::lds_u32*)(uintptr_t)sraw; + const uint32_t l = threadIdx.x; + const uint32_t ntiles = chunked ? 12u : (uint32_t)(chunked == 0 ? 12u : 8u); + (void)ntiles; + // pack pairs per tile + if (chunked == 0 || chunked == 1) + { + const int tiles = chunked == 0 ? 12 : 12; + uint32_t pairs[12 * 4]; + for (int tl = 0; tl < tiles; ++tl) + for (int j = 0; j < 4; ++j) + { + const _Float16 lo = (_Float16)vin[((size_t)l * tiles + tl) * 8 + 2 * j]; + const _Float16 hi = (_Float16)vin[((size_t)l * tiles + tl) * 8 + 2 * j + 1]; + uint16_t bl, bh; + std::memcpy(&bl, &lo, sizeof(bl)); + std::memcpy(&bh, &hi, sizeof(bh)); + pairs[tl * 4 + j] = (uint32_t)bl | ((uint32_t)bh << 16); + } + for (int i = 0; i < 12 * 128; ++i) + lds[i] = 0u; + __builtin_amdgcn_wave_barrier(); + if (chunked == 0) + { + rrswin::swin_dcell_exchange_write(lds, pairs, 2, 6); + __builtin_amdgcn_wave_barrier(); + uint32_t d[6 * 2 * 4]; + rrswin::swin_dcell_exchange_read(lds, d, 6, 4, 6); + for (int i = 0; i < 48; ++i) + out[i * 32 + threadIdx.x] = d[i]; + } + else + { + // the body's chunk shape: value source [2][6][4], d + 16c per chunk + uint32_t xp[16], d[6 * 2 * 4]; + for (int c = 0; c < 3; ++c) + { + for (int mt = 0; mt < 2; ++mt) + for (int l2 = 0; l2 < 2; ++l2) + for (int p = 0; p < 4; ++p) + xp[(mt * 2 + l2) * 4 + p] = pairs[(mt * 6 + 2 * c + l2) * 4 + p]; + rrswin::swin_dcell_exchange_write(lds, xp, 2, 2); + __builtin_amdgcn_wave_barrier(); + rrswin::swin_dcell_exchange_read(lds, d + 16 * c, 2, 4, 2); + __builtin_amdgcn_wave_barrier(); + } + for (int i = 0; i < 48; ++i) + out[i * 32 + threadIdx.x] = d[i]; + } + } + else + { + // eight-tile phase: [2][4][4] pairs -> d[4][2][4]; the value source is laid out + // as [8][8] per lane, at the same (lane, tile, vgpr) offsets as the 12-tile one + uint32_t pairs[8 * 4]; + for (int tl = 0; tl < 8; ++tl) + for (int j = 0; j < 4; ++j) + { + const _Float16 lo = (_Float16)vin[((size_t)l * 8 + tl) * 8 + 2 * j]; + const _Float16 hi = (_Float16)vin[((size_t)l * 8 + tl) * 8 + 2 * j + 1]; + uint16_t bl, bh; + std::memcpy(&bl, &lo, sizeof(bl)); + std::memcpy(&bh, &hi, sizeof(bh)); + pairs[tl * 4 + j] = (uint32_t)bl | ((uint32_t)bh << 16); + } + for (int i = 0; i < 8 * 128; ++i) + lds[i] = 0u; + __builtin_amdgcn_wave_barrier(); + rrswin::swin_dcell_exchange_write(lds, pairs, 2, 4); + __builtin_amdgcn_wave_barrier(); + uint32_t d[4 * 2 * 4]; + rrswin::swin_dcell_exchange_read(lds, d, 4, 4, 4); + for (int i = 0; i < 32; ++i) + out[i * 32 + threadIdx.x] = d[i]; + } +} + +// --------------------------------------------------------------------------- +// step 3 kernel: swin_stage_operand_p's three modes into three images. +__global__ void step3_kernel(const uint32_t* pairs, uint8_t* img0, uint8_t* img1, uint8_t* img2) +{ + __shared__ uint8_t s0[32 * 32], s1[32 * 32], s2[32 * 32]; + for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x) + { + s0[i] = 0; + s1[i] = 0; + s2[i] = 0; + } + __syncthreads(); + const uint32_t* mine = pairs + threadIdx.x * 32; + for (int tile = 0; tile < 2; ++tile) + { + const uint32_t* p = mine + tile * 4; + rrswin::swin_stage_operand_p((lds_u8*)s0, 32, 16 * tile, 0, p, 0); + rrswin::swin_stage_operand_p((lds_u8*)s1, 32, 16 * tile, 0, p, 1); + rrswin::swin_stage_operand_p((lds_u8*)s2, 32, 16 * tile, 0, p, 2); + } + __syncthreads(); + for (int i = threadIdx.x; i < 32 * 32; i += blockDim.x) + { + img0[i] = s0[i]; + img1[i] = s1[i]; + img2[i] = s2[i]; + } +} + +// --------------------------------------------------------------------------- +// step 4 kernel: the integrated engine phase. Per-column f16 C seeds at w+5120, two +// m-tiles x six n-tiles, k32, each tile's pairs rounded once to f16 (enc0_pack). +__global__ void step4_kernel(const uint8_t* wimg, uint8_t* slots_dump, const uint8_t* ximg, + const uint32_t* cseed, uint32_t* out) +{ + __shared__ uint8_t slots[2048]; + __shared__ uint8_t img[96 * 32]; + for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x) + img[i] = ximg[i]; + __syncthreads(); + rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 0, 32); + __syncthreads(); + for (int i = threadIdx.x; i < 2048; i += blockDim.x) + slots_dump[i] = slots[i]; + const uint32_t h = threadIdx.x >> 4; + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + { + // the per-column C seed: c[i] = f32(f16 at w + 2*(i + 8h + 16mt)), §2 + f8v c; + for (int i = 0; i < 8; ++i) + { + // the per-column blob's column 2*(i + 8h) + 32mt sits at byte 2*(...): + // two adjacent f16, in one u32 when the byte offset is even + const uint32_t w32 = cseed[(2 * (i + 8 * h + 16 * mt)) >> 2]; + const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16 + : w32 & 0xffffu); + c[i] = (float)h_of(bits); + } + c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 16 * nt); + uint32_t pairs[4]; + for (int p = 0; p < 4; ++p) + { + const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1]; + uint16_t bl, bh; + std::memcpy(&bl, &lo, sizeof(bl)); + std::memcpy(&bh, &hi, sizeof(bh)); + pairs[p] = (uint32_t)bl | ((uint32_t)bh << 16); + } + for (int p = 0; p < 4; ++p) + out[((mt * 6 + nt) * 4 + p) * 32 + threadIdx.x] = pairs[p]; + } +} + +// --------------------------------------------------------------------------- +// step 5 kernel: the phase1->2 chain. Phase 1: per-column C at w+2048, weights at +// w+4096; its pairs relu-staged (mode 1) into a fresh image; phase 2: per-column C at +// w+5120, weights at w+6144, on the staged image. The C seeds are blobs this test +// generates at w+2048 (phase 1) and w+5120 (phase 2), read column-major (§2). +__global__ void step5_kernel(const uint8_t* wimg, const uint8_t* ximg, const uint32_t* cseed1, + const uint32_t* cseed2, uint8_t* staged, uint32_t* out) +{ + __shared__ uint8_t slots[2048]; + __shared__ uint8_t img[96 * 32]; + __shared__ uint8_t img2[96 * 32]; + for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x) + img[i] = ximg[i]; + for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x) + img2[i] = 0; + __syncthreads(); + rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 4096, 32); + __syncthreads(); + uint32_t pairs1[2][6][4]; + const uint32_t h = threadIdx.x >> 4; + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + { + f8v c; + for (int i = 0; i < 8; ++i) + { + const uint32_t w32 = cseed1[(2 * (i + 8 * h + 16 * mt)) >> 2]; + const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16 + : w32 & 0xffffu); + c[i] = (float)h_of(bits); + } + c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img, 32, 16 * nt); + for (int p = 0; p < 4; ++p) + { + const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1]; + uint16_t bl, bh; + std::memcpy(&bl, &lo, sizeof(bl)); + std::memcpy(&bh, &hi, sizeof(bh)); + pairs1[mt][nt][p] = (uint32_t)bl | ((uint32_t)bh << 16); + } + } + __syncthreads(); + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + rrswin::swin_stage_operand_p((lds_u8*)img2, 32, 16 * nt, mt, pairs1[mt][nt], 1); + __syncthreads(); + for (int i = threadIdx.x; i < 96 * 32; i += blockDim.x) + staged[i] = img2[i]; + __syncthreads(); + rrswin::swin_slots_from_wimg((lds_u8*)slots, wimg, 6144, 32); + __syncthreads(); + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + { + f8v c; + for (int i = 0; i < 8; ++i) + { + const uint32_t w32 = cseed2[(2 * (i + 8 * h + 16 * mt)) >> 2]; + const uint16_t bits = (uint16_t)(((2 * (i + 8 * h + 16 * mt)) & 2) ? w32 >> 16 + : w32 & 0xffffu); + c[i] = (float)h_of(bits); + } + c = rrswin::swin_gemm_tile(c, (const lds_u8*)slots, 16 * mt, 32, (const lds_u8*)img2, 32, 16 * nt); + uint32_t pairs[4]; + for (int p = 0; p < 4; ++p) + { + const _Float16 lo = (_Float16)c[2 * p], hi = (_Float16)c[2 * p + 1]; + uint16_t bl, bh; + std::memcpy(&bl, &lo, sizeof(bl)); + std::memcpy(&bh, &hi, sizeof(bh)); + pairs[p] = (uint32_t)bl | ((uint32_t)bh << 16); + } + for (int p = 0; p < 4; ++p) + out[((mt * 6 + nt) * 4 + p) * 32 + threadIdx.x] = pairs[p]; + } +} + +// pair q's halves are VGPRs 2q, 2q + 1; VGPR v of a lane's half h is row v + 8h. +static int i_vgpr(int q, int half, int h) { return 2 * q + half + 8 * h; } + +// --------------------------------------------------------------------------- +// main --------------------------------------------------------------------------- + +// well-spread e4m3 codes (the k16 test's next_code: exponents 0..7 and mantissas, no +// NaN codes), so a wrong (row, column) mapping shows up as a real mismatch +static uint32_t lcg_main = 0x12345678u; +static uint8_t next_code() +{ + lcg_main = lcg_main * 1664525u + 1013904223u; + uint32_t mag = (lcg_main >> 13) & 0x3fu; // exponents 0..7 and mantissas + if (mag == 0) + mag = 0x21; + const uint32_t sign = (lcg_main >> 7) & 1u; + return (uint8_t)((sign << 7) | mag); +} + +int main() +{ + tab_build(); + + // ---------------------------------------------------------------- step 1 + { + std::vector wimg(4096, 0); + for (auto& v : wimg) + v = next_code(); + std::vector ximg(32 * 32, 0); + for (auto& v : ximg) + v = next_code(); + + uint8_t *dwimg, *dslots, *dximg; + float* dout; + CHECK(hipMalloc(&dwimg, wimg.size())); + CHECK(hipMalloc(&dslots, 2048)); + CHECK(hipMalloc(&dximg, ximg.size())); + CHECK(hipMalloc(&dout, 512 * sizeof(float))); + CHECK(hipMemset(dslots, 0, 2048)); + CHECK(hipMemcpy(dwimg, wimg.data(), wimg.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(dximg, ximg.data(), ximg.size(), hipMemcpyHostToDevice)); + step1_kernel<<<1, 32>>>(dwimg, dslots, dximg, dout); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + std::vector slots(2048, 0); + CHECK(hipMemcpy(slots.data(), dslots, slots.size(), hipMemcpyDeviceToHost)); + std::vector got(512); + CHECK(hipMemcpy(got.data(), dout, got.size() * sizeof(float), hipMemcpyDeviceToHost)); + + // 1a: every slot byte = img[waddr(k, n)] + int bad = 0, checked = 0; + for (int bank = 0; bank < 64; ++bank) + { + // the slot layout is n-indexed: bank n holds k 0..15 (byte j = element + // (j, n), one full waddr each — the image's records do NOT store + // consecutive k contiguously) and bank 32 + n holds k 16..31 + const int n = bank & 31; + const int k0 = bank < 32 ? 0 : 16; + for (int j = 0; j < 16; ++j) + { + const uint8_t have = slots[(size_t)bank * 16 + j]; + const uint8_t want = wimg[waddr_host(k0 + j, n)]; + ++checked; + if (have != want) + { + if (bad < 4) + std::printf(" slot[h%d n%2d j%2d] = 0x%02x, expected 0x%02x\n", + bank >= 32, n, j, have, want); + ++bad; + } + } + } + std::printf("step 1a slot bytes vs waddr: checked %d, mismatched %d\n", checked, bad); + + // 1b: the GEMM through the slots, against the double-precision m16n16 model + // the device mma's accumulator layout: lane l, VGPR i holds D row i + 8h, + // column l & 15; the model's ref is the contract's own double sum for D[wr][tk] + int gbad = 0, gchecked = 0; + for (int mt = 0; mt < 2; ++mt) + for (int l = 0; l < 32; ++l) + { + const int hh = l >> 4, col = l & 15; + for (int i = 0; i < 8; ++i) + { + const int wr = 16 * mt + i + 8 * hh; + double ref = 0.0; + for (int k = 0; k < 32; ++k) + ref += dec(wimg[waddr_host(k, wr)]) * dec(ximg[(size_t)col * 32 + k]); + const float g = got[256 * mt + i * 32 + l]; + ++gchecked; + if (std::fabs((double)g - ref) > 2e-3 * (1.0 + std::fabs(ref))) + { + if (gbad < 4) + std::printf(" step1 GEMM tile %d lane %2d vgpr %d: got %+.6f ref %+.6f\n", + mt, l, i, g, ref); + ++gbad; + } + } + } + std::printf("step 1b GEMM through the slots: checked %d, mismatched %d\n", gchecked, gbad); + CHECK(hipFree(dwimg)); + CHECK(hipFree(dslots)); + CHECK(hipFree(dximg)); + CHECK(hipFree(dout)); + if (bad || gbad) + std::exit(1); + } + + // ---------------------------------------------------------------- step 2 + { + // 12-tile phase: per-lane [12][8] random f32 D values; 8-tile: [8][8] + std::vector v12(32 * 12 * 8), v8(32 * 8 * 8); + uint32_t lcg = 0x0f0f0f0fu; + auto nextr = [&lcg]() { + lcg = lcg * 1664525u + 1013904223u; + return (float)((int)(lcg >> 8) - 8388608) * (3.0f / 8388608.0f); + }; + for (auto& v : v12) + v = nextr() * 60.0f; + for (auto& v : v8) + v = nextr() * 60.0f; + float *dv, *dv8; + uint32_t* dout; + CHECK(hipMalloc(&dv, v12.size() * 4)); + CHECK(hipMalloc(&dv8, v8.size() * 4)); + CHECK(hipMalloc(&dout, 48 * 32 * sizeof(uint32_t))); + CHECK(hipMemcpy(dv, v12.data(), v12.size() * 4, hipMemcpyHostToDevice)); + CHECK(hipMemcpy(dv8, v8.data(), v8.size() * 4, hipMemcpyHostToDevice)); + CHECK(hipMemset(dout, 0, 48 * 32 * sizeof(uint32_t))); + + for (int pass = 0; pass < 3; ++pass) + { + const int chunked = pass == 0 ? 0 : (pass == 1 ? 1 : 2); + step2_kernel<<<1, 32>>>(chunked == 2 ? dv8 : dv, dout, chunked); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + std::vector got(48 * 32); + CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost)); + + const int tiles = chunked == 2 ? 8 : 12; + const float* src = chunked == 2 ? v8.data() : v12.data(); + const int n_nt = chunked == 2 ? 4 : 6; + int bad = 0, checked = 0; + for (int lane = 0; lane < 32; ++lane) + { + const int g = (lane >> 2) & 7, t = lane & 3; + for (int mi = 0; mi < tiles / 2; ++mi) // n_mi + for (int rh = 0; rh < 2; ++rh) + for (int nt = 0; nt < 4; ++nt) // n_nt4 + { + // the wanted cell is D[wr = 8nt + 2t + cb][tk = 16mi + 8rh + g]: + // tile (nt>>1, mi), lane 8rh + g + 16*(nt&1), VGPRs 2t, 2t+1 + const int src_lane = 8 * rh + g + 16 * (nt & 1); + const int src_tile = (nt >> 1) * n_nt + mi; + const float vlo = src[((size_t)src_lane * tiles + src_tile) * 8 + 2 * t]; + const float vhi = src[((size_t)src_lane * tiles + src_tile) * 8 + 2 * t + 1]; + const uint32_t want = (uint32_t)b_of((_Float16)vlo) + | ((uint32_t)b_of((_Float16)vhi) << 16); + const uint32_t have = got[((mi * 2 + rh) * 4 + nt) * 32 + lane]; + ++checked; + if (have != want) + { + if (bad < 4) + std::printf(" dcell lane %2d mi %d rh %d nt %d: got %08x want %08x\n", + lane, mi, rh, nt, have, want); + ++bad; + } + } + } + std::printf("step 2 %s exchange: checked %d, mismatched %d\n", + chunked == 0 ? "twelve-tile whole-buffer" + : (chunked == 1 ? "twelve-tile (2,2)-chunk" : "eight-tile"), + checked, bad); + if (bad) + std::exit(1); + } + CHECK(hipFree(dv)); + CHECK(hipFree(dv8)); + CHECK(hipFree(dout)); + } + + // ---------------------------------------------------------------- step 3 + { + // every f16 of the whole 16-bit space (all 65536, incl. -0.0, negatives, NaN + // halves, values beyond +-2 and subnormals), shuffled into the 64 pairs per lane + // x 4 tiles the kernel stages (mode 0/1's bar is bit exact, so the whole space) + uint16_t allbits[65536]; + for (uint32_t b = 0; b < 0x10000u; ++b) + allbits[b] = (uint16_t)b; + uint32_t lcg = 0x00c0ffeeu; + auto nextr = [&lcg]() { + lcg = lcg * 1664525u + 1013904223u; + return lcg >> 16; + }; + for (uint32_t b = 0xffffu; b >= 1; --b) + { + const uint32_t j = nextr() % (b + 1); + const uint16_t t = allbits[b]; + allbits[b] = allbits[j]; + allbits[j] = t; + } + std::vector pairs(32 * 32); + for (size_t i = 0; i < pairs.size(); ++i) + { + const uint16_t lo = allbits[2 * i]; + const uint16_t hi = allbits[2 * i + 1]; + pairs[i] = (uint32_t)lo | ((uint32_t)hi << 16); + } + uint32_t* dp; + uint8_t *di0, *di1, *di2; + CHECK(hipMalloc(&dp, pairs.size() * 4)); + CHECK(hipMalloc(&di0, 32 * 32)); + CHECK(hipMalloc(&di1, 32 * 32)); + CHECK(hipMalloc(&di2, 32 * 32)); + CHECK(hipMemcpy(dp, pairs.data(), pairs.size() * 4, hipMemcpyHostToDevice)); + CHECK(hipMemset(di0, 0, 32 * 32)); + CHECK(hipMemset(di1, 0, 32 * 32)); + CHECK(hipMemset(di2, 0, 32 * 32)); + step3_kernel<<<1, 32>>>(dp, di0, di1, di2); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + std::vector i0(32 * 32), i1(32 * 32), i2(32 * 32); + CHECK(hipMemcpy(i0.data(), di0, i0.size(), hipMemcpyDeviceToHost)); + CHECK(hipMemcpy(i1.data(), di1, i1.size(), hipMemcpyDeviceToHost)); + CHECK(hipMemcpy(i2.data(), di2, i2.size(), hipMemcpyDeviceToHost)); + + auto within_ulp = [](uint16_t a, uint16_t b) { + // compare as f16 by ordering the sign-magnitude representation + const auto key = [](uint16_t v) -> int { + return (v & 0x8000) ? (int)(0x8000 - (v & 0x7fff)) : (int)(v & 0x7fff); + }; + const int d = key(a) - key(b); + return d >= -1 && d <= 1; + }; + + // mode 0 and mode 1: bit exact through the table + int bad0 = 0, bad1 = 0, checked = 0; + for (int tile = 0; tile < 2; ++tile) + for (int lane = 0; lane < 32; ++lane) + { + const int h = lane >> 4, col = lane & 15; + const uint32_t* p = pairs.data() + (size_t)(lane * 32) + tile * 4; + for (int q = 0; q < 4; ++q) + for (int half = 0; half < 2; ++half) + { + const uint16_t bits = (uint16_t)(half ? p[q] >> 16 : p[q] & 0xffffu); + const uint16_t re = relu_half(bits); + const uint8_t want0 = requant_code(bits); + const uint8_t want1 = requant_code(re); + // the staged byte sits at image row = row_base + (lane & 15), + // byte = 16*chunk + 8h + (VGPR = 2q + half) + const int rowi = 16 * tile + col; + const size_t off = (size_t)rowi * 32 + i_vgpr(q, half, h); + ++checked; + if (i0[off] != want0) + { + if (bad0 < 4) + std::printf(" mode0 img row %2d byte %2d (tile %d lane %2d q %d h%d):" + " got 0x%02x want 0x%02x\n", + rowi, (int)off % 32, tile, lane, q, half, + i0[off], want0); + ++bad0; + } + ++checked; + if (i1[off] != want1) + { + if (bad1 < 4) + std::printf(" mode1 img row %2d byte %2d (tile %d lane %2d q %d h%d):" + " got 0x%02x want 0x%02x\n", + rowi, (int)off % 32, tile, lane, q, half, + i1[off], want1); + ++bad1; + } + } + } + std::printf("step 3a stage modes 0/1 vs the table: checked %d, mode0 %d, mode1 %d mismatched\n", + checked, bad0, bad1); + + // mode 2: the clamped cubic then encode, one ulp per half (the device applies + // the formula with packed f16x2 ops, the scalar _Float16 model can land one ulp + // away -- the k16 test's step 6 bar; a logic error moves whole clamps) + int bad2 = 0, checked2 = 0; + for (int tile = 0; tile < 2; ++tile) + for (int lane = 0; lane < 32; ++lane) + { + const int h = lane >> 4, col = lane & 15; + const uint32_t* p = pairs.data() + (size_t)(lane * 32) + tile * 4; + for (int q = 0; q < 4; ++q) + for (int half = 0; half < 2; ++half) + { + const uint16_t bits = (uint16_t)(half ? p[q] >> 16 : p[q] & 0xffffu); + const uint16_t acted = act_model(bits); + const uint8_t want2 = requant_code(acted); + const int rowi = 16 * tile + col; + const size_t off = (size_t)rowi * 32 + i_vgpr(q, half, h); + ++checked2; + if (!within_ulp(i2[off], want2)) + { + if (bad2 < 4) + std::printf(" mode2 img row %2d byte %2d: got 0x%02x want 0x%02x (act %04x)\n", + rowi, (int)off % 32, i2[off], want2, acted); + ++bad2; + } + } + } + std::printf("step 3b stage mode 2 (1 ulp): checked %d, mismatched %d\n", checked2, bad2); + CHECK(hipFree(dp)); + CHECK(hipFree(di0)); + CHECK(hipFree(di1)); + CHECK(hipFree(di2)); + if (bad0 || bad1 || bad2) + std::exit(1); + } + // ---------------------------------------------------------------- step 4 + { + // random weights in the waddr layout (32 rows x 32 k), random staged image + // (96 rows x 32 bytes), per-column f16 C seeds (32 columns, column c at byte 2c) + std::vector wimg(4096, 0), ximg(96 * 32, 0); + for (int k = 0; k < 32; ++k) + for (int n = 0; n < 32; ++n) + wimg[waddr_host(k, n)] = next_code(); + for (auto& v : ximg) + v = next_code(); + std::vector cseed(64, 0); + for (int c = 0; c < 32; ++c) + { + const uint16_t b = (uint16_t)(0x3c00u - (next_code() & 0xffu)); // +-small f16 + cseed[2 * c] = (uint8_t)(b & 0xffu); + cseed[2 * c + 1] = (uint8_t)(b >> 8); + } + uint8_t *dw, *dx, *dsl; + uint32_t* dseed; + uint32_t* dout; + CHECK(hipMalloc(&dw, wimg.size())); + CHECK(hipMalloc(&dx, ximg.size())); + CHECK(hipMalloc(&dseed, cseed.size())); + CHECK(hipMalloc(&dsl, 2048)); + CHECK(hipMalloc(&dout, 48 * 4 * 32 * 4)); + CHECK(hipMemcpy(dw, wimg.data(), wimg.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(dx, ximg.data(), ximg.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(dseed, cseed.data(), cseed.size(), hipMemcpyHostToDevice)); + CHECK(hipMemset(dout, 0, 48 * 4 * 32 * 4)); + step4_kernel<<<1, 32>>>(dw, dsl, dx, dseed, dout); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + std::vector got(48 * 4 * 32); + CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost)); + + int bad = 0, checked = 0; + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + for (int lane = 0; lane < 32; ++lane) + { + const int h = lane >> 4, col = lane & 15; + for (int p = 0; p < 4; ++p) + for (int half = 0; half < 2; ++half) + { + // the engine cell D[wr = 16mt + 2p + half + 8h][tk = 16nt + col]: + // double sum over k of dec(w[waddr(k, wr)]) * dec(img[tk*32 + k]) + // + f32(C seed f16), rounded once to f16 + const int wr = 16 * mt + 2 * p + half + 8 * h; + const int tk = 16 * nt + col; + double ref = 0.0; + for (int k = 0; k < 32; ++k) + ref += dec(wimg[waddr_host(k, wr)]) * dec(ximg[(size_t)tk * 32 + k]); + const uint16_t seedbits = + (uint16_t)(cseed[2 * wr] | (cseed[2 * wr + 1] << 8)); + ref += (double)h_of(seedbits); + const uint16_t want = b_of((_Float16)ref); + const uint16_t bits = (uint16_t)(half ? got[((mt * 6 + nt) * 4 + p) * 32 + lane] >> 16 + : got[((mt * 6 + nt) * 4 + p) * 32 + lane] & 0xffffu); + ++checked; + if (bits != want) + { + if (bad < 6) + std::printf(" D[wr %2d][tk %2d] (mt %d nt %d lane %2d p %d h%d):" + " got %04x want %04x\n", + wr, tk, mt, nt, lane, p, half, bits, want); + ++bad; + } + } + } + std::printf("step 4 integrated phase with C seeds: checked %d, mismatched %d\n", + checked, bad); + CHECK(hipFree(dw)); + CHECK(hipFree(dx)); + CHECK(hipFree(dseed)); + CHECK(hipFree(dsl)); + CHECK(hipFree(dout)); + if (bad) + std::exit(1); + } + + // ---------------------------------------------------------------- step 5 + { + // phase 1: per-column C at w+2048, weights at w+4096; phase 2: per-column C at + // w+5120, weights at w+6144, off the relu-staged image + std::vector wimg(10240, 0), ximg(96 * 32, 0); + for (auto& v : ximg) + v = next_code(); + for (int k = 0; k < 32; ++k) + for (int n = 0; n < 32; ++n) + wimg[4096 + waddr_host(k, n)] = next_code(); + for (int k = 0; k < 32; ++k) + for (int n = 0; n < 32; ++n) + wimg[6144 + waddr_host(k, n)] = next_code(); + std::vector cs1(64, 0), cs2(64, 0); + for (int c = 0; c < 32; ++c) + { + const uint16_t b = (uint16_t)(0x3c00u - (next_code() & 0xffu)); + cs1[2 * c] = (uint8_t)(b & 0xffu); + cs1[2 * c + 1] = (uint8_t)(b >> 8); + const uint16_t b2 = (uint16_t)(0x3c00u - (next_code() & 0xffu)); + cs2[2 * c] = (uint8_t)(b2 & 0xffu); + cs2[2 * c + 1] = (uint8_t)(b2 >> 8); + } + uint8_t *dw, *dx, *dst; + uint32_t *ds1, *ds2; + uint32_t* dout; + CHECK(hipMalloc(&dw, wimg.size())); + CHECK(hipMalloc(&dx, ximg.size())); + CHECK(hipMalloc(&ds1, cs1.size())); + CHECK(hipMalloc(&ds2, cs2.size())); + CHECK(hipMalloc(&dst, 96 * 32)); + CHECK(hipMalloc(&dout, 48 * 4 * 32 * 4)); + CHECK(hipMemcpy(dw, wimg.data(), wimg.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(dx, ximg.data(), ximg.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(ds1, cs1.data(), cs1.size(), hipMemcpyHostToDevice)); + CHECK(hipMemcpy(ds2, cs2.data(), cs2.size(), hipMemcpyHostToDevice)); + CHECK(hipMemset(dst, 0, 96 * 32)); + CHECK(hipMemset(dout, 0, 48 * 4 * 32 * 4)); + step5_kernel<<<1, 32>>>(dw, dx, ds1, ds2, dst, dout); + CHECK(hipGetLastError()); + CHECK(hipDeviceSynchronize()); + std::vector st(96 * 32, 0); + CHECK(hipMemcpy(st.data(), dst, st.size(), hipMemcpyDeviceToHost)); + std::vector got(48 * 4 * 32); + CHECK(hipMemcpy(got.data(), dout, got.size() * 4, hipMemcpyDeviceToHost)); + + // 5a: the staged image BIT-exactly: D1 cells -> relu per element -> RNE-satfinite + // e4m3 code per element + int sbad = 0, schecked = 0; + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + for (int lane = 0; lane < 32; ++lane) + { + const int h = lane >> 4, col = lane & 15; + for (int p = 0; p < 4; ++p) + for (int half = 0; half < 2; ++half) + { + const int wr = 16 * mt + 2 * p + half + 8 * h; + const int tk = 16 * nt + col; + double ref = 0.0; + for (int k = 0; k < 32; ++k) + ref += dec(wimg[4096 + waddr_host(k, wr)]) * dec(ximg[(size_t)tk * 32 + k]); + const uint16_t seedbits = (uint16_t)(cs1[2 * wr] | (cs1[2 * wr + 1] << 8)); + ref += (double)h_of(seedbits); + uint16_t d1bits = b_of((_Float16)ref); + const uint16_t re = relu_half(d1bits); + const uint8_t want = requant_code(re); + const size_t off = (size_t)tk * 32 + (16 * mt + 2 * p + half + 8 * h); + ++schecked; + if (st[off] != want) + { + if (sbad < 4) + std::printf(" staged[col %2d][k %2d] = 0x%02x, expected 0x%02x\n", + tk, wr, st[off], want); + ++sbad; + } + } + } + std::printf("step 5a staged image (relu requant) bit exact: checked %d, mismatched %d\n", + schecked, sbad); + + // 5b: phase 2's pairs off the staged image, with phase 2's own per-column C + int bad = 0, checked = 0; + for (int mt = 0; mt < 2; ++mt) + for (int nt = 0; nt < 6; ++nt) + for (int lane = 0; lane < 32; ++lane) + { + const int h = lane >> 4, col = lane & 15; + for (int p = 0; p < 4; ++p) + for (int half = 0; half < 2; ++half) + { + const int wr = 16 * mt + 2 * p + half + 8 * h; + const int tk = 16 * nt + col; + // phase 2's GEMM: sum over k of dec(w2[k][wr]) * the staged + // byte's decode (relu(D1 cell (k, tk)) requantised to e4m3, + // the operand the device's GEMM reads) + f32(C seed f16), + // rounded once to f16 + double ref = 0.0; + for (int k = 0; k < 32; ++k) + { + double d1 = 0.0; + for (int j = 0; j < 32; ++j) + d1 += dec(wimg[4096 + waddr_host(j, k)]) * dec(ximg[(size_t)tk * 32 + j]); + const uint16_t seedbits = (uint16_t)(cs1[2 * k] | (cs1[2 * k + 1] << 8)); + d1 += (double)h_of(seedbits); + const uint16_t re = relu_half(b_of((_Float16)d1)); + ref += dec(wimg[6144 + waddr_host(k, wr)]) * dec(requant_code(re)); + } + const uint16_t seed2bits = (uint16_t)(cs2[2 * wr] | (cs2[2 * wr + 1] << 8)); + ref += (double)h_of(seed2bits); + const uint16_t want = b_of((_Float16)ref); + const uint16_t bits = (uint16_t)(half ? got[((mt * 6 + nt) * 4 + p) * 32 + lane] >> 16 + : got[((mt * 6 + nt) * 4 + p) * 32 + lane] & 0xffffu); + ++checked; + if (bits != want) + { + if (bad < 4) + std::printf(" D2[wr %2d][tk %2d] (mt %d nt %d lane %2d p %d h%d):" + " got %04x want %04x\n", + wr, tk, mt, nt, lane, p, half, bits, want); + ++bad; + } + } + } + std::printf("step 5b phase 2 pairs off the staged image: checked %d, mismatched %d\n", + checked, bad); + CHECK(hipFree(dw)); + CHECK(hipFree(dx)); + CHECK(hipFree(ds1)); + CHECK(hipFree(ds2)); + CHECK(hipFree(dst)); + CHECK(hipFree(dout)); + if (sbad || bad) + std::exit(1); + } + return 0; +} diff --git a/packaging/D4R_README.txt b/packaging/D4R_README.txt index ba3c06d..290998a 100644 --- a/packaging/D4R_README.txt +++ b/packaging/D4R_README.txt @@ -85,6 +85,12 @@ Settings - On RDNA4, [Kernels] NativeFp8 is on by default. It selects native FP8 WMMA and the matching gfx12-fp8 kernel folder; set it to false to use f16 widening. The FP8 path has not been tested on a real RDNA4 GPU. +@full - Ray Reconstruction: this zip includes NVIDIA's CUDA-capable denoiser (d4r/nvngx_dlssd.dll +@full 310.7.0). [RayReconstruction] Enable defaults to auto: when the game turns Ray +@full Reconstruction on, d4r runs it through the denoiser. Model = D or E selects a denoiser +@full preset explicitly; auto follows the game. RR presets D/E are validated on RX 9070 XT +@full (gfx1201) with [Kernels] NativeFp8 = on, the packaged default on RDNA4. Keep NativeFp8 on: +@full the scalar FP8 RR path is not validated. - If L/M shows extra trails in motion, try [Kernels] NativeSwinEncoders = false with PreferAccuracy = true. This uses original translated enc1/enc2 layers while retaining other native acceleration. It reduced trails in a captured Townfall sequence on RX 7700 XT and costs @@ -135,6 +141,22 @@ Why the launch options weights that match the kernels d4r runs. +Ray Reconstruction launch options +--------------------------------- +On the tested GE-Proton11-3 setup, Wine's builtin ICU libraries abort inside the denoiser +(u_setMemoryFunctions_65). With Ray Reconstruction (for example Cyberpunk 2077 2.31 with path +tracing and RR preset D or E), use launch options that prefer the game's native ICU DLLs and the +game-local d3d12.dll/d3d12core.dll over Wine's copies: + + PROTON_FORCE_NVAPI=1 DXVK_NVAPI_GPU_ARCH=AD100 WINEDLLOVERRIDES='icuuc=n,b;icuin=n,b;icudt=n,b;d3d12=n,b;d3d12core=n,b' %command% + +Merge that WINEDLLOVERRIDES value with any overrides another mod needs instead of replacing it. +Without Ray Reconstruction, the base launch options above are enough. In the game, select DLSS, +then Ray Reconstruction; d4r/d4r.ini ships [RayReconstruction] Enable = auto and +[Kernels] NativeFp8 = on for RDNA4. Keep the patched d3d12.dll/d3d12core.dll current, including +game-local copies. + + If something goes wrong ----------------------- - d4r's log is d4r/d4r_nvngx.log. It is rewritten at every launch. Lines starting with "d4r:" and @@ -147,8 +169,9 @@ If something goes wrong can be expected with the shim under d4r/; do not move it into the game root to silence that check. - A wrong image: try Model = E in d4r/d4r.ini, and report the problem with both logs. - Do not use d4r in games with anti-cheat. OptiScaler's DLL injection can get an account banned. -- Ray Reconstruction requires the CUDA-capable 3.10.7 nvngx_dlssd.dll beside the shim and the matching - patched d3d12.dll/d3d12core.dll. Enable [Kernels] NativeFp8; the scalar FP8 RR path is not validated. +- Ray Reconstruction needs the CUDA-capable 3.10.7 nvngx_dlssd.dll beside the shim, the matching + patched d3d12.dll/d3d12core.dll, and the ICU launch overrides from the Ray Reconstruction section + above. Enable [Kernels] NativeFp8; the scalar FP8 RR path is not validated. - Not supported: DLSS Frame Generation and DirectX 11. - Development builds expose an experimental native Vulkan NGX path using CPU staging. Youngblood is the first intended game test; gameplay remains unverified. See the diff --git a/scripts/build_d4r_nvngx_shim.sh b/scripts/build_d4r_nvngx_shim.sh index 034fbea..014d3d6 100755 --- a/scripts/build_d4r_nvngx_shim.sh +++ b/scripts/build_d4r_nvngx_shim.sh @@ -21,6 +21,6 @@ PARAM_OBJECT="$ROOT/build/d4r_ngx_param_msvc.obj" # The engine (D4R_ENGINE) reaches Vulkan through a function table: MinGW has no Vulkan import library. "$MINGW_CXX" -std=c++20 -O2 -Wall -Wextra -Wno-missing-field-initializers -DD4R_ENGINE_VK_TABLE -I"$ROOT/build/vulkan-include" -I"$ROOT/build/vulkan-staging" -shared -static -static-libgcc -static-libstdc++ \ "$ROOT/tools/d4r_nvngx_shim.cpp" "$ROOT/tools/d4r_ngx_param_host.cpp" "$PARAM_OBJECT" \ - "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/vk_dispatch.cpp" \ + "$ROOT/engine/runtime.cpp" "$ROOT/engine/game.cpp" "$ROOT/engine/game_m.cpp" "$ROOT/engine/game_rr.cpp" "$ROOT/engine/vk_dispatch.cpp" \ -Wl,--entry,d4r_dll_entry -o "$ROOT/build/d4r_nvngx.dll" printf 'Built %s\n' "$ROOT/build/d4r_nvngx.dll" diff --git a/scripts/run_rr_engine_compare.sh b/scripts/run_rr_engine_compare.sh new file mode 100755 index 0000000..3e37f36 --- /dev/null +++ b/scripts/run_rr_engine_compare.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Runs the D3D12 harness's Ray Reconstruction path twice with identical inputs - once on the CUDA backend and once +# on the native engine (engine/runtime_rr.inc with a model from engine/compile_rr.py) - and compares the frames. +# +# usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR +# environment: as scripts/run_rr_guide_compare.sh (D4R_PROTON_DIR, D4R_RUNTIME_DIR, D4R_VKD3D_DIR, D4R_ZLUDA_DIR, +# D4R_HARNESS_SIZE, D4R_HARNESS_FRAMES); D4R_RR_PRESET 5 = E (default; the model must match) +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +OUTPUT_DIR="$(realpath -m "${1:?usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR}")" +MODEL_DIR="$(realpath "${2:?usage: run_rr_engine_compare.sh OUTPUT_DIR MODEL_DIR}")" +PROTON_DIR="${D4R_PROTON_DIR:-$HOME/.local/share/Steam/compatibilitytools.d/GE-Proton11-3}" +RUNTIME_DIR="${D4R_RUNTIME_DIR:-$HOME/.local/share/d4r-dlss}" +VKD3D_DIR="${D4R_VKD3D_DIR:-$RUNTIME_DIR/vkd3d}" +ZLUDA_DIR="${D4R_ZLUDA_DIR:-$HOME/.cache/d4r-zluda-rr}" +NVAPI_DIR="${D4R_NVAPI_DIR:-$PROTON_DIR/files/lib/wine/nvapi/x86_64-windows}" +FRAMES="${D4R_HARNESS_FRAMES:-6}" +read -r IN_W IN_H OUT_W OUT_H <<< "${D4R_HARNESS_SIZE:-640 360 1280 720}" +APP_DIR="$OUTPUT_DIR/app" +mkdir -p "$APP_DIR" +cp -f "$ROOT/build/d3d12_dlss_harness.exe" "$ROOT/build/d4r_nvngx.dll" \ + "$RUNTIME_DIR/bin/nvngx_dlssd.dll" "$RUNTIME_DIR/bin/nvngx_dlss.dll" \ + "$VKD3D_DIR/d3d12.dll" "$VKD3D_DIR/d3d12core.dll" "$NVAPI_DIR/nvapi64.dll" "$APP_DIR/" + +export D4R_HARNESS_SKIP_BUILD=1 D4R_HARNESS_APP_DIR="$APP_DIR" D4R_HARNESS_NGX_DLL="$APP_DIR/d4r_nvngx.dll" +export D4R_RUNTIME_DIR="$RUNTIME_DIR" D4R_ZLUDA_DIR="$ZLUDA_DIR" +export D4R_ZLUDA_NATIVE_DIR="$RUNTIME_DIR/kernels" D4R_ZLUDA_IMPLICIT_MAX_BLOCK=256 +export D4R_ZLUDA_WMMA=1 D4R_ZLUDA_WMMA_FP8=1 D4R_ZLUDA_WMMA_FP8_NATIVE=1 +export D4R_SHIM_VRAM_INTEROP=1 D4R_RR_ENABLE=1 D4R_RR_PRESET="${D4R_RR_PRESET:-5}" +export D4R_HARNESS_FEATURE_ID=13 D4R_HARNESS_SAVE_FRAMES=1 +export D4R_HARNESS_GUIDE_FORMATS="${D4R_HARNESS_GUIDE_FORMATS:-rgba16f,rgba16f,rgba8unorm,rgba8unorm}" +export WINEDLLOVERRIDES="d3d12=n,b;d3d12core=n,b" + +run_one() { # run_one LABEL ENGINE + printf '\n== %s\n' "$1" >&2 + D4R_ENGINE="$2" D4R_ENGINE_RR_SYNC=1 D4R_ENGINE_MODEL_RR_E_DIR="$MODEL_DIR" D4R_ENGINE_MODEL_RR_D_DIR="$MODEL_DIR" \ + "$ROOT/scripts/run_d3d12_dlss_harness_proton.sh" "$PROTON_DIR" "$OUTPUT_DIR/$1.raw" \ + "$FRAMES" "$IN_W" "$IN_H" "$OUT_W" "$OUT_H" > "$OUTPUT_DIR/$1.log" 2>&1 || true + grep -E 'engine backend|engine frame|Ray Reconstruction|output read back|rror' "$OUTPUT_DIR/$1.log" | head -20 || true +} +run_one cuda 0 +run_one engine 1 +python3 "$ROOT/tools/compare_rr_output.py" --a "$OUTPUT_DIR/cuda.raw" --b "$OUTPUT_DIR/engine.raw" \ + --width "$OUT_W" --height "$OUT_H" --frames "$FRAMES" --same-frames diff --git a/tools/d3d12_dlss_harness.cpp b/tools/d3d12_dlss_harness.cpp index 6e1ee2a..95d8b0c 100644 --- a/tools/d3d12_dlss_harness.cpp +++ b/tools/d3d12_dlss_harness.cpp @@ -1331,6 +1331,38 @@ int main(int argc, char** argv) } upload(alphaOutput, sentinel.data(), alphaWidth * texel, D3D12_RESOURCE_STATE_UNORDERED_ACCESS); } + // D4R_HARNESS_RR_EXTRA=hit,mask: the two optional denoiser inputs RE Requiem passes, in its formats: the specular + // hit distance (R32_FLOAT) and the disocclusion mask (R16_FLOAT), each a pattern that varies over the frame so a + // capture shows where the denoiser reads them. + ID3D12Resource* hitDistance = nullptr; + ID3D12Resource* disocclusionMask = nullptr; + if (const char* extra = reconstruction ? std::getenv("D4R_HARNESS_RR_EXTRA") : nullptr) + { + if (std::strstr(extra, "hit") != nullptr) + { + std::vector distance(static_cast(resourceWidth) * inHeight); + for (UINT y = 0; y < inHeight; ++y) + for (UINT x = 0; x < resourceWidth; ++x) + distance[static_cast(y) * resourceWidth + x] = + 0.5f + 20.0f * static_cast((x / 16 + y / 16) % 7) / 6.0f + 0.01f * static_cast(x % 16); + hitDistance = create_texture(resourceWidth, inHeight, DXGI_FORMAT_R32_FLOAT, D3D12_RESOURCE_FLAG_NONE); + if (!hitDistance) return 1; + upload(hitDistance, distance.data(), resourceWidth * sizeof(float), srv); + } + if (std::strstr(extra, "mask") != nullptr) + { + std::vector mask(static_cast(resourceWidth) * inHeight); + for (UINT y = 0; y < inHeight; ++y) + for (UINT x = 0; x < resourceWidth; ++x) + mask[static_cast(y) * resourceWidth + x] = std::bit_cast(static_cast<_Float16>( + ((x / 24 + y / 24) % 5 == 0) ? 1.0f : ((x / 24 + y / 24) % 5 == 1 ? 0.5f : 0.0f))); + disocclusionMask = create_texture(resourceWidth, inHeight, DXGI_FORMAT_R16_FLOAT, D3D12_RESOURCE_FLAG_NONE); + if (!disocclusionMask) return 1; + upload(disocclusionMask, mask.data(), resourceWidth * sizeof(uint16_t), srv); + } + std::printf("Ray Reconstruction extra inputs:%s%s\n", hitDistance ? " DLSSD.SpecularHitDistance" : "", + disocclusionMask ? " DLSS.DisocclusionMask" : ""); + } ID3D12Resource* colorTexture = create_texture(resourceWidth, inHeight, rgba8 ? DXGI_FORMAT_R8G8B8A8_UNORM : DXGI_FORMAT_R16G16B16A16_FLOAT, D3D12_RESOURCE_FLAG_NONE); ID3D12Resource* depthTexture = create_texture(resourceWidth, inHeight, DXGI_FORMAT_R32_FLOAT, D3D12_RESOURCE_FLAG_NONE); @@ -1464,14 +1496,33 @@ int main(int argc, char** argv) (alphaScenario ? 0x80 : 0)); if (alphaScenario) d4r_ngx_set_int(parameters, "DLSS.Enable.Output.Subrects", 1); - std::array identity{1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + // D4R_HARNESS_RR_MATRICES="w0,...,w15;c0,...,c15": the camera's WorldToView and ViewToClip matrices + // (row-major as the SDK passes them) instead of identities, for Ray Reconstruction. + static std::array worldToView{1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1}; + static std::array viewToClip = worldToView; + if (const char* matrices = std::getenv("D4R_HARNESS_RR_MATRICES")) + { + const char* cursor = matrices; + for (std::array* matrix : {&worldToView, &viewToClip}) + for (float& value : *matrix) + { + char* end = nullptr; + value = std::strtof(cursor, &end); + if (end == cursor) + { + std::fprintf(stderr, "D4R_HARNESS_RR_MATRICES needs 32 numbers\n"); + return 2; + } + cursor = *end != '\0' ? end + 1 : end; + } + } if (reconstruction) { d4r_ngx_set_int(parameters, "DLSS.Denoise.Mode", 1); d4r_ngx_set_int(parameters, "DLSS.Roughness.Mode", 0); d4r_ngx_set_int(parameters, "DLSS.Use.HW.Depth", 0); - d4r_ngx_set_void(parameters, "WorldToViewMatrix", identity.data()); - d4r_ngx_set_void(parameters, "ViewToClipMatrix", identity.data()); + d4r_ngx_set_void(parameters, "WorldToViewMatrix", worldToView.data()); + d4r_ngx_set_void(parameters, "ViewToClipMatrix", viewToClip.data()); const char* preset = std::getenv("D4R_HARNESS_RR_PRESET"); if (preset) for (const char* mode : {"DLAA", "Quality", "Balanced", "Performance", "UltraPerformance", "UltraQuality"}) @@ -1525,6 +1576,10 @@ int main(int argc, char** argv) for (size_t input = 0; input < guides.size(); ++input) d4r_ngx_set_d3d12_resource(parameters, names[input], guides[input]); } + if (hitDistance != nullptr && replayDir == nullptr) + d4r_ngx_set_d3d12_resource(parameters, "DLSSD.SpecularHitDistance", hitDistance); + if (disocclusionMask != nullptr && replayDir == nullptr) + d4r_ngx_set_d3d12_resource(parameters, "DLSS.DisocclusionMask", disocclusionMask); // DLSSD.Alpha is a guide the capture records in its own right, so the synthetic alpha // scenario would double-bind that name; it belongs to the static probe only. if (alphaScenario && replayDir == nullptr) @@ -1915,6 +1970,38 @@ int main(int argc, char** argv) d4r_ngx_set_float(parameters, "DLSS.Pre.Exposure", static_cast(std::atof(value))); if (const char* value = std::getenv("D4R_HARNESS_EXPOSURE_SCALE")) d4r_ngx_set_float(parameters, "DLSS.Exposure.Scale", static_cast(std::atof(value))); + // D4R_HARNESS_NGX_FLOATS="Name=value,Name=value": float evaluation parameters set on every frame after the + // defaults (e.g. MV.Scale.X=2,FrameTimeDeltaInMsec=8), to see how NGX passes them on. + if (const char* list = std::getenv("D4R_HARNESS_NGX_FLOATS")) + { + std::string entries(list); + for (size_t start = 0; start < entries.size();) + { + const size_t end = std::min(entries.find(',', start), entries.size()); + const std::string entry = entries.substr(start, end - start); + const size_t equals = entry.find('='); + if (equals != std::string::npos) + d4r_ngx_set_float(parameters, entry.substr(0, equals).c_str(), + std::strtof(entry.c_str() + equals + 1, nullptr)); + start = end + 1; + } + } + // D4R_HARNESS_RR_MATRIX_STEP="32 numbers": added to D4R_HARNESS_RR_MATRICES' two matrices once per frame + // (frame N uses base + (N - 1) * step), so the camera moves and turns between evaluations. + if (const char* step = std::getenv("D4R_HARNESS_RR_MATRIX_STEP"); step != nullptr && reconstruction) + { + static const std::array baseView = worldToView, baseClip = viewToClip; + const char* cursor = step; + for (int m = 0; m < 2; ++m) + for (int i = 0; i < 16; ++i) + { + char* end = nullptr; + const float value = std::strtof(cursor, &end); + cursor = end != nullptr && *end != '\0' ? end + 1 : end; + (m == 0 ? worldToView : viewToClip)[i] = + (m == 0 ? baseView : baseClip)[i] + static_cast(frame - 1) * value; + } + } // Change exposure after initial frames to exercise temporal/backend transitions. // CUDA supports this scale; the native K adapter currently accepts only a neutral scale. if (const char* changeFrame = std::getenv("D4R_HARNESS_EXPOSURE_CHANGE_FRAME")) diff --git a/tools/d4r_nvngx_shim.cpp b/tools/d4r_nvngx_shim.cpp index 64b6fc1..631d70f 100644 --- a/tools/d4r_nvngx_shim.cpp +++ b/tools/d4r_nvngx_shim.cpp @@ -79,6 +79,7 @@ #include "d4r_vkd3d_interop.h" #include "../engine/game.h" #include "../engine/game_m.h" +#include "../engine/game_rr.h" #include "../engine/vk_dispatch.h" // --- NGX types (declared locally: the core's exports differ from the SDK) -- @@ -290,7 +291,6 @@ struct CoreApi // The core's CUDA entry point takes a CUdevice ordinal, not an IDXGIAdapter* (its // disassembly rejects a negative value in the first argument, which is the ordinal // check). The D3D12 export this file implements takes an adapter and must translate. - NgxResult (*getFeatureRequirements)(int, const void*, NgxFeatureRequirement*) = nullptr; }; // --- DLSS-Denoiser (Ray Reconstruction) ---------------------------------------- @@ -1473,9 +1473,6 @@ static bool load_libraries() ok &= load_export(g.core, "NVSDK_NGX_CUDA_CreateFeature", g.ngx.createFeature); ok &= load_export(g.core, "NVSDK_NGX_CUDA_ReleaseFeature", g.ngx.releaseFeature); ok &= load_export(g.core, "NVSDK_NGX_CUDA_EvaluateFeature", g.ngx.evaluateFeature); - // Optional: the core answers feature 13 by consulting the denoiser it loaded. - if (!load_export(g.core, "NVSDK_NGX_CUDA_GetFeatureRequirements", g.ngx.getFeatureRequirements)) - g.ngx.getFeatureRequirements = nullptr; g.cuda = LoadLibraryA("nvcuda.dll"); if (g.cuda == nullptr) @@ -2538,6 +2535,7 @@ struct Feature std::unique_ptr engineMDynamic; // Vulkan at changing sizes; original HIP resources stay alive unsigned int engineMLastWidth = 0, engineMLastHeight = 0; unsigned int engineMRenderWidth = 0, engineMRenderHeight = 0; // first evaluation's effective extent + std::unique_ptr engineR; // Ray Reconstruction (engine/runtime_rr.inc) // Preset M and preset K with the network on the native HIP layers ([Engine] Network): the buffers shared // with HIP and the bridge's network. The game's list is split around the network, as for the CUDA backend's // split frames. `engineHead` is the network's result (M's expansion / K's final store reads it). @@ -6973,6 +6971,17 @@ D4R_EXPORT NgxResult NVSDK_NGX_D3D12_GetScratchBufferSize(unsigned int, const vo // NVSDK_NGX_FeatureRequirement *OutSupported). The feature is NOT a separate argument: it is // FeatureDiscoveryInfo->FeatureID (defs.h, the field after SDKVersion), and the first argument // is the display adapter, which must never be read as a feature id. +// +// This answer is derived from d4r's own state alone. The earlier version forwarded the query +// to the NGX core's CUDA entry point on the shim's worker thread and, for feature 13, called +// the denoiser's discovery export before Init. Those calls turn a requirements query that +// OptiScaler hooks into a cross-binary handoff: inside a Steam runtime that reorders module +// loads, OptiScaler's follow-up scan of its own support lists then read memory the other +// binary had already released (deterministic EXCEPTION_ACCESS_VIOLATION inside OptiScaler's +// dxgi.dll; see the repo's Cyberpunk notes). The requirements answer carries no state that is +// not already authoritative elsewhere: Ray Reconstruction's real gate is the CreateFeature +// check plus the core's denoiser account reported through GetCapabilityParameters, both of +// which run after their CUDA context exists. D4R_EXPORT NgxResult NVSDK_NGX_D3D12_GetFeatureRequirements(void* adapter, const void* discovery, NgxFeatureRequirement* requirement) { @@ -6984,89 +6993,47 @@ D4R_EXPORT NgxResult NVSDK_NGX_D3D12_GetFeatureRequirements(void* adapter, const if (discovery != nullptr) std::memcpy(&feature, static_cast(discovery) + sizeof(unsigned int), sizeof(feature)); std::memset(requirement, 0, sizeof(*requirement)); - ensure_portable_config(); - // The local switch is checked before the core is asked: with D4R_RR_ENABLE=0 this process - // is not offering Ray Reconstruction, and the core's authoritative "supported" must not - // be passed on to a game that would then create a feature this shim would refuse. + if (feature != NGX_FEATURE_SUPER_SAMPLING && feature != NGX_FEATURE_RAY_RECONSTRUCTION) + { + requirement->FeatureSupported = 1; // not supported + return NGX_SUCCESS; + } + // The local switch is checked before any core appears: with D4R_RR_ENABLE=0 this process + // is not offering Ray Reconstruction, and no core answer may override that. if (feature == NGX_FEATURE_RAY_RECONSTRUCTION && !rr_enabled()) { requirement->FeatureSupported = 1; // not supported logf("Ray Reconstruction is not supported: %s", rr_denial_reason()); return NGX_SUCCESS; } + // CPU-side adapter identity only; safe before Init and without any CUDA context. + if (adapter != nullptr && !g.ngxInitialized) + { + DXGI_ADAPTER_DESC desc = {}; + if (SUCCEEDED(static_cast(adapter)->GetDesc(&desc))) + set_cuda_adapter_luid(desc.AdapterLuid); + } { std::lock_guard lock(g.mutex); - // Discovery can precede Init; CUDA still needs the caller's adapter identity. - if (adapter != nullptr && !g.ngxInitialized) - { - DXGI_ADAPTER_DESC desc = {}; - if (SUCCEEDED(static_cast(adapter)->GetDesc(&desc))) - set_cuda_adapter_luid(desc.AdapterLuid); - } - // Requirements must not depend on a prior Init. Load the provider without - // initializing NGX here: the later Init retains the game's identity and device. + // Load the provider without initializing NGX: creation keeps the game's identity, and + // feature 13's acceptance is decided later from ray_reconstruction_available(). if (ensure_workers_started() != NGX_SUCCESS || (g.core == nullptr && !load_libraries())) return NGX_FAIL_PLATFORM_ERROR; if (feature == NGX_FEATURE_RAY_RECONSTRUCTION) load_denoiser_library(); } - // Otherwise the core's answer is authoritative - it consulted the denoiser it loaded - and - // this file only steps in when the core declines. - if (g.ngx.getFeatureRequirements != nullptr) - { - // The adapter this export receives is not what the core's CUDA entry point wants: it - // takes a CUdevice ordinal. Passing the adapter pointer through would read its low 32 - // bits as an ordinal, so the current context's device is asked for instead. - int device = 0; - const NgxResult result = g.worker.call([&] { - if (g.cu.ctxGetDevice != nullptr && g.cu.ctxGetDevice(&device) != 0) - device = 0; // no current context: device 0 is this process's only CUDA device - return g.ngx.getFeatureRequirements(device, discovery, requirement); - }); - if (result == NGX_SUCCESS) - { - logf("GetFeatureRequirements(feature=%u, cuda device %d) from the NGX core: supported=%u minHW=0x%x " - "'%s'", - feature, device, requirement->FeatureSupported, requirement->MinHWArchitecture, - requirement->MinOSVersion); - return result; - } - logf("GetFeatureRequirements(feature=%u) from the NGX core -> 0x%08x; answering here instead", feature, - result); - } - if (feature == NGX_FEATURE_RAY_RECONSTRUCTION) - { - // Requirements are legal before Init. Ask the denoiser's discovery export, - // not its Init stub, without binding NGX to an application prematurely. - if (!g.ngxInitialized && g.denoiser.cudaCapable) - { - using RequirementsFn = NgxResult (*)(void*, const void*, NgxFeatureRequirement*); - RequirementsFn denoiserRequirements = nullptr; - if (load_export(g.denoiserModule, "NVSDK_NGX_D3D12_GetFeatureRequirements", denoiserRequirements)) - { - const NgxResult result = g.worker.call([&] { - return denoiserRequirements(adapter, discovery, requirement); - }); - logf("GetFeatureRequirements before Init from DLSS-Denoiser -> 0x%08x, supported=%u", - result, requirement->FeatureSupported); - return result; - } - } - const bool supported = ray_reconstruction_available(); - if (!supported) - { - requirement->FeatureSupported = 1; - logf("Ray Reconstruction is not supported: %s", - g.ngxInitialized ? rr_denial_reason() : "no CUDA-capable DLSS-Denoiser for pre-Init discovery"); - return NGX_SUCCESS; - } - } - if (feature != NGX_FEATURE_SUPER_SAMPLING && feature != NGX_FEATURE_RAY_RECONSTRUCTION) + const bool supported = feature == NGX_FEATURE_SUPER_SAMPLING + ? true + : (g.ngxInitialized ? ray_reconstruction_available() : g.denoiser.cudaCapable); + if (!supported) { requirement->FeatureSupported = 1; // not supported + if (feature == NGX_FEATURE_RAY_RECONSTRUCTION) + logf("Ray Reconstruction is not supported: %s", + g.ngxInitialized ? rr_denial_reason() : "no CUDA-capable DLSS-Denoiser for pre-Init discovery"); return NGX_SUCCESS; } - requirement->FeatureSupported = 0; // supported + requirement->FeatureSupported = 0; // supported (NVSDK_NGX_FeatureSupportResult bitfield) requirement->MinHWArchitecture = 0x160; return NGX_SUCCESS; } @@ -7317,10 +7284,12 @@ static NgxResult create_feature(unsigned int featureId, void* parameters, NgxHan feature->subrects = subrects; std::copy(std::begin(presets), std::end(presets), feature->creationPresets.begin()); // The engine records D3D12 evaluations only; the native Vulkan frontend keeps eager CUDA creation. Presets K, L - // and M (the engine's) create the CUDA feature only when a frame falls back: NGX's L kernels take ZLUDA about 20 s - // to compile on a cold cache, which stalled the game inside CreateFeature (Overwatch, 2026-10-11). + // and M and Ray Reconstruction (the engine's) create the CUDA feature only when a frame falls back: NGX's L kernels + // take ZLUDA about 20 s to compile on a cold cache, which stalled the game inside CreateFeature (Overwatch, + // 2026-10-11), and the denoiser's are larger still. feature->cudaCreationDeferred = !nativeVulkan && engine_requested() && - (feature->preset == 11 || feature->preset == 12 || feature->preset == 13); + (rayReconstruction || feature->preset == 11 || feature->preset == 12 || + feature->preset == 13); const NgxResult result = feature->cudaCreationDeferred ? NGX_SUCCESS : create_cuda_feature(*feature); if (feature->cudaCreationDeferred) logf("engine backend: CUDA feature creation deferred until fallback is needed"); @@ -8268,11 +8237,272 @@ static bool engine_evaluate_m(Feature& feature, ID3D12GraphicsCommandList* list, return true; } +// Ray Reconstruction models (engine/compile_rr.py): one per preset, engine\rr-d and engine\rr-e next to this DLL, or +// D4R_ENGINE_MODEL_RR_D_DIR / D4R_ENGINE_MODEL_RR_E_DIR. +struct EngineBackendR +{ + // The model's pipelines take minutes to compile on a cold shader cache, so each preset's model loads on a thread + // of its own; until it is ready the CUDA backend serves the frames. state: 0 not started, 1 loading, 2 ready, + // 3 failed. + std::once_flag once[2]; // [preset E] + std::atomic state[2] = {}; + std::unique_ptr models[2]; +}; +static EngineBackendR g_engineR; + +static void init_engine_backend_rr(bool presetE) +{ + std::string directory = env_string(presetE ? "D4R_ENGINE_MODEL_RR_E_DIR" : "D4R_ENGINE_MODEL_RR_D_DIR"); + if (directory.empty()) + { + char folder[MAX_PATH * 3] = {}; + WideCharToMultiByte(CP_ACP, 0, (g_portable.dir + (presetE ? L"\\engine\\rr-e" : L"\\engine\\rr-d")).c_str(), -1, + folder, sizeof(folder), nullptr, nullptr); + directory = folder; + } + if (g_vk.interop == nullptr || g_vk.lifetime == nullptr || g_vk.device == VK_NULL_HANDLE || d4r::vkTable.CreateBuffer == nullptr) + { + logf("engine backend (Ray Reconstruction): Vulkan interop is unavailable"); + g_engineR.state[presetE] = 3; + return; + } + const std::string path = directory[0] == '/' ? "Z:" + directory : directory; + logf("engine backend (Ray Reconstruction %c): loading the model from %s and compiling its pipelines; the CUDA " + "backend serves frames until it is ready", presetE ? 'E' : 'D', path.c_str()); + const auto start = std::chrono::steady_clock::now(); + try + { + auto model = std::make_unique(g_vk.physical, g_vk.device, path); + if (model->preset() != (presetE ? 5u : 4u)) + throw std::runtime_error(presetE ? "the directory holds a preset D model" : "the directory holds a preset E model"); + g_engineR.models[presetE] = std::move(model); + } + catch (const std::exception& error) + { + logf("engine backend (Ray Reconstruction): cannot load the model from %s: %s", path.c_str(), error.what()); + g_engineR.state[presetE] = 3; + return; + } + g_engineR.state[presetE] = 2; + logf("engine backend (Ray Reconstruction %c): model %s loaded and pipelines compiled in %.0f ms", presetE ? 'E' : 'D', + path.c_str(), std::chrono::duration(std::chrono::steady_clock::now() - start).count()); +} + +// Ray Reconstruction on the engine: d4r's stages in the game's command list, with the denoiser's guides sampled +// straight from the game's textures. False: the CUDA backend serves the feature (reason logged). +static bool engine_evaluate_rr(Feature& feature, ID3D12GraphicsCommandList* list, ID3D12Resource* color, + ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, void* parameters, + const FrameParams& p, uint32_t frame) +{ + auto refuse = [&](const char* reason) { + feature.engineRefused = true; + logf("engine backend (Ray Reconstruction): %s; feature %u uses the CUDA backend from frame %u", reason, + feature.handle.Id, frame); + return false; + }; + // NGX 310.7 runs preset D's kernels when the game names no preset (harness, 2026-10-11). + const bool presetE = feature.preset == 5; + if (feature.preset != 0 && feature.preset != 4 && feature.preset != 5) + return refuse("only presets D and E are implemented"); + std::call_once(g_engineR.once[presetE], [presetE] { + g_engineR.state[presetE] = 1; + // D4R_ENGINE_RR_SYNC=1 (harness comparisons): load in place so the engine serves the first frame + if (env_uint("D4R_ENGINE_RR_SYNC", 0) != 0) + init_engine_backend_rr(presetE); + else + std::thread([presetE] { init_engine_backend_rr(presetE); }).detach(); + }); + const int state = g_engineR.state[presetE].load(); + if (state == 1) + return false; // still compiling: the CUDA backend serves this frame, the engine is tried again next frame + d4r::RModel* const model = g_engineR.models[presetE].get(); + if (state != 2 || model == nullptr) + return refuse(presetE ? "no preset E model (engine\\rr-e)" : "no preset D model (engine\\rr-d)"); + // NVSDK_NGX_DLSS_Feature_Flags: IsHDR 1, MVLowRes 2, MVJittered 4, DepthInverted 8. + if ((feature.flags & 0x3) != 0x3) + return refuse("needs the HDR flag and render-resolution motion vectors"); + if ((feature.flags & 0x4) != 0) + return refuse("jittered motion vectors are not implemented"); + if (!p.auxMatrixSet[0] || !p.auxMatrixSet[1]) + return refuse("the camera matrices are missing"); + if (feature.engineR != nullptr && + (p.renderWidth != feature.engineMRenderWidth || p.renderHeight != feature.engineMRenderHeight)) + return refuse("changing the effective render resolution after initialization is not implemented"); + const char* const guideNames[4][2] = {{"GBuffer.Normals", nullptr}, {"GBuffer.Roughness", nullptr}, + {"DLSS.Input.DiffuseAlbedo", "GBuffer.DiffuseAlbedo"}, + {"DLSS.Input.SpecularAlbedo", "GBuffer.SpecularAlbedo"}}; + ID3D12Resource* guides[4] = {}; + // DLSS.Roughness.Mode 1 (packed): no roughness texture, the roughness is the normals' w + const bool roughnessPacked = get_resource(parameters, "GBuffer.Roughness") == nullptr && + get_uint_or(parameters, "DLSS.Roughness.Mode", 0) == 1; + for (int i = 0; i < 4; ++i) + { + guides[i] = get_resource(parameters, guideNames[i][0]); + if (guides[i] == nullptr && guideNames[i][1] != nullptr) + guides[i] = get_resource(parameters, guideNames[i][1]); + if (guides[i] == nullptr && i == 1 && roughnessPacked) + guides[i] = guides[0]; + if (guides[i] == nullptr) + return refuse("a guide the engine needs (normals, roughness, diffuse or specular albedo) is missing"); + } + // the optional inputs a game may add (RE Requiem passes both): bound when present, like NGX + ID3D12Resource* const hitDistance = get_resource(parameters, "DLSSD.SpecularHitDistance"); + ID3D12Resource* const disocclusion = get_resource(parameters, "DLSS.DisocclusionMask"); + for (ID3D12Resource* input : {color, depth, motion, guides[0], guides[1], guides[2], guides[3], hitDistance, disocclusion}) + if (input == output) + return refuse("input/output texture aliasing is not supported"); + d4r::GameFrameR game; + using Source = EngineFrameSource; + const auto inputState = static_cast(env_uint("D4R_SHIM_INPUT_STATE", 0x40)); + const auto depthState = static_cast(env_uint("D4R_SHIM_DEPTH_STATE", 0x40)); + const auto outputState = static_cast(env_uint("D4R_SHIM_OUTPUT_STATE", 0x8)); + const unsigned int rw = p.renderWidth, rh = p.renderHeight; + std::vector sources = { + {color, &game.color, p.colorBaseX, p.colorBaseY, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {depth, &game.depth, p.depthBaseX, p.depthBaseY, rw, rh, depthState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {motion, &game.motion, p.mvBaseX, p.mvBaseY, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {guides[0], &game.normals, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {guides[1], &game.roughness, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {guides[2], &game.diffuseAlbedo, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {guides[3], &game.specularAlbedo, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}, + {output, &game.output, p.outputBaseX, p.outputBaseY, feature.outWidth, feature.outHeight, outputState, + D3D12_RESOURCE_STATE_COPY_DEST}, + }; + if (hitDistance != nullptr) + sources.push_back({hitDistance, &game.specularHitDistance, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}); + if (disocclusion != nullptr) + sources.push_back({disocclusion, &game.disocclusionMask, 0, 0, rw, rh, inputState, D3D12_RESOURCE_STATE_COPY_SOURCE}); + // a guide with its own subrect origin (DLSS.Input.Normals.Subrect.Base etc.) is not implemented + for (size_t i = 0; i < sizeof(kAuxSubrectNames) / sizeof(kAuxSubrectNames[0]); ++i) + if (p.auxSubrect[i][0] != 0 || p.auxSubrect[i][1] != 0) + return refuse("guide subrect origins are not implemented"); + for (Source& source : sources) + { + UINT64 handle = 0, offset = 0; + D3D12_RESOURCE_DESC desc; + source.resource->GetDesc(&desc); + if (FAILED(g_vk.interop->GetVulkanResourceInfo1(source.resource, &handle, &offset, &source.image->format)) || + handle == 0 || desc.Dimension != D3D12_RESOURCE_DIMENSION_TEXTURE2D || desc.SampleDesc.Count != 1 || + source.x + source.width > desc.Width || source.y + source.height > desc.Height) + return refuse("a texture is not a single-sample 2D image covering its rectangle"); + source.image->image = reinterpret_cast(handle); + source.image->x = source.x; + source.image->y = source.y; + // The stages sample inputs with normalized coordinates over the whole view: a texture larger than the render + // size (dynamic resolution, padded targets) goes through the copy into an exactly sized plane instead. + const bool exact = source.image == &game.output || (desc.Width == source.width && desc.Height == source.height); + prepare_engine_view(feature, source.resource, desc, *source.image, source.image == &game.output, + source.x == 0 && source.y == 0 && exact, source.copyState); + g_vk.interop->GetVulkanImageLayout(source.resource, source.copyState, &source.image->layout); + } + // the same resource can appear twice (a game may pass one texture as two guides): transition it once + std::vector distinct; + for (const Source& source : sources) + if (std::find(distinct.begin(), distinct.end(), source.resource) == distinct.end()) + distinct.push_back(source.resource); + if (distinct.size() != sources.size() - (roughnessPacked ? 1u : 0u)) + return refuse("two inputs share one texture"); + // each distinct resource is transitioned once (the packed roughness repeats the normals) + auto firstUse = [&](size_t i) { + for (size_t j = 0; j < i; ++j) + if (sources[j].resource == sources[i].resource) return false; + return true; + }; + if (game.depth.format == VK_FORMAT_D32_SFLOAT || game.depth.format == VK_FORMAT_D32_SFLOAT_S8_UINT) + game.depth.aspect = VK_IMAGE_ASPECT_DEPTH_BIT; + if (!d4r::gameColorFormat(game.color.format) || !d4r::gameMotionFormat(game.motion.format) || + !d4r::gameDepthFormat(game.depth.format, game.depth.aspect) || !d4r::gameOutputFormat(game.output.format)) + { + logf("engine backend (Ray Reconstruction): Vulkan formats colour %d, depth %d, motion %d, output %d", + game.color.format, game.depth.format, game.motion.format, game.output.format); + return refuse("a texture format is not implemented"); + } + if (frame <= 2) + { + const char* names[8] = {"color", "depth", "motion", "normals", "roughness", "diffuse", "specular", "output"}; + for (int i = 0; i < 8; ++i) + { + D3D12_RESOURCE_DESC desc; + sources[i].resource->GetDesc(&desc); + logf("engine backend (Ray Reconstruction): %s format %d layout %d origin %u,%u size %llux%u view %p aspect %u", + names[i], sources[i].image->format, sources[i].image->layout, sources[i].image->x, sources[i].image->y, + static_cast(desc.Width), desc.Height, reinterpret_cast(sources[i].image->view), + sources[i].image->aspect); + } + } + auto& s = game.settings; + s.jitter[0] = p.jitterX; s.jitter[1] = p.jitterY; + s.motionScale[0] = p.mvScaleX; s.motionScale[1] = p.mvScaleY; + s.preExposure = p.preExposure; + s.frameTimeMs = p.frameTime; + s.reset = p.reset != 0 || feature.engineR == nullptr; + s.depthInverted = (feature.flags & 0x8) != 0; + s.roughnessPacked = roughnessPacked; + if (frame <= 2) + logf("engine backend (Ray Reconstruction): roughness %s, specular hit distance %s, disocclusion mask %s", + roughnessPacked ? "packed in the normals' w" : "texture", hitDistance != nullptr ? "bound" : "none", + disocclusion != nullptr ? "bound" : "none"); + std::copy(std::begin(p.auxMatrix[0]), std::end(p.auxMatrix[0]), s.worldToView); + std::copy(std::begin(p.auxMatrix[1]), std::end(p.auxMatrix[1]), s.viewToClip); + try + { + if (feature.engineR == nullptr) + { + feature.engineR = std::make_unique(*model, g_vk.physical, g_vk.device, feature.outWidth, + feature.outHeight, rw, rh); + feature.engineMRenderWidth = rw; + feature.engineMRenderHeight = rh; + logf("engine backend (Ray Reconstruction %c): feature %u runs in the game's command list (%ux%u -> %ux%u)", + presetE ? 'E' : 'D', feature.handle.Id, rw, rh, feature.outWidth, feature.outHeight); + } + } + catch (const std::exception& error) + { + logf("engine backend (Ray Reconstruction): %s", error.what()); + return refuse("the engine could not be created"); + } + if (FAILED(g_vk.lifetime->RetainExternalResources(list, feature.resources))) + return refuse("the command list cannot retain the engine's resources"); + for (size_t i = 0; i < sources.size(); ++i) + if (firstUse(i)) + transition(list, sources[i].resource, sources[i].state, sources[i].copyState); + VkCommandBuffer command = VK_NULL_HANDLE; + bool recorded = SUCCEEDED(g_vk.interop->BeginVkCommandBufferInterop(list, &command)); + const char* failure = recorded ? nullptr : "BeginVkCommandBufferInterop failed"; + if (recorded) + { + try + { + feature.engineR->record(command, game); + } + catch (const std::exception& error) + { + logf("engine backend (Ray Reconstruction): %s", error.what()); + failure = "the frame could not be recorded"; + recorded = false; + } + if (FAILED(g_vk.interop->EndVkCommandBufferInterop(list)) && recorded) + { + failure = "EndVkCommandBufferInterop failed"; + recorded = false; + } + } + for (size_t i = 0; i < sources.size(); ++i) + if (firstUse(i)) + transition(list, sources[i].resource, sources[i].copyState, sources[i].state); + if (!recorded) + return refuse(failure); + if (frame <= 3 || frame % 600 == 0) + logf("engine frame %u recorded, Ray Reconstruction %c (jitter %.6f,%.6f, reset %d, mvScale %.3f,%.3f)", frame, + presetE ? 'E' : 'D', p.jitterX, p.jitterY, s.reset, p.mvScaleX, p.mvScaleY); + return true; +} + // Records the whole evaluation of `frame` into `list`. False: nothing usable was recorded and the // CUDA backend takes this feature over. static bool engine_evaluate(Feature& feature, ID3D12GraphicsCommandList* list, ID3D12Resource* color, ID3D12Resource* depth, ID3D12Resource* motion, ID3D12Resource* output, - ID3D12Resource* exposure, const FrameParams& p, uint32_t frame) + ID3D12Resource* exposure, void* parameters, const FrameParams& p, uint32_t frame) { if (feature.engineRefused) return false; @@ -8284,6 +8514,8 @@ static bool engine_evaluate(Feature& feature, ID3D12GraphicsCommandList* list, I return false; }; std::call_once(g_engine.vulkanOnce, init_engine_vulkan); + if (feature.rayReconstruction) + return engine_evaluate_rr(feature, list, color, depth, motion, output, parameters, p, frame); if (feature.preset == 13 || feature.preset == 12) return engine_evaluate_m(feature, list, color, depth, motion, output, p, frame); // Without the HDR flag NGX runs its LDR input and output kernels: a model directory of its own. @@ -8659,7 +8891,7 @@ static NgxResult evaluate(ID3D12GraphicsCommandList* list, const NgxHandle* hand } } - if (engine_requested() && engine_evaluate(*feature, list, color, depth, motion, output, exposure, p, frame)) + if (engine_requested() && engine_evaluate(*feature, list, color, depth, motion, output, exposure, parameters, p, frame)) return NGX_SUCCESS; if (feature->cudaCreationDeferred) diff --git a/tools/replay_rr_launch.cpp b/tools/replay_rr_launch.cpp new file mode 100644 index 0000000..cb709c8 --- /dev/null +++ b/tools/replay_rr_launch.cpp @@ -0,0 +1,270 @@ +// Replays one captured Ray Reconstruction (DLSS-D) kernel launch through ZLUDA on Linux, with its textures and +// surfaces: the reference for engine/rr's stage shaders on the GPU the capture came from, and (with debug stores +// added to the PTX) for their intermediates. +// +// usage: replay_rr_launch MODULE.ptx KERNEL MANIFEST +// MANIFEST lines (engine/rr/zluda_replay.py writes them from a capture): +// args FILE the launch's argument block +// grid GX GY GZ / block BX BY BZ +// buf OFFSET FILE DELTA [OUTFILE] upload FILE, store its device address + DELTA at OFFSET, download to OUTFILE +// tex OFFSET FILE W H CHANNELS FORMAT linear|nearest CUarray_format FORMAT, clamp, normalized coordinates +// surf OFFSET W H CHANNELS FORMAT OUTFILE zero-initialised, downloaded after the launch +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using CUresult = int; +using CUdeviceptr = uint64_t; +using CUarray = void*; + +struct ArrayDescriptor3D +{ + size_t Width, Height, Depth; + int Format; + unsigned NumChannels, Flags; +}; +struct Memcpy2D +{ + size_t srcXInBytes, srcY; + int srcMemoryType; + const void* srcHost; + CUdeviceptr srcDevice; + CUarray srcArray; + size_t srcPitch; + size_t dstXInBytes, dstY; + int dstMemoryType; + void* dstHost; + CUdeviceptr dstDevice; + CUarray dstArray; + size_t dstPitch; + size_t WidthInBytes, Height; +}; +struct ResourceDesc +{ + int resType; + union + { + CUarray array; + int reserved[32]; + } res; + unsigned flags; +}; +struct TextureDesc +{ + int addressMode[3]; + int filterMode; + unsigned flags, maxAnisotropy; + int mipmapFilterMode; + float mipmapLevelBias, minMipmapLevelClamp, maxMipmapLevelClamp; + float borderColor[4]; + int reserved[12]; +}; + +static std::vector slurp(const std::string& path) +{ + std::ifstream input(path, std::ios::binary); + if (!input) + { + std::fprintf(stderr, "cannot read %s\n", path.c_str()); + std::exit(1); + } + return {std::istreambuf_iterator(input), std::istreambuf_iterator()}; +} + +static void* g_cuda; +template static F fn(const char* name) +{ + void* raw = dlsym(g_cuda, name); + if (raw == nullptr) + { + std::fprintf(stderr, "libcuda.so has no %s\n", name); + std::exit(1); + } + F f; + std::memcpy(&f, &raw, sizeof(f)); + return f; +} +#define CK(expr) \ + do \ + { \ + CUresult r_ = (expr); \ + if (r_ != 0) \ + { \ + std::fprintf(stderr, "%s:%d: %s -> %d\n", __FILE__, __LINE__, #expr, r_); \ + std::exit(1); \ + } \ + } while (0) + +static size_t texelBytes(int format, unsigned channels) +{ + const size_t component = format == 0x10 ? 2 : format == 0x20 ? 4 : 1; + return format >= 0xb0 ? 4 : component * channels; // the packed int8x4 formats are 4 bytes per texel +} + +int main(int argc, char** argv) +{ + if (argc != 4) + { + std::fprintf(stderr, "usage: %s MODULE.ptx KERNEL MANIFEST\n", argv[0]); + return 2; + } + g_cuda = dlopen("libcuda.so", RTLD_NOW | RTLD_LOCAL); + if (g_cuda == nullptr) + { + std::fprintf(stderr, "dlopen(libcuda.so): %s\n", dlerror()); + return 1; + } + auto init = fn("cuInit"); + auto ctxCreate = fn("cuCtxCreate_v2"); + auto moduleLoad = fn("cuModuleLoadData"); + auto getFunction = fn("cuModuleGetFunction"); + auto memAlloc = fn("cuMemAlloc_v2"); + auto htod = fn("cuMemcpyHtoD_v2"); + auto dtoh = fn("cuMemcpyDtoH_v2"); + auto array3d = fn("cuArray3DCreate_v2"); + auto copy2d = fn("cuMemcpy2D_v2"); + auto texCreate = fn("cuTexObjectCreate"); + auto surfCreate = fn("cuSurfObjectCreate"); + auto launch = fn("cuLaunchKernel"); + auto sync = fn("cuCtxSynchronize"); + + void* context = nullptr; + void* module = nullptr; + void* function = nullptr; + CK(init(0)); + CK(ctxCreate(&context, 0, 0)); + auto ptx = slurp(argv[1]); + ptx.push_back('\0'); + CK(moduleLoad(&module, ptx.data())); + CK(getFunction(&function, module, argv[2])); + + std::vector args; + unsigned grid[3] = {1, 1, 1}, block[3] = {1, 1, 1}; + struct Download + { + std::string path; + CUdeviceptr device = 0; + CUarray array = nullptr; + size_t width = 0, height = 0, texel = 0; + }; + std::vector downloads; + std::vector> uploads; + std::ifstream manifest(argv[3]); + std::string line; + while (std::getline(manifest, line)) + { + std::istringstream in(line); + std::string key; + in >> key; + if (key == "args") + { + std::string file; + in >> file; + args = slurp(file); + } + else if (key == "grid") + in >> grid[0] >> grid[1] >> grid[2]; + else if (key == "block") + in >> block[0] >> block[1] >> block[2]; + else if (key == "buf") + { + size_t offset; + std::string file, out; + int64_t delta; + in >> offset >> file >> delta >> out; + // pointers into one allocation name the same FILE: it is uploaded once + CUdeviceptr device = 0; + for (const auto& [path, address] : uploads) + if (path == file) + device = address; + if (device == 0) + { + auto data = slurp(file); + CK(memAlloc(&device, data.size())); + CK(htod(device, data.data(), data.size())); + uploads.emplace_back(file, device); + if (!out.empty()) + downloads.push_back({out, device, nullptr, data.size(), 1, 1}); + } + const uint64_t value = device + delta; + std::memcpy(args.data() + offset, &value, 8); + } + else if (key == "tex" || key == "surf") + { + size_t offset, width, height; + unsigned channels; + int format; + std::string file, mode, out; + if (key == "tex") + in >> offset >> file >> width >> height >> channels >> format >> mode; + else + in >> offset >> width >> height >> channels >> format >> out; + ArrayDescriptor3D desc{width, height, 0, format, channels, key == "surf" ? 2u : 0u}; + CUarray array = nullptr; + CK(array3d(&array, &desc)); + const size_t texel = texelBytes(format, channels); + std::vector data = key == "tex" ? slurp(file) : std::vector(width * height * texel, 0); + Memcpy2D copy{}; + copy.srcMemoryType = 1; + copy.srcHost = data.data(); + copy.srcPitch = width * texel; + copy.dstMemoryType = 3; + copy.dstArray = array; + copy.WidthInBytes = width * texel; + copy.Height = height; + CK(copy2d(©)); + ResourceDesc resource{}; + resource.resType = 0; + resource.res.array = array; + uint64_t handle = 0; + if (key == "tex") + { + TextureDesc texture{}; + texture.addressMode[0] = texture.addressMode[1] = texture.addressMode[2] = 1; // clamp + texture.filterMode = mode == "linear" ? 1 : 0; + texture.flags = 0x2; // normalized coordinates + CK(texCreate(&handle, &resource, &texture, nullptr)); + } + else + { + CK(surfCreate(&handle, &resource)); + downloads.push_back({out, 0, array, width, height, texel}); + } + std::memcpy(args.data() + offset, &handle, 8); + } + } + size_t argumentBytes = args.size(); + void* extra[] = {reinterpret_cast(1), args.data(), reinterpret_cast(2), &argumentBytes, nullptr}; + CK(launch(function, grid[0], grid[1], grid[2], block[0], block[1], block[2], 0, nullptr, nullptr, extra)); + CK(sync()); + for (const Download& d : downloads) + { + std::vector data(d.width * d.height * d.texel); + if (d.array == nullptr) + CK(dtoh(data.data(), d.device, data.size())); + else + { + Memcpy2D copy{}; + copy.srcMemoryType = 3; + copy.srcArray = d.array; + copy.dstMemoryType = 1; + copy.dstHost = data.data(); + copy.dstPitch = d.width * d.texel; + copy.WidthInBytes = d.width * d.texel; + copy.Height = d.height; + CK(copy2d(©)); + } + std::ofstream(d.path, std::ios::binary).write(data.data(), static_cast(data.size())); + } + std::printf("%s grid %u,%u,%u block %u,%u,%u: done\n", argv[2], grid[0], grid[1], grid[2], block[0], block[1], + block[2]); + return 0; +} diff --git a/tools/wine_nvcuda_bridge.c b/tools/wine_nvcuda_bridge.c index c89cb9c..c07b6e4 100644 --- a/tools/wine_nvcuda_bridge.c +++ b/tools/wine_nvcuda_bridge.c @@ -2795,8 +2795,9 @@ CUresult WINAPI cuSurfObjectCreate(CUsurfObject* object, const void* descriptor) if (redirected && (result = write_redirect_tail(*object, pointer, pitch, 0)) != CUDA_SUCCESS) cuSurfObjectDestroy(*object); } - TRACE_CALL(result, "cuSurfObjectCreate descriptor=%p result=%d object=0x%llx", descriptor, result, - (unsigned long long)(object != NULL ? *object : 0)); + TRACE_CALL(result, "cuSurfObjectCreate descriptor=%p array=%p result=%d object=0x%llx", descriptor, + result == CUDA_SUCCESS && object != NULL ? (void*)find_surface_array(*object) : NULL, result, + (unsigned long long)(object != NULL ? *object : 0)); return result; } diff --git a/vulkan/kvk.cpp b/vulkan/kvk.cpp index 400d5cf..2a01c7c 100644 --- a/vulkan/kvk.cpp +++ b/vulkan/kvk.cpp @@ -10,6 +10,7 @@ // 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) +// inout FILE OUTFILE (storage buffer initialised from FILE, bound like outbuf, written to OUTFILE) // bindings N0 N1 ... (optional: the binding number of each of the above, in that order, for shaders // whose bindings are not ptx2glsl.py's; e.g. engine/layer_m.comp's L stages) #include @@ -75,7 +76,10 @@ int main(int argc, char** argv) 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}; + // FP8 cooperative matrices (engine/rr's translated stages, engine/layer_m.comp) + VkPhysicalDeviceShaderFloat8FeaturesEXT f8{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_FLOAT8_FEATURES_EXT}; + VkPhysicalDeviceCooperativeMatrixFeaturesKHR cm{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_MATRIX_FEATURES_KHR, &f8}; + VkPhysicalDeviceVulkan13Features v13{VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES, &cm}; 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}; @@ -84,7 +88,12 @@ int main(int argc, char** argv) 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; + std::vector exts; + if (cm.cooperativeMatrix) exts.push_back(VK_KHR_COOPERATIVE_MATRIX_EXTENSION_NAME); + if (f8.shaderFloat8) exts.push_back(VK_EXT_SHADER_FLOAT8_EXTENSION_NAME); + if (!cm.cooperativeMatrix) v13.pNext = nullptr; else if (!f8.shaderFloat8) cm.pNext = nullptr; VkDeviceCreateInfo dci{VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO, &f2}; dci.queueCreateInfoCount = 1; dci.pQueueCreateInfos = &qci; + dci.enabledExtensionCount = static_cast(exts.size()); dci.ppEnabledExtensionNames = exts.data(); 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; @@ -126,6 +135,12 @@ int main(int argc, char** argv) 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 == "inout") + { + std::string f, o; ls >> f >> o; 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(o); + } else if (k == "buf") { int off; std::string f; int64_t delta; ls >> off >> f >> delta; @@ -197,6 +212,8 @@ int main(int argc, char** argv) 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}; + // KVK_DISABLE_OPTIMIZATION=1: VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT (RADV: ACO without its optimizer/scheduler) + if (getenv("KVK_DISABLE_OPTIMIZATION")) cpi.flags |= VK_PIPELINE_CREATE_DISABLE_OPTIMIZATION_BIT; 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; diff --git a/vulkan/ptx2glsl.py b/vulkan/ptx2glsl.py index 584bb7e..1df4404 100644 --- a/vulkan/ptx2glsl.py +++ b/vulkan/ptx2glsl.py @@ -287,7 +287,7 @@ class T: def stmt(self, s): if s.startswith('.reg'): - m = re.match(r'\.reg\s+\.(\w+)\s+(.*)$', s) + m = re.match(r'\.reg\s*\.(\w+)\s*(.*)$', s) ty, names = m.groups() if '<' in names: return @@ -714,10 +714,7 @@ class T: raise Err('store ' + '.'.join(o)) # ---------------------------------------------------------------- output - def emit(self, tex_filter): - nb = len(self.blocks) - out = [] - w = out.append + def preamble(self, w, tex_filter): 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.') @@ -741,7 +738,15 @@ class T: 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) + def runtime(self): + return RUNTIME + + def emit(self, tex_filter): + nb = len(self.blocks) + out = [] + w = out.append + self.preamble(w, tex_filter) + w(self.runtime()) w('void main()\n{') if getattr(self, 'local_bytes', 0): w(f' uint loc[{(self.local_bytes + 3) // 4}];') @@ -749,7 +754,7 @@ class T: 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]) + ';') + w(f' {"precise " if (g in ("float", "float16_t", "f16vec2") or getattr(self, "precise_all", False) and g != "bool") 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): @@ -780,26 +785,45 @@ class T: 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) + for h in headers: + g[cont[h]] = [END] # provisional, replaced below + def postorder(): + 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) + return rpo + # A loop may leave through its header as well as through its latches. Its body is the set of nodes that + # reach a latch without passing the header; the loop region ends at the nearest node that post-dominates + # every exit of the body (for a loop with one latch exit this is that exit, as before). + preds = {} + for n in nodes: + for m in succ[n]: + preds.setdefault(m, []).append(n) + pdom = {END: {END}} + for n in postorder(): # successors first + if n != END: + pdom[n] = {n} | set.intersection(*(pdom[m] for m in g[n])) + for h in headers: + body, work = {h}, [n for n, m in back if m == h] + while work: + n = work.pop() + if n not in body: + body.add(n); work += preds.get(n, []) + exits = {m for n in body for m in succ[n] if m not in body} + if not exits: + raise Err('loop without an exit') + common = set.intersection(*(pdom[e] for e in exits)) + g[cont[h]] = [max(common, key=lambda c: len(pdom[c]))] + rpo = postorder() topo = rpo[::-1] # entry first index = {n: i for i, n in enumerate(topo)} ipdom = {END: END} @@ -928,21 +952,10 @@ void sust_b(writeonly uimage2D s, int xb, int y, uint v) } ''' -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]) +def build(make, lut, fast, no_half): + t = make() t.exp_lut = lut - if '--no-half-types' not in args: + if not no_half: t.collect = True t.translate() parent = {} @@ -961,8 +974,24 @@ if __name__ == '__main__': 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 = make(); t.exp_lut = lut2; t.half = half t.fast = fast t.translate() + return t + + +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 = build(lambda: T(src, args[1]), lut, fast, '--no-half-types' in args) 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')