[OpenMP] Only register HeapToShared in the OpenMPOpt module pass

[jhuber6]

Summary

This patch moves registering HeapToShared to only occur in the module pass.
This does nothing if run outside of the module pass because it requires intact
call sites and all the analysis we can do has already been done.

[jdoerfert]

We should instead keep noinline on them till the CGSCC pass, at least for this one as it could trigger late even if it wasn't possible early.

[jhuber6]

In D106687#2902826, @jdoerfert wrote:

We should instead keep noinline on them till the CGSCC pass, at least for this one as it could trigger late even if it wasn't possible early.

But we run the Module pass and CGSCC passes at different optimization levels right now. It's possible that H2S could remove something after more optimizations, but I haven't seen any examples of something failing to be moved to the stack because of optimizations.

[jdoerfert]

In D106687#2902833, @jhuber6 wrote:

In D106687#2902826, @jdoerfert wrote:

We should instead keep noinline on them till the CGSCC pass, at least for this one as it could trigger late even if it wasn't possible early.

But we run the Module pass and CGSCC passes at different optimization levels right now. It's possible that H2S could remove something after more optimizations, but I haven't seen any examples of something failing to be moved to the stack because of optimizations.

not inlining alloc_shared/free_shared is probably not a real concern, it's slow either way.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

IsSPMDRFI.foreachUse(SCC, [&](Use &U, Function &) {
       if (!CI)
         return false;
       A.getOrCreateAAFor<AAFoldRuntimeCall>(
           IRPosition::callsite_returned(*CI), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,
           /* UpdateAfterInit */ false);
       return false;
     });
+
+    auto &GlobalizationRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
+    auto CreateAA = [&](Use &U, Function &F) {
+      A.getOrCreateAAFor<AAHeapToShared>(IRPosition::function(F));
+      return false;
+    };
+    GlobalizationRFI.foreachUse(SCC, CreateAA);
   }
 
   // Create CallSite AA for all Getters.
   for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
     auto ICVInfo = OMPInfoCache.ICVs[static_cast<InternalControlVar>(Idx)];
 
     auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
 
     auto CreateAA = [&](Use &U, Function &Caller) {
       CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &GetterRFI);
       if (!CI)
         return false;
 
       auto &CB = cast<CallBase>(*CI);
 
       IRPosition CBPos = IRPosition::callsite_function(CB);
       A.getOrCreateAAFor<AAICVTracker>(CBPos);
       return false;
     };
 
     GetterRFI.foreachUse(SCC, CreateAA);
   }
-  auto &GlobalizationRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
-  auto CreateAA = [&](Use &U, Function &F) {
-    A.getOrCreateAAFor<AAHeapToShared>(IRPosition::function(F));
-    return false;
-  };
-  GlobalizationRFI.foreachUse(SCC, CreateAA);
 
   // Create an ExecutionDomain AA for every function and a HeapToStack AA for
   // every function if there is a device kernel.
   for (auto *F : SCC) {
     if (!F->isDeclaration())
       A.getOrCreateAAFor<AAExecutionDomain>(IRPosition::function(*F));
     if (isOpenMPDevice(M))
       A.getOrCreateAAFor<AAHeapToStack>(IRPosition::function(*F));