[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

1033

1036

No cast needed

1041

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

1044

No need for it to be a constant

1048

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.

1733

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

1740

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

1073

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

1036

Nit: hoist this out of the lambda.

1061

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

1078

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

2467

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

1065

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/ADT/PostOrderIterator.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
 #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"
-#include "llvm/IR/IntrinsicInst.h"
-#include "llvm/IR/IntrinsicsNVPTX.h"
-#include "llvm/IR/IntrinsicsAMDGPU.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 {
 
 struct AAExecutionDomain
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();
 
+      Changed |= replaceGlobalization();
       if (remarksEnabled())
         analysisGlobalization();
     } else {
       if (PrintICVValues)
         printICVs();
       if (PrintOpenMPKernels)
         printKernels();
 
for (auto &It : BB2PRMap) {
                             << " of function " << BB->getParent()->getName()
                             << "\n";);
         }
 
         MergableCIs.clear();
       }
 
       if (!MergableCIsVector.empty()) {
         Changed = true;
 
         for (auto &MergableCIs : MergableCIsVector)
           Merge(MergableCIs, BB);
         MergableCIsVector.clear();
       }
     }
 
     if (Changed) {
       /// Re-collect use for fork calls, emitted barrier calls, and
       /// 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;
   }
 
+  /// Replace globalization calls in the device with global shared memory if it
+  /// is called by a single thread.
+  bool replaceGlobalization() {
+    auto &RFI = OMPInfoCache.RFIs[OMPRTL___kmpc_alloc_shared];
+    auto &FreeCall = OMPInfoCache.RFIs[OMPRTL___kmpc_free_shared];
+    bool Changed = false;
+
+    auto ReplaceAllocCalls = [&](Use &U, Function &F) {
+      CallBase *CB = OpenMPOpt::getCallIfRegularCall(U, &RFI);
+      if (!CB)
+        return false;
+
+      auto *ED = A.lookupAAFor<AAExecutionDomain>(IRPosition::function(F));
+      if (!ED || !ED->isSingleThreadExecution(*CB))
+        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
+      CallBase *FC = nullptr;
+      for (auto *U : CB->users()) {
+        CallBase *C = dyn_cast<CallBase>(U);
+        if (C && C->getCalledFunction() == FreeCall.Declaration) {
+          if (FC)
+            return false;
+          FC = C;
+        }
+      }
+      FC->eraseFromParent();
+
+      // 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 *NewBuffer =
+          ConstantExpr::getPointerCast(SharedMem, Int8Ty->getPointerTo());
+
+      SharedMem->setAlignment(MaybeAlign(8));
+      CB->replaceAllUsesWith(NewBuffer);
+
+      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>(CB, "OpenMPReplaceGlobalization", Remark);
+
+      CB->eraseFromParent();
+
+      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)
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); }
 
   /// 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.
for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
         return false;
       };
 
       GetterRFI.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.
   {
     Optional<Kernel> &CachedKernel = UniqueKernelMap[&F];
     if (CachedKernel)
       return *CachedKernel;
 
     // TODO: We should use an AA to create an (optimistic and callback
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)

llvm/test/Transforms/OpenMP/replace_globalization.ll

+; RUN: opt -S -passes='openmp-opt' < %s | FileCheck %s
+; RUN: opt -passes=openmp-opt -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"
+
+; CHECK-REMARKS: remark: replace_globalization.c:5:7: Replaced globalized variable with 16 bytes of shared memory
+; CHECK-REMARKS: remark: replace_globalization.c:5:14: Replaced globalized variable with 4 bytes of shared memory
+; 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)