// Preset M runtime; included by runtime.cpp inside namespace d4r (it shares that file's Vulkan helpers). // GCC treats an included file as a header and warns that these types use the including file's anonymous namespace. #pragma GCC diagnostic ignored "-Wsubobject-linkage" namespace { constexpr unsigned MLayers = 10; constexpr const char* MShader[MLayers] = {"m_enc1", "m_enc2", "m_tube", "m_tube", "m_tube", "m_tube", "m_tube", "m_tube", "m_dec2", "m_dec1"}; constexpr int32_t MShift[MLayers] = {2, 2, 2, 0, 2, 0, 2, 0, 0, 0}; constexpr unsigned MScale[MLayers] = {1, 2, 4, 4, 4, 4, 4, 4, 2, 1}; // token grid divisor constexpr unsigned MChannels[MLayers] = {64, 96, 128, 128, 128, 128, 128, 128, 96, 64}; struct MEnc0Params { int32_t size[2], shift[2], tok[2]; float jitter[2], prevJitter[2], mvScale[2], mvOffs[2], rnd2hi[2], hi2rnd[2], hiSize[2]; float changeScale, changeGain; uint32_t depthFlag, reset, frame; float expo, expoPrev; float halfInv[2], renderInv[2]; // CUDA texture coordinates use host-rounded reciprocals }; struct MDec0Params { int32_t size[2], tokW; uint32_t reset; float rndOffs[2], rnd2hi[2], hiSize[2], histOffs[2], sigma, rho, renderInv[2]; }; struct MPostParams { int32_t colOrigin[2], hiSize[2], rndSize[2]; float hi2rnd[2], rndOffs[2], invRnd[2], gain; uint32_t reset; float expo, expoPrev; }; struct MDownParams { int32_t srcSize[2], outOrigin[2], colOrigin[2], colSize[2]; float out2rnd[2], rndOffs[2], expo; }; struct MLayerParams { int32_t w, h, sx, sy; uint32_t o[72]; }; // The grid NGX measures the exposure on (cuda_dldn_engine_luma_convert_kernel's launch, presets M and L alike, // harness captures 2026-10-11): half the output, unless the output is more than twice the render size on that axis // (Ultra Performance), then half the render size but at least 256 (640x360 -> 1920x1080: 320x256, 854x480 -> // 2560x1440: 427x256, 1280x720 -> 3840x2160: 640x360). uint32_t exposureGrid(uint32_t render, uint32_t output) { return output / 2 <= render ? output / 2 : std::max(render / 2, 256u); } static_assert(sizeof(MEnc0Params) == 124 && offsetof(MEnc0Params, expo) == 100 && offsetof(MEnc0Params, halfInv) == 108, "enc0_m.comp and exposure_k.comp (words 25, 26)"); static_assert(sizeof(MDec0Params) == 64, "dec0_m.comp"); static_assert(sizeof(MPostParams) == 64 && offsetof(MPostParams, expo) == 56, "post_m.comp and exposure_k.comp (words 14, 15)"); static_assert(sizeof(MDownParams) == 52 && offsetof(MDownParams, expo) == 48, "down_m.comp and exposure_k.comp (word 12)"); static_assert(sizeof(MLayerParams) == 304, "layer_m.comp"); struct Tex { const Device* device = nullptr; VkImage handle = VK_NULL_HANDLE; VkDeviceMemory memory = VK_NULL_HANDLE; VkImageView view = VK_NULL_HANDLE; ~Tex() { if (!device) return; if (view) vkDestroyImageView(device->handle, view, nullptr); if (handle) vkDestroyImage(device->handle, handle, nullptr); if (memory) vkFreeMemory(device->handle, memory, nullptr); } void create(const Device& d, uint32_t w, uint32_t h, VkFormat format) { device = &d; if (!w || !h || w > d.properties.limits.maxImageDimension2D || h > d.properties.limits.maxImageDimension2D) throw std::runtime_error("image dimensions exceed device limits"); VkImageCreateInfo info{VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO}; info.imageType = VK_IMAGE_TYPE_2D; info.format = format; info.extent = {w, h, 1}; info.mipLevels = info.arrayLayers = 1; info.samples = VK_SAMPLE_COUNT_1_BIT; info.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT; VKCHECK(vkCreateImage(d.handle, &info, nullptr, &handle)); VkMemoryRequirements requirements; vkGetImageMemoryRequirements(d.handle, handle, &requirements); memory = d.allocate(requirements, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT); VKCHECK(vkBindImageMemory(d.handle, handle, memory, 0)); VkImageViewCreateInfo vi{VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO}; vi.image = handle; vi.viewType = VK_IMAGE_VIEW_TYPE_2D; vi.format = format; vi.subresourceRange = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; VKCHECK(vkCreateImageView(d.handle, &vi, nullptr, &view)); } }; void bindRange(VkDevice device, VkDescriptorSet set, uint32_t binding, const Buffer& buffer, VkDeviceSize offset) { VkDescriptorBufferInfo info{buffer.handle, offset, VK_WHOLE_SIZE}; VkWriteDescriptorSet write{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET}; write.dstSet = set; write.dstBinding = binding; write.descriptorCount = 1; write.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; write.pBufferInfo = &info; vkUpdateDescriptorSets(device, 1, &write, 0, nullptr); } } // namespace struct MModel::Impl { Device device; Buffer w8, w16; std::vector bytes8, bytes16; // until recordUpload uint32_t enc0[3]{}, dec0[2]{}; std::array, MLayers> layers{}; std::array network; // one per layer (five distinct shaders) Pipeline input, expand, post, postQuality, postCompact, down, downRgb10; Pipeline inputPlanes, expandPlanes; // token buffers in NGX's plane layout, for an external (HIP) network bool planes = false; std::array exposure; bool uploaded = false, ldr = false; // preset L (offsets.bin D4RL0002, L_NOTES.md): the input and expansion stages are layer_m.comp's KIND 3 and 4, // with their own layer offset blocks; constants: change scale and gain, sigma, rho, the first gate's gain bool presetL = false, dynamicResolution = false; std::array lEnc0{}, lDec0{}; uint32_t dither16 = 0; float lConst[5]{}; Impl(VkPhysicalDevice physical, VkDevice dev) : device(physical, dev) {} }; MModel::MModel(VkPhysicalDevice physical, VkDevice device, const std::string& directory, VkPipelineCache cache) : impl(std::make_unique(physical, device)) { auto& m = *impl; m.ldr = exists(directory + "/ldr"); // the stages were compiled for a game without the HDR flag: no exposure auto head = read(directory + "/offsets.bin"); m.presetL = head.size() >= 8 && (!std::memcmp(head.data(), "D4RL0002", 8) || !std::memcmp(head.data(), "D4RL0001", 8)); m.dynamicResolution = head.size() >= 8 && (!std::memcmp(head.data(), "D4RM0004", 8) || !std::memcmp(head.data(), "D4RL0002", 8)); const size_t expect = m.presetL ? 40 + (MLayers + 2) * 74 * 4 : 16 + 5 * 4 + MLayers * 74 * 4; if (head.size() != expect || (!m.presetL && !m.dynamicResolution && std::memcmp(head.data(), "D4RM0003", 8))) throw std::runtime_error("incompatible M/L model metadata"); uint32_t n8 = 0, n16 = 0; std::memcpy(&n8, head.data() + 8, 4); std::memcpy(&n16, head.data() + 12, 4); auto b8 = read(directory + "/weights8.bin"), b16 = read(directory + "/weights16.bin"); if (b8.size() != n8 || b16.size() != n16 || n16 % 64) throw std::runtime_error("invalid M model sizes"); auto inside = [&](uint32_t off8, uint32_t off16) { if (off8 >= n8 || off16 >= n16 || off8 % 64 || off16 % 64) throw std::runtime_error("invalid M model offsets"); }; if (m.presetL) { // the input stage's block, the ten layers, the expansion stage's block std::memcpy(&m.dither16, head.data() + 16, 4); std::memcpy(m.lConst, head.data() + 20, sizeof(m.lConst)); std::memcpy(m.lEnc0.data(), head.data() + 40, sizeof(m.lEnc0)); std::memcpy(m.layers.data(), head.data() + 40 + sizeof(m.lEnc0), sizeof(m.layers)); std::memcpy(m.lDec0.data(), head.data() + 40 + sizeof(m.lEnc0) + sizeof(m.layers), sizeof(m.lDec0)); inside(m.lEnc0[0], m.lEnc0[1]); inside(m.lDec0[0], m.lDec0[1]); inside(0, m.dither16); for (float c : m.lConst) if (!(c > 0.f && c < 100.f)) throw std::runtime_error("invalid L model constants"); } else { std::memcpy(m.enc0, head.data() + 16, 12); std::memcpy(m.dec0, head.data() + 28, 8); std::memcpy(m.layers.data(), head.data() + 36, sizeof(m.layers)); inside(m.enc0[0], m.enc0[1]); inside(m.dec0[0], m.dec0[1]); inside(0, m.enc0[2]); } for (auto& l : m.layers) inside(l[0], l[1]); const VkBufferUsageFlags usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; m.w8.create(m.device, b8.size(), usage); m.w16.create(m.device, b16.size(), usage); m.bytes8 = b8; m.bytes16 = b16; const auto b = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, t = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, s = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE; for (unsigned i = 0; i < MLayers; ++i) m.network[i].create(device, directory + "/" + MShader[i] + ".spv", {b,b,b,b,b,b,b}, 32, cache); // layer_m.comp's L stages: bindings 0-6 of a layer, then the stage's own (layer_m.comp KIND 3 and 4) const std::vector lIn{b,b,b,b,b,b,b,b,t,t,t,t,t,t,s,s,s,b}, lEx{b,b,b,b,b,b,b,b,s,s,s,s}; const std::string in = m.presetL ? "enc0_l" : "enc0_m", ex = m.presetL ? "dec0_l" : "dec0_m", po = m.presetL ? "post_l" : "post_m"; const auto inTypes = m.presetL ? lIn : std::vector{b,t,t,t,t,t,t,s,s,s,b,b,b,b}; const auto exTypes = m.presetL ? lEx : std::vector{b,t,t,s,s,s,s,b,b,b}; m.input.create(device, directory + "/" + in + ".spv", inTypes, 32, cache); m.expand.create(device, directory + "/" + ex + ".spv", exTypes, 32, cache); if ((m.planes = exists(directory + "/" + in + "_planes.spv") && exists(directory + "/" + ex + "_planes.spv"))) { m.inputPlanes.create(device, directory + "/" + in + "_planes.spv", inTypes, 32, cache); m.expandPlanes.create(device, directory + "/" + ex + "_planes.spv", exTypes, 32, cache); } m.post.create(device, directory + "/" + po + ".spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); m.postQuality.create(device, directory + "/" + po + "_q.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); if (exists(directory + "/" + po + "_c.spv")) m.postCompact.create(device, directory + "/" + po + "_c.spv", {b,t,t,t,t,t,t,t,t,s,s,b}, 64, cache); m.down.create(device, directory + "/down_m.spv", {b,t,s,b}, 64, cache); if (exists(directory + "/down_m_rgb10.spv")) m.downRgb10.create(device, directory + "/down_m_rgb10.spv", {b,t,s,b}, 64, cache); m.exposure[0].create(device, directory + "/exposure_m0.spv", {b,t,b}, 64, cache); m.exposure[1].create(device, directory + "/exposure_m1.spv", {b,b,b,b,b,b}, 64, cache); } MModel::~MModel() = default; bool MModel::uploadRecorded() const { return impl->uploaded; } bool MModel::hasExternalNetworkStages() const { return impl->planes; } bool MModel::hasRgb10Output() const { return impl->downRgb10.handle != VK_NULL_HANDLE; } bool MModel::hasDynamicResolution() const { return impl->dynamicResolution; } bool MModel::isPresetL() const { return impl->presetL; } void MModel::recordUpload(VkCommandBuffer cb) { auto& m = *impl; if (m.uploaded) throw std::runtime_error("model upload already recorded"); recordContents(cb, m.w8, m.bytes8); recordContents(cb, m.w16, m.bytes16); 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 MEngine::Impl { MModel::Impl& model; uint32_t rw, rh, maxRw, maxRh, ow, oh, hw, hh, gw, gh; // render, output, reconstruction, token pixel grid Buffer enc0P, dec0P, postP, downP, expP, expSums, expState, dummy, packed; std::array layerP; Buffer tokens, skip1, merged1, skip2, merged2, tubeA, tubeB, dec2Out, dec1Out; Buffer skipL, enc0LP, dec0LP; // preset L: the input stage's pixel output and the two stages' layer blocks std::array mvA, feat, hist, half; Tex mvB, mvC, gates, mix, cov; Tex output; // the engine's own output: only for frames that do not write the caller's image (created on first use) bool outputNew = false; VkDescriptorPool pool = VK_NULL_HANDLE; std::array layerSets{}; std::array enc0Sets{}, dec0Sets{}, postSets{}, expSets{}; VkDescriptorSet downSet = VK_NULL_HANDLE; VkSampler nearest = VK_NULL_HANDLE, linear = VK_NULL_HANDLE; MEnc0Params enc0{}; MDec0Params dec0{}; MPostParams post{}; MDownParams down{}; ExposureParams exposure{}; struct Bindings { FrameImages images{}; FrameOutput finalOutput{}; VkDescriptorPool pool = VK_NULL_HANDLE; std::array input{}, post{}; VkDescriptorSet exposure = VK_NULL_HANDLE, down = VK_NULL_HANDLE; }; std::vector bindings; float prevJitter[2] = {0, 0}; unsigned current = 0, frame = 0; bool initialized = false, frameSet = false, external = false, frontDone = false, quality = false, compact = false, rgb10Output = false; Impl(MModel::Impl& m) : model(m) {} ~Impl() { auto dev = model.device.handle; for (auto& binding : bindings) if (binding.pool) vkDestroyDescriptorPool(dev, binding.pool, nullptr); if (pool) vkDestroyDescriptorPool(dev, pool, nullptr); if (nearest) vkDestroySampler(dev, nearest, nullptr); if (linear) vkDestroySampler(dev, linear, nullptr); } }; VkDeviceSize MEngine::networkBufferBytes(uint32_t renderWidth, uint32_t renderHeight) { return VkDeviceSize((renderWidth + 31) / 32 * 16) * ((renderHeight + 31) / 32 * 16) * 64 + 64; } MEngine::MEngine(MModel& model, uint32_t renderWidth, uint32_t renderHeight, uint32_t outputWidth, uint32_t outputHeight, const ExternalNetwork* external) : impl(std::make_unique(*model.impl)) { auto& e = *impl; const auto& d = e.model.device; auto dev = d.handle; if (external && (!external->tokens || !external->head)) throw std::runtime_error("external network without its buffers"); if (external && !e.model.planes) throw std::runtime_error("the M model has no stages for an external network"); e.external = external != nullptr; if (renderWidth < 64 || renderHeight < 64 || outputWidth % 2 || outputHeight % 2 || outputWidth > 16384 || outputHeight > 16384) throw std::runtime_error("unsupported preset M dimensions"); e.rw = e.maxRw = renderWidth; e.rh = e.maxRh = renderHeight; e.ow = outputWidth; e.oh = outputHeight; e.hw = outputWidth / 2 * 3; e.hh = outputHeight / 2 * 3; if (renderWidth > outputWidth || renderHeight > outputHeight) throw std::runtime_error("preset M render dimensions must not exceed the output"); e.quality = uint64_t(e.hw) < 3ull * renderWidth || uint64_t(e.hh) < 3ull * renderHeight; // An 11x11 render tile covers a 16x16 reconstruction block when h <= 0.45; keep margin below 7/15.5. e.compact = e.quality && e.model.postCompact.handle && 20ull * renderWidth <= 9ull * e.hw && 20ull * renderHeight <= 9ull * e.hh; e.gw = (renderWidth + 31) / 32 * 32; e.gh = (renderHeight + 31) / 32 * 32; const uint32_t tw = e.gw / 2, th = e.gh / 2; const VkBufferUsageFlags sb = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT, sbt = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT; e.enc0P.create(d, sizeof(MEnc0Params), sbt); e.dec0P.create(d, sizeof(MDec0Params), sbt); e.postP.create(d, sizeof(MPostParams), sbt); e.downP.create(d, sizeof(MDownParams), sbt); e.expP.create(d, sizeof(ExposureParams), sbt); e.expSums.create(d, VkDeviceSize((std::max(outputWidth / 2, 256u) + 31) / 32) * ((std::max(outputHeight / 2, 256u) + 31) / 32) * 4, sb); e.expState.create(d, 8, sbt); e.dummy.create(d, 256, sb); e.packed.create(d, VkDeviceSize(e.hw) * e.hh * 4, sb); auto tok = [&](Buffer& b, unsigned scale, unsigned channels) { b.create(d, VkDeviceSize(tw / scale) * (th / scale) * channels + 64, sb); }; if (external) { // the caller's buffers (shared with its network): not owned, the layers' buffers are not needed e.tokens.handle = external->tokens; e.dec1Out.handle = external->head; e.tokens.size = e.dec1Out.size = networkBufferBytes(renderWidth, renderHeight); } else { tok(e.tokens, 1, 64); tok(e.skip1, 1, 64); tok(e.merged1, 2, 96); tok(e.skip2, 2, 96); tok(e.merged2, 4, 128); tok(e.tubeA, 4, 128); tok(e.tubeB, 4, 128); tok(e.dec2Out, 2, 96); tok(e.dec1Out, 1, 64); } for (auto& t : e.mvA) t.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); e.mvB.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); e.mvC.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16_SFLOAT); for (auto& t : e.feat) t.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_SNORM); e.gates.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_UNORM); e.mix.create(d, renderWidth, renderHeight, VK_FORMAT_R8G8B8A8_UNORM); e.cov.create(d, renderWidth, renderHeight, VK_FORMAT_R16G16B16A16_SFLOAT); for (auto& t : e.hist) t.create(d, e.hw, e.hh, VK_FORMAT_R16G16B16A16_SFLOAT); // The native formatted RGB10 history has RGBA32F backing. UNORM filtering differs by up to three ulps. for (auto& t : e.half) t.create(d, e.hw / 2, e.hh / 2, VK_FORMAT_R32G32B32A32_SFLOAT); VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 128}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 48}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 32}}; VkDescriptorPoolCreateInfo pi{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; pi.maxSets = 24; pi.poolSizeCount = 3; pi.pPoolSizes = sizes; VKCHECK(vkCreateDescriptorPool(dev, &pi, nullptr, &e.pool)); VkSamplerCreateInfo si{VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO}; si.addressModeU = si.addressModeV = si.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE; VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.nearest)); si.magFilter = si.minFilter = VK_FILTER_LINEAR; VKCHECK(vkCreateSampler(dev, &si, nullptr, &e.linear)); auto& m = e.model; // the network: tokens -> enc1 -> enc2 -> six tube blocks (ping-pong) -> dec2 -> dec1 const Buffer* in[MLayers] = {&e.tokens, &e.merged1, &e.merged2, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.dec2Out}; const Buffer* skip[MLayers] = {&e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.skip2, &e.skip1}; const Buffer* out[MLayers] = {&e.skip1, &e.skip2, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.tubeA, &e.tubeB, &e.dec2Out, &e.dec1Out}; const Buffer* merged[MLayers] = {&e.merged1, &e.merged2, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy, &e.dummy}; for (unsigned i = 0; i < MLayers && !external; ++i) { e.layerP[i].create(d, sizeof(MLayerParams), sbt, true); MLayerParams lp{int32_t(tw / MScale[i]), int32_t(th / MScale[i]), MShift[i], MShift[i], {}}; std::memcpy(lp.o, m.layers[i].data() + 2, sizeof(lp.o)); std::memcpy(e.layerP[i].mapped, &lp, sizeof(lp)); auto set = e.layerSets[i] = descriptorSet(dev, e.pool, m.network[i].descriptors); bindBuffer(dev, set, 0, e.layerP[i]); bindRange(dev, set, 1, m.w8, m.layers[i][0]); bindRange(dev, set, 2, m.w16, m.layers[i][1]); bindBuffer(dev, set, 3, *in[i]); bindBuffer(dev, set, 4, *skip[i]); bindBuffer(dev, set, 5, *out[i]); bindBuffer(dev, set, 6, *merged[i]); } if (m.presetL) { // the input stage's pixel-resolution block output (the expansion's skip input) and the two stages' layer // parameter blocks (token pixel grid, window shift, offsets; the shift alternates per frame like the layers') e.skipL.create(d, VkDeviceSize(e.gw) * e.gh * 32 + 64, sb); e.enc0LP.create(d, sizeof(MLayerParams), sbt, true); e.dec0LP.create(d, sizeof(MLayerParams), sbt, true); MLayerParams lp{int32_t(e.gw), int32_t(e.gh), 2, 2, {}}; std::memcpy(lp.o, m.lEnc0.data() + 2, sizeof(lp.o)); std::memcpy(e.enc0LP.mapped, &lp, sizeof(lp)); lp.sx = lp.sy = 0; std::memcpy(lp.o, m.lDec0.data() + 2, sizeof(lp.o)); std::memcpy(e.dec0LP.mapped, &lp, sizeof(lp)); } for (unsigned i = 0; i < 2; ++i) { // i: the index of this frame's history outputs const unsigned prev = 1 - i; if (m.presetL) { auto s0 = e.enc0Sets[i] = descriptorSet(dev, e.pool, (external ? m.inputPlanes : m.input).descriptors); bindBuffer(dev, s0, 0, e.enc0LP); bindRange(dev, s0, 1, m.w8, m.lEnc0[0]); bindRange(dev, s0, 2, m.w16, m.lEnc0[1]); bindBuffer(dev, s0, 3, e.dummy); bindBuffer(dev, s0, 4, e.dummy); bindBuffer(dev, s0, 5, e.skipL); bindBuffer(dev, s0, 6, e.tokens); bindBuffer(dev, s0, 7, e.enc0P); bindImage(dev, s0, 11, e.mvA[prev].view, e.linear); bindImage(dev, s0, 12, e.half[prev].view, e.linear); bindImage(dev, s0, 13, e.feat[prev].view, e.linear); bindImage(dev, s0, 14, e.mvA[i].view); bindImage(dev, s0, 15, e.mvB.view); bindImage(dev, s0, 16, e.mvC.view); bindRange(dev, s0, 17, m.w16, m.dither16); auto s1 = e.dec0Sets[i] = descriptorSet(dev, e.pool, (external ? m.expandPlanes : m.expand).descriptors); bindBuffer(dev, s1, 0, e.dec0LP); bindRange(dev, s1, 1, m.w8, m.lDec0[0]); bindRange(dev, s1, 2, m.w16, m.lDec0[1]); bindBuffer(dev, s1, 3, e.dec1Out); bindBuffer(dev, s1, 4, e.skipL); bindBuffer(dev, s1, 5, e.dummy); bindBuffer(dev, s1, 6, e.dummy); bindBuffer(dev, s1, 7, e.dec0P); bindImage(dev, s1, 8, e.feat[i].view); bindImage(dev, s1, 9, e.gates.view); bindImage(dev, s1, 10, e.mix.view); bindImage(dev, s1, 11, e.cov.view); } else { auto s0 = e.enc0Sets[i] = descriptorSet(dev, e.pool, (external ? m.inputPlanes : m.input).descriptors); bindBuffer(dev, s0, 0, e.enc0P); bindImage(dev, s0, 4, e.mvA[prev].view, e.linear); bindImage(dev, s0, 5, e.half[prev].view, e.linear); bindImage(dev, s0, 6, e.feat[prev].view, e.linear); bindImage(dev, s0, 7, e.mvA[i].view); bindImage(dev, s0, 8, e.mvB.view); bindImage(dev, s0, 9, e.mvC.view); bindRange(dev, s0, 10, m.w8, m.enc0[0]); bindRange(dev, s0, 11, m.w16, m.enc0[1]); bindRange(dev, s0, 12, m.w16, m.enc0[2]); bindBuffer(dev, s0, 13, e.tokens); auto s1 = e.dec0Sets[i] = descriptorSet(dev, e.pool, (external ? m.expandPlanes : m.expand).descriptors); bindBuffer(dev, s1, 0, e.dec0P); bindImage(dev, s1, 1, e.mvA[i].view, e.nearest); bindImage(dev, s1, 2, e.feat[prev].view, e.linear); bindImage(dev, s1, 3, e.feat[i].view); bindImage(dev, s1, 4, e.gates.view); bindImage(dev, s1, 5, e.mix.view); bindImage(dev, s1, 6, e.cov.view); bindBuffer(dev, s1, 7, e.dec1Out); bindRange(dev, s1, 8, m.w8, m.dec0[0]); bindRange(dev, s1, 9, m.w16, m.dec0[1]); } auto s2 = e.postSets[i] = descriptorSet(dev, e.pool, m.post.descriptors); bindBuffer(dev, s2, 0, e.postP); bindImage(dev, s2, 2, e.mix.view, e.nearest); bindImage(dev, s2, 3, e.gates.view, e.nearest); bindImage(dev, s2, 4, e.cov.view, e.nearest); bindImage(dev, s2, 5, e.hist[prev].view, e.linear); bindImage(dev, s2, 6, e.mvA[i].view, e.nearest); bindImage(dev, s2, 7, e.mvB.view, e.nearest); bindImage(dev, s2, 8, e.mvC.view, e.linear); bindImage(dev, s2, 9, e.hist[i].view); bindImage(dev, s2, 10, e.half[i].view); bindBuffer(dev, s2, 11, e.packed); e.expSets[i] = descriptorSet(dev, e.pool, m.exposure[i].descriptors); bindBuffer(dev, e.expSets[i], 0, e.expP); } bindBuffer(dev, e.expSets[0], 2, e.expSums); bindBuffer(dev, e.expSets[1], 1, e.expSums); bindBuffer(dev, e.expSets[1], 2, e.expState); bindBuffer(dev, e.expSets[1], 3, e.enc0P); bindBuffer(dev, e.expSets[1], 4, e.postP); bindBuffer(dev, e.expSets[1], 5, e.downP); e.downSet = descriptorSet(dev, e.pool, m.down.descriptors); bindBuffer(dev, e.downSet, 0, e.downP); bindBuffer(dev, e.downSet, 3, e.packed); // binding 2, the output: setFrame } MEngine::~MEngine() = default; void MEngine::recordInitialize(VkCommandBuffer cb) { auto& e = *impl; if (e.initialized) throw std::runtime_error("engine initialization already recorded"); VkClearColorValue zero{}; VkImageSubresourceRange range{VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1}; for (Tex* t : {&e.mvA[0], &e.mvA[1], &e.mvB, &e.mvC, &e.feat[0], &e.feat[1], &e.gates, &e.mix, &e.cov, &e.hist[0], &e.hist[1], &e.half[0], &e.half[1]}) { imageBarrier(cb, t->handle, 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, t->handle, VK_IMAGE_LAYOUT_GENERAL, &zero, 1, &range); } const float none[2] = {0, 0}; vkCmdUpdateBuffer(cb, e.expState.handle, 0, sizeof(none), none); 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); e.initialized = true; } void MEngine::setFrame(const FrameImages& images, const MFrameSettings& s, FrameOutput finalOutput) { auto& e = *impl; if (!images.color || !images.motion || !images.depth) throw std::runtime_error("frame has a null image view"); if (finalOutput.view && finalOutput.format != VK_FORMAT_R16G16B16A16_SFLOAT && (finalOutput.format != VK_FORMAT_A2B10G10R10_UNORM_PACK32 || !e.model.downRgb10.handle)) throw std::runtime_error("unsupported borrowed output format or missing RGB10A2 shader"); const uint32_t rw = s.renderWidth ? s.renderWidth : e.maxRw; const uint32_t rh = s.renderHeight ? s.renderHeight : e.maxRh; if (rw < 64 || rh < 64 || rw > e.maxRw || rh > e.maxRh) throw std::runtime_error("render size exceeds the engine's capacity or is below 64"); if (!e.model.dynamicResolution && (rw != e.maxRw || rh != e.maxRh)) throw std::runtime_error("recompile the M/L model for dynamic render sizes"); if (e.external && (rw != e.maxRw || rh != e.maxRh)) throw std::runtime_error("external M/L network requires its constructor render size"); const bool reset = s.reset || e.frame == 0 || rw != e.rw || rh != e.rh; e.rw = rw; e.rh = rh; e.gw = (rw + 31) / 32 * 32; e.gh = (rh + 31) / 32 * 32; e.quality = uint64_t(e.hw) < 3ull * rw || uint64_t(e.hh) < 3ull * rh; e.compact = e.quality && e.model.postCompact.handle && 20ull * rw <= 9ull * e.hw && 20ull * rh <= 9ull * e.hh; const bool L = e.model.presetL; const float fx = float(e.hw) / float(e.rw), fy = float(e.hh) / float(e.rh); e.enc0 = {{int32_t(e.rw), int32_t(e.rh)}, {2, 2}, {int32_t(e.gw / 2), int32_t(e.gh / 2)}, {s.jitter[0], s.jitter[1]}, {e.prevJitter[0], e.prevJitter[1]}, {s.motionScale[0] * fx, s.motionScale[1] * fy}, {s.motionOffset[0], s.motionOffset[1]}, {fx, fy}, {1.f / fx, 1.f / fy}, {float(e.hw), float(e.hh)}, L ? e.model.lConst[0] : 9.045065879821777f, L ? e.model.lConst[1] : 0.22356846928596497f, s.depthInverted ? 1u : 0u, reset ? 1u : 0u, e.frame & 7u, 1.f, 0.f, {1.f / float(e.hw), 1.f / float(e.hh)}, {1.f / float(e.maxRw), 1.f / float(e.maxRh)}}; e.dec0 = {{int32_t(e.rw), int32_t(e.rh)}, int32_t(e.gw / 2), reset ? 1u : 0u, {s.jitter[0], s.jitter[1]}, {fx, fy}, {float(e.hw), float(e.hh)}, {e.prevJitter[0], e.prevJitter[1]}, L ? e.model.lConst[2] : 8.f, L ? e.model.lConst[3] : 0.07863261550664902f, {1.f / float(e.maxRw), 1.f / float(e.maxRh)}}; e.post = {{0, 0}, {int32_t(e.hw), int32_t(e.hh)}, {int32_t(e.rw), int32_t(e.rh)}, {1.f / fx, 1.f / fy}, {s.jitter[0], s.jitter[1]}, {1.f / float(e.maxRw), 1.f / float(e.maxRh)}, L ? e.model.lConst[4] : 2.7530112f, reset ? 1u : 0u, 1.f, 0.f}; e.down = {{int32_t(e.hw), int32_t(e.hh)}, {0, 0}, {0, 0}, {int32_t(e.rw), int32_t(e.rh)}, {float(e.rw) / float(e.ow), float(e.rh) / float(e.oh)}, {s.jitter[0], s.jitter[1]}, 1.f}; const uint32_t sw = exposureGrid(e.rw, e.ow), sh = exposureGrid(e.rh, e.oh); const uint32_t bx = (sw + 31) / 32, by = (sh + 31) / 32; e.exposure = {{0, 0}, {int32_t(e.rw), int32_t(e.rh)}, 1.f / (float(sw) * float(sh)), 1.f, 0.23547087609767914f, 1u | (reset ? 2u : 0u), bx * by, bx, {int32_t(sw), int32_t(sh)}}; e.prevJitter[0] = s.jitter[0]; e.prevJitter[1] = s.jitter[1]; auto dev = e.model.device.handle; auto& m = e.model; if (!finalOutput.view && !e.output.handle) { e.output.create(e.model.device, e.ow, e.oh, VK_FORMAT_R16G16B16A16_SFLOAT); e.outputNew = true; } auto found = std::find_if(e.bindings.begin(), e.bindings.end(), [&](const auto& b) { return b.images.color == images.color && b.images.motion == images.motion && b.images.depth == images.depth && b.images.colorLayout == images.colorLayout && b.images.motionLayout == images.motionLayout && b.images.depthLayout == images.depthLayout && b.finalOutput.view == finalOutput.view && b.finalOutput.format == finalOutput.format; }); if (found == e.bindings.end()) { Impl::Bindings b; b.images = images; b.finalOutput = finalOutput; if (e.bindings.empty()) { b.input = e.enc0Sets; b.post = e.postSets; b.exposure = e.expSets[0]; b.down = e.downSet; } else { VkDescriptorPoolSize sizes[] = {{VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 32}, {VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 32}, {VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 16}}; VkDescriptorPoolCreateInfo info{VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO}; info.maxSets = 6; info.poolSizeCount = 3; info.pPoolSizes = sizes; VKCHECK(vkCreateDescriptorPool(dev, &info, nullptr, &b.pool)); } try { if (b.pool) { b.exposure = descriptorSet(dev, b.pool, m.exposure[0].descriptors); b.down = descriptorSet(dev, b.pool, m.down.descriptors); } bindBuffer(dev, b.exposure, 0, e.expP); bindImage(dev, b.exposure, 1, images.color, e.nearest, images.colorLayout); bindBuffer(dev, b.exposure, 2, e.expSums); bindBuffer(dev, b.down, 0, e.downP); bindImage(dev, b.down, 1, images.color, e.nearest, images.colorLayout); bindImage(dev, b.down, 2, finalOutput.view ? finalOutput.view : e.output.view); bindBuffer(dev, b.down, 3, e.packed); for (unsigned i = 0; i < 2; ++i) { if (b.pool) { b.input[i] = descriptorSet(dev, b.pool, (e.external ? m.inputPlanes : m.input).descriptors); b.post[i] = descriptorSet(dev, b.pool, m.post.descriptors); } auto in = b.input[i], post = b.post[i]; if (L) { bindBuffer(dev, in, 0, e.enc0LP); bindRange(dev, in, 1, m.w8, m.lEnc0[0]); bindRange(dev, in, 2, m.w16, m.lEnc0[1]); bindBuffer(dev, in, 3, e.dummy); bindBuffer(dev, in, 4, e.dummy); bindBuffer(dev, in, 5, e.skipL); bindBuffer(dev, in, 6, e.tokens); bindBuffer(dev, in, 7, e.enc0P); bindImage(dev, in, 8, images.color, e.nearest, images.colorLayout); bindImage(dev, in, 9, images.motion, e.nearest, images.motionLayout); bindImage(dev, in, 10, images.depth, e.nearest, images.depthLayout); bindImage(dev, in, 11, e.mvA[1 - i].view, e.linear); bindImage(dev, in, 12, e.half[1 - i].view, e.linear); bindImage(dev, in, 13, e.feat[1 - i].view, e.linear); bindImage(dev, in, 14, e.mvA[i].view); bindImage(dev, in, 15, e.mvB.view); bindImage(dev, in, 16, e.mvC.view); bindRange(dev, in, 17, m.w16, m.dither16); } else { bindBuffer(dev, in, 0, e.enc0P); bindImage(dev, in, 1, images.color, e.nearest, images.colorLayout); bindImage(dev, in, 2, images.motion, e.nearest, images.motionLayout); bindImage(dev, in, 3, images.depth, e.nearest, images.depthLayout); bindImage(dev, in, 4, e.mvA[1 - i].view, e.linear); bindImage(dev, in, 5, e.half[1 - i].view, e.linear); bindImage(dev, in, 6, e.feat[1 - i].view, e.linear); bindImage(dev, in, 7, e.mvA[i].view); bindImage(dev, in, 8, e.mvB.view); bindImage(dev, in, 9, e.mvC.view); bindRange(dev, in, 10, m.w8, m.enc0[0]); bindRange(dev, in, 11, m.w16, m.enc0[1]); bindRange(dev, in, 12, m.w16, m.enc0[2]); bindBuffer(dev, in, 13, e.tokens); } bindBuffer(dev, post, 0, e.postP); bindImage(dev, post, 1, images.color, e.nearest, images.colorLayout); bindImage(dev, post, 2, e.mix.view, e.nearest); bindImage(dev, post, 3, e.gates.view, e.nearest); bindImage(dev, post, 4, e.cov.view, e.nearest); bindImage(dev, post, 5, e.hist[1 - i].view, e.linear); bindImage(dev, post, 6, e.mvA[i].view, e.nearest); bindImage(dev, post, 7, e.mvB.view, e.nearest); bindImage(dev, post, 8, e.mvC.view, e.linear); bindImage(dev, post, 9, e.hist[i].view); bindImage(dev, post, 10, e.half[i].view); bindBuffer(dev, post, 11, e.packed); } e.bindings.push_back(b); } catch (...) { if (b.pool) vkDestroyDescriptorPool(dev, b.pool, nullptr); throw; } found = e.bindings.end() - 1; } e.enc0Sets = found->input; e.postSets = found->post; e.expSets[0] = found->exposure; e.downSet = found->down; e.rgb10Output = finalOutput.view && finalOutput.format == VK_FORMAT_A2B10G10R10_UNORM_PACK32; e.frameSet = true; } namespace { const auto compute = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT; void emptyStage(VkCommandBuffer cb) { vkCmdPipelineBarrier(cb, compute, compute, 0, 0, nullptr, 0, nullptr, 0, nullptr); } void nextStage(VkCommandBuffer cb) { memoryBarrier(cb, compute, compute, VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_SHADER_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); } void openStages(VkCommandBuffer cb) { memoryBarrier(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, compute | VK_PIPELINE_STAGE_TRANSFER_BIT, VK_ACCESS_MEMORY_WRITE_BIT | VK_ACCESS_MEMORY_READ_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT); } void closeStages(VkCommandBuffer cb) { memoryBarrier(cb, compute, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, VK_ACCESS_SHADER_WRITE_BIT, VK_ACCESS_MEMORY_READ_BIT | VK_ACCESS_MEMORY_WRITE_BIT); } } // exposure and the input stage; writes three stamps void MEngine::recordFront(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { auto& e = *impl; auto& m = e.model; if (!e.initialized || !m.uploaded || !e.frameSet || e.frontDone) throw std::runtime_error("initialize/upload/setFrame before recording inference"); if (e.outputNew) { clearNewImage(cb, e.output.handle); e.outputNew = false; } openStages(cb); vkCmdUpdateBuffer(cb, e.enc0P.handle, 0, sizeof(e.enc0), &e.enc0); vkCmdUpdateBuffer(cb, e.dec0P.handle, 0, sizeof(e.dec0), &e.dec0); vkCmdUpdateBuffer(cb, e.postP.handle, 0, sizeof(e.post), &e.post); vkCmdUpdateBuffer(cb, e.downP.handle, 0, sizeof(e.down), &e.down); vkCmdUpdateBuffer(cb, e.expP.handle, 0, sizeof(e.exposure), &e.exposure); if (!e.external) for (unsigned l = 0; l < MLayers; ++l) { const int32_t grid[4] = {int32_t(e.gw / 2 / MScale[l]), int32_t(e.gh / 2 / MScale[l]), MShift[l] ^ int32_t((e.frame & 1u) * 2u), MShift[l]}; vkCmdUpdateBuffer(cb, e.layerP[l].handle, 0, sizeof(grid), grid); } if (m.presetL) { // NGX alternates the horizontal window shift every frame: input stage (2, 2) / (0, 2), expansion (0, 0) / (2, 0) const int32_t in[4] = {int32_t(e.gw), int32_t(e.gh), 2 ^ int32_t((e.frame & 1u) * 2u), 2}; const int32_t ex[4] = {int32_t(e.gw), int32_t(e.gh), int32_t((e.frame & 1u) * 2u), 0}; vkCmdUpdateBuffer(cb, e.enc0LP.handle, 0, sizeof(in), in); vkCmdUpdateBuffer(cb, e.dec0LP.handle, 0, sizeof(ex), ex); } memoryBarrier(cb, VK_PIPELINE_STAGE_TRANSFER_BIT, compute, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT); uint32_t query = first; auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; stamp(); if (!m.ldr) { dispatch(cb, m.exposure[0], e.expSets[0], e.exposure.blocksX, e.exposure.blocks / e.exposure.blocksX); nextStage(cb); dispatch(cb, m.exposure[1], e.expSets[1], 1, 1); nextStage(cb); } stamp(); // (preset L's input stage covers the same blocks: its shift is at most 2 in x and 2 in y) dispatch(cb, e.external ? m.inputPlanes : m.input, e.enc0Sets[e.current], (e.gw + 2 + 7) / 8, (e.gh + 2 + 7) / 8); stamp(); if (e.external) closeStages(cb); else nextStage(cb); e.frontDone = true; } // expansion, reconstruction and downsample; writes four stamps (the first before the expansion) void MEngine::recordBack(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { auto& e = *impl; auto& m = e.model; if (!e.frontDone) throw std::runtime_error("recordFront before recordBack"); if (e.external) openStages(cb); const unsigned i = e.current; uint32_t query = first; auto stamp = [&] { if (stamps) vkCmdWriteTimestamp(cb, VK_PIPELINE_STAGE_ALL_COMMANDS_BIT, stamps, query++); }; stamp(); if (m.presetL) // one block more in x on frames whose expansion shift is 2 dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], (e.gw + (e.frame & 1u) * 2u + 7) / 8, e.gh / 8); else dispatch(cb, e.external ? m.expandPlanes : m.expand, e.dec0Sets[i], e.gw / 8, e.gh / 8); nextStage(cb); stamp(); const bool compact = e.compact && e.post.rndOffs[0] >= -0.5f && e.post.rndOffs[0] <= 0.5f && e.post.rndOffs[1] >= -0.5f && e.post.rndOffs[1] <= 0.5f; dispatch(cb, compact ? m.postCompact : (e.quality ? m.postQuality : m.post), e.postSets[i], (e.hw + 15) / 16, (e.hh + 15) / 16); nextStage(cb); stamp(); // D4RM0004's downsample workgroup covers 8x64 output pixels. dispatch(cb, e.rgb10Output ? m.downRgb10 : m.down, e.downSet, (e.ow + 7) / 8, (e.oh + 63) / 64); stamp(); closeStages(cb); e.current = 1 - e.current; ++e.frame; e.frameSet = false; e.frontDone = false; } void MEngine::recordFrame(VkCommandBuffer cb, VkQueryPool stamps, uint32_t first) { auto& e = *impl; auto& m = e.model; if (e.external) throw std::runtime_error("this engine's network is external: record the front and the back"); recordFront(cb, stamps, first); for (unsigned l = 0; l < MLayers; ++l) { const uint32_t w = e.gw / 2 / MScale[l], h = e.gh / 2 / MScale[l]; dispatch(cb, m.network[l], e.layerSets[l], (w + (MShift[l] ^ ((e.frame & 1u) * 2u)) + 7) / 8, (h + MShift[l] + 7) / 8); emptyStage(cb); } recordBack(cb, stamps, stamps ? first + 3 : 0); } VkImage MEngine::outputImage() const { return impl->output.handle; } std::vector MEngine::debugResources() const { auto& e = *impl; const uint32_t w = e.rw, h = e.rh; std::vector r{{"enc0-params", e.enc0P.handle, VK_NULL_HANDLE, e.enc0P.size, 0, 0}, {"tokens", e.tokens.handle, VK_NULL_HANDLE, e.tokens.size, 0, 0}, {"head", e.dec1Out.handle, VK_NULL_HANDLE, e.dec1Out.size, 0, 0}, {"gates", VK_NULL_HANDLE, e.gates.handle, VkDeviceSize(w) * h * 4, w, h}, {"mix", VK_NULL_HANDLE, e.mix.handle, VkDeviceSize(w) * h * 4, w, h}, {"cov", VK_NULL_HANDLE, e.cov.handle, VkDeviceSize(w) * h * 8, w, h}, {"mvA", VK_NULL_HANDLE, e.mvA[1 - e.current].handle, VkDeviceSize(w) * h * 4, w, h}, {"feat", VK_NULL_HANDLE, e.feat[1 - e.current].handle, VkDeviceSize(w) * h * 4, w, h}}; if (e.model.presetL) r.push_back({"skip", e.skipL.handle, VK_NULL_HANDLE, e.skipL.size, 0, 0}); return r; }