shader: specialize cube image descriptors (#134)

* shader: specialize cube image descriptors

Track whether image descriptors refer to cube maps during resource specialization, and apply the coordinate offset conversion when sampling cube maps as 2D image arrays in SPIR-V emission.

* shader: fix cube array coordinate lowering

---------

Co-authored-by: nmzik <Nmzik@mail.ru>

src/graphics/shader/recompiler/emitter/spirvEmitterImageHelpers.cpp

@@ -2,6 +2,66 @@
 #include "graphics/shader/recompiler/emitter/spirvEmitterInternal.h"
 
 namespace Libs::Graphics::ShaderRecompiler::Spirv::Emitter {
+namespace {
+
+uint32_t EmitCubeAxisF32(EmitterState& state, uint32_t value) {
+	const auto normalized = state.builder.AllocateId();
+	state.builder.AddFunction(
+	    {OpFSub, state.float_type, normalized, value, ConstantF32(state, 0x3f800000u)});
+	return normalized;
+}
+
+uint32_t EmitCubeLayerF32(EmitterState& state, uint32_t face_id) {
+	// Sampled RDNA2 cubemaps encode face_id as slice * 8 + face. The native
+	// 2D-array view stores six contiguous faces per slice, so remove the two
+	// reserved face IDs from every preceding slice.
+	const auto guest_layer = state.builder.AllocateId();
+	const auto slice       = state.builder.AllocateId();
+	const auto padding     = state.builder.AllocateId();
+	const auto host_layer  = state.builder.AllocateId();
+	const auto result      = state.builder.AllocateId();
+	state.builder.AddFunction({OpConvertFToU, state.uint_type, guest_layer, face_id});
+	state.builder.AddFunction(
+	    {OpShiftRightLogical, state.uint_type, slice, guest_layer, ConstantU32(state, 3)});
+	state.builder.AddFunction(
+	    {OpShiftLeftLogical, state.uint_type, padding, slice, ConstantU32(state, 1)});
+	state.builder.AddFunction({OpISub, state.uint_type, host_layer, guest_layer, padding});
+	state.builder.AddFunction({OpConvertUToF, state.float_type, result, host_layer});
+	return result;
+}
+
+uint32_t EmitImageCoordF32Impl(EmitterState& state, const IR::Instruction& inst,
+                               const IR::Operand& address, uint32_t first_component,
+                               uint32_t components) {
+	auto x = EmitImageAddressFloatLoad(state, inst, address, first_component);
+	if (components == 1u) {
+		return x;
+	}
+	auto y = inst.memory.image_address_components > first_component + 1u
+	             ? EmitImageAddressFloatLoad(state, inst, address, first_component + 1u)
+	             : EmitZeroF32(state);
+	if (inst.memory.image_cube) {
+		// RDNA2 sampled cubemap S/T coordinates are biased by +1 relative to
+		// normalized 2D-array coordinates.
+		x = EmitCubeAxisF32(state, x);
+		y = EmitCubeAxisF32(state, y);
+	}
+	const auto coord = state.builder.AllocateId();
+	if (components == 3u) {
+		auto z = inst.memory.image_address_components > first_component + 2u
+		             ? EmitImageAddressFloatLoad(state, inst, address, first_component + 2u)
+		             : EmitZeroF32(state);
+		if (inst.memory.image_cube) {
+			z = EmitCubeLayerF32(state, z);
+		}
+		state.builder.AddFunction({OpCompositeConstruct, state.vec3_float_type, coord, x, y, z});
+	} else {
+		state.builder.AddFunction({OpCompositeConstruct, state.vec2_float_type, coord, x, y});
+	}
+	return coord;
+}
+
+} // namespace
 
 bool HasImageSampleFlag(const IR::Instruction& inst, uint32_t flag) {
 	return (inst.memory.image_sample_flags & flag) != 0;
@@ -21,7 +81,7 @@ ImageSampleLayout MakeImageSampleLayout(const IR::Instruction& inst, ImageViewKi
 	}
 	if (HasImageSampleFlag(inst, Decoder::ImageSampleFlagDerivative)) {
 		const auto components = ImageViewSpatialComponents(view);
-		layout.grad_x = cursor;
+		layout.grad_x         = cursor;
 		cursor += components;
 		layout.grad_y = cursor;
 		cursor += components;
@@ -36,24 +96,8 @@ ImageSampleLayout MakeImageSampleLayout(const IR::Instruction& inst, ImageViewKi
 
 uint32_t EmitImageCoordF32(EmitterState& state, const IR::Instruction& inst,
                            const ImageSampleLayout& layout, ImageViewKind view) {
-	const auto x          = EmitImageAddressFloatLoad(state, inst, inst.src[0], layout.coord);
-	const auto components = ImageViewCoordinateComponents(view);
-	if (components == 1u) {
-		return x;
-	}
-	const auto y = inst.memory.image_address_components > layout.coord + 1u
-	                       ? EmitImageAddressFloatLoad(state, inst, inst.src[0], layout.coord + 1u)
-	                       : EmitZeroF32(state);
-	const auto coord = state.builder.AllocateId();
-	if (components == 3u) {
-		const auto z = inst.memory.image_address_components > layout.coord + 2u
-		                   ? EmitImageAddressFloatLoad(state, inst, inst.src[0], layout.coord + 2u)
-		                   : EmitZeroF32(state);
-		state.builder.AddFunction({OpCompositeConstruct, state.vec3_float_type, coord, x, y, z});
-	} else {
-		state.builder.AddFunction({OpCompositeConstruct, state.vec2_float_type, coord, x, y});
-	}
-	return coord;
+	return EmitImageCoordF32Impl(state, inst, inst.src[0], layout.coord,
+	                             ImageViewCoordinateComponents(view));
 }
 
 uint32_t EmitImageLodF32(EmitterState& state, const IR::Instruction& inst,
@@ -95,10 +139,10 @@ uint32_t EmitImageGradientF32(EmitterState& state, const IR::Instruction& inst,
 	                   : EmitZeroF32(state);
 	const auto grad = state.builder.AllocateId();
 	if (components == 3u) {
-		const auto z = inst.memory.image_address_components > first_component + 2u
-		                   ? EmitImageAddressFloatLoad(state, inst, inst.src[0],
-		                                               first_component + 2u)
-		                   : EmitZeroF32(state);
+		const auto z =
+		    inst.memory.image_address_components > first_component + 2u
+		        ? EmitImageAddressFloatLoad(state, inst, inst.src[0], first_component + 2u)
+		        : EmitZeroF32(state);
 		state.builder.AddFunction({OpCompositeConstruct, state.vec3_float_type, grad, x, y, z});
 	} else {
 		state.builder.AddFunction({OpCompositeConstruct, state.vec2_float_type, grad, x, y});
@@ -120,8 +164,7 @@ uint32_t EmitImagePackedOffsetI32(EmitterState& state, const IR::Instruction& in
 			state.builder.AddFunction(
 			    {OpCompositeConstruct, state.vec3_int_type, ret, zero, zero, zero});
 		} else {
-			state.builder.AddFunction(
-			    {OpCompositeConstruct, state.vec2_int_type, ret, zero, zero});
+			state.builder.AddFunction({OpCompositeConstruct, state.vec2_int_type, ret, zero, zero});
 		}
 		return ret;
 	}
@@ -131,18 +174,18 @@ uint32_t EmitImagePackedOffsetI32(EmitterState& state, const IR::Instruction& in
 	const auto offset_x    = state.builder.AllocateId();
 	state.builder.AddFunction({OpBitcast, state.int_type, packed_i32, packed_bits});
 	state.builder.AddFunction({OpBitFieldSExtract, state.int_type, offset_x, packed_i32,
-	                            ConstantI32(state, 0), ConstantI32(state, 6)});
+	                           ConstantI32(state, 0), ConstantI32(state, 6)});
 	if (components == 1u) {
 		return offset_x;
 	}
 	const auto offset_y = state.builder.AllocateId();
 	const auto offset   = state.builder.AllocateId();
 	state.builder.AddFunction({OpBitFieldSExtract, state.int_type, offset_y, packed_i32,
-	                            ConstantI32(state, 8), ConstantI32(state, 6)});
+	                           ConstantI32(state, 8), ConstantI32(state, 6)});
 	if (components == 3u) {
 		const auto offset_z = state.builder.AllocateId();
 		state.builder.AddFunction({OpBitFieldSExtract, state.int_type, offset_z, packed_i32,
-		                            ConstantI32(state, 16), ConstantI32(state, 6)});
+		                           ConstantI32(state, 16), ConstantI32(state, 6)});
 		state.builder.AddFunction(
 		    {OpCompositeConstruct, state.vec3_int_type, offset, offset_x, offset_y, offset_z});
 	} else {
@@ -158,7 +201,7 @@ uint32_t EmitImageCoordU32(EmitterState& state, const IR::Instruction& inst, Ima
 	if (components == 1u) {
 		return x;
 	}
-	const auto y = inst.memory.image_address_components > 1u
+	const auto y     = inst.memory.image_address_components > 1u
 	                       ? EmitImageAddressValueLoad(state, inst, inst.src[1], 1)
 	                       : ConstantU32(state, 0);
 	const auto coord = state.builder.AllocateId();
@@ -180,7 +223,7 @@ uint32_t EmitImageLoadCoordU32(EmitterState& state, const IR::Instruction& inst,
 	if (components == 1u) {
 		return x;
 	}
-	const auto y = inst.memory.image_address_components > 1u
+	const auto y     = inst.memory.image_address_components > 1u
 	                       ? EmitImageAddressValueLoad(state, inst, inst.src[0], 1)
 	                       : ConstantU32(state, 0);
 	const auto coord = state.builder.AllocateId();
@@ -209,24 +252,8 @@ uint32_t EmitImageMipLodU32(EmitterState& state, const IR::Instruction& inst,
 
 uint32_t EmitImageQueryCoordF32(EmitterState& state, const IR::Instruction& inst,
                                 ImageViewKind view) {
-	const auto x          = EmitImageAddressFloatLoad(state, inst, inst.src[0], 0);
-	const auto components = ImageViewCoordinateComponents(view);
-	if (components == 1u) {
-		return x;
-	}
-	const auto y = inst.memory.image_address_components > 1u
-	                       ? EmitImageAddressFloatLoad(state, inst, inst.src[0], 1)
-	                       : EmitZeroF32(state);
-	const auto coord = state.builder.AllocateId();
-	if (components == 3u) {
-		const auto z = inst.memory.image_address_components > 2u
-		                   ? EmitImageAddressFloatLoad(state, inst, inst.src[0], 2)
-		                   : EmitZeroF32(state);
-		state.builder.AddFunction({OpCompositeConstruct, state.vec3_float_type, coord, x, y, z});
-	} else {
-		state.builder.AddFunction({OpCompositeConstruct, state.vec2_float_type, coord, x, y});
-	}
-	return coord;
+	// OpImageQueryLod takes only the spatial coordinates, even for arrayed images.
+	return EmitImageCoordF32Impl(state, inst, inst.src[0], 0, ImageViewSpatialComponents(view));
 }
 
 uint32_t DmaskComponentIndex(uint32_t dmask, uint32_t component) {

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

@@ -12,7 +12,7 @@ namespace {
 
 constexpr uint64_t AddressMask = 0x0000ffffffffffffull;
 
-Decoder::ImageDimension DescriptorDimension(const DescriptorValue&       descriptor,
+Decoder::ImageDimension DescriptorDimension(const DescriptorValue&  descriptor,
                                             Decoder::ImageDimension requested) {
 	const bool is_array = requested == Decoder::ImageDimension::Dim1DArray ||
 	                      requested == Decoder::ImageDimension::Dim2DArray;
@@ -51,8 +51,7 @@ bool ValidImageDescriptor(const DescriptorValue& descriptor) {
 		const auto base_level = (descriptor.dwords[3] >> 12u) & 0xfu;
 		const auto fragments  = (descriptor.dwords[3] >> 16u) & 0xfu;
 		const auto max_mip    = (descriptor.dwords[5] >> 4u) & 0xfu;
-		return base_level == 0 && fragments >= 1 && fragments <= 3 &&
-		       max_mip == fragments;
+		return base_level == 0 && fragments >= 1 && fragments <= 3 && max_mip == fragments;
 	}
 	return true;
 }
@@ -61,6 +60,11 @@ uint32_t DescriptorImageSwizzle(const DescriptorValue& descriptor) {
 	return descriptor.dwords[3] & 0xfffu;
 }
 
+bool DescriptorIsCube(const DescriptorValue& descriptor) {
+	return static_cast<Prospero::ImageType>((descriptor.dwords[3] >> 28u) & 0xfu) ==
+	       Prospero::ImageType::kCube;
+}
+
 bool DecodeBufferDescriptor(const DescriptorValue& descriptor, ShaderBufferResource& result) {
 	if (descriptor.dword_count != std::size(result.fields)) {
 		return false;
@@ -171,12 +175,13 @@ bool ValidateResourceSpecialization(const Program& program, const ResourceSnapsh
 		const auto& image      = program.info.images[i];
 		const auto& descriptor = snapshot.images[i];
 		if (NullImageDescriptor(descriptor)) {
-			bool canonical_kind = image.kind == ResourceKind::Image ||
-			                      image.kind == ResourceKind::StorageImage;
+			bool canonical_kind =
+			    image.kind == ResourceKind::Image || image.kind == ResourceKind::StorageImage;
 			if (image.atomic) {
 				canonical_kind = image.kind == ResourceKind::StorageImageUint;
 			}
-			if (image.dimension != Decoder::ImageDimension::Dim2D || !canonical_kind) {
+			if (image.dimension != Decoder::ImageDimension::Dim2D || image.cube ||
+			    !canonical_kind) {
 				if (error != nullptr) {
 					*error = fmt::format(
 					    "image descriptor {} no longer matches canonical null specialization", i);
@@ -186,7 +191,8 @@ bool ValidateResourceSpecialization(const Program& program, const ResourceSnapsh
 			continue;
 		}
 		const auto dimension = DescriptorDimension(descriptor, image.dimension);
-		if (dimension == Decoder::ImageDimension::Unknown || dimension != image.dimension) {
+		if (dimension == Decoder::ImageDimension::Unknown || dimension != image.dimension ||
+		    DescriptorIsCube(descriptor) != image.cube) {
 			if (error != nullptr) {
 				*error =
 				    fmt::format("image descriptor {} no longer matches specialized dimension", i);
@@ -361,6 +367,7 @@ bool SpecializeResources(Program& program, const ResourceSnapshot& snapshot, std
 		auto&       image      = next.images[i];
 		if (NullImageDescriptor(descriptor)) {
 			image.dimension = Decoder::ImageDimension::Dim2D;
+			image.cube      = false;
 			switch (image.kind) {
 				case ResourceKind::ImageUint: image.kind = ResourceKind::Image; break;
 				case ResourceKind::StorageImageUint:
@@ -386,6 +393,7 @@ bool SpecializeResources(Program& program, const ResourceSnapshot& snapshot, std
 			return false;
 		}
 		image.dimension = descriptor_dimension;
+		image.cube      = DescriptorIsCube(descriptor);
 		if (image.kind == ResourceKind::StorageImage ||
 		    image.kind == ResourceKind::StorageImageUint) {
 			image.storage_swizzle = DescriptorImageSwizzle(descriptor);
@@ -402,6 +410,7 @@ bool SpecializeResources(Program& program, const ResourceSnapshot& snapshot, std
 		std::reference_wrapper<Instruction> inst;
 		ResourceKind                        kind;
 		Decoder::ImageDimension             dimension;
+		bool                                cube;
 	};
 	std::vector<ImagePatch> patches;
 	for (auto& block: program.blocks) {
@@ -420,13 +429,14 @@ bool SpecializeResources(Program& program, const ResourceSnapshot& snapshot, std
 				return false;
 			}
 			const auto& image = next.images[inst.memory.resource];
-			patches.push_back({std::ref(inst), image.kind, image.dimension});
+			patches.push_back({std::ref(inst), image.kind, image.dimension, image.cube});
 		}
 	}
 	program.info = std::move(next);
 	for (const auto& patch: patches) {
 		patch.inst.get().memory.kind            = patch.kind;
 		patch.inst.get().memory.image_dimension = patch.dimension;
+		patch.inst.get().memory.image_cube      = patch.cube;
 	}
 	return true;
 }

src/graphics/shader/recompiler/ir/ShaderIR.h

@@ -432,6 +432,7 @@ struct MemoryInfo {
 	bool     typed           = false;
 	bool     formatted       = false;
 	bool     image_has_mip   = false;
+	bool     image_cube      = false;
 	bool     glc             = false;
 	bool     slc             = false;
 	bool     idxen           = false;
@@ -607,6 +608,7 @@ struct ImageResource {
 	bool                    written         = false;
 	bool                    atomic          = false;
 	bool                    depth_compare   = false;
+	bool                    cube            = false;
 
 	bool operator==(const ImageResource& other) const = default;
 };

tests/ResourceTrackingTests.cpp

@@ -1217,6 +1217,34 @@ void TestResourceSpecializationIsTypedAndTransactional() {
                 Decoder::ImageDimension::Dim2DArray,
         "array MIMG intent did not produce a 2D-array view");
 
+  Program cube_view;
+  cube_view.stage = ShaderType::Compute;
+  cube_view.blocks.resize(1);
+  cube_view.blocks[0].instructions = {
+      ImageUse(0x24, Opcode::ImageLoad, ResourceKind::Image,
+               Decoder::ImageDimension::Dim2DArray)};
+  Prepare(cube_view);
+  auto cube_snapshot = array_2d_snapshot;
+  cube_snapshot.images[0].dwords[3] =
+      Prospero::GpuEnumValue(Prospero::ImageType::kCube) << 28u;
+  Check(SpecializeResources(cube_view, cube_snapshot, &error) &&
+            ValidateResourceSpecialization(cube_view, cube_snapshot, &error) &&
+            cube_view.info.images[0].cube &&
+            cube_view.blocks[0].instructions[0].memory.image_cube,
+        "cube descriptor identity did not reach the specialized image and IR");
+  auto array_after_cube = cube_snapshot;
+  array_after_cube.images[0].dwords[3] =
+      Prospero::GpuEnumValue(Prospero::ImageType::kColor2DArray) << 28u;
+  Check(!ValidateResourceSpecialization(cube_view, array_after_cube, &error),
+        "2D-array descriptor reused a cube-coordinate specialization");
+  auto null_after_cube = cube_snapshot;
+  null_after_cube.images[0].dwords.fill(0);
+  Check(SpecializeResources(cube_view, null_after_cube, &error) &&
+            ValidateResourceSpecialization(cube_view, null_after_cube, &error) &&
+            !cube_view.info.images[0].cube &&
+            !cube_view.blocks[0].instructions[0].memory.image_cube,
+        "canonical null respecialization retained stale cube-coordinate state");
+
   Program program;
   program.stage = ShaderType::Compute;
   program.blocks.resize(1);

tests/shaderCfgTests.cpp

@@ -3020,6 +3020,37 @@ void TestNewShaderRecompilerImageQueryLowering() {
 	CheckSpirvBinaryValidates(result.spirv);
 }
 
+void TestNewShaderRecompilerCubeSampleCoordinates() {
+	constexpr uint32_t MimgDimCube = 3;
+	const uint32_t shader[] = {
+	    EncodeMimg0(0x20, 0xf, false, MimgDimCube),
+	    EncodeMimg1(0, 0, 1, 0), // image_sample cube
+	    EncodeMimg0(0x60, 0x3, false, MimgDimCube),
+	    EncodeMimg1(8, 0, 1, 4), // image_get_lod cube
+	    0xbf810000u,
+	};
+
+	auto user_data = ImageTestUserData(Prospero::ImageType::kCube);
+	ShaderRecompiler::CompileOptions options;
+	options.stage     = ShaderType::Compute;
+	options.user_data = user_data.data();
+
+	ShaderRecompiler::CompileResult result;
+	std::string                     error;
+	Check(ShaderRecompiler::TryRecompile(shader, options, result, &error), error.c_str());
+	Check(result.program.info.images.size() == 1 && result.program.info.images[0].cube,
+	      "cube descriptor identity was not preserved through compilation");
+	Check(SpirvInstructionOpcodeCount(result.spirv, 131) == 4,
+	      "cube sample/get-lod did not remove the RDNA2 S/T bias");
+	Check(SpirvInstructionOpcodeCount(result.spirv, 109) == 1 &&
+	          SpirvInstructionOpcodeCount(result.spirv, 112) == 1 &&
+	          SpirvInstructionOpcodeCount(result.spirv, 194) == 1 &&
+	          SpirvInstructionOpcodeCount(result.spirv, 196) == 1 &&
+	          SpirvInstructionOpcodeCount(result.spirv, 130) == 1,
+	      "cube sample did not repack its face ID exactly once, or get-lod used an array layer");
+	CheckSpirvBinaryValidates(result.spirv);
+}
+
 void TestNewShaderRecompilerImageSampleVariants() {
 	const uint32_t shader[] = {
 	    EncodeMimg0(0x24, 0xf),
@@ -6957,6 +6988,7 @@ int main() {
 	TestNewShaderRecompilerScalarBitfieldAlu();
 	TestNewShaderRecompilerMemoryFamilyLowering();
 	TestNewShaderRecompilerImageQueryLowering();
+	TestNewShaderRecompilerCubeSampleCoordinates();
 	TestNewShaderRecompilerImageSampleVariants();
 	TestNewShaderRecompilerImageSampleA16SamplerCoords();
 	TestNewShaderRecompilerImageSampleOpcodeAliases();