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

Implements initial support for vk::push_constant.
As is, this allows handling simple push constants, but has one
main issue: layout can be incorrect (See #168401). The layout
issue being not only push-constant related, it's ignored for this PR.

The frontend part of the implementation is straightforward:
 - adding a new attribute
 - when targeting vulkan/spirv, we process it
 - global variables with this attribute gets a new AS:
   hlsl_push_constant

The IR has nothing specific, only some RO globals in this new AS.

On the SPIR-V side, we not convert this AS into a PushConstant storage
class. But this creates some issues: the variables in this storage class
must have a specific set of decoration to define their layout.

Current infra to create the SPIR-V types lacks the context required to
make this decision: no indication on the AS or context around the type
being created. Refactoring this would be a heavy task as it would
require getting this information in every place using the GR for type
creation.

Instead, we do something similar to CBuffers:
 - find all globals with this address space, and change their type to
   a target-specific type.
 - insert a new intrinsic in place of every reference to this global
   variable.

This allow the backend to handle both layout variables loads and type
lowering independently.

Type lowering has nothing specific: when we encounter a target extension
type with spirv.PushConstant, we lower this to the correct SPIR-V type
with the proper offset & block decorations.

As for the intrinsic, it's mostly a no-op, but required since we have
this target-specific type.

Note: this implementation prevents the static declaration of multiple
push constants in a single shader module. The actual specification is
more relaxed: there can be only one **used** push constant block per
entrypoint. To correctly implement this, we'd require to keep some
additional state to determine the list of statically used resources per
entrypoint. This shall be addressed as a follow-up (see #170310)

Author: Keenuts

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

@@ -5164,6 +5164,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

@@ -8794,6 +8794,22 @@ https://github.com/microsoft/hlsl-specs/blob/main/proposals/0011-inline-spirv.md
   }];
 }
 
+def HLSLVkPushConstantDocs : Documentation {
+  let Category = DocCatVariable;
+  let Content = [{
+Vulkan shaders have `PushConstants`
+
+The ``[[vk::push_constant]]`` attribute allows you to declare this
+global variable as a push constant when targeting Vulkan.
+This attribute is ignored otherwise.
+
+This attribute must be applied to the variable, not underlying type.
+The variable type must be a struct, per the requirements of Vulkan, "there
+must be no more than one push constant block statically used per shader entry
+point."
+}];
+}
+
 def AnnotateTypeDocs : Documentation {
   let Category = DocCatType;
   let Heading = "annotate_type";

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -13216,6 +13216,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">;
@@ -13333,6 +13336,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

@@ -188,6 +188,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);
@@ -237,6 +238,8 @@ class SemaHLSL : public SemaBase {
 
   IdentifierInfo *RootSigOverrideIdent = nullptr;
 
+  bool HasDeclaredAPushConstant = false;
+
   // Information about the current subtree being flattened.
   struct SemanticInfo {
     HLSLParsedSemanticAttr *Semantic;

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

@@ -2694,6 +2694,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