[spirv] Add vk::image_format attribute for Buffers, RWBuffers and RWTextures (#3395)

According to Vulkan specification when using `OpImageRead/OpImageWrite`, the `OpTypeImage` (`Buffers`, `RWBuffers`, `RWTextures`) must have a format that matches the format on the API side, unless the StorageImageReadWithoutFormat/StorageImageWriteWithoutFormat is added and `Unknown` is used as the format.

This pull request addressess #2498 for the format part by adding an attribute `[[vk::image_format("<image format as spelled in SPIR-V spec>")]].` Example of the syntax:

```
[[vk::image_format("rgba8")]]
RWBuffer<float4> Buf;

[[vk::image_format("rg16f")]]
RWTexture2D<float2> Tex;

RWTexture2D<float2> Tex2; // Works like before
```

The `image_format` only applies to **global variables** of type `Buffer`, `RWBuffer`, `RWTexture`. For variables and function parameters it is propagated by the inlining pass in legalization. This required a small change to one of the passes in SPIRV-Tools, that should be also checked by someone more familiar with the codebase: KhronosGroup/SPIRV-Tools#4126

Note that this does not fix the handling of unspecified format (that case still works like before, using `R32f`, etc. based on the type in shader), although it should be still fixed to add the
StorageImageReadWithoutFormat and/or StorageImageWriteWithoutFormat and use Undefined. But I think the ability to specify the format is more urgent.

Design note from Jaebaek:
Since the `image_format` attribute only applies to **global variables**, under the DXC architecture
only `DeclResultIdMapper` can check the attribute when it handles `VarDecl`s. It means
we have to pass the `image_format` information to `LowerTypeVisitor` because it cannot access to
`VarDecl`. In order to pass the `image_format`, we use `SpirvContext` that can be accessed by
`SpirvEmitter` and all visitors. We use `SpirvVariable` to `spv::ImageFormat` mapping because the
attribute only applies to **global variables** (not to image types).
See how we use `llvm::DenseMap<const SpirvVariable *, spv::ImageFormat> spvVarToImageFormat`.

Author: seppala2

docs/SPIR-V.rst

@@ -786,6 +786,73 @@ are translated into SPIR-V ``OpTypeImage``, with parameters:
 The meanings of the headers in the above table is explained in ``OpTypeImage``
 of the SPIR-V spec.
 
+Vulkan specific Image Formats
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Since HLSL lacks the syntax for fully specifying image formats for textures in
+SPIR-V, we introduce ``[[vk::image_format("FORMAT")]]`` attribute for texture types.
+For example,
+
+.. code:: hlsl
+  [[vk::image_format("rgba8")]]
+  RWBuffer<float4> Buf;
+
+  [[vk::image_format("rg16f")]]
+  RWTexture2D<float2> Tex;
+
+  RWTexture2D<float2> Tex2; // Works like before
+
+``rgba8`` means ``Rgba8`` `SPIR-V Image Format <https://www.khronos.org/registry/spir-v/specs/unified1/SPIRV.html#_a_id_image_format_a_image_format>`_.
+The following table lists the mapping between ``FORMAT`` of
+``[[vk::image_format("FORMAT")]]`` and its corresponding SPIR-V Image Format.
+
+======================= ============================================
+       FORMAT                   SPIR-V Image Format
+======================= ============================================
+``unknown``             ``Unknown``
+``rgba32f``             ``Rgba32f``
+``rgba16f``             ``Rgba16f``
+``r32f``                ``R32f``
+``rgba8``               ``Rgba8``
+``rgba8snorm``          ``Rgba8Snorm``
+``rg32f``               ``Rg32f``
+``rg16f``               ``Rg16f``
+``r11g11b10f``          ``R11fG11fB10f``
+``r16f``                ``R16f``
+``rgba16``              ``Rgba16``
+``rgb10a2``             ``Rgb10A2``
+``rg16``                ``Rg16``
+``rg8``                 ``Rg8``
+``r16``                 ``R16``
+``r8``                  ``R8``
+``rgba16snorm``         ``Rgba16Snorm``
+``rg16snorm``           ``Rg16Snorm``
+``rg8snorm``            ``Rg8Snorm``
+``r16snorm``            ``R16Snorm``
+``r8snorm``             ``R8Snorm``
+``rgba32i``             ``Rgba32i``
+``rgba16i``             ``Rgba16i``
+``rgba8i``              ``Rgba8i``
+``r32i``                ``R32i``
+``rg32i``               ``Rg32i``
+``rg16i``               ``Rg16i``
+``rg8i``                ``Rg8i``
+``r16i``                ``R16i``
+``r8i``                 ``R8i``
+``rgba32ui``            ``Rgba32ui``
+``rgba16ui``            ``Rgba16ui``
+``rgba8ui``             ``Rgba8ui``
+``r32ui``               ``R32ui``
+``rgb10a2ui``           ``Rgb10a2ui``
+``rg32ui``              ``Rg32ui``
+``rg16ui``              ``Rg16ui``
+``rg8ui``               ``Rg8ui``
+``r16ui``               ``R16ui``
+``r8ui``                ``R8ui``
+``r64ui``               ``R64ui``
+``r64i``                ``R64i``
+======================= ============================================
+
 Constant/Texture/Structured/Byte Buffers
 ----------------------------------------
 

external/SPIRV-Headers

@@ -937,6 +937,32 @@ def ConstantTextureBuffer
                   S->getType()->getAs<RecordType>()->getDecl()->getName() ==
                       "TextureBuffer")}]>;
 
+// Global variable with "RWTexture" type
+def RWTexture
+    : SubsetSubject<
+          Var, [{S->hasGlobalStorage() && S->getType()->getAs<RecordType>() &&
+                 S->getType()->getAs<RecordType>()->getDecl() &&
+                  (S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWTexture1D" ||
+                  S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWTexture1DArray" ||
+                  S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWTexture2D" ||
+                  S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWTexture2DArray" ||
+                  S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWTexture3D")}]>;
+
+// Global variable with "[RW]Buffer" type
+def Buffer
+    : SubsetSubject<
+          Var, [{S->hasGlobalStorage() && S->getType()->getAs<RecordType>() &&
+                 S->getType()->getAs<RecordType>()->getDecl() &&
+                 (S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "Buffer" ||
+                  S->getType()->getAs<RecordType>()->getDecl()->getName() ==
+                      "RWBuffer")}]>;
+
 def VKBuiltIn : InheritableAttr {
   let Spellings = [CXX11<"vk", "builtin">];
   let Subjects = SubjectList<[Function, ParmVar, Field], ErrorDiag>;
@@ -997,6 +1023,29 @@ def VKOffset : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
+def VKImageFormat : InheritableAttr {
+  let Spellings = [CXX11<"vk", "image_format">];
+  let Subjects = SubjectList<[RWTexture, Buffer],
+                             ErrorDiag, "ExpectedRWTextureOrBuffer">;
+  let Args = [EnumArgument<"ImageFormat", "ImageFormatType",
+                           ["unknown", "rgba32f", "rgba16f", "r32f", "rgba8", "rgba8snorm",
+                           "rg32f", "rg16f", "r11g11b10f", "r16f", "rgba16", "rgb10a2",
+                           "rg16", "rg8", "r16", "r8", "rgba16snorm", "rg16snorm", "rg8snorm",
+                           "r16snorm", "r8snorm", "rgba32i", "rgba16i", "rgba8i", "r32i",
+                           "rg32i", "rg16i", "rg8i", "r16i", "r8i", "rgba32ui", "rgba16ui", "rgba8ui",
+                           "r32ui", "rgb10a2ui", "rg32ui", "rg16ui", "rg8ui", "r16ui",
+                           "r8ui", "r64ui", "r64i"],
+                           ["unknown", "rgba32f", "rgba16f", "r32f", "rgba8", "rgba8snorm",
+                           "rg32f", "rg16f", "r11g11b10f", "r16f", "rgba16", "rgb10a2",
+                           "rg16", "rg8", "r16", "r8", "rgba16snorm", "rg16snorm", "rg8snorm",
+                           "r16snorm", "r8snorm", "rgba32i", "rgba16i", "rgba8i", "r32i",
+                           "rg32i", "rg16i", "rg8i", "r16i", "r8i", "rgba32ui", "rgba16ui", "rgba8ui",
+                           "r32ui", "rgb10a2ui", "rg32ui", "rg16ui", "rg8ui", "r16ui",
+                           "r8ui", "r64ui", "r64i"]>];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
 def SubpassInput : SubsetSubject<
     Var,
     [{S->hasGlobalStorage() && S->getType()->getAs<RecordType>() &&

external/SPIRV-Tools

@@ -2332,6 +2332,7 @@ def warn_attribute_wrong_decl_type : Warning<
   "global variables of scalar type|"
   "global variables of struct type|"
   "global variables, cbuffers, and tbuffers|"
+  "RWTextures, Buffers and RWBuffers|"
   "RWStructuredBuffers, AppendStructuredBuffers, and ConsumeStructuredBuffers|"
   "SubpassInput, SubpassInputMS|"
   "cbuffer or ConstantBuffer|"

tools/clang/include/clang/Basic/Attr.td

@@ -235,6 +235,10 @@ class SpirvContext {
                                 ImageType::WithDepth, bool arrayed, bool ms,
                                 ImageType::WithSampler sampled,
                                 spv::ImageFormat);
+  // Get ImageType whose attributes are the same with imageTypeWithUnknownFormat
+  // but it has spv::ImageFormat format.
+  const ImageType *getImageType(const ImageType *imageTypeWithUnknownFormat,
+                                spv::ImageFormat format);
   const SamplerType *getSamplerType() const { return samplerType; }
   const SampledImageType *getSampledImageType(const ImageType *image);
   const HybridSampledImageType *getSampledImageType(QualType image);
@@ -335,6 +339,20 @@ class SpirvContext {
     return currentLexicalScope;
   }
 
+  /// Function to add/get the mapping from a SPIR-V OpVariable to its image
+  /// format.
+  void registerImageFormatForSpirvVariable(const SpirvVariable *spvVar,
+                                           spv::ImageFormat format) {
+    assert(spvVar != nullptr);
+    spvVarToImageFormat[spvVar] = format;
+  }
+  spv::ImageFormat getImageFormatForSpirvVariable(const SpirvVariable *spvVar) {
+    auto itr = spvVarToImageFormat.find(spvVar);
+    if (itr == spvVarToImageFormat.end())
+      return spv::ImageFormat::Unknown;
+    return itr->second;
+  }
+
   /// Function to add/get the mapping from a SPIR-V type to its Decl for
   /// a struct type.
   void registerStructDeclForSpirvType(const SpirvType *spvTy,
@@ -442,6 +460,9 @@ class SpirvContext {
   // Mapping from FunctionDecl to SPIR-V debug function.
   llvm::DenseMap<const FunctionDecl *, SpirvDebugFunction *>
       declToDebugFunction;
+
+  // Mapping from SPIR-V OpVariable to SPIR-V image format.
+  llvm::DenseMap<const SpirvVariable *, spv::ImageFormat> spvVarToImageFormat;
 };
 
 } // end namespace spirv

tools/clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -854,18 +854,19 @@ enum AttributeDeclKind {
   ExpectedStructOrUnionOrTypedef,
   ExpectedStructOrTypedef,
   ExpectedObjectiveCInterfaceOrProtocol,
-  ExpectedKernelFunction
+  ExpectedKernelFunction,
   // SPIRV Change Begins
-  ,ExpectedField
-  ,ExpectedScalarGlobalVar
-  ,ExpectedStructGlobalVar
-  ,ExpectedGlobalVarOrCTBuffer
-  ,ExpectedCounterStructuredBuffer
-  ,ExpectedSubpassInput
-  ,ExpectedCTBuffer
+  ExpectedField,
+  ExpectedScalarGlobalVar,
+  ExpectedStructGlobalVar,
+  ExpectedGlobalVarOrCTBuffer,
+  ExpectedRWTextureOrBuffer,
+  ExpectedCounterStructuredBuffer,
+  ExpectedSubpassInput,
+  ExpectedCTBuffer,
   // SPIRV Change Ends
   // HLSL Change Begins - add attribute decl combinations
-  ,ExpectedVariableOrParam,
+  ExpectedVariableOrParam,
   ExpectedFunctionOrParamOrField,
   ExpectedFunctionOrVariableOrParamOrFieldOrType
   // HLSL Change Ends

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

@@ -398,6 +398,99 @@ SpirvLayoutRule getLayoutRuleForExternVar(QualType type,
   return SpirvLayoutRule::Void;
 }
 
+spv::ImageFormat getSpvImageFormat(const VKImageFormatAttr *imageFormatAttr) {
+  if (imageFormatAttr == nullptr)
+    return spv::ImageFormat::Unknown;
+
+  switch (imageFormatAttr->getImageFormat()) {
+  case VKImageFormatAttr::unknown:
+    return spv::ImageFormat::Unknown;
+  case VKImageFormatAttr::rgba32f:
+    return spv::ImageFormat::Rgba32f;
+  case VKImageFormatAttr::rgba16f:
+    return spv::ImageFormat::Rgba16f;
+  case VKImageFormatAttr::r32f:
+    return spv::ImageFormat::R32f;
+  case VKImageFormatAttr::rgba8:
+    return spv::ImageFormat::Rgba8;
+  case VKImageFormatAttr::rgba8snorm:
+    return spv::ImageFormat::Rgba8Snorm;
+  case VKImageFormatAttr::rg32f:
+    return spv::ImageFormat::Rg32f;
+  case VKImageFormatAttr::rg16f:
+    return spv::ImageFormat::Rg16f;
+  case VKImageFormatAttr::r11g11b10f:
+    return spv::ImageFormat::R11fG11fB10f;
+  case VKImageFormatAttr::r16f:
+    return spv::ImageFormat::R16f;
+  case VKImageFormatAttr::rgba16:
+    return spv::ImageFormat::Rgba16;
+  case VKImageFormatAttr::rgb10a2:
+    return spv::ImageFormat::Rgb10A2;
+  case VKImageFormatAttr::rg16:
+    return spv::ImageFormat::Rg16;
+  case VKImageFormatAttr::rg8:
+    return spv::ImageFormat::Rg8;
+  case VKImageFormatAttr::r16:
+    return spv::ImageFormat::R16;
+  case VKImageFormatAttr::r8:
+    return spv::ImageFormat::R8;
+  case VKImageFormatAttr::rgba16snorm:
+    return spv::ImageFormat::Rgba16Snorm;
+  case VKImageFormatAttr::rg16snorm:
+    return spv::ImageFormat::Rg16Snorm;
+  case VKImageFormatAttr::rg8snorm:
+    return spv::ImageFormat::Rg8Snorm;
+  case VKImageFormatAttr::r16snorm:
+    return spv::ImageFormat::R16Snorm;
+  case VKImageFormatAttr::r8snorm:
+    return spv::ImageFormat::R8Snorm;
+  case VKImageFormatAttr::rgba32i:
+    return spv::ImageFormat::Rgba32i;
+  case VKImageFormatAttr::rgba16i:
+    return spv::ImageFormat::Rgba16i;
+  case VKImageFormatAttr::rgba8i:
+    return spv::ImageFormat::Rgba8i;
+  case VKImageFormatAttr::r32i:
+    return spv::ImageFormat::R32i;
+  case VKImageFormatAttr::rg32i:
+    return spv::ImageFormat::Rg32i;
+  case VKImageFormatAttr::rg16i:
+    return spv::ImageFormat::Rg16i;
+  case VKImageFormatAttr::rg8i:
+    return spv::ImageFormat::Rg8i;
+  case VKImageFormatAttr::r16i:
+    return spv::ImageFormat::R16i;
+  case VKImageFormatAttr::r8i:
+    return spv::ImageFormat::R8i;
+  case VKImageFormatAttr::rgba32ui:
+    return spv::ImageFormat::Rgba32ui;
+  case VKImageFormatAttr::rgba16ui:
+    return spv::ImageFormat::Rgba16ui;
+  case VKImageFormatAttr::rgba8ui:
+    return spv::ImageFormat::Rgba8ui;
+  case VKImageFormatAttr::r32ui:
+    return spv::ImageFormat::R32ui;
+  case VKImageFormatAttr::rgb10a2ui:
+    return spv::ImageFormat::Rgb10a2ui;
+  case VKImageFormatAttr::rg32ui:
+    return spv::ImageFormat::Rg32ui;
+  case VKImageFormatAttr::rg16ui:
+    return spv::ImageFormat::Rg16ui;
+  case VKImageFormatAttr::rg8ui:
+    return spv::ImageFormat::Rg8ui;
+  case VKImageFormatAttr::r16ui:
+    return spv::ImageFormat::R16ui;
+  case VKImageFormatAttr::r8ui:
+    return spv::ImageFormat::R8ui;
+  case VKImageFormatAttr::r64ui:
+    return spv::ImageFormat::R64ui;
+  case VKImageFormatAttr::r64i:
+    return spv::ImageFormat::R64i;
+  }
+  return spv::ImageFormat::Unknown;
+}
+
 } // anonymous namespace
 
 std::string StageVar::getSemanticStr() const {
@@ -847,6 +940,13 @@ SpirvVariable *DeclResultIdMapper::createExternVar(const VarDecl *var) {
       type, storageClass, var->hasAttr<HLSLPreciseAttr>(), name, llvm::None,
       loc);
   varInstr->setLayoutRule(rule);
+
+  // If this variable has [[vk::image_format("..")]] attribute, we have to keep
+  // it in the SpirvContext and use it when we lower the QualType to SpirvType.
+  auto spvImageFormat = getSpvImageFormat(var->getAttr<VKImageFormatAttr>());
+  if (spvImageFormat != spv::ImageFormat::Unknown)
+    spvContext.registerImageFormatForSpirvVariable(varInstr, spvImageFormat);
+
   DeclSpirvInfo info(varInstr);
   astDecls[var] = info;
 

tools/clang/include/clang/Sema/AttributeList.h

@@ -121,6 +121,14 @@ bool LowerTypeVisitor::visitInstruction(SpirvInstruction *instr) {
       if (var->hasBinding() && var->getHlslUserType().empty()) {
         var->setHlslUserType(getHlslResourceTypeName(var->getAstResultType()));
       }
+
+      auto spvImageFormat = spvContext.getImageFormatForSpirvVariable(var);
+      if (spvImageFormat != spv::ImageFormat::Unknown) {
+        if (const auto *imageType = dyn_cast<ImageType>(resultType)) {
+          resultType = spvContext.getImageType(imageType, spvImageFormat);
+          instr->setResultType(resultType);
+        }
+      }
     }
     const SpirvType *pointerType =
         spvContext.getPointerType(resultType, instr->getStorageClass());

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -184,6 +184,17 @@ const SpirvType *SpirvContext::getMatrixType(const SpirvType *elemType,
   return ptr;
 }
 
+const ImageType *
+SpirvContext::getImageType(const ImageType *imageTypeWithUnknownFormat,
+                           spv::ImageFormat format) {
+  return getImageType(imageTypeWithUnknownFormat->getSampledType(),
+                      imageTypeWithUnknownFormat->getDimension(),
+                      imageTypeWithUnknownFormat->getDepth(),
+                      imageTypeWithUnknownFormat->isArrayedImage(),
+                      imageTypeWithUnknownFormat->isMSImage(),
+                      imageTypeWithUnknownFormat->withSampler(), format);
+}
+
 const ImageType *SpirvContext::getImageType(const SpirvType *sampledType,
                                             spv::Dim dim,
                                             ImageType::WithDepth depth,