From 232e2f25938edde573129de56a7b13eeeb467be0 Mon Sep 17 00:00:00 2001 From: nmzik Date: Thu, 15 Oct 2026 12:14:28 +0200 Subject: [PATCH] shader: bound waterfall buffer descriptors with same-lane guards Co-authored-by: fxpw --- .../recompiler/ir/passes/ResourceTracking.cpp | 30 +++-- tests/ResourceTrackingTests.cpp | 123 ++++++++++++++++++ tests/ShaderRecompilerComputeTests.cpp | 110 ++++++++++++++++ 3 files changed, 255 insertions(+), 8 deletions(-) diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp index b55a03fd2..8742d462f 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp +++ b/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 op = image.GetOpcode(); + const auto& access = *use.user; + const auto op = access.GetOpcode(); if (use.operand != 0u || (op != ValueOpcode::ImageSampleRaw && - op != ValueOpcode::ImageGatherRaw && op != ValueOpcode::ImageRead)) return {}; - // The image's own EXEC predicate proves the range for every accessing lane; + op != ValueOpcode::ImageGatherRaw && op != ValueOpcode::ImageRead && + !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; diff --git a/tests/ResourceTrackingTests.cpp b/tests/ResourceTrackingTests.cpp index cf1f28185..91ace92c7 100644 --- a/tests/ResourceTrackingTests.cpp +++ b/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(); + 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{local, key} + : std::initializer_list{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 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 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); diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index d0987ebed..73d1e16e5 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/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 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(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);