shader: resolve nested descriptor selector tables with exact scalar offsets

This commit is contained in:
nmzik
2026-10-06 03:41:59 +04:00
parent 9caec9f2c8
commit c23de62184
5 changed files with 209 additions and 101 deletions
+12 -5
View File
@@ -485,15 +485,22 @@ struct BlockInfo {
struct DescriptorSource {
struct IndirectDescriptor {
uint32_t material_source = UINT32_MAX;
struct SelectorRead {
uint32_t source = UINT32_MAX;
uint32_t stride = 0;
uint32_t offset = 0;
uint32_t immediate = 0;
uint32_t shift = 0;
uint32_t bits = UINT32_MAX;
bool operator==(const SelectorRead& other) const = default;
};
std::vector<SelectorRead> selectors;
uint32_t table_source = 0;
uint32_t selector_stride = 0;
uint32_t selector_offset = 0;
uint32_t table_offset = 0;
uint32_t table_immediate = 0;
uint32_t table_stride = 0;
uint32_t workgroup_axis = UINT32_MAX;
uint32_t selector_shift = 0;
uint32_t selector_bits = UINT32_MAX;
Value key_count;
Value selector_first;
Value selector_mask;
@@ -249,6 +249,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
uint32_t maximum_resources, Normalize&& normalize) {
const auto descriptor_bytes = dword_count * sizeof(uint32_t);
const auto& sources = indirect.sources;
const auto* selector = indirect.selectors.empty() ? nullptr : &indirect.selectors.front();
auto& keys = program.material_keys;
const auto read_keys = [&](const ShaderBufferResource& material, uint64_t first,
uint64_t step, uint64_t count) {
@@ -269,8 +270,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
keys.clear();
DescriptorValue material_value;
DescriptorValue table_value;
if ((indirect.material_source != UINT32_MAX &&
!clean.EvaluateDescriptor(indirect.material_source, material_value)) ||
if ((selector != nullptr &&
!clean.EvaluateDescriptor(selector->source, material_value)) ||
!clean.EvaluateDescriptor(indirect.table_source, table_value)) {
return false;
}
@@ -283,7 +284,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
} else {
return false;
}
if (indirect.material_source == UINT32_MAX) {
if (selector == nullptr) {
uint32_t key_count = 0;
if (indirect.workgroup_axis != UINT32_MAX) {
if (indirect.workgroup_axis >= runtime.workgroup_counts.size()) return false;
@@ -295,7 +296,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
}
if (indirect.table_stride == 0u || key_count > MaxIndirectDescriptorProbes ||
(key_count != 0u && uint64_t {indirect.table_offset} +
(key_count != 0u && uint64_t {indirect.table_offset} + indirect.table_immediate +
uint64_t {key_count - 1u} * indirect.table_stride + descriptor_bytes >
UINT32_MAX + 1ull)) {
return false;
@@ -306,20 +307,20 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
ShaderBufferResource material;
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) || material.Type() != 0u ||
((indirect.selector_shift != 0u || indirect.selector_bits != UINT32_MAX) &&
((selector->shift != 0u || selector->bits != UINT32_MAX) &&
(material.Base48() & 3u) != 0u) ||
material.SwizzleEnabled() ||
material.AddTid() || material.OutOfBounds() != 0u ||
uint64_t {indirect.selector_offset} + 4u > material.Stride() ||
uint64_t {selector->offset} + 4u > material.Stride() ||
!clean.Evaluate(indirect.selector_first, first) ||
!clean.Evaluate(indirect.key_count, count) || count > MaxIndirectDescriptorProbes ||
uint64_t {first} + count > material.NumRecords())
return false;
if (count != 0u && (uint64_t {first} + count - 1u) * material.Stride() +
indirect.selector_offset + 4u >
selector->offset + 4u >
uint64_t {UINT32_MAX} + 1u)
return false;
if (!read_keys(material, uint64_t {first} * material.Stride() + indirect.selector_offset,
if (!read_keys(material, uint64_t {first} * material.Stride() + selector->offset,
material.Stride(), count)) return false;
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
@@ -335,8 +336,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
keys.reserve(std::popcount(mask));
while (mask != 0u) {
const auto index = std::countr_zero(mask);
const auto offset = static_cast<uint64_t>(indirect.selector_offset) +
static_cast<uint64_t>(index) * indirect.selector_stride;
const auto offset = static_cast<uint64_t>(selector->offset) +
static_cast<uint64_t>(index) * selector->stride;
if (offset > UINT32_MAX) return false;
uint32_t key = 0;
if (!ReadScalarTable(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
@@ -351,8 +352,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
}
// The first aligned offset includes the immediate added after shader U32 arithmetic.
const auto step = std::max<uint64_t>(4u,
std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u));
const uint64_t first = indirect.selector_offset;
std::gcd<uint64_t>(selector->stride, uint64_t {1} << 32u));
const uint64_t first = (selector->offset % step & ~uint64_t {3}) +
(selector->immediate & ~3u);
const auto size = material.GetSize();
const auto limit = std::min(first + (uint64_t {1} << 32u) - step, size >= 4u ? size - 4u : 0u);
const auto probe_count = size >= 4u && first <= limit ? (limit - first) / step + 1u : 0u;
@@ -360,13 +362,28 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
return false;
}
if (!read_keys(material, first, step, probe_count)) return false;
if (indirect.selector_shift == 0u) keys.push_back(0u);
if (selector->shift == 0u) keys.push_back(0u);
}
if (indirect.material_source != UINT32_MAX) {
for (auto& key: keys) key = (key >> indirect.selector_shift) & indirect.selector_bits;
if (indirect.selector_shift != 0u || indirect.selector_bits != UINT32_MAX) {
keys.push_back(0u); // An out-of-range material load returns zero.
for (uint32_t step = 0; step < indirect.selectors.size(); ++step) {
const auto& read = indirect.selectors[step];
if (step != 0u) {
DescriptorValue value;
ShaderBufferResource material;
if (!clean.EvaluateDescriptor(read.source, value) || !DecodeBufferDescriptor(value, material))
return false;
// Distinct keys can alias after U32 arithmetic and DWORD alignment.
for (auto& key: keys) key = (key * read.stride + read.offset) & ~3u;
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
for (auto& key: keys) {
const uint64_t offset = uint64_t {key} + (read.immediate & ~3u);
if (!ReadScalarTable(material.Base48(), material.GetSize(), offset, runtime, {&key, 1}))
return false;
}
}
for (auto& key: keys) key = (key >> read.shift) & read.bits;
if (step == 0u && (read.shift != 0u || read.bits != UINT32_MAX))
keys.push_back(0u); // An out-of-range root material load returns zero.
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
@@ -383,8 +400,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
DescriptorValue candidate;
candidate.dword_count = dword_count;
if (sources.empty()) {
const auto table_offset =
static_cast<uint32_t>(key * indirect.table_stride) + indirect.table_offset;
const uint64_t table_offset =
uint64_t {(key * indirect.table_stride + indirect.table_offset) & ~3u} +
(indirect.table_immediate & ~3u);
if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
std::span(candidate.dwords).first(dword_count))) {
return false;
@@ -1010,9 +1028,10 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
if (!source.indirect_descriptor.has_value()) continue;
const auto& indirect = *source.indirect_descriptor;
plan.requires_specialization_memory = true;
capture_indirect_reads |= indirect.material_source != UINT32_MAX || !indirect.sources.empty();
MarkCleanFlatSlots(plan, Source(plan, indirect.material_source), plan.clean_flat_slots,
indirect.selector_mask);
capture_indirect_reads |= !indirect.selectors.empty() || !indirect.sources.empty();
for (const auto& read: indirect.selectors)
MarkCleanFlatSlots(plan, Source(plan, read.source), plan.clean_flat_slots);
MarkCleanFlatSlots(plan, nullptr, plan.clean_flat_slots, indirect.selector_mask);
MarkCleanFlatSlots(plan, nullptr, plan.clean_flat_slots, indirect.selector_first);
MarkCleanFlatSlots(plan, nullptr, plan.clean_flat_slots, indirect.key_count);
if (indirect.sources.empty()) {
@@ -437,6 +437,8 @@ private:
std::array<Value, 8> roots {};
std::array<uint32_t, 8> memory {};
std::array<const Inst*, 8> reads {};
bool split_offsets = false;
bool retain_reads = false;
};
[[noreturn]] void Fail(uint32_t pc, const std::string& reason) const {
@@ -842,10 +844,9 @@ private:
if (current.indirect_descriptor.has_value()) {
const auto& a = *current.indirect_descriptor;
const auto& b = *descriptor.indirect_descriptor;
if (a.material_source != b.material_source || a.table_source != b.table_source ||
a.selector_stride != b.selector_stride || a.selector_offset != b.selector_offset ||
a.selector_shift != b.selector_shift || a.selector_bits != b.selector_bits ||
a.table_offset != b.table_offset || a.table_stride != b.table_stride ||
if (a.selectors != b.selectors || a.table_source != b.table_source ||
a.table_offset != b.table_offset || a.table_immediate != b.table_immediate ||
a.table_stride != b.table_stride ||
a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
!EquivalentValue(m_program, a.key_count, b.key_count) ||
a.selector_first.IsEmpty() != b.selector_first.IsEmpty() ||
@@ -949,7 +950,7 @@ private:
return ValidateSource(descriptor, bad_dword);
}
bool MatchMaterialOffset(Value value, uint32_t& stride, uint32_t& offset) const {
Value MaterialOffset(Value value, uint32_t& stride, uint32_t& offset) const {
stride = 1u;
offset = 0u;
value = value.Resolve();
@@ -974,24 +975,24 @@ private:
value = remaining.Resolve();
}
// Treat the remaining GPU-derived scalar as any U32 value.
return stride != 0u && value.GetType() == Type::U32 && !value.IsImmediate();
return stride != 0u && value.GetType() == Type::U32 && !value.IsImmediate() ? value : Value {};
}
Value MaterialKey(Value key, DescriptorSource::IndirectDescriptor& indirect) const {
indirect.selector_shift = 0u;
indirect.selector_bits = UINT32_MAX;
Value MaterialKey(Value key, DescriptorSource::IndirectDescriptor::SelectorRead& selector) const {
selector.shift = 0u;
selector.bits = UINT32_MAX;
const auto* inst = key.Resolve().TryInstruction();
uint32_t immediate;
if (inst != nullptr && inst->GetOpcode() == ValueOpcode::BitwiseAnd32) {
if (ImmediateU32(inst->Arg(1), immediate)) key = inst->Arg(0);
else if (ImmediateU32(inst->Arg(0), immediate)) key = inst->Arg(1);
else return key;
indirect.selector_bits = immediate;
selector.bits = immediate;
inst = key.Resolve().TryInstruction();
}
if (inst != nullptr && inst->GetOpcode() == ValueOpcode::ShiftRightLogical32 &&
ImmediateU32(inst->Arg(1), immediate) && immediate != 0u && immediate < 32u) {
indirect.selector_shift = immediate;
selector.shift = immediate;
key = inst->Arg(0);
}
return key.Resolve();
@@ -1443,9 +1444,9 @@ private:
static_cast<uint64_t>(UINT32_MAX) + 1u) return false;
const auto mask = BoundedSetBitMask(offset.index, active);
if (mask.IsEmpty()) return false;
indirect.material_source = InternSource(material_source);
indirect.selector_stride = static_cast<uint32_t>(offset.stride);
indirect.selector_offset = static_cast<uint32_t>(offset.offset + memory.offset);
indirect.selectors.push_back({.source = InternSource(material_source),
.stride = static_cast<uint32_t>(offset.stride),
.offset = static_cast<uint32_t>(offset.offset + memory.offset)});
indirect.key_count = Value(32u);
indirect.selector_mask = mask;
return true;
@@ -1606,7 +1607,7 @@ private:
bool MatchDescriptorTable(Inst& handle, const DescriptorSource& descriptor,
IndirectDescriptorPlan& plan,
DescriptorSource& table_source, Value& key, uint32_t& table_offset,
uint32_t& table_stride) {
uint32_t& table_stride, uint32_t& table_immediate) {
Inst* table_handle = nullptr;
for (uint32_t dword = 0; dword < handle.NumArgs(); ++dword) {
auto* read = descriptor.dwords[dword].Resolve().TryInstruction();
@@ -1630,23 +1631,24 @@ private:
(memory->kind == ResourceKind::ScalarAddress && read->Parent() != handle.Parent()) ||
(table_handle != nullptr &&
!EquivalentValue(m_program, Value(table_handle), Value(current_handle))) ||
!MatchTableOffset(read->Arg(1), current_key, offset, stride) ||
memory->offset > UINT32_MAX - offset) {
!MatchTableOffset(read->Arg(1), current_key, offset, stride)) {
return false;
}
offset += memory->offset;
if (dword == 0u) {
key = current_key;
table_offset = offset;
table_immediate = memory->offset;
table_stride = stride;
} else if (table_stride != stride || !EquivalentValue(m_program, key, current_key) ||
static_cast<uint64_t>(table_offset) + dword * sizeof(uint32_t) != offset) {
uint64_t {table_offset} + table_immediate + dword * sizeof(uint32_t) !=
uint64_t {offset} + memory->offset) {
return false;
}
plan.split_offsets |= offset != table_offset;
table_handle = current_handle;
if (handle.NumArgs() == 8u
? WorkgroupAxis(key) == UINT32_MAX && !UsesOnlyImageDescriptors(*read) :
std::ranges::any_of(read->Uses(), [](const Use& use) {
if (handle.NumArgs() == 8u) {
plan.retain_reads |= !UsesOnlyImageDescriptors(*read);
} else if (std::ranges::any_of(read->Uses(), [](const Use& use) {
return use.user->GetOpcode() != ValueOpcode::GetBufferResource ||
std::ranges::any_of(use.user->Uses(), [](const Use& consumer) {
return consumer.user->GetOpcode() != ValueOpcode::StoreBufferU32;
@@ -1669,8 +1671,9 @@ private:
DescriptorSource& material_source) {
const auto* selected = UniformizedMaterialValue(key, image);
if (selected == nullptr) return false;
const auto value = MaterialKey(selected->Arg(1), indirect);
if (indirect.selector_shift == 0u && indirect.selector_bits == UINT32_MAX) return false;
DescriptorSource::IndirectDescriptor::SelectorRead selector;
const auto value = MaterialKey(selected->Arg(1), selector);
if (selector.shift == 0u && selector.bits == UINT32_MAX) return false;
const auto* loaded = value.TryInstruction();
if (loaded == nullptr || loaded->GetOpcode() != ValueOpcode::SelectU32 ||
!EquivalentValue(m_program, selected->Arg(0), loaded->Arg(0))) return false;
@@ -1685,10 +1688,11 @@ private:
!EquivalentValue(m_program, selected->Arg(0), read->Arg(4)) ||
!MakeRuntimeTableSource(*read, material_source)) return false;
// GPU index arithmetic may wrap. Every record in the bounded V# is a candidate.
indirect.material_source = InternSource(material_source);
selector.source = InternSource(material_source);
selector.offset = memory.offset;
indirect.selectors.push_back(selector);
indirect.selector_first = Value(0u);
indirect.key_count = material_source.dwords[2];
indirect.selector_offset = memory.offset;
return true;
}
@@ -1699,17 +1703,19 @@ private:
Value key;
uint32_t table_offset = 0;
uint32_t table_stride = 0;
uint32_t table_immediate = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset,
table_stride)) return false;
table_stride, table_immediate)) return false;
DescriptorSource material_source;
DescriptorSource::IndirectDescriptor indirect;
indirect.table_offset = table_offset;
indirect.table_immediate = table_immediate;
indirect.table_stride = table_stride;
indirect.workgroup_axis = WorkgroupAxis(key);
if (indirect.workgroup_axis != UINT32_MAX) {
// The descriptor also supplies dimensions to shader arithmetic. Keep its reads;
// only the image handle is projected onto the bounded workgroup key.
plan.reads.fill(nullptr);
plan.retain_reads = true;
} else if (table_source.dword_count == 2u) {
if (table_stride != 32u) return false;
const auto* selector = key.Resolve().TryInstruction();
@@ -1729,35 +1735,39 @@ private:
if ((table_offset & 3u) != 0u ||
(bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
} else if (!MatchUniformizedBufferKey(key, handle, indirect, material_source)) {
auto* material_read = MaterialKey(key, indirect).TryInstruction();
uint32_t material_memory_index = 0;
const auto* memory = material_read != nullptr
? ScalarReadMemory(*material_read, material_memory_index) : nullptr;
if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
memory->offset > INT32_MAX || !MemoryIndexBelongsTo(material_memory_index, *material_read)) {
return false;
std::vector<const Inst*> visited;
auto selected = key;
for (;;) {
DescriptorSource::IndirectDescriptor::SelectorRead selector;
const auto* read = MaterialKey(selected, selector).TryInstruction();
uint32_t memory_index = 0;
const auto* memory = read != nullptr ? ScalarReadMemory(*read, memory_index) : nullptr;
DescriptorSource source;
if (memory == nullptr || memory->kind != ResourceKind::ScalarBuffer ||
memory->offset > INT32_MAX || !MemoryIndexBelongsTo(memory_index, *read) ||
!MakeRuntimeTableSource(*read, source)) break;
if (std::ranges::find(visited, read) != visited.end()) return false;
visited.push_back(read);
selected = MaterialOffset(read->Arg(1), selector.stride, selector.offset);
if (selected.IsEmpty()) return false;
selector.source = InternSource(source);
selector.immediate = memory->offset;
indirect.selectors.push_back(selector);
material_source = source;
}
if (!MatchMaterialOffset(material_read->Arg(1), indirect.selector_stride,
indirect.selector_offset)) {
return false;
}
const auto step = std::gcd<uint64_t>(indirect.selector_stride, uint64_t {1} << 32u);
indirect.selector_offset =
(static_cast<uint32_t>(indirect.selector_offset % step) & ~3u) + (memory->offset & ~3u);
const auto* material_handle = material_read->Arg(0).Resolve().TryInstruction();
if (material_handle == nullptr ||
material_handle->GetOpcode() != ValueOpcode::GetBufferResource ||
!MakeRuntimeTableSource(*material_read, material_source)) {
return false;
}
indirect.material_source = InternSource(material_source);
if (indirect.selectors.empty()) return false;
std::ranges::reverse(indirect.selectors);
}
// This plan combines shader byte-offset additions with the scalar immediate.
// A nonzero combined offset is exact only when neither can cross U32 wrap.
const auto maximum_key = (UINT32_MAX >> indirect.selector_shift) & indirect.selector_bits;
if (indirect.workgroup_axis == UINT32_MAX && table_source.dword_count == 4u &&
table_offset != 0u && uint64_t {maximum_key} * table_stride + table_offset + 28u > UINT32_MAX)
return false;
// Separate shader additions may only be merged into a contiguous descriptor
// when the selector cannot wrap between those additions.
if (plan.split_offsets && indirect.workgroup_axis == UINT32_MAX &&
table_source.dword_count == 4u) {
const auto& selector = indirect.selectors.back();
const auto maximum_key = (UINT32_MAX >> selector.shift) & selector.bits;
if (uint64_t {maximum_key} * table_stride + table_offset + table_immediate + 28u >
UINT32_MAX) return false;
}
if (plan.retain_reads) plan.reads.fill(nullptr);
indirect.table_source = InternSource(table_source);
DescriptorSource image_source;
image_source.dword_count = 8u;
@@ -1791,9 +1801,10 @@ private:
Value key;
uint32_t table_offset = 0;
uint32_t table_stride = 0;
uint32_t table_immediate = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset,
table_stride) || table_stride != 16u ||
table_source.dword_count != 4u || table_offset != 0u) return false;
table_stride, table_immediate) || table_stride != 16u ||
table_source.dword_count != 4u || table_offset != 0u || table_immediate != 0u) return false;
const auto* lane_read = key.Resolve().TryInstruction();
if (lane_read == nullptr || lane_read->GetOpcode() != ValueOpcode::ReadFirstLane)
return false;
@@ -1825,9 +1836,9 @@ private:
indirect.table_stride = table_stride;
if (!MakeRuntimeTableSource(*read, material_source) ||
!BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false;
indirect.material_source = InternSource(material_source);
indirect.selectors.push_back({.source = InternSource(material_source),
.offset = memory.offset + lane * 4u});
indirect.table_source = InternSource(table_source);
indirect.selector_offset = memory.offset + lane * 4u;
DescriptorSource source;
source.dword_count = 4u;
source.dwords.fill(Value(0u));