Use IsEntry when LowerStaticGlobalIntoAlloca. (#3353)

include/dxc/DXIL/DxilModule.h

@@ -153,6 +153,9 @@ public:
   // Is an entry function that uses input/output signature conventions?
   // Includes: vs/hs/ds/gs/ps/cs as well as the patch constant function.
   bool IsEntryThatUsesSignatures(const llvm::Function *F) const ;
+  // Is F an entry?
+  // Includes: IsEntryThatUsesSignatures and all ray tracing shaders.
+  bool IsEntry(const llvm::Function *F) const;
 
   // Remove Root Signature from module metadata, return true if changed
   bool StripRootSignatureFromMetadata();

include/dxc/HLSL/HLModule.h

@@ -153,6 +153,9 @@ public:
   // Is an entry function that uses input/output signature conventions?
   // Includes: vs/hs/ds/gs/ps/cs as well as the patch constant function.
   bool IsEntryThatUsesSignatures(llvm::Function *F);
+  // Is F an entry?
+  // Includes: IsEntryThatUsesSignatures and all ray tracing shaders.
+  bool IsEntry(llvm::Function *F);
 
   DxilFunctionAnnotation *GetFunctionAnnotation(llvm::Function *F);
   DxilFunctionAnnotation *AddFunctionAnnotation(llvm::Function *F);

lib/DXIL/DxilModule.cpp

@@ -1253,6 +1253,16 @@ bool DxilModule::IsEntryThatUsesSignatures(const llvm::Function *F) const {
   // Otherwise, return true if patch constant function
   return IsPatchConstantShader(F);
 }
+bool DxilModule::IsEntry(const llvm::Function *F) const {
+  auto propIter = m_DxilEntryPropsMap.find(F);
+  if (propIter != m_DxilEntryPropsMap.end()) {
+    DXASSERT(propIter->second->props.shaderKind != DXIL::ShaderKind::Invalid,
+             "invalid entry props");
+    return true;
+  }
+  // Otherwise, return true if patch constant function
+  return IsPatchConstantShader(F);
+}
 
 bool DxilModule::StripRootSignatureFromMetadata() {
   NamedMDNode *pRootSignatureNamedMD = GetModule()->getNamedMetadata(DxilMDHelper::kDxilRootSignatureMDName);

lib/HLSL/HLModule.cpp

@@ -408,6 +408,16 @@ bool HLModule::IsEntryThatUsesSignatures(llvm::Function *F) {
   // Otherwise, return true if patch constant function
   return IsPatchConstantShader(F);
 }
+bool HLModule::IsEntry(llvm::Function *F) {
+  auto propIter = m_DxilFunctionPropsMap.find(F);
+  if (propIter != m_DxilFunctionPropsMap.end()) {
+    DXASSERT(propIter->second->shaderKind != DXIL::ShaderKind::Invalid,
+             "invalid entry props");
+    return true;
+  }
+  // Otherwise, return true if patch constant function
+  return IsPatchConstantShader(F);
+}
 
 DxilFunctionAnnotation *HLModule::GetFunctionAnnotation(llvm::Function *F) {
   return m_pTypeSystem->GetFunctionAnnotation(F);

lib/Transforms/Scalar/ScalarReplAggregatesHLSL.cpp

@@ -5719,8 +5719,9 @@ public:
         }
       } else {
         for (Function &F : M) {
-          if (!HLM.IsEntryThatUsesSignatures(&F))
+          if (!HLM.IsEntry(&F)) {
             continue;
+          }
           entryAndInitFunctionSet.insert(&F);
         }
       }
@@ -5736,7 +5737,7 @@ public:
         }
       } else {
         for (Function &F : M) {
-          if (!DM.IsEntryThatUsesSignatures(&F))
+          if (!DM.IsEntry(&F))
             continue;
           entryAndInitFunctionSet.insert(&F);
         }

tools/clang/test/HLSLFileCheck/hlsl/intrinsics/createHandleFromHeap/dynamic_res_global_for_lib2.hlsl

@@ -0,0 +1,44 @@
+// RUN: %dxc -T lib_6_6 %s | %FileCheck %s
+
+// Make sure each entry get 2 createHandleFromHeap.
+// CHECK:define void
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+// CHECK:define void
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+// CHECK:define void
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+// CHECK:call %dx.types.Handle @dx.op.createHandleFromHeap(i32 218, i32 %{{.*}}, i1 false, i1 false)
+
+uint ID;
+static const RWBuffer<float>           g_result      = ResourceDescriptorHeap[ID];
+static ByteAddressBuffer         g_rawBuf      = ResourceDescriptorHeap[ID+1];
+static float x = ID + 3;
+
+// TODO: support array.
+//  static Buffer<float> g_bufs[2] = {ResourceDescriptorHeap[ID+2], ResourceDescriptorHeap[ID+3]};
+
+[NumThreads(1, 1, 1)]
+[RootSignature("RootFlags(CBV_SRV_UAV_HEAP_DIRECTLY_INDEXED | SAMPLER_HEAP_DIRECTLY_INDEXED), RootConstants(num32BitConstants=1, b0))")]
+void csmain(uint ix : SV_GroupIndex)
+{
+  g_result[ix] = g_rawBuf.Load<float>(ix);// + g_bufs[0].Load(ix);
+}
+// export foo to make sure init function not removed.
+export float foo(uint i) {
+  return x + i;
+}
+
+[NumThreads(1, 1, 1)]
+[RootSignature("RootFlags(CBV_SRV_UAV_HEAP_DIRECTLY_INDEXED | SAMPLER_HEAP_DIRECTLY_INDEXED), RootConstants(num32BitConstants=1, b0)")]
+void csmain2(uint ix : SV_GroupIndex)
+{
+  g_result[ix] = g_rawBuf.Load<float>(ix+ID);
+}
+
+[shader("raygeneration")]
+void RayGeneration()
+{
+  g_result[0] = g_rawBuf.Load<float>(1);
+}