Lower finite GPU-selected image descriptors through shared indirect bindings

This commit is contained in:
nmzik
2026-10-02 15:37:17 +04:00
parent 82a2f4ad42
commit 7804eef9b6
6 changed files with 482 additions and 95 deletions
@@ -617,7 +617,6 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
}
}
IR::TrackResources(ir, decoded, native_cfg);
IR::EliminateDeadCode(ir.blocks);
TranslateResult result;
result.program = std::move(ir);
if (options.dump_ir) {
@@ -476,6 +476,7 @@ struct DescriptorSource {
uint32_t table_offset = 0;
Value key_count;
Value selector_mask;
std::vector<uint32_t> sources;
bool operator==(const IndirectImage& other) const = default;
};
@@ -227,100 +227,114 @@ bool ReadScalarTable(uint64_t base, uint64_t size, uint64_t dynamic_offset,
bool MaterializeIndirectImage(const ResourcePlan& program,
const DescriptorSource::IndirectImage& indirect,
const DescriptorValue& material_value,
const DescriptorValue& table_value, uint32_t image_index,
uint32_t image_index,
const SrtRuntime& runtime, SrtWalker& clean,
ResourceSnapshot& snapshot,
ResourceSpecialization& specialization) {
const auto& sources = indirect.sources;
auto& keys = program.material_keys;
uint64_t table_base = 0;
uint64_t table_size = UINT64_MAX; // Scalar addresses have no buffer descriptor bounds.
ShaderBufferResource table;
if (table_value.dword_count == 2u) {
table_base = (static_cast<uint64_t>(table_value.dwords[1]) << 32u) | table_value.dwords[0];
} else if (DecodeBufferDescriptor(table_value, table)) {
table_base = table.Base48();
table_size = table.GetSize();
} else {
return false;
}
auto& keys = program.material_keys;
keys.clear();
if (indirect.material_source == UINT32_MAX) {
uint32_t key_count = 0;
const bool evaluated = clean.Evaluate(indirect.key_count, key_count);
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
if (table_value.dword_count != 2u || !evaluated ||
key_count > MaxIndirectImageProbes ||
uint64_t {indirect.table_offset} + uint64_t {key_count} * 32u > UINT32_MAX + 1ull) {
if (sources.empty()) {
keys.clear();
DescriptorValue material_value;
DescriptorValue table_value;
if ((indirect.material_source != UINT32_MAX &&
!clean.EvaluateDescriptor(indirect.material_source, material_value)) ||
!clean.EvaluateDescriptor(indirect.table_source, table_value)) {
return false;
}
keys.resize(key_count);
std::iota(keys.begin(), keys.end(), 0u);
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
uint32_t count = 0;
if (material_value.dword_count != 2u || table_value.dword_count != 2u ||
!clean.Evaluate(indirect.selector_mask, mask) ||
!clean.Evaluate(indirect.key_count, count) || count == 0u || count > 32u) {
ShaderBufferResource table;
if (table_value.dword_count == 2u) {
table_base = (static_cast<uint64_t>(table_value.dwords[1]) << 32u) | table_value.dwords[0];
} else if (DecodeBufferDescriptor(table_value, table)) {
table_base = table.Base48();
table_size = table.GetSize();
} else {
return false;
}
if (count < 32u) mask &= (1u << count) - 1u;
const auto material_base =
(static_cast<uint64_t>(material_value.dwords[1]) << 32u) | material_value.dwords[0];
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;
if (offset > UINT32_MAX) return false;
uint32_t key = 0;
if (!ReadScalarTable(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
runtime, {&key, 1})) return false;
keys.push_back(key);
mask &= mask - 1u;
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
} else {
ShaderBufferResource material;
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
material.Stride() != indirect.selector_stride) {
return false;
}
// 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;
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;
if (probe_count > MaxIndirectImageProbes) {
return false;
}
keys.reserve(static_cast<size_t>(probe_count) + 1u);
keys.push_back(0u);
for (uint64_t probe = 0, offset = first; probe < probe_count; ++probe, offset += step) {
uint32_t key = 0;
if (!ReadScalarTable(material.Base48(), size, offset, runtime, {&key, 1})) {
if (indirect.material_source == UINT32_MAX) {
uint32_t key_count = 0;
const bool evaluated = clean.Evaluate(indirect.key_count, key_count);
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
if (table_value.dword_count != 2u || !evaluated ||
key_count > MaxIndirectImageProbes ||
uint64_t {indirect.table_offset} + uint64_t {key_count} * 32u > UINT32_MAX + 1ull) {
return false;
}
keys.push_back(key);
keys.resize(key_count);
std::iota(keys.begin(), keys.end(), 0u);
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
uint32_t count = 0;
if (material_value.dword_count != 2u || table_value.dword_count != 2u ||
!clean.Evaluate(indirect.selector_mask, mask) ||
!clean.Evaluate(indirect.key_count, count) || count == 0u || count > 32u) {
return false;
}
if (count < 32u) mask &= (1u << count) - 1u;
const auto material_base =
(static_cast<uint64_t>(material_value.dwords[1]) << 32u) | material_value.dwords[0];
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;
if (offset > UINT32_MAX) return false;
uint32_t key = 0;
if (!ReadScalarTable(material_base, UINT64_MAX, static_cast<uint32_t>(offset),
runtime, {&key, 1})) return false;
keys.push_back(key);
mask &= mask - 1u;
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
} else {
ShaderBufferResource material;
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
material.Stride() != indirect.selector_stride) {
return false;
}
// 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;
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;
if (probe_count > MaxIndirectImageProbes) {
return false;
}
keys.reserve(static_cast<size_t>(probe_count) + 1u);
keys.push_back(0u);
for (uint64_t probe = 0, offset = first; probe < probe_count; ++probe, offset += step) {
uint32_t key = 0;
if (!ReadScalarTable(material.Base48(), size, offset, runtime, {&key, 1})) {
return false;
}
keys.push_back(key);
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
const auto children_begin = snapshot.images.size();
const auto mapping_offset = snapshot.flattened_srt.size();
const auto root_image = specialization.images[image_index];
snapshot.flattened_srt.resize(mapping_offset + 1u + keys.size() * 2u);
snapshot.flattened_srt[mapping_offset] = static_cast<uint32_t>(keys.size());
for (uint32_t entry = 0; entry < keys.size(); ++entry) {
const auto key = keys[entry];
const auto key_count = sources.empty() ? keys.size() : sources.size();
snapshot.flattened_srt.resize(mapping_offset + 1u + key_count * 2u);
snapshot.flattened_srt[mapping_offset] = static_cast<uint32_t>(key_count);
for (uint32_t entry = 0; entry < key_count; ++entry) {
const auto key = sources.empty() ? keys[entry] : entry;
DescriptorValue candidate;
candidate.dword_count = 8u;
const auto table_offset = (key << 5u) + indirect.table_offset;
if (!ReadScalarTable(table_base, table_size, table_offset, runtime, candidate.dwords)) {
if (sources.empty()) {
const auto table_offset = (key << 5u) + indirect.table_offset;
if (!ReadScalarTable(table_base, table_size, table_offset, runtime, candidate.dwords)) {
return false;
}
} else if (!clean.EvaluateDescriptor(sources[entry], candidate)) {
return false;
}
if (NullImageDescriptor(candidate) ||
@@ -353,7 +367,7 @@ bool MaterializeIndirectImage(const ResourcePlan& program,
auto& root = specialization.images[image_index];
root.indirect_root = image_index;
root.indirect_mapping_offset = static_cast<uint32_t>(mapping_offset);
root.indirect_search_iterations = std::bit_width(keys.size());
root.indirect_search_iterations = std::bit_width(key_count);
}
return true;
}
@@ -936,20 +950,27 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
plan.uniform_fill.values[i] = Clone(plan.uniform_fill.values[i]);
}
plan.clean_flat_slots.resize(plan.srt_reads.size());
bool masked_image = false;
bool capture_image_reads = false;
for (const auto& image: plan.info.images) {
const auto* source = Source(plan, image.source);
if (source == nullptr || !source->indirect_image.has_value()) {
continue;
}
plan.requires_specialization_memory = true;
masked_image |= !source->indirect_image->selector_mask.IsEmpty();
capture_image_reads |= !source->indirect_image->selector_mask.IsEmpty() ||
!source->indirect_image->sources.empty();
MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->material_source),
plan.clean_flat_slots, source->indirect_image->selector_mask);
MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->table_source),
plan.clean_flat_slots);
if (source->indirect_image->sources.empty()) {
MarkCleanFlatSlots(plan, Source(plan, source->indirect_image->table_source),
plan.clean_flat_slots);
} else {
for (const auto candidate: source->indirect_image->sources) {
MarkCleanFlatSlots(plan, Source(plan, candidate), plan.clean_flat_slots);
}
}
}
plan.capture_specialization_reads = masked_image || !plan.control_flow.empty();
plan.capture_specialization_reads = capture_image_reads || !plan.control_flow.empty();
// Writable descriptor addresses must come from clean backing for the alias proof.
for (const auto& buffer: plan.info.buffers) {
if (buffer.written) MarkCleanFlatSlots(plan, Source(plan, buffer.source), plan.clean_flat_slots);
@@ -957,7 +978,7 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
for (const auto& image: plan.info.images) {
if (image.written) MarkCleanFlatSlots(plan, Source(plan, image.source), plan.clean_flat_slots);
}
if (masked_image) plan.resource_tracking_complete &= !program.has_address_writes;
if (capture_image_reads) plan.resource_tracking_complete &= !program.has_address_writes;
return plan;
}
@@ -1039,13 +1060,7 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
continue;
}
const auto& indirect = *source->indirect_image;
DescriptorValue material;
DescriptorValue table;
if ((indirect.material_source != UINT32_MAX &&
!clean.EvaluateDescriptor(indirect.material_source, material)) ||
!clean.EvaluateDescriptor(indirect.table_source, table) ||
!MaterializeIndirectImage(program, indirect, material, table, i, observed, clean, snapshot,
specialization)) {
if (!MaterializeIndirectImage(program, indirect, i, observed, clean, snapshot, specialization)) {
return false;
}
} else {
@@ -321,8 +321,9 @@ public:
for (uint32_t dword = 5u; dword < plan.roots.size(); dword++) {
plan.handle->SetArg(dword, plan.key);
}
for (const auto index: plan.memory) {
m_program.memory_info[index].planning_only = true;
if (plan.reads[0] != nullptr) {
for (const auto index: plan.memory)
m_program.memory_info[index].planning_only = true;
}
}
m_program.descriptor_sources = std::move(m_sources);
@@ -762,7 +763,7 @@ private:
const auto& b = *descriptor.indirect_image;
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.table_offset != b.table_offset ||
a.table_offset != b.table_offset || a.sources != b.sources ||
!EquivalentValue(m_program, a.key_count, b.key_count) ||
a.selector_mask.IsEmpty() != b.selector_mask.IsEmpty() ||
(!a.selector_mask.IsEmpty() &&
@@ -1474,6 +1475,152 @@ private:
return true;
}
uint32_t SelectorMaximum(Value value) const {
value = value.Resolve();
if (value.IsImmediate() && value.GetType() == Type::U32) return value.U32();
const auto* inst = value.TryInstruction();
if (inst == nullptr) return UINT32_MAX;
if (inst->GetOpcode() == ValueOpcode::FindILsb32)
return NonzeroOnEntry(inst->Arg(0), inst->Parent()) ? 31u : UINT32_MAX;
if (inst->GetOpcode() == ValueOpcode::UMin32)
return std::min(SelectorMaximum(inst->Arg(0)), SelectorMaximum(inst->Arg(1)));
uint32_t shift;
if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 &&
ImmediateU32(inst->Arg(1), shift) && shift < 32u) {
const auto maximum = uint64_t {SelectorMaximum(inst->Arg(0))} << shift;
if (maximum <= UINT32_MAX) return static_cast<uint32_t>(maximum);
}
return UINT32_MAX;
}
bool ImpossibleSwitchEdge(const Block* from, const Block* to) const {
const auto from_it = std::ranges::find(m_program.blocks, from);
const auto to_it = std::ranges::find(m_program.blocks, to);
if (from_it == m_program.blocks.end() || to_it == m_program.blocks.end() ||
m_program.block_info.size() != m_program.blocks.size()) return false;
const auto& info = m_program.block_info[from_it - m_program.blocks.begin()];
const auto& term = info.terminator;
if (term.kind != CFG::TerminatorKind::IndirectBranch ||
term.indirect_selector_code == UINT32_MAX ||
term.indirect_selector_values.empty() ||
term.indirect_selector_values.size() != term.indirect_selector_targets.size()) return false;
const auto target = m_program.block_info[to_it - m_program.blocks.begin()].id;
const auto maximum = SelectorMaximum(info.indirect_target);
bool found = false;
for (size_t i = 0; i < term.indirect_selector_values.size(); ++i) {
if (term.indirect_selector_targets[i] != target) continue;
if (term.indirect_selector_values[i] <= maximum) return false;
found = true;
}
return found;
}
bool TryMakeFiniteImage(Inst& handle, IndirectImagePlan& plan) {
if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u)
return false;
if (std::ranges::any_of(handle.Uses(), [](const Use& use) {
const auto op = use.user->GetOpcode();
return op != ValueOpcode::ImageSampleRaw && op != ValueOpcode::ImageGatherRaw;
})) return false;
bool has_phi = false;
for (size_t word = 0; word < handle.NumArgs(); ++word) has_phi |= handle.Arg(word).Resolve().IsPhi();
if (!has_phi) return false;
// Normalize the eight synchronized descriptor words to one GPU ordinal.
// Build the complete graph before changing IR so an unsupported leaf is transactional.
struct Choice {
DescriptorSource descriptor;
const Inst* branch = nullptr;
std::vector<uint32_t> children;
Value key;
};
std::vector<Choice> choices;
const auto visit = [&](const auto& self, DescriptorSource descriptor) -> uint32_t {
for (auto& word: descriptor.dwords) word = word.Resolve();
for (uint32_t i = 0; i < choices.size(); ++i)
if (choices[i].descriptor.dwords == descriptor.dwords) return i;
const auto index = static_cast<uint32_t>(choices.size());
choices.push_back({descriptor});
uint32_t bad = 0;
if (ValidateSource(descriptor, bad)) return index;
const auto* branch = descriptor.dwords[bad].TryInstruction();
if (branch == nullptr || branch->Parent() == nullptr ||
branch->GetOpcode() != ValueOpcode::Phi ||
branch->NumPhiBlocks() != branch->NumArgs()) return UINT32_MAX;
choices[index].branch = branch;
const auto count = branch->NumArgs();
for (uint32_t arm = 0; arm < count; ++arm) {
// A nonzero bit scan cannot choose its zero-input switch sentinel.
// Preserve the Phi edge; its ordinal is never observed at runtime.
if (ImpossibleSwitchEdge(branch->PhiBlock(arm), branch->Parent())) {
choices[index].children.push_back(UINT32_MAX);
continue;
}
auto child = descriptor;
for (uint32_t word = 0; word < child.dword_count; ++word) {
const auto* selected = descriptor.dwords[word].TryInstruction();
if (selected == nullptr || selected->GetOpcode() != branch->GetOpcode() ||
selected->Parent() != branch->Parent()) {
if (!ValidateRuntimeValue(m_program, descriptor.dwords[word])) return UINT32_MAX;
continue;
}
if (selected->NumArgs() != count || selected->NumPhiBlocks() != count)
return UINT32_MAX;
uint32_t incoming = 0;
while (incoming < count && selected->PhiBlock(incoming) != branch->PhiBlock(arm))
++incoming;
if (incoming == count) return UINT32_MAX;
child.dwords[word] = selected->Arg(incoming);
}
const auto next = self(self, child);
if (next == UINT32_MAX) return UINT32_MAX;
choices[index].children.push_back(next);
}
return index;
};
DescriptorSource root;
root.dword_count = 8u;
for (uint32_t word = 0; word < root.dword_count; ++word) root.dwords[word] = handle.Arg(word);
if (visit(visit, root) == UINT32_MAX || choices.front().branch == nullptr) return false;
DescriptorSource image_source;
image_source.dword_count = 8u;
image_source.dwords.fill(Value(0u));
image_source.indirect_image.emplace();
auto& sources = image_source.indirect_image->sources;
for (auto& choice: choices) {
if (choice.branch != nullptr) continue;
const auto source = InternSource(choice.descriptor);
auto found = std::ranges::find(sources, source);
choice.key = Value(static_cast<uint32_t>(found - sources.begin()));
if (found == sources.end()) sources.push_back(source);
}
if (sources.empty()) return false;
for (auto& choice: choices) {
if (choice.branch == nullptr) continue;
auto* block = choice.branch->Parent();
const auto where = std::ranges::find_if(block->Instructions(), [&](const Inst& inst) {
return &inst == choice.branch;
});
auto& key = *block->PrependNewInst(where, choice.branch->GetOpcode());
key.SetFlags(Type::U32);
choice.key = Value(&key);
}
for (const auto& choice: choices) {
if (choice.branch == nullptr) continue;
auto* key = choice.key.Instruction();
for (uint32_t arm = 0; arm < choice.children.size(); ++arm) {
const auto child = choice.children[arm];
const auto value = child == UINT32_MAX ? Value(0u) : choices[child].key;
key->AddPhiOperand(choice.branch->PhiBlock(arm), value);
}
}
plan.handle = &handle;
plan.source = InternSource(image_source);
plan.key = choices.front().key;
plan.roots = image_source.dwords;
plan.reads.fill(nullptr);
return true;
}
const IndirectImagePlan* FindIndirectImage(const Inst& handle) const {
const auto found =
std::find_if(m_indirect_images.begin(), m_indirect_images.end(),
@@ -1486,7 +1633,8 @@ private:
bool IsIndirectPlanningMemory(uint32_t index) const {
return std::any_of(m_indirect_images.begin(), m_indirect_images.end(),
[&](const IndirectImagePlan& plan) {
return std::ranges::find(plan.memory, index) != plan.memory.end();
return plan.reads[0] != nullptr &&
std::ranges::find(plan.memory, index) != plan.memory.end();
});
}
@@ -1502,7 +1650,7 @@ private:
continue;
}
IndirectImagePlan plan;
if (TryMakeIndirectImage(*handle, plan)) {
if (TryMakeIndirectImage(*handle, plan) || TryMakeFiniteImage(*handle, plan)) {
m_indirect_images.push_back(std::move(plan));
}
}