[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

991

994

No cast needed

999

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

1002

No need for it to be a constant

1006

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.

1708

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

1715

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

1031

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

994

Nit: hoist this out of the lambda.

1019

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

1036

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

2351

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

1023

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

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 {
 
 struct AAICVTracker;
bool run() {
 
     Changed |= rewriteDeviceCodeStateMachine();
 
     Changed |= runAttributor();
 
     // Recollect uses, in case Attributor deleted any.
     OMPInfoCache.recollectUses();
 
+    Changed |= replaceGlobalization();
     Changed |= deleteParallelRegions();
     if (HideMemoryTransferLatency)
       Changed |= hideMemTransfersLatency();
     if (remarksEnabled())
       analysisGlobalization();
     Changed |= deduplicateRuntimeCalls();
     if (EnableParallelRegionMerging) {
       if (mergeParallelRegions()) {
if (Changed) {
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_barrier);
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_master);
       OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_end_master);
     }
 
     return Changed;
   }
 
+  /// Replace globalization calls in the device with shared memory.
+  bool replaceGlobalization() {
+    auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
+    bool Changed = false;
+
+    auto ReplaceAllocCalls = [&](Use &U, Function &F) {
+      auto &FreeCall = OMPInfoCache.RFIs[OMPRTL___kmpc_free_shared];
+      CallBase *CB = OpenMPOpt::getCallIfRegularCall(U, &RFI);
+      if (!CB)
+        return false;
+
+      // Push every globalized variable whose lifetime is the same as the kernel
+      // to global shared memory.
+      // TODO: Determine memory lifetime and parallel level differently to catch
+      // more cases
+      if (!isKernel(F) || !isSPMDMode(F))
+        return false;
+
+      ConstantInt *AllocSize = dyn_cast<ConstantInt>(CB->getArgOperand(0));
+      if (!AllocSize)
+        return false;
+
+      LLVM_DEBUG(dbgs() << TAG << "Replace globalization call in "
+                        << CB->getCaller()->getName() << " with "
+                        << AllocSize->getZExtValue()
+                        << " bytes of shared memory\n");
+
+      // Remove the free call
+      for (auto *U : CB->users()) {
+        CallBase *FC = dyn_cast<CallBase>(U);
+        if (FC && FC->getCalledFunction() == FreeCall.Declaration) {
+          FC->eraseFromParent();
+          break;
+        }
+      }
+
+      // Create a new shared memory buffer of the same size as the allocation
+      // and replace all the uses of the original allocation with it.
+      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 *NullInt = Constant::getNullValue(Type::getInt64Ty(M.getContext()));
+      auto *GEPExpr = ConstantExpr::getGetElementPtr(
+          Int8ArrTy, SharedMem, SmallVector<Constant *>({NullInt, NullInt}));
+
+      auto *NewBuffer = new AddrSpaceCastInst(GEPExpr, Int8Ty->getPointerTo(),
+                                              CB->getName() + "_shared", CB);
+
+      SharedMem->setAlignment(MaybeAlign(8));
+      NewBuffer->setDebugLoc(CB->getDebugLoc());
+      CB->replaceAllUsesWith(NewBuffer);
+      CB->eraseFromParent();
+
+      auto Remark = [&](OptimizationRemark OR) {
+        return OR << "Replaced globalized variable with "
+                  << ore::NV("SharedMemory", AllocSize->getZExtValue())
+                  << ((AllocSize->getZExtValue() != 1) ? " bytes " : " byte ")
+                  << "of shared memory";
+      };
+      emitRemark<OptimizationRemark>(NewBuffer, "OpenMPReplaceGlobalization",
+                                     Remark);
+
+      NumBytesMovedToSharedMemory += AllocSize->getZExtValue();
+      Changed = true;
+
+      return true;
+    };
+    RFI.foreachUse(SCC, ReplaceAllocCalls);
+
+    if (Changed)
+      OMPInfoCache.recollectUsesForFunction(OMPRTL___kmpc_free_shared);
+    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)
void collectGlobalThreadIdArguments(SmallSetVector<Value *, 16> &GTIdArgs) {
     // so we cannot cache the size nor can we use a range based for.
     for (unsigned u = 0; u < GTIdArgs.size(); ++u)
       AddUserArgs(*GTIdArgs[u]);
   }
 
   /// Kernel (=GPU) optimizations and utility functions
   ///
   ///{{
+  enum class AddressSpace : unsigned {
+    Generic = 0,
+    Global = 1,
+    Shared = 3,
+    Constant = 4,
+    Local = 5,
+  };
 
   /// Check if \p F is a kernel, hence entry point for target offloading.
   bool isKernel(Function &F) { return OMPInfoCache.Kernels.count(&F); }
 
+  /// Check if \p F is an SPMD kernel
+  bool isSPMDMode(Function &F) {
+    if (!OMPInfoCache.Kernels.count(&F))
+      return false;
+
+    auto *GVal = M.getNamedValue(F.getName().str() + "_exec_mode");
+    if (auto *GVar = dyn_cast<GlobalVariable>(GVal))
+      return GVar->getInitializer()->isOneValue();
+
+    return false;
+  }
+
   /// Cache to remember the unique kernel for a function.
   DenseMap<Function *, Optional<Kernel>> UniqueKernelMap;
 
   /// Find the unique kernel that will execute \p F, if any.
   Kernel getUniqueKernelFor(Function &F);
 
   /// Find the unique kernel that will execute \p I, if any.
   Kernel getUniqueKernelFor(Instruction &I) {
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
         return false;
       };
 
       GetterRFI.foreachUse(SCC, CreateAA);
     }
   }
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
   {
     Optional<Kernel> &CachedKernel = UniqueKernelMap[&F];
     if (CachedKernel)
       return *CachedKernel;
 
     // TODO: We should use an AA to create an (optimistic and callback
     //       call-aware) call graph. For now we stick to simple patterns that
     //       are less powerful, basically the worst fixpoint.
     if (isKernel(F)) {
       CachedKernel = Kernel(&F);
       return *CachedKernel;
AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
   case IRPosition::IRP_ARGUMENT:
   case IRPosition::IRP_CALL_SITE_ARGUMENT:
     llvm_unreachable("ICVTracker can only be created for function position!");
   case IRPosition::IRP_RETURNED:
     AA = new (A.Allocator) AAICVTrackerFunctionReturned(IRP, A);
     break;
   case IRPosition::IRP_CALL_SITE_RETURNED:
     AA = new (A.Allocator) AAICVTrackerCallSiteReturned(IRP, A);
     break;
   case IRPosition::IRP_CALL_SITE:
     AA = new (A.Allocator) AAICVTrackerCallSite(IRP, A);
     break;
   case IRPosition::IRP_FUNCTION:
     AA = new (A.Allocator) AAICVTrackerFunction(IRP, A);

llvm/test/Transforms/OpenMP/replace_globalization.ll

+; RUN: opt -S -attributor -openmpopt < %s | FileCheck %s
+; RUN: opt -S -passes='attributor,cgscc(openmpopt)' < %s | FileCheck %s
+; RUN: opt -openmpopt -pass-remarks=openmp-opt -disable-output < %s 2>&1 | FileCheck %s -check-prefix=CHECK-REMARKS
+; RUN: opt -passes=openmpopt -pass-remarks=openmp-opt -disable-output < %s 2>&1 | FileCheck %s -check-prefix=CHECK-REMARKS
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64"
+
+@foo_exec_mode = weak constant i8 0
+
+; CHECK-REMARKS: remark: replace_globalization.c:5:7: Replaced globalized variable with 16 bytes of shared memory
+; CHECK: [[SHARED:@.+]] = internal addrspace(3) global [16 x i8] undef
+
+; CHECK: %{{.*}} = call i8* @__kmpc_alloc_shared({{.*}})
+; CHECK: call void @__kmpc_free_shared({{.*}})
+define 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
+}
+
+@bar_exec_mode = weak constant i8 1
+
+; CHECK: %{{.*}} = addrspacecast i8 addrspace(3)* getelementptr inbounds ([16 x i8], [16 x i8] addrspace(3)* [[SHARED]], i64 0, i64 0) to i8*, !dbg !{{[0-9]+}}
+define void @bar() {
+entry:
+  %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)
+  ret void
+}
+
+define void @use(i8* %x) {
+entry:
+  %.addr = alloca i8*
+  store i8* %x, i8** %.addr
+  ret void
+}
+
+; CHECK: declare i8* @__kmpc_alloc_shared(i64)
+declare i8* @__kmpc_alloc_shared(i64)
+
+; CHECK: declare void @__kmpc_free_shared(i8*)
+declare void @__kmpc_free_shared(i8*)
+
+
+!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)