shader: lower FMASK loads before binding images

src/graphics/host_gpu/vulkanCommon.cpp

@@ -65,7 +65,6 @@ constexpr FormatMapping kFormatMappings[] = {
     {Prospero::BufferFormat::k5_6_5UNorm, vk::Format::eB5G6R5UnormPack16},
     {Prospero::BufferFormat::k5_5_5_1UNorm, vk::Format::eR5G5B5A1UnormPack16},
     {Prospero::BufferFormat::k4_4_4_4UNorm, vk::Format::eR4G4B4A4UnormPack16},
-    {Prospero::BufferFormat::kFmask8_S4_F4, vk::Format::eR32Sfloat},
     {Prospero::BufferFormat::kBc1UNorm, vk::Format::eBc1RgbaUnormBlock},
     {Prospero::BufferFormat::kBc1Srgb, vk::Format::eBc1RgbaSrgbBlock},
     {Prospero::BufferFormat::kBc2UNorm, vk::Format::eBc2UnormBlock},

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -687,6 +687,8 @@ CompileResult CompileProgram(TranslateResult translated, const CompileOptions& o
 	const auto emit_begin = std::chrono::steady_clock::now();
 	auto& ir = translated.program;
 	IR::ApplyResourceSpecialization(ir, specialization);
+	IR::RemoveIdentities(ir.blocks);
+	IR::EliminateDeadCode(ir.blocks);
 
 	const ShaderVertexInputInfo*  vertex  = nullptr;
 	const ShaderPixelInputInfo*   pixel   = nullptr;

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

@@ -362,6 +362,28 @@ struct SamplerPlan {
 	uint32_t                                      sampler_count = 0;
 };
 
+struct ImageRemap {
+	explicit ImageRemap(const ResourceSpecialization& specialization) {
+		for (const auto& image: specialization.images) {
+			indices.push_back(image.fmask ? UINT32_MAX : count++);
+		}
+	}
+
+	template <typename T>
+	void Apply(std::vector<T>& images) const {
+		EXIT_IF(images.size() != indices.size());
+		for (uint32_t index = 0; index < indices.size(); index++) {
+			if (indices[index] != UINT32_MAX && indices[index] != index) {
+				images[indices[index]] = std::move(images[index]);
+			}
+		}
+		images.resize(count);
+	}
+
+	std::vector<uint32_t> indices;
+	uint32_t              count = 0;
+};
+
 template <typename Images>
 bool BuildSamplerPlan(const ShaderInfo& base, const Images& images, SamplerPlan& plan);
 
@@ -495,6 +517,15 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, Materialize
 			                static_cast<uint32_t>(format)));
 		}
 		const bool storage      = base.resource_class == ImageResourceClass::Storage;
+		image.fmask             = Prospero::IsFmaskTextureFormat(format);
+		if (image.fmask) {
+			if (storage || base.depth_compare ||
+			    image.indirect_root != ImageResource::NoIndirectImage ||
+			    std::ranges::any_of(program.info.sampled_pairs,
+			                        [&](const auto& pair) { return pair.image == i; })) {
+				return SpecializationFail("FMASK requires a direct image load");
+			}
+		}
 		image.conversion_format = ImageConversionFormat(format);
 		if (storage || image.conversion_format != Prospero::BufferFormat::kInvalid) {
 			image.shader_swizzle = DescriptorImageSwizzle(descriptor);
@@ -585,6 +616,7 @@ static bool BuildResourceSpecialization(const ResourcePlan& program, Materialize
 			next_snapshot.samplers.push_back(next_snapshot.samplers[index]);
 		}
 	}
+	ImageRemap(next_specialization).Apply(next_snapshot.images);
 	specialization       = std::move(next_specialization);
 	specialized_snapshot = std::move(next_snapshot);
 	return true;
@@ -1012,8 +1044,10 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
 	}
 
 	auto memory_info = program.memory_info;
-	for (const auto* block: program.blocks) {
+	const ImageRemap image_remap(specialization);
-		for (const auto& inst: *block) {
+	for (auto* block: program.blocks) {
+		for (auto it = block->begin(); it != block->end(); ++it) {
+			auto& inst = *it;
 			const auto image_opcode = ImageOpcodeInfoOf(inst.GetOpcode());
 			if (image_opcode.access == ImageAccess::None) {
 				continue;
@@ -1023,6 +1057,23 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
 			auto& memory = memory_info[index];
 			EXIT_IF(memory.resource >= images.size());
 			const auto& image = images[memory.resource];
+			if (specialization.images[memory.resource].fmask) {
+				EXIT_IF(inst.GetOpcode() != ValueOpcode::ImageRead || memory.data_bits != 32u);
+				// Vulkan MSAA stores each sample directly; FMASK's four-bit fragment indices
+				// therefore map each coverage sample to the same host sample.
+				constexpr uint32_t indices[] = {0x76543210u, 0xfedcba98u};
+				std::array<Value, 2> fragments;
+				for (uint32_t component = 0; component < fragments.size(); component++) {
+					const auto selected = block->PrependNewInst(
+					    it, ValueOpcode::SelectU32, {inst.Arg(2), Value(indices[component]), Value(0u)});
+					fragments[component] = Value(&*selected);
+				}
+				const auto result = block->PrependNewInst(
+				    it, ValueOpcode::CompositeConstructU32x4,
+				    {fragments[0], fragments[1], Value(0u), Value(0u)});
+				inst.ReplaceUsesWith(Value(&*result));
+				continue;
+			}
 			if (image_opcode.needs_sampler && RequiresPointSampler(image) &&
 			    memory.sampler < program.info.samplers.size()) {
 				EXIT_IF(sampler_plan.point_sampler[memory.sampler] == UINT32_MAX);
@@ -1032,6 +1083,35 @@ void ApplyResourceSpecialization(Program& program, const ResourceSpecialization&
 			        inst.GetOpcode() != ValueOpcode::ImageSampleRaw);
 		}
 	}
+	for (auto* block: program.blocks) {
+		for (auto& inst: *block) {
+			if (inst.GetOpcode() == ValueOpcode::GetImageResource) {
+				inst.SetFlags(image_remap.indices.at(inst.Flags<uint32_t>()));
+			}
+		}
+	}
+	for (auto& memory: memory_info) {
+		if (memory.kind == ResourceKind::Image && !memory.planning_only) {
+			memory.resource = image_remap.indices.at(memory.resource);
+		}
+	}
+	for (auto& buffer: buffers) {
+		if (buffer.image_alias != BufferResource::NoImageAlias) {
+			buffer.image_alias = image_remap.indices.at(buffer.image_alias);
+		}
+	}
+	for (auto& pair: sampled_pairs) {
+		pair.image = image_remap.indices.at(pair.image);
+	}
+	for (auto& image: images) {
+		if (image.indirect_root != ImageResource::NoIndirectImage) {
+			image.indirect_root = image_remap.indices.at(image.indirect_root);
+		}
+		for (auto& resource: image.indirect_resources) {
+			resource = image_remap.indices.at(resource);
+		}
+	}
+	image_remap.Apply(images);
 	program.info.buffers       = std::move(buffers);
 	program.info.images        = std::move(images);
 	program.info.samplers      = std::move(samplers);

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

@@ -25,6 +25,7 @@ struct ResourceSpecialization {
 		uint32_t                      indirect_mapping_offset    = 0;
 		uint32_t                      indirect_search_iterations = 0;
 		bool                          cube                       = false;
+		bool                          fmask                      = false;
 		bool                          operator==(const Image&) const = default;
 	};
 

tests/ResourceTrackingTests.cpp

@@ -825,6 +825,96 @@ void TestSampleAdjustSamplerScratch() {
                 "SampleAdjust canonicalization discarded border-mode bits");
 }
 
+void TestFmaskLoadSpecialization() {
+  namespace Prospero = Libs::Graphics::Prospero;
+  Fixture fixture;
+  std::array<Value, 8> words;
+  for (uint32_t i = 0; i < words.size(); i++) {
+    words[i] = fixture.UserData(i);
+  }
+  const auto fmask = fixture.Image(words, 4);
+  const auto active = fixture.Emit(ValueOpcode::IEqual32,
+                                    {fixture.UserData(8), Value(0u)});
+  MemoryInfo load;
+  load.kind = ResourceKind::Image;
+  load.image_dimension = Decoder::ImageDimension::Dim2D;
+  load.image_address_components = 2;
+  load.dmask = 1;
+  const auto mapping = fixture.Emit(
+      ValueOpcode::ImageRead, {fmask, fixture.ImageAddress(), active},
+      fixture.AddMemory(load, 4));
+  const auto ordinary = fixture.Image(
+      {Value(0x2000u),
+       Value(static_cast<uint32_t>(Prospero::BufferFormat::k8UInt) << 20u),
+       Value(3u | (3u << 14u)),
+       Value(Libs::Graphics::DstSel(4, 5, 6, 7) |
+             (static_cast<uint32_t>(Prospero::ImageType::kColor2D) << 28u)),
+       Value(0u), Value(0u), Value(0u), Value(0u)}, 8);
+  const auto ordinary_flags = fixture.AddMemory(load, 8);
+  const auto color = fixture.Emit(
+      ValueOpcode::ImageRead, {ordinary, fixture.ImageAddress(), Value(true)},
+      ordinary_flags);
+  const auto output = fixture.Buffer(
+      {Value(0x3000u), Value(0u), Value(12u), Value(0u)}, 12);
+  MemoryInfo store;
+  store.kind = ResourceKind::Buffer;
+  Value result;
+  for (uint32_t i = 0; i < 2; i++) {
+    const auto value = fixture.Emit(
+        ValueOpcode::CompositeExtractU32x4,
+        {i == 0 ? mapping : color, Value(0u)});
+    fixture.Emit(ValueOpcode::StoreBufferU32,
+                 {output, Value(0u), Value(i * 4u), Value(0u), value, Value(true)},
+                 fixture.AddMemory(store, 12 + i * 4u));
+    if (i == 0) result = value;
+  }
+  fixture.PlanAndTrack();
+  const auto plan = ExtractResourcePlan(fixture.program);
+  std::array<uint32_t, 9> user_data{
+      0x303ac300u, 0xca100000u, 0x021bc3bfu, 0x91800004u,
+      0u, 0x00700000u, 0u, 0u};
+  ResourceSnapshot snapshot;
+  ResourceSpecialization specialization;
+  Check(MaterializeResources(plan, {.user_data = user_data}, snapshot,
+                             specialization),
+        "FMASK resources did not materialize");
+  ApplyResourceSpecialization(fixture.program, specialization);
+  RemoveIdentities(fixture.program.blocks);
+  EliminateDeadCode(fixture.program.blocks);
+  Check(fixture.program.info.images.size() == 1 && snapshot.images.size() == 1 &&
+            snapshot.images[0].dwords[0] == 0x2000u &&
+            ordinary.Instruction()->Flags<uint32_t>() == 0 &&
+            fixture.program.memory_info[ordinary_flags.index].resource == 0,
+        "FMASK removal did not preserve the remaining image and runtime descriptor");
+  const auto *vector = result.Instruction()->Arg(0).Resolve().TryInstruction();
+  Check(vector != nullptr &&
+            vector->GetOpcode() == ValueOpcode::CompositeConstructU32x4,
+        "FMASK load did not lower to a value vector");
+  result = vector->Arg(0);
+  uint32_t value = 0;
+  Check(EvaluateUniformValues(plan, {&result, 1}, {.user_data = user_data}, {&value, 1}) &&
+            value == 0x76543210u,
+        "FMASK load did not return the native sample-to-fragment mapping");
+  user_data[8] = 1;
+  Check(EvaluateUniformValues(plan, {&result, 1}, {.user_data = user_data}, {&value, 1}) &&
+            value == 0u,
+        "inactive FMASK load did not preserve the execution mask");
+  ShaderComputeInputInfo compute{};
+  CollectShaderInfo(fixture.program, {.compute = &compute});
+  AllocateBindings(fixture.program);
+  const auto kind = DescriptorBindingForImage(fixture.program.info.images[0]);
+  Check(kind.has_value() &&
+            FindBinding(fixture.program.bindings, *kind)->resources ==
+                std::vector<uint32_t>{0},
+        "FMASK allocated an ordinary image descriptor");
+  user_data[8] = 0;
+  user_data[1] = static_cast<uint32_t>(Prospero::BufferFormat::k8UInt) << 20u;
+  ResourceSpecialization rebound;
+  Check(MaterializeResources(plan, {.user_data = user_data}, snapshot, rebound) &&
+            rebound != specialization && snapshot.images.size() == 2,
+        "rebinding FMASK as a texture reused the metadata specialization");
+}
+
 void TestDynamicStorageMipTracking() {
   Fixture fixture;
   std::array<Value, 8> image_words;
@@ -1760,6 +1850,7 @@ int main() {
     Run("runtime unsigned min", TestRuntimeUnsignedMinDescriptor);
     Run("images and samplers", TestImagesSamplersAndAliases);
     Run("SampleAdjust sampler scratch", TestSampleAdjustSamplerScratch);
+    Run("FMASK load specialization", TestFmaskLoadSpecialization);
     Run("dynamic storage mips", TestDynamicStorageMipTracking);
     Run("invariant indirect images", TestInvariantIndirectImageMaterialization);
     Run("SRT runtime", TestSrtFlatteningAndRuntimeMemoization);

tests/ShaderRecompilerComputeTests.cpp

@@ -11597,8 +11597,7 @@ public:
          raw_format <= static_cast<uint32_t>(Prospero::BufferFormat::kBc7Srgb);
          ++raw_format) {
       const auto format = static_cast<Prospero::BufferFormat>(raw_format);
-      // FMASK is synthesized as identity metadata by TextureUploadFmask; it is
+      // FMASK contains sample-to-fragment metadata, not texels for the GPU tiler.
-      // deliberately not a texel surface and must never enter the GPU tiler.
       if (Prospero::IsFmaskTextureFormat(format)) {
         continue;
       }
@@ -12729,9 +12728,9 @@ void CompareGraphicsWords(const GraphicsCase &test,
 
 void RunCase(VulkanHarness *vulkan, const TestCase &test) {
   auto compiled = CompileCase(test, vulkan != nullptr ? vulkan->SubgroupSize() : 64u);
-  if (test.image_descriptor_swizzle != DstSel(4, 5, 6, 7)) {
+  if (test.image_descriptor_swizzle != DstSel(4, 5, 6, 7) &&
+      !compiled.program.info.images.empty()) {
     Require(test.name, "resource specialization",
-            !compiled.program.info.images.empty() &&
             compiled.program.info.images[0].shader_swizzle ==
                     test.image_descriptor_swizzle,
             "storage image descriptor swizzle did not reach the specialized "
@@ -21089,6 +21088,19 @@ TestCase ImageLoadR32UintUsesIntegerSampledImage() {
   return test;
 }
 
+TestCase ImageLoadFmaskUsesNativeSampleMapping() {
+  auto test = ImageLoadR32UintUsesIntegerSampledImage();
+  test.name = "ImageLoadFmaskUsesNativeSampleMapping";
+  test.expected = {0x76543210u};
+  test.user_data = {0x303ac300u, 0xca100000u, 0x021bc3bfu, 0x91800004u,
+                    0u, 0x00700000u, 0u, 0u};
+  test.image_descriptor_swizzle = DstSel(4, 0, 0, 0);
+  test.sampled_image_rgba.clear();
+  test.required_spirv = {"OpConstant %uint 1985229328"};
+  test.forbidden_spirv = {"OpTypeImage", "OpImageFetch"};
+  return test;
+}
+
 TestCase ImageLoadR32SintUsesSignedSampledImage() {
   using O = ShaderOpcode;
 
@@ -22950,6 +22962,7 @@ std::vector<TestCase> MakeCases() {
   AddCase(BufferAtomicFMaxContendedWorkgroup);
   AddCase(ImageLoadVariants);
   AddCase(ImageLoadR32UintUsesIntegerSampledImage);
+  AddCase(ImageLoadFmaskUsesNativeSampleMapping);
   AddCase(ImageLoadR32SintUsesSignedSampledImage);
   AddCase(ImageLoadPackedUintUnpacksAndSwizzles);
   AddCase(ImageSamplePackedUintConvertsSampleAndGather);
@@ -27251,6 +27264,12 @@ int main(int argc, char **argv) {
   std::setvbuf(stdout, nullptr, _IONBF, 0);
   EnsureConfigInitialized();
   CheckLeastRecentlyUsedCacheOrdering();
+  if (argc == 2 && std::strcmp(argv[1], "--fmask-only") == 0) {
+    VulkanHarness vulkan;
+    RunCase(&vulkan, ImageLoadFmaskUsesNativeSampleMapping());
+    RunCase(&vulkan, ImageLoadR32UintUsesIntegerSampledImage());
+    return 0;
+  }
 #if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS
   if (argc == 2 && std::strcmp(argv[1], "--shader-fatal-only") == 0) {
     CheckShaderRecompilerFatalContracts();