shader: keep scalar data native and preserve raw addresses

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

@@ -655,6 +655,7 @@ static std::vector<ResourceBlock> ResourceControlFlow(const Program& program, co
 		}
 	}
 	std::vector<ResourceBlock> blocks(program.blocks.size());
+	std::vector<std::vector<uint32_t>> predecessors(blocks.size());
 	for (uint32_t i = 0; i < blocks.size(); i++) {
 		auto&                 block      = blocks[i];
 		const auto&           info       = program.block_info[i];
@@ -678,6 +679,7 @@ static std::vector<ResourceBlock> ResourceControlFlow(const Program& program, co
 				return {};
 			}
 			block.successors.push_back(found->second);
+			predecessors[found->second].push_back(i);
 		}
 		for (const auto& inst: *program.blocks[i]) {
 			const auto op     = inst.GetOpcode();
@@ -707,7 +709,30 @@ 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);
+	// Host branches only matter when they guard a later resource access. Leave
+	// data-only branches (for example a mutable flag checked before a trap) on the GPU.
+	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: predecessors[index]) {
+			if (resource_reachable[predecessor]) continue;
+			resource_reachable[predecessor] = 1;
+			pending.push_back(predecessor);
+		}
+	}
+	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

@@ -739,15 +739,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];

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

@@ -23,9 +23,6 @@ namespace {
 constexpr uint64_t AddressMask = 0x0000ffffffffffffull;
 
 bool AddSignedAddress(uint64_t base, int64_t offset, uint64_t& result) {
-	if (base > AddressMask) {
-		return false;
-	}
 	if (offset < 0) {
 		const auto magnitude = uint64_t {0} - static_cast<uint64_t>(offset);
 		if (magnitude > base) {
@@ -35,7 +32,7 @@ bool AddSignedAddress(uint64_t base, int64_t offset, uint64_t& result) {
 		return true;
 	}
 	const auto magnitude = static_cast<uint64_t>(offset);
-	if (magnitude > AddressMask - base) {
+	if (magnitude > UINT64_MAX - base) {
 		return false;
 	}
 	result = base + magnitude;
@@ -443,7 +440,7 @@ bool SrtWalker::EvaluateRawRead(const Inst& inst, uint64_t& result) {
 	if (!Arg(*handle, 0, low) || !Arg(*handle, 1, high) || !Arg(inst, 1, offset)) {
 		return false;
 	}
-	const auto base      = ((high << 32u) | static_cast<uint32_t>(low)) & AddressMask;
+	const auto base      = (high << 32u) | static_cast<uint32_t>(low);
 	const auto immediate = static_cast<int64_t>(static_cast<int32_t>(mem.offset));
 	uint64_t   address   = 0;
 	if (inst.GetOpcode() == ValueOpcode::ReadConstBuffer) {
@@ -464,7 +461,7 @@ bool SrtWalker::EvaluateRawRead(const Inst& inst, uint64_t& result) {
 		if (byte_offset > size || size - byte_offset < sizeof(uint32_t)) {
 			return false;
 		}
-		address = (base & ~uint64_t {3}) + byte_offset;
+		address = (base & AddressMask & ~uint64_t {3}) + byte_offset;
 	} else {
 		const auto relative = (immediate & ~int64_t {3}) +
 		                      static_cast<int64_t>(static_cast<uint32_t>(offset) & ~3u);

tests/ResourceMaterializationTests.cpp

@@ -30,7 +30,8 @@ AddValueBlock(Libs::Graphics::ShaderRecompiler::IR::Program &program) {
   return *result;
 }
 
-Libs::Graphics::ShaderRecompiler::IR::ResourcePlan SrtPlan(uint64_t address) {
+Libs::Graphics::ShaderRecompiler::IR::ResourcePlan SrtPlan(uint64_t address,
+                                                       bool buffer = false) {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   Program program;
   program.stage = Libs::Graphics::ShaderType::Compute;
@@ -39,16 +40,22 @@ Libs::Graphics::ShaderRecompiler::IR::ResourcePlan SrtPlan(uint64_t address) {
   auto &value_block = AddValueBlock(program);
 
   MemoryInfo memory;
-  memory.kind = ResourceKind::ScalarAddress;
+  memory.kind = buffer ? ResourceKind::ScalarBuffer : ResourceKind::ScalarAddress;
   memory.planning_only = true;
   program.memory_info.push_back(memory);
   const auto low = Value(static_cast<uint32_t>(address));
   const auto high = Value(static_cast<uint32_t>(address >> 32u));
-  auto &handle =
+  auto &handle = buffer
-      value_block.AppendNewInst(ValueOpcode::GetAddressResource, {low, high});
+                     ? value_block.AppendNewInst(ValueOpcode::GetBufferResource,
-  auto &raw = value_block.AppendNewInst(
+                                                {low, high, Value(4u), Value(0u)})
-      ValueOpcode::LoadAddressU32,
+                     : value_block.AppendNewInst(ValueOpcode::GetAddressResource,
-      {Value(&handle), Value(0u), Value(0u), Value(true)});
+                                                {low, high});
+  auto &raw = buffer
+                  ? value_block.AppendNewInst(ValueOpcode::ReadConstBuffer,
+                                             {Value(&handle), Value(0u)})
+                  : value_block.AppendNewInst(
+                        ValueOpcode::LoadAddressU32,
+                        {Value(&handle), Value(0u), Value(0u), Value(true)});
   raw.SetFlags(MemoryFlags{.index = 0, .pc = 0x40});
   program.srt_reads.push_back({Value(&raw), 0});
 
@@ -163,6 +170,32 @@ void TestMappedSrtUsesDirectReaderByDefault() {
         "cache rematerialization did not use the direct reader by default");
 }
 
+void TestRawScalarAddressPreservesHighBits() {
+  using namespace Libs::Graphics::ShaderRecompiler::IR;
+  constexpr uint64_t address = 0xabcd000156780003ull;
+  for (const bool buffer : {false, true}) {
+    auto plan = SrtPlan(address, buffer);
+    uint64_t observed = 0;
+    const SrtRuntime runtime{
+        .read_memory = [](void *userdata, uint64_t requested,
+                          std::span<uint32_t> words) {
+          if (words.size() != 1) return false;
+          *static_cast<uint64_t *>(userdata) = requested;
+          words[0] = 0x12345678;
+          return true;
+        },
+        .userdata = &observed};
+    ResourceSnapshot snapshot;
+    ResourceSpecialization specialization;
+    Check(MaterializeResources(plan, runtime, snapshot, specialization),
+          "scalar address materialization failed");
+    Check(observed == (buffer ? 0x0000000156780000ull : address & ~uint64_t{3}),
+          "raw pointer high bits and buffer descriptor flags were conflated");
+    Check(snapshot.flattened_srt == std::vector<uint32_t>{0x12345678},
+          "scalar address read did not reach its flattened slot");
+  }
+}
+
 void TestIntegerRuntimeValueFollowsSrtReads() {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   auto plan = SrtPlan(0x10000);
@@ -454,6 +487,7 @@ void DbgExit(int) { std::abort(); }
 
 int main() {
   TestMappedSrtUsesDirectReaderByDefault();
+  TestRawScalarAddressPreservesHighBits();
   TestIntegerRuntimeValueFollowsSrtReads();
   TestUnbasedFlatCacheHitMaterializes();
   TestWrittenDescriptorUsesStrictReaderOnce();

tests/ResourceTrackingTests.cpp

@@ -2927,6 +2927,79 @@ ResourcePlan ConditionalBufferPlan(ConditionalBufferUse use) {
   return ExtractResourcePlan(fixture.program);
 }
 
+void TestNativeScalarAtomicDataDoesNotEnterResourcePlan() {
+  namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
+  Fixture fixture;
+  auto *entry = fixture.block;
+  for (uint32_t i = 1; i < 5; ++i) fixture.AddBlock();
+  fixture.program.block_info[0].terminator = {
+      .kind = CFG::TerminatorKind::ConditionalBranch, .true_block = 1, .false_block = 4};
+  for (uint32_t i = 1; i < 4; ++i) {
+    fixture.program.block_info[i].terminator = {
+        .kind = CFG::TerminatorKind::Branch, .true_block = i + 1};
+  }
+  fixture.program.block_info[4].terminator.kind = CFG::TerminatorKind::Return;
+
+  const auto table = fixture.Address(fixture.UserData(0), fixture.UserData(1));
+  MemoryInfo scalar;
+  scalar.kind = ResourceKind::ScalarAddress;
+  scalar.component_count = 2;
+  scalar.offset = 136;
+  const auto low = fixture.Emit(ValueOpcode::LoadAddressU32,
+      {table, Value(0u), Value(0u), Value(true)}, fixture.AddMemory(scalar, 0x560));
+  scalar.offset = 140;
+  scalar.component_index = 1;
+  const auto high = fixture.Emit(ValueOpcode::LoadAddressU32,
+      {table, Value(0u), Value(0u), Value(true)}, fixture.AddMemory(scalar, 0x560));
+  const auto address = fixture.Address(low, high);
+  scalar.offset = 24;
+  scalar.component_count = 1;
+  scalar.component_index = 0;
+  const auto data_flags = fixture.AddMemory(scalar, 0x584);
+  const auto data = fixture.Emit(ValueOpcode::LoadAddressU32,
+      {address, Value(0u), Value(0u), Value(true)}, data_flags);
+  const auto descriptor = fixture.Buffer({low,
+      fixture.Emit(ValueOpcode::BitwiseOr32, {high, Value(256u << 16u)}),
+      Value(1u), Value(90628u)}, 0x58c);
+  MemoryInfo vector;
+  vector.kind = ResourceKind::Buffer;
+  vector.offset = 24;
+  vector.data_dwords = 2;
+  fixture.Emit(ValueOpcode::BufferAtomicOr64,
+      {descriptor, Value(0u), Value(0u), Value(0u), Value(uint64_t{1}), Value(true)},
+      fixture.AddMemory(vector, 0x58c));
+  const auto flag = fixture.Emit(ValueOpcode::BitwiseAnd32, {data, Value(16u)});
+  fixture.program.block_info[0].condition =
+      fixture.Emit(ValueOpcode::INotEqual32, {flag, Value(0u)});
+  fixture.PlanAndTrack();
+  Check(data.Instruction()->Parent() == entry &&
+            !fixture.program.memory_info[data_flags.index].planning_only &&
+            fixture.program.info.buffers.size() == 1 && fixture.program.info.buffers[0].atomic,
+        "mutable scalar data was replaced by a host snapshot before its native atomic write");
+  auto plan = ExtractResourcePlan(fixture.program);
+  Check(plan.srt_reads.size() == 2 && plan.control_flow.empty() &&
+            !plan.capture_specialization_reads,
+        "resource-free successor chain retained the shader's scalar data predicate");
+  struct Reads { uint32_t descriptors = 0; uint32_t data = 0; } reads;
+  const auto Read = +[](void *userdata, uint64_t address, std::span<uint32_t> words) {
+    auto &reads = *static_cast<Reads *>(userdata);
+    if (address == 0x2018u) { ++reads.data; return false; }
+    if (words.size() != 1 || (address != 0x1088u && address != 0x108cu)) return false;
+    ++reads.descriptors;
+    words[0] = address == 0x1088u ? 0x2000u : 0u;
+    return true;
+  };
+  const std::array<uint32_t, 2> user_data{0x1000u, 0u};
+  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.data == 0 && snapshot.specialization_reads.empty() &&
+            snapshot.buffers[0].dwords[0] == 0x2000u,
+        "resource materialization read mutable scalar data or lost its descriptor pointer");
+}
+
 void TestConditionalBufferMaterialization() {
   auto plan = ConditionalBufferPlan(ConditionalBufferUse::Optional);
   // GTA III leaves packet words in s[12:15] when its scalar control word is zero.
@@ -3483,6 +3556,7 @@ int main() {
     Run("DMA address materialization", TestDmaAddressMaterialization);
     Run("dynamic FLAT address", TestDynamicFlatAddressesUseDma);
     Run("buffer swizzle specialization", TestBufferSwizzleSpecialization);
+    Run("native scalar atomic data", TestNativeScalarAtomicDataDoesNotEnterResourcePlan);
     Run("conditional buffer materialization", TestConditionalBufferMaterialization);
     Run("guarded scalar descriptor reads", TestGuardedScalarDescriptorReads);
     Run("conservative buffer reachability", TestConservativeBufferReachability);

tests/ShaderRecompilerComputeTests.cpp

@@ -25861,15 +25861,18 @@ TestCase ScalarMemoryLoadVariants() {
   std::vector<u32> expected = initial;
   expected.insert(expected.end(), initial.begin(), initial.end());
 
-  return {"ScalarMemoryLoadVariants",
+  TestCase test{"ScalarMemoryLoadVariants",
-          code,
+                code,
-          initial,
+                initial,
-          expected,
+                expected,
-          {O::S_MOV_B32, O::S_LOAD_DWORD, O::S_LOAD_DWORDX2, O::S_LOAD_DWORDX4,
+                {O::S_MOV_B32, O::S_LOAD_DWORD, O::S_LOAD_DWORDX2, O::S_LOAD_DWORDX4,
-           O::S_LOAD_DWORDX8, O::S_LOAD_DWORDX16, O::S_BUFFER_LOAD_DWORD,
+                 O::S_LOAD_DWORDX8, O::S_LOAD_DWORDX16, O::S_BUFFER_LOAD_DWORD,
-           O::S_BUFFER_LOAD_DWORDX2, O::S_BUFFER_LOAD_DWORDX4,
+                 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::S_BUFFER_LOAD_DWORDX8, O::S_BUFFER_LOAD_DWORDX16, O::V_MOV_B32,
-           O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+                 O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+  test.bda_mappings = {{0, 0}};
+  test.required_spirv = {"get_bda_pointer"};
+  return test;
 }
 
 TestCase ScalarBufferOffsetAlignmentAndCarry() {
@@ -25913,11 +25916,15 @@ TestCase ScalarLoadSignedImmediateOffsetAddsSoffset() {
   AppendStoreSgpr(&code, 1, 0);
   AppendEnd(&code);
 
-  return {"ScalarLoadSignedImmediateOffsetAddsSoffset",
+  TestCase test{"ScalarLoadSignedImmediateOffsetAddsSoffset",
-          code,
+                code,
-          {0x11111111u, 0x22222222u},
+                {0x11111111u, 0x22222222u},
-          {0x22222222u, 0x22222222u},
+                {0x22222222u, 0x22222222u},
-          {O::S_MOV_B32, O::S_LOAD_DWORD, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+                {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"};
+  test.forbidden_spirv = {"flattened_srt"};
+  return test;
 }
 
 TestCase BufferLoadStore() {
@@ -28016,23 +28023,27 @@ 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);
   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,
+                code,
-          {0x11111111u, 0x22222222u},
+                {0x11111111u, 0x22222222u},
-          {0x11111111u, 0x22222222u},
+                {0x11111111u, 0x22222222u},
-          {O::S_MOV_B32, O::S_LOAD_DWORD, O::BUFFER_STORE_DWORD, O::S_ENDPGM}};
+                {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"};
+  test.forbidden_spirv = {"flattened_srt"};
+  return test;
 }
 
 TestCase ScalarLoadAlignsDynamicBase() {
@@ -33031,7 +33042,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(ScalarMemoryLoadVariants);
   AddCase(ScalarBufferOffsetAlignmentAndCarry);
   AddCase(ScalarLoadSignedImmediateOffsetAddsSoffset);
-  AddCase(ScalarLoadAlignsComponentsAndMasksAddress);
+  AddCase(ScalarLoadAlignsComponents);
   AddCase(ScalarLoadAlignsDynamicBase);
   cases.push_back(ScalarBufferFromLoopReadlane(32));
   cases.push_back(ScalarBufferFromLoopReadlane(64));
@@ -35909,17 +35920,19 @@ void CheckResourcePlanHandoff() {
     ValidateSpirv(name, compiled.spirv);
     const auto &bindings = compiled.program.bindings;
     Require(name, "GPU data requirements",
-            bindings.user_data_registers.empty() && bindings.memory_offset_count == 1 &&
+            bindings.user_data_registers == (numeric_read ? std::vector<u32>{0, 1}
-                (FindBinding(bindings, DescriptorBindingKind::FlattenedSrt) != nullptr) ==
+                                                          : std::vector<u32>{}) &&
-                    numeric_read,
+                bindings.memory_offset_count == 1 &&
-            "descriptor-only shader retained SRT uploads or live scalar data was removed");
+                FindBinding(bindings, DescriptorBindingKind::FlattenedSrt) == nullptr &&
+                (FindBinding(bindings, DescriptorBindingKind::BdaPagetable) != nullptr) == numeric_read,
+            "descriptor-only shader retained SRT uploads or native scalar data 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() == 4,
             "compiled shader cleanup invalidated or froze the host resource plan");
   }
   std::printf("[host]    %-32s ok\n", name);