Track bounded indirect image tables from signed shader loops

Co-authored-by: TarkusR <90155059+TarkusR@users.noreply.github.com>
This commit is contained in:
nmzik
2026-09-24 01:57:04 +04:00
parent 3dc3fd0197
commit 3b89e36d1e
6 changed files with 288 additions and 17 deletions
@@ -784,11 +784,11 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
? &state.program.descriptor_sources[image.source]
: nullptr;
if (handle == nullptr || source == nullptr || !source->indirect_image.has_value() ||
source->indirect_image->key_arg >= handle->NumArgs()) {
handle->NumArgs() == 0u) {
ctx.Fail(inst, "has invalid indirect image key provenance");
return;
}
const auto key = ctx.Def(handle->Arg(source->indirect_image->key_arg));
const auto key = ctx.Def(handle->Arg(0));
if (state.flattened_srt_variable == 0 || image.indirect_search_iterations == 0u ||
image.indirect_resources.size() < 2u) {
ctx.Fail(inst, "has no indirect image runtime mapping");
+1 -2
View File
@@ -463,9 +463,8 @@ struct DescriptorSource {
uint32_t table_source = 0;
uint32_t selector_stride = 0;
uint32_t selector_offset = 0;
uint32_t key_arg = 0;
uint32_t table_offset = 0;
uint32_t key_count = 0;
Value key_count;
bool operator==(const IndirectImage& other) const = default;
};
@@ -200,6 +200,13 @@ void FoldInstruction(Block& block, Block::iterator instruction,
switch (inst.GetOpcode()) {
case ValueOpcode::Phi: FoldPhi(inst); return;
case ValueOpcode::SelectU1:
if (!FoldSelect(inst) && IsImmediate(Arg(inst, 2), Type::U1) &&
!Arg(inst, 2).U1()) {
auto result = block.PrependNewInst(instruction, ValueOpcode::LogicalAnd,
{Arg(inst, 0), Arg(inst, 1)});
Replace(inst, Value(&*result));
}
return;
case ValueOpcode::SelectU32:
case ValueOpcode::SelectF32: FoldSelect(inst); return;
case ValueOpcode::BitFieldInsert: {
@@ -431,6 +438,10 @@ void FoldInstruction(Block& block, Block::iterator instruction,
return;
case ValueOpcode::ShiftRightLogical32:
if (!FoldU32(inst, [](uint32_t a, uint32_t b) { return a >> (b & 31u); })) {
if (IsImmediate(Arg(inst, 0), Type::U32) && Arg(inst, 0).U32() == 0u) {
Replace(inst, Value(0u));
return;
}
const auto shift = Arg(inst, 1);
if (IsImmediate(shift, Type::U32) && (shift.U32() & 31u) == 0u) {
Replace(inst, Arg(inst, 0));
@@ -573,6 +584,23 @@ void FoldInstruction(Block& block, Block::iterator instruction,
if (!FoldLogical(inst, [](bool a, bool b) { return a && b; })) {
const auto lhs = Arg(inst, 0);
const auto rhs = Arg(inst, 1);
const auto simplify = [&](Value assumption, Value expression) {
const auto* disjunction = expression.TryInstruction();
if (disjunction == nullptr || disjunction->GetOpcode() != ValueOpcode::LogicalOr) {
return false;
}
for (uint32_t i = 0; i < 2u; ++i) {
const auto* inverse = disjunction->Arg(i).Resolve().TryInstruction();
if (inverse != nullptr && inverse->GetOpcode() == ValueOpcode::LogicalNot &&
inverse->Arg(0).Resolve() == assumption) {
inst.SetArg(assumption == lhs ? 1u : 0u,
disjunction->Arg(i ^ 1u));
return true;
}
}
return false;
};
if (simplify(lhs, rhs) || simplify(rhs, lhs)) return;
if (IsImmediate(lhs, Type::U1)) {
Replace(inst, lhs.U1() ? rhs : lhs);
} else if (IsImmediate(rhs, Type::U1)) {
@@ -198,7 +198,8 @@ bool MaterializeIndirectImage(const ResourcePlan& program,
const DescriptorSource::IndirectImage& indirect,
const DescriptorValue& material_value,
const DescriptorValue& table_value, uint32_t image_index,
const SrtRuntime& runtime, ResourceSnapshot& snapshot,
const SrtRuntime& runtime, SrtWalker& clean,
ResourceSnapshot& snapshot,
ResourceSpecialization& specialization) {
uint64_t table_base = 0;
uint64_t table_size = UINT64_MAX; // Scalar addresses have no buffer descriptor bounds.
@@ -214,11 +215,15 @@ bool MaterializeIndirectImage(const ResourcePlan& program,
auto& keys = program.material_keys;
keys.clear();
if (indirect.material_source == UINT32_MAX) {
if (table_value.dword_count != 2u || indirect.key_count == 0u ||
indirect.key_count > MaxIndirectImageProbes) {
uint32_t key_count = 0;
const bool evaluated = clean.Evaluate(indirect.key_count, key_count);
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
if (table_value.dword_count != 2u || !evaluated ||
key_count > MaxIndirectImageProbes ||
uint64_t {indirect.table_offset} + uint64_t {key_count} * 32u > UINT32_MAX + 1ull) {
return false;
}
keys.resize(indirect.key_count);
keys.resize(key_count);
std::iota(keys.begin(), keys.end(), 0u);
} else {
ShaderBufferResource material;
@@ -816,6 +821,9 @@ ResourcePlan ExtractResourcePlan(const Program& program) {
auto& target = plan.descriptor_sources.emplace_back();
target.dword_count = source.dword_count;
target.indirect_image = source.indirect_image;
if (target.indirect_image.has_value()) {
target.indirect_image->key_count = Clone(target.indirect_image->key_count);
}
for (uint32_t dword = 0; dword < source.dword_count; dword++) {
target.dwords[dword] = Clone(source.dwords[dword]);
}
@@ -921,7 +929,7 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
if ((indirect.material_source != UINT32_MAX &&
!clean.EvaluateDescriptor(indirect.material_source, material)) ||
!clean.EvaluateDescriptor(indirect.table_source, table) ||
!MaterializeIndirectImage(program, indirect, material, table, i, runtime, snapshot,
!MaterializeIndirectImage(program, indirect, material, table, i, runtime, clean, snapshot,
specialization)) {
return false;
}
@@ -313,9 +313,17 @@ private:
for (uint32_t candidate = 0; candidate < m_sources.size(); candidate++) {
const auto& current = m_sources[candidate];
if (current.dword_count != descriptor.dword_count ||
current.indirect_image != descriptor.indirect_image) {
current.indirect_image.has_value() != descriptor.indirect_image.has_value()) {
continue;
}
if (current.indirect_image.has_value()) {
const auto& a = *current.indirect_image;
const auto& b = *descriptor.indirect_image;
if (a.material_source != b.material_source || a.table_source != b.table_source ||
a.selector_stride != b.selector_stride || a.selector_offset != b.selector_offset ||
a.table_offset != b.table_offset ||
!EquivalentValue(m_program, a.key_count, b.key_count)) continue;
}
bool same = true;
for (uint32_t i = 0; i < descriptor.dword_count; i++) {
same = same && EquivalentValue(m_program, current.dwords[i], descriptor.dwords[i]);
@@ -504,6 +512,82 @@ private:
}
}
Value BoundedLoopCount(Value key, const Block* use) const {
const auto* phi = key.Resolve().TryInstruction();
if (m_shader_writes || phi == nullptr || phi->GetOpcode() != ValueOpcode::Phi ||
phi->GetType() != Type::U32 || phi->NumArgs() != 2u ||
m_program.blocks.size() != m_program.block_info.size()) {
return {};
}
bool induction = false;
for (uint32_t initial = 0; initial < 2u; ++initial) {
const auto zero = phi->Arg(initial).Resolve();
const auto* step = phi->Arg(initial ^ 1u).Resolve().TryInstruction();
if (!zero.IsImmediate() || zero.GetType() != Type::U32 || zero.U32() != 0u ||
step == nullptr || step->GetOpcode() != ValueOpcode::IAdd32 ||
step->Parent() != phi->PhiBlock(initial ^ 1u)) {
continue;
}
uint32_t increment = 0;
induction = (step->Arg(0).Resolve() == key &&
ImmediateU32(step->Arg(1), increment) && increment == 1u) ||
(step->Arg(1).Resolve() == key &&
ImmediateU32(step->Arg(0), increment) && increment == 1u);
if (induction) break;
}
if (!induction) return {};
const auto find_block = [&](uint32_t id) -> const Block* {
const auto info = std::ranges::find(m_program.block_info, id, &BlockInfo::id);
return info == m_program.block_info.end()
? nullptr
: m_program.blocks[info - m_program.block_info.begin()];
};
const auto reaches = [&](const Block* start, const Block* target, const Block* avoid) {
std::vector<const Block*> pending {start};
std::vector<const Block*> visited;
while (!pending.empty()) {
const auto* block = pending.back();
pending.pop_back();
if (block == nullptr || block == avoid ||
std::ranges::find(visited, block) != visited.end()) continue;
if (block == target) return true;
visited.push_back(block);
for (const auto* next: block->ImmSuccessors()) pending.push_back(next);
}
return false;
};
const auto contains = [&](auto&& self, Value value, const Inst* comparison) -> bool {
value = value.Resolve();
if (value.TryInstruction() == comparison) return true;
const auto* inst = value.TryInstruction();
return inst != nullptr && inst->GetOpcode() == ValueOpcode::LogicalAnd &&
(self(self, inst->Arg(0), comparison) || self(self, inst->Arg(1), comparison));
};
for (const auto& use_of_key: phi->Uses()) {
const auto* compare = use_of_key.user;
if (compare->GetOpcode() != ValueOpcode::SLessThan32 || use_of_key.operand != 0u ||
!ValidateRuntimeValue(m_program, compare->Arg(1))) continue;
for (uint32_t i = 0; i < m_program.block_info.size(); ++i) {
const auto& info = m_program.block_info[i];
const auto* negated = info.condition.Resolve().TryInstruction();
if (info.terminator.kind != CFG::TerminatorKind::ConditionalBranch ||
compare->Parent() != m_program.blocks[i] || negated == nullptr ||
negated->GetOpcode() != ValueOpcode::LogicalNot ||
!contains(contains, negated->Arg(0), compare) ||
use == m_program.blocks[i] ||
reaches(m_program.blocks.front(), use, m_program.blocks[i]) ||
reaches(phi->Parent(), use, m_program.blocks[i]) ||
reaches(find_block(info.terminator.true_block), use, phi->Parent()) ||
!reaches(find_block(info.terminator.false_block), use, phi->Parent())) {
continue;
}
return compare->Arg(1);
}
}
return {};
}
bool TryMakeIndirectImage(Inst& handle, uint32_t pc, IndirectImagePlan& plan) {
if (handle.GetOpcode() != ValueOpcode::GetImageResource || handle.NumArgs() != 8u) {
return false;
@@ -562,13 +646,16 @@ private:
indirect.table_offset = table_offset;
if (table_source.dword_count == 2u) {
const auto* selector = key.Resolve().TryInstruction();
if (m_shader_writes || selector == nullptr ||
selector->GetOpcode() != ValueOpcode::FindILsb32 || selector->NumArgs() != 1u ||
!NonzeroOnEntry(selector->Arg(0), handle.Parent()) ||
(table_offset & 3u) != 0u || table_offset > UINT32_MAX - (32u * 32u - 1u)) {
return false;
const bool bitscan = selector != nullptr && selector->GetOpcode() == ValueOpcode::FindILsb32 &&
selector->NumArgs() == 1u && !m_shader_writes &&
NonzeroOnEntry(selector->Arg(0), handle.Parent());
if (bitscan) {
indirect.key_count = Value(32u);
} else {
indirect.key_count = BoundedLoopCount(key, handle.Parent());
}
indirect.key_count = 32u;
if (indirect.key_count.IsEmpty() || (table_offset & 3u) != 0u ||
(bitscan && table_offset > UINT32_MAX - (32u * 32u - 1u))) return false;
} else {
auto* material_read = key.Resolve().TryInstruction();
uint32_t material_memory_index = 0;
+150 -1
View File
@@ -614,7 +614,8 @@ void TestGuardedDirectImageTable() {
const auto &indirect = fixture.program.descriptor_sources[source].indirect_image;
Check(indirect && indirect->material_source == UINT32_MAX &&
indirect->selector_stride == 0u && indirect->table_offset == 344u &&
indirect->key_count == 32u &&
indirect->key_count.Resolve().IsImmediate() &&
indirect->key_count.Resolve().U32() == 32u &&
fixture.program.descriptor_sources[indirect->table_source].dword_count == 2u,
"guarded direct image table lost its pointer or proven selector range");
return ExtractResourcePlan(fixture.program);
@@ -699,6 +700,153 @@ void TestGuardedDirectImageTable() {
"batched descriptor read crossed the 48-bit endpoint");
}
void TestBoundedComputeImageLoop() {
namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
enum class Variant {
Bounded, WrongGuard, EntryBypass, ExitBypass, GuardBlock, WrongStep
};
const auto make_plan = [](Variant variant) {
Fixture fixture;
auto *entry = fixture.block;
auto *header = fixture.AddBlock();
auto *body = fixture.AddBlock();
auto *latch = fixture.AddBlock();
auto *exit = fixture.AddBlock();
entry->AddBranch(header);
header->AddBranch(exit);
header->AddBranch(body);
body->AddBranch(latch);
latch->AddBranch(header);
if (variant == Variant::EntryBypass) entry->AddBranch(body);
if (variant == Variant::ExitBypass) exit->AddBranch(body);
fixture.program.block_info[0].terminator = {
.kind = variant == Variant::EntryBypass ? CFG::TerminatorKind::ConditionalBranch
: CFG::TerminatorKind::Branch,
.true_block = 1u, .false_block = 2u};
fixture.program.block_info[1].terminator = {
.kind = CFG::TerminatorKind::ConditionalBranch,
.true_block = 4u, .false_block = 2u};
fixture.program.block_info[2].terminator = {
.kind = CFG::TerminatorKind::Branch, .true_block = 3u};
fixture.program.block_info[3].terminator = {
.kind = CFG::TerminatorKind::Branch, .true_block = 1u};
fixture.program.block_info[4].terminator = {
.kind = variant == Variant::ExitBypass ? CFG::TerminatorKind::Branch
: CFG::TerminatorKind::Return,
.true_block = 2u};
auto &phi = header->AppendNewInst(ValueOpcode::Phi, {},
static_cast<uint64_t>(Type::U32));
const auto key = Value(&phi);
const auto count = fixture.UserData(2);
const auto in_range = fixture.Emit(ValueOpcode::SLessThan32,
{variant == Variant::WrongGuard
? Value(0u) : key,
count}, 0, header);
const auto allowed = fixture.Emit(ValueOpcode::LogicalAnd,
{in_range, Value(true)}, 0, header);
fixture.program.block_info[1].condition =
fixture.Emit(ValueOpcode::LogicalNot, {allowed}, 0, header);
const auto step = fixture.Emit(ValueOpcode::IAdd32,
{key, Value(variant == Variant::WrongStep ? 2u : 1u)},
0, latch);
phi.AddPhiOperand(entry, Value(0u));
phi.AddPhiOperand(latch, step);
fixture.block = variant == Variant::GuardBlock ? header : body;
const auto table = fixture.Address(fixture.UserData(0), fixture.UserData(1));
const auto offset = fixture.Emit(ValueOpcode::IAdd32,
{fixture.Emit(ValueOpcode::ShiftLeftLogical32, {key, Value(5u)}),
Value(0x6b0u)});
std::array<Value, 8> words;
for (uint32_t word = 0; word < words.size(); ++word) {
MemoryInfo memory;
memory.kind = ResourceKind::ScalarAddress;
memory.offset = word * sizeof(uint32_t);
words[word] = fixture.Emit(
ValueOpcode::LoadAddressU32,
{table, offset, Value(0u), Value(true)},
fixture.AddMemory(memory, 0x29c));
}
const auto image = fixture.Image(words, 0x29c);
const auto sampler = fixture.Sampler({Value(0u), Value(0u), Value(0u), Value(0u)});
MemoryInfo sample;
sample.kind = ResourceKind::Image;
sample.image_dimension = Decoder::ImageDimension::Dim2D;
fixture.Emit(ValueOpcode::ImageSampleRaw,
{image, sampler, fixture.ImageAddress()},
fixture.AddMemory(sample, 0x29c));
fixture.PlanAndTrack();
const auto source = fixture.program.info.images[0].source;
const auto &indirect = fixture.program.descriptor_sources[source].indirect_image;
Check(indirect && indirect->material_source == UINT32_MAX &&
indirect->table_offset == 0x6b0u &&
indirect->key_count.Resolve() == count.Resolve(),
"bounded compute loop lost its runtime image count");
return ExtractResourcePlan(fixture.program);
};
auto plan = make_plan(Variant::Bounded);
CheckFatal([&] { make_plan(Variant::WrongGuard); },
"not a valid runtime value",
"compute image loop accepted an unrelated guard");
CheckFatal([&] { make_plan(Variant::EntryBypass); },
"not a valid runtime value",
"compute image loop accepted an entry bypass");
CheckFatal([&] { make_plan(Variant::ExitBypass); },
"not a valid runtime value",
"compute image loop accepted an exit bypass");
CheckFatal([&] { make_plan(Variant::GuardBlock); },
"not a valid runtime value",
"compute image loop accepted a descriptor read before the guard");
CheckFatal([&] { make_plan(Variant::WrongStep); },
"not a valid runtime value",
"compute image loop accepted a two-step induction");
LinearTestMemory memory;
const auto table = 0x1800u + 0x6b0u;
for (uint32_t key = 0; key < 3u; ++key) {
const auto word = (table - memory.base) / 4u + key * 8u;
memory.words[word] = 0x100u + key;
memory.words[word + 1u] = static_cast<uint32_t>(
Libs::Graphics::Prospero::BufferFormat::k32_32_32_32Float) << 20u;
memory.words[word + 3u] = Libs::Graphics::DstSel(4, 5, 6, 7) |
(static_cast<uint32_t>(Libs::Graphics::Prospero::ImageType::kColor2D) << 28u);
}
std::array<uint32_t, 3> user_data{0x1800u, 0u, 2u};
SrtRuntime runtime{.user_data = user_data,
.read_memory = ReadLinearTestMemory,
.userdata = &memory,
.read_specialization_memory = ReadLinearTestMemory};
ResourceSnapshot snapshot;
ResourceSpecialization specialization;
for (const uint32_t count : {2u, 3u, 2u}) {
user_data[2] = count;
Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
snapshot.images.size() == count &&
specialization.images.size() == count &&
snapshot.flattened_srt[
specialization.images[0].indirect_mapping_offset] == count &&
snapshot.images.back().dwords[0] == 0x100u + count - 1u,
"compute image table did not refresh for a changed loop bound");
}
for (const uint32_t count : {0u, UINT32_MAX}) {
user_data[2] = count;
Check(MaterializeResources(plan, runtime, snapshot, specialization) &&
snapshot.images.size() == 1u &&
specialization.images.size() == 1u &&
std::ranges::all_of(snapshot.images[0].dwords,
[](uint32_t word) { return word == 0u; }) &&
specialization.images[0].indirect_root ==
ImageResource::NoIndirectImage &&
specialization.images[0].indirect_search_iterations == 0u,
"empty compute loop bound retained unreachable image candidates");
}
user_data[2] = 65537u;
Check(!MaterializeResources(plan, runtime, snapshot, specialization),
"oversized compute loop bound was accepted for image enumeration");
}
void TestImageDescriptorFields() {
constexpr std::array<std::pair<uint32_t, uint32_t>, 5> reserved{
{{1u, 0x20000000u}, {2u, 0x70003000u}, {4u, 0xe000e000u},
@@ -2506,6 +2654,7 @@ int main() {
Run("dynamic storage mips", TestDynamicStorageMipTracking);
Run("invariant indirect images", TestInvariantIndirectImageMaterialization);
Run("guarded direct image table", TestGuardedDirectImageTable);
Run("bounded compute image loop", TestBoundedComputeImageLoop);
Run("image descriptor fields", TestImageDescriptorFields);
Run("draw-uniform scalar image", TestUniformScalarBufferImage);
Run("SRT runtime", TestSrtFlatteningAndRuntimeMemoization);