shader: preserve affine image heap candidate domains

Co-authored-by: fxpw <fxpw@outlook.com>
This commit is contained in:
nmzik
2026-10-15 03:34:07 +02:00
parent fcab4e5d0d
commit 5916a184cc
4 changed files with 116 additions and 34 deletions
@@ -410,8 +410,10 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
(table.Stride() & 3u) != 0u || table.Stride() < indirect.table_record_bytes)))
return false;
const bool affine = indirect.table_stride != 0u;
const bool image_heap = dword_count == 8u && !indirect.pointer_offset;
const uint64_t step = indirect.table_maximum ? indirect.table_stride
: affine ? uint64_t {1} << std::countr_zero(indirect.table_stride) : 4u;
: affine ? std::max(uint64_t {1} << std::countr_zero(indirect.table_stride), uint64_t {4})
: 4u;
const uint64_t first = ((indirect.table_maximum ? indirect.table_offset
: indirect.table_offset % step) & ~uint64_t {3}) +
(indirect.table_immediate & ~3u);
@@ -420,10 +422,10 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
const auto count = table_size >= first + 4u
? std::min((table_size - first - 4u) / step + 1u, domain)
: 0u;
if (affine && count > MaxIndirectDescriptorProbes) return false;
if (!image_heap && count > MaxIndirectDescriptorProbes) return false;
// Affine SMEM offsets wrap in U32 before adding the aligned immediate.
// Read only their congruence class; unrestricted image tables retain a bulk read.
if (!affine) {
// Image heaps retain one bulk snapshot and their absolute byte-offset mapping.
if (image_heap) {
keys.resize(table_size / 4u);
if (!ReadScalarTable(table_base, table_size, 0u, runtime, keys)) return false;
}
@@ -432,7 +434,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
snapshot.flattened_srt.resize(mapping_offset + 2u, 0u);
for (uint32_t word = 0; word < count; ++word) {
DescriptorValue candidate {.dword_count = dword_count};
const auto table_offset = affine || indirect.pointer_offset
const auto table_offset = affine
? first + word * step : uint64_t {word} * 4u;
if (indirect.pointer_offset) {
std::array<uint32_t, 2> address {};
@@ -440,11 +442,13 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
!ReadScalarTable((uint64_t {address[1]} << 32u) | address[0], UINT64_MAX,
*indirect.pointer_offset, runtime,
std::span(candidate.dwords).first(read_dwords))) return false;
} else if (affine) {
} else if (!image_heap) {
if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
std::span(candidate.dwords).first(read_dwords))) return false;
} else {
std::copy_n(keys.begin() + word, std::min<size_t>(read_dwords, keys.size() - word),
const auto table_word = table_offset / 4u;
std::copy_n(keys.begin() + table_word,
std::min<size_t>(read_dwords, keys.size() - table_word),
candidate.dwords.begin());
}
if (!normalize(candidate))
@@ -453,8 +457,9 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
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;
const auto mapping_word = image_heap ? table_offset / 4u : word;
snapshot.flattened_srt.resize(mapping_offset + 2u + mapping_word, 0u);
snapshot.flattened_srt[mapping_offset + 1u + mapping_word] = ordinal;
}
if (descriptors.size() == children_begin) {
snapshot.flattened_srt.resize(mapping_offset);
@@ -846,7 +851,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
if (i >= program.info.images.size() && source != nullptr &&
source->indirect_descriptor &&
IsBoundedDescriptorTable(program, *source->indirect_descriptor) &&
source->indirect_descriptor->table_stride == 0u) {
!source->indirect_descriptor->pointer_offset) {
const auto& indirect = *source->indirect_descriptor;
const size_t mapping = specialization.images[base_index].indirect_mapping_offset;
const auto& flat = snapshot.flattened_srt;
@@ -972,6 +972,7 @@ private:
uint32_t InternSource(const DescriptorSource& descriptor) {
for (uint32_t candidate = 0; candidate < m_sources.size(); candidate++) {
auto& current = m_sources[candidate];
bool shared_image_heap = false;
if (current.dword_count != descriptor.dword_count ||
current.indirect_descriptor.has_value() != descriptor.indirect_descriptor.has_value()) {
continue;
@@ -979,11 +980,16 @@ private:
if (current.indirect_descriptor.has_value()) {
const auto& a = *current.indirect_descriptor;
const auto& b = *descriptor.indirect_descriptor;
// An indexed image heap already covers every DWORD start, so its domain
// subsumes an affine scalar use of the same table and shares its snapshot.
shared_image_heap = descriptor.dword_count == 8u &&
(a.table_record_bytes != 0u || b.table_record_bytes != 0u);
if (a.selector != b.selector || a.table_maximum != b.table_maximum ||
a.pointer_offset != b.pointer_offset ||
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.indexed != b.indexed ||
(!shared_image_heap && (a.table_offset != b.table_offset ||
a.table_immediate != b.table_immediate ||
a.table_stride != b.table_stride)) || a.indexed != b.indexed ||
a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
a.signed_count != b.signed_count ||
!EquivalentValue(m_program, a.key_count, b.key_count) ||
@@ -1000,6 +1006,11 @@ private:
}
if (same) {
if (current.indirect_descriptor) {
if (shared_image_heap) {
current.indirect_descriptor->table_stride = 0u;
current.indirect_descriptor->table_offset = 0u;
current.indirect_descriptor->table_immediate = 0u;
}
current.indirect_descriptor->table_scalar |= descriptor.indirect_descriptor->table_scalar;
current.indirect_descriptor->table_record_bytes = std::max(
current.indirect_descriptor->table_record_bytes,
@@ -2920,16 +2931,19 @@ private:
}
if ((table_offset & 3u) != 0u) return false;
} else {
// A bounded V# supplies the complete image table. Leave every GPU selector
// and descriptor read in the shader; the host only translates table bytes.
// Keep the GPU selector and descriptor reads, and preserve the proven scalar
// offset domain within the bounded image table.
if (sampler || plan.split_offsets) return false;
indirect.table_stride = 0u;
indirect.table_offset = 0u;
indirect.table_immediate = 0u;
indirect.workgroup_axis = UINT32_MAX;
plan.table_read = plan.reads[0];
if (plan.table_indexed) indirect.table_record_bytes = table_immediate + plan.read_dwords * 4u;
else indirect.table_scalar = true;
if (plan.table_indexed) {
indirect.table_stride = 0u;
indirect.table_offset = 0u;
indirect.table_immediate = 0u;
indirect.table_record_bytes = table_immediate + plan.read_dwords * 4u;
} else {
indirect.table_scalar = true;
}
plan.retain_reads = true;
}