shader: read image descriptors indexed by workgroup

Co-authored-by: Çağatay Güçlü <33038942+defektu@users.noreply.github.com>
This commit is contained in:
nmzik
2026-10-04 22:06:58 +04:00
parent 0bef3fc0d3
commit 0ee8800905
9 changed files with 187 additions and 29 deletions
@@ -299,11 +299,14 @@ struct PipelineCache::ProgramCache {
lookup_key.code_size = static_cast<uint32_t>(params.code.size());
BuildStageStaticKey(input_info, lookup_key.static_state);
auto entry = programs.find(lookup_key);
const ShaderRecompiler::IR::SrtRuntime runtime {
ShaderRecompiler::IR::SrtRuntime runtime {
.user_data = user_data,
.shader_base = params.Base(),
.read_specialization_memory = ReadShaderGuestMemory,
};
if constexpr (std::is_same_v<InputInfo, ShaderComputeInputInfo>) {
runtime.workgroup_counts = input_info.workgroup_counts;
}
if (entry != programs.end()) {
EXIT_IF(!ShaderRecompiler::IR::MaterializeResources(
entry->second.resource_plan, runtime, entry->second.resources,
@@ -254,6 +254,18 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer,
ShaderComputeInputInfo input_info {};
const bool use_thread_dimensions = (mode & DISPATCH_INITIATOR_USE_THREAD_DIMENSIONS) != 0;
input_info.dispatch_thread_dimensions = use_thread_dimensions;
input_info.workgroup_counts[0] = thread_group_x;
input_info.workgroup_counts[1] = thread_group_y;
input_info.workgroup_counts[2] = thread_group_z;
if (use_thread_dimensions) {
const uint32_t group_sizes[] = {cs_regs.cs_regs.num_thread_x, cs_regs.cs_regs.num_thread_y,
cs_regs.cs_regs.num_thread_z};
for (uint32_t axis = 0; axis < 3u; ++axis) {
const auto size = std::max(group_sizes[axis], 1u);
const auto threads = input_info.workgroup_counts[axis];
input_info.workgroup_counts[axis] = threads / size + (threads % size != 0u);
}
}
const auto compute_program =
m_context.GetPipelineCache().GetComputeProgram(cs_regs, sh_regs, input_info);
if (use_thread_dimensions) {
@@ -273,18 +285,12 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer,
}
if (use_thread_dimensions) {
auto groups_from_threads = [](uint32_t threads, uint32_t group_size) {
return (threads == 0
? 0u
: (threads + std::max(group_size, 1u) - 1u) / std::max(group_size, 1u));
};
const uint32_t old_x = thread_group_x;
const uint32_t old_y = thread_group_y;
const uint32_t old_z = thread_group_z;
thread_group_x = groups_from_threads(thread_group_x, cs_regs.cs_regs.num_thread_x);
thread_group_y = groups_from_threads(thread_group_y, cs_regs.cs_regs.num_thread_y);
thread_group_z = groups_from_threads(thread_group_z, cs_regs.cs_regs.num_thread_z);
thread_group_x = input_info.workgroup_counts[0];
thread_group_y = input_info.workgroup_counts[1];
thread_group_z = input_info.workgroup_counts[2];
static std::atomic<uint32_t> log_count {0};
if (log_count.fetch_add(1, std::memory_order_relaxed) < 32) {
@@ -488,6 +488,8 @@ struct DescriptorSource {
uint32_t selector_stride = 0;
uint32_t selector_offset = 0;
uint32_t table_offset = 0;
uint32_t table_stride = 0;
uint32_t workgroup_axis = UINT32_MAX;
uint32_t selector_shift = 0;
Value key_count;
Value selector_first;
@@ -285,12 +285,19 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
}
if (indirect.material_source == UINT32_MAX) {
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 > MaxIndirectDescriptorProbes ||
uint64_t {indirect.table_offset} + uint64_t {key_count} * descriptor_bytes >
UINT32_MAX + 1ull) {
if (indirect.workgroup_axis != UINT32_MAX) {
if (indirect.workgroup_axis >= runtime.workgroup_counts.size()) return false;
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))
return false;
if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
}
if (indirect.table_stride == 0u || key_count > MaxIndirectDescriptorProbes ||
(key_count != 0u && uint64_t {indirect.table_offset} +
uint64_t {key_count - 1u} * indirect.table_stride + descriptor_bytes >
UINT32_MAX + 1ull)) {
return false;
}
keys.resize(key_count);
@@ -376,7 +383,7 @@ bool MaterializeIndirectDescriptor(const ResourcePlan& p
candidate.dword_count = dword_count;
if (sources.empty()) {
const auto table_offset =
static_cast<uint32_t>(key * descriptor_bytes) + indirect.table_offset;
static_cast<uint32_t>(key * indirect.table_stride) + indirect.table_offset;
if (!ReadScalarTable(table_base, table_size, table_offset, runtime,
std::span(candidate.dwords).first(dword_count))) {
return false;
@@ -834,7 +834,8 @@ private:
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.selector_shift != b.selector_shift ||
a.table_offset != b.table_offset || a.sources != b.sources ||
a.table_offset != b.table_offset || a.table_stride != b.table_stride ||
a.workgroup_axis != b.workgroup_axis || a.sources != b.sources ||
!EquivalentValue(m_program, a.key_count, b.key_count) ||
a.selector_first.IsEmpty() != b.selector_first.IsEmpty() ||
(!a.selector_first.IsEmpty() &&
@@ -1010,7 +1011,17 @@ private:
return false;
}
bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t shift = 5u) const {
uint32_t WorkgroupAxis(Value key) const {
const auto* builtin = key.Resolve().TryInstruction();
uint32_t axis = 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)) ||
!ImmediateU32(builtin->Arg(1), axis) || axis >= 3u) return UINT32_MAX;
return axis;
}
bool MatchTableOffset(Value value, Value& key, uint32_t& offset, uint32_t& stride) const {
offset = 0;
for (;;) {
const auto* inst = value.Resolve().TryInstruction();
@@ -1019,10 +1030,17 @@ private:
}
uint32_t immediate;
if (inst->GetOpcode() == ValueOpcode::ShiftLeftLogical32 &&
ImmediateU32(inst->Arg(1), immediate) && immediate == shift) {
ImmediateU32(inst->Arg(1), immediate) && immediate < 32u) {
key = inst->Arg(0).Resolve();
stride = 1u << immediate;
return key.GetType() == Type::U32;
}
if (inst->GetOpcode() == ValueOpcode::IMul32) {
if (ImmediateU32(inst->Arg(0), stride)) key = inst->Arg(1).Resolve();
else if (ImmediateU32(inst->Arg(1), stride)) key = inst->Arg(0).Resolve();
else return false;
return stride != 0u && key.GetType() == Type::U32;
}
if (inst->GetOpcode() != ValueOpcode::IAdd32) {
return false;
}
@@ -1563,7 +1581,8 @@ private:
bool MatchDescriptorTable(Inst& handle, const DescriptorSource& descriptor,
IndirectDescriptorPlan& plan,
DescriptorSource& table_source, Value& key, uint32_t& table_offset) {
DescriptorSource& table_source, Value& key, uint32_t& table_offset,
uint32_t& table_stride) {
Inst* table_handle = nullptr;
for (uint32_t dword = 0; dword < handle.NumArgs(); ++dword) {
auto* read = descriptor.dwords[dword].Resolve().TryInstruction();
@@ -1579,6 +1598,7 @@ private:
auto* current_handle = read->Arg(0).Resolve().TryInstruction();
Value current_key;
uint32_t offset = 0;
uint32_t stride = 0;
if (current_handle == nullptr ||
current_handle->GetOpcode() != (memory->kind == ResourceKind::ScalarAddress
? ValueOpcode::GetAddressResource
@@ -1586,7 +1606,7 @@ private:
(memory->kind == ResourceKind::ScalarAddress && read->Parent() != handle.Parent()) ||
(table_handle != nullptr &&
!EquivalentValue(m_program, Value(table_handle), Value(current_handle))) ||
!MatchTableOffset(read->Arg(1), current_key, offset, handle.NumArgs() == 4u ? 4u : 5u) ||
!MatchTableOffset(read->Arg(1), current_key, offset, stride) ||
memory->offset > UINT32_MAX - offset) {
return false;
}
@@ -1594,12 +1614,14 @@ private:
if (dword == 0u) {
key = current_key;
table_offset = offset;
} else if (!EquivalentValue(m_program, key, current_key) ||
table_stride = stride;
} else if (table_stride != stride || !EquivalentValue(m_program, key, current_key) ||
static_cast<uint64_t>(table_offset) + dword * sizeof(uint32_t) != offset) {
return false;
}
table_handle = current_handle;
if (handle.NumArgs() == 8u ? !UsesOnlyImageDescriptors(*read) :
if (handle.NumArgs() == 8u
? WorkgroupAxis(key) == UINT32_MAX && !UsesOnlyImageDescriptors(*read) :
std::ranges::any_of(read->Uses(), [](const Use& use) {
return use.user->GetOpcode() != ValueOpcode::GetBufferResource ||
std::ranges::any_of(use.user->Uses(), [](const Use& consumer) {
@@ -1654,11 +1676,21 @@ private:
DescriptorSource table_source;
Value key;
uint32_t table_offset = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset)) return false;
uint32_t table_stride = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset,
table_stride)) return false;
DescriptorSource material_source;
DescriptorSource::IndirectDescriptor indirect;
indirect.table_offset = table_offset;
if (table_source.dword_count == 2u) {
indirect.table_stride = table_stride;
indirect.workgroup_axis = WorkgroupAxis(key);
if (indirect.workgroup_axis != UINT32_MAX) {
// The descriptor also supplies dimensions to shader arithmetic. Keep its reads;
// only the image handle is projected onto the bounded workgroup key.
plan.reads.fill(nullptr);
} else if (table_stride != 32u) {
return false;
} else if (table_source.dword_count == 2u) {
const auto* selector = key.Resolve().TryInstruction();
const bool bitscan = selector != nullptr && selector->GetOpcode() == ValueOpcode::FindILsb32 &&
selector->NumArgs() == 1u && !m_shader_writes &&
@@ -1733,7 +1765,9 @@ private:
DescriptorSource table_source;
Value key;
uint32_t table_offset = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset) ||
uint32_t table_stride = 0;
if (!MatchDescriptorTable(handle, descriptor, plan, table_source, key, table_offset,
table_stride) || table_stride != 16u ||
table_source.dword_count != 4u || table_offset != 0u) return false;
const auto* lane_read = key.Resolve().TryInstruction();
if (lane_read == nullptr || lane_read->GetOpcode() != ValueOpcode::ReadFirstLane)
@@ -1763,6 +1797,7 @@ private:
!EquivalentValue(m_program, selected->Arg(0), read->Arg(4))) return false;
DescriptorSource material_source;
DescriptorSource::IndirectDescriptor indirect;
indirect.table_stride = table_stride;
if (!MakeRuntimeTableSource(*read, material_source) ||
!BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false;
indirect.material_source = InternSource(material_source);
@@ -17,6 +17,7 @@ struct SrtRuntime {
SrtMemoryReader read_memory = nullptr;
void* userdata = nullptr;
SrtMemoryReader read_specialization_memory = nullptr;
std::span<const uint32_t> workgroup_counts;
};
enum class RuntimeValueType { Any, Integer };
+2
View File
@@ -571,9 +571,11 @@ static void ShaderGetStaticInputInfoCS(const HW::ComputeShaderInfo& regs,
const ShaderMappedData& data, ShaderComputeInputInfo& info) {
const bool dispatch_thread_dimensions = info.dispatch_thread_dimensions;
const auto host_subgroup_size = info.host_subgroup_size;
const auto workgroup_counts = std::to_array(info.workgroup_counts);
info = {};
info.dispatch_thread_dimensions = dispatch_thread_dimensions;
info.host_subgroup_size = host_subgroup_size;
std::ranges::copy(workgroup_counts, info.workgroup_counts);
info.threads_num[0] = regs.cs_regs.num_thread_x;
info.threads_num[1] = regs.cs_regs.num_thread_y;
info.threads_num[2] = regs.cs_regs.num_thread_z;
+1
View File
@@ -148,6 +148,7 @@ struct ShaderVertexInputInfo {
struct ShaderComputeInputInfo: ShaderWorkgroupInputInfo {
uint8_t float_mode = 0xc0;
uint32_t dispatch_threads_num[3] = {0, 0, 0};
uint32_t workgroup_counts[3] = {0, 0, 0};
bool group_id[3] = {false, false, false};
bool dispatch_thread_dimensions = false;
int thread_ids_num = 0;