diff --git a/core/libs/prx/libSceAgcDriver/tests/DenormalFlushElimination.cpp b/core/libs/prx/libSceAgcDriver/tests/DenormalFlushElimination.cpp index cc952f5ea999..5dd5cfc4ba8b 100644 --- a/core/libs/prx/libSceAgcDriver/tests/DenormalFlushElimination.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/DenormalFlushElimination.cpp @@ -110,7 +110,8 @@ std::vector Recompiled(std::span code) { request.target.maxWorkgroupSharedMemoryBytes = 49152; request.layout = {0, 0, 0, 128}; request.useCache = false; - const auto& words = Recompile(request).spirv.Words(); + const auto result = Recompile(request); + const auto& words = result.spirv.Words(); return {words.begin(), words.end()}; } diff --git a/core/libs/prx/libSceAgcDriver/tests/RedundantValueElimination.cpp b/core/libs/prx/libSceAgcDriver/tests/RedundantValueElimination.cpp index cc88c42cfeaf..2aec9adc4f01 100644 --- a/core/libs/prx/libSceAgcDriver/tests/RedundantValueElimination.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/RedundantValueElimination.cpp @@ -208,7 +208,8 @@ std::vector Recompiled(std::span code) { request.target.maxWorkgroupSharedMemoryBytes = 49152; request.layout = {0, 0, 0, 128}; request.useCache = false; - const auto& words = Recompile(request).spirv.Words(); + const auto result = Recompile(request); + const auto& words = result.spirv.Words(); return {words.begin(), words.end()}; } diff --git a/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp b/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp index f015fad6b4a7..69d35f12821c 100644 --- a/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/ShaderMemory.cpp @@ -406,6 +406,36 @@ void verifyFrontendPair() { } } +void verifySpecializedModuleRelease() { + using namespace ShaderRecompiler; + const std::array color{0x3f800000u}; + const std::vector code{0xbf800003u, 0x7e000200u, 0xf800180fu, 0u, 0xbf810000u}; + auto pixel = FrontendRequest(code, color); + pixel.shader.stage = ShaderStage::Fragment; + pixel.context.compute.reset(); + pixel.context.pixel = ShaderPixelStageInfo{}; + pixel.context.pixel->wave32 = true; + pixel.context.pixel->targetOutputMode[0] = 9u; + pixel.context.waveSize = 32u; + pixel.useCache = true; + SharedSpirv::Weak released; + std::uint64_t specializationId = 0; + std::vector words; + { + const auto first = Recompile(pixel); + require(first.specializationId != 0u, "specialized module release: the pixel fixture was not specialized"); + released = first.spirv.Weaken(); + specializationId = first.specializationId; + words = first.spirv.Copy(); + } + require(released.Expired(), "specialized module release: the module caches kept the words of a module no result uses"); + const auto again = Recompile(pixel); + require(again.specializationId == specializationId, "specialized module release: a module built again from its key got another specialization ID"); + require(std::ranges::equal(again.spirv.Words(), words), "specialized module release: a module built again from its key has other words"); + const auto shared = Recompile(pixel); + require(shared.spirv.data() == again.spirv.data(), "specialized module release: a second result of a live module copied its words"); +} + ShaderRecompiler::RecompileRequest LayoutPairRequest(std::span code, std::span userData, bool registration, std::uint32_t wordLimit) { using namespace ShaderRecompiler; auto request = FrontendRequest(code, userData); @@ -8466,6 +8496,7 @@ int main(int argc, char** argv) { verifyTargetFeatureKey(); verifyImageTableKeys(); verifyImageTables(); + verifySpecializedModuleRelease(); verifyImageTableRegistration(); verifyScalarAddressRule(); verifyPointerImageTable(); diff --git a/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp b/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp index 05d848cdf25c..956897cb51be 100644 --- a/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/SpirvModule.cpp @@ -460,6 +460,20 @@ void testInternedSpirvSharesEqualWords() { check(edited[6] == spv::CapabilityFloat64 && first[6] == spv::CapabilityShader && again.data() == first.data(), "a write to one holder of interned SPIR-V changed the others"); } +void testInternedSpirvReleasesUnusedWords() { + const std::vector words{spv::MagicNumber, 0x00010400u, 0u, 7u, 0u, (2u << 16u) | spv::OpCapability, spv::CapabilityGeometry}; + SharedSpirv::Weak released; + { + const auto first = InternSpirv(words); + const auto second = InternSpirv(words); + check(first.data() == second.data(), "two equal specialized modules kept separate copies of their words"); + released = first.Weaken(); + } + check(released.Expired(), "the SPIR-V interning table kept words no holder uses"); + const auto again = InternSpirv(words); + check(std::ranges::equal(again.Words(), words) && released.Lock().empty(), "SPIR-V interned again after its release changed its words"); +} + int main() { testFoldingLeavesNoDeadBlocksOrConstants(); testIdsBeyondTheBoundAreRejected(); @@ -477,6 +491,7 @@ int main() { testEmittedWordsAreStable(); testBlockBallotsOfOnePredicateShareOneBallot(); testInternedSpirvSharesEqualWords(); + testInternedSpirvReleasesUnusedWords(); if (failures != 0) { std::fprintf(stderr, "%d check(s) failed\n", failures); return 1; diff --git a/core/shader/recompiler/Recompiler.cpp b/core/shader/recompiler/Recompiler.cpp index 8c90d564880b..f959e013c1ce 100644 --- a/core/shader/recompiler/Recompiler.cpp +++ b/core/shader/recompiler/Recompiler.cpp @@ -746,12 +746,19 @@ CompiledVariant compileVariant(const RecompileRequest& request, IrProgram progra struct SpecializedModule { std::uint64_t specializationId = 0; - SharedSpirv spirv; std::vector bindings; bool pushData = false; + SpirvContentHash hash; + mutable std::mutex mutex; + mutable SharedSpirv::Weak spirv; +}; + +struct SpecializedSpirv { + std::shared_ptr module; + SharedSpirv spirv; }; -std::shared_ptr buildSpecializedModule(const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { +std::vector specializeWords(const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { auto source = SpecializeVertexInputTypes(artifact, classes); std::map supplied; for (const auto& constant : constants) { @@ -845,6 +852,11 @@ std::shared_ptr buildSpecializedModule(const CompiledSh #if ANYPS5_ENABLE_SPIRV_TOOLS materialized = ValidateAndOptimizeSpirv(materialized, target.vulkanVersion, target.spirvVersion, target.nonConstantImageOffsets, false); #endif + return materialized; +} + +SpecializedSpirv buildSpecializedModule(const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { + auto materialized = specializeWords(artifact, classes, constants, target); auto module = std::make_shared(); std::map bindingNumbers; std::unordered_set functionLdsBuffer; @@ -859,9 +871,20 @@ std::shared_ptr buildSpecializedModule(const CompiledSh } cursor += count; } - module->spirv = InternSpirv(std::move(materialized)); + auto spirv = InternSpirv(std::move(materialized)); + module->hash = spirv.ContentHash(); + module->spirv = spirv.Weaken(); module->specializationId = constants.empty() ? 0u : nextVariantId(); - return module; + return {std::move(module), std::move(spirv)}; +} + +SharedSpirv specializedWords(const SpecializedModule& module, const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { + std::lock_guard lock(module.mutex); + if (auto spirv = module.spirv.Lock(); !spirv.empty()) return spirv; + auto spirv = InternSpirv(specializeWords(artifact, classes, constants, target)); + if (spirv.ContentHash() != module.hash) throw std::runtime_error("ShaderRecompiler: a specialized module built again from its key differs from its first build"); + module.spirv = spirv.Weaken(); + return spirv; } struct SpecializedModuleEntry { @@ -891,7 +914,12 @@ struct RegistrationModuleEntry { std::shared_ptr module; }; -std::shared_ptr registrationModule(std::uint64_t variantId, std::span moduleKey, const std::function()>& specialize) { +SharedSpirv lockedWords(const SpecializedModule& module) { + std::lock_guard lock(module.mutex); + return module.spirv.Lock(); +} + +SpecializedSpirv registrationModule(std::uint64_t variantId, std::span moduleKey, const std::function& specialize) { std::vector key{variantId}; key.insert(key.end(), moduleKey.begin(), moduleKey.end()); static PrimedSharedMutex mutex; @@ -906,15 +934,20 @@ std::shared_ptr registrationModule(std::uint64_t varian const auto found = modules.find(key); entry = found != modules.end() ? found->second : modules.emplace(key, std::make_shared()).first->second; } + SharedSpirv built; std::call_once(entry->ready, [&] { try { - entry->module = specialize(); + auto specialized = specialize(); + entry->module = std::move(specialized.module); + built = std::move(specialized.spirv); } catch (...) { entry->failure = std::current_exception(); } }); if (entry->failure) std::rethrow_exception(entry->failure); - return entry->module; + if (built.empty()) built = lockedWords(*entry->module); + if (built.empty()) built = specialize().spirv; + return {entry->module, std::move(built)}; } std::shared_ptr preparedBindingPlan(const CompiledVariant& variant, const ResourceSnapshot& snapshot, std::span exports, bool fragmentDescriptorSet) { @@ -967,7 +1000,7 @@ std::shared_ptr preparedBindingPlan(const CompiledVariant& return plan; } -std::shared_ptr specializeModule(const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { +SpecializedSpirv specializeModule(const CompiledShaderArtifact& artifact, std::span classes, std::span constants, const SpirvTarget& target) { struct ModuleKeyStorage {}; auto& key = HostThreadLocal, ModuleKeyStorage>(); key.clear(); @@ -985,15 +1018,19 @@ std::shared_ptr specializeModule(const CompiledShaderAr const auto found = modules.find(key); entry = found != modules.end() ? found->second : modules.emplace(key, std::make_shared()).first->second; } + SharedSpirv built; std::call_once(entry->ready, [&] { try { - entry->module = buildSpecializedModule(artifact, classes, constants, target); + auto specialized = buildSpecializedModule(artifact, classes, constants, target); + entry->module = std::move(specialized.module); + built = std::move(specialized.spirv); } catch (...) { entry->failure = std::current_exception(); } }); if (entry->failure) std::rethrow_exception(entry->failure); - return entry->module; + if (built.empty()) built = specializedWords(*entry->module, artifact, classes, constants, target); + return {entry->module, std::move(built)}; } const ResourceSnapshot& mergeTargetImage(const CompiledVariant& variant, const RecompileRequest& request, const ResourceSnapshot& snapshot, bool registrationDefaults, ResourceSnapshot& merged) { @@ -1115,7 +1152,7 @@ RecompileResult materializeResult(const CompiledVariant& variant, const Recompil const auto found = plan->modules.find(moduleKey); entry = found != plan->modules.end() ? found->second : plan->modules.emplace(moduleKey, std::make_shared()).first->second; } - const auto prepare = [&] { + const auto specialize = [&] { auto constants = plan->bindings.specialization; std::size_t index = 0; if (variant.bindings.layout.UsesPushData()) constants.push_back({PipelineSpecialization::PushDataOffset, moduleKey[index++]}); @@ -1140,11 +1177,17 @@ RecompileResult materializeResult(const CompiledVariant& variant, const Recompil constants.push_back({first + 5u, kind}); } } - const auto module = artifact.registrationModule && plan->bindings.specialization.empty() ? registrationModule(artifact.variantId, moduleKey, [&] { return specializeModule(variant.artifact, vertexClasses, constants, request.target); }) : specializeModule(variant.artifact, vertexClasses, constants, request.target); + return artifact.registrationModule && plan->bindings.specialization.empty() ? registrationModule(artifact.variantId, moduleKey, [&] { return specializeModule(variant.artifact, vertexClasses, constants, request.target); }) : specializeModule(variant.artifact, vertexClasses, constants, request.target); + }; + SharedSpirv built; + const auto prepare = [&] { + auto specialized = specialize(); + const auto& module = specialized.module; if (variant.bindings.layout.UsesPushData() && !module->pushData) throw std::runtime_error("specialization removed the prepared push constant interface"); auto selected = DescriptorBindingBuilder{}.Select(plan->bindings, module->bindings); entry->bindings = std::move(selected); entry->module = module; + built = std::move(specialized.spirv); }; std::call_once(entry->ready, [&] { try { @@ -1155,8 +1198,10 @@ RecompileResult materializeResult(const CompiledVariant& variant, const Recompil }); if (entry->failure) std::rethrow_exception(entry->failure); const auto& module = entry->module; + if (built.empty()) built = lockedWords(*module); + if (built.empty()) built = specialize().spirv; result.specializationId = module->specializationId; - result.spirv = module->spirv; + result.spirv = std::move(built); bindingPlan = &entry->bindings; } else { result.spirv = artifact.spirv; @@ -1214,13 +1259,16 @@ SharedSpirv InternSpirv(std::vector words) { SharedSpirv spirv(std::move(words)); const auto hash = spirv.ContentHash(); static std::mutex mutex; - static std::unordered_map> interned; + static std::unordered_map> interned; std::lock_guard lock(mutex); - auto& candidates = interned[hash.low ^ (hash.high * 0x9e3779b97f4a7c15ull)]; - for (const auto& candidate : candidates) { + const auto bucket = hash.low ^ (hash.high * 0x9e3779b97f4a7c15ull); + auto& candidates = interned[bucket]; + std::erase_if(candidates, [](const SharedSpirv::Weak& candidate) { return candidate.Expired(); }); + for (const auto& weak : candidates) { + const auto candidate = weak.Lock(); if (candidate.ContentHash() == hash && candidate == spirv) return candidate; } - candidates.push_back(spirv); + candidates.push_back(spirv.Weaken()); return spirv; } diff --git a/core/shader/recompiler/Recompiler.hpp b/core/shader/recompiler/Recompiler.hpp index 9462ba07215c..212e6b2f8d8b 100644 --- a/core/shader/recompiler/Recompiler.hpp +++ b/core/shader/recompiler/Recompiler.hpp @@ -603,6 +603,8 @@ inline std::uint64_t HashSpirvWords(std::span words, std::u // reference, or the per-dispatch copy this class removes comes back without a compiler hint. // ContentHash is computed once per shared holder, on first use, and dropped by a write. class SharedSpirv { + struct Holder; + public: SharedSpirv() = default; SharedSpirv(std::vector words) : holder(std::make_shared(std::move(words), nullptr, std::span{})) {} @@ -639,6 +641,25 @@ class SharedSpirv { std::vector::iterator insert(std::vector::const_iterator where, std::initializer_list values) { return Mutable().insert(where, values); } friend bool operator==(const SharedSpirv& left, const SharedSpirv& right) { return std::ranges::equal(left.Words(), right.Words()); } + class Weak { + public: + [[nodiscard]] SharedSpirv Lock() const { + SharedSpirv spirv; + spirv.holder = holder.lock(); + return spirv; + } + [[nodiscard]] bool Expired() const { return holder.expired(); } + + private: + friend class SharedSpirv; + std::weak_ptr holder; + }; + [[nodiscard]] Weak Weaken() const { + Weak weak; + weak.holder = holder; + return weak; + } + private: struct Holder { Holder(std::vector words, std::shared_ptr owner, std::span view) : words(std::move(words)), owner(std::move(owner)), view(view) {}