Files
d4r/engine/hip_net.cpp
T
ayo 7c5d3c9157 feat: add d4r inference engines and Windows parity tooling
Add native K/M/RR inference engines, HIP kernels, replay and validation
harnesses, and Windows DLSS setup and runtime support. Include GPU
profiling tools, build and packaging scripts, platform integration
patches, and implementation documentation.
2026-10-16 08:54:47 -04:00

333 lines
18 KiB
C++

// d4r engine: the DLSS network on d4r's native HIP layer kernels. See hip_net.h.
//
// The backend is shared by the two presets and only the plan differs: preset M runs ten layers with the
// kernels/m parameter block (kernels/m/swin_common.h) in two window-alignment phases, preset K runs eleven
// layers with kernels/k/pwin_common.h's PwinParams (the bottleneck layer as its four phase kernels) in the
// eight-frame window cycle of makeKWindows. Everything else - the stream, the immutable captured graphs, the
// events and the per-evaluation frame counter - is the same code.
#define __HIP_PLATFORM_AMD__
#include "hip_net.h"
#include <hip/hip_runtime_api.h>
#include <algorithm>
#include <cstring>
#include <fstream>
#include <stdexcept>
#include <vector>
namespace d4r {
namespace {
void check(hipError_t e, const char* what) {
if (e != hipSuccess) throw std::runtime_error(std::string("HIP: ") + what + " failed: " + hipGetErrorString(e));
}
#define HIPCHECK(x) check((x), #x)
std::vector<char> slurp(const std::string& path) {
std::ifstream f(path, std::ios::binary);
if (!f) throw std::runtime_error("cannot read " + path);
return std::vector<char>((std::istreambuf_iterator<char>(f)), {});
}
// preset M's kernels' parameter block (kernels/m/swin_common.h: CommonParams, TubeParams). Both blocks are
// plain aggregates: every parameter set is value-initialized (all bytes zero) before its fields are filled.
struct LayerParamsM {
void* w; int32_t sx, sy, tw, th; void* in; void* skip; void* out; void* merged;
int32_t r69, r70, r71, r72; void* inFlags; void* outFlags;
};
static_assert(sizeof(LayerParamsM) == 88 && offsetof(LayerParamsM, in) == 24, "M param block");
// preset K's kernels' parameter block (kernels/k/pwin_layer.h: PwinParams)
struct LayerParamsK {
int32_t W, H; void* in; void* skip; void* out24; void* out32; uint64_t pad40, pad48; int32_t sx, sy; void* w;
uint8_t rest[104];
};
static_assert(sizeof(LayerParamsK) == 176 && offsetof(LayerParamsK, in) == 8 && offsetof(LayerParamsK, w) == 64, "K param block");
union Params {
LayerParamsM m;
LayerParamsK k;
uint8_t raw[176];
};
// one code object of a layer: the main entry, the weight-image preparation, and the K bottleneck's post phases
struct Kernel {
hipModule_t module = nullptr;
hipFunction_t main = nullptr, prep = nullptr;
hipFunction_t post[8] = {};
uint32_t prepBlocks = 0, blockZ = 4, gridX = 0, gridY = 0;
uint32_t postGridX[8] = {}, postGridY[8] = {}, postBlockX[8] = {}, postBlockY[8] = {}, postBlockZ[8] = {};
int postCount = 0;
};
// one immutable launch: entry point, grid, block and the (frozen) parameter block
struct Node {
hipFunction_t function = nullptr;
uint32_t gx = 1, gy = 1, bx = 32, by = 1, bz = 1;
Params params{};
};
uint32_t global32(hipModule_t module, const char* name, uint32_t fallback) {
hipDeviceptr_t at = nullptr; size_t bytes = 0; uint32_t v = fallback;
if (hipModuleGetGlobal(&at, &bytes, module, name) == hipSuccess && bytes == 4) {
if (hipMemcpyDtoH(&v, at, 4) != hipSuccess || v == 0) v = fallback;
}
return v;
}
void* import(int fd, uint64_t bytes, hipExternalMemory_t& memory) {
hipExternalMemoryHandleDesc desc{}; desc.type = hipExternalMemoryHandleTypeOpaqueFd; desc.handle.fd = fd; desc.size = bytes;
HIPCHECK(hipImportExternalMemory(&memory, &desc));
hipExternalMemoryBufferDesc buffer{}; buffer.offset = 0; buffer.size = bytes;
void* at = nullptr; HIPCHECK(hipExternalMemoryGetMappedBuffer(&at, memory, &buffer));
return at;
}
// makeKWindows (engine/runtime.cpp): dec5 moves once per frame around an eight-frame cycle and every
// shallower layer follows the one below it as (2 * shift + 4) mod 8. Role 0 = dec5, 1 = enc4/dec4,
// 2 = enc3/dec3; every other layer keeps shift 4.
constexpr int32_t KCycle[8][2] = {{0, 2}, {5, 6}, {4, 0}, {1, 4}, {7, 5}, {2, 1}, {3, 7}, {6, 3}};
void kShift(int role, uint32_t frame, int32_t& sx, int32_t& sy) {
if (role < 0) { sx = sy = 4; return; }
int32_t x = KCycle[frame % 8][0], y = KCycle[frame % 8][1];
for (int depth = 0; depth < role; ++depth) { x = (2 * x + 4) & 7; y = (2 * y + 4) & 7; }
sx = x; sy = y;
}
// preset K's eleven layers in network order: token-grid divisor, input channels, core channels, the channels
// written to out24 (the patch-merged skip for encoders, the block output for the decoders, the 40-channel head
// for dec0) and the window-schedule role.
struct KLayer {
const char* name; uint32_t divisor, cin, c, out24ch; int role;
};
constexpr KLayer KPlan[11] = {
{"enc0", 1, 16, 32, 64, -1}, {"enc1", 2, 64, 64, 64, -1}, {"enc2", 4, 64, 64, 96, -1},
{"enc3", 8, 96, 96, 128, 2}, {"enc4", 16, 128, 128, 160, 1}, {"dec5", 32, 160, 160, 160, 0},
{"dec4", 16, 160, 128, 128, 1}, {"dec3", 8, 128, 96, 96, 2}, {"dec2", 4, 96, 64, 64, -1},
{"dec1", 2, 64, 64, 64, -1}, {"dec0", 1, 64, 32, 40, -1},
};
constexpr uint32_t KLayers = 11, MLayers = 10, MaxGraphs = 8;
} // namespace
struct HipNetwork::Impl {
std::vector<Kernel> kernels;
std::vector<Node> preps;
std::vector<Node> nodes[MaxGraphs];
std::vector<void*> owned;
hipExternalMemory_t shared[2]{};
hipStream_t stream = nullptr;
hipGraph_t graph[MaxGraphs]{};
hipGraphExec_t executable[MaxGraphs]{};
uint32_t graphs = 2, frame = 0;
hipEvent_t start = nullptr, stop = nullptr;
bool pending = false;
~Impl() {
if (stream) (void)hipStreamSynchronize(stream);
for (auto e : executable) if (e) (void)hipGraphExecDestroy(e);
for (auto g : graph) if (g) (void)hipGraphDestroy(g);
for (void* p : owned) (void)hipFree(p);
for (auto& k : kernels) if (k.module) (void)hipModuleUnload(k.module);
for (auto m : shared) if (m) (void)hipDestroyExternalMemory(m);
if (start) (void)hipEventDestroy(start);
if (stop) (void)hipEventDestroy(stop);
if (stream) (void)hipStreamDestroy(stream);
}
void run(const Node& n) {
void* args[] = {const_cast<Params*>(&n.params)};
HIPCHECK(hipModuleLaunchKernel(n.function, n.gx, n.gy, 1, n.bx, n.by, n.bz, 0, stream, args, nullptr));
}
};
HipNetwork::HipNetwork(const std::string& directory, uint32_t width, uint32_t height, int tokensFd, int outputFd,
uint64_t tokensBytes, uint64_t outputBytes) : impl(std::make_unique<Impl>()) {
auto& n = *impl;
HIPCHECK(hipInit(0));
int leastPriority = 0, greatestPriority = 0;
HIPCHECK(hipDeviceGetStreamPriorityRange(&leastPriority, &greatestPriority));
HIPCHECK(hipStreamCreateWithPriority(&n.stream, hipStreamNonBlocking, greatestPriority));
HIPCHECK(hipEventCreate(&n.start)); HIPCHECK(hipEventCreate(&n.stop));
const auto file = slurp(directory + "/hipnet.bin");
if (file.size() < 16) throw std::runtime_error("incompatible hipnet.bin");
const bool kPreset = std::memcmp(file.data(), "D4RKHIP1", 8) == 0;
if (!kPreset && std::memcmp(file.data(), "D4RMHIP1", 8)) throw std::runtime_error("incompatible hipnet.bin");
uint32_t count = 0;
std::memcpy(&count, file.data() + 8, 4);
const uint32_t layers = kPreset ? KLayers : MLayers;
if (count != layers || file.size() < 16 + 64ull * count) throw std::runtime_error("invalid hipnet.bin");
void* tokens = import(tokensFd, tokensBytes, n.shared[0]);
void* output = import(outputFd, outputBytes, n.shared[1]);
// Both dimensions describe the layer-0 token grid: preset M's render dimensions, preset K's output
// dimensions (makeKShape: max 256, align 32 of the output quarter resolution).
uint32_t tw = 0, th = 0;
if (kPreset) {
if (width < 256 || height < 256 || width > 8192 || height > 8192)
throw std::runtime_error("standard K output dimensions must be 256..8192");
tw = std::max(256u, ((width + 3) / 4 + 31) / 32 * 32);
th = std::max(256u, ((height + 3) / 4 + 31) / 32 * 32);
} else {
tw = (width + 31) / 32 * 16; th = (height + 31) / 32 * 16;
}
// Kernel lookup by name; the M tube shares one code object across six layers, the K layer names are unique.
std::vector<std::string> names;
auto kernelOf = [&](const char* name) -> uint32_t {
size_t k = 0;
while (k < names.size() && names[k] != name) ++k;
if (k < names.size()) return uint32_t(k);
names.emplace_back(name); n.kernels.emplace_back();
auto& kernel = n.kernels.back();
const auto code = slurp(directory + "/hip/" + name + ".hsaco");
HIPCHECK(hipModuleLoadData(&kernel.module, code.data()));
HIPCHECK(hipModuleGetFunction(&kernel.main, kernel.module, name));
const std::string base(name);
if (hipModuleGetFunction(&kernel.prep, kernel.module, (base + "_prep").c_str()) != hipSuccess) kernel.prep = nullptr;
kernel.prepBlocks = global32(kernel.module, "d4r_prep_blocks", 0);
kernel.blockZ = global32(kernel.module, "d4r_block_z", 4);
kernel.gridX = global32(kernel.module, "d4r_grid_x", 0);
kernel.gridY = global32(kernel.module, "d4r_grid_y", 0);
if (kernel.prep && !kernel.prepBlocks) throw std::runtime_error(base + ": prep kernel without d4r_prep_blocks");
// The deep layers' code objects also carry their post phases; their grids are compile-time.
for (int j = 1; j <= 8; ++j) {
const std::string tag = base + "_post" + std::to_string(j);
if (hipModuleGetFunction(&kernel.post[kernel.postCount], kernel.module, tag.c_str()) != hipSuccess) break;
const std::string suffix = "_post" + std::to_string(j);
kernel.postGridX[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_grid_x").c_str(), 0);
kernel.postGridY[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_grid_y").c_str(), 0);
kernel.postBlockX[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_x").c_str(), 32);
kernel.postBlockY[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_y").c_str(), 1);
kernel.postBlockZ[kernel.postCount] = global32(kernel.module, ("d4r" + suffix + "_block_z").c_str(), 1);
if (!kernel.postGridX[kernel.postCount]) kernel.postGridX[kernel.postCount] = kernel.gridX;
++kernel.postCount;
}
// a module keeps one prepared weight image per weights pointer while its slots last; the M tube needs six
if (global32(kernel.module, "d4r_prep_key_slots", 0) < 6 && base.find("tube") != std::string::npos)
throw std::runtime_error(base + ": fewer than six weight image slots");
return uint32_t(names.size() - 1);
};
auto allocate = [&](uint64_t bytes) {
void* p = nullptr; HIPCHECK(hipMalloc(&p, size_t(bytes) + 4096)); n.owned.push_back(p); return p;
};
std::vector<uint32_t> layerKernel(layers), divisor(layers);
std::vector<int32_t> shiftX(layers), shiftY(layers);
std::vector<void*> weights(layers);
for (uint32_t i = 0; i < count; ++i) {
const char* e = file.data() + 16 + 64 * i;
char name[33]{}; std::memcpy(name, e, 32);
int32_t sx, sy; uint32_t scale; uint64_t at, bytes;
std::memcpy(&sx, e + 32, 4); std::memcpy(&sy, e + 36, 4); std::memcpy(&scale, e + 40, 4);
std::memcpy(&at, e + 48, 8); std::memcpy(&bytes, e + 56, 8);
if (!scale || at > file.size() || bytes > file.size() - at) throw std::runtime_error("invalid hipnet.bin layer");
if (kPreset && std::strcmp(name, ("dltss_pwin_" + std::string(KPlan[i].name) + "_layer").c_str()))
throw std::runtime_error(std::string("unexpected K layer order: ") + name);
layerKernel[i] = kernelOf(name);
void* w = allocate(bytes);
HIPCHECK(hipMemcpyHtoD(w, file.data() + at, bytes));
weights[i] = w; divisor[i] = scale; shiftX[i] = sx; shiftY[i] = sy;
}
// The layer plan: activation buffers and the frozen parameter block of every layer at frame/phase 0.
std::vector<Params> base(layers);
if (kPreset) {
// NGX packs the encoder skips, the merged outputs and the decoder outputs into one scratch allocation;
// separate buffers only have to respect the channel counts and the wiring, not the original aliasing.
std::vector<void*> act(layers), full(5, nullptr);
for (uint32_t i = 0; i < layers; ++i) {
const auto& p = KPlan[i];
const uint32_t w = tw / p.divisor, h = th / p.divisor;
// an encoder's out24 is the patch-merged 2x2 output, half the resolution of its input
const uint32_t ow = i < 5 ? tw / (2 * p.divisor) : w, oh = i < 5 ? th / (2 * p.divisor) : h;
// dec0 writes the head straight into the caller's output buffer
if (i + 1 != layers) act[i] = allocate(uint64_t(ow) * oh * p.out24ch * 2);
if (i < 5) full[i] = allocate(uint64_t(w) * h * p.c * 2);
}
for (uint32_t i = 0; i < layers; ++i) {
const auto& p = KPlan[i];
auto& q = base[i].k;
q.W = int32_t(tw / p.divisor); q.H = int32_t(th / p.divisor);
q.in = i == 0 ? tokens : act[i - 1];
q.skip = i >= 6 ? full[10 - i] : nullptr;
q.out24 = i + 1 == layers ? output : act[i];
q.out32 = i < 5 ? full[i] : nullptr;
kShift(p.role, 0, q.sx, q.sy);
q.w = weights[i];
}
} else {
void* skip1 = allocate(uint64_t(tw) * th * 64);
void* merged1 = allocate(uint64_t(tw / 2) * (th / 2) * 96);
void* skip2 = allocate(uint64_t(tw / 2) * (th / 2) * 96);
void* merged2 = allocate(uint64_t(tw / 4) * (th / 4) * 128);
void* tubeA = allocate(uint64_t(tw / 4) * (th / 4) * 128);
void* tubeB = allocate(uint64_t(tw / 4) * (th / 4) * 128);
void* dec2 = allocate(uint64_t(tw / 2) * (th / 2) * 96);
// tokens -> enc1 -> enc2 -> six tube blocks (ping-pong) -> dec2 -> dec1
void* in[MLayers] = {tokens, merged1, merged2, tubeA, tubeB, tubeA, tubeB, tubeA, tubeB, dec2};
void* skip[MLayers] = {nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, skip2, skip1};
void* out[MLayers] = {skip1, skip2, tubeA, tubeB, tubeA, tubeB, tubeA, tubeB, dec2, output};
void* merged[MLayers] = {merged1, merged2, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr};
for (uint32_t i = 0; i < layers; ++i) {
auto& q = base[i].m;
q.w = weights[i]; q.sx = shiftX[i]; q.sy = shiftY[i];
q.tw = int32_t(tw / divisor[i]); q.th = int32_t(th / divisor[i]);
q.in = in[i]; q.skip = skip[i]; q.out = out[i]; q.merged = merged[i];
}
}
// Immutable graphs: preset M has two horizontal window-alignment phases, preset K the eight-frame cycle.
n.graphs = kPreset ? MaxGraphs : 2;
for (uint32_t phase = 0; phase < n.graphs; ++phase) {
auto& nodes = n.nodes[phase];
for (uint32_t i = 0; i < layers; ++i) {
const auto& kernel = n.kernels[layerKernel[i]];
Params params = base[i];
if (kPreset) kShift(KPlan[i].role, phase, params.k.sx, params.k.sy);
else params.m.sx = shiftX[i] ^ int32_t(phase * 2); // NGX alternates horizontal alignment every evaluation
if (kernel.postCount) {
Node node; node.function = kernel.main; node.params = params;
node.gx = kernel.gridX; node.gy = kernel.gridY ? kernel.gridY : 1; node.bz = kernel.blockZ;
nodes.push_back(node);
for (int j = 0; j < kernel.postCount; ++j) {
Node post; post.function = kernel.post[j]; post.params = params;
post.gx = kernel.postGridX[j]; post.gy = kernel.postGridY[j] ? kernel.postGridY[j] : 1;
post.bx = kernel.postBlockX[j]; post.by = kernel.postBlockY[j]; post.bz = kernel.postBlockZ[j];
nodes.push_back(post);
}
} else {
Node node; node.function = kernel.main; node.params = params;
if (kPreset) {
node.gx = (params.k.W + params.k.sx + 7) / 8; node.gy = (params.k.H + params.k.sy + 7) / 8;
} else {
node.gx = (params.m.tw + params.m.sx + 7) / 8; node.gy = (params.m.th + params.m.sy + 7) / 8;
}
node.bz = kernel.blockZ;
nodes.push_back(node);
}
}
}
// The weights never change: every layer's weight image is prepared once, here, and found by its pointer.
// A shared code object (the M tube) gets one preparation per layer, i.e. per distinct weights pointer.
for (uint32_t i = 0; i < layers; ++i) {
const auto& kernel = n.kernels[layerKernel[i]];
if (!kernel.prep) continue;
Node prep; prep.function = kernel.prep; prep.gx = kernel.prepBlocks; prep.bx = 128; prep.params = base[i];
n.preps.push_back(prep);
}
for (const auto& prep : n.preps) n.run(prep);
HIPCHECK(hipStreamSynchronize(n.stream));
for (uint32_t phase = 0; phase < n.graphs; ++phase) {
HIPCHECK(hipStreamBeginCapture(n.stream, hipStreamCaptureModeRelaxed));
try {
for (const auto& node : n.nodes[phase]) n.run(node);
} catch (...) {
(void)hipStreamEndCapture(n.stream, &n.graph[phase]);
throw;
}
HIPCHECK(hipStreamEndCapture(n.stream, &n.graph[phase]));
HIPCHECK(hipGraphInstantiate(&n.executable[phase], n.graph[phase], nullptr, nullptr, 0));
HIPCHECK(hipGraphUpload(n.executable[phase], n.stream));
}
HIPCHECK(hipStreamSynchronize(n.stream));
}
HipNetwork::~HipNetwork() = default;
void HipNetwork::launch() {
auto& n = *impl;
HIPCHECK(hipEventRecord(n.start, n.stream));
HIPCHECK(hipGraphLaunch(n.executable[n.frame % n.graphs], n.stream));
++n.frame;
HIPCHECK(hipEventRecord(n.stop, n.stream));
n.pending = true;
}
bool HipNetwork::done() {
auto& n = *impl;
if (!n.pending) return true;
const hipError_t e = hipEventQuery(n.stop);
if (e == hipErrorNotReady) return false;
check(e, "hipEventQuery"); n.pending = false; return true;
}
void HipNetwork::wait() { HIPCHECK(hipStreamSynchronize(impl->stream)); impl->pending = false; }
float HipNetwork::lastMilliseconds() { float ms = 0; HIPCHECK(hipEventElapsedTime(&ms, impl->start, impl->stop)); return ms; }
} // namespace d4r