[OpenMPOpt][SplitMemTransfer] Grouping the setup instructions for the runtime call.

[hamax97]

Summary

This is just a proposal on how to group the "issue" of the memory transfer __tgt_target_data_begin_mapper.

Based on the information previously stored in OffloadArrays, that is the last stores made to each of the offload array, the structure MemoryTransfer groups these setup instructions into Issue.

[hamax97]

• NIT changes to comments

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

937

We should dump MT or do something to test this.

Also, do we store MT somewhere? Otherwise we might be able to just use a object that cannot be copied.

973

Commit this part as NFC before this patch.

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

937

Sure, I'll use LLVM_DEBUG to print MT.Issue and I'll create another regression test that checks for that debug information.

Sure, I'll do it. MT is intended to be used for the next step in the algorithm, that is, determining where MT.Issue can be moved. But not for long term storage.

973

Before I can commit this we need to commit https://reviews.llvm.org/D86300.

[hamax97]

• Adding regression test

[sstefan1]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

611

I think it would be useful to make these constants with proper names, eg. constexpr int OffloadPtrsIdx = 0; or something similar. WDYT?

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

611

I thought abot that but I didn't do it because it might've been too descriptive. Perhaps as attributes of MemoryTransfer would be even better, WDYT?.

[hamax97]

• Added constants to ease the access to OffloadArrays in the function detectIssue.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

bool isFilled() {
     const unsigned NumValues = StoredValues.size();
     for (unsigned I = 0; I < NumValues; ++I) {
       if (!StoredValues[I] || !LastAccesses[I])
         return false;
     }
 
     return true;
   }
+
+  /// Position of the offload arrays arguments in the runtime call:
+  /// __tgt_target_data_begin_mapper.
+  static const unsigned BasePtrsArgNum = 2;
+  static const unsigned PtrsArgNum = 3;
+  static const unsigned SizesArgNum = 4;
 };
 
 using OffloadArrayPtr = std::unique_ptr<OffloadArray>;
 
+struct MemoryTransfer;
+using MemoryTransferPtr = std::unique_ptr<MemoryTransfer>;
+
+/// Used to store the instructions that serve as setup ("issue") for the
+/// offload arrays that the runtime call __tgt_target_data_begin_mapper has.
+/// __tgt_target_data_begin(...,
+///   i8** %offload_baseptrs, i8** %offload_ptrs, i64* %offload_sizes,
+/// ...)
+struct MemoryTransfer {
+  /// Call to __tgt_target_data_begin_mapper.
+  CallInst &RuntimeCall;
+  /// 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<OffloadArrayPtr> OffloadArrays;
+  /// 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.
+  static MemoryTransferPtr initialize(CallInst &RuntimeCall,
+                                      ArrayRef<OffloadArrayPtr> OffloadArrays) {
+
+    assert(OffloadArrays.size() > 2 && OffloadArrays[0] && OffloadArrays[1] &&
+           "No offload arrays to look at!");
+
+    auto MT = std::make_unique<MemoryTransfer>(RuntimeCall, OffloadArrays);
+    if (!MT->detectIssue())
+      return nullptr;
+    return MT;
+  }
+
+  MemoryTransfer(CallInst &RuntimeCall, ArrayRef<OffloadArrayPtr> OffloadArrays)
+      : RuntimeCall(RuntimeCall), OffloadArrays(OffloadArrays) {}
+
+private:
+  /// Groups into Issue the instructions that compose the argument setup for
+  /// RuntimeCall.
+  bool detectIssue() {
+
+    // Get setup instruction for %offload_baseptrs.
+    if (!getSetupInstructions(*OffloadArrays[0]))
+      return false;
+
+    auto *BasePtrsGEP = cast<Instruction>(
+        RuntimeCall.getArgOperand(OffloadArray::BasePtrsArgNum));
+    Issue.insert(BasePtrsGEP);
+
+    // Get setup instruction for %offload_ptrs.
+    if (!getSetupInstructions(*OffloadArrays[1]))
+      return false;
+    auto *PtrsGEP =
+        cast<Instruction>(RuntimeCall.getArgOperand(OffloadArray::PtrsArgNum));
+    Issue.insert(PtrsGEP);
+
+    // Get setup instruction for %offload_sizes.
+    if (OffloadArrays[2]) {
+      if (!getSetupInstructions(*OffloadArrays[2]))
+        return false;
+
+      auto *SizesGEP = cast<Instruction>(
+          RuntimeCall.getArgOperand(OffloadArray::SizesArgNum));
+      Issue.insert(SizesGEP);
+    }
+
+    return true;
+  }
+
+  /// Gets the setup instructions for each of the values in \p OA. These
+  /// instructions are stored into Issue.
+  bool getSetupInstructions(OffloadArray &OA) {
+    for (auto *S : OA.LastAccesses) {
+      if (!getValueSetupInstructions(*S))
+        return false;
+
+      if (!getPointerSetupInstructions(*S))
+        return false;
+
+      Issue.insert(S);
+    }
+    return true;
+  }
+
+  /// Gets the setup instructions for the value operand of \p S.
+  bool getValueSetupInstructions(StoreInst &S) {
+    auto *V = S.getValueOperand();
+    // Auxiliary storage to later insert the found instructions in the order
+    // needed.
+    SmallVector<Instruction *, 8> TempStorage;
+    bool Success = false;
+    unsigned MaxLookup = 6;
+    for (unsigned I = 0; I < MaxLookup; ++I) {
+      if (isa<AllocaInst>(V) || isa<Argument>(V) || isa<GlobalValue>(V) ||
+          isa<Constant>(V)) {
+        Success = true;
+        break;
+      }
+
+      if (!isa<Instruction>(V)) {
+        Success = false;
+        break;
+      }
+
+      auto *Inst = cast<Instruction>(V);
+      TempStorage.push_back(Inst);
+
+      // FIXME: Inst might depend on more instructions through its second operand.
+      V = Inst->getOperand(0);
+    }
+
+    if (Success)
+      while (!TempStorage.empty())
+        Issue.insert(TempStorage.pop_back_val());
+
+    return Success;
+  }
+
+  /// Gets the setup instructions for the pointer operand of \p S.
+  bool getPointerSetupInstructions(StoreInst &S) {
+    auto *P = S.getPointerOperand();
+
+    // TODO: P might be a global value. Make it general.
+    if (!isa<Instruction>(P))
+      return false;
+
+    auto *DstInst = cast<Instruction>(P);
+    if (isa<GetElementPtrInst>(DstInst)) {
+      Issue.insert(DstInst);
+
+    } else if (DstInst->isCast()) {
+      auto *Casted = DstInst->getOperand(0);
+
+      // TODO: Casted might be a global value. Make it general.
+      if (!isa<Instruction>(Casted))
+        return false;
+
+      if (auto *GEP = dyn_cast<GetElementPtrInst>(Casted))
+        Issue.insert(GEP);
+
+      Issue.insert(DstInst);
+    }
+
+    return true;
+  }
+};
+
 struct OpenMPOpt {
 
   using OptimizationRemarkGetter =
       function_ref<OptimizationRemarkEmitter &(Function *)>;
 
   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
             OptimizationRemarkGetter OREGetter,
             OMPInformationCache &OMPInfoCache, Attributor &A)
auto SplitMemTransfers = [&](Use &U, Function &Decl) {
       if (!RTCall)
         return false;
 
       OffloadArrayPtr OffloadArrays[3];
       if (!getValuesInOffloadArrays(*RTCall, OffloadArrays))
         return false;
       debugValuesInOffloadArrays(OffloadArrays);
 
-      // TODO: Check if can be moved upwards.
+      MemoryTransferPtr MT = MemoryTransfer::initialize(*RTCall, OffloadArrays);
+      if (!MT)
+        return false;
+
+      // TODO: Check if MT can be moved upwards.
       bool WasSplit = false;
       Instruction *WaitMovementPoint = canBeMovedDownwards(*RTCall);
       if (WaitMovementPoint)
         WasSplit = splitTargetDataBeginRTC(*RTCall, *WaitMovementPoint);
 
       Changed |= WasSplit;
       return WasSplit;
     };
bool getValuesInOffloadArrays(CallInst &RuntimeCall,
     // A runtime call that involves memory offloading looks something like:
     // call void @__tgt_target_data_begin_mapper(arg0, arg1,
     //   i8** %offload_baseptrs, i8** %offload_ptrs, i64* %offload_sizes,
     // ...)
     // So, the idea is to access the allocas that allocate space for these
     // offload arrays, offload_baseptrs, offload_ptrs, offload_sizes.
     // Therefore:
     // i8** %offload_baseptrs.
-    const unsigned BasePtrsArgNum = 2;
-    Value *BasePtrsArg = RuntimeCall.getArgOperand(BasePtrsArgNum);
+    Value *BasePtrsArg =
+        RuntimeCall.getArgOperand(OffloadArray::BasePtrsArgNum);
     // i8** %offload_ptrs.
-    const unsigned PtrsArgNum = 3;
-    Value *PtrsArg = RuntimeCall.getArgOperand(PtrsArgNum);
+    Value *PtrsArg =
+        RuntimeCall.getArgOperand(OffloadArray::PtrsArgNum);
     // i8** %offload_sizes.
-    const unsigned SizesArgNum = 4;
-    Value *SizesArg = RuntimeCall.getArgOperand(SizesArgNum);
+    Value *SizesArg =
+        RuntimeCall.getArgOperand(OffloadArray::SizesArgNum);
 
     // Get values stored in **offload_baseptrs.
     auto *V = getUnderlyingObject(BasePtrsArg);
     if (!isa<AllocaInst>(V))
       return false;
     auto *BasePtrsArray = cast<AllocaInst>(V);
     OAs[0] = OffloadArray::initialize(*BasePtrsArray, RuntimeCall);
     if (!OAs[0])