shader: read image descriptors indexed by workgroup

Co-authored-by: Çağatay Güçlü <33038942+defektu@users.noreply.github.com>

src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp

@@ -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,

src/graphics/host_gpu/renderer/renderCompute.cpp

@@ -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_x       = input_info.workgroup_counts[0];
-		thread_group_y       = groups_from_threads(thread_group_y, cs_regs.cs_regs.num_thread_y);
+		thread_group_y       = input_info.workgroup_counts[1];
-		thread_group_z       = groups_from_threads(thread_group_z, cs_regs.cs_regs.num_thread_z);
+		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) {

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -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;

src/graphics/shader/recompiler/ir/passes/ResourceMaterialization.cpp

@@ -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 (indirect.workgroup_axis != UINT32_MAX) {
-			if (std::bit_cast<int32_t>(key_count) <= 0) key_count = 0;
+				if (indirect.workgroup_axis >= runtime.workgroup_counts.size()) return false;
-			if (table_value.dword_count != 2u || !evaluated ||
+				key_count = runtime.workgroup_counts[indirect.workgroup_axis];
-			    key_count > MaxIndirectDescriptorProbes ||
+				if (key_count == 0u) return false;
-			    uint64_t {indirect.table_offset} + uint64_t {key_count} * descriptor_bytes >
+			} else {
-			        UINT32_MAX + 1ull) {
+				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;

src/graphics/shader/recompiler/ir/passes/ResourceTracking.cpp

@@ -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);

src/graphics/shader/recompiler/ir/passes/SrtWalker.h

@@ -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 };

src/graphics/shader/shader.cpp

@@ -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;

src/graphics/shader/shader.h

@@ -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;

tests/ShaderRecompilerComputeTests.cpp

@@ -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();