diff --git a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp index 04052d48..cb6689be 100644 --- a/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp @@ -1481,7 +1481,8 @@ private: } const Inst* BoundedInduction(Value key, const Block* use, const auto& accepts_step, - const auto& accepts_bound) const { + const auto& accepts_bound, + LaneQuantifier lanes = LaneQuantifier::Any) const { const auto* phi = key.Resolve().TryInstruction(); if (phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi || phi->GetType() != Type::U32 || phi->NumArgs() != 2u || phi->NumPhiBlocks() != 2u) @@ -1523,7 +1524,8 @@ private: if (use_of_key.operand != 0u || !accepts_bound(*compare, initial_arm)) continue; if (!GuardedOnEntry(use != nullptr ? use : increment_block, stops, [&](const EdgePredicate& edge) { - return edge.positive && Implies(edge.condition, Value(use_of_key.user)); + return edge.positive && (lanes == LaneQuantifier::Any || edge.lanes == lanes) && + Implies(edge.condition, Value(use_of_key.user)); })) continue; LoopBoundProof proof(m_program, *phi, *compare); // The image bound and the increment guard are separate obligations: a @@ -1535,14 +1537,15 @@ private: return {}; } - const Inst* BoundedLoop(Value key, const Block* use, const auto& accepts_bound) const { + const Inst* BoundedLoop(Value key, const Block* use, const auto& accepts_bound, + LaneQuantifier lanes = LaneQuantifier::Any) const { return BoundedInduction(key, use, [](uint32_t initial, uint32_t step) { return initial == 0u && step == 1u; }, [&](const Inst& compare, uint32_t initial) { return (compare.GetOpcode() == ValueOpcode::SLessThan32 || compare.GetOpcode() == ValueOpcode::ULessThan32) && accepts_bound(compare, initial); - }); + }, lanes); } const Inst* DescendingLoop(Value key, const Block* use, uint32_t& maximum) const { @@ -1830,6 +1833,12 @@ private: uint32_t offset = 0, stride = 0; if (!MatchTableOffset(read.Arg(1), element, offset, stride) || offset != 0u || stride != 4u) return false; + const auto* lane = element.TryInstruction(); + const bool selected_lane = lane != nullptr && lane->GetOpcode() == ValueOpcode::ReadFirstLane; + if (selected_lane) { + if (!HasActiveLane(lane->Arg(1), lane->Parent())) return false; + element = ActiveValue(lane->Arg(0), lane->Arg(1)); + } const auto* sum = element.TryInstruction(); if (sum == nullptr || sum->GetOpcode() != ValueOpcode::IAdd32) return false; for (uint32_t arm = 0; arm < 2u; ++arm) { @@ -1845,20 +1854,28 @@ private: if (memory == nullptr || memory->offset != 0u || !MatchTableOffset(count->Arg(1), row, count_offset, count_stride) || count_offset != 0u || count_stride != 4u) return false; - const auto affine = AffineIndex<3>(read.Arg(1), [](Value value) { return value.Resolve(); }, + if (selected_lane) { + const auto* row_lane = row.TryInstruction(); + if (row_lane == nullptr || row_lane->GetOpcode() != ValueOpcode::ReadFirstLane || + !ValidateWaveUniformValue(m_program, row_lane->Arg(0))) return false; + // Different EXEC masks may select different lanes. A globally uniform row + // keeps their count and key addresses correlated; i is bounded on every lane. + row = row_lane->Arg(0).Resolve(); + } + const auto affine = AffineIndex<3>(element, [](Value value) { return value.Resolve(); }, [&](Value value) -> std::optional { if (value == induction) return 1u; if (EquivalentValue(m_program, value, row)) return 2u; return {}; }); - if (!affine || (*affine)[0] != 0u || (*affine)[1] != 4u || - (*affine)[2] == 0u || ((*affine)[2] & 3u) != 0u || + if (!affine || (*affine)[0] != 0u || (*affine)[1] != 1u || + (*affine)[2] == 0u || (*affine)[2] > UINT32_MAX / 4u || !MakeRuntimeTableSource(*count, count_source)) return false; // Count addresses identify rows modulo 2^30. A DWORD-aligned row // spacing preserves those aliases in the U32 selector address too. - row_stride = static_cast((*affine)[2]); + row_stride = static_cast((*affine)[2] * 4u); return true; - }) != nullptr) return true; + }, selected_lane ? LaneQuantifier::All : LaneQuantifier::Any) != nullptr) return true; } return false; } diff --git a/src/graphics/shader/recompiler/ir/passes/SrtWalker.cpp b/src/graphics/shader/recompiler/ir/passes/SrtWalker.cpp index c9c49c55..75f5c13b 100644 --- a/src/graphics/shader/recompiler/ir/passes/SrtWalker.cpp +++ b/src/graphics/shader/recompiler/ir/passes/SrtWalker.cpp @@ -178,8 +178,9 @@ bool IsRuntimeUniformOp(ValueOpcode op) { class RuntimeValidator { public: - explicit RuntimeValidator(const ResourcePlan& program, RuntimeValueType type) - : m_program(program), m_type(type) {} + explicit RuntimeValidator(const ResourcePlan& program, RuntimeValueType type, + bool wave_uniform = false) + : m_program(program), m_type(type), m_wave_uniform(wave_uniform) {} bool Run(Value value) { return Validate(value); } @@ -278,6 +279,12 @@ private: } return finish(true); } + if (m_wave_uniform && op == ValueOpcode::GetBuiltin) { + const auto axis = inst->NumArgs() == 2u ? inst->Arg(1).Resolve() : Value {}; + return finish(m_program.stage == ShaderType::Compute && inst->NumArgs() == 2u && + inst->Arg(0) == Value(static_cast(StageInputKind::WorkgroupId)) && + axis.IsImmediate() && axis.GetType() == Type::U32 && axis.U32() < 3u); + } if (op == ValueOpcode::Phi) { if (m_type != RuntimeValueType::Any && !ValidateArguments(*inst, false)) { return finish(false); @@ -357,6 +364,7 @@ private: const ResourcePlan& m_program; RuntimeValueType m_type; + bool m_wave_uniform; Value m_active_mask; std::unordered_set m_visiting; std::unordered_set m_validated_dependencies; @@ -1153,5 +1161,9 @@ bool ValidateRuntimeValue(const ResourcePlan& program, Value value, RuntimeValue return RuntimeValidator(program, type).Run(value); } +bool ValidateWaveUniformValue(const ResourcePlan& program, Value value) { + return RuntimeValidator(program, RuntimeValueType::Any, true).Run(value); +} + } // namespace Libs::Graphics::ShaderRecompiler::IR diff --git a/src/graphics/shader/recompiler/ir/passes/SrtWalker.h b/src/graphics/shader/recompiler/ir/passes/SrtWalker.h index 2c78b8ca..5aa8d70d 100644 --- a/src/graphics/shader/recompiler/ir/passes/SrtWalker.h +++ b/src/graphics/shader/recompiler/ir/passes/SrtWalker.h @@ -46,6 +46,9 @@ enum class RuntimeValueType { Any, Integer, ImmutableInteger }; bool ValidateRuntimeValue(const ResourcePlan& program, Value value, RuntimeValueType type = RuntimeValueType::Any); +// Uniform across a compute wave, including values derived from its workgroup ID. +// This does not imply that the host can evaluate the value. +bool ValidateWaveUniformValue(const ResourcePlan& program, Value value); void CompileSrtWalker(ResourcePlan& program); // Uses the strict reader for values that affect shader specialization. SrtRuntime CleanRuntime(SrtRuntime runtime); diff --git a/tests/ResourceTrackingTests.cpp b/tests/ResourceTrackingTests.cpp index d9bc0aaf..a6e4213b 100644 --- a/tests/ResourceTrackingTests.cpp +++ b/tests/ResourceTrackingTests.cpp @@ -2055,19 +2055,22 @@ void TestScalarBufferSelectorDomain() { void TestCountedRowImageSelector() { namespace CFG = Libs::Graphics::ShaderRecompiler::CFG; enum class Variant { Valid, Signed, WrongStep, Bypass, WrongRow, CountStride, - KeyStride, CountImmediate, KeyImmediate }; + KeyStride, CountImmediate, KeyImmediate, Masked, EmptyCountMask, VaryingRow, + EmptyMask, UnprovenMask, WrongMask, AnyBound, MaskedWrongStep }; const auto make_plan = [](Variant variant = Variant::Valid, uint32_t capacity = 16u) { Fixture fixture; + const bool masked = variant >= Variant::Masked; auto *entry = fixture.block; auto *header = fixture.AddBlock(); auto *body = fixture.AddBlock(); auto *latch = fixture.AddBlock(); auto *exit = fixture.AddBlock(); + auto *preheader = masked ? fixture.AddBlock() : entry; const auto branch = [](Block *from, Block *to) { from->AddBranch(to); from->terminator = {.kind = CFG::TerminatorKind::Branch, .true_block = to}; }; - branch(entry, header); + branch(preheader, header); branch(body, latch); branch(latch, header); const auto keys = fixture.Buffer({fixture.UserData(0), fixture.UserData(1), @@ -2076,15 +2079,45 @@ void TestCountedRowImageSelector() { fixture.UserData(6), fixture.UserData(7)}); const auto material = fixture.Buffer({fixture.UserData(8), fixture.UserData(9), fixture.UserData(10), fixture.UserData(11)}); - const auto row = fixture.Emit(ValueOpcode::GetBuiltin, + auto row = fixture.Emit(ValueOpcode::GetBuiltin, {Value(static_cast(StageInputKind::WorkgroupId)), Value(0u)}); + MemoryInfo scalar; + scalar.kind = ResourceKind::ScalarBuffer; + Value active(true); + if (masked) { + const auto local = fixture.Emit(ValueOpcode::GetBuiltin, + {Value(static_cast(StageInputKind::LocalInvocationId)), Value(0u)}); + row = fixture.Emit(ValueOpcode::ReadConstBuffer, + {keys, fixture.Emit(ValueOpcode::ShiftLeftLogical32, {row, Value(2u)})}, + fixture.AddMemory(scalar, 0x04)); + row = fixture.Emit(ValueOpcode::IMul32, + {row, fixture.Emit(ValueOpcode::ConvertU32F32, + {fixture.Emit(ValueOpcode::BitCastF32U32, {fixture.UserData(12u)})})}); + Check(ValidateWaveUniformValue(fixture.program, row) && + !ValidateRuntimeValue(fixture.program, row), + "wave-uniform workgroup data escaped the host evaluation boundary"); + if (variant == Variant::VaryingRow) + row = fixture.Emit(ValueOpcode::IAdd32, {row, local}); + active = fixture.Emit(ValueOpcode::INotEqual32, {local, Value(0u)}); + entry->condition = fixture.Emit(ValueOpcode::ConditionRef, {active}, + CFG::BranchCondition::ExecNonZero); + if (variant == Variant::UnprovenMask) { + entry->condition = fixture.Emit(ValueOpcode::INotEqual32, + {fixture.UserData(13u), Value(0u)}); + } + entry->AddBranch(preheader); + entry->AddBranch(exit); + entry->terminator = {.kind = CFG::TerminatorKind::ConditionalBranch, + .true_block = preheader, .false_block = exit}; + if (variant == Variant::EmptyMask) active = Value(false); + } auto count_row = row; + if (masked) count_row = fixture.Emit(ValueOpcode::ReadFirstLane, + {row, Value(variant != Variant::EmptyCountMask)}); if (variant == Variant::WrongRow) count_row = fixture.Emit(ValueOpcode::IAdd32, {row, Value(1u)}); const auto count_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32, {count_row, Value(variant == Variant::CountStride ? 3u : 2u)}); - MemoryInfo scalar; - scalar.kind = ResourceKind::ScalarBuffer; scalar.offset = variant == Variant::CountImmediate ? 4u : 0u; const auto count = fixture.Emit(ValueOpcode::ReadConstBuffer, {counts, count_offset}, fixture.AddMemory(scalar, 0x10)); @@ -2095,7 +2128,9 @@ void TestCountedRowImageSelector() { header->AddBranch(body); auto *outside = variant == Variant::Bypass ? latch : exit; header->AddBranch(outside); - header->condition = guard; + header->condition = variant == Variant::AnyBound + ? fixture.Emit(ValueOpcode::ConditionRef, {guard}, CFG::BranchCondition::ExecNonZero, header) + : guard; header->terminator = {.kind = CFG::TerminatorKind::ConditionalBranch, .true_block = body, .false_block = outside}; fixture.block = body; @@ -2103,8 +2138,14 @@ void TestCountedRowImageSelector() { auto index = Value(&induction); if (variant == Variant::KeyStride) index = fixture.Emit(ValueOpcode::ShiftLeftLogical32, {index, Value(1u)}); - const auto key_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32, - {fixture.Emit(ValueOpcode::IAdd32, {row_start, index}), Value(2u)}); + auto element = fixture.Emit(ValueOpcode::IAdd32, {row_start, index}); + if (masked) { + const auto write_mask = variant == Variant::WrongMask + ? fixture.Emit(ValueOpcode::LogicalNot, {active}) : active; + element = fixture.Emit(ValueOpcode::ReadFirstLane, + {fixture.Emit(ValueOpcode::SelectU32, {write_mask, element, Value(0xdeadbeefu)}), active}); + } + const auto key_offset = fixture.Emit(ValueOpcode::ShiftLeftLogical32, {element, Value(2u)}); scalar.offset = variant == Variant::KeyImmediate ? 4u : 0u; const auto key = fixture.Emit(ValueOpcode::ReadConstBuffer, {keys, key_offset}, fixture.AddMemory(scalar, 0x20)); @@ -2123,25 +2164,32 @@ void TestCountedRowImageSelector() { {fixture.Image(image_words), fixture.Sampler({Value(0u), Value(0u), Value(0u), Value(0u)}), fixture.ImageAddress(), Value(true)}, fixture.AddMemory(sample, 0x40)); const auto step = fixture.Emit(ValueOpcode::IAdd32, - {Value(&induction), Value(variant == Variant::WrongStep ? 2u : 1u)}, 0, latch); - induction.AddPhiOperand(entry, Value(0u)); + {Value(&induction), Value(variant == Variant::WrongStep || + variant == Variant::MaskedWrongStep ? 2u : 1u)}, 0, latch); + induction.AddPhiOperand(preheader, Value(0u)); induction.AddPhiOperand(latch, step); fixture.PlanAndTrack(); const auto &source = fixture.program.descriptor_sources[fixture.program.info.images[0].source]; Check(source.indirect_descriptor.has_value(), "counted-row image lost its indirect source"); const auto &selector = source.indirect_descriptor->selector; - Check((selector && selector->rows.has_value()) == (variant == Variant::Valid), + const bool valid = variant == Variant::Valid || variant == Variant::Masked || + variant == Variant::EmptyCountMask; + Check((selector && selector->rows.has_value()) == valid, "counted-row selector lost its proof or accepted mismatched row/induction bounds"); - if (variant == Variant::Valid) + if (valid) Check(selector->stride == 4u && selector->offset == 0u && selector->rows->stride == uint64_t(capacity) * 4u, "counted-row selector changed native row or element byte spacing"); return ExtractResourcePlan(fixture.program); }; auto plan = make_plan(); + make_plan(Variant::Masked); + make_plan(Variant::EmptyCountMask); for (const auto variant : {Variant::Signed, Variant::WrongStep, Variant::Bypass, Variant::WrongRow, Variant::CountStride, Variant::KeyStride, - Variant::CountImmediate, Variant::KeyImmediate}) + Variant::CountImmediate, Variant::KeyImmediate, Variant::VaryingRow, + Variant::EmptyMask, Variant::UnprovenMask, Variant::WrongMask, + Variant::AnyBound, Variant::MaskedWrongStep}) make_plan(variant); constexpr uint32_t CountBase = 0x1000u, KeyBase = 0x2000u, MaterialBase = 0x4000u; diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 2f262672..f0464bcc 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -36182,13 +36182,17 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) { constexpr u32 poison = 0x345d1746u; const auto alignment = static_cast(vulkan.RuntimeContext().GetPhysicalDeviceProperties() .limits.minStorageBufferOffsetAlignment); - const auto slot = std::max(256u, alignment); - const u32 srt = slot, counts = slot * 2u, keys = slot * 3u, material = slot * 4u; + for (const bool masked : {false, true}) for (const u32 wave : {32u, 64u}) { + const auto output_words = masked ? 48u * wave : 64u; + const auto slot = (output_words * 4u + alignment - 1u) / alignment * alignment; + const u32 srt = slot, counts = slot * 2u, keys = slot * 3u, material = slot * 4u; + const auto tile_list = slot * 5u; TestCase test; - test.name = wave == 32u ? "CountedRowImageSelector32" : "CountedRowImageSelector64"; - test.initial.resize((material + 396u * 4u) / 4u, 0u); - std::fill_n(test.initial.begin(), 64u, 0xdeadbeefu); + test.name = masked ? (wave == 32u ? "MaskedRowImageSelector32" : "MaskedRowImageSelector64") + : (wave == 32u ? "CountedRowImageSelector32" : "CountedRowImageSelector64"); + test.initial.resize((masked ? tile_list + 12u : material + 396u * 4u) / 4u, 0u); + std::fill_n(test.initial.begin(), output_words, 0xdeadbeefu); const std::array descriptors{counts, 4u << 16u, 3u, 0x14004u, keys, 4u << 16u, 48u, 0x14004u, material, 396u << 16u, 4u, 0x14004u}; std::copy(descriptors.begin(), descriptors.end(), test.initial.begin() + srt / 4u); @@ -36210,19 +36214,32 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) { put_image(key * 396u + 184u, 0x100u + key, F::k32_32_32_32Float); put_image(32u, 0x104u, F::k32_32_32_32Float); put_image(256u, 0x200u, F::k8SScaled); + if (masked) { + // The scalar tile list reverses workgroups; row selection must remain on the GPU. + test.initial[tile_list / 4u] = 2u; + test.initial[tile_list / 4u + 1u] = 1u; + test.user_data[52] = tile_list; + test.user_data[53] = 4u << 16u; + test.user_data[54] = 3u; + test.user_data[55] = 0x14004u; + } test.expected = test.initial; - test.expected[0] = 0x3f800002u; - test.expected[1] = 0x3f800001u; - test.expected[2] = 0x3f800004u; - test.expected[32] = 0x3f800003u; + for (const auto [index, bits] : std::array, 4>{{ + {0u, 0x3f800002u}, {1u, 0x3f800001u}, {2u, 0x3f800004u}, {32u, 0x3f800003u}}}) { + if (masked) { + for (u32 lane = 1u; lane < wave; lane += 2u) + test.expected[index * wave + lane] = bits; + } else test.expected[index] = bits; + } test.user_data[0] = srt; - test.user_data[50] = 256u; + test.user_data[50] = output_words * 4u; test.user_data[51] = 3u << 28u; test.has_user_data = true; test.has_compute_info = true; test.compute_info.wave_size = wave; - test.compute_info.threads_num[0] = test.compute_info.threads_num[1] = - test.compute_info.threads_num[2] = 1u; + test.compute_info.threads_num[0] = masked ? wave : 1u; + test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1u; + test.compute_info.thread_ids_num = masked ? 1 : 0; test.compute_info.group_id[0] = true; test.compute_info.workgroup_register = 44u; test.dispatch_x = 3u; @@ -36230,32 +36247,58 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) { code.insert(code.end(), {EncodeSmem0(0x03, 8), EncodeSmem1(0, 125), EncodeSmem0(0x02, 16), EncodeSmem1(32, 125)}); code.push_back(EncodeSop2(0x1e, 22, 44, InlineU32(2))); + if (masked) { + code.push_back(EncodeVop1(0x01, 2, Vgpr(0))); // Preserve the lane for disjoint output addresses. + code.insert(code.end(), {EncodeSmem0(0x08, 23, 26), EncodeSmem1(0, 22)}); + code.push_back(EncodeVop1(0x01, 18, 23)); + code.push_back(EncodeVop1(0x02, 22, Vgpr(18))); + code.push_back(EncodeSop2(0x1e, 22, 22, InlineU32(2))); + } code.insert(code.end(), {EncodeSmem0(0x08, 21, 4), EncodeSmem1(0, 22)}); AppendSMovLiteral(&code, 20, 0u); AppendVMovU32(&code, 0, 0x3f000000u); AppendVMovU32(&code, 1, 0x3f000000u); + const auto masked_exit_branch = masked ? code.size() + 2u : 0u; + if (masked) { + code.push_back(EncodeVop2(0x1b, 6, InlineU32(1), 2)); + code.push_back(EncodeVopc(0xd5, InlineU32(0), 6)); + code.push_back(0u); // Skip the loop when the new EXEC is empty. + } const auto header = code.size(); code.push_back(EncodeSopc(0x0a, 20, 21)); const auto exit_branch = code.size(); code.push_back(0u); - code.push_back(EncodeSop2(0x1e, 22, 44, InlineU32(4))); - code.push_back(EncodeSop2(0x02, 22, 22, 20)); + if (masked) { + code.insert(code.end(), {EncodeVop3Word0(0x346, 3), + EncodeVop3Word1(Vgpr(18), InlineU32(4), 20)}); + code.push_back(EncodeVop1(0x02, 22, Vgpr(3))); + } else { + code.push_back(EncodeSop2(0x1e, 22, 44, InlineU32(4))); + code.push_back(EncodeSop2(0x02, 22, 22, 20)); + } code.push_back(EncodeSop2(0x1e, 22, 22, InlineU32(2))); code.insert(code.end(), {EncodeSmem0(0x08, 23, 6), EncodeSmem1(0, 22)}); code.push_back(EncodeSop2(0x26, 24, 23, 255)); code.push_back(396u); code.push_back(EncodeSop2(0x02, 24, 24, 255)); code.push_back(184u); code.insert(code.end(), {EncodeSmem0(0x0b, 32, 8), EncodeSmem1(0, 24), EncodeMimg0(0x27, 1), EncodeMimg1(4, 0, 8, 10)}); - code.push_back(EncodeVop1(0x01, 30, 22)); + if (masked) { + code.push_back(EncodeVop2(0x1a, 5, InlineU32(2), 2)); + code.insert(code.end(), {EncodeVop3Word0(0x346, 30), + EncodeVop3Word1(22, InlineU32(std::countr_zero(wave)), Vgpr(5))}); + } else code.push_back(EncodeVop1(0x01, 30, 22)); AppendBufferStoreDword(&code, 4, 30); code.push_back(EncodeSop2(0x02, 20, 20, InlineU32(1))); code.push_back(EncodeSopp(0x02, header - code.size() - 1u)); code[exit_branch] = EncodeSopp(0x04, code.size() - exit_branch - 1u); + if (masked) code[masked_exit_branch] = EncodeSopp(0x08, code.size() - masked_exit_branch - 1u); AppendEnd(&code); test.opcodes = {O::S_LOAD_DWORDX8, O::S_LOAD_DWORDX4, O::S_BUFFER_LOAD_DWORD, O::S_BUFFER_LOAD_DWORDX8, O::S_MOV_B32, O::S_LSHL_B32, O::S_ADD_I32, O::S_MUL_I32, O::S_CMP_LT_U32, O::S_CBRANCH_SCC0, O::S_BRANCH, O::V_MOV_B32, O::IMAGE_SAMPLE, O::BUFFER_STORE_DWORD, O::S_ENDPGM}; + if (masked) test.opcodes.insert(test.opcodes.end(), {O::V_READFIRSTLANE_B32, + O::V_LSHL_ADD_U32, O::V_LSHLREV_B32, O::V_AND_B32, O::V_CMPX_NE_U32, O::S_CBRANCH_EXECZ}); const auto compiled = CompileCase(test, vulkan.SubgroupSize()); Require(test.name, "exact active images", compiled.resources.images.size() == 4u && std::ranges::all_of(compiled.resources.images, [](const auto &image) {