[SPIR-V] Fix image write to unknown format (microsoft#6984)

By default, texture format is guessed from the type. But the
`vk::image_format` attribute can be added to specify it.
The `unknown` value was used to convey `no format selected`, and require
the guess to happen.
This made selecting `unknown` as value impossible.

Once the optional added, we have a new issue:
- The format setup relied on the legalizer.
- load/stores are done using the default format, then we fix it, then we
use the legalizer to propagate the format fix to all users.

The capability visitor is running before the legalizer, meaning it can
look at a non-fixed load, which still carries the old type.
The way to solve this is to remove this logic, and move the capability
addition/deletion to the capability_trimming pass, run after
legalization.

Fixes microsoft#6981

---------

Signed-off-by: Nathan Gauër <[email protected]>

Author: Keenuts

tools/clang/include/clang/SPIRV/SpirvContext.h

@@ -11,6 +11,7 @@
 
 #include <array>
 #include <limits>
+#include <optional>
 
 #include "dxc/DXIL/DxilShaderModel.h"
 #include "clang/AST/DeclTemplate.h"
@@ -138,7 +139,7 @@ struct FunctionTypeMapInfo {
 struct VkImageFeatures {
   // True if it is a Vulkan "Combined Image Sampler".
   bool isCombinedImageSampler;
-  spv::ImageFormat format; // SPIR-V image format.
+  std::optional<spv::ImageFormat> format; // SPIR-V image format.
 };
 
 // A struct that contains the information of a resource that will be used to
@@ -378,7 +379,7 @@ class SpirvContext {
   getVkImageFeaturesForSpirvVariable(const SpirvVariable *spvVar) {
     auto itr = spvVarToVkImageFeatures.find(spvVar);
     if (itr == spvVarToVkImageFeatures.end())
-      return {false, spv::ImageFormat::Unknown};
+      return {false, std::nullopt};
     return itr->second;
   }
 

tools/clang/lib/SPIRV/CapabilityVisitor.cpp

@@ -439,36 +439,13 @@ bool CapabilityVisitor::visit(SpirvImageSparseTexelsResident *instr) {
   return true;
 }
 
-namespace {
-bool isImageOpOnUnknownFormat(const SpirvImageOp *instruction) {
-  if (!instruction->getImage() || !instruction->getImage()->getResultType()) {
-    return false;
-  }
-
-  const ImageType *imageType =
-      dyn_cast<ImageType>(instruction->getImage()->getResultType());
-  if (!imageType || imageType->getImageFormat() != spv::ImageFormat::Unknown) {
-    return false;
-  }
-
-  return imageType->getImageFormat() == spv::ImageFormat::Unknown;
-}
-} // namespace
-
 bool CapabilityVisitor::visit(SpirvImageOp *instr) {
   addCapabilityForType(instr->getResultType(), instr->getSourceLocation(),
                        instr->getStorageClass());
   if (instr->hasOffset() || instr->hasConstOffsets())
     addCapability(spv::Capability::ImageGatherExtended);
   if (instr->isSparse())
     addCapability(spv::Capability::SparseResidency);
-
-  if (isImageOpOnUnknownFormat(instr)) {
-    addCapability(instr->isImageWrite()
-                      ? spv::Capability::StorageImageWriteWithoutFormat
-                      : spv::Capability::StorageImageReadWithoutFormat);
-  }
-
   return true;
 }
 
@@ -864,6 +841,8 @@ bool CapabilityVisitor::visit(SpirvModule *, Visitor::Phase phase) {
   // supports only some capabilities. This list should be expanded to match the
   // supported capabilities.
   addCapability(spv::Capability::MinLod);
+  addCapability(spv::Capability::StorageImageWriteWithoutFormat);
+  addCapability(spv::Capability::StorageImageReadWithoutFormat);
 
   addExtensionAndCapabilitiesIfEnabled(
       Extension::EXT_fragment_shader_interlock,

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -517,9 +517,10 @@ SpirvLayoutRule getLayoutRuleForExternVar(QualType type,
   return SpirvLayoutRule::Void;
 }
 
-spv::ImageFormat getSpvImageFormat(const VKImageFormatAttr *imageFormatAttr) {
+std::optional<spv::ImageFormat>
+getSpvImageFormat(const VKImageFormatAttr *imageFormatAttr) {
   if (imageFormatAttr == nullptr)
-    return spv::ImageFormat::Unknown;
+    return std::nullopt;
 
   switch (imageFormatAttr->getImageFormat()) {
   case VKImageFormatAttr::unknown:
@@ -1235,14 +1236,13 @@ SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var,
   VkImageFeatures vkImgFeatures = {
       var->getAttr<VKCombinedImageSamplerAttr>() != nullptr,
       getSpvImageFormat(var->getAttr<VKImageFormatAttr>())};
-  if (vkImgFeatures.format != spv::ImageFormat::Unknown) {
+  if (vkImgFeatures.format) {
     // Legalization is needed to propagate the correct image type for
     // instructions in addition to cases where the resource is assigned to
     // another variable or function parameter
     needsLegalization = true;
   }
-  if (vkImgFeatures.isCombinedImageSampler ||
-      vkImgFeatures.format != spv::ImageFormat::Unknown) {
+  if (vkImgFeatures.isCombinedImageSampler || vkImgFeatures.format) {
     spvContext.registerVkImageFeaturesForSpvVariable(varInstr, vkImgFeatures);
   }
 
@@ -1256,8 +1256,9 @@ SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var,
   }
 
   if (hlsl::IsHLSLResourceType(type)) {
-    if (!areFormatAndTypeCompatible(vkImgFeatures.format,
-                                    hlsl::GetHLSLResourceResultType(type))) {
+    if (!areFormatAndTypeCompatible(
+            vkImgFeatures.format.value_or(spv::ImageFormat::Unknown),
+            hlsl::GetHLSLResourceResultType(type))) {
       emitError("The image format and the sampled type are not compatible.\n"
                 "For the table of compatible types, see "
                 "https://docs.vulkan.org/spec/latest/appendices/"

tools/clang/lib/SPIRV/LowerTypeVisitor.cpp

@@ -198,15 +198,17 @@ bool LowerTypeVisitor::visitInstruction(SpirvInstruction *instr) {
       }
 
       auto vkImgFeatures = spvContext.getVkImageFeaturesForSpirvVariable(var);
-      if (vkImgFeatures.format != spv::ImageFormat::Unknown) {
+      if (vkImgFeatures.format) {
         if (const auto *imageType = dyn_cast<ImageType>(resultType)) {
-          resultType = spvContext.getImageType(imageType, vkImgFeatures.format);
+          resultType =
+              spvContext.getImageType(imageType, *vkImgFeatures.format);
           instr->setResultType(resultType);
         } else if (const auto *arrayType = dyn_cast<ArrayType>(resultType)) {
           if (const auto *imageType =
                   dyn_cast<ImageType>(arrayType->getElementType())) {
-            auto newImgType =
-                spvContext.getImageType(imageType, vkImgFeatures.format);
+            auto newImgType = spvContext.getImageType(
+                imageType,
+                vkImgFeatures.format.value_or(spv::ImageFormat::Unknown));
             resultType = spvContext.getArrayType(newImgType,
                                                  arrayType->getElementCount(),
                                                  arrayType->getStride());

tools/clang/test/CodeGenSPIRV/op.buffer.access.hlsl

@@ -1,8 +1,7 @@
 // RUN: %dxc -T ps_6_0 -E main -fcgl  %s -spirv | FileCheck %s
 
 // CHECK: OpCapability ImageBuffer
-// CHECK: OpCapability StorageImageReadWithoutFormat
-
+// CHECK-NOT: OpCapability StorageImageReadWithoutFormat
 
 // CHECK: %type_buffer_image = OpTypeImage %int Buffer 2 0 0 1 R32i
 // CHECK: %type_buffer_image_0 = OpTypeImage %uint Buffer 2 0 0 1 R32ui

tools/clang/test/CodeGenSPIRV/vk.attribute.image-format.unknown.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -T cs_6_7 -E main -spirv -fspv-target-env=vulkan1.3 %s
+
+// CHECK: OpCapability StorageImageWriteWithoutFormat
+
+// CHECK: %[[#image:]] = OpTypeImage %float 3D 2 0 0 2 Unknown
+[[vk::image_format("unknown")]] RWTexture3D<float32_t2> untypedImage;
+
+[numthreads(8,8,8)]
+void main(uint32_t3 gl_GlobalInvocationID : SV_DispatchThreadID)
+{
+// CHECK: %[[#tmp:]] = OpLoad %[[#image]] %[[#]]
+// CHECK:              OpImageWrite [[#tmp]] %[[#]] %[[#]] None
+    untypedImage[gl_GlobalInvocationID] = float32_t2(4,5);
+}

tools/clang/test/CodeGenSPIRV/vk.attribute.image-format.unknown.read.hlsl

@@ -0,0 +1,16 @@
+// RUN: %dxc -T cs_6_7 -E main -spirv -fspv-target-env=vulkan1.3 %s
+
+// CHECK: OpCapability StorageImageReadWithoutFormat
+
+// CHECK: %[[#image:]] = OpTypeImage %float 1D 2 0 0 2 Unknown
+[[vk::image_format("unknown")]] RWTexture1D<float32_t2> untypedImage;
+RWStructuredBuffer<float32_t2> output;
+
+[numthreads(8,8,8)]
+void main()
+{
+// CHECK: %[[#tmp:]] = OpLoad %[[#image]] %[[#]]
+// CHECK:              OpImageRead %[[#]] [[#tmp]] %[[#]] None
+    output[0] = untypedImage[0];
+}
+