shader: share bounded material-key scans across scalar and vector loads

This commit is contained in:
nmzik
2026-10-04 07:15:51 +04:00
parent 919c058fee
commit b183cb1011
3 changed files with 118 additions and 94 deletions
@@ -249,6 +249,19 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
const auto descriptor_bytes = dword_count * sizeof(uint32_t);
const auto& sources = indirect.sources;
auto& keys = program.material_keys;
const auto read_keys = [&](const ShaderBufferResource& material, uint64_t first,
uint64_t step, uint64_t count) {
keys.resize(count);
if (step == 4u) {
return ReadScalarTable(material.Base48(), material.GetSize(), first, runtime, keys);
}
for (auto& key: keys) {
if (!ReadScalarTable(material.Base48(), material.GetSize(), first, runtime, {&key, 1}))
return false;
first += step;
}
return true;
};
uint64_t table_base = 0;
uint64_t table_size = UINT64_MAX; // Scalar addresses have no buffer descriptor bounds.
if (sources.empty()) {
@@ -297,20 +310,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
indirect.selector_offset + 4u >
uint64_t {UINT32_MAX} + 1u)
return false;
if (material.Stride() == 4u && indirect.selector_offset == 0u) {
keys.resize(count);
if (!ReadScalarTable(material.Base48(), material.GetSize(), uint64_t {first} * 4u,
runtime, keys)) return false;
} else {
keys.reserve(count);
for (uint64_t index = first; index < uint64_t {first} + count; ++index) {
uint32_t key = 0;
const auto offset = index * material.Stride() + indirect.selector_offset;
if (!ReadScalarTable(material.Base48(), material.GetSize(), offset, runtime,
{&key, 1})) return false;
keys.push_back(key);
}
}
if (!read_keys(material, uint64_t {first} * material.Stride() + indirect.selector_offset,
material.Stride(), count)) return false;
} else if (!indirect.selector_mask.IsEmpty()) {
uint32_t mask = 0;
uint32_t count = 0;
@@ -336,8 +337,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
}
} else {
ShaderBufferResource material;
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u ||
material.Stride() != indirect.selector_stride) {
if (!DecodeBufferDescriptor(material_value, material) || table_value.dword_count != 4u) {
return false;
}
// The first aligned offset includes the immediate added after shader U32 arithmetic.
@@ -350,15 +350,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
if (probe_count > MaxIndirectDescriptorProbes) {
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);
}
if (!read_keys(material, first, step, probe_count)) return false;
if (indirect.selector_shift == 0u) keys.push_back(0u);
}
if (indirect.material_source != UINT32_MAX) {
if (indirect.selector_shift != 0u) {
@@ -1011,8 +1004,7 @@ 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.selector_mask.IsEmpty() ||
!indirect.selector_first.IsEmpty() || !indirect.sources.empty();
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);
MarkCleanFlatSlots(plan, nullptr, plan.clean_flat_slots, indirect.selector_first);
@@ -923,38 +923,39 @@ private:
return ValidateSource(descriptor, bad_dword);
}
bool MatchMaterialOffset(Value value, Value& selector, uint32_t& stride,
uint32_t& offset) const {
value = value.Resolve();
offset = 0;
auto* candidate = value.TryInstruction();
if (candidate != nullptr && candidate->GetOpcode() == ValueOpcode::IAdd32 &&
candidate->NumArgs() == 2u) {
uint32_t immediate = 0;
if (ImmediateU32(candidate->Arg(0), immediate)) {
value = candidate->Arg(1).Resolve();
} else if (ImmediateU32(candidate->Arg(1), immediate)) {
value = candidate->Arg(0).Resolve();
bool MatchMaterialOffset(Value value, uint32_t& stride, uint32_t& offset) const {
stride = 1u;
offset = 0u;
value = value.Resolve();
while (const auto* inst = value.TryInstruction()) {
if (inst->NumArgs() != 2u) break;
uint32_t constant = 0;
Value remaining;
if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 &&
ImmediateU32(inst->Arg(1), constant) && constant < 32u) {
stride *= 1u << constant;
remaining = inst->Arg(0);
} else if (inst->GetOpcode() == ValueOpcode::IAdd32 ||
inst->GetOpcode() == ValueOpcode::IMul32) {
if (ImmediateU32(inst->Arg(0), constant)) remaining = inst->Arg(1);
else if (ImmediateU32(inst->Arg(1), constant)) remaining = inst->Arg(0);
else break;
if (inst->GetOpcode() == ValueOpcode::IAdd32) offset += stride * constant;
else stride *= constant;
} else {
return false;
break;
}
offset = immediate;
value = remaining.Resolve();
}
const auto* multiply = value.TryInstruction();
if (multiply == nullptr || multiply->GetOpcode() != ValueOpcode::IMul32 ||
multiply->NumArgs() != 2u) {
return false;
}
if (ImmediateU32(multiply->Arg(0), stride)) {
selector = multiply->Arg(1).Resolve();
} else if (ImmediateU32(multiply->Arg(1), stride)) {
selector = multiply->Arg(0).Resolve();
} else {
return false;
}
const auto* selector_inst = selector.TryInstruction();
return stride != 0u && selector_inst != nullptr &&
selector_inst->GetOpcode() == ValueOpcode::ReadFirstLane;
// Treat the remaining GPU-derived scalar as any U32 value.
return stride != 0u && value.GetType() == Type::U32 && !value.IsImmediate();
}
const Inst* MaterialKeyShift(Value key, uint32_t& amount) const {
const auto* shift = key.Resolve().TryInstruction();
return shift != nullptr && shift->GetOpcode() == ValueOpcode::ShiftRightLogical32 &&
ImmediateU32(shift->Arg(1), amount) && amount != 0u && amount < 32u
? shift : nullptr;
}
enum class LaneQuantifier { Any, All };
@@ -1558,10 +1559,9 @@ private:
DescriptorSource& material_source) {
const auto* selected = UniformizedMaterialValue(key, image);
if (selected == nullptr) return false;
const auto* shift = selected->Arg(1).Resolve().TryInstruction();
uint32_t amount = 0;
if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftRightLogical32 ||
!ImmediateU32(shift->Arg(1), amount) || amount == 0u || amount >= 32u) return false;
const auto* shift = MaterialKeyShift(selected->Arg(1), amount);
if (shift == nullptr) return false;
const auto* loaded = shift->Arg(0).Resolve().TryInstruction();
if (loaded == nullptr || loaded->GetOpcode() != ValueOpcode::SelectU32 ||
!EquivalentValue(m_program, selected->Arg(0), loaded->Arg(0))) return false;
@@ -1611,7 +1611,9 @@ 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 = key.Resolve().TryInstruction();
const auto* key_shift = MaterialKeyShift(key, indirect.selector_shift);
auto* material_read =
(key_shift != nullptr ? key_shift->Arg(0) : key).Resolve().TryInstruction();
uint32_t material_memory_index = 0;
const auto* memory = material_read != nullptr
? ScalarReadMemory(*material_read, material_memory_index) : nullptr;
@@ -1619,8 +1621,7 @@ private:
memory->offset > INT32_MAX || !MemoryIndexBelongsTo(material_memory_index, *material_read)) {
return false;
}
Value selector;
if (!MatchMaterialOffset(material_read->Arg(1), selector, indirect.selector_stride,
if (!MatchMaterialOffset(material_read->Arg(1), indirect.selector_stride,
indirect.selector_offset)) {
return false;
}
@@ -1628,8 +1629,11 @@ private:
indirect.selector_offset =
(static_cast<uint32_t>(indirect.selector_offset % step) & ~3u) + (memory->offset & ~3u);
const auto* shift = plan.reads[0]->Arg(1).Resolve().TryInstruction();
const std::array<const Inst*, 1> material_users {shift};
if (!UsesOnly(*material_read, material_users) || !UsesOnly(*shift, plan.reads)) {
const std::array<const Inst*, 1> material_users {key_shift != nullptr ? key_shift : shift};
const std::array<const Inst*, 1> key_users {shift};
if (!UsesOnly(*material_read, material_users) ||
(key_shift != nullptr && !UsesOnly(*key_shift, key_users)) ||
!UsesOnly(*shift, plan.reads)) {
return false;
}
const auto* material_handle = material_read->Arg(0).Resolve().TryInstruction();