[OpenMPOpt][SplitMemTransfer] Is it possible/worth it to split the runtime call?

[hamax97]

Summary

• Fix to function canBeMovedUpwards.
• Declaration of function canBeMovedDownwards.

[sstefan1]

Can you add parent revisions for this and maybe other related patches? It will be easier to know the order.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1005–1006

If I'm not wrong canBeMovedDownwards() never returns nullptr?

1061

I'd put Instruction * here.
Also, Maybe put it in the while condition?

[hamax97]

In D85840#2213676, @sstefan1 wrote:

Can you add parent revisions for this and maybe other related patches? It will be easier to know the order.

Yes, I already did it. Sorry for not doing it before.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1005–1006

Yes you're right. I'll it.

1061

Sure, I'll change it.

[hamax97]

• Fix to canBeMovedDownwards.

[hamax97]

This comment was removed by hamax97.

[hamax97]

Sorry. Error updating last patch.

• Fix to canBeMovedDownwards

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 #ifndef OMP_STRUCT_TYPE
 #define OMP_STRUCT_TYPE(VarName, StructName, ...)
 #endif
 
 #define __OMP_STRUCT_TYPE(VarName, Name, ...)                                  \
   OMP_STRUCT_TYPE(VarName, "struct." #Name, __VA_ARGS__)
 
 __OMP_STRUCT_TYPE(Ident, ident_t, Int32, Int32, Int32, Int32, Int8Ptr)
+__OMP_STRUCT_TYPE(AsyncInfo, __tgt_async_info, Int8Ptr)
 
 #undef __OMP_STRUCT_TYPE
 #undef OMP_STRUCT_TYPE
 
 #ifndef OMP_FUNCTION_TYPE
 #define OMP_FUNCTION_TYPE(VarName, IsVarArg, ReturnType, ...)
 #endif
 
__OMP_RTL(__tgt_target_teams, false, Int32, Int64, VoidPtr, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr, Int32, Int32)
 __OMP_RTL(__tgt_target_teams_nowait, false, Int32, Int64, VoidPtr, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr, Int32, Int32)
 __OMP_RTL(__tgt_register_requires, false, Void, Int64)
 __OMP_RTL(__tgt_target_data_begin, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr)
 __OMP_RTL(__tgt_target_data_begin_nowait, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr)
+__OMP_RTL(__tgt_target_data_begin_issue, false, AsyncInfo, Int64, Int32, VoidPtrPtr,
+          VoidPtrPtr, Int64Ptr, Int64Ptr)
+__OMP_RTL(__tgt_target_data_begin_wait, false, Void, Int64, AsyncInfo)
 __OMP_RTL(__tgt_target_data_end, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr)
 __OMP_RTL(__tgt_target_data_end_nowait, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr)
 __OMP_RTL(__tgt_target_data_update, false, Void, Int64, Int32, VoidPtrPtr,
           VoidPtrPtr, Int64Ptr, Int64Ptr)
 __OMP_RTL(__tgt_target_data_update_nowait, false, Void, Int64, Int32,
           VoidPtrPtr, VoidPtrPtr, Int64Ptr, Int64Ptr)

llvm/include/llvm/Transforms/IPO/OpenMPOpt.h

private:
       bool isFilled();
     };
 
     CallInst *RuntimeCall; /// Call that involves a memotry transfer.
     OMPInformationCache &InfoCache;
 
     /// These help mapping the values in offload_baseptrs, offload_ptrs, and
     /// offload_sizes, respectively.
+    const unsigned BasePtrsArgNum = 2;
     std::unique_ptr<OffloadArray> BasePtrs = nullptr;
+    const unsigned PtrsArgNum = 3;
     std::unique_ptr<OffloadArray> Ptrs = nullptr;
+    const unsigned SizesArgNum = 4;
     std::unique_ptr<OffloadArray> Sizes = nullptr;
 
     /// Set of instructions that compose the argument setup for the call
     /// RuntimeCall.
     SetVector<Instruction *> Issue;
 
     MemoryTransfer(CallInst *RuntimeCall, OMPInformationCache &InfoCache) :
         RuntimeCall{RuntimeCall}, InfoCache{InfoCache}
struct MemoryTransfer {
     /// Groups the instructions that compose the argument setup for the call
     /// RuntimeCall.
     bool detectIssue();
 
     /// Returns true if \p I might modify some of the values in the
     /// offload arrays.
     bool mayBeModifiedBy(Instruction *I);
 
+    /// Splits this object into its "issue" and "wait" corresponding runtime
+    /// calls. The "issue" is moved after \p After and the "wait" is moved
+    /// before \p Before.
+    bool split(Instruction *After, Instruction *Before);
+
   private:
     /// Gets the setup instructions for each of the values in \p OA. These
     /// instructions are stored into Issue.
     bool getSetupInstructions(std::unique_ptr<OffloadArray> &OA);
     /// Gets the setup instructions for the pointer operand of \p S.
     bool getPointerSetupInstructions(StoreInst *S);
     /// Gets the setup instructions for the value operand of \p S.
     bool getValueSetupInstructions(StoreInst *S);

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
     cl::desc("Disable OpenMP specific optimizations."), cl::Hidden,
     cl::init(false));
 
 static cl::opt<bool> PrintICVValues("openmp-print-icv-values", cl::init(false),
                                     cl::Hidden);
 
+static cl::opt<bool> SplitMemoryTransfers(
+    "openmp-split-memtransfers",
+    cl::desc("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");
bool MemoryTransfer::detectIssue() {
   bool Success = getSetupInstructions(BasePtrs);
   if (!Success) {
     LLVM_DEBUG(dbgs() << TAG << "Couldn't get setup instructions of "
                       << "offload_baseptrs. In call to "
                       << RuntimeCall->getName() << " in function "
                       << RuntimeCall->getCaller()->getName() << "\n");
     return false;
   }
+  auto *BasePtrsGEP =
+      cast<Instruction>(RuntimeCall->getArgOperand(BasePtrsArgNum));
+  Issue.insert(BasePtrsGEP);
 
   Success = getSetupInstructions(Ptrs);
   if (!Success) {
     LLVM_DEBUG(dbgs() << TAG << "Couldn't get setup instructions of "
                       << "offload_ptrs. In call to "
                       << RuntimeCall->getName() << " in function "
                       << RuntimeCall->getCaller()->getName() << "\n");
     return false;
   }
+  auto *PtrsGEP =
+      cast<Instruction>(RuntimeCall->getArgOperand(PtrsArgNum));
+  Issue.insert(PtrsGEP);
 
   if (Sizes) {
     Success = getSetupInstructions(Sizes);
     if (!Success) {
       LLVM_DEBUG(dbgs() << TAG << "Couldn't get setup instructions of "
                         << "offload_sizes. In call to "
                         << RuntimeCall->getName() << " in function "
                         << RuntimeCall->getCaller()->getName() << "\n");
       return false;
     }
+    auto *SizesGEP =
+        cast<Instruction>(RuntimeCall->getArgOperand(SizesArgNum));
+    Issue.insert(SizesGEP);
   }
 
   return true;
 }
 
 bool MemoryTransfer::getSetupInstructions(std::unique_ptr<OffloadArray> &OA) {
   for (auto *S : OA->LastAccesses) {
     if (!getValueSetupInstructions(S))
for (auto *V : Values) {
       if (isModSet(ModRefResult))
         return true;
     }
   }
 
   return true;
 }
 
+bool MemoryTransfer::split(Instruction *After, Instruction *Before) {
+  assert((After || Before) &&
+         "Must have a place to move the split runtime call");
+
+  auto *M = RuntimeCall->getModule();
+  auto &IRBuilder = InfoCache.OMPBuilder;
+  // Add "issue" runtime call declaration.
+  // declare %struct.tgt_async_info @__tgt_target_data_begin_issue(i64, i32,
+  //   i8**, i8**, i64*, i64*)
+  FunctionCallee IssueDecl = IRBuilder.getOrCreateRuntimeFunction(
+      *M, OMPRTL___tgt_target_data_begin_issue);
+
+  // Change RuntimeCall callsite for its asynchronous version.
+  SmallVector<Value *, 6> Args;
+  Args.reserve(RuntimeCall->getNumArgOperands());
+  for (auto &Arg : RuntimeCall->args())
+    Args.push_back(Arg.get());
+
+  CallInst *IssueCallsite = CallInst::Create(
+      IssueDecl, ArrayRef<Value *>(Args), "handle", RuntimeCall);
+  RuntimeCall->removeFromParent();
+  RuntimeCall->deleteValue();
+  Issue.insert(IssueCallsite);
+
+  // Add "wait" runtime call declaration.
+  // declare void @__tgt_target_data_begin_wait(i64, %struct.__tgt_async_info)
+  FunctionCallee WaitDecl = IRBuilder.getOrCreateRuntimeFunction(
+      *M, OMPRTL___tgt_target_data_begin_wait);
+
+  // Add "wait" call site.
+  const unsigned WaitNumParams = 2;
+  Value *WaitParams[] = {
+      IssueCallsite->getArgOperand(0), // device_id.
+      IssueCallsite // returned handle.
+  };
+  CallInst::Create(
+      WaitDecl, ArrayRef<Value*>(WaitParams, WaitNumParams), /*NameStr=*/"",
+      /*InsertBefore=*/(Instruction *)nullptr);
+
+  // TODO: Move the "issue" after After and the "wait" before Before.
+  return true;
+}
+
 std::unique_ptr<OffloadArray> OffloadArray::initialize(
     AllocaInst &Array, Instruction &Before, InformationCache &InfoCache) {
   if (!Array.getAllocatedType()->isArrayTy()) {
     LLVM_DEBUG(dbgs() << TAG << "Allocated type is not array.\n");
     return nullptr;
   }
 
   auto OA = std::make_unique<OffloadArray>(Array, InfoCache);
for (Function *F : OMPInfoCache.ModuleSlice) {
         emitRemarkOnFunction(F, "OpenMPICVTracker", Remark);
       }
     }
   }
 
   Changed |= runAttributor();
   Changed |= deduplicateRuntimeCalls();
   Changed |= deleteParallelRegions();
+  if (SplitMemoryTransfers)
     Changed |= hideMemTransfersLatency();
 
   return Changed;
 }
 
 CallInst *OpenMPOpt::getCallIfRegularCall(
     Use &U, OMPInformationCache::RuntimeFunctionInfo *RFI) {
   CallInst *CI = dyn_cast<CallInst>(U.getUser());
   if (CI && CI->isCallee(&U) && !CI->hasOperandBundles() &&
auto SplitDataTransfer = [&] (Use &U, Function &Decl) {
 
     Success = MT.detectIssue();
     if (!Success) {
       LLVM_DEBUG(dbgs() << TAG << "Couldn't detect issue in call to "
                         << MT.RuntimeCall->getName() << " in function "
                         << MT.RuntimeCall->getCaller()->getName() << "\n");
       return false;
     }
 
-    if (canBeMovedUpwards(MT) || canBeMovedDownwards(MT)) {
-      // TODO: Split runtime call.
-    }
-
-    return false;
+    auto *After = canBeMovedUpwards(MT);
+    auto *Before = canBeMovedDownwards(MT);
+    return (After || Before) && MT.split(After, Before);
   };
 
   RFI.foreachUse(SplitDataTransfer);
   return Changed;
 }
 
 Instruction *OpenMPOpt::canBeMovedUpwards(MemoryTransfer &MT) {
   assert(MT.Issue.size() > 0 && "There's not set of instructions to be moved!");
Instruction *OpenMPOpt::canBeMovedDownwards(MemoryTransfer &MT) {
   assert(MT.Issue.size() > 0 && "There's not set of instructions to be moved!");
 
   // FIXME: This traverses only the BasicBlock where MT is. Make it traverse
   //        the CFG.
   GlobalValue *TgtTargetDecl = M.getNamedValue("__tgt_target");
   GlobalValue *TgtTargetTeamsDecl = M.getNamedValue("__tgt_target_teams");
   GlobalValue *TgtTargetDataEndDecl = M.getNamedValue("__tgt_target_data_end");
 
   CallInst *RC = MT.RuntimeCall;
   Instruction *Inst = RC;
   bool IsWorthIt = false;
   while ((Inst = Inst->getNextNode())) {
     if (auto *C = dyn_cast<CallInst>(Inst)) {
       auto *Callee = C->getCalledFunction();
       if (Callee == TgtTargetDecl || Callee == TgtTargetTeamsDecl ||
           Callee == TgtTargetDataEndDecl) {
         if (IsWorthIt)