shader: bind bounded GPU-selected buffer store destinations

Co-authored-by: Ali Almohaya <AliAlmohaya@gmail.com>

tests/ShaderRecompilerComputeTests.cpp

@@ -9202,7 +9202,7 @@ public:
               rendering.width == width && rendering.height == height &&
               rendering.num_layers == view_layers &&
               rendering.num_color_attachments == 1,
-          "SDK dimension=0 slice range did not create a Vulkan 1D-array "
+          "dimension=0 slice range did not create a Vulkan 1D-array "
           "attachment view");
 
       RenderExecutorTestAccess::ResetBindings(executor);
@@ -30569,6 +30569,117 @@ TestCase ImageD16StoreUnpacksHalfPairs() {
   return test;
 }
 
+void CheckIndirectBufferStore(VulkanHarness &vulkan) {
+  using namespace ShaderRecompiler::IR;
+  TestCase test;
+  test.name = "IndirectBufferStore";
+  test.initial.assign(64, 0xdeadbeefu);
+  test.storage_buffer_range_dwords = 16;
+  test.storage_buffer_offsets = {12, 108};
+  test.required_spirv = {"OpSwitch"};
+  test.forbidden_spirv = {"PhysicalStorageBuffer"};
+
+  CompiledShader compiled;
+  auto &program = compiled.program;
+  program.stage = ShaderType::Compute;
+  program.wave_size = 32;
+  program.srt_plan_complete = true;
+  program.resource_tracking_complete = true;
+  program.shader_info_complete = true;
+  program.block_storage.push_back(std::make_unique<Block>());
+  auto *block = program.block_storage.back().get();
+  program.blocks.push_back(block);
+  program.block_info.push_back({.id = 0});
+
+  auto &lane = block->AppendNewInst(ValueOpcode::LaneId);
+  auto &selector = block->AppendNewInst(ValueOpcode::BitwiseAnd32,
+                                      {Value(&lane), Value(7u)});
+  auto &key = block->AppendNewInst(ValueOpcode::IAdd32,
+                                 {Value(&selector), Value(8u)});
+  auto &handle = block->AppendNewInst(
+      ValueOpcode::GetBufferResource,
+      {Value(&key), Value(0u), Value(0u), Value(0u)});
+  handle.SetFlags<uint32_t>(0u);
+  auto &offset = block->AppendNewInst(ValueOpcode::IMul32,
+                                    {Value(&lane), Value(4u)});
+  auto &data = block->AppendNewInst(ValueOpcode::IAdd32,
+                                  {Value(&lane), Value(0x1000u)});
+  auto &exec = block->AppendNewInst(ValueOpcode::INotEqual32,
+                                  {Value(&lane), Value(4u)});
+  block->AppendNewInst(ValueOpcode::StoreBufferU32,
+                      {Value(&handle), Value(0u), Value(&offset), Value(0u),
+                       Value(&data), Value(&exec)});
+  program.memory_info.push_back({.kind = ResourceKind::Buffer, .offen = true});
+
+  BufferResource root;
+  root.packed_stride = 1;
+  root.written = true;
+  root.indirect_root = 0;
+  root.indirect_mapping_offset = 4;
+  root.indirect_search_iterations = std::bit_width(3u);
+  root.indirect_resources = {0u, 1u};
+  auto candidate = root;
+  candidate.indirect_search_iterations = 0;
+  candidate.indirect_resources.clear();
+  program.descriptor_sources.resize(1);
+  program.descriptor_sources[0].indirect_descriptor.emplace();
+  ShaderComputeInputInfo compute;
+  compute.wave_size = 32;
+  compute.host_subgroup_size = vulkan.SubgroupSize();
+  compute.threads_num[0] = 32;
+  compute.threads_num[1] = compute.threads_num[2] = 1;
+  for (const u32 null_buffers : {0u, 1u, 2u}) {
+    program.info.buffers = {root, candidate};
+    if (null_buffers != 0u) program.info.buffers[1].packed_stride = 0;
+    if (null_buffers == 2u) {
+      // Materialization deduplicates all null destinations into one unmapped root.
+      program.info.buffers.resize(1);
+      program.info.buffers[0].packed_stride = 0;
+      program.info.buffers[0].indirect_root = BufferResource::NoIndirectBuffer;
+      program.info.buffers[0].indirect_search_iterations = 0;
+      program.info.buffers[0].indirect_resources.clear();
+    }
+    program.binding_layout_complete = false;
+    AllocateBindings(program);
+    const auto *buffers = FindBinding(program.bindings, DescriptorBindingKind::Buffers);
+    Require(test.name, "candidate bindings",
+            buffers != nullptr && buffers->resources ==
+                (null_buffers == 2u ? std::vector<u32>{0u} : std::vector<u32>{0u, 1u}),
+            "indirect store lost a buffer binding");
+    compiled.packed_user_data.resize(program.bindings.ShaderDataDwords());
+    compiled.packed_user_data[program.bindings.memory_offset_dword] = 12u | (108u << 8u);
+    compiled.spirv = ShaderRecompiler::Spirv::EmitProgram(program, {.compute = &compute});
+    ValidateSpirv(test.name, compiled.spirv);
+    if (null_buffers == 2u) {
+      test.required_spirv.clear();
+      test.forbidden_spirv.push_back("OpSwitch");
+    }
+    CheckSpirvText(test, compiled.spirv);
+
+    for (const bool swapped : {false, true}) {
+      if (null_buffers != 0u && swapped) continue;
+      // Repeated keys select two distinct logical buffers within one native allocation.
+      // Keys below, between and above the live keys must not select either buffer.
+      compiled.resources.flattened_srt = {0u, 0u, 0u, 0u, 3u,
+                                         10u, swapped ? 0u : 1u,
+                                         12u, swapped ? 1u : 0u,
+                                         14u, swapped ? 0u : 1u};
+      test.expected = test.initial;
+      for (const u32 lane : {2u, 6u, 10u, 12u, 14u}) {
+        const u32 resource = (lane & 7u) == 4u ? swapped : !swapped;
+        if (null_buffers == 2u || (null_buffers == 1u && resource == 1u)) continue;
+        test.expected[test.storage_buffer_offsets[resource] / 4u + lane] = 0x1000u + lane;
+      }
+      auto output = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
+      vulkan.Dispatch(test, compiled, output);
+      const auto actual = vulkan.ReadBuffer(test.name, output, test.expected.size());
+      vulkan.DestroyBuffer(&output);
+      CompareWords(test, "key selection, EXEC, bounds and null descriptors", test.expected, actual);
+    }
+  }
+  std::printf("[compute] %-32s ok\n", test.name);
+}
+
 void CheckIndirectImageKeySwitch() {
   constexpr const char *name = "IndirectImageKeySwitch";
   constexpr uint32_t mapping_capacity = 1793u;
@@ -30621,8 +30732,8 @@ void CheckIndirectImageKeySwitch() {
 
   program.descriptor_sources.resize(2);
   program.descriptor_sources[0].dword_count = 8;
-  program.descriptor_sources[0].indirect_image =
-      DescriptorSource::IndirectImage{0u, 0u, 224u, 12u, 0u};
+  program.descriptor_sources[0].indirect_descriptor =
+      DescriptorSource::IndirectDescriptor{0u, 0u, 224u, 12u, 0u};
   program.descriptor_sources[1].dword_count = 4;
 
   ImageResource root{};
@@ -30656,7 +30767,7 @@ void CheckIndirectImageKeySwitch() {
           text.find("OpSwitch") != std::string::npos &&
               text.find("OpPhi") != std::string::npos &&
               CountText(text, "OpImageSampleExplicitLod") == 2 &&
-              CountText(text, "OpIEqual") == 11,
+              CountText(text, "OpIEqual") == 1,
           "dynamic image key did not use a compact two-sample switch");
 
   program.memory_info[0].image_dimension =
@@ -37292,6 +37403,7 @@ int main(int argc, char **argv) {
   }
   if (argc == 2 && std::strcmp(argv[1], "--indirect-buffer-only") == 0) {
     VulkanHarness vulkan;
+    CheckIndirectBufferStore(vulkan);
     RunCase(&vulkan, BufferLoadsGpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadDwordx3GpuSelectedDescriptors());
     RunCase(&vulkan, BufferLoadDwordx4SnapshotsOverlappingAddress());
@@ -37857,6 +37969,7 @@ int main(int argc, char **argv) {
   CheckTessellationPrograms();
   CheckPixelParameterAliases();
   CheckRectListShaders();
+  CheckIndirectBufferStore(vulkan);
   CheckIndirectImageKeySwitch();
   CheckWave64WholeWaveResults();
   CheckPs5GameExampleImageClearRuntimeShape();