shader: materialize exact buffer capacity through shared clean reads

Co-authored-by: shadowbeat070 <70843316+shadowbeat070@users.noreply.github.com>

tests/ResourceMaterializationTests.cpp

@@ -1,8 +1,10 @@
 #include "graphics/shader/recompiler/ir/ShaderIR.h"
 #include "graphics/shader/recompiler/ir/passes/ResourceMaterialization.h"
 
+#include <bit>
 #include <cstdio>
 #include <cstdlib>
+#include <limits>
 #include <memory>
 
 namespace {
@@ -196,45 +198,6 @@ void TestRawScalarAddressPreservesHighBits() {
   }
 }
 
-void TestDispatchBufferKeepsOnlyBindingLayoutOnHost() {
-  using namespace Libs::Graphics::ShaderRecompiler::IR;
-  auto plan = UserDataBufferPlan();
-  plan.info.buffers[0].dispatch_stride = 64;
-  plan.info.buffers[0].max_byte_extent = 20;
-  auto &source = plan.descriptor_sources[0];
-  source.dwords[1] = Value(20u << 16u);
-  source.dwords[3] = Value(0x16204u);
-  std::array<uint32_t, 1> user_data{0x1000u};
-  uint32_t reads = 0;
-  const SrtRuntime runtime{.user_data = user_data,
-                           .read_memory = RejectSpecializationRead,
-                           .userdata = &reads,
-                           .read_specialization_memory = RejectSpecializationRead};
-  ResourceSnapshot snapshot;
-  ResourceSpecialization specialization;
-  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
-            snapshot.buffers[0].dwords[2] == 0 &&
-            specialization.buffers[0].packed_stride == 20 && reads == 0,
-        "native record count was required to materialize its stable binding layout");
-  for (const auto words : {std::array<uint32_t, 2>{0u, 0x16204u},
-                           {16u << 16u, 0x16204u},
-                           {(20u << 16u) | (1u << 31u), 0x16204u},
-                           {20u << 16u, 0x16204u | (1u << 23u)},
-                           {20u << 16u, 0x16204u | (1u << 28u)}}) {
-    source.dwords[1] = Value(words[0]);
-    source.dwords[3] = Value(words[1]);
-    Check(!MaterializeResources(plan, runtime, snapshot, specialization),
-          "dispatch buffer accepted a layout outside its native bounds proof");
-  }
-  user_data[0] = 0;
-  source.dwords[1] = Value(0u);
-  source.dwords[3] = Value(0u);
-  Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
-            specialization.buffers[0].packed_stride == 0,
-        "inactive dispatch buffer did not preserve null binding semantics");
-  Check(reads == 0, "dispatch buffer layout validation read GPU data on the host");
-}
-
 void TestIntegerRuntimeValueFollowsSrtReads() {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   auto plan = SrtPlan(0x10000);
@@ -284,6 +247,113 @@ void TestIntegerRuntimeValueFollowsSrtReads() {
         "cyclic SRT read-first-lane dependency was accepted");
 }
 
+void TestUniformVectorDescriptorRead() {
+  using namespace Libs::Graphics::ShaderRecompiler::IR;
+  Program program;
+  program.stage = Libs::Graphics::ShaderType::Compute;
+  program.srt_plan_complete = program.resource_tracking_complete = true;
+  auto &block = AddValueBlock(program);
+  auto &handle = block.AppendNewInst(ValueOpcode::GetBufferResource,
+      {Value(0x1000u), Value(4u << 16u), Value(1u), Value(0x16204u)});
+  program.memory_info.push_back({.kind = ResourceKind::Buffer});
+  auto &count = block.AppendNewInst(ValueOpcode::LoadBufferU32,
+      {Value(&handle), Value(0u), Value(0u), Value(0u), Value(true)});
+  count.SetFlags(MemoryFlags{.index = 0});
+  // The captured indirect kernel shares one read across sibling scalar lane reads,
+  // enclosed by a different EXEC mask. Its resource plan must share that read too.
+  auto &lane = block.AppendNewInst(ValueOpcode::GetBuiltin,
+      {Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId)), Value(0u)});
+  auto &active = block.AppendNewInst(ValueOpcode::ULessThan32, {Value(&lane), Value(32u)});
+  auto &first = block.AppendNewInst(ValueOpcode::ReadFirstLane, {Value(&count), Value(true)});
+  auto &next = block.AppendNewInst(ValueOpcode::IAdd32, {Value(&count), Value(1u)});
+  auto &second = block.AppendNewInst(ValueOpcode::ReadFirstLane, {Value(&next), Value(true)});
+  auto &sum = block.AppendNewInst(ValueOpcode::IAdd32, {Value(&first), Value(&second)});
+  auto &small = block.AppendNewInst(ValueOpcode::ULessThanEqual32, {Value(&count), Value(72u)});
+  auto &selected = block.AppendNewInst(ValueOpcode::SelectU32, {Value(&small), Value(&sum), Value(&count)});
+  auto &masked = block.AppendNewInst(ValueOpcode::SelectU32, {Value(&active), Value(&selected), Value(0u)});
+  auto &records = block.AppendNewInst(ValueOpcode::ReadFirstLane, {Value(&masked), Value(&active)});
+  DescriptorSource source;
+  source.dword_count = 4;
+  source.dwords = {Value(0x2000u), Value(4u << 16u), Value(&records), Value(0x16204u)};
+  program.descriptor_sources.push_back(source);
+  program.info.buffers.push_back({.source = 0, .written = true});
+  Check(ValidateRuntimeValue(program, Value(&count)), "uniform DWORD count was rejected");
+  struct Reads { uint32_t value = 72; uint32_t strict = 0; uint32_t ordinary = 0; bool clean = true; } reads;
+  const SrtRuntime runtime{
+      .read_memory = [](void *data, uint64_t, std::span<uint32_t> words) {
+        ++static_cast<Reads *>(data)->ordinary;
+        words[0] = 999;
+        return true;
+      },
+      .userdata = &reads,
+      .read_specialization_memory = [](void *data, uint64_t address, std::span<uint32_t> words) {
+        auto &reads = *static_cast<Reads *>(data);
+        ++reads.strict;
+        if (!reads.clean || address != 0x1000u || words.size() != 1) return false;
+        words[0] = reads.value;
+        return true;
+      }};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  for (const bool written : {false, true}) {
+    program.info.buffers[0].written = written;
+    auto plan = ExtractResourcePlan(program);
+    Check(plan.control_flow.empty() && plan.requires_specialization_memory &&
+              plan.capture_specialization_reads,
+          "vector descriptor read depended on incidental control-flow capture");
+    reads = {};
+    Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
+              snapshot.buffers[0].dwords[2] == 145 && reads.strict == 1 && reads.ordinary == 0 &&
+              snapshot.specialization_reads ==
+                  std::vector<std::pair<uint64_t, uint64_t>>{{0x1000u, 4u}},
+          "vector descriptor input was not read and captured exactly once");
+    reads.clean = false;
+    reads.strict = 0;
+    Check(!MaterializeResources(plan, runtime, snapshot, specialization) &&
+              reads.strict == 1 && reads.ordinary == 0,
+          "dirty vector descriptor input fell back to an ordinary memory read");
+  }
+  count.SetArg(4, Value(false));
+  auto &inactive = block.AppendNewInst(ValueOpcode::ReadFirstLane, {Value(&count), Value(false)});
+  reads.strict = 0;
+  uint32_t result = 99;
+  Check(SrtWalker(program, runtime).Evaluate(Value(&inactive), result) &&
+            result == 0 && reads.strict == 0 && reads.ordinary == 0,
+        "literal false EXEC read vector memory");
+  count.SetArg(4, Value(true));
+  handle.SetArg(3, Value(0x204u));
+  Check(SrtWalker(program, runtime).Evaluate(Value(&count), result) &&
+            result == 0 && reads.strict == 0 && reads.ordinary == 0,
+        "invalid vector buffer format read memory");
+  handle.SetArg(3, Value(0x16204u));
+  count.SetArg(1, Value(&lane));
+  Check(!ValidateRuntimeValue(program, Value(&count)),
+        "varying vector address was treated as a uniform descriptor read");
+}
+
+void TestExactReciprocalDescriptorArithmetic() {
+  using namespace Libs::Graphics::ShaderRecompiler::IR;
+  Program program;
+  auto &block = AddValueBlock(program);
+  for (const float divisor : {64.f, 128.f, 256.f, 512.f,
+                              std::numeric_limits<float>::min(),
+                              std::bit_cast<float>(253u << 23u)}) {
+    auto &reciprocal = block.AppendNewInst(ValueOpcode::FPRecipIFlag32, {Value::F32(divisor)});
+    uint32_t result = 0;
+    Check(SrtWalker(program, {}).Evaluate(Value(&reciprocal), result) &&
+              result == std::bit_cast<uint32_t>(1.f / divisor),
+          "power-of-two reciprocal was not exact");
+  }
+  for (const float divisor : {0.f, 3.f, -64.f, std::numeric_limits<float>::infinity(),
+                              std::numeric_limits<float>::denorm_min(),
+                              std::bit_cast<float>(254u << 23u)}) {
+    auto &reciprocal = block.AppendNewInst(ValueOpcode::FPRecipIFlag32, {Value::F32(divisor)});
+    uint32_t result = 0;
+    Check(!SrtWalker(program, {}).Evaluate(Value(&reciprocal), result),
+          "unsupported reciprocal rounding or exceptional input was accepted");
+  }
+}
+
 void TestUnbasedFlatCacheHitMaterializes() {
   using namespace Libs::Graphics::ShaderRecompiler::IR;
   auto plan = UnbasedFlatPlan();
@@ -527,8 +597,9 @@ void DbgExit(int) { std::abort(); }
 int main() {
   TestMappedSrtUsesDirectReaderByDefault();
   TestRawScalarAddressPreservesHighBits();
-  TestDispatchBufferKeepsOnlyBindingLayoutOnHost();
   TestIntegerRuntimeValueFollowsSrtReads();
+  TestUniformVectorDescriptorRead();
+  TestExactReciprocalDescriptorArithmetic();
   TestUnbasedFlatCacheHitMaterializes();
   TestWrittenDescriptorUsesStrictReaderOnce();
   TestFailedMaterializationRejectsStage();

tests/ResourceTrackingTests.cpp

@@ -2728,6 +2728,33 @@ void TestInvariantLoopPhi() {
         "loop-invariant descriptor phi was not evaluated through typed SSA");
 }
 
+void TestBufferStoreUsesItsOwnActiveValue() {
+  Fixture fixture;
+  const auto lane = fixture.Emit(ValueOpcode::GetBuiltin,
+      {Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId)), Value(0u)});
+  const auto guard = fixture.Emit(ValueOpcode::ULessThan32, {lane, Value(16u)});
+  const auto other_guard = fixture.Emit(ValueOpcode::ULessThan32, {lane, Value(32u)});
+  const auto inner = fixture.Emit(ValueOpcode::SelectU32, {guard, Value(9u), lane});
+  const auto data = fixture.Emit(ValueOpcode::SelectU32, {guard, inner, lane});
+  const auto handle = fixture.Buffer({Value(0x1000u), Value(4u << 16u),
+                                       Value(64u), Value(0x16204u)});
+  MemoryInfo memory;
+  memory.kind = ResourceKind::Buffer;
+  const auto store = [&](Value active) {
+    return fixture.Emit(ValueOpcode::StoreBufferU32,
+        {handle, Value(0u), Value(0u), Value(0u), data, active},
+        fixture.AddMemory(memory, 0));
+  };
+  const auto matching = store(guard);
+  const auto different = store(other_guard);
+  ConstantPropagationPass(fixture.program.blocks);
+  Check(matching.Instruction()->Arg(4).Resolve() == Value(9u),
+        "buffer store retained data from its inactive lanes");
+  Check(different.Instruction()->Arg(4).Resolve() == data.Resolve() &&
+            data.Resolve().Instruction()->GetOpcode() == ValueOpcode::SelectU32,
+        "buffer store changed a shared value outside its own EXEC mask");
+}
+
 void TestBoundedRelativeRegisterWrites() {
   namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
   enum class Variant { Bounded, BoundClobber, DescriptorClobber, EntryBypass };
@@ -3639,6 +3666,7 @@ int main() {
     Run("finite image bit scan sentinel", TestFiniteImageBitScanSentinel);
     Run("runtime-rooted loop", TestLoopCycleEnteredThroughRuntimeValue);
     Run("invariant loop phi", TestInvariantLoopPhi);
+    Run("buffer store active value", TestBufferStoreUsesItsOwnActiveValue);
     Run("bounded relative register writes", TestBoundedRelativeRegisterWrites);
     Run("DMA address materialization", TestDmaAddressMaterialization);
     Run("dynamic FLAT address", TestDynamicFlatAddressesUseDma);

tests/ShaderRecompilerComputeTests.cpp

@@ -31282,11 +31282,11 @@ void CheckComputeThreadDimensions(VulkanHarness &vulkan) {
   }
 }
 
-void CheckGpuComputedBufferRecords(VulkanHarness &vulkan) {
+void CheckRuntimeBufferRecords(VulkanHarness &vulkan) {
   using namespace ShaderRecompiler::IR;
   TestCase test;
-  test.name = "GpuComputedBufferRecords";
-  // Captured b71ef7c579835a0d: read a GPU counter, divide it into chunks,
+  test.name = "RuntimeBufferRecords";
+  // Captured b71ef7c579835a0d: read an input count, divide it into chunks,
   // construct a structured descriptor, then clear each live record's last DWORD.
   test.code = {
     0xbfa00003u, 0xd7460000u, 0x04010c0eu, 0x8805ff0du, 0x00040000u, 0xbe84030cu,
@@ -31303,17 +31303,17 @@ void CheckGpuComputedBufferRecords(VulkanHarness &vulkan) {
     0xe0700000u, 0x80000500u, 0xbefe046au, 0x7e0a0280u, 0x7e000502u, 0x7e020501u,
     0x7e040504u, 0x7e060503u, 0xe0702010u, 0x80000500u, 0xbf810000u,
   };
-  constexpr u32 guest_base = 0x1000;
-  test.user_data[0] = guest_base + 16;
+  constexpr u32 count_address = 4;
+  test.user_data[0] = 16;
   test.user_data[1] = 4u << 16;
   test.user_data[2] = 1;
   test.user_data[3] = 0x5204;
-  test.user_data[8] = guest_base + 20;
-  test.user_data[12] = guest_base;
+  test.user_data[8] = 20;
+  test.user_data[12] = count_address;
   test.has_user_data = true;
-  test.storage_buffer_offsets = {0, 16, 20};
+  test.storage_buffer_offsets = {count_address, 16, 20};
   test.initial.assign(5 + 5 * 130, 0xdeadbeefu);
-  test.initial[0] = 72;
+  test.initial[1] = 72;
   test.has_compute_info = true;
   test.compute_info.threads_num[0] = 64;
   test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1;
@@ -31326,23 +31326,62 @@ void CheckGpuComputedBufferRecords(VulkanHarness &vulkan) {
   test.compute_info.dispatch_thread_dimensions = true;
   test.compute_info.dispatch_threads_num[0] = 2;
   test.compute_info.dispatch_threads_num[1] = test.compute_info.dispatch_threads_num[2] = 1;
-  test.required_spirv = {"OpFDiv", "OpUMulExtended"};
-  test.forbidden_spirv = {"get_bda_pointer", "flattened_srt"};
+  test.forbidden_spirv = {"get_bda_pointer", "flattened_srt", "OpFDiv", "OpUMulExtended"};
   const auto compiled = CompileCase(test, vulkan.SubgroupSize());
-  Require(test.name, "native record count",
+  Require(test.name, "runtime record count",
           compiled.program.info.buffers.size() == 3 &&
-              compiled.program.info.buffers[2].dispatch_stride == 64 &&
-              compiled.resources.buffers[2].dwords[2] == 0 &&
-              compiled.resources.specialization_reads.empty(),
-          "GPU count entered the host descriptor snapshot or lost its dispatch bound");
-  // Reuse the same compiled shader as both GPU data and dispatch coverage change.
+              compiled.resources.buffers[2].dwords[2] == 2 &&
+              compiled.resources.specialization_reads ==
+                  std::vector<std::pair<uint64_t, uint64_t>>{{count_address, 4}},
+          "descriptor capacity lost its exact clean input read");
+  ShaderRecompiler::CompileOptions options;
+  options.stage = ShaderType::Compute;
+  options.wave_size = 64;
+  options.user_data = test.user_data;
+  options.input_info.compute = &test.compute_info;
+  auto translated = ShaderRecompiler::TranslateProgram(test.code, options);
+  auto plan = ExtractResourcePlan(translated.program);
+  struct CountMemory { u32 value = 72; u32 reads = 0; bool clean = true; } memory;
+  SrtRuntime runtime{.user_data = test.user_data, .userdata = &memory};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Require(test.name, "strict reader required",
+          !MaterializeResources(plan, runtime, snapshot, specialization),
+          "vector descriptor input bypassed the clean-memory reader");
+  runtime.read_specialization_memory = [](void *data, uint64_t address, std::span<u32> words) {
+    auto &memory = *static_cast<CountMemory *>(data);
+    ++memory.reads;
+    if (!memory.clean || address != count_address || words.size() != 1) return false;
+    words[0] = memory.value;
+    return true;
+  };
+  memory.clean = false;
+  Require(test.name, "dirty input rejected",
+          !MaterializeResources(plan, runtime, snapshot, specialization),
+          "GPU-dirty count was accepted for host descriptor evaluation");
+  memory.clean = true;
+  ResourceSpecialization initial_specialization;
+  // Reuse the same resource plan and compiled shader as the count and coverage change.
   for (const u32 count : {0u, 72u, 512u, 8192u, 262144u}) {
+    memory.value = count;
+    memory.reads = 0;
+    const u32 chunk = count < 512 ? 64 : count < 8192 ? 128 : count < 262144 ? 256 : 512;
+    const u32 records = (count + chunk - 1) / chunk;
+    Require(test.name, "cached descriptor refresh",
+            MaterializeResources(plan, runtime, snapshot, specialization) && memory.reads == 1 &&
+                snapshot.buffers[2].dwords[2] == records &&
+                snapshot.specialization_reads ==
+                    std::vector<std::pair<uint64_t, uint64_t>>{{count_address, 4}},
+            "count refresh repeated a memory read or retained an old descriptor capacity");
+    if (count == 0) initial_specialization = specialization;
+    Require(test.name, "capacity is not specialization",
+            specialization == initial_specialization,
+            "changing record capacity would recompile the cached shader");
+    test.initial.assign(5 + 5 * records, 0xdeadbeefu);
     for (const u32 threads : {2u, 3u, 65u}) {
-      test.initial[0] = count;
+      test.initial[1] = count;
       test.compute_info.dispatch_threads_num[0] = threads;
       test.dispatch_x = (threads + 63) / 64;
-      const u32 chunk = count < 512 ? 64 : count < 8192 ? 128 : count < 262144 ? 256 : 512;
-      const u32 records = (count + chunk - 1) / chunk;
       test.expected = test.initial;
       test.expected[4] = 0;
       for (u32 index = 0; index < std::min(threads, records); ++index) {
@@ -31352,7 +31391,7 @@ void CheckGpuComputedBufferRecords(VulkanHarness &vulkan) {
       vulkan.Dispatch(test, compiled, output);
       const auto actual = vulkan.ReadBuffer(test.name, output, test.expected.size());
       vulkan.DestroyBuffer(&output);
-      CompareWords(test, "GPU records and padded invocation bounds", test.expected, actual);
+      CompareWords(test, "record capacity and padded invocation bounds", test.expected, actual);
     }
   }
   std::printf("[compute] %-32s ok\n", test.name);
@@ -38454,7 +38493,7 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ScalarBufferFromLoopReadlane(32));
     RunCase(&vulkan, ScalarBufferFromLoopReadlane(64));
     CheckIndirectBufferStore(vulkan);
-    CheckGpuComputedBufferRecords(vulkan);
+    CheckRuntimeBufferRecords(vulkan);
     RunCase(&vulkan, BufferLoadsGpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadDwordGpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadDwordx3GpuSelectedDescriptors());
@@ -39030,7 +39069,7 @@ int main(int argc, char **argv) {
   CheckPixelParameterAliases();
   CheckRectListShaders();
   CheckIndirectBufferStore(vulkan);
-  CheckGpuComputedBufferRecords(vulkan);
+  CheckRuntimeBufferRecords(vulkan);
   CheckComputeThreadDimensions(vulkan);
   CheckIndirectImageKeySwitch();
   CheckWave64WholeWaveResults();