shader: preserve native types in indirect image tables

src/graphics/shader/recompiler/backend/spirv/spirvEmitterImage.cpp

@@ -743,6 +743,87 @@ uint32_t EmitImageAccess(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t r
 	                             result_type, EmitRun);
 }
 
+template <typename Emit>
+uint32_t EmitImageResult(ValueEmitContext& ctx, const IR::Inst& inst, const Emit& emit) {
+	auto& state = ctx.state;
+	const auto& mem = ctx.Memory(inst);
+	const auto& image = state.program.info.images[mem.resource];
+	const bool gather = inst.GetOpcode() == IR::ValueOpcode::ImageGatherRaw;
+	const bool dref = HasFlag(mem, Decoder::ImageSampleFlagCompare);
+	const bool scalar = dref && !gather;
+	const bool mixed = std::ranges::any_of(image.indirect_resources, [&](uint32_t resource) {
+		return state.program.info.images[resource].numeric_class != image.numeric_class;
+	});
+	const auto native_type = [&](uint32_t resource) {
+		return scalar ? TypeF32(state) : ImageVectorType(
+		    state, dref ? Prospero::TextureNumericClass::Float
+		                : state.program.info.images[resource].numeric_class, 4);
+	};
+	const auto result_bits = [&](uint32_t resource, uint32_t value) {
+		auto candidate_mem = mem;
+		candidate_mem.resource = resource;
+		if (!scalar) {
+			value = gather ? UnpackImageGather(ctx, candidate_mem, value)
+			               : UnpackImageTexel(ctx, candidate_mem, value);
+		}
+		return ResultVector(ctx, value,
+		                    dref ? Prospero::TextureNumericClass::Float
+		                         : state.program.info.images[resource].numeric_class,
+		                    scalar, candidate_mem, gather);
+	};
+	// RDNA returns register bits. Convert per typed run only when the join mixes
+	// numeric classes; homogeneous accesses keep their single conversion after it.
+	const auto value = EmitImageAccess(
+	    ctx, inst, mixed ? TypeU32Vector(state, 4) : native_type(mem.resource),
+	    [&](uint32_t resource, uint32_t array_index) {
+		    const auto native = emit(resource, array_index, native_type(resource));
+		    return mixed ? result_bits(resource, native) : native;
+	    });
+	return mixed ? value : result_bits(mem.resource, value);
+}
+
+struct ImageSampler {
+	uint32_t id = 0;
+	bool dynamic = false;
+};
+
+std::array<ImageSampler, 2> LoadImageSamplers(ValueEmitContext& ctx, const IR::Inst& inst) {
+	auto& state = ctx.state;
+	const auto& mem = ctx.Memory(inst);
+	const auto& samplers = state.program.info.samplers;
+	const auto& root = samplers.at(mem.sampler);
+	std::array<ImageSampler, 2> result {};
+	const auto load = [&](const IR::ImageResource& image) {
+		const bool integer = image.numeric_class != Prospero::TextureNumericClass::Float;
+		auto& sampler = result[integer];
+		if (sampler.id != 0u) return;
+		uint32_t index = mem.sampler;
+		if (root.integer_border != integer) {
+			const auto found = std::ranges::find_if(samplers, [&](const auto& candidate) {
+				return candidate.snapshot_index == root.snapshot_index &&
+				       candidate.integer_border == integer;
+			});
+			if (found == samplers.end()) ctx.Fail(inst, "has no typed sampler variant");
+			index = static_cast<uint32_t>(found - samplers.begin());
+		}
+		uint32_t array_index = 0;
+		const auto& candidates = samplers[index].indirect_resources;
+		if (!candidates.empty()) {
+			const auto* handle = inst.Arg(1).ResolveInstruction();
+			if (handle == nullptr || handle->GetOpcode() != IR::ValueOpcode::GetSamplerResource)
+				ctx.Fail(inst, "has invalid finite sampler key provenance");
+			array_index = EmitIndexSwitch(
+			    state, ctx.Def(handle->Arg(0)), static_cast<uint32_t>(candidates.size()),
+			    TypeU32(state), [&](uint32_t ordinal) { return ConstantU32(state, candidates[ordinal]); });
+		}
+		sampler = {LoadSamplerDescriptor(state, index, array_index), array_index != 0u};
+	};
+	const auto& image = state.program.info.images[mem.resource];
+	load(image);
+	for (const auto candidate: image.indirect_resources) load(state.program.info.images[candidate]);
+	return result;
+}
+
 } // namespace
 
 void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
@@ -800,16 +881,14 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 		return;
 	}
 	if (op == IR::ValueOpcode::ImageRead) {
-		const auto numeric_class = image.numeric_class;
 		const auto condition = ctx.Arg(inst, 2);
 		ctx.Define(
 		    inst,
 		    EmitValueOrDefaultIfCondition(
 		        state, condition, TypeU32Vector(state, 4), ConstantU32CompositeZero(state, 4),
 		        [&]() {
-			        const auto color = EmitImageAccess(
+			        return EmitImageResult(
-			            ctx, inst, ImageVectorType(state, numeric_class, 4),
+			            ctx, inst, [&](uint32_t resource, uint32_t array_index, uint32_t result_type) {
-			            [&](uint32_t resource, uint32_t array_index) {
 				        const auto& candidate = state.program.info.images[resource];
 				        const auto dimension = candidate.dimension;
 				        const auto& dimension_info = ImageDimensionInfoFor(dimension);
@@ -818,19 +897,17 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 				        const auto coord = CoordU32(ctx, mem, *address, dimension);
 				        if (dimension_info.multisampled != 0u) {
 					        state.builder.AddFunction(
-					            spv::OpImageFetch, ImageVectorType(state, candidate.numeric_class, 4), fetched,
+					            spv::OpImageFetch, result_type, fetched,
 					            descriptor, coord, spv::ImageOperandsSampleMask,
 					            AddressU32(ctx, mem, *address, dimension_info.coordinate_components));
 				        } else {
 					        state.builder.AddFunction(
-					            spv::OpImageFetch, ImageVectorType(state, candidate.numeric_class, 4), fetched,
+					            spv::OpImageFetch, result_type, fetched,
 					            descriptor, coord, spv::ImageOperandsLodMask,
 					            LodU32(ctx, mem, *address, mem.image_dimension));
 				        }
 				        return fetched;
 			            });
-			        return ResultVector(ctx, UnpackImageTexel(ctx, mem, color), numeric_class,
-			                            false, mem);
 		        }));
 		return;
 	}
@@ -885,9 +962,6 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 				ctx.Fail(inst, "has an unsupported 1D-array gather");
 				return;
 			}
-			const auto result_numeric_class = dref ? Prospero::TextureNumericClass::Float
-			                                       : numeric_class;
-			const auto result_type = ImageVectorType(state, result_numeric_class, 4);
 			uint32_t component_or_dref;
 			if (dref) {
 				component_or_dref = layout.dref != NoImageComponent
@@ -896,9 +970,10 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 			} else {
 				component_or_dref = ConstantU32(state, ImageGatherSource(state, mem));
 			}
-			const auto sampler_id = LoadSamplerDescriptor(state, mem.sampler);
+			const auto samplers = LoadImageSamplers(ctx, inst);
-			const auto EmitGather = [&](uint32_t resource, uint32_t array_index) {
+			const auto EmitGather = [&](uint32_t resource, uint32_t array_index, uint32_t result_type) {
 				const auto& candidate = state.program.info.images[resource];
+				const auto sampler = samplers[candidate.numeric_class != Prospero::TextureNumericClass::Float];
 				const auto dimension = candidate.dimension;
 				const auto coord = CoordF32(
 				    ctx, mem, *address, layout.coord,
@@ -912,7 +987,8 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 					offset = PackedOffset(ctx, offset_value, dimension, operand_mask);
 				}
 				const auto EmitMip = [&](uint32_t mip) {
-					const auto sampled = MakeSampledImage(state, resource, sampler_id, mip, array_index);
+					const auto sampled = MakeSampledImage(state, resource, sampler.id, mip, array_index,
+					                                      sampler.dynamic);
 					const auto sample = state.builder.AllocateId();
 					const std::array operands {operand_mask, offset};
 					state.builder.AddFunction(
@@ -927,9 +1003,7 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 				                 candidate.mip_count, result_type, EmitMip)
 				           : EmitMip(0u);
 			};
-			const auto sample = EmitImageAccess(ctx, inst, result_type, EmitGather);
+			ctx.Define(inst, EmitImageResult(ctx, inst, EmitGather));
-			ctx.Define(inst, ResultVector(ctx, UnpackImageGather(ctx, mem, sample),
-			                              result_numeric_class, false, mem, true));
 			return;
 		}
 		const bool explicit_lod = HasFlag(mem, Decoder::ImageSampleFlagDerivative) ||
@@ -942,10 +1016,8 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 		} else if (dref) {
 			opcode = spv::OpImageSampleDrefImplicitLod;
 		}
-		uint32_t result_type = ImageVectorType(state, numeric_class, 4);
 		uint32_t dref_value  = 0;
 		if (dref) {
-			result_type = TypeF32(state);
 			dref_value  = ZeroF32(state);
 			if (layout.dref != NoImageComponent) {
 				dref_value = AddressF32(ctx, mem, *address, layout.dref);
@@ -971,23 +1043,14 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 			operands[1] = AddressF32(ctx, mem, *address, layout.bias);
 			operand_count = 2;
 		}
-		const auto& samplers = state.program.info.samplers.at(mem.sampler).indirect_resources;
-		uint32_t sampler_index = 0;
-		if (!samplers.empty()) {
-			const auto* handle = inst.Arg(1).ResolveInstruction();
-			if (handle == nullptr || handle->GetOpcode() != IR::ValueOpcode::GetSamplerResource)
-				ctx.Fail(inst, "has invalid finite sampler key provenance");
-			sampler_index = EmitIndexSwitch(
-			    state, ctx.Def(handle->Arg(0)), static_cast<uint32_t>(samplers.size()),
-			    TypeU32(state), [&](uint32_t ordinal) { return ConstantU32(state, samplers[ordinal]); });
-		}
 		if (layout.offset != NoImageComponent &&
 		    !offset_value.IsImmediate() && !state.maintenance8_enabled) {
 			ctx.Fail(inst, "requires maintenance8 for a dynamic sample texel offset");
 		}
-		const auto sampler_id = LoadSamplerDescriptor(state, mem.sampler, sampler_index);
+		const auto samplers = LoadImageSamplers(ctx, inst);
-		const auto EmitSample = [&](uint32_t resource, uint32_t array_index) {
+		const auto EmitSample = [&](uint32_t resource, uint32_t array_index, uint32_t result_type) {
 			const auto& candidate = state.program.info.images[resource];
+			const auto sampler = samplers[candidate.numeric_class != Prospero::TextureNumericClass::Float];
 			auto sample_operands = operands;
 			auto sample_operand_count = operand_count;
 			if (layout.offset != NoImageComponent) {
@@ -999,19 +1062,15 @@ void EmitImage(ValueEmitContext& ctx, const IR::Inst& inst) {
 			    CoordF32(ctx, mem, *address, layout.coord,
 			             ImageDimensionInfoFor(candidate.dimension).coordinate_components,
 			             candidate.cube);
-			const auto sampled = MakeSampledImage(state, resource, sampler_id, 0u, array_index,
+			const auto sampled = MakeSampledImage(state, resource, sampler.id, 0u, array_index,
-			                                      sampler_index != 0u);
+			                                      sampler.dynamic);
 			const auto sample = state.builder.AllocateId();
 			state.builder.AddFunction(opcode, result_type, sample, sampled, coord,
 			                          std::span(&dref_value, dref ? 1u : 0u),
 			                          std::span<const uint32_t>(sample_operands).first(sample_operand_count));
 			return sample;
 		};
-		auto result = EmitImageAccess(ctx, inst, result_type, EmitSample);
+		ctx.Define(inst, EmitImageResult(ctx, inst, EmitSample));
-		if (!dref) {
-			result = UnpackImageTexel(ctx, mem, result);
-		}
-		ctx.Define(inst, ResultVector(ctx, result, numeric_class, dref, mem));
 		return;
 	}
 	if (image_info.access == IR::ImageAccess::Atomic) {

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

@@ -785,8 +785,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, ResourceSna
 			}
 			const bool same_coordinates = image.dimension == image_class.dimension &&
 			                              image.cube == image_class.cube;
-			if (image.numeric_class != image_class.numeric_class ||
+			if ((!same_coordinates && !(is_2d(image.dimension) && is_2d(image_class.dimension))) ||
-			    (!same_coordinates && !(is_2d(image.dimension) && is_2d(image_class.dimension))) ||
 			    image.mip_count != image_class.mip_count ||
 			    image.conversion_format != image_class.conversion_format ||
 			    image.shader_swizzle != image_class.shader_swizzle) {
@@ -811,6 +810,10 @@ bool BuildSamplerPlan(const ShaderInfo& base, SamplerPlan& plan) {
 			return false;
 		}
 		usage[pair.sampler] |= 1u << static_cast<uint32_t>(ClassifySampler(base.images[pair.image]));
+		for (const auto candidate: base.images[pair.image].indirect_resources) {
+			if (candidate >= base.images.size()) return false;
+			usage[pair.sampler] |= 1u << static_cast<uint32_t>(ClassifySampler(base.images[candidate]));
+		}
 		for (const auto candidate: base.samplers[pair.sampler].indirect_resources) {
 			if (candidate >= base.samplers.size()) return false;
 			usage[candidate] |= usage[pair.sampler];

tests/ResourceMaterializationTests.cpp

@@ -974,6 +974,63 @@ void TestMixedSamplerVariantsShareRuntimeDescriptor() {
         "sampler variants retained stale or duplicated descriptors after refresh");
 }
 
+void TestIndirectImageSamplerClasses(bool mixed) {
+  using namespace Libs::Graphics::ShaderRecompiler::IR;
+  using Libs::Graphics::Prospero::BufferFormat;
+  using Libs::Graphics::Prospero::TextureNumericClass;
+  auto program = MixedSamplerProgram();
+  auto &first = program.descriptor_sources[0];
+  first.dwords[0] = Value(0x100u);
+  first.dwords[1] = Value(static_cast<uint32_t>(BufferFormat::k8UNorm) << 20u);
+  first.dwords[3] = program.descriptor_sources[1].dwords[3];
+  program.descriptor_sources[1].dwords[1] = Value(static_cast<uint32_t>(
+      mixed ? BufferFormat::k8UInt : BufferFormat::k8UNorm) << 20u);
+  DescriptorSource root;
+  root.dword_count = 8;
+  root.dwords.fill(Value(0u));
+  root.indirect_descriptor.emplace().sources = {0, 1, 0};
+  program.info.images.resize(1);
+  program.info.images[0].source = static_cast<uint32_t>(program.descriptor_sources.size());
+  program.descriptor_sources.push_back(root);
+  program.info.samplers.resize(1);
+  program.info.sampled_pairs.resize(1);
+  auto plan = ExtractResourcePlan(program);
+  std::array<uint32_t, 1> user_data{0x11111111u};
+  uint32_t reads = 0;
+  const SrtRuntime runtime{.user_data = user_data, .userdata = &reads,
+                           .read_specialization_memory = RejectSpecializationRead};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization) && reads == 0 &&
+            snapshot.images.size() == 2 && snapshot.images[0].dwords[0] == 0x100u &&
+            snapshot.images[1].dwords[0] == 1u &&
+            specialization.images[0].numeric_class == TextureNumericClass::Float &&
+            specialization.images[1].numeric_class ==
+                (mixed ? TextureNumericClass::Uint : TextureNumericClass::Float),
+        "indirect image materialization rejected or changed a typed candidate");
+  const auto offset = specialization.images[0].indirect_mapping_offset;
+  Check(snapshot.flattened_srt[offset] == 3u &&
+            snapshot.flattened_srt[offset + 2u] == 0u &&
+            snapshot.flattened_srt[offset + 4u] == 1u &&
+            snapshot.flattened_srt[offset + 6u] == 0u,
+        "mixed image classes changed selector ordinals or duplicated candidates");
+  ApplyResourceSpecialization(program, specialization);
+  Check(program.info.images[0].indirect_resources == std::vector<uint32_t>({0, 1}) &&
+            snapshot.samplers.size() == 1 && program.info.samplers.size() == (mixed ? 2u : 1u),
+        "indirect images lost candidates or created unnecessary sampler variants");
+  user_data[0] = 0x33333333u;
+  Check(MaterializeResources(plan, runtime, snapshot, specialization),
+        "mixed image sampler refresh failed");
+  for (uint32_t index = 0; index < program.info.samplers.size(); ++index) {
+    const auto &sampler = program.info.samplers[index];
+    Check(sampler.snapshot_index == 0u && sampler.source == plan.info.samplers[0].source &&
+              sampler.force_point_filtering == (index != 0u) &&
+              sampler.integer_border == (index != 0u) &&
+              snapshot.samplers[sampler.snapshot_index].dwords[0] == user_data[0],
+          "typed sampler variants lost filtering, border class, or shared live descriptor");
+  }
+}
+
 } // namespace
 
 namespace Common {
@@ -1007,6 +1064,7 @@ int main() {
     TestEmptyTraversalSkipsUnmappedDescriptor(independent_mask);
   TestCyclicControlPredicatesRemainUnknown();
   TestMixedSamplerVariantsShareRuntimeDescriptor();
+  for (const bool mixed : {false, true}) TestIndirectImageSamplerClasses(mixed);
   std::puts("ResourceMaterializationTests: all cases passed");
   return 0;
 }

tests/ShaderRecompilerComputeTests.cpp

@@ -37279,6 +37279,46 @@ void CheckIndirectImageKeySwitch(VulkanHarness &vulkan) {
       const auto actual = vulkan.ReadBuffer(test.name, buffer, test.expected.size());
       vulkan.DestroyBuffer(&buffer);
       CompareWords(test, "root, native slots, array runs and mapping bounds", test.expected, actual);
+      if (wave_size == 32u && !homogeneous) {
+        // A single integer candidate must keep raw register bits across the
+        // same descriptor runs used by samples, gathers, and image fetches.
+        const auto coordinate_conversions = CountText(text, "OpConvertUToF");
+        auto &typed = compiled.program.info.images[3];
+        typed.numeric_class = Prospero::TextureNumericClass::Uint;
+        if (!fetch) {
+          auto integer_sampler = compiled.program.info.samplers[0];
+          integer_sampler.force_point_filtering = true;
+          integer_sampler.integer_border = true;
+          compiled.program.info.samplers.push_back(integer_sampler);
+        }
+        compiled.program.shader_info_complete = false;
+        CollectShaderInfo(compiled.program, {.compute = &compute});
+        const auto mixed_spirv = ShaderRecompiler::Spirv::EmitProgram(
+            compiled.program, {.compute = &compute});
+        ValidateSpirv("IndirectImageMixedNumeric", mixed_spirv);
+        u32 unsigned_type = 0u, raw_vector = 0u;
+        bool raw_join = false;
+        for (size_t offset = 5; offset < mixed_spirv.size();) {
+          const auto words = std::span<const u32>(mixed_spirv).subspan(
+              offset, mixed_spirv[offset] >> 16u);
+          const auto opcode = static_cast<spv::Op>(words[0] & 0xffffu);
+          if (opcode == spv::OpTypeInt && words[2] == 32u && words[3] == 0u)
+            unsigned_type = words[1];
+          if (opcode == spv::OpTypeVector && words[2] == unsigned_type && words[3] == 4u)
+            raw_vector = words[1];
+          if (opcode == spv::OpPhi && words[1] == raw_vector) raw_join = true;
+          offset += words.size();
+        }
+        Require(test.name, "mixed numeric image result", raw_join &&
+                    tools.Disassemble(mixed_spirv, &text) &&
+                    CountText(text, image_opcode) == 4u * halves &&
+                    CountText(text, "OpConvertUToF") == coordinate_conversions,
+                "mixed image classes lost raw result bits or expanded per image");
+        typed.numeric_class = Prospero::TextureNumericClass::Float;
+        if (!fetch) compiled.program.info.samplers.pop_back();
+        compiled.program.shader_info_complete = false;
+        CollectShaderInfo(compiled.program, {.compute = &compute});
+      }
     }
     std::printf("[compute] %-32s ok\n", test.name);
   }