[HLSL][SPIR-V] Implement vk::ext_builtin_input attribute (#138530)

This variable attribute is used in HLSL to add Vulkan specific builtins
in a shader.
The attribute is documented here:

https://github.com/microsoft/hlsl-specs/blob/17727e88fd1cb09013cb3a144110826af05f4dd5/proposals/0011-inline-spirv.md

Those variable, even if marked as `static` are externally initialized by
the pipeline/driver/GPU. This is handled by moving them to a specific
address space `hlsl_input`, also added by this commit.

The design for input variables in Clang can be found here:
https://github.com/llvm/wg-hlsl/blob/355771361ef69259fef39a65caef8bff9cb4046d/proposals/0019-spirv-input-builtin.md


Co-authored-by: Justin Bogner <[email protected]>

Author: Keenuts

clang/include/clang/Basic/AddressSpaces.h

@@ -61,6 +61,7 @@ enum class LangAS : unsigned {
   hlsl_constant,
   hlsl_private,
   hlsl_device,
+  hlsl_input,
 
   // Wasm specific address spaces.
   wasm_funcref,

clang/include/clang/Basic/Attr.td

@@ -140,6 +140,12 @@ def SharedVar : SubsetSubject<Var,
                               [{S->hasGlobalStorage() && !S->getTLSKind()}],
                               "global variables">;
 
+def HLSLInputBuiltin
+    : SubsetSubject<Var, [{S->hasGlobalStorage() &&
+                           S->getStorageClass() == StorageClass::SC_Static &&
+                           S->getType().isConstQualified()}],
+                    "static const globals">;
+
 def GlobalVar : SubsetSubject<Var,
                              [{S->hasGlobalStorage()}], "global variables">;
 
@@ -4951,6 +4957,14 @@ def HLSLWaveSize: InheritableAttr {
   let Documentation = [WaveSizeDocs];
 }
 
+def HLSLVkExtBuiltinInput : InheritableAttr {
+  let Spellings = [CXX11<"vk", "ext_builtin_input">];
+  let Args = [UnsignedArgument<"BuiltIn">];
+  let Subjects = SubjectList<[HLSLInputBuiltin], ErrorDiag>;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLVkExtBuiltinInputDocs];
+}
+
 def RandomizeLayout : InheritableAttr {
   let Spellings = [GCC<"randomize_layout">];
   let Subjects = SubjectList<[Record]>;

clang/include/clang/Basic/AttrDocs.td

@@ -8508,6 +8508,29 @@ and copied back to the argument after the callee returns.
   }];
 }
 
+def HLSLVkExtBuiltinInputDocs : Documentation {
+  let Category = DocCatVariable;
+  let Content = [{
+Vulkan shaders have `Input` builtins. Those variables are externally
+initialized by the driver/pipeline, but each copy is private to the current
+lane.
+
+Those builtins can be declared using the `[[vk::ext_builtin_input]]` attribute
+like follows:
+
+.. code-block:: c++
+
+  [[vk::ext_builtin_input(/* WorkgroupId */ 26)]]
+  static const uint3 groupid;
+
+This variable will be lowered into a module-level variable, with the `Input`
+storage class, and the `BuiltIn 26` decoration.
+
+The full documentation for this inline SPIR-V attribute can be found here:
+https://github.com/microsoft/hlsl-specs/blob/main/proposals/0011-inline-spirv.md
+  }];
+}
+
 def AnnotateTypeDocs : Documentation {
   let Category = DocCatType;
   let Heading = "annotate_type";

clang/include/clang/Basic/AttributeCommonInfo.h

@@ -69,7 +69,7 @@ class AttributeCommonInfo {
     IgnoredAttribute,
     UnknownAttribute,
   };
-  enum class Scope { NONE, CLANG, GNU, MSVC, OMP, HLSL, GSL, RISCV };
+  enum class Scope { NONE, CLANG, GNU, MSVC, OMP, HLSL, VK, GSL, RISCV };
   enum class AttrArgsInfo {
     None,
     Optional,

clang/include/clang/Sema/SemaHLSL.h

@@ -131,6 +131,8 @@ class SemaHLSL : public SemaBase {
   void handleParamModifierAttr(Decl *D, const ParsedAttr &AL);
   bool handleResourceTypeAttr(QualType T, const ParsedAttr &AL);
 
+  void handleVkExtBuiltinInputAttr(Decl *D, const ParsedAttr &AL);
+
   bool CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall);
   QualType ProcessResourceTypeAttributes(QualType Wrapped);
   HLSLAttributedResourceLocInfo

clang/lib/AST/Type.cpp

@@ -100,6 +100,7 @@ bool Qualifiers::isTargetAddressSpaceSupersetOf(LangAS A, LangAS B,
          // address spaces to default to work around this problem.
          (A == LangAS::Default && B == LangAS::hlsl_private) ||
          (A == LangAS::Default && B == LangAS::hlsl_device) ||
+         (A == LangAS::Default && B == LangAS::hlsl_input) ||
          // Conversions from target specific address spaces may be legal
          // depending on the target information.
          Ctx.getTargetInfo().isAddressSpaceSupersetOf(A, B);

clang/lib/AST/TypePrinter.cpp

@@ -2667,6 +2667,8 @@ std::string Qualifiers::getAddrSpaceAsString(LangAS AS) {
     return "hlsl_private";
   case LangAS::hlsl_device:
     return "hlsl_device";
+  case LangAS::hlsl_input:
+    return "hlsl_input";
   case LangAS::wasm_funcref:
     return "__funcref";
   default:

clang/lib/Basic/Attributes.cpp

@@ -210,6 +210,7 @@ getScopeFromNormalizedScopeName(StringRef ScopeName) {
       .Case("gnu", AttributeCommonInfo::Scope::GNU)
       .Case("gsl", AttributeCommonInfo::Scope::GSL)
       .Case("hlsl", AttributeCommonInfo::Scope::HLSL)
+      .Case("vk", AttributeCommonInfo::Scope::VK)
       .Case("msvc", AttributeCommonInfo::Scope::MSVC)
       .Case("omp", AttributeCommonInfo::Scope::OMP)
       .Case("riscv", AttributeCommonInfo::Scope::RISCV);

clang/lib/Basic/TargetInfo.cpp

@@ -49,6 +49,8 @@ static const LangASMap FakeAddrSpaceMap = {
     13, // hlsl_groupshared
     14, // hlsl_constant
     15, // hlsl_private
+    16, // hlsl_device
+    17, // hlsl_input
     20, // wasm_funcref
 };
 

clang/lib/Basic/Targets/AArch64.h

@@ -47,6 +47,7 @@ static const unsigned ARM64AddrSpaceMap[] = {
     0, // hlsl_constant
     0, // hlsl_private
     0, // hlsl_device
+    0, // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/Basic/Targets/AMDGPU.cpp

@@ -62,6 +62,7 @@ const LangASMap AMDGPUTargetInfo::AMDGPUDefIsGenMap = {
     // will break loudly.
     llvm::AMDGPUAS::PRIVATE_ADDRESS, // hlsl_private
     llvm::AMDGPUAS::GLOBAL_ADDRESS,  // hlsl_device
+    llvm::AMDGPUAS::PRIVATE_ADDRESS, // hlsl_input
 };
 
 const LangASMap AMDGPUTargetInfo::AMDGPUDefIsPrivMap = {
@@ -89,6 +90,7 @@ const LangASMap AMDGPUTargetInfo::AMDGPUDefIsPrivMap = {
     llvm::AMDGPUAS::CONSTANT_ADDRESS, // hlsl_constant
     llvm::AMDGPUAS::PRIVATE_ADDRESS,  // hlsl_private
     llvm::AMDGPUAS::GLOBAL_ADDRESS,   // hlsl_device
+    llvm::AMDGPUAS::PRIVATE_ADDRESS,  // hlsl_input
 };
 } // namespace targets
 } // namespace clang

clang/lib/Basic/Targets/DirectX.h

@@ -45,6 +45,7 @@ static const unsigned DirectXAddrSpaceMap[] = {
     2, // hlsl_constant
     0, // hlsl_private
     0, // hlsl_device
+    0, // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/Basic/Targets/NVPTX.h

@@ -49,6 +49,7 @@ static const unsigned NVPTXAddrSpaceMap[] = {
     0, // hlsl_constant
     0, // hlsl_private
     0, // hlsl_device
+    0, // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/Basic/Targets/SPIR.h

@@ -50,6 +50,7 @@ static const unsigned SPIRDefIsPrivMap[] = {
     12, // hlsl_constant
     10, // hlsl_private
     11, // hlsl_device
+    7,  // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref
@@ -85,6 +86,7 @@ static const unsigned SPIRDefIsGenMap[] = {
     0,  // hlsl_constant
     10, // hlsl_private
     11, // hlsl_device
+    7,  // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/Basic/Targets/SystemZ.h

@@ -45,6 +45,7 @@ static const unsigned ZOSAddressMap[] = {
     0, // hlsl_constant
     0, // hlsl_private
     0, // hlsl_device
+    0, // hlsl_input
     0  // wasm_funcref
 };
 

clang/lib/Basic/Targets/TCE.h

@@ -54,6 +54,7 @@ static const unsigned TCEOpenCLAddrSpaceMap[] = {
     0, // hlsl_constant
     0, // hlsl_private
     0, // hlsl_device
+    0, // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/Basic/Targets/WebAssembly.h

@@ -45,6 +45,7 @@ static const unsigned WebAssemblyAddrSpaceMap[] = {
     0,  // hlsl_constant
     0,  // hlsl_private
     0,  // hlsl_device
+    0,  // hlsl_input
     20, // wasm_funcref
 };
 

clang/lib/Basic/Targets/X86.h

@@ -49,6 +49,7 @@ static const unsigned X86AddrSpaceMap[] = {
     0,   // hlsl_constant
     0,   // hlsl_private
     0,   // hlsl_device
+    0,   // hlsl_input
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -624,6 +624,20 @@ void CGHLSLRuntime::initializeBufferFromBinding(const HLSLBufferDecl *BufDecl,
   }
 }
 
+void CGHLSLRuntime::handleGlobalVarDefinition(const VarDecl *VD,
+                                              llvm::GlobalVariable *GV) {
+  if (auto Attr = VD->getAttr<HLSLVkExtBuiltinInputAttr>()) {
+    LLVMContext &Ctx = GV->getContext();
+    IRBuilder<> B(GV->getContext());
+    MDNode *Operands = MDNode::get(
+        Ctx, {ConstantAsMetadata::get(
+                  B.getInt32(/* Spirv::Decoration::BuiltIn */ 11)),
+              ConstantAsMetadata::get(B.getInt32(Attr->getBuiltIn()))});
+    MDNode *Decoration = MDNode::get(Ctx, {Operands});
+    GV->addMetadata("spirv.Decorations", *Decoration);
+  }
+}
+
 llvm::Instruction *CGHLSLRuntime::getConvergenceToken(BasicBlock &BB) {
   if (!CGM.shouldEmitConvergenceTokens())
     return nullptr;

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -157,6 +157,7 @@ class CGHLSLRuntime {
 
   void emitEntryFunction(const FunctionDecl *FD, llvm::Function *Fn);
   void setHLSLFunctionAttributes(const FunctionDecl *FD, llvm::Function *Fn);
+  void handleGlobalVarDefinition(const VarDecl *VD, llvm::GlobalVariable *Var);
 
   llvm::Instruction *getConvergenceToken(llvm::BasicBlock &BB);
 

clang/lib/CodeGen/CodeGenModule.cpp

@@ -5763,7 +5763,17 @@ void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D,
     getCUDARuntime().handleVarRegistration(D, *GV);
   }
 
-  GV->setInitializer(Init);
+  if (LangOpts.HLSL && GetGlobalVarAddressSpace(D) == LangAS::hlsl_input) {
+    // HLSL Input variables are considered to be set by the driver/pipeline, but
+    // only visible to a single thread/wave.
+    GV->setExternallyInitialized(true);
+  } else {
+    GV->setInitializer(Init);
+  }
+
+  if (LangOpts.HLSL)
+    getHLSLRuntime().handleGlobalVarDefinition(D, GV);
+
   if (emitter)
     emitter->finalize(GV);
 
@@ -5806,6 +5816,12 @@ void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D,
       !D->hasAttr<ConstInitAttr>())
     Linkage = llvm::GlobalValue::InternalLinkage;
 
+  // HLSL variables in the input address space maps like memory-mapped
+  // variables. Even if they are 'static', they are externally initialized and
+  // read/write by the hardware/driver/pipeline.
+  if (LangOpts.HLSL && GetGlobalVarAddressSpace(D) == LangAS::hlsl_input)
+    Linkage = llvm::GlobalValue::ExternalLinkage;
+
   GV->setLinkage(Linkage);
   if (D->hasAttr<DLLImportAttr>())
     GV->setDLLStorageClass(llvm::GlobalVariable::DLLImportStorageClass);

clang/lib/Sema/SemaDecl.cpp

@@ -14419,6 +14419,12 @@ void Sema::ActOnUninitializedDecl(Decl *RealDecl) {
     if (getLangOpts().HLSL && HLSL().ActOnUninitializedVarDecl(Var))
       return;
 
+    // HLSL input variables are expected to be externally initialized, even
+    // when marked `static`.
+    if (getLangOpts().HLSL &&
+        Var->getType().getAddressSpace() == LangAS::hlsl_input)
+      return;
+
     // C++03 [dcl.init]p9:
     //   If no initializer is specified for an object, and the
     //   object is of (possibly cv-qualified) non-POD class type (or

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_HLSLVkExtBuiltinInput:
+    S.HLSL().handleVkExtBuiltinInputAttr(D, AL);
+    break;
   case ParsedAttr::AT_HLSLSV_GroupThreadID:
     S.HLSL().handleSV_GroupThreadIDAttr(D, AL);
     break;

clang/lib/Sema/SemaHLSL.cpp

@@ -1117,6 +1117,14 @@ void SemaHLSL::handleWaveSizeAttr(Decl *D, const ParsedAttr &AL) {
     D->addAttr(NewAttr);
 }
 
+void SemaHLSL::handleVkExtBuiltinInputAttr(Decl *D, const ParsedAttr &AL) {
+  uint32_t ID;
+  if (!SemaRef.checkUInt32Argument(AL, AL.getArgAsExpr(0), ID))
+    return;
+  D->addAttr(::new (getASTContext())
+                 HLSLVkExtBuiltinInputAttr(getASTContext(), AL, ID));
+}
+
 bool SemaHLSL::diagnoseInputIDType(QualType T, const ParsedAttr &AL) {
   const auto *VT = T->getAs<VectorType>();
 
@@ -3158,6 +3166,14 @@ void SemaHLSL::deduceAddressSpace(VarDecl *Decl) {
     return;
 
   QualType Type = Decl->getType();
+
+  if (Decl->hasAttr<HLSLVkExtBuiltinInputAttr>()) {
+    LangAS ImplAS = LangAS::hlsl_input;
+    Type = SemaRef.getASTContext().getAddrSpaceQualType(Type, ImplAS);
+    Decl->setType(Type);
+    return;
+  }
+
   if (Type->isSamplerT() || Type->isVoidType())
     return;
 

clang/test/CodeGenHLSL/vk-input-builtin.hlsl

@@ -0,0 +1,15 @@
+// RUN: %clang_cc1 -finclude-default-header -x hlsl -triple \
+// RUN:   spirv-unknown-vulkan1.3-compute %s -emit-llvm -O3 -o - | FileCheck %s
+
+[[vk::ext_builtin_input(/* WorkgroupId */ 26)]]
+static const uint3 groupid;
+// CHECK: @_ZL7groupid = external local_unnamed_addr addrspace(7) externally_initialized constant <3 x i32>, align 16, !spirv.Decorations [[META0:![0-9]+]]
+
+RWStructuredBuffer<int> output : register(u1, space0);
+
+[numthreads(1, 1, 1)]
+void main() {
+  output[0] = groupid;
+}
+// CHECK: [[META0]] = !{[[META1:![0-9]+]]}
+// CHECK: [[META1]] = !{i32 11, i32 26}

clang/test/SemaHLSL/vk-ext-input-builtin.hlsl

@@ -0,0 +1,29 @@
+// RUN: %clang_cc1 -triple spirv-unkown-vulkan1.3-compute -x hlsl -hlsl-entry foo  -finclude-default-header -o - %s -verify
+
+// expected-error@+1 {{'ext_builtin_input' attribute only applies to static const globals}}
+[[vk::ext_builtin_input(/* WorkgroupId */ 26)]]
+const uint3 groupid1;
+
+// expected-error@+1 {{'ext_builtin_input' attribute only applies to static const globals}}
+[[vk::ext_builtin_input(/* WorkgroupId */ 26)]]
+static uint3 groupid2;
+
+// expected-error@+1 {{'ext_builtin_input' attribute takes one argument}}
+[[vk::ext_builtin_input()]]
+// expected-error@+1 {{default initialization of an object of const type 'const hlsl_private uint3' (aka 'const hlsl_private vector<uint, 3>')}}
+static const uint3 groupid3;
+
+// expected-error@+1 {{'ext_builtin_input' attribute requires an integer constant}}
+[[vk::ext_builtin_input(0.4f)]]
+// expected-error@+1 {{default initialization of an object of const type 'const hlsl_private uint3' (aka 'const hlsl_private vector<uint, 3>')}}
+static const uint3 groupid4;
+
+// expected-error@+1 {{'ext_builtin_input' attribute only applies to static const globals}}
+[[vk::ext_builtin_input(1)]]
+void some_function() {
+}
+
+[numthreads(1,1,1)]
+void foo() {
+}
+

clang/test/SemaTemplate/address_space-dependent.cpp

@@ -43,7 +43,7 @@ void neg() {
 
 template <long int I>
 void tooBig() {
-  __attribute__((address_space(I))) int *bounds; // expected-error {{address space is larger than the maximum supported (8388583)}}
+  __attribute__((address_space(I))) int *bounds; // expected-error {{address space is larger than the maximum supported (8388582)}}
 }
 
 template <long int I>
@@ -101,7 +101,7 @@ int main() {
   car<1, 2, 3>(); // expected-note {{in instantiation of function template specialization 'car<1, 2, 3>' requested here}}
   HasASTemplateFields<1> HASTF;
   neg<-1>(); // expected-note {{in instantiation of function template specialization 'neg<-1>' requested here}}
-  correct<0x7FFFE7>();
+  correct<0x7FFFE6>();
   tooBig<8388650>(); // expected-note {{in instantiation of function template specialization 'tooBig<8388650L>' requested here}}
 
   __attribute__((address_space(1))) char *x;