[HLSL][SPIR-V] Implement vk::location for inputs

clang/include/clang/Basic/Attr.td

@@ -5172,6 +5172,14 @@ def HLSLVkConstantId : InheritableAttr {
   let Documentation = [VkConstantIdDocs];
 }
 
+def HLSLVkLocation : HLSLAnnotationAttr {
+  let Spellings = [CXX11<"vk", "location">];
+  let Args = [IntArgument<"Location">];
+  let Subjects = SubjectList<[ParmVar, Field, Function], ErrorDiag>;
+  let LangOpts = [HLSL];
+  let Documentation = [HLSLVkLocationDocs];
+}
+
 def RandomizeLayout : InheritableAttr {
   let Spellings = [GCC<"randomize_layout">];
   let Subjects = SubjectList<[Record]>;

clang/include/clang/Basic/AttrDocs.td

@@ -8981,6 +8981,18 @@ The descriptor set is optional and defaults to 0 if not provided.
   }];
 }
 
+def HLSLVkLocationDocs : Documentation {
+  let Category = DocCatVariable;
+  let Content = [{
+Attribute used for specifying the location number for the stage input/output
+variables. Allowed on function parameters, function returns, and struct
+fields. This parameter has no effect when used outside of an entrypoint
+parameter/parameter field/return value.
+
+This attribute maps to the 'Location' SPIR-V decoration.
+  }];
+}
+
 def WebAssemblyFuncrefDocs : Documentation {
   let Category = DocCatType;
   let Content = [{

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -13234,6 +13234,9 @@ def err_hlsl_semantic_index_overlap : Error<"semantic index overlap %0">;
 def err_hlsl_semantic_unsupported_iotype_for_stage
     : Error<"semantic %0 is unsupported in %2 shaders as %1, requires one of "
             "the following: %3">;
+def err_hlsl_semantic_partial_explicit_indexing
+    : Error<"partial explicit stage input location assignment via "
+            "vk::location(X) unsupported">;
 
 def warn_hlsl_user_defined_type_missing_member: Warning<"binding type '%select{t|u|b|s|c}0' only applies to types containing %select{SRV resources|UAV resources|constant buffer resources|sampler state|numeric types}0">, InGroup<LegacyConstantRegisterBinding>;
 def err_hlsl_binding_type_mismatch: Error<"binding type '%select{t|u|b|s|c}0' only applies to %select{SRV resources|UAV resources|constant buffer resources|sampler state|numeric variables in the global scope}0">;

clang/include/clang/Sema/SemaHLSL.h

@@ -168,6 +168,7 @@ class SemaHLSL : public SemaBase {
   void handleWaveSizeAttr(Decl *D, const ParsedAttr &AL);
   void handleVkConstantIdAttr(Decl *D, const ParsedAttr &AL);
   void handleVkBindingAttr(Decl *D, const ParsedAttr &AL);
+  void handleVkLocationAttr(Decl *D, const ParsedAttr &AL);
   void handlePackOffsetAttr(Decl *D, const ParsedAttr &AL);
   void handleShaderAttr(Decl *D, const ParsedAttr &AL);
   void handleResourceBindingAttr(Decl *D, const ParsedAttr &AL);
@@ -240,6 +241,8 @@ class SemaHLSL : public SemaBase {
     HLSLParsedSemanticAttr *Semantic;
     std::optional<uint32_t> Index;
   };
+  std::optional<bool> InputUsesExplicitVkLocations = std::nullopt;
+  std::optional<bool> OutputUsesExplicitVkLocations = std::nullopt;
 
   enum IOType {
     In = 0b01,

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -582,20 +582,22 @@ static llvm::Value *createSPIRVLocationLoad(IRBuilder<> &B, llvm::Module &M,
   return B.CreateLoad(Ty, GV);
 }
 
-llvm::Value *
-CGHLSLRuntime::emitSPIRVUserSemanticLoad(llvm::IRBuilder<> &B, llvm::Type *Type,
-                                         HLSLAppliedSemanticAttr *Semantic,
-                                         std::optional<unsigned> Index) {
+llvm::Value *CGHLSLRuntime::emitSPIRVUserSemanticLoad(
+    llvm::IRBuilder<> &B, llvm::Type *Type, const clang::DeclaratorDecl *Decl,
+    HLSLAppliedSemanticAttr *Semantic, std::optional<unsigned> Index) {
   Twine BaseName = Twine(Semantic->getAttrName()->getName());
   Twine VariableName = BaseName.concat(Twine(Index.value_or(0)));
 
   unsigned Location = SPIRVLastAssignedInputSemanticLocation;
+  if (auto *L = Decl->getAttr<HLSLVkLocationAttr>())
+    Location = L->getLocation();
 
   // DXC completely ignores the semantic/index pair. Location are assigned from
   // the first semantic to the last.
   llvm::ArrayType *AT = dyn_cast<llvm::ArrayType>(Type);
   unsigned ElementCount = AT ? AT->getNumElements() : 1;
   SPIRVLastAssignedInputSemanticLocation += ElementCount;
+
   return createSPIRVLocationLoad(B, CGM.getModule(), Type, Location,
                                  VariableName.str());
 }
@@ -616,10 +618,14 @@ static void createSPIRVLocationStore(IRBuilder<> &B, llvm::Module &M,
 
 void CGHLSLRuntime::emitSPIRVUserSemanticStore(
     llvm::IRBuilder<> &B, llvm::Value *Source,
-    HLSLAppliedSemanticAttr *Semantic, std::optional<unsigned> Index) {
+    const clang::DeclaratorDecl *Decl, HLSLAppliedSemanticAttr *Semantic,
+    std::optional<unsigned> Index) {
   Twine BaseName = Twine(Semantic->getAttrName()->getName());
   Twine VariableName = BaseName.concat(Twine(Index.value_or(0)));
+
   unsigned Location = SPIRVLastAssignedOutputSemanticLocation;
+  if (auto *L = Decl->getAttr<HLSLVkLocationAttr>())
+    Location = L->getLocation();
 
   // DXC completely ignores the semantic/index pair. Location are assigned from
   // the first semantic to the last.
@@ -671,7 +677,7 @@ llvm::Value *CGHLSLRuntime::emitUserSemanticLoad(
     IRBuilder<> &B, llvm::Type *Type, const clang::DeclaratorDecl *Decl,
     HLSLAppliedSemanticAttr *Semantic, std::optional<unsigned> Index) {
   if (CGM.getTarget().getTriple().isSPIRV())
-    return emitSPIRVUserSemanticLoad(B, Type, Semantic, Index);
+    return emitSPIRVUserSemanticLoad(B, Type, Decl, Semantic, Index);
 
   if (CGM.getTarget().getTriple().isDXIL())
     return emitDXILUserSemanticLoad(B, Type, Semantic, Index);
@@ -684,7 +690,7 @@ void CGHLSLRuntime::emitUserSemanticStore(IRBuilder<> &B, llvm::Value *Source,
                                           HLSLAppliedSemanticAttr *Semantic,
                                           std::optional<unsigned> Index) {
   if (CGM.getTarget().getTriple().isSPIRV())
-    return emitSPIRVUserSemanticStore(B, Source, Semantic, Index);
+    return emitSPIRVUserSemanticStore(B, Source, Decl, Semantic, Index);
 
   if (CGM.getTarget().getTriple().isDXIL())
     return emitDXILUserSemanticStore(B, Source, Semantic, Index);

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -278,6 +278,7 @@ class CGHLSLRuntime {
                                    HLSLResourceBindingAttr *RBA);
 
   llvm::Value *emitSPIRVUserSemanticLoad(llvm::IRBuilder<> &B, llvm::Type *Type,
+                                         const clang::DeclaratorDecl *Decl,
                                          HLSLAppliedSemanticAttr *Semantic,
                                          std::optional<unsigned> Index);
   llvm::Value *emitDXILUserSemanticLoad(llvm::IRBuilder<> &B, llvm::Type *Type,
@@ -289,6 +290,7 @@ class CGHLSLRuntime {
                                     std::optional<unsigned> Index);
 
   void emitSPIRVUserSemanticStore(llvm::IRBuilder<> &B, llvm::Value *Source,
+                                  const clang::DeclaratorDecl *Decl,
                                   HLSLAppliedSemanticAttr *Semantic,
                                   std::optional<unsigned> Index);
   void emitDXILUserSemanticStore(llvm::IRBuilder<> &B, llvm::Value *Source,

clang/lib/Sema/SemaDeclAttr.cpp

@@ -7703,6 +7703,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, const ParsedAttr &AL,
   case ParsedAttr::AT_HLSLUnparsedSemantic:
     S.HLSL().handleSemanticAttr(D, AL);
     break;
+  case ParsedAttr::AT_HLSLVkLocation:
+    S.HLSL().handleVkLocationAttr(D, AL);
+    break;
 
   case ParsedAttr::AT_AbiTag:
     handleAbiTagAttr(S, D, AL);

clang/lib/Sema/SemaHLSL.cpp

@@ -771,6 +771,22 @@ void SemaHLSL::ActOnTopLevelFunction(FunctionDecl *FD) {
   }
 }
 
+static bool isPipelineBuiltin(const ASTContext &AstContext, FunctionDecl *FD,
+                              HLSLAppliedSemanticAttr *Semantic) {
+  if (AstContext.getTargetInfo().getTriple().getOS() != llvm::Triple::Vulkan)
+    return false;
+
+  const auto *ShaderAttr = FD->getAttr<HLSLShaderAttr>();
+  assert(ShaderAttr && "Entry point has no shader attribute");
+  llvm::Triple::EnvironmentType ST = ShaderAttr->getType();
+  auto SemanticName = Semantic->getSemanticName().upper();
+
+  if (ST == llvm::Triple::Pixel && SemanticName == "SV_POSITION")
+    return true;
+
+  return false;
+}
+
 bool SemaHLSL::determineActiveSemanticOnScalar(FunctionDecl *FD,
                                                DeclaratorDecl *OutputDecl,
                                                DeclaratorDecl *D,
@@ -800,6 +816,22 @@ bool SemaHLSL::determineActiveSemanticOnScalar(FunctionDecl *FD,
 
   unsigned Location = ActiveSemantic.Index.value_or(0);
 
+  if (!isPipelineBuiltin(getASTContext(), FD, A)) {
+    bool HasVkLocation = false;
+    if (auto *A = D->getAttr<HLSLVkLocationAttr>()) {
+      HasVkLocation = true;
+      Location = A->getLocation();
+    }
+
+    auto &UsesExplicitVkLocations =
+        IsInput ? InputUsesExplicitVkLocations : OutputUsesExplicitVkLocations;
+    if (UsesExplicitVkLocations.value_or(HasVkLocation) != HasVkLocation) {
+      Diag(D->getLocation(), diag::err_hlsl_semantic_partial_explicit_indexing);
+      return false;
+    }
+    UsesExplicitVkLocations = HasVkLocation;
+  }
+
   const ConstantArrayType *AT = dyn_cast<ConstantArrayType>(D->getType());
   unsigned ElementCount = AT ? AT->getZExtSize() : 1;
   ActiveSemantic.Index = Location + ElementCount;
@@ -1757,6 +1789,15 @@ void SemaHLSL::handleVkBindingAttr(Decl *D, const ParsedAttr &AL) {
                  HLSLVkBindingAttr(getASTContext(), AL, Binding, Set));
 }
 
+void SemaHLSL::handleVkLocationAttr(Decl *D, const ParsedAttr &AL) {
+  uint32_t Location;
+  if (!SemaRef.checkUInt32Argument(AL, AL.getArgAsExpr(0), Location))
+    return;
+
+  D->addAttr(::new (getASTContext())
+                 HLSLVkLocationAttr(getASTContext(), AL, Location));
+}
+
 bool SemaHLSL::diagnoseInputIDType(QualType T, const ParsedAttr &AL) {
   const auto *VT = T->getAs<VectorType>();
 

clang/test/CodeGenHLSL/semantics/semantic.explicit-location-output-struct.hlsl

@@ -0,0 +1,37 @@
+// RUN: %clang_cc1 -triple spirv-pc-vulkan1.3-pixel -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK,CHECK-SPIRV
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-pixel -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK,CHECK-DXIL
+
+// CHECK-SPIRV: @SV_Position = external hidden thread_local addrspace(7) externally_initialized constant <4 x float>, !spirv.Decorations ![[#MD_0:]]
+// CHECK-SPIRV: @SV_Target0 = external hidden thread_local addrspace(8) global <4 x float>, !spirv.Decorations ![[#MD_2:]]
+
+struct Output {
+  [[vk::location(2)]] float4 field : SV_Target;
+};
+
+// CHECK: define void @main() {{.*}} {
+Output main(float4 p : SV_Position) {
+  // CHECK:   %[[#OUT:]] = alloca %struct.Output, align 16
+
+  // CHECK-SPIRV:    %[[#IN:]] = load <4 x float>, ptr addrspace(7) @SV_Position, align 16
+  // CHECK-SPIRV:                call spir_func void @_Z4mainDv4_f(ptr %[[#OUT]], <4 x float> %[[#IN]])
+
+  // CHECK-DXIL:                 call void @_Z4mainDv4_f(ptr %[[#OUT]], <4 x float> %SV_Position0)
+
+  // CHECK:   %[[#TMP:]] = load %struct.Output, ptr %[[#OUT]], align 16
+  // CHECK: %[[#FIELD:]] = extractvalue %struct.Output %[[#TMP]], 0
+
+  // CHECK-SPIRV:                store <4 x float> %[[#FIELD]], ptr addrspace(8) @SV_Target0, align 16
+  // CHECK-DXIL:                 call void @llvm.dx.store.output.v4f32(i32 4, i32 0, i32 0, i8 0, i32 poison, <4 x float> %[[#FIELD]])
+  Output o;
+  o.field = p;
+  return o;
+}
+
+// CHECK-SPIRV-DAG: ![[#MD_0]] = !{![[#MD_1:]]}
+// CHECK-SPIRV-DAG: ![[#MD_1]] = !{i32 11, i32 15}
+//                                      |       `-> BuiltIn 'FragCoord'
+//                                      `-> SPIR-V decoration 'BuiltIn'
+// CHECK-SPIRV-DAG: ![[#MD_2]] = !{![[#MD_3:]]}
+// CHECK-SPIRV-DAG: ![[#MD_3]] = !{i32 30, i32 2}
+//                                      |       `-> Location index
+//                                      `-> SPIR-V decoration 'Location'

clang/test/CodeGenHLSL/semantics/semantic.explicit-location.hlsl

@@ -0,0 +1,19 @@
+// RUN: %clang_cc1 -triple spirv-pc-vulkan1.3-pixel -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefix=CHECK-SPIRV
+// RUN: %clang_cc1 -triple dxil-pc-shadermodel6.3-pixel -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefix=CHECK-DXIL
+
+// CHECK-SPIRV: @SV_Target0 = external hidden thread_local addrspace(8) global <4 x float>, !spirv.Decorations ![[#MD_2:]]
+
+// CHECK: define void @main() {{.*}} {
+[[vk::location(2)]] float4 main(float4 p : SV_Position) : SV_Target {
+  // CHECK-SPIRV: %[[#R:]] = call spir_func <4 x float> @_Z4mainDv4_f(<4 x float> %[[#]])
+  // CHECK-SPIRV:            store <4 x float> %[[#R]], ptr addrspace(8) @SV_Target0, align 16
+
+  // CHECK-DXIL:    %[[#TMP:]] = call <4 x float> @_Z4mainDv4_f(<4 x float> %SV_Position0)
+  // CHECK-DXIL:                 call void @llvm.dx.store.output.v4f32(i32 4, i32 0, i32 0, i8 0, i32 poison, <4 x float> %[[#TMP]])
+  return p;
+}
+
+// CHECK-SPIRV-DAG: ![[#MD_2]] = !{![[#MD_3:]]}
+// CHECK-SPIRV-DAG: ![[#MD_3]] = !{i32 30, i32 2}
+//                                      |       `-> Location index
+//                                      `-> SPIR-V decoration 'Location'

clang/test/CodeGenHLSL/semantics/semantic.explicit-mix-builtin.hlsl

@@ -0,0 +1,39 @@
+// RUN: %clang_cc1 -triple spirv-linux-vulkan-pixel -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK,CHECK-SPIRV
+
+// The following code is allowed because the `SV_Position` semantic is here
+// translated into a SPIR-V builtin. Meaning there is no implicit `Location`
+// assignment.
+
+struct S2 {
+  float4 a;
+  float4 b;
+};
+
+struct S1 {
+  float4 position : SV_Position;
+  [[vk::location(3)]] float4 color0 : COLOR0;
+};
+
+// CHECK-SPIRV: @SV_Position = external hidden thread_local addrspace(7) externally_initialized constant <4 x float>, !spirv.Decorations ![[#MD_0:]]
+// CHECK-SPIRV: @COLOR0 = external hidden thread_local addrspace(7) externally_initialized constant <4 x float>, !spirv.Decorations ![[#MD_2:]]
+// CHECK-SPIRV: @SV_Target0 = external hidden thread_local addrspace(8) global <4 x float>, !spirv.Decorations ![[#MD_4:]]
+
+[shader("pixel")]
+float4 main(S1 p) : SV_Target {
+  return p.position + p.color0;
+}
+// CHECK-SPIRV:    %[[#SV_POS:]] = load <4 x float>, ptr addrspace(7) @SV_Position, align 16
+// CHECK:            %[[#TMP1:]] = insertvalue %struct.S1 poison, <4 x float> %[[#SV_POS]], 0
+// CHECK-SPIRV:        %[[#A0:]] = load <4 x float>, ptr addrspace(7) @COLOR0, align 16
+// CHECK:            %[[#TMP2:]] = insertvalue %struct.S1 %[[#TMP1]], <4 x float> %[[#A0]], 1
+// CHECK:               %[[#P:]] = alloca %struct.S1, align 16
+// CHECK:                          store %struct.S1 %[[#TMP2]], ptr %[[#P]], align 16
+// CHECK-SPIRV:         %[[#R:]] = call spir_func <4 x float> @_Z4main2S1(ptr %[[#P]])
+// CHECK-SPIRV:                    store <4 x float> %[[#R]], ptr addrspace(8) @SV_Target0, align 16
+
+// CHECK-SPIRV: ![[#MD_0]] = !{![[#MD_1:]]}
+// CHECK-SPIRV: ![[#MD_1]] = !{i32 11, i32 15}
+// CHECK-SPIRV: ![[#MD_2]] = !{![[#MD_3:]]}
+// CHECK-SPIRV: ![[#MD_3]] = !{i32 30, i32 3}
+// CHECK-SPIRV: ![[#MD_4]] = !{![[#MD_5:]]}
+// CHECK-SPIRV: ![[#MD_5]] = !{i32 30, i32 0}

clang/test/Misc/pragma-attribute-supported-attributes-list.test

@@ -89,6 +89,7 @@
 // CHECK-NEXT: FunctionReturnThunks (SubjectMatchRule_function)
 // CHECK-NEXT: GNUInline (SubjectMatchRule_function)
 // CHECK-NEXT: HIPManaged (SubjectMatchRule_variable)
+// CHECK-NEXT: HLSLVkLocation (SubjectMatchRule_variable_is_parameter, SubjectMatchRule_field, SubjectMatchRule_function)
 // CHECK-NEXT: Hot (SubjectMatchRule_function)
 // CHECK-NEXT: HybridPatchable (SubjectMatchRule_function)
 // CHECK-NEXT: IBAction (SubjectMatchRule_objc_method_is_instance)

clang/test/SemaHLSL/Semantics/semantic.explicit-mix-builtin-vs.hlsl

@@ -0,0 +1,16 @@
+// RUN: %clang_cc1 -triple spirv-linux-vulkan-vertex -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s -verify -verify-ignore-unexpected=note
+
+// This is almost the same as semantic.explicit-mix-builtin.hlsl, except this
+// time we build a vertex shader. This means the SV_Position semantic is not
+// a BuiltIn anymore, but a Location decorated variable. This means we mix
+// implicit and explicit location assignment.
+struct S1 {
+  float4 position : SV_Position;
+  [[vk::location(3)]] float4 color : A;
+  // expected-error@-1 {{partial explicit stage input location assignment via vk::location(X) unsupported}}
+};
+
+[shader("vertex")]
+float4 main1(S1 p) : A {
+  return p.position + p.color;
+}

clang/test/SemaHLSL/Semantics/semantic.explicit-mix-location-2.hlsl

@@ -0,0 +1,15 @@
+// RUN: %clang --driver-mode=dxc %s -T ps_6_8 -E main1 -O3 -spirv -Xclang -verify -Xclang -verify-ignore-unexpected=note
+
+// The following code is not legal: both semantics A and B will be lowered
+// into a Location decoration. And mixing implicit and explicit Location
+// assignment is not supported.
+struct S1 {
+  [[vk::location(3)]] float4 color : B;
+  float4 position : A;
+  // expected-error@-1 {{partial explicit stage input location assignment via vk::location(X) unsupported}}
+};
+
+[shader("pixel")]
+float4 main1(S1 p) : SV_Target {
+  return p.position + p.color;
+}

clang/test/SemaHLSL/Semantics/semantic.explicit-mix-location.hlsl

@@ -0,0 +1,15 @@
+// RUN: %clang --driver-mode=dxc %s -T ps_6_8 -E main1 -O3 -spirv -Xclang -verify -Xclang -verify-ignore-unexpected=note
+
+// The following code is not legal: both semantics A and B will be lowered
+// into a Location decoration. And mixing implicit and explicit Location
+// assignment is not supported.
+struct S1 {
+  float4 position : A;
+  [[vk::location(3)]] float4 color : B;
+  // expected-error@-1 {{partial explicit stage input location assignment via vk::location(X) unsupported}}
+};
+
+[shader("pixel")]
+float4 main1(S1 p) : SV_Target {
+  return p.position + p.color;
+}