[OpenMPOpt][SplitMemTransfer] Can move setup instructions upwards?

[hamax97]

Summary

Checks if the setup (issue) instructions for __tgt_target_data_begin_mapper can be moved upwards.

[jdoerfert]

Tests?

[hamax97]

In D87260#2260098, @jdoerfert wrote:

Tests?

The regression test for this is hide_mem_transfer_latency.ll. But there's still the issue with the AAManager, so the test has the fixmes of where the set up instructions should be moved.

[jdoerfert]

In D87260#2261006, @hamax97 wrote:

In D87260#2260098, @jdoerfert wrote:

Tests?

The regression test for this is hide_mem_transfer_latency.ll. But there's still the issue with the AAManager, so the test has the fixmes of where the set up instructions should be moved.

But it seems not to include any changes to the test. why not? because this doesn't have an impact without resolving the issue you mentioned?

[hamax97]

In D87260#2261012, @jdoerfert wrote:

In D87260#2261006, @hamax97 wrote:

In D87260#2260098, @jdoerfert wrote:

Tests?

The regression test for this is hide_mem_transfer_latency.ll. But there's still the issue with the AAManager, so the test has the fixmes of where the set up instructions should be moved.

But it seems not to include any changes to the test. why not? because this doesn't have an impact without resolving the issue you mentioned?

Exactly, the mayModify function always returns true at the instruction immediately above the issue, so there's no movement.

[hamax97]

• Fixed bug when creating the wait runtime call.

[jdoerfert]

tests?

[hamax97]

In D87260#2275287, @jdoerfert wrote:

tests?

Recall our previous discussion. Given the issue with the AAManager there's no movement in any of the test cases. Anyway, we need first to review D88029 before accepting any of this.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.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/CallGraphUpdater.h"
 #include "llvm/Analysis/ValueTracking.h"
+#include "llvm/Analysis/MemorySSA.h"
 
 using namespace llvm;
 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
struct MemoryTransfer {
   /// The values stored in %offload_baseptrs, %offload_ptrs, and %offload_sizes
   /// (before RuntimeCall is issued) are mapped here, in the first, second, and
   /// third position, respectively.
   ArrayRef<OffloadArray> OffloadArrays;
   /// Used to access OffloadArrays.
   static const unsigned BasePtrsIdx = 0;
   static const unsigned PtrsIdx = 1;
   static const unsigned SizesIdx = 2;
+
   /// Set of instructions that compose the argument setup for RuntimeCall.
   SetVector<Instruction *> Issue;
 
   /// Factory function that initializes the MemoryTrasfer by getting the
   /// instructions that compose the arguments setup ("issue") of RuntimeCall.
   /// This should be used instead of the constructor.
   bool initialize(CallInst &RuntimeCall, ArrayRef<OffloadArray> OffloadArrays) {
     assert(OffloadArrays.size() == 3 &&
            OffloadArrays[BasePtrsIdx].isInitialized() &&
            OffloadArrays[PtrsIdx].isInitialized() &&
            "No offload arrays to look at!");
 
     if (!detectIssue(RuntimeCall, OffloadArrays))
       return false;
 
     this->RuntimeCall = &RuntimeCall;
     this->OffloadArrays = OffloadArrays;
     return true;
   }
 
+  bool isInitialized() { return RuntimeCall != nullptr; }
+
+  void moveIssue(Instruction &After) {
+    Instruction *Before = After.getNextNode();
+    for (auto *I : Issue) {
+      I->removeFromParent();
+      I->insertBefore(Before);
+    }
+  }
+
+  const OffloadArray &getBasePtrs() { return OffloadArrays[BasePtrsIdx]; }
+  const OffloadArray &getPtrs() { return OffloadArrays[PtrsIdx]; }
+  const OffloadArray &getSizes() { return OffloadArrays[SizesIdx]; }
+
   MemoryTransfer() = default;
   MemoryTransfer(const MemoryTransfer &) = delete;
   MemoryTransfer &operator=(const MemoryTransfer &) = delete;
 
 private:
   /// Groups into Issue the instructions that compose the argument setup for
   /// RuntimeCall.
   bool detectIssue(CallInst &RuntimeCall,
auto SplitMemTransfers = [&](Use &U, Function &Decl) {
         return false;
       LLVM_DEBUG(dumpValuesInOffloadArrays(OffloadArrays));
 
       MemoryTransfer MT;
       if (!MT.initialize(*RTCall, OffloadArrays))
         return false;
       LLVM_DEBUG(dumpMemoryTransferSetupInstructions(MT));
 
-      // TODO: Check if can be moved upwards.
       bool WasSplit = false;
-      Instruction *WaitMovementPoint = canBeMovedDownwards(*RTCall);
-      if (WaitMovementPoint)
-        WasSplit = splitTargetDataBeginRTC(*RTCall, *WaitMovementPoint);
+      Instruction *IssueNewLocation = canBeMovedUpwards(MT);
+      Instruction *WaitNewLocation = canBeMovedDownwards(*RTCall);
+      if (IssueNewLocation || WaitNewLocation)
+        WasSplit = splitTargetDataBeginRTC(
+            MT, IssueNewLocation, WaitNewLocation);
 
       Changed |= WasSplit;
       return WasSplit;
     };
     RFI.foreachUse(SCC, SplitMemTransfers);
 
     return Changed;
   }
void dumpMemoryTransferSetupInstructions(MemoryTransfer &MT) {
 
     for (auto *I : MT.Issue) {
       I->print(Printer);
       Printer << "\n";
     }
     LLVM_DEBUG(dbgs() << Printer.str());
   }
 
+  Instruction *canBeMovedUpwards(MemoryTransfer &MT) {
+    assert(MT.isInitialized() && "Initialize MemoryTransfer first!");
+
+    CallInst *RC = MT.RuntimeCall;
+    auto *MSSAResult =
+        OMPInfoCache.getAnalysisResultForFunction<MemorySSAAnalysis>(
+            *RC->getCaller());
+    if (!MSSAResult)
+      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 (mayBeModifiedBy(MT, *MemInst)) {
+        // If MemInst is not the instruction immediately before the setup
+        // instructions, that is, MT.Issue, the movement is worth it.
+        if (!MT.Issue.count(MemInst->getNextNode()))
+          return MemInst;
+
+        return nullptr;
+      }
+
+      MemAccess = MSSAWalker->getClobberingMemoryAccess(MemAccess);
+    }
+
+    return nullptr;
+  }
+
+  bool mayBeModifiedBy(MemoryTransfer &MT, Instruction &I) {
+    if (MT.Issue.count(&I))
+      return false;
+
+    if (mayModify(I, MT.getBasePtrs().StoredValues))
+      return true;
+    if (mayModify(I, MT.getPtrs().StoredValues))
+      return true;
+    auto &OffloadSizes = MT.getSizes();
+    if (OffloadSizes.isInitialized()) {
+      if (mayModify(I, OffloadSizes.StoredValues))
+        return true;
+    }
+
+    return false;
+  }
+
+  bool mayModify(Instruction &I, const SmallVectorImpl<Value *> &Values) {
+    auto *AAResults = OMPInfoCache.getAnalysisResultForFunction<AAManager>(
+        *I.getFunction());
+
+    if (!AAResults)
+      return true;
+
+    if (isa<StoreInst>(&I)) {
+      auto *Dst = getUnderlyingObject(I.getOperand(1));
+      for (auto *V : Values) {
+        if (Dst == V) {
+          return true;
+        }
+      }
+    } else if (isa<CallInst>(&I)) {
+      for (auto *V : Values) {
+        // FIXME: This usage of the AAResults is not working properly. It always
+        //        returns that the call instruction I may modify a value V.
+        //        For example:
+        //        define i32 @func(double* noalias %a) {
+        //        ...
+        //        %1 = call i32 @rand()
+        //        ...
+        //        }
+        //        The getModRefInfo always returns that rand() modifies %a, even
+        //        though it has the noalias attribute.
+        auto ModRefResult =
+            AAResults->getModRefInfo(
+                &I, MemoryLocation(
+                        V, LocationSize::precise(
+                               V->getType()->getPrimitiveSizeInBits()))
+        );
+        if (isModSet(ModRefResult))
+          return true;
+      }
+    }
+
+    return true;
+  }
+
   /// Returns the instruction where the "wait" counterpart \p RuntimeCall can be
   /// moved. Returns nullptr if the movement is not possible, or not worth it.
   Instruction *canBeMovedDownwards(CallInst &RuntimeCall) {
     // FIXME: This traverses only the BasicBlock where RuntimeCall is.
     //  Make it traverse the CFG.
 
     Instruction *CurrentI = &RuntimeCall;
     bool IsWorthIt = false;
while ((CurrentI = CurrentI->getNextNode())) {
       //  is worth it.
       IsWorthIt = true;
     }
 
     // Return end of BasicBlock.
     return RuntimeCall.getParent()->getTerminator();
   }
 
-  /// Splits \p RuntimeCall into its "issue" and "wait" counterparts.
-  bool splitTargetDataBeginRTC(CallInst &RuntimeCall,
-                               Instruction &WaitMovementPoint) {
+  /// Splits \p MT into its "issue" and "wait" counterparts.
+  bool splitTargetDataBeginRTC(MemoryTransfer &MT,
+                               Instruction *After, Instruction *Before) {
+    assert(MT.isInitialized() && "Must initialize MemoryTransfer first!");
+    assert(After || Before && "No place to move the split runtime call!");
+
     auto &IRBuilder = OMPInfoCache.OMPBuilder;
+    CallInst &RuntimeCall = *MT.RuntimeCall;
     // 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_mapper_issue);
 
     // Change RuntimeCall call site for its asynchronous version.
     SmallVector<Value *, 8> Args;
     for (auto &Arg : RuntimeCall.args())
       Args.push_back(Arg.get());
 
     CallInst *IssueCallsite =
         CallInst::Create(IssueDecl, Args, "handle", &RuntimeCall);
     RuntimeCall.eraseFromParent();
+    MT.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_mapper_wait);
 
     // Add call site to WaitDecl.
     Value *WaitParams[2] = {
         IssueCallsite->getArgOperand(0), // device_id.
         IssueCallsite // returned handle.
     };
-    CallInst::Create(WaitDecl, WaitParams, /*NameStr=*/"", &WaitMovementPoint);
+    CallInst *WaitCallsite =
+        CallInst::Create(WaitDecl, WaitParams, /*NameStr=*/"",
+                         /*InsertBefore=*/ (Instruction *)nullptr);
+
+    if (!Before)
+      WaitCallsite->insertAfter(IssueCallsite);
+    else
+      WaitCallsite->insertBefore(Before);
+
+    if (After)
+      MT.moveIssue(*After);
 
     return true;
   }
 
   static Value *combinedIdentStruct(Value *CurrentIdent, Value *NextIdent,
                                     bool GlobalOnly, bool &SingleChoice) {
     if (CurrentIdent == NextIdent)
       return CurrentIdent;