[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

948–949

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

1005

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

948–949

Yes you're right. I'll it.

1005

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/Transforms/IPO/OpenMPOpt.h

private:
       /// made. Returns -1 if the index can't be deduced.
       int32_t getAccessedIdx(StoreInst &S);
 
       /// Returns true if all values in StoredValues and
       /// LastAccesses are not nullptrs.
       bool isFilled();
     };
 
-    CallBase *RuntimeCall; /// Call that involves a memotry transfer.
-    InformationCache &InfoCache;
+    CallInst *RuntimeCall; /// Call that involves a memotry transfer.
+    OMPInformationCache &InfoCache;
 
     /// These help mapping the values in offload_baseptrs, offload_ptrs, and
     /// offload_sizes, respectively.
     std::unique_ptr<OffloadArray> BasePtrs = nullptr;
     std::unique_ptr<OffloadArray> Ptrs = nullptr;
     std::unique_ptr<OffloadArray> Sizes = nullptr;
 
     /// Set of instructions that compose the argument setup for the call
     /// RuntimeCall.
     SetVector<Instruction *> Issue;
 
-    MemoryTransfer(CallBase *RuntimeCall, InformationCache &InfoCache) :
+    MemoryTransfer(CallInst *RuntimeCall, OMPInformationCache &InfoCache) :
         RuntimeCall{RuntimeCall}, InfoCache{InfoCache}
     {}
 
     /// Maps the values physically (the IR) stored in the offload arrays
     /// offload_baseptrs, offload_ptrs, offload_sizes to their corresponding
     /// members, BasePtrs, Ptrs, Sizes.
     /// Returns false if one of the arrays couldn't be processed or some of the
     /// values couldn't be found.
private:
   /// asynchronously, returning a handle. The "wait" waits in the returned
   /// handle for the memory transfer to finish.
   bool hideMemTransfersLatency();
 
   /// Returns a pointer to the instruction where the "issue" of \p MT can be
   /// moved. Returns nullptr if the movement is not possible, or not worth it.
   Instruction *canBeMovedUpwards(MemoryTransfer &MT);
 
+  /// Returns a pointer to the instruction where the "wait" of \p MT can be
+  /// moved. Returns nullptr if the movement is not possible, or not worth it.
+  Instruction *canBeMovedDownwards(MemoryTransfer &MT);
+
   static Value *combinedIdentStruct(Value *CurrentIdent, Value *NextIdent,
                                     bool GlobalOnly, bool &SingleChoice);
 
   /// Return an `struct ident_t*` value that represents the ones used in the
   /// calls of \p RFI inside of \p F. If \p GlobalOnly is true, we will not
   /// return a local `struct ident_t*`. For now, if we cannot find a suitable
   /// return value we create one from scratch. We also do not yet combine
   /// information, e.g., the source locations, see combinedIdentStruct.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

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 (auto *I = canBeMovedUpwards(MT)) {
-      // TODO: Split call and move "issue" below I.
+    if (canBeMovedUpwards(MT) || canBeMovedDownwards(MT)) {
+      // TODO: Split runtime call.
     }
+
     return false;
   };
 
   RFI.foreachUse(SplitDataTransfer);
   return Changed;
 }
 
 Instruction *OpenMPOpt::canBeMovedUpwards(MemoryTransfer &MT) {
   assert(MT.Issue.size() > 0 && "There's not set of instructions to be moved!");
 
-  CallBase *RC = MT.RuntimeCall;
+  CallInst *RC = MT.RuntimeCall;
   auto *MSSAResult =
       OMPInfoCache.getAnalysisResultForFunction<MemorySSAAnalysis>(
           *RC->getCaller());
   if (!MSSAResult) {
     LLVM_DEBUG(dbgs() << TAG << "Couldn't get MemorySSAAnalysis in function "
                       << RC->getCaller()->getName() << "\n");
     return nullptr;
   }
 
   auto &MSSA = MSSAResult->getMSSA();
   auto *MSSAWalker = MSSA.getWalker();
   const auto *LiveOnEntry = MSSA.getLiveOnEntryDef();
   auto *MemAccess = MSSAWalker->getClobberingMemoryAccess(RC);
 
   while (MemAccess != LiveOnEntry) {
     if (!isa<MemoryDef>(MemAccess))
       continue;
 
     auto *MemInst = (cast<MemoryDef>(MemAccess))->getMemoryInst();
-    if (MT.mayBeModifiedBy(MemInst))
+    if (MT.mayBeModifiedBy(MemInst)) {
+      // If MemInst is not the instruction immediately before the Issue.
+      if (!MT.Issue.count(MemInst->getNextNode()))
         return MemInst;
 
+      return nullptr;
+    }
+
     MemAccess = MSSAWalker->getClobberingMemoryAccess(MemAccess);
   }
 
   return nullptr;
 }
 
+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)
+          return Inst;
+
+        return nullptr;
+      }
+
+      // If it's a call, it is worth the movement below the call.
+      IsWorthIt = true;
+    } else if (isa<StoreInst>(Inst))
+      // If we go over a StoreInst the movement is worth it too.
+      IsWorthIt = true;
+  }
+
+  // Return end of BasicBlock.
+  return &*(RC->getParent()->end());
+}
+
 Value *OpenMPOpt::combinedIdentStruct(Value *CurrentIdent, Value *NextIdent,
     bool GlobalOnly, bool &SingleChoice) {
   if (CurrentIdent == NextIdent)
     return CurrentIdent;
 
   // TODO: Figure out how to actually combine multiple debug locations. For
   //       now we just keep an existing one if there is a single choice.
   if (!GlobalOnly || isa<GlobalValue>(NextIdent)) {