perf(agc): release unused registered shader resource plans

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp

@@ -725,6 +725,7 @@ void BuildRegisteredAbiKey(const QueueState& state, const VulkanDevice& device,
 
 struct RegisteredPreparation {
     std::size_t codeOffset = 0;
+    bool registration = false;
     ShaderRecompiler::SwappcInfo swappc{};
     std::vector<std::uint32_t> userData;
     std::array<ShaderRecompiler::MemoryRegion, 2> memory{};
@@ -770,6 +771,7 @@ std::unique_ptr<RegisteredPreparation> PlanRegistered(const ShaderSnapshot& snap
     auto plan = std::make_unique<RegisteredPreparation>();
     plan->codeOffset = codeOffset;
     ShaderRecompiler::RdnaProgram decoded;
+    plan->registration = registration;
     try {
         decoded = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(code);
     } catch (const ShaderRecompiler::UnsupportedInstructionError& error) {
@@ -878,6 +880,11 @@ std::vector<PreparedShaders::Entry> PrepareRegistered(RegisteredPreparation& pla
         request.layout.pushConstantSizeBytes = 0;
         append();
     }
+    if (plan.registration) {
+        for (const auto& entry : entries) {
+            if (entry.handle != nullptr) ShaderRecompiler::ReleaseUndrawnResourcePlan(*entry.handle);
+        }
+    }
     return entries;
 }
 

core/libs/prx/libSceAgcDriver/tests/ShaderDiskCache.cpp

@@ -1860,6 +1860,69 @@ void verifyPlanRoundTrip() {
         std::cout << "plan round trip: " << fixture.name << ", " << plan->valueStorage.Size() << " values, " << ShaderDiskCache::EncodePlanEntry(key, *plan).size() << " bytes per entry, " << succeeded << " of 1000 runtimes materialized\n";
         ++index;
     }
+    auto releaseFixture = fixtures.front();
+    releaseFixture.waveSize = 32u;
+    auto registeredRequest = releaseFixture.Request(true);
+    const auto registeredHandle = PrepareShader(registeredRequest);
+    const auto registeredPlan = GetResourcePlan(registeredRequest);
+    ShaderDiskCache::Flush();
+    ReleaseUndrawnResourcePlan(*registeredHandle);
+    const auto beforeReload = ShaderDiskCache::Totals();
+    const auto reloadedPlan = GetResourcePlan(registeredRequest);
+    require(ShaderDiskCache::Totals().planHits == beforeReload.planHits + 1u, "a released registration plan was not reloaded from the disk cache");
+    require(samePlan(*registeredPlan, *reloadedPlan), "a released registration plan changed when reloaded");
+
+    auto reloadFixture = fixtures.front();
+    reloadFixture.waveSize = 32u;
+    reloadFixture.code.insert(reloadFixture.code.begin(), 0xbf800004u);
+    const auto reloadRequest = reloadFixture.Request(true);
+    const auto reloadHandle = PrepareShader(reloadRequest);
+    static_cast<void>(GetResourcePlan(reloadRequest));
+    ShaderDiskCache::Flush();
+    std::vector<std::byte> reloadKey;
+    ShaderDiskCache::BuildPlanKey(reloadRequest, 32u, reloadKey);
+    const auto reloadPath = ShaderDiskCache::EntryDirectory() / ShaderDiskCache::EntryName(reloadKey);
+    require(std::filesystem::is_regular_file(reloadPath), "the reload regression plan was not persisted");
+    const auto preservedPath = reloadPath.string() + ".reload-preserved";
+    std::filesystem::rename(reloadPath, preservedPath);
+    ReleaseUndrawnResourcePlan(*reloadHandle);
+    static_cast<void>(GetResourcePlan(reloadRequest));
+    auto secondLayout = reloadRequest;
+    secondLayout.layout.pushConstantSizeBytes = 0u;
+    ShaderPreparationContext preparation;
+    static_cast<void>(PrepareShader(secondLayout, &preparation));
+    require(preparation.ResourcePrograms() == 1u, "a plan reload retained full IR after an artifact for that layout was already available");
+
+    auto relocatedFixture = fixtures.front();
+    relocatedFixture.waveSize = 32u;
+    relocatedFixture.code.insert(relocatedFixture.code.begin(), 0xbf800005u);
+    const auto relocatedRequest = relocatedFixture.Request(true);
+    const auto relocatedHandle = PrepareShader(relocatedRequest);
+    static_cast<void>(GetResourcePlan(relocatedRequest));
+    ShaderDiskCache::Flush();
+    std::vector<std::byte> relocatedKey;
+    ShaderDiskCache::BuildPlanKey(relocatedRequest, 32u, relocatedKey);
+    const auto relocatedPath = ShaderDiskCache::EntryDirectory() / ShaderDiskCache::EntryName(relocatedKey);
+    std::filesystem::rename(relocatedPath, relocatedPath.string() + ".reload-preserved");
+    ReleaseUndrawnResourcePlan(*relocatedHandle);
+    auto drawn = relocatedRequest;
+    drawn.layout.pushConstantOffsetBytes = 16u;
+    drawn.layout.pushConstantSizeBytes = 112u;
+    const auto relocatedInvocation = PreparedShaderInvocation::TryCreate(drawn, relocatedHandle);
+    require(relocatedInvocation.has_value(), "the prepared artifact refused a relocated push layout");
+    RequestMemoryView drawnMemory(drawn.context.memory);
+    static_cast<void>(relocatedInvocation->Materialize(*relocatedInvocation->Capture(drawnMemory.MakeRuntime(drawn.context.userData, drawn.shader.codeAddress))));
+    auto relocatedProbe = relocatedRequest;
+    relocatedProbe.layout.pushConstantSizeBytes = 0u;
+    ShaderPreparationContext relocatedPreparation;
+    static_cast<void>(PrepareShader(relocatedProbe, &relocatedPreparation));
+    require(relocatedPreparation.ResourcePrograms() == 1u, "a plan reload retained full IR after the prepared artifact served a relocated push layout");
+
+    RequestMemoryView registeredMemory(registeredRequest.context.memory);
+    const auto runtimeCapture = CaptureResources(registeredRequest, registeredMemory.MakeRuntime(registeredRequest.context.userData, registeredRequest.shader.codeAddress), *registeredHandle);
+    const auto beforeKeep = ShaderDiskCache::Totals();
+    ReleaseUndrawnResourcePlan(*registeredHandle);
+    require(GetResourcePlan(registeredRequest) == runtimeCapture->plan && ShaderDiskCache::Totals().planHits == beforeKeep.planHits, "a runtime-captured plan was released");
     const auto covers = [&](const std::function<bool(const IrResourcePlan&)>& predicate) {
         return std::any_of(plans.begin(), plans.end(), [&](const std::shared_ptr<const IrResourcePlan>& plan) { return predicate(*plan); });
     };

core/shader/recompiler/CMakeLists.txt

@@ -248,6 +248,7 @@ function(add_shader_recompiler target runtimeTarget)
             ResolveSource
             ResourceMaterializer::EvaluateNanoseconds
             ResourceMaterializer::MaterializeNanoseconds
+            ReleaseUndrawnResourcePlan
             ResourceMaterializer::SpecializationNanoseconds
             RuntimeAbi::RequireVersion
             RuntimeAbi::HeapCapacity

core/shader/recompiler/Optimization/include/Optimization/ResourceProgram.hpp

@@ -80,6 +80,7 @@ void BuildPreparedShaderKey(const RecompileRequest& request, std::vector<std::ui
 [[nodiscard]] std::shared_ptr<const RecompileResult> RegistrationResult(const RecompileRequest& request, const SourceHandle& handle);
 [[nodiscard]] std::shared_ptr<const RecompileResult> MaterializeShader(const RecompileRequest& request, const ResourceCapture& capture, const SourceHandle& handle);
 [[nodiscard]] std::shared_ptr<const SourceHandle> ResolveSource(const RecompileRequest& request);
+void ReleaseUndrawnResourcePlan(const SourceHandle& handle);
 [[nodiscard]] std::shared_ptr<const ResourceCapture> CaptureResources(const RecompileRequest& request, const SrtRuntime& runtime, const SourceHandle& handle);
 
 }

core/shader/recompiler/Recompiler.cpp

@@ -326,6 +326,7 @@ struct SourceEntry : std::enable_shared_from_this<SourceEntry> {
     std::vector<std::uint32_t> code;
     std::vector<std::uint32_t> linkedCallFallbacks;
     std::shared_ptr<const IrResourcePlan> plan;
+    bool runtimeCaptured = false;
     // A plan build that threw (an unsupported resource chain or control flow) is remembered and
     // rethrown: the front end ran every pass before failing, ~13 ms per dispatch of a shader the
     // title issues every frame (0x1048947300 at the intro video). APS5_NO_FAILURE_MEMO=1 rebuilds.
@@ -396,7 +397,9 @@ bool diskPlans() {
     return enabled && ShaderDiskCache::Enabled() && !DebugProbeActive();
 }
 
-std::shared_ptr<SourceEntry> getSource(const RecompileRequest& request, ShaderPreparationContext* preparation = nullptr, bool reclaim = true) {
+bool sameLayout(const BindingLayout& left, const BindingLayout& right);
+
+std::shared_ptr<SourceEntry> getSource(const RecompileRequest& request, ShaderPreparationContext* preparation = nullptr, bool reclaim = true, const SourceHandle* preparedHandle = nullptr) {
     static PrimedSharedMutex mutex;
     // Entries whose code hashes alike share a bucket; the code comparison picks the right one.
     static std::unordered_map<std::vector<std::uint64_t>, std::vector<std::shared_ptr<SourceEntry>>, SourceKeyHash> sources;
@@ -439,6 +442,8 @@ std::shared_ptr<SourceEntry> getSource(const RecompileRequest& request, ShaderPr
     {
         std::lock_guard lock(source->mutex);
         if (source->plan == nullptr) {
+            const bool hasPreparedArtifact = preparedHandle != nullptr && preparedHandle->artifact != nullptr;
+            if (hasPreparedArtifact && (preparedHandle->source != source || !MatchesPreparedShader(request, *preparedHandle))) throw std::runtime_error("ShaderRecompiler: prepared artifact does not match the resource plan source");
             if (FailureMemo() && source->planFailure) std::rethrow_exception(source->planFailure);
             const bool disk = diskPlans();
             std::vector<std::byte> planKey;
@@ -454,7 +459,9 @@ std::shared_ptr<SourceEntry> getSource(const RecompileRequest& request, ShaderPr
                 try {
                     ResourceProgram resource(request, preparation, reclaim);
                     source->plan = std::move(resource.plan);
-                    source->program = std::move(resource.program);
+                    if (!hasPreparedArtifact && std::ranges::none_of(source->variants, [&](const auto& variant) { return sameLayout(variant->layout, request.layout); })) {
+                        source->program = std::move(resource.program);
+                    }
                 } catch (...) {
                     if (FailureMemo()) source->planFailure = std::current_exception();
                     throw;
@@ -466,6 +473,25 @@ std::shared_ptr<SourceEntry> getSource(const RecompileRequest& request, ShaderPr
     return source;
 }
 
+std::shared_ptr<const IrResourcePlan> acquirePlan(const RecompileRequest& request, const std::shared_ptr<SourceEntry>& source, bool runtimeCapture, const SourceHandle* preparedHandle = nullptr) {
+    for (;;) {
+        {
+            std::lock_guard lock(source->mutex);
+            if (source->plan != nullptr) {
+                if (runtimeCapture) source->runtimeCaptured = true;
+                return source->plan;
+            }
+        }
+        if (getSource(request, nullptr, true, preparedHandle) != source) throw std::runtime_error("ShaderRecompiler: resource plan source changed while reloading");
+    }
+}
+
+void releaseUndrawnPlan(const SourceHandle& handle) {
+    if (handle.source == nullptr || !diskPlans()) return;
+    std::lock_guard lock(handle.source->mutex);
+    if (!handle.source->runtimeCaptured) handle.source->plan.reset();
+}
+
 std::array<std::uint32_t, 3> partialThreads(const RecompileRequest& request) {
     return request.context.compute ? request.context.compute->partialThreads : std::array<std::uint32_t, 3>{};
 }
@@ -1104,7 +1130,7 @@ RecompileResult RecompileImpl(const RecompileRequest& request) {
         return materializeResult(variant, request, snapshot);
     }
     const auto source = getSource(request);
-    materializer.Materialize(*source->plan, runtime, snapshot);
+    materializer.Materialize(*acquirePlan(request, source, true), runtime, snapshot);
     return materializeVariant(*source, request, snapshot);
 }
 
@@ -1312,10 +1338,17 @@ auto recompileReporting(const RecompileRequest& request, Impl&& impl) -> decltyp
 
 std::shared_ptr<const IrResourcePlan> GetResourcePlan(const RecompileRequest& request) {
     static_cast<void>(RequestInputInfo(request));
-    if (request.useCache) return getSource(request)->plan;
+    if (request.useCache) {
+        const auto source = getSource(request);
+        return acquirePlan(request, source, false);
+    }
     return makeResourcePlan(request);
 }
 
+void ReleaseUndrawnResourcePlan(const SourceHandle& handle) {
+    releaseUndrawnPlan(handle);
+}
+
 namespace {
 
 // The capture's materialization; with pure flat slots in the plan the walk's read addresses are
@@ -1344,7 +1377,7 @@ std::shared_ptr<const ResourceCapture> CaptureResources(const RecompileRequest&
     auto capture = std::make_shared<ResourceCapture>();
     if (request.useCache) {
         capture->source = getSource(request);
-        capture->plan = capture->source->plan;
+        capture->plan = acquirePlan(request, capture->source, true);
     } else {
         capture->plan = makeResourcePlan(request);
     }
@@ -1428,13 +1461,13 @@ std::optional<PreparedShaderInvocation> PreparedShaderInvocation::TryCreate(cons
 std::shared_ptr<const ResourceCapture> PreparedShaderInvocation::Capture(const SrtRuntime& runtime) const {
     auto capture = std::make_shared<ResourceCapture>();
     capture->source = handle->source;
-    capture->plan = handle->source->plan;
+    capture->plan = acquirePlan(request, handle->source, true, handle.get());
     materializeCapture(*capture, runtime);
     return capture;
 }
 
 std::shared_ptr<const RecompileResult> PreparedShaderInvocation::Materialize(const ResourceCapture& capture) const {
-    if (capture.source != handle->source || capture.plan != handle->source->plan) throw std::runtime_error("ShaderRecompiler: resource capture belongs to another prepared shader");
+    if (capture.source != handle->source) throw std::runtime_error("ShaderRecompiler: resource capture belongs to another prepared shader");
     if (request.useCache && ResultMemo()) return materializePreparedMemoized(*handle->source, handle->artifact, request, capture.snapshot, true, nullptr);
     auto result = std::make_shared<RecompileResult>(materializeResult(*handle->artifact, request, capture.snapshot));
     result->cacheHit = true;
@@ -1462,7 +1495,7 @@ std::shared_ptr<const RecompileResult> RegistrationResult(const RecompileRequest
 
 std::shared_ptr<const RecompileResult> MaterializeShader(const RecompileRequest& request, const ResourceCapture& capture, const SourceHandle& handle) {
     if (!MatchesPreparedShader(request, handle)) throw std::runtime_error("ShaderRecompiler: prepared artifact does not match the static ABI");
-    if (capture.source != handle.source || capture.plan != handle.source->plan) throw std::runtime_error("ShaderRecompiler: resource capture belongs to another prepared shader");
+    if (capture.source != handle.source) throw std::runtime_error("ShaderRecompiler: resource capture belongs to another prepared shader");
     if (request.useCache && ResultMemo()) return materializePreparedMemoized(*handle.source, handle.artifact, request, capture.snapshot, true, nullptr);
     auto result = std::make_shared<RecompileResult>(materializeResult(*handle.artifact, request, capture.snapshot));
     result->cacheHit = true;
@@ -1478,7 +1511,7 @@ std::shared_ptr<const ResourceCapture> CaptureResources(const RecompileRequest&
     if (request.shader.stage != ShaderStage::Compute && request.shader.stage != ShaderStage::Fragment) static_cast<void>(RequestInputInfo(request));
     auto capture = std::make_shared<ResourceCapture>();
     capture->source = handle.source;
-    capture->plan = handle.source->plan;
+    capture->plan = acquirePlan(request, handle.source, true, &handle);
     if (profile) capture->sourceNanoseconds = static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - started).count());
     materializeCapture(*capture, runtime);
     return capture;