shader: replace CPU selector chains with direct image-table lookup

This commit is contained in:
Claxten
2026-10-06 04:25:06 +04:00
committed by nmzik
parent c23de62184
commit cd06601022
8 changed files with 467 additions and 637 deletions
@@ -249,8 +249,24 @@ 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();
const auto* selector = indirect.selector ? &*indirect.selector : nullptr;
auto& keys = program.material_keys;
const auto children_begin = descriptors.size();
const auto root_resource = specializations[resource_index];
const auto intern_candidate = [&](const DescriptorValue& candidate, uint32_t& ordinal) {
ordinal = 0u;
if (candidate == descriptors[resource_index]) return true;
const auto found = std::find(descriptors.begin() + children_begin, descriptors.end(), candidate);
ordinal = static_cast<uint32_t>(found - descriptors.begin() - children_begin + 1u);
if (found == descriptors.end()) {
if (descriptors.size() >= maximum_resources) return false;
descriptors.push_back(candidate);
auto child = root_resource;
child.indirect_root = resource_index;
specializations.push_back(child);
}
return true;
};
const auto read_keys = [&](const ShaderBufferResource& material, uint64_t first,
uint64_t step, uint64_t count) {
keys.resize(count);
@@ -284,6 +300,48 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
} else {
return false;
}
if (selector == nullptr && table_value.dword_count == 4u &&
indirect.workgroup_axis == UINT32_MAX) {
if (!indirect.table_scalar && table.Format() == Prospero::BufferFormat::kInvalid) {
descriptors[resource_index] = {.dword_count = dword_count};
return true;
}
if (table_size > UINT32_MAX ||
(indirect.table_record_bytes != 0u &&
(table.Type() != 0u || table.SwizzleEnabled() || table.AddTid() ||
table.OutOfBounds() != 0u || (table_base & 3u) != 0u ||
(table.Stride() & 3u) != 0u || table.Stride() < indirect.table_record_bytes)))
return false;
// Descriptor bytes are refreshed together. The GPU selects the DWORD address;
// no material selector or unused table field is interpreted on the CPU.
keys.resize(table_size / 4u);
if (!ReadScalarTable(table_base, table_size, 0u, runtime, keys)) return false;
const auto mapping_offset = snapshot.flattened_srt.size();
descriptors[resource_index] = {.dword_count = dword_count};
snapshot.flattened_srt.resize(mapping_offset + 2u, 0u);
for (uint32_t word = 0; word < keys.size(); ++word) {
DescriptorValue candidate {.dword_count = dword_count};
std::copy_n(keys.begin() + word, std::min<size_t>(dword_count, keys.size() - word),
candidate.dwords.begin());
if (!normalize(candidate)) return false;
uint32_t ordinal;
if (!intern_candidate(candidate, ordinal)) return false;
if (ordinal == 0u) continue;
snapshot.flattened_srt.resize(mapping_offset + 2u + word, 0u);
snapshot.flattened_srt[mapping_offset + 1u + word] = ordinal;
}
if (descriptors.size() == children_begin) {
snapshot.flattened_srt.resize(mapping_offset);
} else {
snapshot.flattened_srt[mapping_offset] =
static_cast<uint32_t>(snapshot.flattened_srt.size() - mapping_offset - 1u);
auto& root = specializations[resource_index];
root.indirect_root = resource_index;
root.indirect_mapping_offset = static_cast<uint32_t>(mapping_offset);
root.indirect_search_iterations = 0u;
}
return true;
}
if (selector == nullptr) {
uint32_t key_count = 0;
if (indirect.workgroup_axis != UINT32_MAX) {
@@ -307,8 +365,6 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
ShaderBufferResource material;
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) || material.Type() != 0u ||
((selector->shift != 0u || selector->bits != UINT32_MAX) &&
(material.Base48() & 3u) != 0u) ||
material.SwizzleEnabled() ||
material.AddTid() || material.OutOfBounds() != 0u ||
uint64_t {selector->offset} + 4u > material.Stride() ||
@@ -345,53 +401,12 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
keys.push_back(key);
mask &= mask - 1u;
}
} else {
ShaderBufferResource material;
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u) {
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>(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;
if (probe_count > MaxIndirectDescriptorProbes) {
return false;
}
if (!read_keys(material, first, step, probe_count)) return false;
if (selector->shift == 0u) keys.push_back(0u);
}
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());
}
} else return false;
std::ranges::sort(keys);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
const auto children_begin = descriptors.size();
const auto mapping_offset = snapshot.flattened_srt.size();
const auto root_resource = specializations[resource_index];
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);
@@ -414,19 +429,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
uint32_t ordinal = 0;
if (entry == 0) {
descriptors[resource_index] = candidate;
} else if (descriptors[resource_index] != candidate) {
const auto found =
std::find(descriptors.begin() + children_begin, descriptors.end(), candidate);
ordinal = static_cast<uint32_t>(found - descriptors.begin() - children_begin + 1u);
if (found == descriptors.end()) {
if (descriptors.size() >= maximum_resources) {
return false;
}
descriptors.push_back(candidate);
auto child = root_resource;
child.indirect_root = resource_index;
specializations.push_back(child);
}
} else if (!intern_candidate(candidate, ordinal)) {
return false;
}
snapshot.flattened_srt[mapping_offset + 1u + entry * 2u] = key;
snapshot.flattened_srt[mapping_offset + 2u + entry * 2u] = ordinal;
@@ -584,9 +588,9 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
const auto key_count = root.indirect_mapping_offset < snapshot.flattened_srt.size()
? snapshot.flattened_srt[root.indirect_mapping_offset]
: 0u;
if (root.indirect_search_iterations == 0u || key_count < 2u ||
if (key_count == 0u ||
static_cast<size_t>(root.indirect_mapping_offset) + 1u +
static_cast<size_t>(key_count) * 2u >
static_cast<size_t>(key_count) * (root.indirect_search_iterations == 0u ? 1u : 2u) >
snapshot.flattened_srt.size()) {
return SpecializationFail("indirect image specialization has an invalid key mapping");
}
@@ -1028,9 +1032,9 @@ 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.selectors.empty() || !indirect.sources.empty();
for (const auto& read: indirect.selectors)
MarkCleanFlatSlots(plan, Source(plan, read.source), plan.clean_flat_slots);
capture_indirect_reads = true;
if (indirect.selector)
MarkCleanFlatSlots(plan, Source(plan, indirect.selector->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);