[HLSL][SPIR-V] Implement vk::push_constant

clang/include/clang/Basic/AddressSpaces.h

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

clang/include/clang/Basic/Attr.td

@@ -5146,6 +5146,14 @@ def HLSLVkExtBuiltinInput : InheritableAttr {
   let Documentation = [HLSLVkExtBuiltinInputDocs];
 }
 
+def HLSLVkPushConstant : InheritableAttr {
+  let Spellings = [CXX11<"vk", "push_constant">];
+  let Args = [];
+  let Subjects = SubjectList<[GlobalVar], ErrorDiag>;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLVkPushConstantDocs];
+}
+
 def HLSLVkConstantId : InheritableAttr {
   let Spellings = [CXX11<"vk", "constant_id">];
   let Args = [IntArgument<"Id">];

clang/include/clang/Basic/AttrDocs.td

@@ -8777,6 +8777,11 @@ https://github.com/microsoft/hlsl-specs/blob/main/proposals/0011-inline-spirv.md
   }];
 }
 
+def HLSLVkPushConstantDocs : Documentation {
+  let Category = DocCatVariable;
+  let Content = [{ FIXME }];
+}
+
 def AnnotateTypeDocs : Documentation {
   let Category = DocCatType;
   let Heading = "annotate_type";

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -13183,6 +13183,9 @@ def err_hlsl_attr_invalid_type : Error<
    "attribute %0 only applies to a field or parameter of type '%1'">;
 def err_hlsl_attr_invalid_ast_node : Error<
    "attribute %0 only applies to %1">;
+def err_hlsl_attr_incompatible
+    : Error<"%0 attribute is not compatible with %1 attribute">;
+
 def err_hlsl_entry_shader_attr_mismatch : Error<
    "%0 attribute on entry function does not match the target profile">;
 def err_hlsl_numthreads_argument_oor : Error<"argument '%select{X|Y|Z}0' to numthreads attribute cannot exceed %1">;
@@ -13294,6 +13297,9 @@ def err_hlsl_incomplete_resource_array_in_function_param: Error<
 def err_hlsl_assign_to_global_resource: Error<
   "assignment to global resource variable %0 is not allowed">;
 
+def err_hlsl_push_constant_unique
+    : Error<"cannot have more than one push constant block">;
+
 // Layout randomization diagnostics.
 def err_non_designated_init_used : Error<
   "a randomized struct can only be initialized with a designated initializer">;

clang/include/clang/Basic/HLSLRuntime.h

@@ -14,6 +14,7 @@
 #ifndef CLANG_BASIC_HLSLRUNTIME_H
 #define CLANG_BASIC_HLSLRUNTIME_H
 
+#include "clang/Basic/AddressSpaces.h"
 #include "clang/Basic/LangOptions.h"
 #include <cstdint>
 
@@ -30,6 +31,10 @@ getStageFromEnvironment(const llvm::Triple::EnvironmentType &E) {
   return static_cast<ShaderStage>(Pipeline);
 }
 
+constexpr bool isInitializedByPipeline(LangAS AS) {
+  return AS == LangAS::hlsl_input || AS == LangAS::hlsl_push_constant;
+}
+
 #define ENUM_COMPARE_ASSERT(Value)                                             \
   static_assert(                                                               \
       getStageFromEnvironment(llvm::Triple::Value) == ShaderStage::Value,      \

clang/include/clang/Sema/SemaHLSL.h

@@ -190,6 +190,7 @@ class SemaHLSL : public SemaBase {
   void handleSemanticAttr(Decl *D, const ParsedAttr &AL);
 
   void handleVkExtBuiltinInputAttr(Decl *D, const ParsedAttr &AL);
+  void handleVkPushConstantAttr(Decl *D, const ParsedAttr &AL);
 
   bool CheckBuiltinFunctionCall(unsigned BuiltinID, CallExpr *TheCall);
   QualType ProcessResourceTypeAttributes(QualType Wrapped);
@@ -239,6 +240,8 @@ class SemaHLSL : public SemaBase {
 
   IdentifierInfo *RootSigOverrideIdent = nullptr;
 
+  bool HasDeclaredAPushConstant = false;
+
   struct SemanticInfo {
     HLSLParsedSemanticAttr *Semantic;
     std::optional<uint32_t> Index;

clang/lib/AST/Type.cpp

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

clang/lib/AST/TypePrinter.cpp

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

clang/lib/Basic/TargetInfo.cpp

@@ -52,6 +52,7 @@ static const LangASMap FakeAddrSpaceMap = {
     15, // hlsl_private
     16, // hlsl_device
     17, // hlsl_input
+    18, // hlsl_push_constant
     20, // wasm_funcref
 };
 

clang/lib/Basic/Targets/AArch64.h

@@ -48,6 +48,7 @@ static const unsigned ARM64AddrSpaceMap[] = {
     0, // hlsl_private
     0, // hlsl_device
     0, // hlsl_input
+    0, // hlsl_push_constant
     // 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

@@ -63,6 +63,7 @@ const LangASMap AMDGPUTargetInfo::AMDGPUDefIsGenMap = {
     llvm::AMDGPUAS::PRIVATE_ADDRESS, // hlsl_private
     llvm::AMDGPUAS::GLOBAL_ADDRESS,  // hlsl_device
     llvm::AMDGPUAS::PRIVATE_ADDRESS, // hlsl_input
+    llvm::AMDGPUAS::GLOBAL_ADDRESS,  // hlsl_push_constant
 };
 
 const LangASMap AMDGPUTargetInfo::AMDGPUDefIsPrivMap = {
@@ -91,6 +92,7 @@ const LangASMap AMDGPUTargetInfo::AMDGPUDefIsPrivMap = {
     llvm::AMDGPUAS::PRIVATE_ADDRESS,  // hlsl_private
     llvm::AMDGPUAS::GLOBAL_ADDRESS,   // hlsl_device
     llvm::AMDGPUAS::PRIVATE_ADDRESS,  // hlsl_input
+    llvm::AMDGPUAS::GLOBAL_ADDRESS,   // hlsl_push_constant
 };
 } // namespace targets
 } // namespace clang

clang/lib/Basic/Targets/DirectX.h

@@ -46,6 +46,7 @@ static const unsigned DirectXAddrSpaceMap[] = {
     0, // hlsl_private
     0, // hlsl_device
     0, // hlsl_input
+    0, // hlsl_push_constant
     // 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

@@ -50,6 +50,7 @@ static const unsigned NVPTXAddrSpaceMap[] = {
     0, // hlsl_private
     0, // hlsl_device
     0, // hlsl_input
+    0, // hlsl_push_constant
     // 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

@@ -51,6 +51,7 @@ static const unsigned SPIRDefIsPrivMap[] = {
     10, // hlsl_private
     11, // hlsl_device
     7,  // hlsl_input
+    13, // hlsl_push_constant
     // Wasm address space values for this target are dummy values,
     // as it is only enabled for Wasm targets.
     20, // wasm_funcref
@@ -87,6 +88,7 @@ static const unsigned SPIRDefIsGenMap[] = {
     10, // hlsl_private
     11, // hlsl_device
     7,  // hlsl_input
+    13, // hlsl_push_constant
     // 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

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

clang/lib/Basic/Targets/TCE.h

@@ -55,6 +55,7 @@ static const unsigned TCEOpenCLAddrSpaceMap[] = {
     0, // hlsl_private
     0, // hlsl_device
     0, // hlsl_input
+    0, // hlsl_push_constant
     // 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

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

clang/lib/Basic/Targets/X86.h

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

clang/lib/CodeGen/CodeGenModule.cpp

@@ -6049,9 +6049,11 @@ void CodeGenModule::EmitGlobalVarDefinition(const VarDecl *D,
     getCUDARuntime().handleVarRegistration(D, *GV);
   }
 
-  if (LangOpts.HLSL && GetGlobalVarAddressSpace(D) == LangAS::hlsl_input) {
+  if (LangOpts.HLSL &&
+      hlsl::isInitializedByPipeline(GetGlobalVarAddressSpace(D))) {
     // HLSL Input variables are considered to be set by the driver/pipeline, but
-    // only visible to a single thread/wave.
+    // only visible to a single thread/wave. Push constants are also externally
+    // initialized, but constant, hence cross-wave visibility is not relevant.
     GV->setExternallyInitialized(true);
   } else {
     GV->setInitializer(Init);
@@ -6102,10 +6104,11 @@ 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)
+  // HLSL variables in the input or push-constant address space maps are like
+  // memory-mapped variables. Even if they are 'static', they are externally
+  // initialized and read/write by the hardware/driver/pipeline.
+  if (LangOpts.HLSL &&
+      hlsl::isInitializedByPipeline(GetGlobalVarAddressSpace(D)))
     Linkage = llvm::GlobalValue::ExternalLinkage;
 
   GV->setLinkage(Linkage);

clang/lib/Sema/SemaDecl.cpp

@@ -30,6 +30,7 @@
 #include "clang/AST/Type.h"
 #include "clang/Basic/Builtins.h"
 #include "clang/Basic/DiagnosticComment.h"
+#include "clang/Basic/HLSLRuntime.h"
 #include "clang/Basic/PartialDiagnostic.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Basic/TargetInfo.h"
@@ -14559,10 +14560,10 @@ 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`.
+    // HLSL input & push-constant variables are expected to be externally
+    // initialized, even when marked `static`.
     if (getLangOpts().HLSL &&
-        Var->getType().getAddressSpace() == LangAS::hlsl_input)
+        hlsl::isInitializedByPipeline(Var->getType().getAddressSpace()))
       return;
 
     // C++03 [dcl.init]p9:

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7614,6 +7614,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_HLSLVkExtBuiltinInput:
     S.HLSL().handleVkExtBuiltinInputAttr(D, AL);
     break;
+  case ParsedAttr::AT_HLSLVkPushConstant:
+    S.HLSL().handleVkPushConstantAttr(D, AL);
+    break;
   case ParsedAttr::AT_HLSLVkConstantId:
     S.HLSL().handleVkConstantIdAttr(D, AL);
     break;

clang/lib/Sema/SemaHLSL.cpp

@@ -1667,6 +1667,11 @@ void SemaHLSL::handleVkExtBuiltinInputAttr(Decl *D, const ParsedAttr &AL) {
                  HLSLVkExtBuiltinInputAttr(getASTContext(), AL, ID));
 }
 
+void SemaHLSL::handleVkPushConstantAttr(Decl *D, const ParsedAttr &AL) {
+  D->addAttr(::new (getASTContext())
+                 HLSLVkPushConstantAttr(getASTContext(), AL));
+}
+
 void SemaHLSL::handleVkConstantIdAttr(Decl *D, const ParsedAttr &AL) {
   uint32_t Id;
   if (!SemaRef.checkUInt32Argument(AL, AL.getArgAsExpr(0), Id))
@@ -3837,12 +3842,15 @@ QualType SemaHLSL::getInoutParameterType(QualType Ty) {
   return Ty;
 }
 
-static bool IsDefaultBufferConstantDecl(VarDecl *VD) {
+static bool IsDefaultBufferConstantDecl(const ASTContext &Ctx, VarDecl *VD) {
+  bool IsVulkan =
+      Ctx.getTargetInfo().getTriple().getOS() == llvm::Triple::Vulkan;
+  bool IsVKPushConstant = IsVulkan && VD->hasAttr<HLSLVkPushConstantAttr>();
   QualType QT = VD->getType();
   return VD->getDeclContext()->isTranslationUnit() &&
          QT.getAddressSpace() == LangAS::Default &&
          VD->getStorageClass() != SC_Static &&
-         !VD->hasAttr<HLSLVkConstantIdAttr>() &&
+         !VD->hasAttr<HLSLVkConstantIdAttr>() && !IsVKPushConstant &&
          !isInvalidConstantBufferLeafElementType(QT.getTypePtr());
 }
 
@@ -3863,6 +3871,19 @@ void SemaHLSL::deduceAddressSpace(VarDecl *Decl) {
     return;
   }
 
+  bool IsVulkan = getASTContext().getTargetInfo().getTriple().getOS() ==
+                  llvm::Triple::Vulkan;
+  if (IsVulkan && Decl->hasAttr<HLSLVkPushConstantAttr>()) {
+    if (HasDeclaredAPushConstant)
+      SemaRef.Diag(Decl->getLocation(), diag::err_hlsl_push_constant_unique);
+
+    LangAS ImplAS = LangAS::hlsl_push_constant;
+    Type = SemaRef.getASTContext().getAddrSpaceQualType(Type, ImplAS);
+    Decl->setType(Type);
+    HasDeclaredAPushConstant = true;
+    return;
+  }
+
   if (Type->isSamplerT() || Type->isVoidType())
     return;
 
@@ -3895,7 +3916,7 @@ void SemaHLSL::ActOnVariableDeclarator(VarDecl *VD) {
     // Global variables outside a cbuffer block that are not a resource, static,
     // groupshared, or an empty array or struct belong to the default constant
     // buffer $Globals (to be created at the end of the translation unit).
-    if (IsDefaultBufferConstantDecl(VD)) {
+    if (IsDefaultBufferConstantDecl(getASTContext(), VD)) {
       // update address space to hlsl_constant
       QualType NewTy = getASTContext().getAddrSpaceQualType(
           VD->getType(), LangAS::hlsl_constant);
@@ -4196,8 +4217,11 @@ void SemaHLSL::processExplicitBindingsOnDecl(VarDecl *VD) {
 
   bool HasBinding = false;
   for (Attr *A : VD->attrs()) {
-    if (isa<HLSLVkBindingAttr>(A))
+    if (isa<HLSLVkBindingAttr>(A)) {
       HasBinding = true;
+      if (auto PA = VD->getAttr<HLSLVkPushConstantAttr>())
+        Diag(PA->getLoc(), diag::err_hlsl_attr_incompatible) << A << PA;
+    }
 
     HLSLResourceBindingAttr *RBA = dyn_cast<HLSLResourceBindingAttr>(A);
     if (!RBA || !RBA->hasRegisterSlot())

clang/test/AST/HLSL/vk.pushconstant.hlsl

@@ -0,0 +1,12 @@
+// RUN: %clang_cc1 -triple spirv-unknown-vulkan1.3-compute -x hlsl -ast-dump -o - %s | FileCheck %s
+
+struct S {
+  int value;
+};
+
+[[vk::push_constant]] S PC;
+// CHECK:      VarDecl 0x[[A:[0-9a-f]+]] <line:7:23, col:25> col:25 PC 'hlsl_push_constant S'
+// CHECK-NEXT: HLSLVkPushConstantAttr 0x[[A:[0-9a-f]+]] <col:3, col:7>
+
+[numthreads(1, 1, 1)]
+void main() { }

clang/test/CodeGenHLSL/vk-features/vk.pushconstant.access.bitfield.hlsl

@@ -0,0 +1,20 @@
+// RUN: %clang_cc1 -triple spirv-pc-vulkan-compute -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s
+
+struct S {
+  uint32_t a : 1;
+  uint32_t b : 1;
+};
+// CHECK: %struct.S = type { i8 }
+
+[[vk::push_constant]] S buffer;
+// CHECK: @buffer = external hidden addrspace(13) externally_initialized global %struct.S, align 1
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint32_t v = buffer.b;
+// CHECK:  %bf.load = load i8, ptr addrspace(13) @buffer, align 1
+// CHECK:  %bf.lshr = lshr i8 %bf.load, 1
+// CHECK: %bf.clear = and i8 %bf.lshr, 1
+// CHECK:  %bf.cast = zext i8 %bf.clear to i32
+// CHECK:             store i32 %bf.cast
+}

clang/test/CodeGenHLSL/vk-features/vk.pushconstant.access.hlsl

@@ -0,0 +1,16 @@
+// RUN: %clang_cc1 -triple spirv-pc-vulkan-compute -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s
+
+struct S {
+  uint value;
+};
+// CHECK: %struct.S = type { i32 }
+
+[[vk::push_constant]] S buffer;
+// CHECK: @buffer = external hidden addrspace(13) externally_initialized global %struct.S, align 1
+
+[numthreads(1, 1, 1)]
+void main() {
+  uint32_t v = buffer.value;
+// CHECK:  %[[#REG:]] = load i32, ptr addrspace(13) @buffer, align 1
+// CHECK:             store i32 %[[#REG]], ptr %v, align 4
+}