shader: give A16 1D and late 3D gradients their own address VGPRs

Co-authored-by: shadowbeat070 <70843316+shadowbeat070@users.noreply.github.com>

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

@@ -38,7 +38,8 @@ bool HasFlag(const IR::MemoryInfo& mem, uint32_t flag) {
 
 uint32_t AddressU32(ValueEmitContext& ctx, const IR::MemoryInfo& mem, const IR::Inst& address,
                     uint32_t component) {
-	const auto layout = Decoder::ImageAddressComponentLayout(mem.image_sample_flags, component);
+	const auto layout = Decoder::ImageAddressComponentLayout(mem.image_sample_flags,
+	                                                         mem.image_dimension, component);
 	const auto packed = layout.bit_offset / 32u;
 	if (packed >= address.NumArgs()) return ConstantU32(ctx.state, 0);
 	auto value = ctx.Def(address.Arg(packed));
@@ -55,10 +56,11 @@ 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 Decoder::ImageAddressComponentLayout(mem.image_sample_flags, component).bit_width == 16u
-	           ? EmitF16BitsToF32(ctx.state, value)
-	           : Unary(ctx.state, spv::OpBitcast, TypeF32(ctx.state), value);
+	const auto value  = AddressU32(ctx, mem, address, component);
+	const auto layout = Decoder::ImageAddressComponentLayout(mem.image_sample_flags,
+	                                                         mem.image_dimension, component);
+	return layout.bit_width == 16u ? EmitF16BitsToF32(ctx.state, value)
+	                               : Unary(ctx.state, spv::OpBitcast, TypeF32(ctx.state), value);
 }
 
 ImageSampleLayout Layout(const IR::MemoryInfo& mem) {

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

@@ -319,21 +319,31 @@ uint32_t CountDmaskComponents(uint32_t dmask) {
 
 } // namespace
 
-ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, uint32_t component) {
+ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, ImageDimension dimension,
+                                                  uint32_t component) {
 	if ((flags & ImageSampleFlagA16) == 0u) return {component * 32u, 32u};
 	// A16 offsets, bias and PCF reference each take a whole VGPR, the bias as a half in its low
-	// 16 bits; the components after them pack two per VGPR.
+	// 16 bits. The other components pack two per VGPR, except that each 1D derivative and the last
+	// two 3D ones take the low half of their own VGPR.
 	const bool     has_bias    = (flags & ImageSampleFlagBias) != 0u;
 	const uint32_t bias_index  = (flags & ImageSampleFlagOffset) != 0u ? 1u : 0u;
 	const uint32_t whole_vgprs = std::popcount(
 	    flags & (ImageSampleFlagOffset | ImageSampleFlagBias | ImageSampleFlagCompare));
-	if (component >= whole_vgprs) return {(whole_vgprs + component) * 16u, 16u};
-	return {component * 32u, has_bias && component == bias_index ? 16u : 32u};
+	if (component < whole_vgprs) {
+		return {component * 32u, has_bias && component == bias_index ? 16u : 32u};
+	}
+	const uint32_t axes =
+	    (flags & ImageSampleFlagDerivative) != 0u ? ImageGradientComponents(dimension) : 0u;
+	const uint32_t unpaired = axes % 2u * 2u;
+	const uint32_t paired   = axes * 2u - unpaired;
+	const uint32_t index    = component - whole_vgprs;
+	const uint32_t half     = index + std::clamp(index, paired, paired + unpaired) - paired;
+	return {whole_vgprs * 32u + half * 16u, 16u};
 }
 
-uint32_t ImageAddressDwordCount(uint32_t flags, uint32_t components) {
+uint32_t ImageAddressDwordCount(uint32_t flags, ImageDimension dimension, uint32_t components) {
 	if (components == 0u) return 0u;
-	const auto last = ImageAddressComponentLayout(flags, components - 1u);
+	const auto last = ImageAddressComponentLayout(flags, dimension, components - 1u);
 	return (last.bit_offset + last.bit_width + 31u) / 32u;
 }
 

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

@@ -15,8 +15,9 @@ struct ImageAddressComponent {
 	uint32_t bit_width;
 };
 
-ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, uint32_t component);
-uint32_t ImageAddressDwordCount(uint32_t flags, uint32_t components);
+ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, ImageDimension dimension,
+                                                  uint32_t component);
+uint32_t ImageAddressDwordCount(uint32_t flags, ImageDimension dimension, uint32_t components);
 
 } // namespace Libs::Graphics::ShaderRecompiler::Decoder
 

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

@@ -322,8 +322,8 @@ IR::Value Translator::MakeImageAddress(const Decoder::Instruction& inst,
                                        const Decoder::Operand&     base) {
 	std::array<IR::Value, 13> components {};
 	components.fill(IR::Value(0u));
-	const auto count =
-	    Decoder::ImageAddressDwordCount(inst.image_sample_flags, inst.image_address_components);
+	const auto count = Decoder::ImageAddressDwordCount(
+	    inst.image_sample_flags, inst.image_dimension, inst.image_address_components);
 	EXIT_IF(count > components.size());
 	const auto nsa_components =
 	    std::min(inst.image_nsa_dwords * 4u, Decoder::MaxImageNsaAddressComponents);