[HLSL][SPIRV] Add vk::constant_id attribute.

clang/include/clang/Basic/Attr.td

@@ -4951,6 +4951,14 @@ def HLSLWaveSize: InheritableAttr {
   let Documentation = [WaveSizeDocs];
 }
 
+def HLSLVkConstantId : InheritableAttr {
+  let Spellings = [CXX11<"vk", "constant_id">];
+  let Args = [IntArgument<"Id">];
+  let Subjects = SubjectList<[Var]>;
+  let LangOpts = [HLSL];
+  let Documentation = [VkConstantIdDocs];
+}
+
 def RandomizeLayout : InheritableAttr {
   let Spellings = [GCC<"randomize_layout">];
   let Subjects = SubjectList<[Record]>;

clang/include/clang/Basic/AttrDocs.td

@@ -8195,6 +8195,21 @@ and https://microsoft.github.io/hlsl-specs/proposals/0013-wave-size-range.html
   }];
 }
 
+def VkConstantIdDocs : Documentation {
+  let Category = DocCatFunction;
+  let Content = [{
+The ``vk::constant_id`` attribute specify the id for a SPIR-V specialization
+constant. The attribute applies to const global scalar variables. The variable must be initialized with a C++11 constexpr.
+In SPIR-V, the
+variable will be replaced with an `OpSpecConstant` with the given id.
+The syntax is:
+
+.. code-block:: text
+
+  ``[[vk::constant_id(<Id>)]] const T Name = <Init>``
+}];
+}
+
 def RootSignatureDocs : Documentation {
   let Category = DocCatFunction;
   let Content = [{

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -12891,6 +12891,21 @@ def err_spirv_enum_not_int : Error<
 def err_spirv_enum_not_valid : Error<
    "invalid value for %select{storage class}0 argument">;
 
+def err_specialization_const_lit_init
+    : Error<"variable with 'vk::constant_id' attribute cannot have an "
+            "initializer that is not a constexpr">;
+def err_specialization_const_is_not_externally_visible
+    : Error<"variable with 'vk::constant_id' attribute must be externally "
+            "visible">;
+def err_specialization_const_missing_initializer
+    : Error<
+          "variable with 'vk::constant_id' attribute must have an initializer">;
+def err_specialization_const_missing_const
+    : Error<"variable with 'vk::constant_id' attribute must be const">;
+def err_specialization_const_is_not_int_or_float
+    : Error<"variable with 'vk::constant_id' attribute must be an enum, bool, "
+            "integer, or floating point value">;
+
 // errors of expect.with.probability
 def err_probability_not_constant_float : Error<
    "probability argument to __builtin_expect_with_probability must be constant "

clang/include/clang/Sema/SemaHLSL.h

@@ -98,6 +98,8 @@ class SemaHLSL : public SemaBase {
   HLSLWaveSizeAttr *mergeWaveSizeAttr(Decl *D, const AttributeCommonInfo &AL,
                                       int Min, int Max, int Preferred,
                                       int SpelledArgsCount);
+  HLSLVkConstantIdAttr *
+  mergeVkConstantIdAttr(Decl *D, const AttributeCommonInfo &AL, int Id);
   HLSLShaderAttr *mergeShaderAttr(Decl *D, const AttributeCommonInfo &AL,
                                   llvm::Triple::EnvironmentType ShaderType);
   HLSLParamModifierAttr *
@@ -122,6 +124,7 @@ class SemaHLSL : public SemaBase {
   void handleRootSignatureAttr(Decl *D, const ParsedAttr &AL);
   void handleNumThreadsAttr(Decl *D, const ParsedAttr &AL);
   void handleWaveSizeAttr(Decl *D, const ParsedAttr &AL);
+  void handleVkConstantIdAttr(Decl *D, const ParsedAttr &AL);
   void handleSV_DispatchThreadIDAttr(Decl *D, const ParsedAttr &AL);
   void handleSV_GroupThreadIDAttr(Decl *D, const ParsedAttr &AL);
   void handleSV_GroupIDAttr(Decl *D, const ParsedAttr &AL);

clang/lib/AST/ExprConstant.cpp

@@ -3570,6 +3570,11 @@ static bool evaluateVarDeclInit(EvalInfo &Info, const Expr *E,
   if (E->isValueDependent())
     return false;
 
+  // The initializer on a specialization constant is only its default value
+  // when it is not externally initialized. This value cannot be evaluated.
+  if (VD->hasAttr<HLSLVkConstantIdAttr>())
+    return false;
+
   // Dig out the initializer, and use the declaration which it's attached to.
   // FIXME: We should eventually check whether the variable has a reachable
   // initializing declaration.

clang/lib/Basic/Attributes.cpp

@@ -212,7 +212,8 @@ getScopeFromNormalizedScopeName(StringRef ScopeName) {
       .Case("hlsl", AttributeCommonInfo::Scope::HLSL)
       .Case("msvc", AttributeCommonInfo::Scope::MSVC)
       .Case("omp", AttributeCommonInfo::Scope::OMP)
-      .Case("riscv", AttributeCommonInfo::Scope::RISCV);
+      .Case("riscv", AttributeCommonInfo::Scope::RISCV)
+      .Case("vk", AttributeCommonInfo::Scope::HLSL);
 }
 
 unsigned AttributeCommonInfo::calculateAttributeSpellingListIndex() const {

clang/lib/CodeGen/CodeGenModule.cpp

@@ -5190,6 +5190,29 @@ CodeGenModule::GetOrCreateLLVMGlobal(StringRef MangledName, llvm::Type *Ty,
     if (const auto *CMA = D->getAttr<CodeModelAttr>())
       GV->setCodeModel(CMA->getModel());
 
+    if (const auto *ConstIdAttr = D->getAttr<HLSLVkConstantIdAttr>()) {
+      const Expr *Init = D->getInit();
+      APValue InitValue;
+      bool IsConstExpr = Init->isCXX11ConstantExpr(getContext(), &InitValue);
+      assert(IsConstExpr &&
+             "HLSLVkConstantIdAttr requires a constant initializer");
+      llvm::SmallString<10> InitString;
+      switch (InitValue.getKind()) {
+      case APValue::ValueKind::Int:
+        InitValue.getInt().toString(InitString);
+        break;
+      case APValue::ValueKind::Float:
+        InitValue.getFloat().toString(InitString);
+        break;
+      default:
+        llvm_unreachable(
+            "HLSLVkConstantIdAttr requires an int or float initializer");
+      }
+      std::string ConstIdStr =
+          (llvm::Twine(ConstIdAttr->getId()) + "," + InitString).str();
+      GV->addAttribute("spirv-constant-id", ConstIdStr);
+    }
+
     // Check if we a have a const declaration with an initializer, we may be
     // able to emit it as available_externally to expose it's value to the
     // optimizer.

clang/lib/Sema/SemaDecl.cpp

@@ -2889,6 +2889,8 @@ static bool mergeDeclAttribute(Sema &S, NamedDecl *D,
     NewAttr = S.HLSL().mergeWaveSizeAttr(D, *WS, WS->getMin(), WS->getMax(),
                                          WS->getPreferred(),
                                          WS->getSpelledArgsCount());
+  else if (const auto *CI = dyn_cast<HLSLVkConstantIdAttr>(Attr))
+    NewAttr = S.HLSL().mergeVkConstantIdAttr(D, *CI, CI->getId());
   else if (const auto *SA = dyn_cast<HLSLShaderAttr>(Attr))
     NewAttr = S.HLSL().mergeShaderAttr(D, *SA, SA->getType());
   else if (isa<SuppressAttr>(Attr))
@@ -13757,6 +13759,14 @@ void Sema::AddInitializerToDecl(Decl *RealDecl, Expr *Init, bool DirectInit) {
     return;
   }
 
+  if (VDecl->hasAttr<HLSLVkConstantIdAttr>()) {
+    if (!Init->isCXX11ConstantExpr(Context)) {
+      Diag(VDecl->getLocation(), diag::err_specialization_const_lit_init);
+      VDecl->setInvalidDecl();
+      return;
+    }
+  }
+
   // Get the decls type and save a reference for later, since
   // CheckInitializerTypes may change it.
   QualType DclT = VDecl->getType(), SavT = DclT;
@@ -14217,6 +14227,14 @@ void Sema::ActOnUninitializedDecl(Decl *RealDecl) {
       }
     }
 
+    // HLSL variable with the `vk::constant_id` attribute must be initialized.
+    if (!Var->isInvalidDecl() && Var->hasAttr<HLSLVkConstantIdAttr>()) {
+      Diag(Var->getLocation(),
+           diag::err_specialization_const_missing_initializer);
+      Var->setInvalidDecl();
+      return;
+    }
+
     if (!Var->isInvalidDecl() && RealDecl->hasAttr<LoaderUninitializedAttr>()) {
       if (Var->getStorageClass() == SC_Extern) {
         Diag(Var->getLocation(), diag::err_loader_uninitialized_extern_decl)

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7510,6 +7510,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_HLSLWaveSize:
     S.HLSL().handleWaveSizeAttr(D, AL);
     break;
+  case ParsedAttr::AT_HLSLVkConstantId:
+    S.HLSL().handleVkConstantIdAttr(D, AL);
+    break;
   case ParsedAttr::AT_HLSLSV_GroupThreadID:
     S.HLSL().handleSV_GroupThreadIDAttr(D, AL);
     break;

clang/lib/Sema/SemaHLSL.cpp

@@ -505,6 +505,7 @@ static CXXRecordDecl *createHostLayoutStruct(Sema &S,
 // - empty structs
 // - zero-sized arrays
 // - non-variable declarations
+// - SPIR-V specialization constants
 // The layout struct will be added to the HLSLBufferDecl declarations.
 void createHostLayoutStructForBuffer(Sema &S, HLSLBufferDecl *BufDecl) {
   ASTContext &AST = S.getASTContext();
@@ -520,7 +521,8 @@ void createHostLayoutStructForBuffer(Sema &S, HLSLBufferDecl *BufDecl) {
   for (Decl *D : BufDecl->buffer_decls()) {
     VarDecl *VD = dyn_cast<VarDecl>(D);
     if (!VD || VD->getStorageClass() == SC_Static ||
-        VD->getType().getAddressSpace() == LangAS::hlsl_groupshared)
+        VD->getType().getAddressSpace() == LangAS::hlsl_groupshared ||
+        VD->hasAttr<HLSLVkConstantIdAttr>())
       continue;
     const Type *Ty = VD->getType()->getUnqualifiedDesugaredType();
     if (FieldDecl *FD =
@@ -607,6 +609,54 @@ HLSLWaveSizeAttr *SemaHLSL::mergeWaveSizeAttr(Decl *D,
   return Result;
 }
 
+HLSLVkConstantIdAttr *
+SemaHLSL::mergeVkConstantIdAttr(Decl *D, const AttributeCommonInfo &AL,
+                                int Id) {
+
+  auto &TargetInfo = getASTContext().getTargetInfo();
+  if (TargetInfo.getTriple().getArch() != llvm::Triple::spirv) {
+    Diag(AL.getLoc(), diag::warn_attribute_ignored) << AL;
+    return nullptr;
+  }
+
+  auto *VD = cast<VarDecl>(D);
+
+  if (!VD->getType()->isIntegerType() && !VD->getType()->isFloatingType()) {
+    Diag(VD->getLocation(), diag::err_specialization_const_is_not_int_or_float);
+    return nullptr;
+  }
+
+  if (VD->getStorageClass() != StorageClass::SC_None &&
+      VD->getStorageClass() != StorageClass::SC_Extern) {
+    Diag(VD->getLocation(),
+         diag::err_specialization_const_is_not_externally_visible);
+    return nullptr;
+  }
+
+  if (VD->isLocalVarDecl()) {
+    Diag(VD->getLocation(),
+         diag::err_specialization_const_is_not_externally_visible);
+    return nullptr;
+  }
+
+  if (!VD->getType().isConstQualified()) {
+    Diag(VD->getLocation(), diag::err_specialization_const_missing_const);
+    return nullptr;
+  }
+
+  if (HLSLVkConstantIdAttr *CI = D->getAttr<HLSLVkConstantIdAttr>()) {
+    if (CI->getId() != Id) {
+      Diag(CI->getLocation(), diag::err_hlsl_attribute_param_mismatch) << AL;
+      Diag(AL.getLoc(), diag::note_conflicting_attribute);
+    }
+    return nullptr;
+  }
+
+  HLSLVkConstantIdAttr *Result =
+      ::new (getASTContext()) HLSLVkConstantIdAttr(getASTContext(), AL, Id);
+  return Result;
+}
+
 HLSLShaderAttr *
 SemaHLSL::mergeShaderAttr(Decl *D, const AttributeCommonInfo &AL,
                           llvm::Triple::EnvironmentType ShaderType) {
@@ -1117,6 +1167,15 @@ void SemaHLSL::handleWaveSizeAttr(Decl *D, const ParsedAttr &AL) {
     D->addAttr(NewAttr);
 }
 
+void SemaHLSL::handleVkConstantIdAttr(Decl *D, const ParsedAttr &AL) {
+  uint32_t Id;
+  if (!SemaRef.checkUInt32Argument(AL, AL.getArgAsExpr(0), Id))
+    return;
+  HLSLVkConstantIdAttr *NewAttr = mergeVkConstantIdAttr(D, AL, Id);
+  if (NewAttr)
+    D->addAttr(NewAttr);
+}
+
 bool SemaHLSL::diagnoseInputIDType(QualType T, const ParsedAttr &AL) {
   const auto *VT = T->getAs<VectorType>();
 

clang/test/AST/HLSL/vk.spec-constnat.usage.hlsl

@@ -0,0 +1,12 @@
+// RUN: %clang_cc1 -finclude-default-header -triple spirv-unknown-vulkan-compute -x hlsl -ast-dump -o - %s | FileCheck %s
+
+// CHECK: VarDecl {{.*}} specConst 'const hlsl_constant int' cinit
+// CHECK-NEXT: IntegerLiteral {{.*}} 'int' 12
+// CHECK-NEXT: HLSLVkConstantIdAttr {{.*}} 10    
+[[vk::constant_id(10)]]
+const int specConst = 12;
+
+// CHECK: CXXRecordDecl {{.*}} implicit struct __cblayout_$Globals definition
+// CHECK-NOT: FieldDecl {{.*}} specConst 'int'
+
+

clang/test/CodeGenHLSL/vk-features/vk.spec-constant.hlsl

@@ -0,0 +1,44 @@
+// RUN: %clang_cc1 -finclude-default-header -x hlsl -triple \
+// RUN:   spirv-unknown-vulkan-compute %s -emit-llvm -disable-llvm-passes \
+// RUN:   -o - | FileCheck %s
+
+
+// CHECK-DAG: @_ZL3sc0 = external addrspace(12) constant i32, align 4 [[A0:#[0-9]+]]
+// CHECK-DAG: attributes [[A0]] = { "spirv-constant-id"="0,1" }                                                                                                                                                                                                                                                
+[[vk::constant_id(0)]]
+const bool sc0 = true;
+
+// CHECK-DAG: @_ZL3sc1 = external addrspace(12) constant i32, align 4 [[A1:#[0-9]+]]
+// CHECK-DAG: attributes [[A1]] = { "spirv-constant-id"="1,10" }                                                                                                                                                                                                                                               
+[[vk::constant_id(1)]]
+const int sc1 = 10;
+
+// CHECK-DAG: @_ZL3sc2 = external addrspace(12) constant i32, align 4 [[A2:#[0-9]+]]
+// CHECK-DAG: attributes [[A2]] = { "spirv-constant-id"="2,-20" }                                                                                                                                                                                                                                              
+[[vk::constant_id(2)]]
+const int sc2 = 10-30;
+
+// CHECK-DAG: @_ZL3sc3 = external addrspace(12) constant float, align 4 [[A3:#[0-9]+]]
+// CHECK-DAG: attributes [[A3]] = { "spirv-constant-id"="3,0.25" }
+[[vk::constant_id(3)]]
+const float sc3 = 0.5*0.5;
+
+// CHECK-DAG: @_ZL3sc4 = external addrspace(12) constant i32, align 4 [[A4:#[0-9]+]]
+// CHECK-DAG: attributes [[A4]] = { "spirv-constant-id"="4,2" }
+enum E {
+    A,
+    B,
+    C
+};
+
+[[vk::constant_id(4)]]
+const E sc4 = E::C;
+
+[numthreads(1,1,1)]
+void main() {
+    bool b = sc0;
+    int i = sc1;
+    int j = sc2;
+    float f = sc3;
+    E e = sc4;
+}

clang/test/SemaHLSL/vk.spec-constant.error.hlsl

@@ -0,0 +1,37 @@
+// RUN: %clang_cc1 -finclude-default-header -triple spirv-pc-vulkan1.3-compute -verify %s
+// RUN: %clang_cc1 -finclude-default-header -triple dxil-pc-shadermodel6.8-compute -verify %s
+
+#ifndef __spirv__
+// expected-warning@+2{{'constant_id' attribute ignored}}
+#endif
+[[vk::constant_id(0)]]
+const bool sc0 = true;
+
+#ifdef __spirv__
+// expected-error@+2{{variable with 'vk::constant_id' attribute cannot have an initializer that is not a constexpr}}
+[[vk::constant_id(1)]]
+const bool sc1 = sc0; // error
+
+// expected-error@+2{{variable with 'vk::constant_id' attribute must be externally visible}}
+[[vk::constant_id(2)]]
+static const bool sc2 = false; // error
+
+// expected-error@+2{{variable with 'vk::constant_id' attribute must have an initializer}}
+[[vk::constant_id(3)]]
+const bool sc3; // error
+
+// expected-error@+2{{variable with 'vk::constant_id' attribute must be const}}
+[[vk::constant_id(4)]]
+bool sc4 = false; // error
+
+// expected-error@+2{{variable with 'vk::constant_id' attribute must be an enum, bool, integer, or floating point value}}
+[[vk::constant_id(5)]]
+const int2 sc5 = {0,0}; // error
+
+[numthreads(1,1,1)]
+void main() {
+  // expected-error@+2{{variable with 'vk::constant_id' attribute must be externally visible}}
+  [[vk::constant_id(6)]]
+  const bool sc6 = false; // error
+}
+#endif