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;
+103 -2
View File
@@ -9921,6 +9921,101 @@ public:
std::vector<u32>(8, image_value + count),
"raw buffer descriptor materialization overwrote the newer native image");
}
// PPSA24156 loads an image sharp from a 148-byte record selected by
// WorkGroupID.z, then uses the same descriptor words for dimensions.
constexpr auto parameters = base + 0x80000;
constexpr auto table_address = parameters + 0x100;
std::vector<u32> table_shader{
EncodeSop2(0x26, 8, 4, 255), 148,
EncodeSmem0(0x02, 16), EncodeSmem1(0, 125),
EncodeSmem0(0x0b, 32, 8), EncodeSmem1(0, 8),
EncodeSop2(0x1e, 8, 34, InlineU32(2)),
EncodeSop2(0x20, 9, 33, InlineU32(30)),
EncodeSop2(0x0e, 8, 8, 255), 0xfffcu,
EncodeSop2(0x10, 8, 8, 9),
EncodeSop2(0x27, 9, 34, 255), 0x000e000eu,
EncodeSop2(0x02, 8, 8, 9),
EncodeSop2(0x02, 8, 8, InlineU32(2)),
EncodeVop1(0x01, 0, 4), EncodeVop1(0x01, 1, InlineU32(0)),
EncodeVop1(0x01, 2, 8),
EncodeMimg0(0x08, 1), EncodeMimg1(2, 0, 8)};
AppendEnd(&table_shader);
ShaderMapUserData(reinterpret_cast<uint64_t>(table_shader.data()),
{.type = Prospero::ShaderBinaryType::kCs,
.code_size_bytes = static_cast<u32>(table_shader.size() * sizeof(u32))});
shaders.SetCsShader({.data_addr = reinterpret_cast<uint64_t>(table_shader.data()),
.num_thread_x = 1, .num_thread_y = 1, .num_thread_z = 1,
.wave_size = 64, .user_sgpr = 2, .tgid_x_en = true,
.tgid_y_en = true, .tgid_z_en = true});
shaders.SetCsUserSgpr(0, static_cast<u32>(parameters), HW::UserSgprType::Unknown);
shaders.SetCsUserSgpr(1, static_cast<u32>(parameters >> 32u), HW::UserSgprType::Unknown);
const ShaderBufferResource table_buffer{{static_cast<u32>(table_address),
static_cast<u32>(table_address >> 32u) | (148u << 16u), 2, 0x5204u}};
LibKernel::Memory::WriteBacking(parameters, table_buffer.fields, sizeof(table_buffer.fields));
struct TableCase { u32 groups, width, height, image; };
uint64_t table_program_id = 0;
for (const auto test : {TableCase{1, 4, 2, 0}, TableCase{2, 4, 2, 0},
TableCase{2, 8, 3, 1}, TableCase{1, 4, 2, 0}}) {
const auto address = base + 0x90000 + test.image * 0x10000;
const ShaderTextureResource sharp{{static_cast<u32>(address >> 8u),
(static_cast<u32>(Prospero::BufferFormat::k32UInt) << 20u) |
(((test.width - 1u) & 3u) << 30u),
((test.width - 1u) >> 2u) | ((test.height - 1u) << 14u),
DstSel(4, 5, 6, 7) |
(static_cast<u32>(Prospero::TileMode::kLinear) << 20u) |
(static_cast<u32>(Prospero::ImageType::kColor2D) << 28u),
0, 0x00700000u, 0, 0}};
std::array<u32, 74> table{};
for (u32 group = 0; group < test.groups; group++) {
std::copy_n(sharp.fields, 8, table.begin() + group * 37);
}
Require(name, "descriptor table refresh", context.InvalidateMemory(table_address, sizeof(table)),
"the descriptor table escaped its mapped allocation");
LibKernel::Memory::WriteBacking(table_address, table.data(), sizeof(table));
ShaderComputeInputInfo input{};
input.workgroup_counts[0] = input.workgroup_counts[1] = 1;
input.workgroup_counts[2] = test.groups;
const auto program = context.GetPipelineCache().GetComputeProgram(
shaders.GetCs(), processor.GetCtx().GetShaderRegisters(), input);
Require(name, "workgroup table refresh and coalescing",
input.workgroup_counts[2] == test.groups &&
input.stage.program->info.images.size() == 1 &&
input.stage.program->info.images[0].indirect_resources.empty() &&
input.stage.resources->images.size() == 1 &&
std::equal(std::begin(sharp.fields), std::end(sharp.fields),
input.stage.resources->images[0].dwords.begin()) &&
!input.stage.program->info.buffers.empty(),
"workgroup bounds, refreshed sharp, or live descriptor dimension reads were lost");
Require(name, "workgroup table program reuse",
table_program_id == 0 || program.id == table_program_id,
"changing the workgroup count or identical sharp values recompiled the shader");
table_program_id = program.id;
const auto *snapshot = input.stage.resources;
const auto *image_storage = snapshot->images.data();
const auto flattened_capacity = snapshot->flattened_srt.capacity();
const auto repeated = context.GetPipelineCache().GetComputeProgram(
shaders.GetCs(), processor.GetCtx().GetShaderRegisters(), input);
Require(name, "workgroup table cache-hit storage reuse",
repeated.id == program.id && input.stage.resources == snapshot &&
snapshot->images.data() == image_storage &&
snapshot->flattened_srt.capacity() == flattened_capacity,
"a repeated descriptor table hit replaced its snapshot or storage");
auto binding = RenderExecutorTestAccess::ResolveTexture(
context.GetRenderExecutor(), input.stage.program->info.images[0],
input.stage.resources->images[0]);
const auto target = textures.FindImage(binding.desc);
TextureCacheTestAccess::ClearImage(textures, scheduler.Current(), target,
{vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1}, {});
processor.DispatchDirect(1, 1, test.groups, 0x41u);
const auto pixels = ReadCachedTexel(name, context, target, {},
{test.width, test.height, 1});
for (u32 pixel = 0; pixel < pixels.size(); pixel++) {
const auto expected = pixel < test.groups ? test.width + test.height : 0u;
Require(name, "workgroup table GPU dimension reads", pixels[pixel] == expected,
"a retained descriptor load used stale dimensions or the wrong record stride");
}
}
context.UnmapMemory(base, allocation_size);
scheduler.Finish();
});
@@ -32161,7 +32256,7 @@ void CheckIndirectBufferStore(VulkanHarness &vulkan) {
candidate.indirect_search_iterations = 0;
candidate.indirect_resources.clear();
program.descriptor_sources.resize(1);
program.descriptor_sources[0].indirect_descriptor.emplace();
program.descriptor_sources[0].indirect_descriptor.emplace().table_stride = 16;
ShaderComputeInputInfo compute;
compute.wave_size = 32;
compute.host_subgroup_size = vulkan.SubgroupSize();
@@ -32272,7 +32367,8 @@ void CheckIndirectImageKeySwitch() {
program.descriptor_sources.resize(2);
program.descriptor_sources[0].dword_count = 8;
program.descriptor_sources[0].indirect_descriptor =
DescriptorSource::IndirectDescriptor{0u, 0u, 224u, 12u, 0u};
DescriptorSource::IndirectDescriptor{.material_source = 0, .table_source = 0,
.selector_stride = 224, .selector_offset = 12, .table_stride = 32};
program.descriptor_sources[1].dword_count = 4;
ImageResource root{};
@@ -39609,6 +39705,11 @@ int main(int argc, char **argv) {
vulkan.CheckRenderExecutorStencilBindingDiscovery();
return 0;
}
if (argc == 2 && std::strcmp(argv[1], "--native-indirect-only") == 0) {
VulkanHarness vulkan;
vulkan.CheckNativeIndirectDispatch();
return 0;
}
if (argc == 2 && std::strcmp(argv[1], "--layered-image-only") == 0) {
VulkanHarness vulkan;
vulkan.CheckUnifiedImageViewCache();