fix(recompiler): leave out descriptor sources behind a branch the capture decides

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

@@ -610,6 +610,89 @@ void verifyPureFlatSlots() {
     require(Detail::ComputePureFlatSlots(plan) == Pure{0, 0, 0, 0}, "pure flat slots: an incomplete plan was classified");
 }
 
+void verifyGuardedDescriptorSources() {
+    using namespace ShaderRecompiler;
+    constexpr std::uint32_t Format8888UNorm = 56;
+    constexpr std::uint32_t Type2D = 9;
+    struct alignas(256) Texture { std::array<std::uint8_t, 256> bytes{}; };
+    static Texture texture;
+    const auto base = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(texture.bytes.data()));
+    const std::array<std::uint32_t, 8> valid{static_cast<std::uint32_t>(base >> 8u), static_cast<std::uint32_t>((base >> 40u) & 0xffu) | (Format8888UNorm << 20u) | (3u << 30u), 3u << 14u, 0xfacu | (Type2D << 28u), 0u, 0u, 0u, 0u};
+    const std::array<std::uint32_t, 8> invalid{0x12345678u, 0xffffffffu, 0xffffffffu, 0u, 0u, 0u, 0u, 0u};
+    const std::array<std::uint32_t, 4> sampler{0u, 0u, 0u, 0u};
+    std::array<std::uint32_t, 4> output{};
+    const auto outputAddress = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(output.data()));
+    const std::array<std::uint32_t, 4> outputV{static_cast<std::uint32_t>(outputAddress), static_cast<std::uint32_t>((outputAddress >> 32u) & 0xffffu), 16u, 0xfacu};
+    std::array<std::uint32_t, 32> srt{};
+    const auto fill = [&](const std::array<std::uint32_t, 8>& image, std::uint32_t mode) {
+        std::copy(image.begin(), image.end(), srt.begin());
+        std::copy(sampler.begin(), sampler.end(), srt.begin() + 8);
+        std::copy(outputV.begin(), outputV.end(), srt.begin() + 16);
+        srt[20] = mode;
+    };
+    const auto srtAddress = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(srt.data()));
+    const std::array<std::uint32_t, 2> userData{static_cast<std::uint32_t>(srtAddress), static_cast<std::uint32_t>(srtAddress >> 32u)};
+    const std::array<std::uint32_t, 4> capabilities{29u, spv::CapabilityInt64, spv::CapabilityPhysicalStorageBufferAddresses, spv::CapabilityStorageBuffer8BitAccess};
+    const std::array<std::string_view, 2> extensions{"SPV_KHR_physical_storage_buffer", "SPV_KHR_8bit_storage"};
+    RecompileRequest request{};
+    const std::vector<std::uint32_t> code{0xf4000400u, 0xfa000050u, 0xf4080500u, 0xfa000020u, 0xf4080600u, 0xfa000040u, 0xbf8cc07fu, 0xbf068010u, 0xbf850006u, 0xf40c0200u, 0xfa000000u, 0xbf8cc07fu, 0xf09c0f08u, 0x00a20000u, 0xbf8c3f70u, 0xe0700000u, 0x80060000u, 0xbf810000u};
+    request.shader = {ShaderStage::Compute, 0x22000u, code, 0, {}};
+    request.context.waveSize = 32;
+    request.context.userDataBaseRegister = 0;
+    request.context.userData = userData;
+    request.context.compute = ShaderComputeStageInfo{{32u, 1u, 1u}, 0u, {false, false, false}, false, 1u};
+    request.target.vulkanVersion = 0x00401000u;
+    request.target.spirvVersion = 0x00010300u;
+    request.target.subgroupSize = 32;
+    request.target.supportedCapabilities = capabilities;
+    request.target.supportedExtensions = extensions;
+    request.target.bdaAbiVersion = BdaAbi::Version;
+    request.target.fragmentShaderBarycentricEnabled = false;
+    request.layout.pushConstantSizeBytes = 128;
+
+    fill(invalid, 0u);
+    const auto plan = GetResourcePlan(request);
+    require(plan->info.images.size() == 1u && plan->info.samplers.size() == 1u, "guarded sources: the image and its sampler were not planned");
+    require(!plan->controlFlow.empty(), "guarded sources: the branch around the image was not planned");
+    {
+        AgcDriver::ShaderMemory memory({});
+        const auto capture = memory.Capture(request);
+        require(capture->snapshot.images.size() == 1u && std::ranges::all_of(capture->snapshot.images[0].dwords, [](std::uint32_t word) { return word == 0u; }), "guarded sources: the skipped T# was not left null");
+    }
+    fill(invalid, 1u);
+    {
+        AgcDriver::ShaderMemory memory({});
+        expectFailure([&] { static_cast<void>(memory.Capture(request)); }, "runtime image descriptor is invalid", "guarded sources: an invalid T# behind a taken branch was accepted");
+    }
+    fill(valid, 1u);
+    {
+        AgcDriver::ShaderMemory memory({});
+        const auto capture = memory.Capture(request);
+        require(capture->snapshot.images.size() == 1u && std::equal(valid.begin(), valid.end(), capture->snapshot.images[0].dwords.begin()), "guarded sources: the T# behind a taken branch was not read");
+    }
+    fill(valid, 0u);
+    {
+        AgcDriver::ShaderMemory memory({});
+        const auto capture = memory.Capture(request);
+        require(capture->snapshot.images.size() == 1u && std::ranges::all_of(capture->snapshot.images[0].dwords, [](std::uint32_t word) { return word == 0u; }), "guarded sources: a valid T# behind a skipped branch was read");
+    }
+
+    const std::vector<std::uint32_t> storeCode{0xf4000400u, 0xfa000050u, 0xbf8cc07fu, 0xbf068010u, 0xbf850006u, 0xf40c0100u, 0xfa000000u, 0xbf8cc07fu, 0xf0200f08u, 0x00010004u, 0xbf820005u, 0xf40c0100u, 0xfa000000u, 0xbf8cc07fu, 0xf0200f10u, 0x00010004u, 0xbf810000u};
+    auto storeRequest = request;
+    storeRequest.shader = {ShaderStage::Compute, 0x23000u, storeCode, 0, {}};
+    const auto compileStore = [&](std::uint32_t mode) {
+        fill(valid, mode);
+        AgcDriver::ShaderMemory memory({});
+        const auto capture = memory.Capture(storeRequest);
+        storeRequest.context.memory = memory.Regions();
+        return Recompile(storeRequest, *capture);
+    };
+    const auto storePlan = GetResourcePlan(storeRequest);
+    require(storePlan->info.images.size() == 2u && storePlan->descriptorSources.size() == 1u, "guarded images: the two stores were not planned as two images of one T#");
+    static_cast<void>(compileStore(1u));
+    expectFailure([&] { static_cast<void>(compileStore(0u)); }, "incompatible with the static runtime image interface", "guarded images: a store of the wrong dimension behind a taken branch was accepted");
+}
+
 // A compute program sampling a T# loaded from a table buffer at a runtime key (a bindless image
 // table): mode M enumerates the keys from the material records, mode T binds the whole table.
 void verifyBindlessTable() {
@@ -3097,6 +3180,7 @@ void verifyGuardedNullPointers() {
     const auto tableAddress = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(table.data()));
     vertexRoot[2] = static_cast<std::uint32_t>(tableAddress);
     vertexRoot[3] = static_cast<std::uint32_t>(tableAddress >> 32u);
+    vertexRoot[0] = 1u;
     AgcDriver::ShaderMemory mappedVertexMemory({});
     const auto mappedVertex = mappedVertexMemory.Capture(vertex.request);
     require(mappedVertex->snapshot.srtPoison.empty() && mappedVertex->snapshot.buffers.size() == 1u && mappedVertex->snapshot.buffers[0].dwords[0] == table[4], "guarded pointer: a vertex V# behind a mapped pointer was poisoned or not read");
@@ -3188,6 +3272,7 @@ int main(int argc, char** argv) {
         verifyOutOfRangeConstantReads();
         verifySignedMinimum();
         verifyPureFlatSlots();
+        verifyGuardedDescriptorSources();
         verifyBindlessTable();
         verifyDescriptorPhis();
         verifyLaneSelectedDescriptors();

core/shader/recompiler/IntermediateRepresentation/include/IntermediateRepresentation/IrMetadata/ControlFlowInfo.hpp

@@ -62,6 +62,7 @@ struct ResourceBlock {
     IrValue* condition = nullptr;
     std::vector<std::uint32_t> successors;
     std::vector<std::uint32_t> sources;
+    std::vector<std::uint32_t> images;
 };
 
 }

core/shader/recompiler/Optimization/src/ResourceMaterializer.cpp

@@ -487,6 +487,11 @@ void materializeSnapshot(const IrResourcePlan& plan, const SrtRuntime& runtime,
         if (descriptor.dwordCount != 8u) {
             throw std::runtime_error("image descriptor has an invalid width");
         }
+        const auto activity = plan.descriptorSources.size() + i;
+        if (activity < activeSources.size() && activeSources[activity] == 0u) {
+            descriptor = DescriptorValue{};
+            descriptor.dwordCount = 8u;
+        }
         if (!validImageDescriptor(descriptor, image.r128) && !nullImageDescriptor(descriptor)) {
             throw std::runtime_error("runtime image descriptor is invalid");
         }

core/shader/recompiler/Optimization/src/ResourceTracker.cpp

@@ -114,11 +114,21 @@ class Tracker {
         }
         SplitDescriptorPhis();
         PlanIndirectImages();
+        std::unordered_map<const IrBlock*, std::uint32_t> blockIndices;
+        for (const IrBlock* block : m_program.BlockOrder()) {
+            blockIndices.emplace(block, static_cast<std::uint32_t>(blockIndices.size()));
+        }
+        m_blockSources.assign(blockIndices.size(), {});
+        m_blockImages.assign(blockIndices.size(), {});
         for (auto& block : m_program.Blocks()) {
+            const auto found = blockIndices.find(block.get());
+            m_currentBlock = found == blockIndices.end() ? std::numeric_limits<std::uint32_t>::max() : found->second;
             for (IrValue* inst : block->Instructions()) {
                 Collect(*inst);
             }
         }
+        m_currentBlock = std::numeric_limits<std::uint32_t>::max();
+        PlanResourceControlFlow();
         if (m_bdaTraces > bdaTraceLimit) std::fprintf(stderr, "[bda] %u more address accesses in this program not shown\n", m_bdaTraces - bdaTraceLimit);
         LinkImageAliases();
         for (const auto& patch : m_handlePatches) {
@@ -241,13 +251,104 @@ class Tracker {
                 same = same && EquivalentValue(m_program.Resources(), current.dwords[i], descriptor.dwords[i]);
             }
             if (same) {
+                NoteBlockSource(candidate);
                 return candidate;
             }
         }
         m_sources.push_back(descriptor);
+        NoteBlockSource(static_cast<std::uint32_t>(m_sources.size() - 1));
         return static_cast<std::uint32_t>(m_sources.size() - 1);
     }
 
+    void NoteBlockImage(std::uint32_t image) {
+        if (m_currentBlock >= m_blockImages.size()) {
+            return;
+        }
+        auto& images = m_blockImages[m_currentBlock];
+        if (std::ranges::find(images, image) == images.end()) {
+            images.push_back(image);
+        }
+    }
+
+    void NoteBlockSource(std::uint32_t source) {
+        if (m_currentBlock >= m_blockSources.size()) {
+            return;
+        }
+        auto& sources = m_blockSources[m_currentBlock];
+        if (std::ranges::find(sources, source) == sources.end()) {
+            sources.push_back(source);
+        }
+    }
+
+    void PlanResourceControlFlow() {
+        const auto& order = m_program.BlockOrder();
+        const auto& info = m_program.Metadata().blockInfo;
+        if (order.empty() || order.size() != info.size()) {
+            return;
+        }
+        std::unordered_map<const IrBlock*, std::uint32_t> indexOf;
+        std::unordered_map<std::uint32_t, std::uint32_t> idToIndex;
+        for (std::uint32_t index = 0; index < order.size(); index++) {
+            indexOf.emplace(order[index], index);
+            idToIndex.emplace(info[index].id, index);
+        }
+        std::vector<ResourceBlock> controlFlow(order.size());
+        for (std::uint32_t index = 0; index < order.size(); index++) {
+            auto& block = controlFlow[index];
+            const auto& terminator = info[index].terminator;
+            const auto target = [&](std::uint32_t id) { return idToIndex.find(id); };
+            const auto trueEdge = target(terminator.trueBlock);
+            const auto falseEdge = target(terminator.falseBlock);
+            if (terminator.kind == TerminatorKind::ConditionalBranch && !terminator.loopHeader && info[index].condition != nullptr && trueEdge != idToIndex.end() && falseEdge != idToIndex.end()) {
+                block.condition = info[index].condition;
+                block.successors = {trueEdge->second, falseEdge->second};
+            } else if (terminator.kind != TerminatorKind::Return) {
+                for (const IrBlock* successor : order[index]->Successors()) {
+                    const auto found = indexOf.find(successor);
+                    if (found == indexOf.end()) {
+                        return;
+                    }
+                    block.successors.push_back(found->second);
+                }
+            }
+            block.sources = m_blockSources[index];
+            block.images = m_blockImages[index];
+        }
+        std::vector<std::uint8_t> unconditional(order.size());
+        for (std::uint32_t index = 0; index < order.size() && unconditional[index] == 0u;) {
+            unconditional[index] = 1u;
+            if (controlFlow[index].successors.size() != 1u) {
+                break;
+            }
+            index = controlFlow[index].successors.front();
+        }
+        std::vector<std::uint8_t> reached(m_sources.size());
+        std::vector<std::uint8_t> reachedImages(m_info.images.size());
+        for (std::uint32_t index = 0; index < order.size(); index++) {
+            if (unconditional[index] == 0u) {
+                continue;
+            }
+            for (const auto source : controlFlow[index].sources) {
+                reached[source] = 1u;
+            }
+            for (const auto image : controlFlow[index].images) {
+                reachedImages[image] = 1u;
+            }
+        }
+        bool guarded = false;
+        for (const auto& block : controlFlow) {
+            for (const auto source : block.sources) {
+                guarded = guarded || reached[source] == 0u;
+            }
+            for (const auto image : block.images) {
+                guarded = guarded || reachedImages[image] == 0u;
+            }
+        }
+        if (guarded) {
+            m_program.Resources().controlFlow = std::move(controlFlow);
+        }
+    }
+
     static bool immediateU32(IrValue* value, std::uint32_t& result) {
         value = value->Resolve();
         if (!value->HasImmediate() || value->Type() != IrType::U32) {
@@ -1330,6 +1431,7 @@ class Tracker {
             GetHandle(inst.Argument(0), IrOpcode::GetImageResource, 8, handle, source);
         }
         resource = AddImage(source, memory, op, flags.pc);
+        NoteBlockImage(resource);
         if (resource == std::numeric_limits<std::uint32_t>::max()) {
             fail("image resource limit exceeded");
         }
@@ -1384,6 +1486,9 @@ class Tracker {
     ShaderInfo m_info;
     IrBuilder m_builder;
     std::vector<DescriptorSource> m_sources;
+    std::vector<std::vector<std::uint32_t>> m_blockSources;
+    std::vector<std::vector<std::uint32_t>> m_blockImages;
+    std::uint32_t m_currentBlock = std::numeric_limits<std::uint32_t>::max();
     std::vector<HandlePatch> m_handlePatches;
     std::vector<MemoryPatch> m_memoryPatches;
     // APS5_TRACE_BDA: address accesses seen by Collect (the first bdaTraceLimit are printed).

core/shader/recompiler/Optimization/src/SrtWalker/SrtDescriptorEvaluation.cpp

@@ -63,16 +63,23 @@ bool EvaluateRuntimeSourcesImpl(const IrResourcePlan& program, std::span<const s
     Evaluator evaluator(program, runtime, cleanFlatSlots, &cleanEvaluator);
     std::vector<std::uint8_t> active;
     if (evaluateFlat) {
-        active.assign(program.descriptorSources.size(), 1u);
+        active.assign(program.descriptorSources.size() + program.info.images.size(), 1u);
     }
     if (evaluateFlat && !program.controlFlow.empty()) {
         for (const auto& block : program.controlFlow) {
             for (const auto source : block.sources) {
                 active.at(source) = 0u;
             }
+            for (const auto image : block.images) {
+                active.at(program.descriptorSources.size() + image) = 0u;
+            }
         }
         std::vector<std::uint8_t> visited(program.controlFlow.size());
         std::vector<std::uint32_t> pending {0};
+        InaccessibleRead conditionRead;
+        std::uint32_t conditionNullRoots = 0;
+        cleanEvaluator.ReportInaccessibleReads(&conditionRead);
+        cleanEvaluator.ReportNullRootReads(&conditionNullRoots);
         while (!pending.empty()) {
             const auto index = pending.back();
             pending.pop_back();
@@ -84,14 +91,20 @@ bool EvaluateRuntimeSourcesImpl(const IrResourcePlan& program, std::span<const s
             for (const auto source : block.sources) {
                 active[source] = 1u;
             }
+            for (const auto image : block.images) {
+                active[program.descriptorSources.size() + image] = 1u;
+            }
             std::uint32_t condition = 0;
-            const bool cleanEvaluable = block.condition != nullptr && runtime.readSpecializationMemory != nullptr && cleanEvaluator.Evaluate(block.condition, condition);
+            const auto nullRootsBefore = conditionNullRoots;
+            const bool cleanEvaluable = block.condition != nullptr && runtime.readSpecializationMemory != nullptr && cleanEvaluator.Evaluate(block.condition, condition) && conditionNullRoots == nullRootsBefore;
             if (cleanEvaluable) {
                 pending.push_back(block.successors[condition != 0u ? 0u : 1u]);
             } else {
                 pending.insert(pending.end(), block.successors.begin(), block.successors.end());
             }
         }
+        cleanEvaluator.ReportInaccessibleReads(nullptr);
+        cleanEvaluator.ReportNullRootReads(nullptr);
     }
     InaccessibleRead inaccessible;
     std::uint32_t nullRoots = 0;