shader: prove finite masked selectors for writable buffer tables

Co-authored-by: fxpw <fxpw@outlook.com>

tests/ResourceTrackingTests.cpp

@@ -4393,6 +4393,80 @@ void TestBoundedFormattedBufferStore() {
                "finite store accepted an unproved induction or indexed byte offset");
 }
 
+void TestFiniteBufferSelection() {
+  enum class Variant { Valid, WiderExec, Unbounded, Cycle, TooMany };
+  const auto make = [](Variant variant) {
+    auto fixture = std::make_unique<Fixture>();
+    auto& f = *fixture;
+    auto* entry = f.block;
+    auto* left = f.AddBlock();
+    auto* right = f.AddBlock();
+    auto* merge = f.AddBlock();
+    entry->AddBranch(left);
+    (variant == Variant::Cycle ? merge : entry)->AddBranch(right);
+    left->AddBranch(merge);
+    right->AddBranch(merge);
+    const auto table = f.Address(f.UserData(0u), f.UserData(1u));
+    const auto lane = f.Emit(ValueOpcode::LaneId);
+    const auto active = f.Emit(ValueOpcode::ULessThan32, {lane, Value(16u)});
+    auto selected = Value(0u);
+    for (uint32_t key = 1; key < 5u; ++key) {
+      const auto match = f.Emit(ValueOpcode::IEqual32, {lane, Value(key)}, 0, right);
+      selected = f.Emit(ValueOpcode::SelectU32, {match, Value(key), selected}, 0, right);
+    }
+    auto& phi = merge->AppendNewInst(ValueOpcode::Phi, {}, uint64_t(Type::U32));
+    phi.AddPhiOperand(left, Value(0u));
+    phi.AddPhiOperand(right, variant == Variant::Cycle
+        ? f.Emit(ValueOpcode::IAdd32, {Value(&phi), Value(1u)}, 0, right)
+        : variant == Variant::TooMany ? Value(MaxIndirectDescriptorProbes) : selected);
+    f.block = merge;
+    const auto stale = f.Emit(ValueOpcode::UndefU32);
+    const auto masked = f.Emit(ValueOpcode::SelectU32,
+        {active, variant == Variant::Unbounded ? lane : Value(&phi), stale});
+    const auto key = f.Emit(ValueOpcode::ReadFirstLane, {masked, active});
+    const auto offset = f.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(4u)});
+    std::array<Value, 4> words;
+    for (uint32_t word = 0; word < words.size(); ++word)
+      words[word] = f.Emit(ValueOpcode::LoadAddressU32,
+          {table, offset, Value(0u), Value(true)},
+          f.AddMemory({.kind = ResourceKind::ScalarAddress, .offset = 32u + word * 4u,
+                       .component_index = word, .component_count = 4u}, 0x50cu));
+    f.Emit(ValueOpcode::StoreBufferU32,
+        {f.Buffer(words), lane, Value(0u), Value(0u), Value(42u),
+         variant == Variant::WiderExec ? Value(true) : active},
+        f.AddMemory({.kind = ResourceKind::Buffer, .idxen = true, .coherent = true}, 0x530u));
+    return fixture;
+  };
+  auto fixture = make(Variant::Valid);
+  fixture->PlanAndTrack();
+  const auto plan = ExtractResourcePlan(fixture->program);
+  const auto& source = plan.descriptor_sources[plan.info.buffers[0].source];
+  Check(source.indirect_descriptor && source.indirect_descriptor->key_count == Value(5u),
+        "masked finite selection did not retain all five descriptor candidates");
+  EliminateDeadCode(fixture->program.blocks);
+  LinearTestMemory memory;
+  for (uint32_t key = 0; key < 5u; ++key) {
+    const auto word = 8u + key * 4u;
+    memory.words[word] = 0x100000u + key * 0x10000u;
+    memory.words[word + 1u] = 4u << 16u;
+    memory.words[word + 2u] = 64u;
+    memory.words[word + 3u] = 0x5204u;
+  }
+  memory.fail_address = memory.base + 32u + 5u * 16u;
+  const std::array<uint32_t, 2> user_data{static_cast<uint32_t>(memory.base), 0u};
+  SrtRuntime runtime{.user_data = user_data, .read_memory = ReadLinearTestMemory,
+                     .userdata = &memory, .read_specialization_memory = ReadLinearTestMemory};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) && snapshot.buffers.size() == 5u,
+        "finite buffer selection lost table roots after DCE or probed outside its domain");
+  for (const auto variant : {Variant::WiderExec, Variant::Unbounded, Variant::Cycle, Variant::TooMany}) {
+    auto invalid = make(variant);
+    CheckFatal([&] { invalid->PlanAndTrack(); }, "not a valid runtime value",
+               "finite buffer selection accepted an unproved EXEC or selector domain");
+  }
+}
+
 void TestPackedScalarBufferSelection() {
   const auto make = [](uint32_t shift_mask = 24u, bool cycle = false, uint32_t store_offset = 0u,
                        bool byte_store = false) {
@@ -5814,6 +5888,7 @@ int main() {
     Run("dynamic SRT", TestDynamicSrtReadRemainsExplicit);
     Run("writable descriptor phi", TestWritableDescriptorPhi);
     Run("bounded formatted buffer store", TestBoundedFormattedBufferStore);
+    Run("finite buffer selection", TestFiniteBufferSelection);
     Run("packed scalar buffer selection", TestPackedScalarBufferSelection);
     Run("conditional sampler phi", TestConditionalSamplerPhi);
     Run("finite image phi cycle", TestFiniteImagePhiCycle);

tests/ShaderRecompilerComputeTests.cpp

@@ -36368,6 +36368,70 @@ void CheckBoundedFormattedBufferCopy(VulkanHarness &vulkan) {
   std::printf("[compute] %-32s ok\n", "BoundedFormattedBufferCopy");
 }
 
+void CheckFiniteBufferSelection(VulkanHarness &vulkan) {
+  for (const u32 wave : {32u, 64u}) {
+    TestCase test;
+    test.name = wave == 32 ? "FiniteBufferSelection32" : "FiniteBufferSelection64";
+    const u32 first_lane = wave == 32 ? 1u : 33u;
+    // Water Lighting's masked V_CNDMASK -> READFIRSTLANE -> S_LOAD_DWORDX4 -> store.
+    test.code = {EncodeSMovB32(126, InlineU32(wave == 32 ? 6u : 0u)),
+                 EncodeSMovB32(127, InlineU32(wave == 64 ? 6u : 0u)),
+                 EncodeMubuf0(0x0cu, 0, true, false), EncodeMubuf1(20, 2, 0),
+                 EncodeSopp(0x0cu, 0)};
+    AppendVMovU32(&test.code, 4, 0);
+    test.code.push_back(EncodeSopc(0x0au, 2, InlineU32(1)));
+    const auto skip = test.code.size();
+    test.code.push_back(0);
+    for (u32 key = 1; key < 5u; ++key) {
+      AppendVMovU32(&test.code, 22, key);
+      test.code.push_back(EncodeVopc(0xc2u, InlineU32(key), 20));
+      test.code.push_back(EncodeVop2(0x01u, 4, Vgpr(4), 22));
+    }
+    test.code[skip] = EncodeSopp(0x05u, static_cast<u32>(test.code.size() - skip - 1u));
+    test.code.insert(test.code.end(), {
+        EncodeVop1(0x02u, 3, Vgpr(4)), EncodeSop2(0x31u, 12, 3, InlineU32(32)),
+        EncodeSmem0(0x02u, 4), EncodeSmem1(0, 12), EncodeSopp(0x0cu, 0),
+        EncodeVop2(0x25u, 1, 0xffu, 0), 0x1000u,
+        EncodeMubuf0(0x1cu, 0, true, false, true), EncodeMubuf1(1, 1, 0)});
+    AppendEnd(&test.code);
+    constexpr u32 table = 256u, input = 512u, output = 1024u, slot = 512u;
+    test.initial.assign((output + 5u * slot) / 4u, 0xdeadbeefu);
+    test.has_user_data = true;
+    test.user_data[0] = table;
+    test.user_data[2] = 1u;
+    const std::array<u32, 4> input_descriptor{input, 4u << 16u, wave, 0x5204u};
+    std::copy(input_descriptor.begin(), input_descriptor.end(), test.user_data.begin() + 8u);
+    for (u32 key = 0; key < 5u; ++key) {
+      const std::array<u32, 4> descriptor{output + key * slot, 4u << 16u, wave, 0x5204u};
+      std::copy(descriptor.begin(), descriptor.end(), test.initial.begin() + table / 4u + 8u + key * 4u);
+    }
+    test.has_compute_info = true;
+    test.compute_info.wave_size = wave;
+    test.compute_info.threads_num[0] = wave;
+    test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1;
+    test.compute_info.lds_size_dwords = 0;
+    test.compute_info.thread_ids_num = 1;
+    const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+    Require(test.name, "finite destination count", compiled.resources.buffers.size() == 6u,
+            "selector lost a destination or retained a dynamic descriptor load");
+    for (const auto& descriptor : compiled.resources.buffers)
+      test.storage_buffer_ranges.emplace_back(descriptor.dwords[0], descriptor.dwords[2] * 4u);
+    for (u32 key = 0; key < 5u; ++key) {
+      for (u32 lane = 0; lane < wave; ++lane)
+        test.initial[input / 4u + lane] = lane == first_lane ? key : (key + 1u) % 5u;
+      test.expected = test.initial;
+      for (const auto lane : {first_lane, first_lane + 1u})
+        test.expected[(output + key * slot) / 4u + lane] = 0x1000u + lane;
+      auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
+      vulkan.Dispatch(test, compiled, buffer);
+      const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size());
+      vulkan.DestroyBuffer(&buffer);
+      CompareWords(test, "first active lane, masked writes and destination canaries", test.expected, actual);
+    }
+    std::printf("[compute] %-32s ok\n", test.name);
+  }
+}
+
 void CheckIndirectBufferStore(VulkanHarness &vulkan) {
   using namespace ShaderRecompiler::IR;
   TestCase test;
@@ -45083,6 +45147,7 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ScalarBufferFromLoopReadlane(32));
     RunCase(&vulkan, ScalarBufferFromLoopReadlane(64));
     CheckBoundedFormattedBufferCopy(vulkan);
+    CheckFiniteBufferSelection(vulkan);
     CheckIndirectBufferStore(vulkan);
     CheckRuntimeBufferRecords(vulkan);
     RunCase(&vulkan, BufferLoadsGpuSelectedDescriptors());
@@ -45781,6 +45846,7 @@ int main(int argc, char **argv) {
   CheckRectListShaders();
   CheckBoundedFormattedBufferCopy(vulkan);
   CheckIndirectBufferStore(vulkan);
+  CheckFiniteBufferSelection(vulkan);
   CheckRuntimeBufferRecords(vulkan);
   CheckComputeThreadDimensions(vulkan);
   CheckCountedRowImageSelector(vulkan);