shader: correlate masked first-lane material rows

This commit is contained in:
nmzik 2026-10-12 22:35:48 +02:00
parent 8ee31abd9d
commit 68d7bc94e2
5 changed files with 162 additions and 39 deletions

View File

@ -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<size_t> {
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<uint32_t>((*affine)[2]);
row_stride = static_cast<uint32_t>((*affine)[2] * 4u);
return true;
}) != nullptr) return true;
}, selected_lane ? LaneQuantifier::All : LaneQuantifier::Any) != nullptr) return true;
}
return false;
}

View File

@ -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<uint32_t>(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<const Inst*> m_visiting;
std::unordered_set<const Inst*> 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

View File

@ -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);

View File

@ -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<uint32_t>(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<uint32_t>(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;

View File

@ -36182,13 +36182,17 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) {
constexpr u32 poison = 0x345d1746u;
const auto alignment = static_cast<u32>(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<u32, 12> 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<std::pair<u32, u32>, 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) {