fix shader kind in CG

Author: Keenuts

clang/lib/CodeGen/CGHLSLRuntime.cpp

@@ -699,8 +699,9 @@ void CGHLSLRuntime::emitUserSemanticStore(IRBuilder<> &B, llvm::Value *Source,
 }
 
 llvm::Value *CGHLSLRuntime::emitSystemSemanticLoad(
-    IRBuilder<> &B, llvm::Type *Type, const clang::DeclaratorDecl *Decl,
-    HLSLAppliedSemanticAttr *Semantic, std::optional<unsigned> Index) {
+    IRBuilder<> &B, const FunctionDecl *FD, llvm::Type *Type,
+    const clang::DeclaratorDecl *Decl, HLSLAppliedSemanticAttr *Semantic,
+    std::optional<unsigned> Index) {
 
   std::string SemanticName = Semantic->getAttrName()->getName().upper();
   if (SemanticName == "SV_GROUPINDEX") {
@@ -736,8 +737,12 @@ llvm::Value *CGHLSLRuntime::emitSystemSemanticLoad(
     return buildVectorInput(B, GroupIDIntrinsic, Type);
   }
 
+  const auto *ShaderAttr = FD->getAttr<HLSLShaderAttr>();
+  assert(ShaderAttr && "Entry point has no shader attribute");
+  llvm::Triple::EnvironmentType ST = ShaderAttr->getType();
+
   if (SemanticName == "SV_POSITION") {
-    if (CGM.getTriple().getEnvironment() == Triple::EnvironmentType::Pixel) {
+    if (ST == Triple::EnvironmentType::Pixel) {
       if (CGM.getTarget().getTriple().isSPIRV())
         return createSPIRVBuiltinLoad(B, CGM.getModule(), Type,
                                       Semantic->getAttrName()->getName(),
@@ -746,7 +751,7 @@ llvm::Value *CGHLSLRuntime::emitSystemSemanticLoad(
         return emitDXILUserSemanticLoad(B, Type, Semantic, Index);
     }
 
-    if (CGM.getTriple().getEnvironment() == Triple::EnvironmentType::Vertex) {
+    if (ST == Triple::EnvironmentType::Vertex) {
       return emitUserSemanticLoad(B, Type, Decl, Semantic, Index);
     }
   }
@@ -804,7 +809,7 @@ llvm::Value *CGHLSLRuntime::handleScalarSemanticLoad(
 
   std::optional<unsigned> Index = Semantic->getSemanticIndex();
   if (Semantic->getAttrName()->getName().starts_with_insensitive("SV_"))
-    return emitSystemSemanticLoad(B, Type, Decl, Semantic, Index);
+    return emitSystemSemanticLoad(B, FD, Type, Decl, Semantic, Index);
   return emitUserSemanticLoad(B, Type, Decl, Semantic, Index);
 }
 

clang/lib/CodeGen/CGHLSLRuntime.h

@@ -179,7 +179,8 @@ class CGHLSLRuntime {
 protected:
   CodeGenModule &CGM;
 
-  llvm::Value *emitSystemSemanticLoad(llvm::IRBuilder<> &B, llvm::Type *Type,
+  llvm::Value *emitSystemSemanticLoad(llvm::IRBuilder<> &B,
+                                      const FunctionDecl *FD, llvm::Type *Type,
                                       const clang::DeclaratorDecl *Decl,
                                       HLSLAppliedSemanticAttr *Semantic,
                                       std::optional<unsigned> Index);

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

@@ -0,0 +1,40 @@
+// RUN: %clang_cc1 -triple spirv-linux-vulkan-library -x hlsl -emit-llvm -finclude-default-header -disable-llvm-passes -o - %s | FileCheck %s --check-prefixes=CHECK
+
+// The followiong file contains both implicit and explicit vk::location, but
+// because each entrypoint has only one kind, this is allowed.
+
+[shader("vertex")]
+float4 vs_main(float4 p : SV_Position) : A {
+  return p;
+}
+
+[shader("pixel")]
+float4 ps_main([[vk::location(0)]] float4 p : A) : SV_Target {
+  return p;
+}
+
+// The following function is not marked as being a shader entrypoint, this
+// means the semantics and [[vk::location]] attributes are ignored.
+// Otherwise, the partial explicit location assignment would be illegal.
+float4 not_an_entry([[vk::location(0)]] float4 a : A, float4 b : B) : C {
+  return a + b;
+}
+
+// CHECK: @SV_Position0 = external hidden thread_local addrspace(7) externally_initialized constant <4 x float>, !spirv.Decorations ![[#MD_0:]]
+// CHECK: @A0 = external hidden thread_local addrspace(8) global <4 x float>, !spirv.Decorations ![[#MD_0:]]
+// CHECK: @A0.1 = external hidden thread_local addrspace(7) externally_initialized constant <4 x float>, !spirv.Decorations ![[#MD_0:]]
+// CHECK: @SV_Target0 = external hidden thread_local addrspace(8) global <4 x float>, !spirv.Decorations ![[#MD_2:]]
+
+
+// CHECK: define void @vs_main()
+// CHECK: %[[#]] = load <4 x float>, ptr addrspace(7) @SV_Position0, align 16
+// CHECK: store <4 x float> %[[#]], ptr addrspace(8) @A0, align 16
+
+// CHECK: define void @ps_main()
+// CHECK: %[[#]] = load <4 x float>, ptr addrspace(7) @A0.1, align 16
+// CHECK: store <4 x float> %[[#]], ptr addrspace(8) @SV_Target0, align 16
+
+// CHECK: ![[#MD_0]] = !{![[#MD_1:]]}
+// CHECK: ![[#MD_1]] = !{i32 30, i32 0}
+// CHECK: ![[#MD_2]] = !{![[#MD_3:]]}
+// CHECK: ![[#MD_3]] = !{i32 30, i32 1}