shader: resolve workgroup-selected descriptor prefixes from dispatch arguments

Co-authored-by: fxpw <fxpw@outlook.com>

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

@@ -347,6 +347,7 @@ struct PipelineCache::ProgramCache {
 		};
 		if constexpr (std::is_same_v<InputInfo, ShaderComputeInputInfo>) {
 			runtime.workgroup_counts = input_info.workgroup_counts;
+			runtime.indirect_args_address = input_info.indirect_args_address;
 		}
 		bool source_changed = false;
 		if (entry != programs.end()) {

src/graphics/host_gpu/renderer/renderCompute.cpp

@@ -457,6 +457,7 @@ void RenderExecutor::DispatchIndirect(uint64_t submit_id, CommandBuffer& buffer,
 	}
 	ShaderComputeInputInfo input_info {};
 	input_info.async_compute = async_compute;
+	input_info.indirect_args_address = args_addr;
 	const auto compute_program = m_context.GetPipelineCache().GetComputeProgram(
 	    cs_regs, buffer.GetRegisters().GetShaderRegisters(), input_info);
 	buffer.EndRendering();

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

@@ -307,9 +307,10 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 	const auto root_resource = specializations[resource_index];
 	const auto evaluate_count = [&](uint32_t& count) {
 		if (indirect.workgroup_axis != UINT32_MAX) {
-			if (indirect.workgroup_axis >= runtime.workgroup_counts.size()) return false;
-			count = runtime.workgroup_counts[indirect.workgroup_axis];
-			return count != 0u;
+			if (runtime.workgroup_counts.size() != 3u || indirect.workgroup_axis >= 3u) return false;
+			count = std::ranges::find(runtime.workgroup_counts, 0u) != runtime.workgroup_counts.end()
+			    ? 0u : runtime.workgroup_counts[indirect.workgroup_axis];
+			return true;
 		}
 		if (!clean.Evaluate(indirect.key_count, count)) return false;
 		if (indirect.signed_count && std::bit_cast<int32_t>(count) < 0) count = 0;
@@ -337,14 +338,28 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 	const auto read_keys = [&](uint64_t base, uint64_t size, uint64_t first,
 	                           uint64_t step, uint64_t count) {
 		keys.resize(count);
+		if (count == 0u) return true;
 		if (step == 4u) {
 			return read_selector(base, size, first, keys);
 		}
-		for (auto& key: keys) {
-			if (!read_selector(base, size, first, {&key, 1}))
+		if (step < 4u) return false;
+		for (uint64_t copied = 0; copied < count;) {
+			const auto offset = first + copied * step;
+			// Scalar reads align each strided offset independently.
+			const auto bias = selector->global ? 0u : offset & 3u;
+			const auto length = std::min(count - copied,
+			    1u + (uint64_t {MaxIndirectDescriptorProbes - 1u} * 4u - bias) / step);
+			keys.resize(count + (bias + (length - 1u) * step) / 4u + 1u);
+			const auto batch = std::span(keys).subspan(count);
+			if (!read_selector(base, size, offset, batch)) {
+				keys.resize(count);
 				return false;
-			first += step;
+			}
+			for (uint64_t index = 0; index < length; ++index)
+				keys[copied + index] = batch[(bias + index * step) / 4u];
+			copied += length;
 		}
+		keys.resize(count);
 		return true;
 	};
 	uint64_t table_base = 0;
@@ -564,7 +579,8 @@ bool MaterializeIndirectDescriptor(const ResourcePlan&                         p
 			     uint64_t {first} + count > material.NumRecords())) return false;
 			const auto end = count != 0u
 			    ? (uint64_t {first} + count - 1u) * step + selector->offset + 4u : 0u;
-			if (end > uint64_t {UINT32_MAX} + 1u || end > size)
+			if (end > uint64_t {UINT32_MAX} + 1u ||
+			    (indirect.workgroup_axis == UINT32_MAX && end > size))
 				return false;
 			if (!read_keys(base, size, uint64_t {first} * step + selector->offset, step, count))
 				return false;
@@ -1331,11 +1347,16 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 	     runtime.read_specialization_memory == nullptr)) {
 		return false;
 	}
-	const bool capture_reads = program.capture_specialization_reads || !program.source_reads.empty();
+	const auto uses_workgroup_count = [](const DescriptorSource& source) {
+		return source.indirect_descriptor && source.indirect_descriptor->workgroup_axis != UINT32_MAX;
+	};
+	const bool capture_reads = program.capture_specialization_reads || !program.source_reads.empty() ||
+	    (runtime.indirect_args_address != 0u &&
+	     std::ranges::any_of(program.descriptor_sources, uses_workgroup_count));
 	auto& reads = snapshot.specialization_reads;
 	reads = program.source_reads;
 	SrtReadCapture capture {runtime, reads};
-	const auto observed = capture_reads ? capture.ObservedRuntime() : runtime;
+	auto observed = capture_reads ? capture.ObservedRuntime() : runtime;
 	SrtWalker clean(program, CleanRuntime(observed));
 	SrtWalker walker(program, observed,
 	                 capture_reads || program.requires_specialization_memory ? &clean : nullptr);
@@ -1343,6 +1364,18 @@ bool MaterializeResources(const ResourcePlan& program, const SrtRuntime& runtime
 		return false;
 	}
 	const auto active = std::span<const uint8_t>(program.active_sources);
+	std::array<uint32_t, 3> workgroup_counts {};
+	if (runtime.indirect_args_address != 0u) {
+		for (uint32_t source = 0; source < program.descriptor_sources.size(); ++source) {
+			if ((!active.empty() && !active[source]) ||
+			    !uses_workgroup_count(program.descriptor_sources[source])) continue;
+			if ((runtime.indirect_args_address & 3u) != 0u ||
+			    !ReadTable(runtime.indirect_args_address, sizeof(workgroup_counts), 0u,
+			               observed, workgroup_counts)) return false;
+			observed.workgroup_counts = workgroup_counts;
+			break;
+		}
+	}
 	snapshot.uniform_fill = {};
 	const auto& fill = program.uniform_fill;
 	const auto words = fill.fill.words;

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

@@ -2204,6 +2204,7 @@ private:
 		const auto* memory = ScalarReadMemory(*read, memory_index);
 		Value index;
 		Value count;
+		uint32_t workgroup_axis = UINT32_MAX;
 		DescriptorSource count_source;
 		uint32_t row_stride = 0;
 		uint32_t offset = 0, stride = 0;
@@ -2212,19 +2213,21 @@ private:
 		if (MatchTableOffset(read->Arg(1), index, offset, stride) &&
 		    ((offset | stride | memory->offset) & 3u) == 0u &&
 		    uint64_t {offset} + memory->offset <= UINT32_MAX) {
+			workgroup_axis = WorkgroupAxis(index);
 			const auto* bound = BoundedLoop(index, read->Parent());
 			if (bound != nullptr && bound->GetOpcode() == ValueOpcode::ULessThan32) {
 				count = bound->Arg(1);
 			}
 		}
 		// Raw pointers have no V# bounds for enumerating an unbounded selector domain.
-		if (memory->kind == ResourceKind::ScalarAddress && count.IsEmpty()) return false;
-		const bool rows = count.IsEmpty() && memory->offset == 0u &&
+		if (memory->kind == ResourceKind::ScalarAddress && count.IsEmpty() &&
+		    workgroup_axis == UINT32_MAX) return false;
+		const bool rows = count.IsEmpty() && workgroup_axis == UINT32_MAX && memory->offset == 0u &&
 		    MatchCountedRows(*read, count_source, row_stride);
 		if (rows) {
 			stride = 4u;
 			offset = 0u;
-		} else if (count.IsEmpty()) {
+		} else if (count.IsEmpty() && workgroup_axis == UINT32_MAX) {
 			auto address = read->Arg(1);
 			auto active = Value(true);
 			const auto* lane = address.Resolve().TryInstruction();
@@ -2255,17 +2258,18 @@ private:
 		}
 		// A bounded image heap already supports GPU selection. Only enumerate a
 		// narrower proven selector domain instead of scanning unrelated material data.
-		if (require_bound && count.IsEmpty() && !rows) return false;
+		if (require_bound && count.IsEmpty() && workgroup_axis == UINT32_MAX && !rows) return false;
 		if (!MakeRuntimeTableSource(*read, material_source)) return false;
 		selector.source = InternSource(material_source);
 		selector.stride = stride;
 		selector.offset = offset + (memory->offset & ~3u);
 		indirect.selector = std::move(selector);
 		indirect.key_count = count;
+		indirect.workgroup_axis = workgroup_axis;
 		if (rows) indirect.selector->rows.emplace(
 		    DescriptorSource::IndirectDescriptor::SelectorRead::CountedRows {
 		        .source = InternSource(count_source), .stride = row_stride});
-		if (!count.IsEmpty()) indirect.selector_first = Value(0u);
+		if (!count.IsEmpty() || workgroup_axis != UINT32_MAX) indirect.selector_first = Value(0u);
 		return true;
 	}
 
@@ -3025,6 +3029,11 @@ private:
 			indirect.table_scalar = true;
 			return finish();
 		}
+		DescriptorSource material_source;
+		if (!plan.split_offsets && MatchScalarMaterialKey(key, indirect, material_source, true)) {
+			plan.retain_reads = true;
+			return finish();
+		}
 		if (!legacy_store) {
 			if (table_source.dword_count == 2u && !plan.split_offsets) {
 				const auto maximum = SelectorMaximum(key);
@@ -3082,7 +3091,6 @@ private:
 		    !ImmediateU32(read->Arg(2), offset) || offset != 0u ||
 		    !ImmediateU32(read->Arg(3), scalar) || scalar != 0u ||
 		    !EquivalentValue(m_program, selected->Arg(0), read->Arg(4))) return false;
-		DescriptorSource material_source;
 		if (!MakeRuntimeTableSource(*read, material_source) ||
 		    !BoundedSelectedIndex(read->Arg(1), read->Arg(4), indirect)) return false;
 		indirect.selector.emplace(DescriptorSource::IndirectDescriptor::SelectorRead{.source = InternSource(material_source),

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

@@ -19,6 +19,7 @@ struct SrtRuntime {
 	void*                     userdata                   = nullptr;
 	SrtMemoryReader           read_specialization_memory = nullptr;
 	std::span<const uint32_t> workgroup_counts;
+	uint64_t                  indirect_args_address = 0;
 };
 
 class SrtReadCapture {

src/graphics/shader/shader.cpp

@@ -578,11 +578,13 @@ static void ShaderGetStaticInputInfoCS(const HW::ComputeShaderInfo& regs,
 	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);
+	const auto indirect_args_address      = info.indirect_args_address;
 	info                                  = {};
 	info.async_compute                    = async_compute;
 	info.dispatch_thread_dimensions       = dispatch_thread_dimensions;
 	info.host_subgroup_size               = host_subgroup_size;
 	std::ranges::copy(workgroup_counts, info.workgroup_counts);
+	info.indirect_args_address            = indirect_args_address;
 	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

@@ -150,6 +150,7 @@ 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};
+	uint64_t           indirect_args_address       = 0;
 	bool               group_id[3]                = {false, false, false};
 	bool               dispatch_thread_dimensions = false;
 	bool               lds_storage                = false;