shader: fix packed image address reads

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

@@ -1,5 +1,6 @@
 #include "graphics/guest_gpu/gpu_format.h"
 #include "graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h"
+#include "graphics/shader/recompiler/frontend/decode/ImageOps.h"
 
 #include <algorithm>
 #include <bit>
@@ -47,45 +48,14 @@ bool HasFlag(const IR::MemoryInfo& mem, uint32_t flag) {
 	return (mem.image_sample_flags & flag) != 0u;
 }
 
-bool UsesA16(const IR::MemoryInfo& mem) {
-	return HasFlag(mem, Decoder::ImageSampleFlagA16);
-}
-
-bool ComponentUsesA16(const IR::MemoryInfo& mem, uint32_t component) {
-	if (!UsesA16(mem)) return false;
-	uint32_t cursor = 0;
-	if (HasFlag(mem, Decoder::ImageSampleFlagOffset)) {
-		if (component == cursor) return false;
-		cursor++;
-	}
-	if (HasFlag(mem, Decoder::ImageSampleFlagBias)) {
-		if (component == cursor) return true;
-		cursor++;
-	}
-	if (HasFlag(mem, Decoder::ImageSampleFlagCompare) && component == cursor) return false;
-	return true;
-}
-
-uint32_t HalfComponent(const IR::MemoryInfo& mem, uint32_t component) {
-	if (!UsesA16(mem)) return component * 2u;
-	uint32_t half = 0;
-	for (uint32_t index = 0; index < component; index++) {
-		const auto width = ComponentUsesA16(mem, index) ? 1u : 2u;
-		if (width == 2u && (half & 1u) != 0u) half++;
-		half += width;
-	}
-	if (!ComponentUsesA16(mem, component) && (half & 1u) != 0u) half++;
-	return half;
-}
-
 uint32_t AddressU32(ValueEmitContext& ctx, const IR::MemoryInfo& mem, const IR::Inst& address,
                     uint32_t component) {
-	const auto half   = HalfComponent(mem, component);
+	const auto layout = Decoder::ImageAddressComponentLayout(mem.image_sample_flags, component);
-	const auto packed = half / 2u;
+	const auto packed = layout.bit_offset / 32u;
 	if (packed >= address.NumArgs()) return ConstantU32(ctx.state, 0);
 	auto value = ctx.Def(address.Arg(packed));
-	if (ComponentUsesA16(mem, component)) {
+	if (layout.bit_width == 16u) {
-		if ((half & 1u) != 0u) {
+		if ((layout.bit_offset & 31u) != 0u) {
 			value = Binary(ctx.state, OpShiftRightLogical, TypeU32(ctx.state), value,
 			               ConstantU32(ctx.state, 16));
 		}
@@ -98,7 +68,7 @@ uint32_t AddressU32(ValueEmitContext& ctx, const IR::MemoryInfo& mem, const IR::
 uint32_t AddressF32(ValueEmitContext& ctx, const IR::MemoryInfo& mem, const IR::Inst& address,
                     uint32_t component) {
 	const auto value = AddressU32(ctx, mem, address, component);
-	return ComponentUsesA16(mem, component)
+	return Decoder::ImageAddressComponentLayout(mem.image_sample_flags, component).bit_width == 16u
 	           ? EmitF16BitsToF32(ctx.state, value)
 	           : Unary(ctx.state, OpBitcast, TypeF32(ctx.state), value);
 }

src/graphics/shader/recompiler/frontend/decode/ImageOps.cpp

@@ -274,6 +274,7 @@ uint32_t DecodeMimgAddressComponents(uint32_t opcode, ImageDimension dimension,
 	}
 
 	switch (opcode) {
+		case 0x0eu: return 1u;
 		case 0x01u:
 		case 0x09u: return ImageCoordComponents(dimension) + 1u;
 		case 0x00u:
@@ -297,6 +298,36 @@ bool IsSingleDmaskBit(uint32_t dmask) {
 
 } // namespace
 
+ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, uint32_t component) {
+	const auto width = [flags](uint32_t index) {
+		if ((flags & ImageSampleFlagA16) == 0u) return 32u;
+		uint32_t cursor = 0;
+		if ((flags & ImageSampleFlagOffset) != 0u) {
+			if (index == cursor++) return 32u;
+		}
+		if ((flags & ImageSampleFlagBias) != 0u) {
+			if (index == cursor++) return 16u;
+		}
+		if ((flags & ImageSampleFlagCompare) != 0u && index == cursor) return 32u;
+		return 16u;
+	};
+	uint32_t offset = 0;
+	for (uint32_t index = 0; index < component; index++) {
+		const auto bits = width(index);
+		if (bits == 32u) offset = (offset + 31u) & ~31u;
+		offset += bits;
+	}
+	const auto bits = width(component);
+	if (bits == 32u) offset = (offset + 31u) & ~31u;
+	return {offset, bits};
+}
+
+uint32_t ImageAddressDwordCount(uint32_t flags, uint32_t components) {
+	if (components == 0u) return 0u;
+	const auto last = ImageAddressComponentLayout(flags, components - 1u);
+	return (last.bit_offset + last.bit_width + 31u) / 32u;
+}
+
 void DecodeMimg(uint32_t pc, std::span<const uint32_t> code, uint32_t word_index,
                 Instruction& inst) {
 	const uint32_t word0      = code[word_index];

src/graphics/shader/recompiler/frontend/decode/ImageOps.h

@@ -10,6 +10,14 @@ void DecodeMimg(uint32_t pc, std::span<const uint32_t> code, uint32_t word_index
 
 const char* MimgSampleOpcodeName(uint32_t opcode);
 
+struct ImageAddressComponent {
+	uint32_t bit_offset;
+	uint32_t bit_width;
+};
+
+ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, uint32_t component);
+uint32_t ImageAddressDwordCount(uint32_t flags, uint32_t components);
+
 } // namespace Libs::Graphics::ShaderRecompiler::Decoder
 
 #endif /* EMULATOR_INCLUDE_EMULATOR_GRAPHICS_SHADER_RECOMPILER_IMAGEOPS_H_ */

src/graphics/shader/recompiler/frontend/translate/Memory.cpp

@@ -1,4 +1,5 @@
 #include "graphics/shader/recompiler/frontend/translate/Translator.h"
+#include "graphics/shader/recompiler/frontend/decode/ImageOps.h"
 
 #include <algorithm>
 #include <array>
@@ -330,11 +331,14 @@ IR::Value Translator::MakeImageAddress(const Decoder::Instruction& inst,
                                        const Decoder::Operand&     base) {
 	const auto                memory = MemoryInfoFromDecoded(inst);
 	std::array<IR::Value, 13> components {};
-	components[0] = ReadRawU32(PlainOperand(base));
+	components.fill(IR::Value(0u));
+	const auto count =
+	    Decoder::ImageAddressDwordCount(memory.image_sample_flags, memory.image_address_components);
+	EXIT_IF(count > components.size());
 	const auto nsa_components =
 	    std::min(memory.image_nsa_dwords * 4u, Decoder::MaxImageNsaAddressComponents);
-	for (uint32_t index = 1; index < components.size(); index++) {
+	for (uint32_t index = 0; index < count; index++) {
-		if (index - 1u < nsa_components) {
+		if (index != 0u && index - 1u < nsa_components) {
 			components[index] =
 			    ir.GetVectorReg(static_cast<IR::VectorReg>(memory.image_nsa_addr[index - 1u]));
 		} else {

tests/shaderCfgTests.cpp

@@ -4720,6 +4720,65 @@ void TestNewShaderRecompilerCubeSampleCoordinates() {
   CheckSpirvBinaryValidates(result.spirv);
 }
 
+void TestImageAddressOperands() {
+  using namespace ShaderRecompiler;
+  struct Case {
+    uint32_t opcode;
+    uint32_t dimension;
+    bool a16;
+    bool nsa;
+    uint32_t logical_components;
+    uint32_t dwords;
+  };
+  constexpr std::array cases{
+      Case{0x27, 1, false, false, 2, 2}, // xy, explicit level zero
+      Case{0x22, 2, false, false, 9, 9}, // 3D gradients and xyz
+      Case{0x20, 1, true, false, 2, 1},  // packed xy
+      Case{0x3d, 1, true, false, 5, 4},  // offset, half bias, full dref, packed xy
+      Case{0x20, 2, true, true, 3, 2},   // packed xy/z in nonconsecutive VGPRs
+      Case{0x0e, 1, false, false, 1, 1}, // resource query mip
+  };
+  for (const auto &test : cases) {
+    std::vector<uint32_t> shader;
+    for (uint32_t reg = 0; reg < 16; reg++) {
+      shader.push_back(EncodeVop1(0x01, reg, 129 + reg));
+    }
+    shader.push_back(EncodeMimg0(test.opcode, 0x1, false, test.dimension) |
+                     (test.nsa ? 1u << 1u : 0u));
+    shader.push_back(EncodeMimg1(48, 0, 4, 0, test.a16));
+    if (test.nsa) shader.push_back(0x0f0b0703u);
+    shader.push_back(EncodeSopp(0x01));
+    Decoder::Program decoded;
+    Decoder::DecodeProgram(shader, decoded);
+    Check(decoded.instructions[16].image_address_components == test.logical_components,
+          "MIMG decoded address count does not match its logical operands");
+    ShaderComputeInputInfo compute = RegressionComputeInputInfo();
+    Frontend::TranslateOptions options{};
+    options.stage = ShaderType::Compute;
+    options.compute = &compute;
+    auto program = Frontend::TranslateProgram(decoded, CFG::BuildGraph(decoded), options);
+    IR::RewriteToSsa(program.blocks);
+    IR::ConstantPropagationPass(program.blocks);
+    uint32_t addresses = 0;
+    for (const auto *block : program.blocks) {
+      for (const auto &inst : *block) {
+        if (inst.GetOpcode() != IR::ValueOpcode::MakeImageAddress) continue;
+        addresses++;
+        for (uint32_t index = 0; index < inst.NumArgs(); index++) {
+          const auto value = inst.Arg(index).Resolve();
+          const auto reg = test.nsa && index > 0 ? 3u + (index - 1u) * 4u : index;
+          const auto expected = index < test.dwords ? reg + 1u : 0u;
+          Check(value.IsImmediate() && value.GetType() == IR::Type::U32 &&
+                    value.U32() == expected,
+                "MIMG address read an unused VGPR or lost a packed/NSA operand");
+        }
+      }
+    }
+    Check(addresses == 1, "MIMG address regression did not reach translation");
+  }
+}
+
+
 void TestNewShaderRecompilerImageSampleVariants() {
   const uint32_t shader[] = {
       EncodeMimg0(0x24, 0xf),
@@ -11939,6 +11998,7 @@ int main() {
   // ShaderRecompilerComputeTests; keep the distinct decoder contract checks
   // here.
   TestNewShaderDecoderArchitecture();
+  TestImageAddressOperands();
   TestSopkCompareImmediateExtension();
   TestNewShaderRecompilerCapturedVopcSdwaCmpxClass();
   TestNewShaderRecompilerIrLookupMissFailsExplicitly();