[OpenMP] Replace GPU globalization calls with shared memory in the middle-end

[jhuber6]

Summary

The changes introduced in D97680 create a simpler interface to code that needs to be globalized. This interface is used to simplify the globalization calls in the middle end. We can check any globalization call that is only called by a single thread in the team and replace it with a static shared memory buffer.

Depends on D102700 D97680

[jdoerfert]

some initial comments.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1015

1018

No cast needed

1023

prefer early exit, if (!...) return false;

1026

No need for it to be a constant

1030

Don't assume an order. Check all users, one should be a free, others can be whatever. If you find bitcast users, remember the type, if they all agree, use that for the alloca.

1622

no need to go over the free calls. they need to be users of the alloc and we remove them with the alloc.

1629

No need, use *CI below.

[jdoerfert]

We can reuse Attributor HeapToStack for this :)

[jhuber6]

Changing this optimization to replace the globalization calls with shared memory. Removing them will be done by the attributor using HeapToStack once we add the allocation calls and improve the attributor.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

68

I don't think we need this after all.

1010

For now, check isKernel(F) and only do this for kernel functions. Later we can be more aggressive but for now that should limit it properly, also with regards to the lifetime of those allocations.

1068

You cannot recollect while looping.
Move the recollect for free after the loop, delete the other one. Return true from this function when it deleted the use, that will update the alloc use list.

1072

return true;

[jhuber6]

Changing the test to simply check if we are in a non-SPMD kernel function. A more advanced approach can be used in the future.

[jdoerfert]

LGTM, one nit.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1055

Can you use a enum here or variable instead of "3".

[jhuber6]

Making changes.

[jhuber6]

Adding enum for address space constant.

[jhuber6]

Changing the method for determining valid allocation calls to use AAExecutionDomain introduced in D101578. Currently, this doesn't work with the code Clang generates but will soon.

[jhuber6]

Adding a pattern match check so this picks up on the code Clang currently generates. This is a temporary solution to improve performance until a full runtime rewrite can be landed.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1018

Nit: hoist this out of the lambda.

1043

You have to check that a unique free call was found or give up otherwise.

1060

You might be able to replace these two with just: ConstantExpr::getPointerCast

2367

Check the Pred as well. Maybe try to split this in multiple matches or somehow to make it easier to read.

[jhuber6]

Addressing comments.

[jdoerfert]

I looked over this again because I think we will later move the logic into an AbstractAttribute. The reason is that we want to reuse some of the analysis in AAKernelInfo to determine if we can use SPMD mode. Basically, if the shared memory is only used in a nice way and fees into an actual parallel51, it is compatible with SPMD mode. If we go to SPMD mode we have (or better should) eliminate the shared memory indirection. Maybe we move all of this into AAHeapToStack and make a more general AAAllocationInfo out of it. Let's see.

That said, I noticed a problem below. Please add a test that misses the free, and/or that has a phi before the free, to make sure we do not globalize.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1047

if (!FC) return false;

Also move the debug below the exists.

[jhuber6]

Change the implementation to use the Attributor. This is because the attributor after registering HeapToStack in D102197 would change the call sites. Putting this pass inside an attributor instance allows both to run and now globalization will only be replaced if HeapToStack doesn't remove it first.

[jhuber6]

Changing alignment of new global variable, alignment of 8 caused errors in some applications.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 
+using namespace llvm::PatternMatch;
 using namespace llvm;
 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
     cl::desc("Disable OpenMP specific optimizations."), cl::Hidden,
static cl::opt<bool> HideMemoryTransferLatency(
     "openmp-hide-memory-transfer-latency",
     cl::desc("[WIP] Tries to hide the latency of host to device memory"
              " transfers"),
     cl::Hidden, cl::init(false));
 
 STATISTIC(NumOpenMPRuntimeCallsDeduplicated,
           "Number of OpenMP runtime calls deduplicated");
 STATISTIC(NumOpenMPParallelRegionsDeleted,
           "Number of OpenMP parallel regions deleted");
 STATISTIC(NumOpenMPRuntimeFunctionsIdentified,
           "Number of OpenMP runtime functions identified");
 STATISTIC(NumOpenMPRuntimeFunctionUsesIdentified,
           "Number of OpenMP runtime function uses identified");
 STATISTIC(NumOpenMPTargetRegionKernels,
           "Number of OpenMP target region entry points (=kernels) identified");
 STATISTIC(
     NumOpenMPParallelRegionsReplacedInGPUStateMachine,
     "Number of OpenMP parallel regions replaced with ID in GPU state machines");
 STATISTIC(NumOpenMPParallelRegionsMerged,
           "Number of OpenMP parallel regions merged");
+STATISTIC(NumBytesMovedToSharedMemory,
+          "Amount of memory pushed to shared memory");
 
 #if !defined(NDEBUG)
 static constexpr auto TAG = "[" DEBUG_TYPE "]";
 #endif
 
 namespace {
 
+enum class AddressSpace : unsigned {
+  Generic = 0,
+  Global = 1,
+  Shared = 3,
+  Constant = 4,
+  Local = 5,
+};
+
+struct AAHeapToShared;
+
 struct AAICVTracker;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
 struct OMPInformationCache : public InformationCache {
   OMPInformationCache(Module &M, AnalysisGetter &AG,
                       BumpPtrAllocator &Allocator, SetVector<Function *> &CGSCC,
                       SmallPtrSetImpl<Kernel> &Kernels)
bool run(bool IsModulePass) {
 
     LLVM_DEBUG(dbgs() << TAG << "Run on SCC with " << SCC.size()
                       << " functions in a slice with "
                       << OMPInfoCache.ModuleSlice.size() << " functions\n");
 
     if (IsModulePass) {
       Changed |= runAttributor();
 
+      // Recollect uses, in case Attributor deleted any.
+      OMPInfoCache.recollectUses();
+
       if (remarksEnabled())
         analysisGlobalization();
     } else {
       if (PrintICVValues)
         printICVs();
       if (PrintOpenMPKernels)
         printKernels();
 
if (Changed) {
       /// any emitted master/end_master calls.
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_fork_call);
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_barrier);
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_master);
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_end_master);
     }
 
     return Changed;
   }
 
   /// Try to delete parallel regions if possible.
   bool deleteParallelRegions() {
     const unsigned CallbackCalleeOperand = 2;
 
     OMPInformationCache::RuntimeFunctionInfo &RFI =
         OMPInfoCache.RFIs[OMPRTL___kmpc_fork_call];
 
     if (!RFI.Declaration)
       return false;
 
     bool Changed = false;
     auto DeleteCallCB = [&](Use &U, Function &) {
       CallInst *CI = getCallIfRegularCall(U);
       if (!CI)
         return false;
       auto *Fn = dyn_cast<Function>(
           CI->getArgOperand(CallbackCalleeOperand)->stripPointerCasts());
       if (!Fn)
         return false;
       if (!Fn->onlyReadsMemory())
         return false;
       if (!Fn->hasFnAttribute(Attribute::WillReturn))
         return false;
 
       LLVM_DEBUG(dbgs() << TAG << "Delete read-only parallel region in "
                         << CI->getCaller()->getName() << "\n");
 
       auto Remark = [&](OptimizationRemark OR) {
         return OR << "Parallel region in "
                   << ore::NV("OpenMPParallelDelete", CI->getCaller()->getName())
                   << " deleted";
       };
       emitRemark<OptimizationRemark>(CI, "OpenMPParallelRegionDeletion",
                                      Remark);
 
       CGUpdater.removeCallSite(*CI);
       CI->eraseFromParent();
       Changed = true;
       ++NumOpenMPParallelRegionsDeleted;
       return true;
     };
 
     RFI.foreachUse(SCC, DeleteCallCB);
 
     return Changed;
   }
 
   /// Try to eliminate runtime calls by reusing existing ones.
   bool deduplicateRuntimeCalls() {
     bool Changed = false;
 
     RuntimeFunction DeduplicableRuntimeCallIDs[] = {
         OMPRTL_omp_get_num_threads,
         OMPRTL_omp_in_parallel,
         OMPRTL_omp_get_cancellation,
         OMPRTL_omp_get_thread_limit,
         OMPRTL_omp_get_supported_active_levels,
         OMPRTL_omp_get_level,
         OMPRTL_omp_get_ancestor_thread_num,
         OMPRTL_omp_get_team_size,
         OMPRTL_omp_get_active_level,
         OMPRTL_omp_in_final,
         OMPRTL_omp_get_proc_bind,
         OMPRTL_omp_get_num_places,
         OMPRTL_omp_get_num_procs,
         OMPRTL_omp_get_place_num,
         OMPRTL_omp_get_partition_num_places,
         OMPRTL_omp_get_partition_place_nums};
 
auto SplitMemTransfers = [&](Use &U, Function &Decl) {
       return WasSplit;
     };
     RFI.foreachUse(SCC, SplitMemTransfers);
 
     return Changed;
   }
 
   void analysisGlobalization() {
-    RuntimeFunction GlobalizationRuntimeIDs[] = {OMPRTL___kmpc_alloc_shared,
-                                                 OMPRTL___kmpc_free_shared};
-
-    for (const auto GlobalizationCallID : GlobalizationRuntimeIDs) {
-      auto &RFI = OMPInfoCache.RFIs[GlobalizationCallID];
+    auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
 
-      auto CheckGlobalization = [&](Use &U, Function &Decl) {
-        if (CallInst *CI = getCallIfRegularCall(U, &RFI)) {
-          auto Remark = [&](OptimizationRemarkAnalysis ORA) {
-            return ORA
-                   << "Found thread data sharing on the GPU. "
-                   << "Expect degraded performance due to data globalization.";
-          };
-          emitRemark<OptimizationRemarkAnalysis>(CI, "OpenMPGlobalization",
-                                                 Remark);
-        }
 
-        return false;
-      };
 
-      RFI.foreachUse(SCC, CheckGlobalization);
-    }
-  }
 
   /// Maps the values stored in the offload arrays passed as arguments to
   /// \p RuntimeCall into the offload arrays in \p OAs.
   bool getValuesInOffloadArrays(CallInst &RuntimeCall,
                                 MutableArrayRef<OffloadArray> OAs) {
     assert(OAs.size() == 3 && "Need space for three offload arrays!");
 
     // A runtime call that involves memory offloading looks something like:
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
 
         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);
 
     for (auto &F : M) {
       if (!F.isDeclaration())
         A.getOrCreateAAFor<AAExecutionDomain>(IRPosition::function(F));
     }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
if (!A.checkForAllCallSites(PredForCallSite, *this,
                               /* RequiresAllCallSites */ true,
                               AllCallSitesKnown))
     SingleThreadedBBs.erase(&F->getEntryBlock());
 
   // Check if the edge into the successor block compares a thread-id function to
   // a constant zero.
   // TODO: Use AAValueSimplify to simplify and propogate constants.
   // TODO: Check more than a single use for thread ID's.
-  auto IsSingleThreadOnly = [&](BranchInst *Edge, BasicBlock *SuccessorBB) {
+  auto IsInitialThreadOnly = [&](BranchInst *Edge, BasicBlock *SuccessorBB) {
     if (!Edge || !Edge->isConditional())
       return false;
     if (Edge->getSuccessor(0) != SuccessorBB)
       return false;
 
     auto *Cmp = dyn_cast<CmpInst>(Edge->getCondition());
     if (!Cmp || !Cmp->isTrueWhenEqual() || !Cmp->isEquality())
       return false;
 
+    // Temporarily match the pattern generated by clang for teams regions.
+    // TODO: Remove this once the new runtime is in place.
+    ConstantInt *One, *NegOne;
+    CmpInst::Predicate Pred;
+    auto &&m_ThreadID = m_Intrinsic<Intrinsic::nvvm_read_ptx_sreg_tid_x>();
+    auto &&m_WarpSize = m_Intrinsic<Intrinsic::nvvm_read_ptx_sreg_warpsize>();
+    auto &&m_BlockSize = m_Intrinsic<Intrinsic::nvvm_read_ptx_sreg_ntid_x>();
+    if (match(Cmp, m_Cmp(Pred, m_ThreadID,
+                         m_And(m_Sub(m_BlockSize, m_ConstantInt(One)),
+                               m_Xor(m_Sub(m_WarpSize, m_ConstantInt(One)),
+                                     m_ConstantInt(NegOne))))))
+      if (One->isOne() && NegOne->isMinusOne() &&
+          Pred == CmpInst::Predicate::ICMP_EQ)
+        return true;
+
     ConstantInt *C = dyn_cast<ConstantInt>(Cmp->getOperand(1));
     if (!C || !C->isZero())
       return false;
 
     if (auto *II = dyn_cast<IntrinsicInst>(Cmp->getOperand(0)))
       if (II->getIntrinsicID() == Intrinsic::nvvm_read_ptx_sreg_tid_x)
         return true;
     if (auto *II = dyn_cast<IntrinsicInst>(Cmp->getOperand(0)))
       if (II->getIntrinsicID() == Intrinsic::amdgcn_workitem_id_x)
         return true;
 
     return false;
   };
 
   // Merge all the predecessor states into the current basic block. A basic
   // block is executed by a single thread if all of its predecessors are.
   auto MergePredecessorStates = [&](BasicBlock *BB) {
     if (pred_begin(BB) == pred_end(BB))
       return SingleThreadedBBs.contains(BB);
 
-    bool IsSingleThreaded = true;
+    bool IsInitialThread = true;
     for (auto PredBB = pred_begin(BB), PredEndBB = pred_end(BB);
          PredBB != PredEndBB; ++PredBB) {
-      if (!IsSingleThreadOnly(dyn_cast<BranchInst>((*PredBB)->getTerminator()),
+      if (!IsInitialThreadOnly(dyn_cast<BranchInst>((*PredBB)->getTerminator()),
                               BB))
-        IsSingleThreaded &= SingleThreadedBBs.contains(*PredBB);
+        IsInitialThread &= SingleThreadedBBs.contains(*PredBB);
     }
 
-    return IsSingleThreaded;
+    return IsInitialThread;
   };
 
   for (auto *BB : RPOT) {
     if (!MergePredecessorStates(BB))
       SingleThreadedBBs.erase(BB);
   }
 
   return (NumSingleThreadedBBs == SingleThreadedBBs.size())
              ? ChangeStatus::UNCHANGED
              : ChangeStatus::CHANGED;
 }
 
+/// Try to replace memory allocation calls called by a single thread with a
+/// static buffer of shared memory.
+struct AAHeapToShared : public StateWrapper<BooleanState, AbstractAttribute> {
+  using Base = StateWrapper<BooleanState, AbstractAttribute>;
+  AAHeapToShared(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
+
+  /// Create an abstract attribute view for the position \p IRP.
+  static AAHeapToShared &createForPosition(const IRPosition &IRP,
+                                           Attributor &A);
+
+  /// See AbstractAttribute::getName().
+  const std::string getName() const override { return "AAHeapToShared"; }
+
+  /// See AbstractAttribute::getIdAddr().
+  const char *getIdAddr() const override { return &ID; }
+
+  /// This function should return true if the type of the \p AA is
+  /// AAHeapToShared.
+  static bool classof(const AbstractAttribute *AA) {
+    return (AA->getIdAddr() == &ID);
+  }
+
+  /// Unique ID (due to the unique address)
+  static const char ID;
+};
+
+struct AAHeapToSharedFunction : public AAHeapToShared {
+  AAHeapToSharedFunction(const IRPosition &IRP, Attributor &A)
+      : AAHeapToShared(IRP, A) {}
+
+  const std::string getAsStr() const override {
+    return "[AAHeapToShared] " + std::to_string(MallocCalls.size()) +
+           " malloc calls eligible.";
+  }
+
+  /// See AbstractAttribute::trackStatistics().
+  void trackStatistics() const override {}
+
+  void initialize(Attributor &A) override {
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
+
+    for (User *U : RFI.Declaration->users())
+      if (CallBase *CB = dyn_cast<CallBase>(U))
+        MallocCalls.insert(CB);
+  }
+
+  ChangeStatus manifest(Attributor &A) override {
+    if (MallocCalls.empty())
+      return ChangeStatus::UNCHANGED;
+
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    auto &FreeCall = OMPInfoCache.RFIs[OMPRTL___kmpc_free_shared];
+
+    Function *F = getAnchorScope();
+    auto *HS = A.lookupAAFor<AAHeapToStack>(IRPosition::function(*F), this,
+                                            DepClassTy::OPTIONAL);
+
+    ChangeStatus Changed = ChangeStatus::UNCHANGED;
+    for (CallBase *CB : MallocCalls) {
+      // Skip replacing this if HeapToStack has already claimed it.
+      if (HS && HS->isKnownHeapToStack(*CB))
+        continue;
+
+      // Find the unique free call to remove it.
+      SmallVector<CallBase *, 4> FreeCalls;
+      for (auto *U : CB->users()) {
+        CallBase *C = dyn_cast<CallBase>(U);
+        if (C && C->getCalledFunction() == FreeCall.Declaration)
+          FreeCalls.push_back(C);
+      }
+      if (FreeCalls.size() != 1)
+        continue;
+
+      ConstantInt *AllocSize = dyn_cast<ConstantInt>(CB->getArgOperand(0));
+
+      LLVM_DEBUG(dbgs() << TAG << "Replace globalization call in "
+                        << CB->getCaller()->getName() << " with "
+                        << AllocSize->getZExtValue()
+                        << " bytes of shared memory\n");
+
+      // Create a new shared memory buffer of the same size as the allocation
+      // and replace all the uses of the original allocation with it.
+      Module *M = CB->getModule();
+      Type *Int8Ty = Type::getInt8Ty(M->getContext());
+      Type *Int8ArrTy = ArrayType::get(Int8Ty, AllocSize->getZExtValue());
+      auto *SharedMem = new GlobalVariable(
+          *M, Int8ArrTy, /* IsConstant */ false, GlobalValue::InternalLinkage,
+          UndefValue::get(Int8ArrTy), CB->getName(), nullptr,
+          GlobalValue::NotThreadLocal,
+          static_cast<unsigned>(AddressSpace::Shared));
+      auto *NewBuffer =
+          ConstantExpr::getPointerCast(SharedMem, Int8Ty->getPointerTo());
+
+      SharedMem->setAlignment(MaybeAlign(32));
+
+      A.changeValueAfterManifest(*CB, *NewBuffer);
+      A.deleteAfterManifest(*CB);
+      A.deleteAfterManifest(*FreeCalls.front());
+
+      NumBytesMovedToSharedMemory += AllocSize->getZExtValue();
+      Changed = ChangeStatus::CHANGED;
+    }
+
+    return Changed;
+  }
+
+  ChangeStatus updateImpl(Attributor &A) override {
+    auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
+    auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
+    Function *F = getAnchorScope();
+
+    auto NumMallocCalls = MallocCalls.size();
+
+    // Only consider malloc calls executed by a single thread with a constant.
+    for (User *U : RFI.Declaration->users()) {
+      const auto &ED = A.getAAFor<AAExecutionDomain>(
+          *this, IRPosition::function(*F), DepClassTy::REQUIRED);
+      if (CallBase *CB = dyn_cast<CallBase>(U))
+        if (!dyn_cast<ConstantInt>(CB->getArgOperand(0)) ||
+            !ED.isExecutedByInitialThreadOnly(*CB))
+          MallocCalls.erase(CB);
+    }
+
+    if (NumMallocCalls != MallocCalls.size())
+      return ChangeStatus::CHANGED;
+
+    return ChangeStatus::UNCHANGED;
+  }
+
+  /// Collection of all malloc calls in a function.
+  SmallPtrSet<CallBase *, 4> MallocCalls;
+};
+
 } // namespace
 
 const char AAICVTracker::ID = 0;
 const char AAExecutionDomain::ID = 0;
+const char AAHeapToShared::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
   case IRPosition::IRP_ARGUMENT:
AAExecutionDomain &AAExecutionDomain::createForPosition(const IRPosition &IRP,
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAExecutionDomainFunction(IRP, A);
     break;
   }
 
   return *AA;
 }
 
+AAHeapToShared &AAHeapToShared::createForPosition(const IRPosition &IRP,
+                                                  Attributor &A) {
+  AAHeapToSharedFunction *AA = nullptr;
+  switch (IRP.getPositionKind()) {
+  case IRPosition::IRP_INVALID:
+  case IRPosition::IRP_FLOAT:
+  case IRPosition::IRP_ARGUMENT:
+  case IRPosition::IRP_CALL_SITE_ARGUMENT:
+  case IRPosition::IRP_RETURNED:
+  case IRPosition::IRP_CALL_SITE_RETURNED:
+  case IRPosition::IRP_CALL_SITE:
+    llvm_unreachable(
+        "AAHeapToShared can only be created for function position!");
+  case IRPosition::IRP_FUNCTION:
+    AA = new (A.Allocator) AAHeapToSharedFunction(IRP, A);
+    break;
+  }
+
+  return *AA;
+}
+
 PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
   if (!containsOpenMP(M, OMPInModule))
     return PreservedAnalyses::all();
 
   if (DisableOpenMPOptimizations)
     return PreservedAnalyses::all();
 
   // Look at every function definition in the Module.

llvm/test/Transforms/OpenMP/replace_globalization.ll

+; RUN: opt -S -passes='openmp-opt' < %s | FileCheck %s
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64"
+
+; CHECK: [[SHARED_X:@.+]] = internal addrspace(3) global [16 x i8] undef
+; CHECK: [[SHARED_Y:@.+]] = internal addrspace(3) global [4 x i8] undef
+
+; CHECK: %{{.*}} = call i8* @__kmpc_alloc_shared({{.*}})
+; CHECK: call void @__kmpc_free_shared({{.*}})
+define dso_local void @foo() {
+entry:
+  %x = call i8* @__kmpc_alloc_shared(i64 4)
+  %x_on_stack = bitcast i8* %x to i32*
+  %0 = bitcast i32* %x_on_stack to i8*
+  call void @use(i8* %0)
+  call void @__kmpc_free_shared(i8* %x)
+  ret void
+}
+
+define void @bar() {
+  call void @baz()
+  call void @qux()
+  ret void
+}
+
+; CHECK: %{{.*}} = bitcast i8* addrspacecast (i8 addrspace(3)* getelementptr inbounds ([16 x i8], [16 x i8] addrspace(3)* [[SHARED_X]], i32 0, i32 0) to i8*) to [4 x i32]*
+define internal void @baz() {
+entry:
+  %tid = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+  %cmp = icmp eq i32 %tid, 0
+  br i1 %cmp, label %master, label %exit
+master:
+  %x = call i8* @__kmpc_alloc_shared(i64 16), !dbg !9
+  %x_on_stack = bitcast i8* %x to [4 x i32]*
+  %0 = bitcast [4 x i32]* %x_on_stack to i8*
+  call void @use(i8* %0)
+  call void @__kmpc_free_shared(i8* %x)
+  br label %exit
+exit:
+  ret void
+}
+
+; CHECK: %{{.*}} = bitcast i8* addrspacecast (i8 addrspace(3)* getelementptr inbounds ([4 x i8], [4 x i8] addrspace(3)* [[SHARED_Y]], i32 0, i32 0) to i8*) to [4 x i32]*
+define internal void @qux() {
+entry:
+  %tid = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+  %ntid = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
+  %warpsize = call i32 @llvm.nvvm.read.ptx.sreg.warpsize()
+  %0 = sub nuw i32 %warpsize, 1
+  %1 = sub nuw i32 %ntid, 1
+  %2 = xor i32 %0, -1
+  %master_tid = and i32 %1, %2
+  %3 = icmp eq i32 %tid, %master_tid
+  br i1 %3, label %master, label %exit
+master:
+  %y = call i8* @__kmpc_alloc_shared(i64 4), !dbg !10
+  %y_on_stack = bitcast i8* %y to [4 x i32]*
+  %4 = bitcast [4 x i32]* %y_on_stack to i8*
+  call void @use(i8* %4)
+  call void @__kmpc_free_shared(i8* %y)
+  br label %exit
+exit:
+  ret void
+}
+
+
+define void @use(i8* %x) {
+entry:
+  %addr = alloca i8*
+  store i8* %x, i8** %addr
+  ret void
+}
+
+declare i8* @__kmpc_alloc_shared(i64)
+
+declare void @__kmpc_free_shared(i8*)
+
+declare i32 @llvm.nvvm.read.ptx.sreg.tid.x()
+
+declare i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
+
+declare i32 @llvm.nvvm.read.ptx.sreg.warpsize()
+
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!3, !4}
+!nvvm.annotations = !{!5, !6}
+
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C99, file: !1, producer: "clang version 12.0.0", isOptimized: false, runtimeVersion: 0, emissionKind: FullDebug, enums: !2, splitDebugInlining: false, nameTableKind: None)
+!1 = !DIFile(filename: "replace_globalization.c", directory: "/tmp/replace_globalization.c")
+!2 = !{}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{void ()* @foo, !"kernel", i32 1}
+!6 = !{void ()* @bar, !"kernel", i32 1}
+!7 = distinct !DISubprogram(name: "bar", scope: !1, file: !1, line: 1, type: !8, scopeLine: 1, flags: DIFlagPrototyped, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0, retainedNodes: !2)
+!8 = !DISubroutineType(types: !2)
+!9 = !DILocation(line: 5, column: 7, scope: !7)
+!10 = !DILocation(line: 5, column: 14, scope: !7)