shader: keep scalar payload and terminal conditions on the GPU

src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp

@@ -741,7 +741,31 @@ static std::vector<ResourceBlock> ResourceControlFlow(const Program& program, co
 		block.sources.erase(std::unique(block.sources.begin(), block.sources.end()),
 		                    block.sources.end());
 	}
-	for (auto& block: blocks) block.condition = predicate(block.condition);
+	// Only branches that can change future resource access belong in the host plan.
+	std::vector<uint8_t> resource_reachable(blocks.size());
+	std::vector<uint32_t> pending;
+	for (uint32_t i = 0; i < blocks.size(); ++i) {
+		if (!blocks[i].sources.empty() || !blocks[i].srt_reads.empty()) {
+			resource_reachable[i] = 1;
+			pending.push_back(i);
+		}
+	}
+	while (!pending.empty()) {
+		const auto index = pending.back();
+		pending.pop_back();
+		for (const auto* predecessor: program.blocks[index]->ImmPredecessors()) {
+			const auto found = indices.find(predecessor);
+			if (found == indices.end()) return {};
+			if (resource_reachable[found->second]) continue;
+			resource_reachable[found->second] = 1;
+			pending.push_back(found->second);
+		}
+	}
+	for (auto& block: blocks) {
+		block.condition = std::ranges::any_of(block.successors, [&](uint32_t successor) {
+			return resource_reachable[successor] != 0;
+		}) ? predicate(block.condition) : Value {};
+	}
 	if (std::ranges::none_of(
 	        blocks, [](const ResourceBlock& block) { return !block.condition.IsEmpty(); })) {
 		return {};

src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp

@@ -807,15 +807,6 @@ private:
 				}
 			}
 		}
-		for (auto* block: m_program.blocks) {
-			for (auto& inst: *block) {
-				uint32_t index = 0;
-				if (inst.GetOpcode() == ValueOpcode::LoadAddressU32 &&
-				    ScalarReadMemory(inst, index) != nullptr && inst.Arg(1).Resolve().IsImmediate() &&
-				    ValidateRuntimeValue(m_program, Value(&inst)))
-					CollectScalarRead(Value(&inst), inst.Flags<MemoryFlags>().pc);
-			}
-		}
 		for (auto* read: m_scalar_reads) {
 			const auto flags = read->Flags<MemoryFlags>();
 			auto& memory = m_program.memory_info[flags.index];

tests/ResourceMaterializationTests.cpp

@@ -435,7 +435,7 @@ void TestUnbasedFlatCacheHitMaterializes() {
         "unbased FLAT plan produced unexpected descriptors");
 }
 
-void TestWrittenDescriptorPredicateReads(bool memory_condition) {
+void TestWrittenDescriptorPredicateReads(bool memory_condition, bool loop) {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   Program program;
   program.stage = Libs::Graphics::ShaderType::Compute;
@@ -467,12 +467,26 @@ void TestWrittenDescriptorPredicateReads(bool memory_condition) {
       Libs::Graphics::ShaderRecompiler::CFG::TerminatorKind::ConditionalBranch;
   program.blocks[0]->terminator.true_block = program.blocks[1];
   program.blocks[0]->terminator.false_block = program.blocks[2];
+
   program.blocks[1]->id = 1;
   program.blocks[1]->terminator.kind =
       Libs::Graphics::ShaderRecompiler::CFG::TerminatorKind::Return;
   program.blocks[2]->id = 2;
   program.blocks[2]->terminator.kind =
       Libs::Graphics::ShaderRecompiler::CFG::TerminatorKind::Return;
+  if (loop) {
+    // The predicate reaches the store through a continue block, then loops back.
+    auto &continued = AddValueBlock(program);
+    program.blocks[0]->terminator.true_block = &continued;
+    continued.terminator.kind = Libs::Graphics::ShaderRecompiler::CFG::TerminatorKind::Branch;
+    continued.terminator.true_block = &store_block;
+    continued.AddBranch(&store_block);
+    store_block.terminator.kind = Libs::Graphics::ShaderRecompiler::CFG::TerminatorKind::Branch;
+    store_block.terminator.true_block = █
+    store_block.AddBranch(&block);
+  }
+  block.AddBranch(program.blocks[0]->terminator.true_block);
+  block.AddBranch(program.blocks[2]);
   program.memory_info.push_back({.kind = ResourceKind::Buffer, .resource = 0});
   auto &output = store_block.AppendNewInst(ValueOpcode::GetBufferResource,
       {source.dwords[0], source.dwords[1], source.dwords[2], source.dwords[3]});
@@ -677,7 +691,8 @@ int main() {
   TestExactReciprocalDescriptorArithmetic();
   TestUnbasedFlatCacheHitMaterializes();
   for (const bool memory_condition : {false, true})
-    TestWrittenDescriptorPredicateReads(memory_condition);
+    for (const bool loop : {false, true})
+      TestWrittenDescriptorPredicateReads(memory_condition, loop);
   TestFailedMaterializationRejectsStage();
   TestFiniteImageRefreshReusesScalarReads();
   TestMixedSamplerVariantsShareRuntimeDescriptor();

tests/ShaderRecompilerComputeTests.cpp

@@ -28387,7 +28387,7 @@ TestCase ScalarMemoryLoadVariants() {
   std::vector<u32> expected = initial;
   expected.insert(expected.end(), initial.begin(), initial.end());
 
-  return {"ScalarMemoryLoadVariants",
+  TestCase test{"ScalarMemoryLoadVariants",
           code,
           initial,
           expected,
@@ -28396,6 +28396,9 @@ TestCase ScalarMemoryLoadVariants() {
            O::S_BUFFER_LOAD_DWORDX2, O::S_BUFFER_LOAD_DWORDX4,
            O::S_BUFFER_LOAD_DWORDX8, O::S_BUFFER_LOAD_DWORDX16, O::V_MOV_B32,
            O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+  test.bda_mappings = {{0, 0}};
+  test.required_spirv = {"get_bda_pointer"};
+  return test;
 }
 
 TestCase ScalarBufferOffsetAlignmentAndCarry() {
@@ -28439,11 +28442,43 @@ TestCase ScalarLoadSignedImmediateOffsetAddsSoffset() {
   AppendStoreSgpr(&code, 1, 0);
   AppendEnd(&code);
 
-  return {"ScalarLoadSignedImmediateOffsetAddsSoffset",
+  TestCase test{"ScalarLoadSignedImmediateOffsetAddsSoffset",
           code,
           {0x11111111u, 0x22222222u},
           {0x22222222u, 0x22222222u},
           {O::S_MOV_B32, O::S_LOAD_DWORD, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+  test.bda_mappings = {{0, 0}};
+  test.required_spirv = {"get_bda_pointer"};
+  return test;
+}
+
+TestCase ScalarPayloadBeforeAliasedAtomic() {
+  using O = ShaderOpcode;
+  constexpr u32 GuestBase = 0x10000u;
+  std::vector<u32> code;
+  AppendSMovLiteral(&code, 8, GuestBase);
+  AppendSMovLiteral(&code, 9, 0u);
+  code.insert(code.end(), {EncodeSmem0(0x00, 12, 4), EncodeSmem1(24, 125),
+                           EncodeSopp(0x0c, 0xc07f)});
+  AppendVMovU32(&code, 0, 1u << 20u);
+  AppendVMovU32(&code, 1, 0u);
+  // Same 64-bit OR and byte offset as SAROS b62b494cb2567011.
+  code.insert(code.end(), {0xe1680018u, 0x80000000u});
+  AppendStoreSgpr(&code, 12, 0);
+  AppendEnd(&code);
+  TestCase test{"ScalarPayloadBeforeAliasedAtomic", code,
+                {0u, 0u, 0u, 0u, 0u, 0u, 0x10u, 0x11223344u},
+                {0x10u, 0u, 0u, 0u, 0u, 0u, 0x100010u, 0x11223344u},
+                {O::S_MOV_B32, O::S_LOAD_DWORD, O::S_WAITCNT, O::V_MOV_B32,
+                 O::BUFFER_ATOMIC_OR_X2, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+  test.bda_mappings = {{GuestBase, 0}};
+  test.required_spirv = {"get_bda_pointer"};
+  test.forbidden_spirv = {"flattened_srt"};
+  test.compute_info.threads_num[0] = 1;
+  test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1;
+  test.compute_info.workgroup_register = 4;
+  test.has_compute_info = true;
+  return test;
 }
 
 TestCase BufferLoadStore() {
@@ -31034,23 +31069,26 @@ TestCase BranchVccnzUsesCarryProducedWaveMask() {
   return test;
 }
 
-TestCase ScalarLoadAlignsComponentsAndMasksAddress() {
+TestCase ScalarLoadAlignsComponents() {
   using O = ShaderOpcode;
 
   std::vector<u32> code;
   AppendSMovLiteral(&code, 0, 1u);
   AppendSMovLiteral(&code, 2, 3u);
-  AppendSMovLiteral(&code, 3, 0xffff0000u);
+  AppendSMovLiteral(&code, 3, 0u); // S_LOAD uses the full base; only buffer descriptors mask it.
   code.push_back(EncodeSmem0(0x00, 1, 1));
   code.push_back(EncodeSmem1(3, 0));
   AppendStoreSgpr(&code, 1, 0);
   AppendEnd(&code);
 
-  return {"ScalarLoadAlignsComponentsAndMasksAddress",
+  TestCase test{"ScalarLoadAlignsComponents",
           code,
           {0x11111111u, 0x22222222u},
           {0x11111111u, 0x22222222u},
           {O::S_MOV_B32, O::S_LOAD_DWORD, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+  test.bda_mappings = {{0, 0}};
+  test.required_spirv = {"get_bda_pointer"};
+  return test;
 }
 
 TestCase ScalarLoadAlignsDynamicBase() {
@@ -36691,7 +36729,8 @@ std::vector<TestCase> MakeCases() {
   AddCase(ScalarMemoryLoadVariants);
   AddCase(ScalarBufferOffsetAlignmentAndCarry);
   AddCase(ScalarLoadSignedImmediateOffsetAddsSoffset);
-  AddCase(ScalarLoadAlignsComponentsAndMasksAddress);
+  AddCase(ScalarPayloadBeforeAliasedAtomic);
+  AddCase(ScalarLoadAlignsComponents);
   AddCase(ScalarLoadAlignsDynamicBase);
   cases.push_back(ScalarBufferFromLoopReadlane(32));
   cases.push_back(ScalarBufferFromLoopReadlane(64));
@@ -39615,17 +39654,19 @@ void CheckResourcePlanHandoff() {
     ValidateSpirv(name, compiled.spirv);
     const auto &bindings = compiled.program.bindings;
     Require(name, "GPU data requirements",
-            compiled.program.info.user_data_registers.empty() && bindings.descriptor_counts[0] == 1 &&
+            compiled.program.info.user_data_registers ==
-                (bindings.descriptor_counts[static_cast<size_t>(DescriptorBindingKind::FlattenedSrt)] != 0) ==
+                (numeric_read ? std::vector<u32>{0, 1} : std::vector<u32>{}) &&
+                bindings.descriptor_counts[0] == 1 &&
+                bindings.descriptor_counts[static_cast<size_t>(DescriptorBindingKind::FlattenedSrt)] == 0 &&
+                (bindings.descriptor_counts[static_cast<size_t>(DescriptorBindingKind::BdaPagetable)] != 0) ==
                     numeric_read,
-            "descriptor-only shader retained SRT uploads or live scalar data was removed");
+            "descriptor-only shader retained SRT uploads or native scalar payload lost its address");
 
     memory[0] = 0x2000u;
-    memory[4] = 13u;
     Require(name, "independent host plan",
             MaterializeResources(plan, runtime, snapshot, specialization) &&
                 snapshot.buffers[0].dwords[0] == 0x2000u &&
-                (!numeric_read || snapshot.flattened_srt.back() == 13u),
+                snapshot.flattened_srt.size() == 4u,
             "compiled shader cleanup invalidated or froze the host resource plan");
   }
   std::printf("[host]    %-32s ok\n", name);
@@ -42591,6 +42632,10 @@ int main(int argc, char **argv) {
     VulkanHarness vulkan;
     RunCase(&vulkan, BufferOffsetsUsePackedWords(false));
     RunCase(&vulkan, BufferOffsetsUsePackedWords(true));
+    RunCase(&vulkan, ScalarMemoryLoadVariants());
+    RunCase(&vulkan, ScalarLoadSignedImmediateOffsetAddsSoffset());
+    RunCase(&vulkan, ScalarLoadAlignsComponents());
+    RunCase(&vulkan, ScalarPayloadBeforeAliasedAtomic());
     return 0;
   }
   if (argc == 2 && std::strcmp(argv[1], "--buffer-publication-only") == 0) {

tests/shaderCfgTests.cpp

@@ -13804,6 +13804,59 @@ void TestUniformSelectedDescriptorLoadAddress() {
   }
 }
 
+void TestNativeScalarAtomicPayloadStaysOnGpu() {
+  using namespace ShaderRecompiler::IR;
+  // Exact final block from SAROS b62b494cb2567011: read a mutable flag,
+  // atomically update that address, then test the loaded flag before S_TRAP.
+  const uint32_t shader[] = {
+      0x7e0002ffu, 0x00100000u, 0x7e020280u, 0xf4040100u,
+      0xfa000018u, 0xbf8cc07fu, 0x8801ff05u, 0x01000000u,
+      0xbe800304u, 0xbe820381u, 0xbe8303ffu, 0x00016204u,
+      0xf4001a82u, 0xfa000018u, 0xe1680018u, 0x80000000u,
+      0xbf8cc07fu, 0xbf0d846au, 0xbf840001u, 0xbf920001u,
+      0xbf810000u,
+  };
+  const std::array<uint32_t, 2> user_data{0x1000u, 0u};
+  auto options = MakeCompileOptions(ShaderType::Compute);
+  options.user_data = user_data;
+  auto translated = ShaderRecompiler::TranslateProgram(shader, options);
+  const auto &program = translated.program;
+  const Inst *payload = nullptr;
+  for (const auto *block : program.blocks) {
+    for (const auto &inst : *block) {
+      if (inst.GetOpcode() == ValueOpcode::LoadAddressU32 &&
+          program.memory_info[inst.Flags<MemoryFlags>().index].kind == ResourceKind::ScalarAddress)
+        payload = &inst;
+    }
+  }
+  Check(payload != nullptr && payload->Parent() != nullptr &&
+            !program.memory_info[payload->Flags<MemoryFlags>().index].planning_only &&
+            program.info.uses_dma && program.info.buffers.size() == 1 &&
+            program.info.buffers[0].atomic,
+        "native scalar payload was replaced by a host snapshot before its atomic write");
+  auto plan = ExtractResourcePlan(program);
+  Check(plan.srt_reads.size() == 2 && plan.control_flow.empty() &&
+            !plan.capture_specialization_reads,
+        "terminal native shader assertion entered resource planning");
+  struct Reads { uint32_t descriptors = 0; uint32_t payload = 0; } reads;
+  const auto read = +[](void *data, uint64_t address, std::span<uint32_t> words) {
+    auto &reads = *static_cast<Reads *>(data);
+    if (address == 0x2018u) { ++reads.payload; return false; }
+    if (words.size() != 1 || (address != 0x1018u && address != 0x101cu)) return false;
+    ++reads.descriptors;
+    words[0] = address == 0x1018u ? 0x2000u : 0u;
+    return true;
+  };
+  const SrtRuntime runtime{.user_data = user_data, .read_memory = read,
+                           .userdata = &reads, .read_specialization_memory = read};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+            reads.descriptors == 2 && reads.payload == 0 &&
+            snapshot.specialization_reads.empty() && snapshot.buffers[0].dwords[0] == 0x2000u,
+        "native resource materialization read the mutable flag or lost its atomic descriptor");
+}
+
 void TestTypedDescriptorRealCarryAndScalarLoads() {
   const uint32_t carry_shader[] = {
       EncodeSop1(0x1f, 0, 0),      // s_getpc_b64 s[0:1]
@@ -15277,6 +15330,7 @@ int main() {
   TestGpuProducedWritableDescriptor();
   TestUniformSelectedWritableDescriptor();
   TestUniformSelectedDescriptorLoadAddress();
+  TestNativeScalarAtomicPayloadStaysOnGpu();
   TestImmutableDescriptorPredicate();
   TestTypedDescriptorRealCarryAndScalarLoads();
   TestSrtWalkerRealSmemTranslation();