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);
+	const auto value  = AddressU32(ctx, mem, address, component);
-	return Decoder::ImageAddressComponentLayout(mem.image_sample_flags, component).bit_width == 16u
+	const auto layout = Decoder::ImageAddressComponentLayout(mem.image_sample_flags,
-	           ? EmitF16BitsToF32(ctx.state, value)
+	                                                         mem.image_dimension, component);
-	           : Unary(ctx.state, spv::OpBitcast, TypeF32(ctx.state), value);
+	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};
+	if (component < whole_vgprs) {
-	return {component * 32u, has_bias && component == bias_index ? 16u : 32u};
+		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);
+ImageAddressComponent ImageAddressComponentLayout(uint32_t flags, ImageDimension dimension,
-uint32_t ImageAddressDwordCount(uint32_t flags, uint32_t components);
+                                                  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 =
+	const auto count = Decoder::ImageAddressDwordCount(
-	    Decoder::ImageAddressDwordCount(inst.image_sample_flags, inst.image_address_components);
+	    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);

tests/ShaderRecompilerComputeTests.cpp

@@ -40055,6 +40055,51 @@ TestCase ImageA16BiasTakesWholeVgpr() {
   return test;
 }
 
+// Each A16 1D gradient, and each of the last two 3D ones, takes the low half of its own VGPR. Their
+// high halves hold 0.0625, which would move the coordinates if read as part of the address.
+template <u32 dim> TestCase ImageA16GradientsTakeOwnVgpr() {
+  using O = ShaderOpcode;
+  constexpr bool volume = dim == 2u;
+  constexpr std::array names{"ImageA16GradientsTakeOwnVgpr1D", "ImageA16GradientsTakeOwnVgpr3D",
+                             "ImageA16GradientsTakeOwnVgpr1DArray"};
+  constexpr std::array types{Prospero::ImageType::kColor1D, Prospero::ImageType::kColor3D,
+                             Prospero::ImageType::kColor1DArray};
+  constexpr std::array views{vk::ImageViewType::e1D, vk::ImageViewType::e3D,
+                             vk::ImageViewType::e1DArray};
+  TestCase test;
+  test.name = names[dim / 2u];
+  u32 vaddr = 20;
+  if constexpr (volume) {
+    AppendVMovU32(&test.code, vaddr++, 0); // dh.x, dh.y
+    AppendVMovU32(&test.code, vaddr++, 0); // dh.z, dv.x
+  }
+  AppendVMovLiteral(&test.code, vaddr++, 0x2c000000u); // 1D dh or 3D dv.y
+  AppendVMovLiteral(&test.code, vaddr++, 0x2c000000u); // 1D dv or 3D dv.z
+  if constexpr (volume) {
+    AppendVMovLiteral(&test.code, vaddr++, 0x36003900u); // x=0.625, y=0.375: texel (2, 1)
+    AppendVMovLiteral(&test.code, vaddr, 0x2c003800u);   // z=0.5
+  } else {
+    AppendVMovLiteral(&test.code, vaddr, 0x3c003900u); // x=0.625, slice 1: texel 2 of layer 1
+  }
+  test.code.push_back(EncodeMimg0(0x22, 0x1, 0, false, dim)); // image_sample_d
+  test.code.push_back(EncodeMimg1(0, 20, 0, 0, true));
+  AppendStoreVgpr(&test.code, 0, 0);
+  AppendEnd(&test.code);
+  test.opcodes = {O::V_MOV_B32, O::IMAGE_SAMPLE, O::BUFFER_STORE_DWORD, O::S_ENDPGM};
+  test.expected = {RampTexel(2, dim == 0u ? 0u : 1u)};
+  // Row y of the ramp is 1D-array layer y.
+  test.sampled_image_rgba = MakeRampImage();
+  test.image_height = volume ? 4u : 1u;
+  test.sampled_image_type = volume ? vk::ImageType::e3D : vk::ImageType::e1D;
+  test.sampled_image_view_type = views[dim / 2u];
+  test.sampled_image_layers = dim == 4u ? 2u : 1u;
+  test.user_data = MakeSampledTextureData(Prospero::BufferFormat::k32_32_32_32Float);
+  test.user_data[3] = static_cast<u32>(types[dim / 2u]) << 28u;
+  test.has_user_data = true;
+  test.required_spirv = {"OpImageSampleExplicitLod", "Grad"};
+  return test;
+}
+
 TestCase ImageGatherAddressVariants() {
   using O = ShaderOpcode;
   struct Gather {
@@ -45957,6 +46002,9 @@ std::vector<TestCase> MakeCases() {
   AddCase(ImageSampleOffsetsReadFourNeighbors<64>);
   AddCase(ImageSampleA16CompareBiasRdna2AddressOrder);
   AddCase(ImageA16BiasTakesWholeVgpr);
+  AddCase(ImageA16GradientsTakeOwnVgpr<0>);
+  AddCase(ImageA16GradientsTakeOwnVgpr<4>);
+  AddCase(ImageA16GradientsTakeOwnVgpr<2>);
   AddCase(ImageGatherCompareOpcodes);
   AddCase(ImageStoreVariants);
   AddCase(ImageD16StoreUnpacksHalfPairs);
@@ -52436,6 +52484,9 @@ int main(int argc, char **argv) {
     RunCase(&vulkan, ImageSampleAndGather());
     RunCase(&vulkan, ImageGatherAddressVariants());
     RunCase(&vulkan, ImageA16BiasTakesWholeVgpr());
+    RunCase(&vulkan, ImageA16GradientsTakeOwnVgpr<0>());
+    RunCase(&vulkan, ImageA16GradientsTakeOwnVgpr<4>());
+    RunCase(&vulkan, ImageA16GradientsTakeOwnVgpr<2>());
     RunCase(nullptr, ImageGatherCompareOpcodes());
     return 0;
   }

tests/shaderCfgTests.cpp

@@ -5754,6 +5754,9 @@ void TestImageAddressOperands() {
       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{0x22, 4, true, false, 4, 3},  // a VGPR per 1D gradient, then packed x/slice
+      Case{0x22, 1, true, false, 6, 3},  // packed 2D gradients and xy
+      Case{0x22, 2, true, false, 9, 6},  // packed first four 3D gradients, a VGPR per last two
       Case{0x4c, 1, false, true, 4, 4},  // full dref, xy, and LOD in NSA VGPRs
       Case{0x0e, 1, false, false, 1, 1}, // resource query mip
   };