shader: track image table keys divided from a runtime base plus the workgroup ID

This commit is contained in:
Claxten
2026-10-26 16:02:19 +01:00
committed by nmzik
parent 0a140911ad
commit de3ed905c8
6 changed files with 121 additions and 20 deletions
@@ -169,6 +169,8 @@ bool EquivalentSelectorDomain(const ResourcePlan& program,
const DescriptorSource::IndirectDescriptor& right) {
return left.selector == right.selector && left.signed_count == right.signed_count &&
left.workgroup_axis == right.workgroup_axis &&
left.workgroup_multiplier == right.workgroup_multiplier &&
left.workgroup_shift == right.workgroup_shift &&
EquivalentValue(program, left.key_count, right.key_count) &&
EquivalentValue(program, left.selector_first, right.selector_first) &&
EquivalentValue(program, left.selector_mask, right.selector_mask);
@@ -159,7 +159,8 @@ KYTY_IO_FIELDS(ShaderInfo, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p, q, r,
KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead::CountedRows, a, b)
KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::SelectorRead, a, b, c, d, e, f, g, h, i, j)
KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor::PointerTable, a, b, c, d)
KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p)
KYTY_IO_FIELDS(DescriptorSource::IndirectDescriptor, a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p,
q, r)
KYTY_IO_FIELDS(DescriptorSource, a, b, c)
KYTY_IO_FIELDS(SrtRead, a, b)
KYTY_IO_FIELDS(ResourceBlock, a, b, c, d)
+5 -1
View File
@@ -606,8 +606,12 @@ struct DescriptorSource {
bool indexed = false;
bool signed_count = false;
uint32_t workgroup_axis = UINT32_MAX;
// Nonzero shift: workgroup keys are (selector_first + ID) * multiplier >> shift.
uint32_t workgroup_multiplier = 0;
uint32_t workgroup_shift = 0;
Value key_count;
// First selector index, or first descriptor key for a direct or bitmap range.
// First selector index, workgroup key base, or first descriptor key for a direct or bitmap
// range.
Value selector_first;
Value selector_mask;
std::vector<uint32_t> sources;
@@ -756,6 +756,12 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
std::iota(keys.begin(), keys.end(), first_key);
if (!keys.empty() && keys.back() < keys.front())
std::rotate(keys.begin(), keys.begin() + (UINT32_MAX - first_key + 1ull), keys.end());
if (indirect.workgroup_shift != 0u) {
for (auto& key: keys)
key = static_cast<uint32_t>(uint64_t {key} * indirect.workgroup_multiplier >>
indirect.workgroup_shift);
keys.erase(std::unique(keys.begin(), keys.end()), keys.end());
}
} else {
if (!indirect.selector_mask.IsEmpty() && table_value.dword_count != 2u) return false;
if (!selector_keys->resolved) {
@@ -1014,7 +1014,9 @@ private:
(!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.workgroup_axis != b.workgroup_axis ||
a.workgroup_multiplier != b.workgroup_multiplier ||
a.workgroup_shift != b.workgroup_shift || 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() ||
@@ -1158,9 +1160,21 @@ private:
return guarded && witnessed;
}
uint32_t WorkgroupAxis(Value key) const {
// The workgroup ID axis of key; with first, also of a runtime first + workgroup ID sum.
uint32_t WorkgroupAxis(Value key, Value* first = nullptr) const {
const auto* builtin = key.Resolve().TryInstruction();
uint32_t axis = UINT32_MAX;
if (first != nullptr && builtin != nullptr && builtin->GetOpcode() == ValueOpcode::IAdd32) {
for (uint32_t arm = 0; arm < 2u; ++arm) {
axis = WorkgroupAxis(builtin->Arg(arm));
const auto base = builtin->Arg(arm ^ 1u);
if (axis == UINT32_MAX ||
!ValidateRuntimeValue(m_program, base, RuntimeValueType::Integer)) continue;
*first = base;
return axis;
}
return UINT32_MAX;
}
if (m_program.stage != ShaderType::Compute || builtin == nullptr ||
builtin->GetOpcode() != ValueOpcode::GetBuiltin || builtin->NumArgs() != 2u ||
builtin->Arg(0) != Value(static_cast<uint32_t>(StageInputKind::WorkgroupId)) ||
@@ -1168,6 +1182,27 @@ private:
return axis;
}
// Matches key = (first + workgroup ID) * multiplier >> shift with shift >= 32; first and the
// division are optional.
bool MatchWorkgroupKey(Value key, DescriptorSource::IndirectDescriptor& indirect) const {
uint32_t bits = 0, multiplier = 0;
const auto* shift = key.Resolve().TryInstruction();
if (shift != nullptr && shift->GetOpcode() == ValueOpcode::ShiftRightLogical32) {
const auto* high = shift->Arg(0).Resolve().TryInstruction();
if (high != nullptr && high->GetOpcode() == ValueOpcode::UMulHi &&
ImmediateU32(shift->Arg(1), bits) && bits < 32u &&
ImmediateU32(high->Arg(1), multiplier)) {
key = high->Arg(0);
bits += 32u;
}
}
indirect.workgroup_axis = WorkgroupAxis(key, &indirect.selector_first);
if (indirect.workgroup_axis == UINT32_MAX) return false;
indirect.workgroup_multiplier = multiplier;
indirect.workgroup_shift = bits;
return true;
}
bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t& stride) const {
offset = 0;
for (;;) {
@@ -2304,21 +2339,7 @@ private:
((offset | stride | memory->offset) & 3u) == 0u &&
uint64_t {offset} + memory->offset <= UINT32_MAX;
if (affine) {
workgroup_axis = WorkgroupAxis(index);
if (workgroup_axis == UINT32_MAX) {
const auto* sum = index.Resolve().TryInstruction();
if (sum != nullptr && sum->GetOpcode() == ValueOpcode::IAdd32) {
for (uint32_t arm = 0; arm < 2u; ++arm) {
const auto axis = WorkgroupAxis(sum->Arg(arm));
const auto base = sum->Arg(arm ^ 1u);
if (axis == UINT32_MAX ||
!ValidateRuntimeValue(m_program, base, RuntimeValueType::Integer)) continue;
workgroup_axis = axis;
first = base;
break;
}
}
}
workgroup_axis = WorkgroupAxis(index, &first);
const auto* bound = BoundedLoop(index, read->Parent());
if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
count = bound->Arg(1);
@@ -3080,7 +3101,7 @@ private:
if (indirect.key_count.IsEmpty() &&
!MatchUniformizedMaterialKey(key, handle, indirect, material_source, false) &&
!MatchUniformizedMaterialKey(key, handle, indirect, material_source, true) &&
!MatchReversedImageRange(key, indirect)) {
!MatchReversedImageRange(key, indirect) && !MatchWorkgroupKey(key, indirect)) {
const auto maximum = SelectorMaximum(key);
if (maximum >= MaxIndirectDescriptorProbes) return false;
indirect.key_count = Value(maximum + 1u);