[SPIRV] Add support of [[vk::ext_decorate]] (#4035)

related to
issue:#3919

Add support to these two attributes:
 [[vk::ext_decorate(decoration, … int literal)]]
 [[vk::ext_decorate_string(decoration, … string literals)]]

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

@@ -1026,6 +1026,22 @@ def VKCounterBinding : InheritableAttr {
   let Documentation = [Undocumented];
 }
 
+def VKDecorateExt : InheritableAttr {
+  let Spellings = [CXX11<"vk", "ext_decorate">];
+  let Subjects = SubjectList<[Function, Var, ParmVar, Field, TypedefName], ErrorDiag>;
+  let Args = [UnsignedArgument<"decorate">, VariadicUnsignedArgument<"literal">];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
+def VKDecorateStringExt : InheritableAttr {
+  let Spellings = [CXX11<"vk", "ext_decorate_string">];
+  let Subjects = SubjectList<[Function, Var, ParmVar, Field, TypedefName], ErrorDiag>;
+  let Args = [UnsignedArgument<"decorate">, StringArgument<"literals">];
+  let LangOpts = [SPIRV];
+  let Documentation = [Undocumented];
+}
+
 def VKExtensionExt : InheritableAttr {
   let Spellings = [CXX11<"vk", "ext_extension">];
   let Subjects = SubjectList<[Function], ErrorDiag>;

tools/clang/include/clang/SPIRV/SpirvBuilder.h

@@ -660,6 +660,16 @@ class SpirvBuilder {
                        llvm::StringRef name, spv::LinkageType linkageType,
                        SourceLocation);
 
+  /// \brief Decorates the given target with information from VKDecorateExt
+  void decorateLiterals(SpirvInstruction *targetInst, unsigned decorate,
+                        unsigned *literal, unsigned literalSize,
+                        SourceLocation);
+
+  /// \brief Decorates the given target with the given string.
+  void decorateString(SpirvInstruction *target, unsigned decorate,
+                      llvm::StringRef strLiteral,
+                      llvm::Optional<uint32_t> memberIdx = llvm::None);                    
+
   /// --- Constants ---
   /// Each of these methods can acquire a unique constant from the SpirvContext,
   /// and add the context to the list of constants in the module.

tools/clang/include/clang/SPIRV/SpirvInstruction.h

@@ -461,9 +461,12 @@ class SpirvModuleProcessed : public SpirvInstruction {
 /// \brief OpDecorate(Id) and OpMemberDecorate instructions
 class SpirvDecoration : public SpirvInstruction {
 public:
+  // OpDecorate/OpMemberDecorate
   SpirvDecoration(SourceLocation loc, SpirvInstruction *target,
                   spv::Decoration decor, llvm::ArrayRef<uint32_t> params = {},
                   llvm::Optional<uint32_t> index = llvm::None);
+
+  // OpDecorateString/OpMemberDecorateString
   SpirvDecoration(SourceLocation loc, SpirvInstruction *target,
                   spv::Decoration decor, llvm::StringRef stringParam,
                   llvm::Optional<uint32_t> index = llvm::None);
@@ -499,6 +502,7 @@ class SpirvDecoration : public SpirvInstruction {
 private:
   spv::Op getDecorateOpcode(spv::Decoration,
                             const llvm::Optional<uint32_t> &memberIndex);
+  spv::Op getDecorateStringOpcode(bool isMemberDecoration);
 
 private:
   SpirvInstruction *target;

tools/clang/lib/SPIRV/DeclResultIdMapper.cpp

@@ -1753,9 +1753,11 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
   // Returns false if the given StageVar is an input/output variable without
   // explicit location assignment. Otherwise, returns true.
   const auto locAssigned = [forInput, this](const StageVar &v) {
-    if (forInput == isInputStorageClass(v))
+    if (forInput == isInputStorageClass(v)) {
       // No need to assign location for builtins. Treat as assigned.
-      return v.isSpirvBuitin() || v.getLocationAttr() != nullptr;
+      return v.isSpirvBuitin() || v.hasLocOrBuiltinDecorateAttr() ||
+             v.getLocationAttr() != nullptr;
+    }
     // For the ones we don't care, treat as assigned.
     return true;
   };
@@ -1773,8 +1775,10 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
 
     for (const auto &var : stageVars) {
       // Skip builtins & those stage variables we are not handling for this call
-      if (var.isSpirvBuitin() || forInput != isInputStorageClass(var))
+      if (var.isSpirvBuitin() || var.hasLocOrBuiltinDecorateAttr() ||
+          forInput != isInputStorageClass(var)) {
         continue;
+      }
 
       const auto *attr = var.getLocationAttr();
       const auto loc = attr->getNumber();
@@ -1815,8 +1819,10 @@ bool DeclResultIdMapper::finalizeStageIOLocations(bool forInput) {
   LocationSet locSet;
 
   for (const auto &var : stageVars) {
-    if (var.isSpirvBuitin() || forInput != isInputStorageClass(var))
+    if (var.isSpirvBuitin() || var.hasLocOrBuiltinDecorateAttr() ||
+        forInput != isInputStorageClass(var)) {
       continue;
+    }
 
     if (var.getLocationAttr()) {
       // We have checked that not all of the stage variables have explicit
@@ -2420,6 +2426,23 @@ bool DeclResultIdMapper::createStageVars(
         stageVar.getStorageClass() == spv::StorageClass::Output) {
       stageVar.setEntryPoint(entryFunction);
     }
+
+    if (decl->hasAttr<VKDecorateExtAttr>()) {
+      auto checkBuiltInLocationDecoration =
+          [&stageVar](const VKDecorateExtAttr *decoAttr) {
+            auto decorate =
+                static_cast<spv::Decoration>(decoAttr->getDecorate());
+            if (decorate == spv::Decoration::BuiltIn ||
+                decorate == spv::Decoration::Location) {
+              // This information will be used to avoid
+              // assigning multiple location decorations
+              // in finalizeStageIOLocations()
+              stageVar.setIsLocOrBuiltinDecorateAttr();
+            }
+          };
+      decorateWithIntrinsicAttrs(decl, varInstr,
+                                 checkBuiltInLocationDecoration);
+    }
     stageVars.push_back(stageVar);
 
     // Emit OpDecorate* instructions to link this stage variable with the HLSL
@@ -3962,5 +3985,31 @@ DeclResultIdMapper::createHullMainOutputPatch(const ParmVarDecl *param,
   return hullMainOutputPatch;
 }
 
+
+template <typename Functor>
+void DeclResultIdMapper::decorateWithIntrinsicAttrs(const NamedDecl *decl,
+                                                    SpirvVariable *varInst,
+                                                    Functor func) {
+  if (!decl->hasAttrs())
+    return;
+
+  for (auto &attr : decl->getAttrs()) {
+    if (auto decoAttr = dyn_cast<VKDecorateExtAttr>(attr)) {
+      func(decoAttr);
+      spvBuilder.decorateLiterals(
+          varInst, decoAttr->getDecorate(), decoAttr->literal_begin(),
+          decoAttr->literal_size(), varInst->getSourceLocation());
+    } else if (auto decoAttr = dyn_cast<VKDecorateStringExtAttr>(attr)) {
+      spvBuilder.decorateString(varInst, decoAttr->getDecorate(),
+                                decoAttr->getLiterals());
+    }
+  }
+}
+
+void DeclResultIdMapper::decorateVariableWithIntrinsicAttrs(
+    const NamedDecl *decl, SpirvVariable *varInst) {
+  decorateWithIntrinsicAttrs(decl, varInst, [](VKDecorateExtAttr *) {});
+}
+
 } // end namespace spirv
 } // end namespace clang

tools/clang/lib/SPIRV/DeclResultIdMapper.h

@@ -56,7 +56,8 @@ class StageVar {
       : sigPoint(sig), semanticInfo(std::move(semaInfo)), builtinAttr(builtin),
         type(astType), value(nullptr), isBuiltin(false),
         storageClass(spv::StorageClass::Max), location(nullptr),
-        locationCount(locCount), entryPoint(nullptr) {
+        locationCount(locCount), entryPoint(nullptr),
+        locOrBuiltinDecorateAttr(false) {
     isBuiltin = builtinAttr != nullptr;
   }
 
@@ -87,6 +88,8 @@ class StageVar {
 
   SpirvFunction *getEntryPoint() const { return entryPoint; }
   void setEntryPoint(SpirvFunction *entry) { entryPoint = entry; }
+  bool hasLocOrBuiltinDecorateAttr() const { return locOrBuiltinDecorateAttr; }
+  void setIsLocOrBuiltinDecorateAttr() { locOrBuiltinDecorateAttr = true; }
 
 private:
   /// HLSL SigPoint. It uniquely identifies each set of parameters that may be
@@ -113,6 +116,7 @@ class StageVar {
   /// Entry point for this stage variable. If this stage variable is not
   /// specific for an entry point e.g., built-in, it must be nullptr.
   SpirvFunction *entryPoint;
+  bool locOrBuiltinDecorateAttr;
 };
 
 /// \brief The struct containing information of stage variable's location and
@@ -604,6 +608,10 @@ class DeclResultIdMapper {
     return value;
   }
 
+  /// \brief Decorate variable with spirv intrinsic attributes
+  void decorateVariableWithIntrinsicAttrs(const NamedDecl *decl,
+                                          SpirvVariable *varInst);
+
 private:
   /// \brief Wrapper method to create a fatal error message and report it
   /// in the diagnostic engine associated with this consumer.
@@ -825,6 +833,12 @@ class DeclResultIdMapper {
   bool getImplicitRegisterType(const ResourceVar &var,
                                char *registerTypeOut) const;
 
+  /// Decorate with spirv intrinsic attributes with lamda function variable
+  /// check
+  template <typename Functor>
+  void decorateWithIntrinsicAttrs(const NamedDecl *decl, SpirvVariable *varInst,
+                                  Functor func);
+
 private:
   SpirvBuilder &spvBuilder;
   SpirvEmitter &theEmitter;

tools/clang/lib/SPIRV/SpirvBuilder.cpp

@@ -1441,6 +1441,27 @@ void SpirvBuilder::decorateLinkage(SpirvInstruction *targetInst,
   mod->addDecoration(decor);
 }
 
+void SpirvBuilder::decorateLiterals(SpirvInstruction *targetInst,
+                                    unsigned decorate, unsigned *literal,
+                                    unsigned literalSize,
+                                    SourceLocation srcLoc) {
+  SmallVector<uint32_t, 2> operands(literal, literal + literalSize);
+  SpirvDecoration *decor = new (context) SpirvDecoration(
+      srcLoc, targetInst, static_cast<spv::Decoration>(decorate), operands);
+  assert(decor != nullptr);
+  mod->addDecoration(decor);
+}
+
+void SpirvBuilder::decorateString(SpirvInstruction *target, unsigned decorate,
+                                  llvm::StringRef strLiteral,
+                                  llvm::Optional<uint32_t> memberIdx) {
+
+  auto *decor = new (context) SpirvDecoration(
+      target->getSourceLocation(), target,
+      static_cast<spv::Decoration>(decorate), strLiteral, memberIdx);
+  mod->addDecoration(decor);
+}
+
 SpirvConstant *SpirvBuilder::getConstantInt(QualType type, llvm::APInt value,
                                             bool specConst) {
   // We do not reuse existing constant integers. Just create a new one.

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -1654,6 +1654,10 @@ void SpirvEmitter::doVarDecl(const VarDecl *decl) {
       needsLegalization = true;
   }
 
+  if (var != nullptr) {
+    declIdMapper.decorateVariableWithIntrinsicAttrs(decl, var);
+  }
+
   // All variables that are of opaque struct types should request legalization.
   if (!needsLegalization && isOpaqueStructType(decl->getType()))
     needsLegalization = true;

tools/clang/lib/SPIRV/SpirvInstruction.cpp

@@ -227,18 +227,18 @@ SpirvDecoration::SpirvDecoration(SourceLocation loc,
                                  llvm::Optional<uint32_t> idx)
     : SpirvInstruction(IK_Decoration, getDecorateOpcode(decor, idx),
                        /*type*/ {}, loc),
-      target(targetInst), targetFunction(nullptr), decoration(decor), index(idx),
-      params(p.begin(), p.end()), idParams() {}
+      target(targetInst), targetFunction(nullptr), decoration(decor),
+      index(idx), params(p.begin(), p.end()), idParams() {}
 
 SpirvDecoration::SpirvDecoration(SourceLocation loc,
                                  SpirvInstruction *targetInst,
                                  spv::Decoration decor,
                                  llvm::StringRef strParam,
                                  llvm::Optional<uint32_t> idx)
-    : SpirvInstruction(IK_Decoration, getDecorateOpcode(decor, idx),
+    : SpirvInstruction(IK_Decoration, getDecorateStringOpcode(idx.hasValue()),
                        /*type*/ {}, loc),
-      target(targetInst), targetFunction(nullptr), decoration(decor), index(idx),
-      params(), idParams() {
+      target(targetInst), targetFunction(nullptr), decoration(decor),
+      index(idx), params(), idParams() {
   const auto &stringWords = string::encodeSPIRVString(strParam);
   params.insert(params.end(), stringWords.begin(), stringWords.end());
 }
@@ -266,11 +266,15 @@ spv::Op SpirvDecoration::getDecorateOpcode(
       decoration == spv::Decoration::UserTypeGOOGLE)
     return memberIndex.hasValue() ? spv::Op::OpMemberDecorateStringGOOGLE
                                   : spv::Op::OpDecorateStringGOOGLE;
-
   return memberIndex.hasValue() ? spv::Op::OpMemberDecorate
                                 : spv::Op::OpDecorate;
 }
 
+spv::Op SpirvDecoration::getDecorateStringOpcode(bool isMemberDecoration) {
+  return isMemberDecoration ? spv::Op::OpMemberDecorateString
+                            : spv::Op::OpDecorateString;
+}
+
 bool SpirvDecoration::operator==(const SpirvDecoration &that) const {
   return target == that.target && decoration == that.decoration &&
          params == that.params && idParams == that.idParams &&

tools/clang/lib/Sema/SemaHLSL.cpp

@@ -12068,6 +12068,28 @@ void hlsl::HandleDeclAttributeForHLSL(Sema &S, Decl *D, const AttributeList &A,
     declAttr = ::new (S.Context) VKReferenceExtAttr(
         A.getRange(), S.Context, A.getAttributeSpellingListIndex());
     break;
+  case AttributeList::AT_VKDecorateExt: {
+    // Note that `llvm::SmallVector<unsigned, 3> args` will be destroyed at
+    // the end of this function. However, VKDecorateExtAttr() constructor
+    // allocate a new integer array internally for literal and passing
+    // `&args[1]` is used just once for the initialization. It does not
+    // create a dangling pointer.
+    llvm::SmallVector<unsigned, 3> args;
+    auto argNum = A.getNumArgs();
+    for (unsigned i = 0; i < argNum; ++i) {
+      args.push_back(unsigned(ValidateAttributeIntArg(S, A, i)));
+    }
+    unsigned *literal = (argNum > 1) ? &args[1] : nullptr;
+    declAttr = ::new (S.Context)
+        VKDecorateExtAttr(A.getRange(), S.Context, args[0], literal, argNum - 1,
+                          A.getAttributeSpellingListIndex());
+  } break;
+  case AttributeList::AT_VKDecorateStringExt:
+    declAttr = ::new (S.Context) VKDecorateStringExtAttr(
+        A.getRange(), S.Context, unsigned(ValidateAttributeIntArg(S, A)),
+        ValidateAttributeStringArg(S, A, nullptr, 1),
+        A.getAttributeSpellingListIndex());
+    break;
   default:
     Handled = false;
     return;

tools/clang/test/CodeGenSPIRV/spv.intrinsicDecorate.hlsl

@@ -0,0 +1,39 @@
+// Run: %dxc -T ps_6_0 -E main -fcgl -Vd -spirv
+
+[[vk::ext_decorate(1, 0)]]
+bool b0;
+
+[[vk::ext_decorate(30, 23)]]
+float4 main(
+
+//CHECK: OpDecorate {{%\w+}} SpecId 0
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordNoPerspAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordNoPerspCentroidAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordNoPerspSampleAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordSmoothAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordSmoothCentroidAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordSmoothSampleAMD
+//CHECK: OpDecorate {{%\w+}} BuiltIn BaryCoordPullModelAMD
+//CHECK: OpDecorate {{%\w+}} ExplicitInterpAMD
+//CHECK: OpDecorate {{%\w+}} Location 23
+//CHECK: OpDecorateString {{%\w+}} UserSemantic "return variable"
+//CHECK: OpDecorate {{%\w+}} FPRoundingMode RTE
+
+    // spv::Decoration::builtIn = 11
+    [[vk::ext_decorate(11, 4992)]] float2 b0 : COLOR0,
+    [[vk::ext_decorate(11, 4993)]] float2 b1 : COLOR1,
+    [[vk::ext_decorate(11, 4994)]] float2 b2 : COLOR2,
+    [[vk::ext_decorate(11, 4995)]] float2 b3 : COLOR3,
+    [[vk::ext_decorate(11, 4996)]] float2 b4 : COLOR4,
+    [[vk::ext_decorate(11, 4997)]] float2 b5 : COLOR5,
+    [[vk::ext_decorate(11, 4998)]] float2 b6 : COLOR6,
+    // ExplicitInterpAMD
+    [[vk::location(12), vk::ext_decorate(4999)]] float2 b7 : COLOR7
+            ) : SV_Target {
+
+    // spv::Decoration::FPRoundingMode = 39, RTZ=0
+    [[vk::ext_decorate(39, 0), vk::ext_decorate_string(5635, "return variable")]] float4 ret = 1.0;
+    ret.xy = b0 * b1 + b2 + b3 + b4;
+    ret.zw = b5 + b6 + b7;
+    return ret;
+}

tools/clang/unittests/SPIRV/CodeGenSpirvTest.cpp

@@ -1342,6 +1342,7 @@ TEST_F(FileTest, IntrinsicsVkQueueFamilyScope) {
 TEST_F(FileTest, IntrinsicsSpirv) {
   runFileTest("spv.intrinsicInstruction.hlsl");
   runFileTest("spv.intrinsicLiteral.hlsl");
+  runFileTest("spv.intrinsicDecorate.hlsl", Expect::Success, false);
   runFileTest("spv.intrinsic.reference.error.hlsl", Expect::Failure);
 }
 TEST_F(FileTest, IntrinsicsVkReadClock) {