shader: specialize bounded pointer and packed image tables

This commit is contained in:
nmzik
2026-10-13 01:12:19 +02:00
parent c9ec82e32e
commit 1fa79d5cac
5 changed files with 219 additions and 24 deletions
@@ -604,8 +604,20 @@ void FoldInstruction(Block& block, Block::iterator instruction,
});
return;
case ValueOpcode::BitwiseAnd32:
if (!FoldU32(inst, [](uint32_t a, uint32_t b) { return a & b; })) {
ReplaceBinaryIdentity(inst, Type::U32, 0xffffffffu);
if (FoldU32(inst, [](uint32_t a, uint32_t b) { return a & b; }) ||
ReplaceBinaryIdentity(inst, Type::U32, UINT32_MAX)) return;
for (uint32_t arm = 0; arm < 2u; ++arm) {
const auto mask = Arg(inst, arm);
const auto* shifted = Arg(inst, arm ^ 1u).TryInstruction();
if (!IsImmediate(mask, Type::U32) || shifted == nullptr ||
shifted->GetOpcode() != ValueOpcode::ShiftLeftLogical32 || !HasOneUse(*shifted)) continue;
const auto shift = Arg(*shifted, 1);
if (!IsImmediate(shift, Type::U32) || shift.U32() >= 32u) continue;
const auto masked = worklist.Prepend(block, instruction, ValueOpcode::BitwiseAnd32,
{Arg(*shifted, 0), Value(mask.U32() >> shift.U32())});
Replace(inst, Value(&*worklist.Prepend(block, instruction, ValueOpcode::ShiftLeftLogical32,
{Value(&*masked), shift})));
return;
}
return;
case ValueOpcode::BitwiseAnd64:
@@ -282,14 +282,14 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
}
return true;
};
const auto read_keys = [&](const ShaderBufferResource& material, uint64_t first,
const auto read_keys = [&](uint64_t base, uint64_t size, uint64_t first,
uint64_t step, uint64_t count) {
keys.resize(count);
if (step == 4u) {
return ReadScalarTable(material.Base48(), material.GetSize(), first, runtime, keys);
return ReadScalarTable(base, size, first, runtime, keys);
}
for (auto& key: keys) {
if (!ReadScalarTable(material.Base48(), material.GetSize(), first, runtime, {&key, 1}))
if (!ReadScalarTable(base, size, first, runtime, {&key, 1}))
return false;
first += step;
}
@@ -468,9 +468,17 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
if (selected[key]) keys.push_back(key);
} else if (!indirect.selector_first.IsEmpty()) {
ShaderBufferResource material;
uint64_t base = 0, size = UINT64_MAX;
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) ||
!clean.Evaluate(indirect.selector_first, first) ||
if (material_value.dword_count == 2u) {
if (selector->stride == 0u) return false;
base = (uint64_t {material_value.dwords[1]} << 32u) | material_value.dwords[0];
} else {
if (!DecodeBufferDescriptor(material_value, material)) return false;
base = material.Base48();
size = material.GetSize();
}
if (!clean.Evaluate(indirect.selector_first, first) ||
!evaluate_count(count) || count > MaxIndirectDescriptorProbes)
return false;
// SMEM uses its instruction's byte stride; V# stride only supplies its bounds.
@@ -481,9 +489,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
uint64_t {first} + count > material.NumRecords())) return false;
const auto end = count != 0u
? (uint64_t {first} + count - 1u) * step + selector->offset + 4u : 0u;
if (end > uint64_t {UINT32_MAX} + 1u || end > material.GetSize())
if (end > uint64_t {UINT32_MAX} + 1u || end > size)
return false;
if (!read_keys(material, uint64_t {first} * step + selector->offset, step, count))
if (!read_keys(base, size, uint64_t {first} * step + selector->offset, step, count))
return false;
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
@@ -521,7 +529,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
const auto count = std::min(domain_count,
available < 4u ? 0u : (available - 4u) / selector->stride + 1u);
if (count > MaxIndirectDescriptorProbes ||
!read_keys(material, selector->offset, selector->stride, count)) return false;
!read_keys(material.Base48(), material.GetSize(), selector->offset, selector->stride, count)) return false;
if (count < domain_count) keys.push_back(0u); // An out-of-bounds scalar read returns zero.
}
if (selector != nullptr && selector->byte_ranges.empty()) {
@@ -1933,7 +1933,7 @@ private:
bool MatchScalarMaterialKey(Value key, DescriptorSource::IndirectDescriptor& indirect,
DescriptorSource& material_source) {
const auto* read = key.Resolve().TryInstruction();
if (read == nullptr || read->GetOpcode() != ValueOpcode::ReadConstBuffer) return false;
if (read == nullptr) return false;
uint32_t memory_index = 0;
const auto* memory = ScalarReadMemory(*read, memory_index);
Value index;
@@ -1941,8 +1941,8 @@ private:
DescriptorSource count_source;
uint32_t row_stride = 0;
uint32_t offset = 0, stride = 0;
if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read)) return false;
if (memory == nullptr || memory->offset > INT32_MAX ||
!MemoryIndexBelongsTo(memory_index, *read)) return false;
if (MatchTableOffset(read->Arg(1), index, offset, stride) &&
((offset | stride | memory->offset) & 3u) == 0u &&
uint64_t {offset} + memory->offset <= UINT32_MAX) {
@@ -1951,6 +1951,8 @@ private:
count = bound->Arg(1);
}
}
// Raw pointers have no V# bounds for enumerating an unbounded selector domain.
if (memory->kind == ResourceKind::ScalarAddress && count.IsEmpty()) return false;
const bool rows = count.IsEmpty() && memory->offset == 0u &&
MatchCountedRows(*read, count_source, row_stride);
if (rows) {
@@ -2476,7 +2478,9 @@ private:
if (indirect.key_count.IsEmpty() &&
!MatchUniformizedMaterialKey(key, handle, indirect, material_source) &&
!MatchReversedImageRange(key, indirect)) {
return false;
const auto maximum = SelectorMaximum(key);
if (maximum >= MaxIndirectDescriptorProbes) return false;
indirect.key_count = Value(maximum + 1u);
}
if ((table_offset & 3u) != 0u) return false;
} else {
+72 -10
View File
@@ -570,7 +570,7 @@ void TestGuardedDirectImageTable() {
Finite, FiniteZero, FiniteMissing, FiniteWrong, FiniteSigned, FiniteAdd, FiniteXor,
FiniteShared, FiniteSharedBypass, FiniteLate
};
const auto make_plan = [](Guard guard, uint32_t operand_order = 0u) {
const auto make_plan = [](Guard guard, uint32_t operand_order = 0u, bool packed_key = false) {
const bool finite = guard >= Guard::Finite;
const bool bounded = guard >= Guard::Bounded && !finite;
const bool exec_bound = bounded && guard >= Guard::BoundExec;
@@ -751,11 +751,15 @@ void TestGuardedDirectImageTable() {
fixture.AddMemory(memory, 0x20));
}
fixture.block = descriptor_block;
if (!bounded && !finite) key = fixture.Emit(ValueOpcode::FindILsb32, {mask});
if (!bounded && !finite) key = fixture.Emit(packed_key ? ValueOpcode::ShiftRightLogical32
: ValueOpcode::FindILsb32,
packed_key ? std::initializer_list<Value>{mask, Value(12u)}
: std::initializer_list<Value>{mask});
auto scaled = bounded ? fixture.Emit(ValueOpcode::IMul32, {key, Value(table_stride)})
: fixture.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)});
if (packed_key) scaled = fixture.Emit(ValueOpcode::BitwiseAnd32, {scaled, Value(992u)});
const auto offset = fixture.Emit(ValueOpcode::IAdd32,
{bounded ? fixture.Emit(ValueOpcode::IMul32, {key, Value(table_stride)})
: fixture.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)}),
Value(table_offset)});
{scaled, Value(table_offset)});
std::array<Value, 8> words;
for (uint32_t group = 0; group < 2u; ++group) {
// Separate equivalent pointer handles mirror the two scalar x4 loads.
@@ -810,7 +814,7 @@ void TestGuardedDirectImageTable() {
fixture.Emit(ValueOpcode::ImageWrite, {output, fixture.ImageAddress(), texel, Value(true)},
fixture.AddMemory(memory, 0x130));
}
if (exec_bound) ConstantPropagationPass(fixture.program.blocks);
if (exec_bound || packed_key) ConstantPropagationPass(fixture.program.blocks);
if (guard == Guard::BoundExec || guard == Guard::BoundExecFactored) {
std::vector<const Inst *> pending{active.ResolveInstruction()};
uint32_t conjunctions = 0;
@@ -876,6 +880,23 @@ void TestGuardedDirectImageTable() {
.read_specialization_memory = ReadLinearTestMemory};
ResourceSnapshot snapshot;
ResourceSpecialization specialization;
auto packed = make_plan(Guard::Plain, 0u, true);
memory.fail_address = table + 32u * 32u;
Check(MaterializeResources(packed, runtime, snapshot, specialization) &&
memory.descriptor_reads == 32u && snapshot.images.size() == 32u,
"packed five-bit selector did not enumerate exactly its finite domain");
for (uint32_t key = 0; key < 32u; ++key)
memory.words[(table - memory.base) / 4u + key * 8u] = 0x100u;
memory.descriptor_reads = 0u;
Check(MaterializeResources(packed, runtime, snapshot, specialization) &&
memory.descriptor_reads == 32u && snapshot.images.size() == 1u &&
specialization.images[0].indirect_root == ImageResource::NoIndirectImage &&
specialization.images[0].indirect_search_iterations == 0u,
"packed five-bit selector retained duplicate image bindings or lost keys");
memory.fail_address = UINT64_MAX;
memory.reads = 0u;
memory.descriptor_reads = 0u;
fill_table(memory, table);
Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
snapshot.images.size() == 32u && specialization.images.size() == 32u &&
snapshot.flattened_srt[specialization.images[0].indirect_mapping_offset] == 32u &&
@@ -1026,7 +1047,7 @@ void TestBoundedComputeImageLoop() {
Masked, MaskedWrongGuard, MaskResurrection, StatusOverwrite,
GuardedDiamond, GuardedDiamondBypass
};
const auto make_plan = [](Variant variant) {
const auto make_plan = [](Variant variant, bool material = false) {
const bool diamond = variant >= Variant::GuardedDiamond;
Fixture fixture(diamond ? ShaderType::Vertex : ShaderType::Compute);
fixture.program.wave_size = 64u;
@@ -1102,7 +1123,7 @@ void TestBoundedComputeImageLoop() {
{Value(static_cast<uint32_t>(diamond ? StageInputKind::VertexIndex
: StageInputKind::LocalInvocationId)), Value(0u)});
const auto count = variant == Variant::DivergentBound ? local : fixture.UserData(2);
const auto in_range = fixture.Emit(ValueOpcode::SLessThan32,
const auto in_range = fixture.Emit(material ? ValueOpcode::ULessThan32 : ValueOpcode::SLessThan32,
{variant == Variant::WrongGuard
? Value(0u) : key,
count}, 0, compare);
@@ -1249,8 +1270,17 @@ void TestBoundedComputeImageLoop() {
fixture.block = variant == Variant::GuardBlock ? header : body;
const auto table = fixture.Address(fixture.UserData(0), fixture.UserData(1));
auto image_key = key;
if (material) {
MemoryInfo memory;
memory.kind = ResourceKind::ScalarAddress;
memory.offset = 2756u;
image_key = fixture.Emit(ValueOpcode::LoadAddressU32,
{table, fixture.Emit(ValueOpcode::IMul32, {key, Value(48u)}), Value(0u), Value(true)},
fixture.AddMemory(memory, 0x2ff0));
}
const auto offset = fixture.Emit(ValueOpcode::IAdd32,
{fixture.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)}),
{fixture.Emit(ValueOpcode::ShiftLeftLogical32, {image_key, Value(5u)}),
Value(0x6b0u)});
std::array<Value, 8> words;
for (uint32_t word = 0; word < words.size(); ++word) {
@@ -1273,10 +1303,14 @@ void TestBoundedComputeImageLoop() {
fixture.PlanAndTrack();
const auto source = fixture.program.info.images[0].source;
const auto &indirect = fixture.program.descriptor_sources[source].indirect_descriptor;
Check(indirect && !indirect->selector.has_value() &&
Check(indirect && indirect->selector.has_value() == material &&
indirect->table_offset == 0x6b0u &&
indirect->key_count.Resolve() == count.Resolve(),
"bounded compute loop lost its runtime image count");
if (material) Check(indirect->selector->stride == 48u &&
indirect->selector->offset == 2756u &&
fixture.program.descriptor_sources[indirect->selector->source].dword_count == 2u,
"bounded scalar selector lost its pointer or record layout");
return ExtractResourcePlan(fixture.program);
};
@@ -1360,6 +1394,34 @@ void TestBoundedComputeImageLoop() {
user_data[2] = 65537u;
Check(!MaterializeResources(plan, runtime, snapshot, specialization),
"oversized compute loop bound was accepted for image enumeration");
auto material = make_plan(Variant::Bounded, true);
for (const auto variant : {Variant::WrongGuard, Variant::EntryBypass, Variant::WrongStep})
CheckFatal([&] { make_plan(variant, true); }, "not a valid runtime value",
"raw pointer selector accepted an unbounded source");
const auto first_key = uint64_t{user_data[0]} + 2756u;
memory.words[(first_key - memory.base) / 4u] = 0u;
memory.words[(first_key + 48u - memory.base) / 4u] = 2u;
user_data[2] = 0u;
memory.reads = 0u;
memory.fail_address = first_key;
Check(MaterializeResources(material, runtime, snapshot, specialization) && memory.reads == 0u &&
snapshot.images.size() == 1u && snapshot.images[0].dwords[0] == 0u,
"empty raw pointer selector read an unreachable record");
user_data[2] = 2u;
memory.fail_address = first_key + 96u;
memory.descriptor_reads = 0u;
Check(MaterializeResources(material, runtime, snapshot, specialization) &&
memory.descriptor_reads == 2u && snapshot.images.size() == 2u &&
snapshot.images[0].dwords[0] == 0x100u && snapshot.images[1].dwords[0] == 0x102u,
"bounded raw pointer selector lost sparse keys or read an extra record");
memory.fail_address = first_key + 48u;
Check(!MaterializeResources(material, runtime, snapshot, specialization),
"raw pointer selector accepted an unreadable selected record");
memory.fail_address = UINT64_MAX;
memory.words[(first_key - memory.base) / 4u] = UINT32_MAX - 1000u;
Check(!MaterializeResources(material, runtime, snapshot, specialization),
"raw pointer selector silently discarded an invalid negative key");
}
void TestDescendingImageRange() {
+109
View File
@@ -36505,6 +36505,113 @@ void CheckCountedRowImageSelector(VulkanHarness &vulkan) {
}
}
void CheckPointerImageTable(VulkanHarness &vulkan) {
using O = ShaderOpcode;
constexpr u32 base = 256u, table_offset = 24364u;
for (const u32 variant : {0u, 1u, 2u})
for (const u32 wave : {32u, 64u}) {
const bool packed = variant == 2u;
const u32 count = variant * 2u;
TestCase test;
test.name = packed ? "PackedPointerImageTable" : "CountedPointerImageTable";
test.dispatch_x = packed ? 4u : 1u;
test.bda_mappings = {{0u, 0u}};
test.initial.resize((base + table_offset + 32u * 32u) / 4u);
std::fill_n(test.initial.begin(), test.dispatch_x, 0xdeadbeefu);
test.initial[(base + 492u) / 4u] = count;
for (u32 row = 0; row < 32u; ++row) {
auto image = MakeSampledTextureData(Prospero::BufferFormat::k32Float);
image[0] = row + 1u;
image[3] = DstSel(4, 0, 0, 1) |
(static_cast<u32>(Prospero::ImageType::kColor2D) << 28u);
std::copy_n(image.begin(), 8,
test.initial.begin() + (base + table_offset + row * 32u) / 4u);
}
test.initial[(base + 2756u) / 4u] = 1u;
test.initial[(base + 2756u + 48u) / 4u] = 31u;
constexpr std::array fields{31u, 0u, 7u, 1u};
for (u32 row = 0; row < fields.size(); ++row)
test.initial[(base + 5008u + row * 60u) / 4u] = 0xffe00abcu | (fields[row] << 12u);
test.user_data[0] = base;
const auto output = MakeStructuredStorageBufferData(0u, test.dispatch_x * 4u);
std::copy_n(output.begin(), 4, test.user_data.begin() + 48);
test.has_user_data = 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.group_id[0] = true;
test.compute_info.workgroup_register = 44u;
auto &code = test.code;
AppendVMovLiteral(&code, 0, 0u);
AppendVMovLiteral(&code, 20, 0x3f000000u);
AppendVMovLiteral(&code, 21, 0x3f000000u);
size_t loop = 0, exit = 0, skip = 0;
if (packed) {
code.insert(code.end(), {EncodeSop2(0x26, 23, 44, InlineU32(60)),
EncodeSop2(0x02, 23, 23, 255), 5008u,
EncodeSmem0(0x00, 24), EncodeSmem1(0, 23),
EncodeSop2(0x20, 24, 24, InlineU32(12)),
EncodeSop2(0x1e, 23, 24, InlineU32(5)),
EncodeSop2(0x0e, 23, 23, 255), 992u});
} else {
code.insert(code.end(), {EncodeSmem0(0x00, 21), EncodeSmem1(492, 125),
EncodeSMovB32(20, InlineU32(0))});
loop = code.size();
code.push_back(EncodeSopc(0x0a, 20, 21));
exit = code.size(); code.push_back(0u);
code.insert(code.end(), {EncodeSop2(0x26, 23, 20, InlineU32(48)),
EncodeSop2(0x02, 23, 23, 255), 2756u,
EncodeSmem0(0x00, 24), EncodeSmem1(0, 23),
EncodeSopc(0x03, 24, InlineU32(0))});
skip = code.size(); code.push_back(0u);
code.push_back(EncodeSop2(0x1e, 23, 24, InlineU32(5)));
}
code.insert(code.end(), {EncodeSop2(0x02, 23, 23, 255), table_offset,
EncodeSmem0(0x03, 8), EncodeSmem1(0, 23),
EncodeMimg0(0x27, 1), EncodeMimg1(4, 20, 2, 10),
EncodeVop2(0x03, 0, Vgpr(4), 0)});
if (!packed) {
code[skip] = EncodeSopp(0x04, code.size() - skip - 1u);
code.push_back(EncodeSop2(0x02, 20, 20, InlineU32(1)));
code.push_back(EncodeSopp(0x02, loop - code.size() - 1u));
code[exit] = EncodeSopp(0x04, code.size() - exit - 1u);
}
code.push_back(EncodeSop2(0x1e, 25, 44, InlineU32(2)));
code.push_back(EncodeVop1(0x01, 30, 25));
AppendBufferStoreDword(&code, 0, 30);
AppendEnd(&code);
test.opcodes = {O::S_LOAD_DWORD, O::S_MUL_I32, O::S_ADD_I32, O::S_LSHL_B32,
O::S_LOAD_DWORDX8, O::IMAGE_SAMPLE, O::V_ADD_F32, O::V_MOV_B32,
O::BUFFER_STORE_DWORD, O::S_ENDPGM};
if (packed) test.opcodes.insert(test.opcodes.end(), {O::S_LSHR_B32, O::S_AND_B32});
else test.opcodes.insert(test.opcodes.end(), {O::S_MOV_B32, O::S_CMP_LT_U32,
O::S_CMP_GE_I32, O::S_CBRANCH_SCC0, O::S_BRANCH});
const auto compiled = CompileCase(test, vulkan.SubgroupSize());
Require(test.name, "bounded image candidates",
packed ? compiled.resources.images.size() == 32u
: compiled.resources.images.size() == std::max(count, 1u),
"descriptor planning escaped its selector domain");
std::vector<VulkanHarness::Image> images;
for (const auto &descriptor : compiled.resources.images)
images.push_back(vulkan.CreateImage2D(test.name, 1, 1, vk::Format::eR32Sfloat,
vk::ImageUsageFlagBits::eSampled,
{std::bit_cast<u32>(static_cast<float>(descriptor.dwords[0]))}, 1,
vk::ImageLayout::eShaderReadOnlyOptimal));
auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size(), true);
auto sampler = vulkan.CreateSampler(test.name, 1);
vulkan.Dispatch(test, compiled, buffer, nullptr, nullptr, nullptr, nullptr, sampler, images);
std::vector<u32> expected(test.dispatch_x, std::bit_cast<u32>(count == 0u ? 0.0f : 34.0f));
if (packed) for (u32 row = 0; row < fields.size(); ++row)
expected[row] = std::bit_cast<u32>(static_cast<float>(fields[row] + 1u));
CompareWords(test, "bounded pointer selectors and packed fields", expected,
vulkan.ReadBuffer(test.name, buffer, expected.size()));
vulkan.DestroyBuffer(&buffer);
vulkan.Device().destroySampler(sampler);
for (auto &image : images) vulkan.DestroyImage(&image);
std::printf("[compute] %s wave=%u count=%u ok\n", test.name, wave, count);
}
}
void CheckGuardedRecordImageTable(VulkanHarness &vulkan) {
using O = ShaderOpcode;
constexpr u32 table = 256u;
@@ -45134,6 +45241,7 @@ int main(int argc, char **argv) {
CheckImageSamplerSpecialization();
VulkanHarness vulkan;
CheckCountedRowImageSelector(vulkan);
CheckPointerImageTable(vulkan);
CheckGuardedRecordImageTable(vulkan);
CheckIndirectImageKeySwitch(vulkan);
CheckFiniteInlineSamplerPhi(vulkan);
@@ -45301,6 +45409,7 @@ int main(int argc, char **argv) {
CheckRuntimeBufferRecords(vulkan);
CheckComputeThreadDimensions(vulkan);
CheckCountedRowImageSelector(vulkan);
CheckPointerImageTable(vulkan);
CheckGuardedRecordImageTable(vulkan);
CheckIndirectImageKeySwitch(vulkan);
CheckDynamicSampleOffsets(vulkan, 32u);