Only save root signature for entry function. (#86)

tools/clang/lib/CodeGen/CGHLSLMS.cpp

@@ -111,6 +111,8 @@ private:
   // List for functions with clip plane.
   std::vector<Function *> clipPlaneFuncList;
   std::unordered_map<Value *, DebugLoc> debugInfoMap;
+
+  DxilRootSignatureVersion  rootSigVer;
   
   Value *EmitHLSLMatrixLoad(CGBuilderTy &Builder, Value *Ptr, QualType Ty);
   void EmitHLSLMatrixStore(CGBuilderTy &Builder, Value *Val, Value *DestPtr,
@@ -309,6 +311,19 @@ CGMSHLSLRuntime::CGMSHLSLRuntime(CodeGenModule &CGM)
   // set entry name
   m_pHLModule->SetEntryFunctionName(CGM.getCodeGenOpts().HLSLEntryFunction);
 
+  // set root signature version.
+  if (CGM.getLangOpts().RootSigMinor == 0) {
+    rootSigVer = hlsl::DxilRootSignatureVersion::Version_1_0;
+  }
+  else {
+    DXASSERT(CGM.getLangOpts().RootSigMinor == 1,
+      "else CGMSHLSLRuntime Constructor needs to be updated");
+    rootSigVer = hlsl::DxilRootSignatureVersion::Version_1_1;
+  }
+
+  DXASSERT(CGM.getLangOpts().RootSigMajor == 1,
+           "else CGMSHLSLRuntime Constructor needs to be updated");
+
   // add globalCB
   unique_ptr<HLCBuffer> CB = std::make_unique<HLCBuffer>();
   std::string globalCBName = "$Globals";
@@ -5185,6 +5200,10 @@ void CGMSHLSLRuntime::EmitHLSLFlatConversionToAggregate(CodeGenFunction &CGF,
 void CGMSHLSLRuntime::EmitHLSLRootSignature(CodeGenFunction &CGF,
                                             HLSLRootSignatureAttr *RSA,
                                             Function *Fn) {
+  // Only parse root signature for entry function.
+  if (Fn != EntryFunc)
+    return;
+
   StringRef StrRef = RSA->getSignatureName();
   DiagnosticsEngine &Diags = CGF.getContext().getDiagnostics();
   SourceLocation SLoc = RSA->getLocation();
@@ -5192,19 +5211,7 @@ void CGMSHLSLRuntime::EmitHLSLRootSignature(CodeGenFunction &CGF,
   raw_string_ostream OS(OSStr);
   hlsl::DxilVersionedRootSignatureDesc *D = nullptr;
 
-  DXASSERT(CGF.getLangOpts().RootSigMajor == 1,
-           "else EmitHLSLRootSignature needs to be updated");
-  hlsl::DxilRootSignatureVersion Ver;
-  if (CGF.getLangOpts().RootSigMinor == 0) {
-    Ver = hlsl::DxilRootSignatureVersion::Version_1_0;
-  }
-  else {
-    DXASSERT(CGF.getLangOpts().RootSigMinor == 1,
-      "else EmitHLSLRootSignature needs to be updated");
-    Ver = hlsl::DxilRootSignatureVersion::Version_1_1;
-  }
-
-  if (ParseHLSLRootSignature(StrRef.data(), StrRef.size(), Ver, &D, SLoc,
+  if (ParseHLSLRootSignature(StrRef.data(), StrRef.size(), rootSigVer, &D, SLoc,
                              Diags)) {
     CComPtr<IDxcBlob> pSignature;
     CComPtr<IDxcBlobEncoding> pErrors;
@@ -5214,10 +5221,8 @@ void CGMSHLSLRuntime::EmitHLSLRootSignature(CodeGenFunction &CGF,
       ReportHLSLRootSigError(Diags, SLoc,
         (char *)pErrors->GetBufferPointer(), pErrors->GetBufferSize());
       hlsl::DeleteRootSignature(D);
-    }
-    else {
-      llvm::Module *pModule = Fn->getParent();
-      pModule->GetHLModule().GetRootSignature().Assign(D, pSignature);
+    } else {
+      m_pHLModule->GetRootSignature().Assign(D, pSignature);
     }
   }
 }

tools/clang/test/CodeGenHLSL/rootSigEntry.hlsl

@@ -0,0 +1,13 @@
+// RUN: %dxc -E main -T ps_6_0 %s
+
+float a;
+float4 main() : SV_Target {
+   return a;
+}
+
+
+[RootSignature("SRV(t0)")]
+float4 main2() : SV_Target {
+   return a;
+}
+

tools/clang/unittests/HLSL/CompilerTest.cpp

@@ -577,6 +577,7 @@ public:
   TEST_METHOD(CodeGenResourceInCBV2)
   TEST_METHOD(CodeGenResourceInTB2)
   TEST_METHOD(CodeGenResourceInTBV2)
+  TEST_METHOD(CodeGenRootSigEntry)
   TEST_METHOD(CodeGenCBufferStructArray)
   TEST_METHOD(PreprocessWhenValidThenOK)
   TEST_METHOD(WhenSigMismatchPCFunctionThenFail)
@@ -2951,6 +2952,10 @@ TEST_F(CompilerTest, CodeGenResourceInTBV2) {
   CodeGenTestCheck(L"..\\CodeGenHLSL\\resource-in-tbv2.hlsl");
 }
 
+TEST_F(CompilerTest, CodeGenRootSigEntry) {
+  CodeGenTest(L"..\\CodeGenHLSL\\rootSigEntry.hlsl");
+}
+
 TEST_F(CompilerTest, CodeGenCBufferStructArray) {
   CodeGenTestCheck(L"..\\CodeGenHLSL\\cbuffer-structarray.hlsl");
 }