[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

847

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.

878

Commit this part as NFC before this patch.

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

847

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.

878

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

511

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

511

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 initialize(AllocaInst &Array, Instruction &Before) {
 
     if (!getValues(Array, Before))
       return false;
 
     this->Array = &Array;
     return true;
   }
 
+  bool isInitialized() const { return Array != nullptr; }
+
   static const unsigned BasePtrsArgNum = 2;
   static const unsigned PtrsArgNum = 3;
   static const unsigned SizesArgNum = 4;
 
 private:
   /// Traverses the BasicBlock where \p Array is, collecting the stores made to
   /// \p Array, leaving StoredValues with the values stored before the
   /// instruction \p Before is reached.
for (unsigned I = 0; I < NumValues; ++I) {
       if (!StoredValues[I] || !LastAccesses[I])
         return false;
     }
 
     return true;
   }
 };
 
+/// 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 = nullptr;
+  /// 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;
+  }
+
+  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,
+                   ArrayRef<OffloadArray> OffloadArrays) {
+
+    // Get setup instruction for %offload_baseptrs.
+    if (!getSetupInstructions(OffloadArrays[BasePtrsIdx]))
+      return false;
+
+    auto *BasePtrsGEP = cast<Instruction>(
+        RuntimeCall.getArgOperand(OffloadArray::BasePtrsArgNum));
+    Issue.insert(BasePtrsGEP);
+
+    // Get setup instruction for %offload_ptrs.
+    if (!getSetupInstructions(OffloadArrays[PtrsIdx]))
+      return false;
+    auto *PtrsGEP =
+        cast<Instruction>(RuntimeCall.getArgOperand(OffloadArray::PtrsArgNum));
+    Issue.insert(PtrsGEP);
+
+    // Get setup instruction for %offload_sizes.
+    // %offload_sizes may be a constant global array, so no need to analyze it.
+    if (!OffloadArrays[SizesIdx].isInitialized())
+      return true;
+    if (!getSetupInstructions(OffloadArrays[SizesIdx]))
+      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(const 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 for later popping out the found instructions in the
+    // needed order.
+    const unsigned MaxLookup = 6;
+    SmallVector<Instruction *, MaxLookup> TempStorage;
+    bool Success = false;
+    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)
bool hideMemTransfersLatency() {
     auto SplitMemTransfers = [&](Use &U, Function &Decl) {
       auto *RTCall = getCallIfRegularCall(U, &RFI);
       if (!RTCall)
         return false;
 
       OffloadArray OffloadArrays[3];
       if (!getValuesInOffloadArrays(*RTCall, OffloadArrays))
         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);
 
       Changed |= WasSplit;
       return WasSplit;
     };
     RFI.foreachUse(SCC, SplitMemTransfers);
 
     return Changed;
bool getValuesInOffloadArrays(CallInst &RuntimeCall,
     // Therefore:
     // i8** %offload_baseptrs.
     Value *BasePtrsArg =
         RuntimeCall.getArgOperand(OffloadArray::BasePtrsArgNum);
     // i8** %offload_ptrs.
     Value *PtrsArg = RuntimeCall.getArgOperand(OffloadArray::PtrsArgNum);
     // i8** %offload_sizes.
     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);
     if (!OAs[0].initialize(*BasePtrsArray, RuntimeCall))
       return false;
 
     // Get values stored in **offload_baseptrs.
     V = getUnderlyingObject(PtrsArg);
     if (!isa<AllocaInst>(V))
       return false;
     auto *PtrsArray = cast<AllocaInst>(V);
     if (!OAs[1].initialize(*PtrsArray, RuntimeCall))
       return false;
 
     // Get values stored in **offload_sizes.
     V = getUnderlyingObject(SizesArg);
+
     // If it's a [constant] global array don't analyze it.
     if (isa<GlobalValue>(V))
       return isa<Constant>(V);
     if (!isa<AllocaInst>(V))
       return false;
 
     auto *SizesArray = cast<AllocaInst>(V);
     if (!OAs[2].initialize(*SizesArray, RuntimeCall))
void dumpValuesInOffloadArrays(ArrayRef<OffloadArray> OAs) {
 
     for (auto *S : OAs[2].StoredValues) {
       S->print(Printer);
       Printer << Separator;
     }
     LLVM_DEBUG(dbgs() << "\t\toffload_sizes: " << Printer.str() << "\n");
   }
 
+  void dumpMemoryTransferSetupInstructions(MemoryTransfer &MT) {
+    LLVM_DEBUG(dbgs() << TAG << " Successfully got set up instructions for: "
+                      << MT.RuntimeCall->getCaller()->getName() << "\n");
+    std::string ValuesStr;
+    raw_string_ostream Printer(ValuesStr);
+
+    for (auto *I : MT.Issue) {
+      I->print(Printer);
+      Printer << "\n";
+    }
+    LLVM_DEBUG(dbgs() << Printer.str());
+  }
+
   /// 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;

llvm/test/Transforms/OpenMP/runtime_call_setup_instructions.ll

+; RUN: opt -S -passes=openmpopt -aa-pipeline=basic-aa -openmp-hide-memory-transfer-latency -debug-only=openmp-opt < %s 2>&1 | FileCheck %s
+; REQUIRES: asserts
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+@.offload_maptypes. = private unnamed_addr constant [2 x i64] [i64 35, i64 35]
+@.__omp_offloading_heavyComputation.region_id = weak constant i8 0
+@.offload_sizes. = private unnamed_addr constant [2 x i64] [i64 4, i64 0]
+@.offload_maptypes.2 = private unnamed_addr constant [2 x i64] [i64 800, i64 544]
+
+; CHECK-LABEL: {{[^@]+}}Successfully got set up instructions for: heavyComputation
+; CHECK-NEXT: %2 = bitcast [2 x i8*]* %.offload_baseptrs to double**
+; CHECK-NEXT: store double* %a, double** %2, align 8
+; CHECK-NEXT: %6 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs, i64 0, i64 1
+; CHECK-NEXT: %7 = bitcast i8** %6 to i32**
+; CHECK-NEXT: store i32* %size.addr, i32** %7, align 8
+; CHECK-NEXT: %1 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs, i64 0, i64 0
+; CHECK-NEXT: %4 = bitcast [2 x i8*]* %.offload_ptrs to double**
+; CHECK-NEXT: store double* %a, double** %4, align 8
+; CHECK-NEXT: %8 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs, i64 0, i64 1
+; CHECK-NEXT: %9 = bitcast i8** %8 to i32**
+; CHECK-NEXT: store i32* %size.addr, i32** %9, align 8
+; CHECK-NEXT: %3 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs, i64 0, i64 0
+; CHECK-NEXT: %conv = zext i32 %size to i64
+; CHECK-NEXT: %0 = shl nuw nsw i64 %conv, 3
+; CHECK-NEXT: %5 = getelementptr inbounds [2 x i64], [2 x i64]* %.offload_sizes, i64 0, i64 0
+; CHECK-NEXT: store i64 %0, i64* %5, align 8
+; CHECK-NEXT: %10 = getelementptr inbounds [2 x i64], [2 x i64]* %.offload_sizes, i64 0, i64 1
+; CHECK-NEXT: store i64 4, i64* %10, align 8
+
+;int heavyComputation(double* restrict a, unsigned size) {
+;  int random = rand() % 7;
+;
+;  //#pragma omp target data map(a[0:size], size)
+;  void* args[2];
+;  args[0] = &a;
+;  args[1] = &size;
+;  __tgt_target_data_begin(..., args, ...)
+;
+;  #pragma omp target teams
+;  for (int i = 0; i < size; ++i) {
+;    a[i] = ++a[i] * 3.141624;
+;  }
+;
+;  return random;
+;}
+define dso_local i32 @heavyComputation(double* noalias %a, i32 %size) {
+entry:
+  %size.addr = alloca i32, align 4
+  %.offload_baseptrs = alloca [2 x i8*], align 8
+  %.offload_ptrs = alloca [2 x i8*], align 8
+  %.offload_sizes = alloca [2 x i64], align 8
+  %.offload_baseptrs2 = alloca [2 x i8*], align 8
+  %.offload_ptrs3 = alloca [2 x i8*], align 8
+  store i32 %size, i32* %size.addr, align 4
+
+  %call = tail call i32 (...) @rand()
+
+  %conv = zext i32 %size to i64
+  %0 = shl nuw nsw i64 %conv, 3
+  %1 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs, i64 0, i64 0
+  %2 = bitcast [2 x i8*]* %.offload_baseptrs to double**
+  store double* %a, double** %2, align 8
+  %3 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs, i64 0, i64 0
+  %4 = bitcast [2 x i8*]* %.offload_ptrs to double**
+  store double* %a, double** %4, align 8
+  %5 = getelementptr inbounds [2 x i64], [2 x i64]* %.offload_sizes, i64 0, i64 0
+  store i64 %0, i64* %5, align 8
+  %6 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs, i64 0, i64 1
+  %7 = bitcast i8** %6 to i32**
+  store i32* %size.addr, i32** %7, align 8
+  %8 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs, i64 0, i64 1
+  %9 = bitcast i8** %8 to i32**
+  store i32* %size.addr, i32** %9, align 8
+  %10 = getelementptr inbounds [2 x i64], [2 x i64]* %.offload_sizes, i64 0, i64 1
+  store i64 4, i64* %10, align 8
+  call void @__tgt_target_data_begin_mapper(i64 -1, i32 2, i8** nonnull %1, i8** nonnull %3, i64* nonnull %5, i64* getelementptr inbounds ([2 x i64], [2 x i64]* @.offload_maptypes., i64 0, i64 0), i8** null)
+
+  %11 = load i32, i32* %size.addr, align 4
+  %size.casted = zext i32 %11 to i64
+  %12 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs2, i64 0, i64 0
+  %13 = bitcast [2 x i8*]* %.offload_baseptrs2 to i64*
+  store i64 %size.casted, i64* %13, align 8
+  %14 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs3, i64 0, i64 0
+  %15 = bitcast [2 x i8*]* %.offload_ptrs3 to i64*
+  store i64 %size.casted, i64* %15, align 8
+  %16 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_baseptrs2, i64 0, i64 1
+  %17 = bitcast i8** %16 to double**
+  store double* %a, double** %17, align 8
+  %18 = getelementptr inbounds [2 x i8*], [2 x i8*]* %.offload_ptrs3, i64 0, i64 1
+  %19 = bitcast i8** %18 to double**
+  store double* %a, double** %19, align 8
+
+  %20 = call i32 @__tgt_target_teams_mapper(i64 -1, i8* nonnull @.__omp_offloading_heavyComputation.region_id, i32 2, i8** nonnull %12, i8** nonnull %14, i64* getelementptr inbounds ([2 x i64], [2 x i64]* @.offload_sizes., i64 0, i64 0), i64* getelementptr inbounds ([2 x i64], [2 x i64]* @.offload_maptypes.2, i64 0, i64 0), i8** null, i32 0, i32 0)
+  %.not = icmp eq i32 %20, 0
+  br i1 %.not, label %omp_offload.cont, label %omp_offload.failed
+
+omp_offload.failed:                               ; preds = %entry
+  call void @heavyComputationFallBack(i64 %size.casted, double* %a)
+  br label %omp_offload.cont
+
+omp_offload.cont:                                 ; preds = %omp_offload.failed, %entry
+  %rem = srem i32 %call, 7
+  call void @__tgt_target_data_end_mapper(i64 -1, i32 2, i8** nonnull %1, i8** nonnull %3, i64* nonnull %5, i64* getelementptr inbounds ([2 x i64], [2 x i64]* @.offload_maptypes., i64 0, i64 0), i8** null)
+  ret i32 %rem
+}
+
+define internal void @heavyComputationFallBack(i64 %size, double* %a) {
+entry:
+  ; Fallback for offloading function heavyComputation3.
+  ret void
+}
+
+declare void @__tgt_target_data_begin_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
+declare i32 @__tgt_target_teams_mapper(i64, i8*, i32, i8**, i8**, i64*, i64*, i8**, i32, i32)
+declare void @__tgt_target_data_end_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
+
+declare dso_local i32 @rand(...)