Remove globalopt dependency on the name of the entrypoint (#2350)

GlobalOpt has an explicit test for a function called "main" (eww...) to enable an optimization. This updates the test to match the entry point, independent of its name.

include/dxc/DXIL/DxilModule.h

@@ -79,6 +79,7 @@ public:
   llvm::Function *GetPatchConstantFunction();
   const llvm::Function *GetPatchConstantFunction() const;
   void SetPatchConstantFunction(llvm::Function *pFunc);
+  bool IsEntryOrPatchConstantFunction(const llvm::Function* pFunc) const;
 
   // Flags.
   unsigned GetGlobalFlags() const;

include/llvm/IR/Module.h

@@ -698,7 +698,7 @@ public:
   }
   bool HasHLModule() const { return TheHLModule != nullptr; }
   void SetHLModule(hlsl::HLModule *pValue) { TheHLModule = pValue; }
-  hlsl::HLModule &GetHLModule() { return *TheHLModule; }
+  hlsl::HLModule &GetHLModule() const { return *TheHLModule; }
   hlsl::HLModule &GetOrCreateHLModule(bool skipInit = false);
   ResetModuleCallback pfnResetHLModule = nullptr;
   void ResetHLModule() { if (pfnResetHLModule) (*pfnResetHLModule)(this); }

lib/DXIL/DxilModule.cpp

@@ -275,6 +275,10 @@ void DxilModule::SetPatchConstantFunction(llvm::Function *patchConstantFunc) {
   }
 }
 
+bool DxilModule::IsEntryOrPatchConstantFunction(const llvm::Function* pFunc) const {
+  return pFunc == GetEntryFunction() || pFunc == GetPatchConstantFunction();
+}
+
 unsigned DxilModule::GetGlobalFlags() const {
   unsigned Flags = m_ShaderFlags.GetGlobalFlags();
   return Flags;

lib/Transforms/IPO/GlobalOpt.cpp

@@ -42,6 +42,7 @@
 #include "llvm/Transforms/Utils/CtorUtils.h"
 #include "llvm/Transforms/Utils/GlobalStatus.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
+#include "dxc/DXIL/DxilModule.h" // HLSL Change - Entrypoint testing
 #include <algorithm>
 #include <deque>
 using namespace llvm;
@@ -1722,12 +1723,23 @@ bool GlobalOpt::ProcessGlobal(GlobalVariable *GV,
   return ProcessInternalGlobal(GV, GVI, GS);
 }
 
+// HLSL Change Begin
+static bool isEntryPoint(const llvm::Function* Func) {
+  const llvm::Module* Mod = Func->getParent();
+  return Mod->HasDxilModule()
+    ? Mod->GetDxilModule().IsEntryOrPatchConstantFunction(Func)
+    : Func->getName() == "main"; // Original logic for non-HLSL
+}
+// HLSL Change End
+
 /// ProcessInternalGlobal - Analyze the specified global variable and optimize
 /// it if possible.  If we make a change, return true.
 bool GlobalOpt::ProcessInternalGlobal(GlobalVariable *GV,
                                       Module::global_iterator &GVI,
                                       const GlobalStatus &GS) {
   auto &DL = GV->getParent()->getDataLayout();
+
+
   // If this is a first class global and has only one accessing function
   // and this function is main (which we know is not recursive), we replace
   // the global with a local alloca in this function.
@@ -1739,7 +1751,7 @@ bool GlobalOpt::ProcessInternalGlobal(GlobalVariable *GV,
   if (!GS.HasMultipleAccessingFunctions &&
       GS.AccessingFunction && !GS.HasNonInstructionUser &&
       GV->getType()->getElementType()->isSingleValueType() &&
-      GS.AccessingFunction->getName() == "main" &&
+      isEntryPoint(GS.AccessingFunction) && // HLSL Change - Generalize entrypoint testing
       GS.AccessingFunction->hasExternalLinkage() &&
       GV->getType()->getAddressSpace() == 0) {
     DEBUG(dbgs() << "LOCALIZING GLOBAL: " << *GV);

tools/clang/test/CodeGenHLSL/batch/passes/globalopt/entrypoint_name_invariant_main.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -T vs_6_0 -E main %s | FileCheck %s
+
+// Regression test for #2318, where the globalopt had behavior
+// conditional on the name of the main function.
+
+static uint g[1] = { 0 };
+uint main() : OUT
+{
+  // CHECK-NOT: load i32
+  // CHECK-NOT: store i32
+  g[0]++;
+  // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 1)
+  return g[0];
+}

tools/clang/test/CodeGenHLSL/batch/passes/globalopt/entrypoint_name_invariant_main2.hlsl

@@ -0,0 +1,14 @@
+// RUN: %dxc -T vs_6_0 -E main2 %s | FileCheck %s
+
+// Regression test for #2318, where the globalopt had behavior
+// conditional on the name of the main function.
+
+static uint g[1] = { 0 };
+uint main2() : OUT
+{
+  // CHECK-NOT: load i32
+  // CHECK-NOT: store i32
+  g[0]++;
+  // CHECK: call void @dx.op.storeOutput.i32(i32 5, i32 0, i32 0, i8 0, i32 1)
+  return g[0];
+}