shader: lower bounded UInt formatted buffer copies to shared stores

This commit is contained in:
nmzik
2026-10-12 21:24:00 +02:00
parent a7a734b5d3
commit 5dc6914688
5 changed files with 300 additions and 10 deletions
@@ -521,6 +521,7 @@ struct DescriptorSource {
uint32_t table_record_bytes = 0;
bool table_scalar = false;
bool indexed_store = false;
bool signed_count = false;
uint32_t workgroup_axis = UINT32_MAX;
Value key_count;
// First selector index, or first descriptor key for a direct or bitmap range.
@@ -190,7 +190,7 @@ bool ReadScalarTable(uint64_t base, uint64_t size, uint64_t dynamic_offset,
runtime.read_specialization_memory(runtime.userdata, address, prefix);
}
bool NormalizeIndirectStoreBuffer(DescriptorValue& value, bool indexed) {
bool NormalizeIndirectStoreBuffer(DescriptorValue& value, bool indexed, bool formatted) {
ShaderBufferResource descriptor;
if (!DecodeBufferDescriptor(value, descriptor)) return false;
if (descriptor.Type() != 0u || descriptor.Format() == Prospero::BufferFormat::kInvalid ||
@@ -199,6 +199,10 @@ bool NormalizeIndirectStoreBuffer(DescriptorValue& value, bool indexed) {
} else if (descriptor.Stride() != (indexed ? 4u : 1u) || descriptor.SwizzleEnabled() ||
descriptor.AddTid() || descriptor.OutOfBounds() != (indexed ? 0u : 3u)) {
return SpecializationFail("indirect buffer stores require exact linear DWORD or byte bounds");
} else if (formatted && (descriptor.Format() != Prospero::BufferFormat::k32UInt ||
descriptor.DstSelXYZW() != DstSel(4, 0, 0, 0) ||
(descriptor.Base48() & 3u) != 0u)) {
return SpecializationFail("indirect formatted X stores require an aligned 32-bit UInt identity mapping");
}
return true;
}
@@ -225,6 +229,11 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
auto& keys = program.material_keys;
const auto children_begin = descriptors.size();
const auto root_resource = specializations[resource_index];
const auto evaluate_count = [&](uint32_t& count) {
if (!clean.Evaluate(indirect.key_count, count)) return false;
if (indirect.signed_count && std::bit_cast<int32_t>(count) < 0) count = 0;
return true;
};
const auto intern_candidate = [&](const DescriptorValue& candidate, uint32_t& ordinal) {
ordinal = 0u;
if (candidate == descriptors[resource_index]) return true;
@@ -321,9 +330,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
key_count = runtime.workgroup_counts[indirect.workgroup_axis];
if (key_count == 0u) return false;
} else {
if (table_value.dword_count != 2u || !clean.Evaluate(indirect.key_count, key_count))
if (table_value.dword_count != 2u || !evaluate_count(key_count))
return false;
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
}
if (!indirect.selector_first.IsEmpty() && !clean.Evaluate(indirect.selector_first, first_key))
return false;
@@ -344,7 +352,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) ||
!clean.Evaluate(indirect.selector_first, first) ||
!clean.Evaluate(indirect.key_count, count) || count != masks.size() * 32u ||
!evaluate_count(count) || count != masks.size() * 32u ||
selector->stride != sizeof(masks) || selector->offset != 0u ||
material.Type() != 0u || material.Stride() != sizeof(masks) ||
material.SwizzleEnabled() || material.AddTid() || material.OutOfBounds() != 0u ||
@@ -385,7 +393,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
uint32_t first = 0, count = 0;
if (!DecodeBufferDescriptor(material_value, material) ||
!clean.Evaluate(indirect.selector_first, first) ||
!clean.Evaluate(indirect.key_count, count) || count > MaxIndirectDescriptorProbes)
!evaluate_count(count) || count > MaxIndirectDescriptorProbes)
return false;
// SMEM uses its instruction's byte stride; V# stride only supplies its bounds.
const auto step = selector->stride != 0u ? selector->stride : material.Stride();
@@ -404,7 +412,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
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) {
!evaluate_count(count) || count == 0u || count > 32u) {
return false;
}
if (count < 32u) mask &= (1u << count) - 1u;
@@ -1152,7 +1160,8 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
program, *source->indirect_descriptor, i, 4u, observed, clean, snapshot,
snapshot.buffers, specialization.buffers, ShaderInfo::MaxBuffers,
[&](DescriptorValue& value) {
return NormalizeIndirectStoreBuffer(value, source->indirect_descriptor->indexed_store);
return NormalizeIndirectStoreBuffer(
value, source->indirect_descriptor->indexed_store, base.formatted);
}))
return false;
}
@@ -1356,6 +1365,16 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
inst.SetFlags(image_remap[inst.Flags<uint32_t>()]);
continue;
}
if (inst.GetOpcode() == ValueOpcode::StoreBufferU32) {
auto& memory = memory_info[inst.Flags<MemoryFlags>().index];
if (memory.kind == ResourceKind::Buffer && memory.formatted &&
program.descriptor_sources[buffers[memory.resource].source].indirect_descriptor) {
// Candidate validation proved an index-only k32UInt identity store.
// Reuse finite raw-store dispatch and its original-index bounds check.
memory.formatted = false;
}
continue;
}
if (BufferAccessOf(inst.GetOpcode()) == BufferAccess::Read) {
const auto& memory = memory_info[inst.Flags<MemoryFlags>().index];
if (memory.kind == ResourceKind::Buffer &&
@@ -914,6 +914,7 @@ private:
a.table_offset != b.table_offset || a.table_immediate != b.table_immediate ||
a.table_stride != b.table_stride || a.indexed_store != b.indexed_store ||
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) ||
a.selector_first.IsEmpty() != b.selector_first.IsEmpty() ||
(!a.selector_first.IsEmpty() &&
@@ -1974,6 +1975,7 @@ private:
if (!proof.SelectsRange(*selected, active, *candidate, *carried, mask, exits)) continue;
indirect.selector_first = first;
indirect.key_count = bound->Arg(1);
indirect.signed_count = bound->GetOpcode() == ValueOpcode::SLessThan32;
return true;
}
}
@@ -2347,8 +2349,10 @@ private:
if (bitscan) {
indirect.key_count = Value(32u);
} else if (!m_shader_writes) {
if (const auto* bound = BoundedLoop(key, handle.Parent()))
if (const auto* bound = BoundedLoop(key, handle.Parent())) {
indirect.key_count = bound->Arg(1);
indirect.signed_count = bound->GetOpcode() == ValueOpcode::SLessThan32;
}
}
if (indirect.key_count.IsEmpty() &&
!MatchUniformizedMaterialKey(key, handle, indirect, material_source) &&
@@ -2449,7 +2453,8 @@ private:
if (inst.GetOpcode() != ValueOpcode::StoreBufferU32) return false;
const auto& memory = m_program.memory_info[inst.Flags<MemoryFlags>().index];
uint32_t scalar = 1u;
if (memory.kind != ResourceKind::Buffer || memory.typed || memory.formatted ||
if (memory.kind != ResourceKind::Buffer || memory.typed ||
(memory.formatted && (!indexed || memory.data_bits != 32u || memory.data_dwords != 1u)) ||
memory.idxen != indexed || memory.offen == indexed ||
(indexed && (memory.offset != 0u || inst.Arg(2).Resolve() != Value(0u))) ||
!ImmediateU32(inst.Arg(3), scalar) || scalar != 0u)
@@ -2479,6 +2484,13 @@ private:
return true;
};
if (table_source.dword_count == 2u && MatchPackedScalarKey(key, indirect)) return finish();
if (table_source.dword_count == 2u) {
const auto* bound = BoundedLoop(key, handle.Parent());
if (bound == nullptr) return false;
indirect.key_count = bound->Arg(1);
indirect.signed_count = bound->GetOpcode() == ValueOpcode::SLessThan32;
return finish();
}
if (table_source.dword_count != 4u || table_offset != 0u || table_immediate != 0u || indexed)
return false;
const auto* lane_read = key.Resolve().TryInstruction();