shader: bound waterfall buffer descriptors with same-lane guards

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

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

@@ -1956,18 +1956,25 @@ private:
 		});
 	}
 
-	Value UnsignedImageBound(Value key, const Inst& handle) const {
+	Value UnsignedResourceBound(Value key, const Inst& handle) const {
 		Value bound;
 		for (const auto& use: handle.Uses()) {
-			const auto& image = *use.user;
+			const auto& access = *use.user;
-			const auto op = image.GetOpcode();
+			const auto op = access.GetOpcode();
 			if (use.operand != 0u || (op != ValueOpcode::ImageSampleRaw &&
-			    op != ValueOpcode::ImageGatherRaw && op != ValueOpcode::ImageRead)) return {};
+			    op != ValueOpcode::ImageGatherRaw && op != ValueOpcode::ImageRead &&
-			// The image's own EXEC predicate proves the range for every accessing lane;
+			    !IndirectBufferAccess(access))) return {};
+			// The access's own EXEC predicate proves the range for every accessing lane;
 			// an EXEC-nonzero branch alone only proves that some lane is active.
-			const auto active = image.Arg(image.NumArgs() - 1u);
+			const auto active = access.Arg(access.NumArgs() - 1u);
 			auto candidate = UnsignedBound(active, key);
-			if (candidate.IsEmpty()) candidate = UnsignedBoundOnEntry(key, image.Parent(), false, active);
+			if (candidate.IsEmpty()) candidate = UnsignedBoundOnEntry(key, access.Parent(), false, active);
+			if (candidate.IsEmpty()) {
+				// A waterfall key inherits the local value's bound only where this same
+				// EXEC predicate proves both their equality and the local range.
+				const auto local = EqualLocalKey(active, key);
+				if (!local.IsEmpty()) candidate = UnsignedBound(active, local, true, false, active);
+			}
 			if (candidate.IsEmpty() ||
 			    !ValidateRuntimeValue(m_program, candidate, RuntimeValueType::Integer)) return {};
 			if (bound.IsEmpty()) bound = candidate;
@@ -2881,7 +2888,7 @@ private:
 			plan.retain_reads |= table_source.dword_count == 4u;
 		} else if (table_source.dword_count == 2u) {
 			auto guarded_count = UnsignedBoundOnEntry(key, handle.Parent());
-			if (guarded_count.IsEmpty()) guarded_count = UnsignedImageBound(key, handle);
+			if (guarded_count.IsEmpty()) guarded_count = UnsignedResourceBound(key, handle);
 			// A runtime count can describe fewer rows than the mask's static bit range.
 			if (guarded_count.IsEmpty()) guarded_count = UnsignedBoundOnEntry(key, handle.Parent(), true);
 			if (!guarded_count.IsEmpty()) {
@@ -3104,6 +3111,13 @@ private:
 			return finish();
 		}
 		if (!legacy_store) {
+			if (!plan.split_offsets) {
+				const auto count = UnsignedResourceBound(key, handle);
+				if (!count.IsEmpty()) {
+					indirect.key_count = count;
+					return finish();
+				}
+			}
 			if (table_source.dword_count == 2u && !plan.split_offsets) {
 				const auto maximum = SelectorMaximum(key);
 				if (maximum >= MaxIndirectDescriptorProbes) return false;

tests/ResourceTrackingTests.cpp

@@ -5381,6 +5381,128 @@ void TestBoundedBufferRecords(uint32_t d32_components) {
   }
 }
 
+void TestGuardedWaterfallBufferRecords() {
+  namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
+  namespace Prospero = Libs::Graphics::Prospero;
+  using Libs::Graphics::DstSel;
+  enum class Guard { Exact, Reversed, NoRange, NoEquality, OtherLocal, UnguardedUse, SeparateVotes };
+  const auto make = [](Guard guard) {
+    auto f = std::make_unique<Fixture>();
+    const auto count = f->UserData(2u);
+    const auto table = f->Buffer({f->UserData(0u), f->UserData(1u), count, f->UserData(3u)});
+    const auto lane = f->Emit(ValueOpcode::LaneId);
+    const auto input = f->Buffer({f->UserData(4u), f->UserData(5u), f->UserData(6u), f->UserData(7u)});
+    const auto local = f->Emit(ValueOpcode::LoadBufferU32,
+        {input, lane, Value(0u), Value(0u), Value(true)},
+        f->AddMemory({.kind = ResourceKind::Buffer, .idxen = true}, 0x200u));
+    const auto key = f->Emit(ValueOpcode::ReadLane, {local, Value(0u)});
+    const auto initial = f->Emit(ValueOpcode::ULessThan32, {lane, Value(32u)});
+    const auto range_local = guard == Guard::OtherLocal
+        ? f->Emit(ValueOpcode::IAdd32, {local, Value(1u)}) : local;
+    // Match the captured ANDN1_SAVEEXEC predicate, including its initial-EXEC witness.
+    const auto excluded = f->Emit(ValueOpcode::LogicalAnd,
+        {initial, f->Emit(ValueOpcode::ULessThanEqual32, {count, range_local})});
+    const auto range = f->Emit(ValueOpcode::LogicalAnd,
+        {initial, f->Emit(ValueOpcode::LogicalNot, {excluded})});
+    const auto equal = f->Emit(ValueOpcode::IEqual32,
+        guard == Guard::Reversed ? std::initializer_list<Value>{local, key}
+                                 : std::initializer_list<Value>{key, local});
+    auto active = f->Emit(ValueOpcode::LogicalAnd, {range, equal});
+    if (guard == Guard::NoRange) active = equal;
+    if (guard == Guard::NoEquality) active = range;
+    if (guard == Guard::SeparateVotes)
+      active = f->Emit(ValueOpcode::LogicalAnd,
+          {f->Emit(ValueOpcode::ConditionRef, {range}, CFG::BranchCondition::ExecNonZero),
+           f->Emit(ValueOpcode::ConditionRef, {equal}, CFG::BranchCondition::ExecNonZero)});
+    const auto offset = f->Emit(ValueOpcode::IMul32, {key, Value(488u)});
+    std::array<Value, 4> words;
+    for (uint32_t word = 0; word < 4u; ++word)
+      words[word] = f->Emit(ValueOpcode::ReadConstBuffer, {table, offset},
+          f->AddMemory({.kind = ResourceKind::ScalarBuffer, .offset = 216u + word * 4u,
+                        .component_index = word, .component_count = 4u}, 0x304u));
+    const auto handle = f->Buffer(words);
+    for (uint32_t use = 0; use < 2u; ++use) {
+      const auto loaded = f->Emit(ValueOpcode::LoadBufferU32x3,
+          {handle, lane, Value(0u), Value(0u),
+           guard == Guard::UnguardedUse && use == 1u ? Value(true) : active},
+          f->AddMemory({.kind = ResourceKind::Buffer, .data_dwords = 3u, .data_bits = 32u,
+                        .component_count = 3u, .formatted = true, .idxen = true}, 0x360u));
+      f->Emit(ValueOpcode::ReferenceU32,
+          {f->Emit(ValueOpcode::CompositeExtractU32x3, {loaded, Value(use)})});
+    }
+    f->PlanAndTrack();
+    return f;
+  };
+  LinearTestMemory memory;
+  constexpr uint32_t table_offset = 256u;
+  for (uint32_t row = 0; row < 2u; ++row) {
+    const std::array descriptor{0x10000u + row * 0x1000u, 12u << 16u, 32u,
+        (uint32_t(Prospero::BufferFormat::k32_32_32Float) << 12u) | DstSel(4, 5, 6, 1)};
+    std::copy(descriptor.begin(), descriptor.end(),
+              memory.words.begin() + (table_offset + row * 488u + 216u) / 4u);
+  }
+  // An unbounded key 0x029f79b5 wraps to this real congruence slot, outside key < N.
+  constexpr std::array poison{0x012b6890u, 0x501502f9u, 0x493ba30eu, 0x3ea2f352u};
+  std::copy(poison.begin(), poison.end(), memory.words.begin() + (table_offset + 480u) / 4u);
+  std::array<uint32_t, 8> user_data{uint32_t(memory.base) + table_offset, 488u << 16u,
+      2u, 0x5204u, 0x9000u, 4u << 16u, 32u, 0x5204u};
+  SrtRuntime runtime{.user_data = user_data, .read_memory = ReadLinearTestMemory,
+      .userdata = &memory, .read_specialization_memory = ReadLinearTestMemory};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  for (const auto guard : {Guard::Exact, Guard::Reversed, Guard::NoRange, Guard::NoEquality,
+                          Guard::OtherLocal, Guard::UnguardedUse, Guard::SeparateVotes}) {
+    auto fixture = make(guard);
+    const auto plan = ExtractResourcePlan(fixture->program);
+    const auto root = std::ranges::find_if(plan.info.buffers, [&](const auto& buffer) {
+      return plan.descriptor_sources[buffer.source].indirect_descriptor.has_value();
+    });
+    Check(root != plan.info.buffers.end(), "waterfall buffer lost its indirect resource");
+    const bool proven = guard == Guard::Exact || guard == Guard::Reversed;
+    const auto& domain = *plan.descriptor_sources[root->source].indirect_descriptor;
+    Check(domain.key_count.IsEmpty() != proven && !domain.selector &&
+              domain.table_stride == 488u && domain.table_immediate == 216u,
+          "waterfall bound mixed different lanes, values or resource uses");
+    const auto materialize = [&] { return MaterializeResources(plan, runtime, snapshot, specialization); };
+    if (!proven) {
+      Check(!materialize(), "unproved waterfall skipped a reachable unsupported descriptor");
+      continue;
+    }
+    for (const uint32_t count : {2u, 0u, 1u, 2u}) {
+      user_data[2] = count;
+      memory.reads = 0;
+      Check(materialize() && memory.reads == count &&
+                snapshot.buffers.size() == 1u + std::max(count, 1u),
+            "waterfall count refresh read GPU keys or unrelated record fields");
+      for (uint32_t row = 0; row < count; ++row)
+        Check(std::ranges::find(snapshot.specialization_reads,
+                  std::pair{memory.base + table_offset + row * 488u + 216u, uint64_t{16u}}) !=
+                  snapshot.specialization_reads.end(), "waterfall descriptor read was not captured");
+    }
+    // SMEM reads the complete DWORD prefix and zeroes the missing format word.
+    user_data[1] = 359u << 16u;
+    memory.reads = 0;
+    Check(materialize() && memory.reads == 2u &&
+              std::ranges::find(snapshot.specialization_reads,
+                  std::pair{memory.base + table_offset + 704u, uint64_t{12u}}) !=
+                  snapshot.specialization_reads.end() &&
+              std::ranges::find(snapshot.specialization_reads,
+                  std::pair{memory.base + table_offset + 704u, uint64_t{16u}}) ==
+                  snapshot.specialization_reads.end() &&
+              std::ranges::all_of(snapshot.buffers.back().dwords, [](auto word) { return word == 0u; }),
+          "waterfall partial descriptor overread or retained its missing format word");
+    user_data[1] = 488u << 16u;
+    memory.fail_address = memory.base + table_offset + 704u;
+    Check(!materialize(), "waterfall descriptor bypassed a coherent read failure");
+    memory.fail_address = UINT64_MAX;
+    Check(materialize(), "waterfall descriptor did not recover after a failed refresh");
+    user_data[2] = MaxIndirectDescriptorProbes + 1u;
+    memory.reads = 0;
+    Check(!materialize() && memory.reads == 0u, "waterfall accepted an excessive runtime bound");
+    user_data[2] = 2u;
+  }
+}
+
 void TestGpuSelectedBufferRecords() {
   namespace Prospero = Libs::Graphics::Prospero;
   using Libs::Graphics::DstSel;
@@ -7677,6 +7799,7 @@ int main() {
     Run("bounded formatted buffer store", TestBoundedFormattedBufferStore);
     Run("bounded buffer records", [] {
       for (const uint32_t components : {0u, 2u, 3u, 4u}) TestBoundedBufferRecords(components);
+      TestGuardedWaterfallBufferRecords();
     });
     Run("GPU-selected buffer records", TestGpuSelectedBufferRecords);
     Run("workgroup-selected buffer records", TestWorkgroupSelectedBufferRecords);

tests/ShaderRecompilerComputeTests.cpp

@@ -37982,6 +37982,114 @@ void CheckGpuSelectedRecordBufferAccess(VulkanHarness &vulkan) {
   std::printf("[compute] %-32s ok\n", "GpuSelectedRecordBufferAccess");
 }
 
+void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
+  const u32 alignment = vulkan.RuntimeContext().GetPhysicalDeviceProperties()
+                            .limits.minStorageBufferOffsetAlignment;
+  const u32 slot = (4096u + alignment - 1u) / alignment * alignment;
+  constexpr u32 selectors = 256u, table = 512u;
+  for (const auto [wave, upper_only] : {std::pair{32u, false}, std::pair{64u, false},
+                                      std::pair{64u, true}}) {
+    TestCase test;
+    test.name = upper_only ? "GuardedWaterfallRecordsUpper64" : wave == 32u
+        ? "GuardedWaterfallRecords32" : "GuardedWaterfallRecords64";
+    test.initial.assign((slot * 3u + wave * 16u + 16u) / 4u, 0xdeadbeefu);
+    std::fill_n(test.initial.begin() + table / 4u, 976u / 4u, 0u);
+    constexpr std::array poison{0x012b6890u, 0x501502f9u, 0x493ba30eu, 0x3ea2f352u};
+    std::copy(poison.begin(), poison.end(), test.initial.begin() + (table + 480u) / 4u);
+    test.has_user_data = test.has_compute_info = true;
+    const std::array<u32, 8> descriptors{table, 488u << 16u, 2u, 0x5204u,
+                                        selectors, 4u << 16u, wave, 0x5204u};
+    std::copy(descriptors.begin(), descriptors.end(), test.user_data.begin());
+    const std::array output{slot * 3u, 16u << 16u, wave, 0x5204u};
+    std::copy(output.begin(), output.end(), test.user_data.begin() + 16u);
+    for (u32 row = 0; row < 2u; ++row) {
+      const std::array descriptor{slot * (row + 1u), 12u << 16u, wave,
+          (u32(Prospero::BufferFormat::k32_32_32Float) << 12u) | DstSel(4, 5, 6, 1)};
+      std::copy(descriptor.begin(), descriptor.end(),
+                test.initial.begin() + (table + row * 488u + 216u) / 4u);
+      for (u32 lane = 0; lane < wave; ++lane)
+        for (u32 component = 0; component < 3u; ++component)
+          test.initial[slot * (row + 1u) / 4u + lane * 3u + component] =
+              0x3f800000u + (row << 20u) + (lane << 12u) + (component << 8u);
+    }
+    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] = 1u;
+    test.compute_info.thread_ids_num = 1;
+    test.compute_info.lds_size_dwords = 0;
+    auto &code = test.code;
+    code.push_back(EncodeVop1(0x01, 30, Vgpr(0)));
+    if (upper_only) {
+      AppendSMovLiteral(&code, 126, 0u);
+      AppendSMovLiteral(&code, 127, UINT32_MAX);
+    }
+    code.insert(code.end(), {EncodeSop1(0x04, 20, 126),
+        EncodeMubuf0(0x0c, 0, true, false), EncodeMubuf1(5, 1, 30), EncodeSopp(0x0c, 0),
+        EncodeVopc(0xc3, 2, 5), EncodeSop1(0x37, 52, 106)}); // Exclude N <= local key.
+    const auto empty = code.size();
+    code.push_back(0u);
+    code.push_back(EncodeSop1(0x04, 54, 126));
+    const auto loop = code.size();
+    code.push_back(EncodeSop1(0x14, 34, 54));
+    AppendVop3(&code, 0x360, 32, Vgpr(5), 34);
+    code.insert(code.end(), {EncodeVopc(0xc2, 32, 5), EncodeSop1(0x24, 58, 106)});
+    const auto skip = code.size();
+    code.push_back(0u);
+    code.insert(code.end(), {EncodeSop2(0x26, 33, 32, 255), 488u,
+        EncodeSmem0(0x0a, 12, 0), EncodeSmem1(216u, 33), EncodeSopp(0x0c, 0),
+        EncodeMubuf0(0x02, 0, true, false, true), EncodeMubuf1(6, 3, 30),
+        EncodeSopp(0x0c, 0), EncodeMubuf0(0x1f, 0, true, false), EncodeMubuf1(6, 4, 30)});
+    code[skip] = EncodeSopp(0x08, code.size() - skip - 1u);
+    code.insert(code.end(), {EncodeSop2(0x15, 54, 54, 106), EncodeSop1(0x04, 126, 58)});
+    code.push_back(EncodeSopp(0x05, static_cast<int16_t>(loop - code.size() - 1u)));
+    code[empty] = EncodeSopp(0x08, code.size() - empty - 1u);
+    code.push_back(EncodeSop1(0x04, 126, 20));
+    AppendVMovLiteral(&code, 9, 0x76543210u);
+    code.insert(code.end(), {EncodeMubuf0(0x1c, 12u, true, false), EncodeMubuf1(9, 4, 30)});
+    AppendEnd(&code);
+    test.required_spirv = {"OpLoopMerge", "OpGroupNonUniformShuffle", "Volatile"};
+    test.forbidden_spirv = {"load_indirect_buffer_", "indirect_buffer_format_geometry"};
+    const auto compiled = CompileCase(test, vulkan.SubgroupSize());
+    for (const auto &descriptor : compiled.resources.buffers) {
+      ShaderBufferResource resource{};
+      std::copy_n(descriptor.dwords.begin(), 4u, resource.fields);
+      test.storage_buffer_ranges.emplace_back(resource.Base48(), std::max(4u, u32(resource.GetSize())));
+    }
+    TestCase producer;
+    producer.name = "WaterfallRecordKeyProducer";
+    producer.initial = test.initial;
+    producer.has_compute_info = true;
+    producer.compute_info.wave_size = 32u;
+    producer.compute_info.threads_num[0] = producer.compute_info.threads_num[1] =
+        producer.compute_info.threads_num[2] = 1u;
+    producer.compute_info.lds_size_dwords = 0;
+    test.expected = test.initial;
+    for (u32 lane = 0; lane < wave; ++lane) {
+      // This excluded key would wrap to the unsupported descriptor at byte480.
+      const u32 key = lane % 3u == 2u ? 0x029f79b5u : lane % 3u;
+      AppendVMovLiteral(&producer.code, 1, key);
+      AppendVMovLiteral(&producer.code, 2, selectors + lane * 4u);
+      AppendBufferStoreDword(&producer.code, 1, 2);
+      test.expected[selectors / 4u + lane] = key;
+      if (upper_only && lane < 32u) continue;
+      const u32 out = slot * 3u / 4u + lane * 4u;
+      if (key < 2u)
+        std::copy_n(test.initial.begin() + slot * (key + 1u) / 4u + lane * 3u, 3u,
+                    test.expected.begin() + out);
+      test.expected[out + 3u] = 0x76543210u;
+    }
+    AppendEnd(&producer.code);
+    const auto producer_shader = CompileCase(producer, vulkan.SubgroupSize());
+    auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
+    vulkan.Dispatch(producer, producer_shader, buffer);
+    vulkan.Dispatch(test, compiled, buffer);
+    CompareWords(test, "same-lane bound, GPU keys, wrapped poison, EXEC restoration and canaries",
+                 test.expected, vulkan.ReadBuffer(test.name, buffer, test.expected.size()));
+    vulkan.DestroyBuffer(&buffer);
+  }
+  std::printf("[compute] %-32s ok\n", "GuardedWaterfallBufferRecords");
+}
+
 void CheckGpuSelectedFormattedRecordStore(VulkanHarness &vulkan) {
   constexpr u32 selectors = 128, table = 256, stride = 592, member = 448;
   const u32 alignment = vulkan.RuntimeContext().GetPhysicalDeviceProperties()
@@ -47491,6 +47599,7 @@ int main(int argc, char **argv) {
     CheckIndirectBufferStore(vulkan);
     CheckBoundedRecordBufferAccess(vulkan);
     CheckGpuSelectedRecordBufferAccess(vulkan);
+    CheckGuardedWaterfallBufferRecords(vulkan);
     CheckGpuSelectedFormattedRecordStore(vulkan);
     CheckRuntimeBufferRecords(vulkan);
     CheckSharedIndirectBufferLoads(vulkan);
@@ -48198,6 +48307,7 @@ int main(int argc, char **argv) {
   CheckIndirectBufferStore(vulkan);
   CheckBoundedRecordBufferAccess(vulkan);
   CheckGpuSelectedRecordBufferAccess(vulkan);
+  CheckGuardedWaterfallBufferRecords(vulkan);
   CheckGpuSelectedFormattedRecordStore(vulkan);
   CheckFiniteBufferSelection(vulkan);
   CheckRuntimeBufferRecords(vulkan);