[OpenMPOpt][SplitMemTransfer][WIP] Getting values stored in offload arrays

[hamax97]

Summary

Initial approach to get the values stored in the offloading arrays passed to the runtime calls.

The main idea is to split the runtime calls that involve host to device memory offloading into an "issue" and "wait" runtime calls. The "issue" is the asynchronous version of the original runtime call, but now it returns a handle. That handle is used by "wait" to wait for the transfer to complete. The objective is to forward move the "issue" as much as possible, issuing the memory transfer and continuing with computation that does not need it. The "wait" is moved downwards as much as possible, until another runtime call uses that data previously transferred. Hence, trying to use the processor in other stuff while the transfer is being made, hopefully not having to wait for the transfer and then continuing with the computation.

[sstefan1]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

573

Is this supposed to return anything else other that EXIT_FAILURE or EXIT_SUCCESS? If not, why not just use bool?

Same for getValuesInOfflArray()?

903

I think you should be able to use getAnalysisResultForFunction() from OMPInformationCache since analysis getters are available in Attributor's InformationCache. (see Attributor.h)

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

573

I thought this would be more understandable. Also, it allows to return some other code and probably make a different decision. But I don't mind changing it

903

Ohh thanks. I'll use that then.

[hamax97]

Changing from returning status code to bool + Getting MemorySSAAnalysis through OMPInfoCache

[sstefan1]

One more thing, looks like you included only the last change, not the complete diff.

[hamax97]

Including all diff

[jdoerfert]

As discussed on the call, we need a unit test to verify this is (now and in the future) doing what we expect it to :)

[hamax97]

[WIP]

• Creating the first unit test for OpenMPOpt. This creation led me to move the declaration of OpenMPOpt and OMPInformationCache to OpenMPOpt.h.
• I propose refactoring the optimizations to separate structures so that they are easily maintainable and testable.

[hamax97]

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

285

IMHO it would nice to have the optimizations outside OpenMPOpt. It's easier testing them like this, and a good practice according to Googletest, Testing private code.
It's also easier adding new features to them.

llvm/unittests/Transforms/IPO/OpenMPOpt/HideMemTransferLatency.cpp

13 ↗

(On Diff #275514)

Here are the unittests for this specific optimization.

23 ↗

(On Diff #275514)

I don't know if this is a good approach. The uses of the runtime call are first stored in the MemTransferCalls, and later used in the test fixture. I could've included this in the test itself, but that would mean writing it over and over again.

llvm/unittests/Transforms/IPO/OpenMPOpt/OpenMPOptTest.h

22

This can be the base class for testing each separate optimization. The CallGraphSCC can be moved as a class attribute, it's just that I don't use it for now.

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

548

This is where the structures that handles the optimization would be used.

[sstefan1]

I think you should split this in 2 patches. One being the refactoring. (this happened to me with the ICV patch)

But lets hear from Johannes as well.

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

285

I'm ok with having this here. I guess it would be good to add some documentation comments in the .h files.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

241–422

Not sure we need these 'end of' comments.

[jdoerfert]

Can you extract the NFC code movement changes. It is impossible to determine what is new here and what is just moved/rewritten.

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

285

I'm not convinced that we need to expose an optimization that runs as part of OpenMPOpt (only) to the world.

Why do we need all this code in the (public) header? If you really want to separate it, I guess a private header in the lib folder would suffice, right? Or do you expect other code to interact with this?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

12

Yay! :)

243

I don't think they help

[hamax97]

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

285

Well, the idea is to make this easily testable. So, if having a private header, we will need to include it in the unittests. If that is not a problem, then it is a better idea having the private header. What do you think then?

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

241–422

Sure, I'll remove them.

[hamax97]

• Exposing getValuesInOfflArrays() to make it unit testable.

[hamax97]

Finished getting values in offload arrays

[jdoerfert]

Some style and minor remarks together with some proposals.

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

150

Do we expect to call this with a nullptr? If not, we should make it a reference instead.

152

Nit: I guess adding the missing characters of Offload will only help ;).

StoresTo or do you mean stores into pointers that are in the offload array?

175

We need documentation for all the members and methods. The comment on the struct should also be expanded, e.g., list all the thing we store in here.

268

Comment broken?

295

Feel free to elaborate here a bit what is going to happen. Or in the file comment.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

231

Just commit this change.

249

Style: Maybe put a comment in the beginning explaining how the function declaration looks. I also prefer comments in their own lines, generally as full sentences.

253

Nit: Write the types here, usually helps.

292

Are we OK if any of these is not an AllocaInst? Or should we also bail in that situations? Do we have a test for the problematic cases?

298

You need to verify before that the type is an array type.

302

Maybe merge initialize in this method and call this method initialize. Or, take the runtime call do it all in the constructor. Afterwards you can ask if the object is "valid".

312

I guess for now we could assume the stores are in the same block as Before? We could iterate from Before to the basic block begin, WDYT?

314

We should never rely on the order of use lists or other things. We should also never modify the order ;)

Assume users are in any order, there is no guarantee otherwise anyway.

317

Unexpected users should be reason to give up.

759

You clobber the outer Changed variable here. So even if it was true at some point, you might return false. You should also be careful because the foreachUse callback expects you to return true only if the use was deleted.

[hamax97]

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

150

The idea is to make that lower bound in the function optional. So, if it's not given, then the values returned would be the last values stored. This, assuming that there might be multiple stores to the offload array.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

292

We can handle ArrayTypes too, I'll fix that.
We don't have a test for non alloca arrays?. Do we need it if we handle ArrayTypes too?

298

The problem with isArrayAllocation() is that it returns false if the array has size 1.

302

Sure. That sounds nicer.

312

Yes we can assume they are in the same BasicBlock. But why would we want to iterate from Before's BasicBlock to the entry BasicBlock?, I don't fully understand your point.

314

Mmmmm. The problem would be knowing when to stop. The idea here is to go from top to bottom and stop if we go over Before. But if there's no sense of direction how would you know how to stop?. I thought using MemorySSA, we can traverse upwards the definition chain starting in the call site.

[jdoerfert]

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

150

I see. Could we make it less generic and more specific for now. This is by nature a complex thing and I would prefer the first version to be rather targeted.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

292

I would suggest to *only* handle what clang generates right now. And to test for it specifically. No extra uses, no weird types, nothing.

298

There are two levels of "array" here. AllocaInst can allocate an array of the given type and the given type can be an array. check for what you expect (I think the latter).

312

That was just a suggestion. Any (simple) way to identify what we need to is fine with me.

314

Oh, I see. That is where my basic block traversal comes in. If you want to know if it is "before Before", actually verify that, e.g., by traversing the basic block from Before to the beginning. All the instructions you see are "before Before. The use list is not order in any specific way.

[ye-luo]

Could you describe what "SplitMemTransfer" is? The current patch summary doesn't provide sufficient explanation. What kind of offloading arrays is considered by this optimization?

[hamax97]

• Refactors discussed made.
• OpenMPOpt unit tests infrastructure merged.

[JonChesterfield]

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

255

This looks wrong. If V is always a ConstantInt, we want cast. If it isn't, then dyn_cast will return null and this will crash on the dereference.

285

If stuff needs to be a header in order to poke at it from unit tests then so be it.

OpenMPOpt.h is the header developers will look at to see how to invoke the optimisation, so it's better to keep implementation cruft out of it.

So a second header is probably the right compromise.

[hamax97]

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

255

You're right. I'll fix it. It seems weird though, what should I return then?. Probably not having this function is a better idea

285

Do you mean a private header inside lib/ or a public one inside include/?.

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

 #include "llvm/Transforms/IPO/Attributor.h"
 #include "llvm/Transforms/Utils/CallGraphUpdater.h"
 #include "llvm/Analysis/CGSCCPassManager.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/Analysis/MemorySSA.h"
 
 namespace llvm {
-
 namespace omp {
 
 using namespace types;
 
 /// OpenMP specific information. For now, stores RFIs and ICVs also needed for
 /// Attributor runs.
 struct OMPInformationCache : public InformationCache {
   OMPInformationCache(Module &M, AnalysisGetter &AG,
void foreachUse(function_ref<bool(Use &, Function &)> CB, Function *F,
                     UseVector *Uses = nullptr);
 
   private:
     /// Map from functions to all uses of this runtime function contained in
     /// them.
     DenseMap<Function *, std::shared_ptr<UseVector>> UsesMap;
   };
 
+  /// Used to store information about a runtime call that involves
+  /// host to device memory offloading. For example:
+  /// __tgt_target_data_begin(...,
+  ///   i8** %offload_baseptrs, i8** %offload_ptrs, i64* %offload_sizes,
+  /// ...)
+  struct MemoryTransfer {
+
+    /// Used to map the values physically (in the IR) stored in an offload
+    /// array, to a vector in memory.
+    struct OffloadArray {
+      AllocaInst &Array; /// Physical array (in the IR).
+      SmallVector<Value *, 8> StoredValues; /// Mapped values.
+      InformationCache &InfoCache;
+
+      /// Factory function for creating and initializing the OffloadArray with
+      /// the values stored in \p Array before the instruction \p Before is
+      /// reached.
+      /// This MUST be used instead of the constructor.
+      static std::unique_ptr<OffloadArray> initialize(
+          AllocaInst &Array,
+          Instruction &Before,
+          InformationCache &InfoCache);
+
+      /// Use the factory function initialize(...) instead.
+      OffloadArray(AllocaInst &Array, InformationCache &InfoCache)
+          : Array(Array), InfoCache(InfoCache) {}
+
+    private:
+      /// Traverses the BasicBlocks collecting the stores made to
+      /// OffloadArray::Array, leaving OffloadArray::StoredValues with the
+      /// values stored before the instruction \p Before is reached.
+      bool getValues(Instruction &Before);
+
+      /// Returns the index of OffloadArray::Array where the store is being
+      /// made. Returns -1 if the index can't be deduced.
+      int32_t getAccessedIdx(StoreInst &S);
+
+      /// Returns true all values in \p V are not nullptrs.
+      static bool isFilled(const SmallVectorImpl<Value *> &V);
+    };
+
+    CallBase *RuntimeCall; /// Call that involves a memotry transfer.
+    InformationCache &InfoCache;
+
+    /// These help mapping the values in offload_baseptrs, offload_ptrs, and
+    /// offload_sizes, respectively.
+    std::unique_ptr<OffloadArray> BasePtrs;
+    std::unique_ptr<OffloadArray> Ptrs;
+    std::unique_ptr<OffloadArray> Sizes;
+
+    MemoryTransfer(CallBase *RuntimeCall, InformationCache &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, MemoryTransfer::BasePtrs, MemoryTransfer::Ptrs,
+    /// MemoryTransfer::Sizes.
+    /// Returns false if one of the arrays couldn't be processed or some of the
+    /// values couldn't be found.
+    bool getValuesInOffloadArrays();
+  };
+
   /// The slice of the module we are allowed to look at.
   SmallPtrSetImpl<Function *> &ModuleSlice;
 
   /// An OpenMP-IR-Builder instance
   OpenMPIRBuilder OMPBuilder;
 
   /// Map from runtime function kind to the runtime function description.
   EnumeratedArray<RuntimeFunctionInfo, RuntimeFunction,
struct OMPInformationCache : public InformationCache {
   /// function types, that is, return type \p RTFRetType, and argument types
   /// \p RTFArgTypes.
   static bool declMatchesRTFTypes(Function *F, Type *RTFRetType,
                                   SmallVector<Type *, 8> &RTFArgTypes);
 };
 
 struct OpenMPOpt {
 
+  using MemoryTransfer = OMPInformationCache::MemoryTransfer;
   using OptimizationRemarkGetter =
   function_ref<OptimizationRemarkEmitter &(Function *)>;
 
   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
             OptimizationRemarkGetter OREGetter,
             OMPInformationCache &OMPInfoCache, Attributor &A)
       : M(*(*SCC.begin())->getParent()), SCC(SCC), CGUpdater(CGUpdater),
         OREGetter(OREGetter), OMPInfoCache(OMPInfoCache), A(A) {}
 
   /// Run all OpenMP optimizations on the underlying SCC/ModuleSlice.
   bool run();
 
   /// Return the call if \p U is a callee use in a regular call. If \p RFI is
   /// given it has to be the callee or a nullptr is returned.
   static CallInst *getCallIfRegularCall(
       Use &U, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr);
 
   /// Return the call if \p V is a regular call. If \p RFI is given it has to be
   /// the callee or a nullptr is returned.
   static CallInst *getCallIfRegularCall(
       Value &V, OMPInformationCache::RuntimeFunctionInfo *RFI = nullptr);
 
+  /// Returns the integer representation of \p V.
+  static int64_t getIntLiteral(const Value *V) {
+    assert(V && "Getting Integer value of nullptr");
+    return (dyn_cast<ConstantInt>(V))->getZExtValue();
+  }
+
 private:
   /// Try to delete parallel regions if possible.
   bool deleteParallelRegions();
 
   /// Try to eliminiate runtime calls by reusing existing ones.
   bool deduplicateRuntimeCalls();
 
+  /// Tries to hide the latency of runtime calls that involve host to
+  /// device memory transfers by splitting them into their "issue" and "wait".
+  /// versions. The "issue" is moved upwards as much as possible. The "wait" is
+  /// moved downards as much as possible. The "issue" issues the memory transfer
+  /// asynchronously, returning a handle. The "wait" waits in the returned
+  /// handle for the memory transfer to finish.
+  bool hideMemTransfersLatency();
+
   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.
   Value *
   getCombinedIdentFromCallUsesIn(OMPInformationCache::RuntimeFunctionInfo &RFI,
                                  Function &F, bool GlobalOnly);
 
   /// Try to eliminiate calls of \p RFI in \p F by reusing an existing one or
   /// \p ReplVal if given.
   bool deduplicateRuntimeCalls(Function &F,
                                OMPInformationCache::RuntimeFunctionInfo &RFI,
                                Value *ReplVal = nullptr);
 
   /// Collect arguments that represent the global thread id in \p GTIdArgs.
   void collectGlobalThreadIdArguments(SmallSetVector<Value *, 16> &GTIdArgs);
 
   /// Emit a remark generically
   ///
   /// This template function can be used to generically emit a remark. The
   /// RemarkKind should be one of the following:
   ///   - OptimizationRemark to indicate a successful optimization attempt
   ///   - OptimizationRemarkMissed to report a failed optimization attempt
   ///   - OptimizationRemarkAnalysis to provide additional information about an
   ///     optimization attempt

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 //===-- IPO/OpenMPOpt.cpp - Collection of OpenMP specific optimizations ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // OpenMP specific optimizations:
 //
 // - Deduplication of runtime calls, e.g., omp_get_thread_num.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/OpenMPOpt.h"
 
 #include "llvm/ADT/EnumeratedArray.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/CallGraphSCCPass.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #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;
 using namespace types;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(

 };
 
 } // end namespace llvm
 
 //===----------------------------------------------------------------------===//
 // Definitions of the OMPInformationCache helper structure.
 //===----------------------------------------------------------------------===//
 
+using MemoryTransfer = OMPInformationCache::MemoryTransfer;
+using OffloadArray = MemoryTransfer::OffloadArray;
+
 void OMPInformationCache::RuntimeFunctionInfo::foreachUse(
     function_ref<bool(Use &, Function &)> CB, Function *F, UseVector *Uses) {
   SmallVector<unsigned, 8> ToBeDeleted;
   ToBeDeleted.clear();
 
   unsigned Idx = 0;
   UseVector &UV = Uses ? *Uses : getOrCreateUseVector(F);
 
bool OMPInformationCache::declMatchesRTFTypes(
 
   if (!F)
     return false;
   if (F->getReturnType() != RTFRetType)
     return false;
   if (F->arg_size() != RTFArgTypes.size())
     return false;
 
   auto RTFTyIt = RTFArgTypes.begin();
   for (Argument &Arg : F->args()) {
     if (Arg.getType() != *RTFTyIt)
       return false;
 
     ++RTFTyIt;
   }
 
   return true;
 }
 
 //===----------------------------------------------------------------------===//
+// Definitions of the MemoryTransfer helper structure.
+//===----------------------------------------------------------------------===//
+
+bool MemoryTransfer::getValuesInOffloadArrays() {
+  // A runtime call that involves memory offloading looks something like:
+  // call void @__tgt_target_data_begin(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;
+  Use *BasePtrsArg = RuntimeCall->arg_begin() + BasePtrsArgNum;
+  // i8** %offload_ptrs.
+  const unsigned PtrsArgNum = 3;
+  Use *PtrsArg = RuntimeCall->arg_begin() + PtrsArgNum;
+  // i8** %offload_sizes.
+  const unsigned SizesArgNum = 4;
+  Use *SizesArg = RuntimeCall->arg_begin() + SizesArgNum;
+
+  const DataLayout &DL = InfoCache.getDL();
+
+  // Get values stored in **offload_baseptrs.
+  auto *V = GetUnderlyingObject(BasePtrsArg->get(), DL);
+  if (!isa<AllocaInst>(V)) {
+    LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_baseptrs, only"
+                      << " alloca arrays supported. In call to "
+                      << RuntimeCall->getName() << " in function "
+                      << RuntimeCall->getCaller()->getName() << "\n");
+    return false;
+  }
+
+  auto *Array = cast<AllocaInst>(V);
+  BasePtrs = OffloadArray::initialize(*Array, *RuntimeCall, InfoCache);
+  if (!BasePtrs) {
+    LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_baseptrs in call to "
+                      << RuntimeCall->getName() << " in function "
+                      << RuntimeCall->getCaller()->getName() << "\n");
+    return false;
+  }
+
+  // Get values stored in **offload_ptrs.
+  V = GetUnderlyingObject(PtrsArg->get(), DL);
+  if (!isa<AllocaInst>(V)) {
+    LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_ptrs, only"
+                      << " alloca arrays supported. In call to "
+                      << RuntimeCall->getName() << " in function "
+                      << RuntimeCall->getCaller()->getName() << "\n");
+    return false;
+  }
+  Array = cast<AllocaInst>(V);
+  Ptrs = OffloadArray::initialize(*Array, *RuntimeCall, InfoCache);
+  if (!Ptrs) {
+    LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_ptrs in call to "
+                      << RuntimeCall->getName() << " in function "
+                      << RuntimeCall->getCaller()->getName() << "\n");
+    return false;
+  }
+
+  // Get values stored in **offload_sizes.
+  V = GetUnderlyingObject(SizesArg->get(), DL);
+  // Sometimes the frontend generates this array as a constant global array.
+  if (!isa<GlobalValue>(V)) {
+    if (!isa<AllocaInst>(V)) {
+      LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_sizes, only"
+                        << " alloca arrays supported. In call to "
+                        << RuntimeCall->getName() << " in function "
+                        << RuntimeCall->getCaller()->getName() << "\n");
+      return false;
+    }
+
+    Array = cast<AllocaInst>(V);
+    Sizes = OffloadArray::initialize(*Array, *RuntimeCall, InfoCache);
+    if (!Sizes) {
+      LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload_sizes in call to "
+                        << RuntimeCall->getName() << " in function "
+                        << RuntimeCall->getCaller()->getName() << "\n");
+      return false;
+    }
+  }
+
+  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);
+  bool Success = OA->getValues(Before);
+  if (!Success) {
+    LLVM_DEBUG(dbgs() << TAG << "Error getting values in offload array.\n");
+    return nullptr;
+  }
+
+  return OA;
+}
+
+bool OffloadArray::getValues(Instruction &Before) {
+  // Initialize container.
+  const uint64_t NumValues =
+      Array.getAllocatedType()->getArrayNumElements();
+  StoredValues.assign(NumValues, nullptr);
+
+  // TODO: This assumes the instruction \p Before is in the same BasicBlock
+  //       as OffloadArray::Array. Make it general, for any control flow graph.
+  auto *BB = Array.getParent();
+  if (BB != Before.getParent()) {
+    LLVM_DEBUG(dbgs() << TAG << "The lower limit instruction is in a"
+                      << " different BasicBlock.\n");
+    return false;
+  }
+
+  const DataLayout &DL = InfoCache.getDL();
+  for (auto &I : *BB) {
+    if (&I == &Before) break;
+
+    if (isa<StoreInst>(&I)) {
+      auto *Dst = GetUnderlyingObject(I.getOperand(1), DL);
+
+      if (Dst == &Array) {
+        int32_t AccessedIdx = getAccessedIdx(*cast<StoreInst>(&I));
+        if (AccessedIdx < 0) {
+          LLVM_DEBUG(dbgs() << TAG << "Unexpected StoreInst\n");
+          return false;
+        }
+
+        StoredValues[AccessedIdx] = GetUnderlyingObject(I.getOperand(0), DL);
+      }
+    }
+  }
+
+  return isFilled(StoredValues);
+}
+
+int32_t OffloadArray::getAccessedIdx(StoreInst &S) {
+  auto *Dst = S.getOperand(1);
+  if (!isa<Instruction>(Dst)) {
+    LLVM_DEBUG(dbgs() << TAG << "Unrecognized store pattern.\n");
+    return -1;
+  }
+  auto *DstInst = cast<Instruction>(Dst);
+
+  Value *Access = DstInst;
+  if (DstInst->isCast()) {
+    Access = DstInst->getOperand(0);
+
+    // Direct cast from the AllocaInst, which means a store to the
+    // first position of the array.
+    if (Access == &Array) return 0;
+  }
+
+  if (!isa<GetElementPtrInst>(Access)) {
+    LLVM_DEBUG(dbgs() << TAG << "Unrecognized store pattern.\n");
+    return -1;
+  }
+  auto *GEPInst = cast<GetElementPtrInst>(Access);
+
+  auto *ArrayIdx = GEPInst->idx_begin() + 1;
+  if (ArrayIdx == GEPInst->idx_end()) {
+    LLVM_DEBUG(dbgs() << TAG << "Unrecognized store pattern.\n");
+    return -1;
+  }
+
+  return OpenMPOpt::getIntLiteral(ArrayIdx->get());
+}
+
+bool OffloadArray::isFilled(const SmallVectorImpl<Value *> &V) {
+  for (auto *E : V)
+    if (!E)
+      return false;
+
+  return true;
+}
+
+//===----------------------------------------------------------------------===//
 // Declarations and definitions of AAICVTracker.
 //===----------------------------------------------------------------------===//
 namespace {
 
 /// Abstract Attribute for tracking ICV values.
 struct AAICVTracker : public StateWrapper<BooleanState, AbstractAttribute> {
   using Base = StateWrapper<BooleanState, AbstractAttribute>;
   AAICVTracker(const IRPosition &IRP, Attributor &A) : Base(IRP) {}
Value *getReplacementValue(InternalControlVar ICV, const Instruction *I,
     const BasicBlock *CurrBB = I->getParent();
 
     auto &ValuesSet = ICVValuesMap[ICV];
     auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
     auto &GetterRFI = OMPInfoCache.RFIs[OMPInfoCache.ICVs[ICV].Getter];
 
     for (const auto &ICVVal : ValuesSet) {
       if (CurrBB == ICVVal.Inst->getParent()) {
         if (!ICVVal.Inst->comesBefore(I))
           continue;
 
         // both instructions are in the same BB and at \p I we know the ICV
         // value.
         while (I != ICVVal.Inst) {
           // we don't yet know if a call might update an ICV.
           // TODO: check callsite AA for value.
           if (const auto *CB = dyn_cast<CallBase>(I))
             if (CB->getCalledFunction() != GetterRFI.Declaration)
               return nullptr;
 
           I = I->getPrevNode();
         }
 
         // No call in between, return the value.
         return ICVVal.TrackedValue;
       }
     }
 
     // No value was tracked.
     return nullptr;
   }
 };
 } // namespace
 
 const char AAICVTracker::ID = 0;
 
 AAICVTracker &AAICVTracker::createForPosition(const IRPosition &IRP,
                                               Attributor &A) {
   AAICVTracker *AA = nullptr;
   switch (IRP.getPositionKind()) {
   case IRPosition::IRP_INVALID:
   case IRPosition::IRP_FLOAT:
for (Function *F : OMPInfoCache.ModuleSlice) {
         emitRemarkOnFunction(F, "OpenMPICVTracker", Remark);
       }
     }
   }
 
   Changed |= runAttributor();
   Changed |= deduplicateRuntimeCalls();
   Changed |= deleteParallelRegions();
+  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() &&
for (Argument &Arg : F->args())
       }
     Changed |= deduplicateRuntimeCalls(
         *F, OMPInfoCache.RFIs[OMPRTL___kmpc_global_thread_num], GTIdArg);
   }
 
   return Changed;
 }
 
+bool OpenMPOpt::hideMemTransfersLatency() {
+  OMPInformationCache::RuntimeFunctionInfo &RFI =
+      OMPInfoCache.RFIs[OMPRTL___tgt_target_data_begin];
+
+  bool Changed = false;
+  auto SplitDataTransfer = [&] (Use &U, Function &Decl) {
+    auto *RTCall = getCallIfRegularCall(U, &RFI);
+    if (!RTCall)
+      return false;
+
+    MemoryTransfer MT(RTCall, OMPInfoCache);
+    bool Success = MT.getValuesInOffloadArrays();
+    if (!Success) {
+      LLVM_DEBUG(dbgs() << TAG << "Couldn't get offload arrays in call to "
+                        << MT.RuntimeCall->getName() << " in function "
+                        << MT.RuntimeCall->getCaller()->getName() << "\n");
+      return false;
+    }
+    return false;
+  };
+
+  RFI.foreachUse(SplitDataTransfer);
+  return Changed;
+}
+
 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)) {
auto ReplaceAndDeleteCB = [&](Use &U, Function &Caller) {
     return true;
   };
   RFI.foreachUse(ReplaceAndDeleteCB);
 
   return Changed;
 }
 
 void OpenMPOpt::collectGlobalThreadIdArguments(
     SmallSetVector<Value *, 16> &GTIdArgs) {
   // TODO: Below we basically perform a fixpoint iteration with a pessimistic
   //       initialization. We could define an AbstractAttribute instead and
   //       run the Attributor here once it can be run as an SCC pass.
 
   // Helper to check the argument \p ArgNo at all call sites of \p F for
   // a GTId.
   auto CallArgOpIsGTId = [&](Function &F, unsigned ArgNo, CallInst &RefCI) {
     if (!F.hasLocalLinkage())

llvm/unittests/Transforms/IPO/CMakeLists.txt

 set(LLVM_LINK_COMPONENTS
   Core
   Support
   IPO
   )
 
 add_llvm_unittest(IPOTests
   LowerTypeTests.cpp
   WholeProgramDevirt.cpp
   )
+
+add_subdirectory(OpenMPOpt)

llvm/unittests/Transforms/IPO/OpenMPOpt/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS
+  ipo
+  Core
+  Analysis
+  AsmParser
+  FrontendOpenMP
+  TransformUtils
+  )
+
+add_llvm_unittest(OpenMPOptUnitTests
+  HideMemTransferLatencyTest.cpp
+  )
 No newline at end of file

llvm/unittests/Transforms/IPO/OpenMPOpt/HideMemTransferLatencyTest.cpp

+//===- HideMemTransferLatencyTest.cpp -------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "OpenMPOptTest.h"
+#include "llvm/IR/ValueSymbolTable.h"
+
+namespace {
+
+class HideMemTransferLatencyTest : public OpenMPOptTest {
+protected:
+  NiceMock<MockModulePass> MMP;
+  MockSCCPass MSCCP;
+  ModulePassManager MPM;
+
+  FunctionAnalysisManager FAM;
+  ModuleAnalysisManager MAM;
+  CGSCCAnalysisManager CGAM;
+
+  HideMemTransferLatencyTest()
+      : OpenMPOptTest(), FAM(true), MAM(true), CGAM(true) {
+    MAM.registerPass([] { return LazyCallGraphAnalysis(); });
+    FAM.registerPass([] { return TargetLibraryAnalysis(); });
+    FAM.registerPass([]{ return AAManager(); });
+    FAM.registerPass([]{ return DominatorTreeAnalysis(); });
+    FAM.registerPass([]{ return MemorySSAAnalysis(); });
+    CGAM.registerPass([&] { return FunctionAnalysisManagerCGSCCProxy(); });
+
+    // Register required pass instrumentation analysis.
+    FAM.registerPass([] { return PassInstrumentationAnalysis(); });
+    MAM.registerPass([] { return PassInstrumentationAnalysis(); });
+    CGAM.registerPass([] {return PassInstrumentationAnalysis(); });
+
+    // Cross-register proxies.
+    MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+    MAM.registerPass([&] { return CGSCCAnalysisManagerModuleProxy(CGAM); });
+    CGAM.registerPass([&] { return ModuleAnalysisManagerCGSCCProxy(MAM); });
+    FAM.registerPass([&] { return CGSCCAnalysisManagerFunctionProxy(CGAM); });
+    FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
+
+    CGSCCPassManager CGPM;
+    CGPM.addPass(MSCCP.getPass());
+
+    MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(std::move(CGPM)));
+  }
+};
+
+TEST_F(HideMemTransferLatencyTest, GetValuesInOfflArrays) {
+  const char *ModuleString =
+      "@.offload_maptypes = private unnamed_addr constant [1 x i64] [i64 33]\n"
+      "define dso_local i32 @dataTransferOnly(double* noalias %a, i32 %size) {\n"
+      "entry:\n"
+      "  %.offload_baseptrs = alloca [1 x i8*], align 8\n"
+      "  %.offload_ptrs = alloca [1 x i8*], align 8\n"
+      "  %.offload_sizes = alloca [1 x i64], align 8\n"
+
+      "  %call = call i32 @rand()\n"
+
+      "  %conv = zext i32 %size to i64\n"
+      "  %0 = shl nuw nsw i64 %conv, 3\n"
+      "  %1 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_baseptrs, i64 0, i64 0\n"
+      "  %2 = bitcast [1 x i8*]* %.offload_baseptrs to double**\n"
+      "  store double* %a, double** %2, align 8\n"
+      "  %3 = getelementptr inbounds [1 x i8*], [1 x i8*]* %.offload_ptrs, i64 0, i64 0\n"
+      "  %4 = bitcast [1 x i8*]* %.offload_ptrs to double**\n"
+      "  store double* %a, double** %4, align 8\n"
+      "  %5 = getelementptr inbounds [1 x i64], [1 x i64]* %.offload_sizes, i64 0, i64 0\n"
+      "  store i64 %0, i64* %5, align 8\n"
+      "  call void @__tgt_target_data_begin(i64 -1, i32 1, i8** nonnull %1, i8** nonnull %3, i64* nonnull %5, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes, i64 0, i64 0))\n"
+
+      "  %rem = urem i32 %call, %size\n"
+
+      "  call void @__tgt_target_data_end(i64 -1, i32 1, i8** nonnull %1, i8** nonnull %3, i64* nonnull %5, i64* getelementptr inbounds ([1 x i64], [1 x i64]* @.offload_maptypes, i64 0, i64 0))\n"
+      "  ret i32 %rem\n"
+      "}\n"
+
+      "declare dso_local void @__tgt_target_data_begin(i64, i32, i8**, i8**, i64*, i64*)\n"
+      "declare dso_local void @__tgt_target_data_end(i64, i32, i8**, i8**, i64*, i64*)\n"
+      "declare dso_local i32 @rand()";
+
+  std::unique_ptr<Module> M = parseModuleString(ModuleString);
+
+  EXPECT_CALL(MSCCP, run(_, _, _, _))
+      .WillRepeatedly(Invoke([](LazyCallGraph::SCC &C, CGSCCAnalysisManager &AM,
+                                LazyCallGraph &CG, CGSCCUpdateResult &UR) {
+        SmallPtrSet<Function *, 16> ModuleSlice;
+        SmallVector<Function *, 16> SCC;
+        for (LazyCallGraph::Node &N : C) {
+          SCC.push_back(&N.getFunction());
+          ModuleSlice.insert(SCC.back());
+        }
+
+        EXPECT_FALSE(SCC.empty());
+
+        FunctionAnalysisManager &FAM =
+            AM.getResult<FunctionAnalysisManagerCGSCCProxy>(C, CG).getManager();
+
+        AnalysisGetter AG(FAM);
+        CallGraphUpdater CGUpdater;
+        CGUpdater.initialize(CG, C, AM, UR);
+
+        SetVector<Function *> Functions(SCC.begin(), SCC.end());
+        BumpPtrAllocator Allocator;
+        OMPInformationCache InfoCache(*(Functions.back()->getParent()), AG, Allocator,
+            /*CGSCC*/ &Functions, ModuleSlice);
+        Attributor A(Functions, InfoCache, CGUpdater);
+
+        auto &RFI = InfoCache.RFIs[OMPRTL___tgt_target_data_begin];
+        auto GetValuesInOfflArrays = [&](Use &U, Function &Decl) {
+          auto *RTCall = OpenMPOpt::getCallIfRegularCall(U, &RFI);
+          EXPECT_TRUE(RTCall);
+
+          OMPInformationCache::MemoryTransfer MT(RTCall, InfoCache);
+          bool Success = MT.getValuesInOffloadArrays();
+          EXPECT_TRUE(Success);
+
+          Module &M = *RTCall->getModule();
+          Function &F = *M.getFunction("dataTransferOnly");
+          auto *Arg = F.getArg(0);
+
+          // Check for **offload_baseptrs.
+          auto &BasePtrsValues = MT.BasePtrs->StoredValues;
+          EXPECT_EQ(BasePtrsValues.size(), (size_t) 1);
+          EXPECT_EQ(BasePtrsValues[0], Arg);
+
+          // Check for **offload_ptrs.
+          auto &PtrsValues = MT.Ptrs->StoredValues;
+          EXPECT_EQ(PtrsValues.size(), (size_t) 1);
+          EXPECT_EQ(PtrsValues[0], Arg);
+
+          // Check for **offload_sizes.
+          auto &SizesValues = MT.Sizes->StoredValues;
+          EXPECT_EQ(SizesValues.size(), (size_t) 1);
+          std::string ValueName;
+          raw_string_ostream OS(ValueName);
+          for (auto &I : F.getEntryBlock()) {
+            I.print(OS);
+            if (OS.str() == "  %0 = shl nuw nsw i64 %conv, 3") {
+              EXPECT_EQ(&I, SizesValues[0]);
+              break;
+            }
+            ValueName.clear();
+          }
+
+          return false;
+        };
+        RFI.foreachUse(GetValuesInOfflArrays);
+
+        return PreservedAnalyses::all();
+      }));
+
+  MPM.run(*M, MAM);
+}
+
+} // end anonymous namespace
 No newline at end of file

llvm/unittests/Transforms/IPO/OpenMPOpt/OpenMPOptTest.h

+//===- OpenMPOptTest.h - Base file for OpenMPOpt unittests ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_UNITTESTS_TRANSFORMS_IPO_OPENMPOPT_H
+#define LLVM_UNITTESTS_TRANSFORMS_IPO_OPENMPOPT_H
+
+#include "llvm/Transforms/IPO/OpenMPOpt.h"
+#include "llvm/AsmParser/Parser.h"
+#include "gtest/gtest.h"
+#include "gmock/gmock.h"
+
+using namespace llvm;
+using namespace omp;
+
+namespace {
+
+using testing::NiceMock;
+using testing::Matcher;
+using testing::Return;
+using testing::Invoke;
+using testing::_;
+
+template <typename DerivedT, typename IRUnitT,
+    typename AnalysisManagerT = AnalysisManager<IRUnitT>,
+    typename... ExtraArgTs>
+class MockPassBase {
+public:
+  class Pass : public PassInfoMixin<Pass> {
+    friend MockPassBase;
+
+    DerivedT *Handle;
+
+    Pass(DerivedT &Handle) : Handle(&Handle) {
+      static_assert(std::is_base_of<MockPassBase, DerivedT>::value,
+                    "Must pass the derived type to this template!");
+    }
+
+  public:
+    PreservedAnalyses run(IRUnitT &IR, AnalysisManagerT &AM,
+                          ExtraArgTs... ExtraArgs) {
+      return Handle->run(IR, AM, ExtraArgs...);
+    }
+  };
+
+  Pass getPass() { return Pass(static_cast<DerivedT &>(*this)); }
+
+protected:
+  /// Derived classes should call this in their constructor to set up default
+  /// mock actions. (We can't do this in our constructor because this has to
+  /// run after the DerivedT is constructed.)
+  void setDefaults() {
+    ON_CALL(static_cast<DerivedT &>(*this),
+            run(_, _, testing::Matcher<ExtraArgTs>(_)...))
+        .WillByDefault(Return(PreservedAnalyses::all()));
+  }
+};
+
+struct MockFunctionPass
+    : MockPassBase<MockFunctionPass, Function> {
+  MOCK_METHOD2(run, PreservedAnalyses(Function &, FunctionAnalysisManager &));
+
+  MockFunctionPass() { setDefaults(); }
+};
+
+struct MockModulePass : MockPassBase<MockModulePass, Module> {
+  MOCK_METHOD2(run, PreservedAnalyses(Module &, ModuleAnalysisManager &));
+
+  MockModulePass() { setDefaults(); }
+};
+
+
+struct MockSCCPass : MockPassBase<MockSCCPass, LazyCallGraph::SCC,
+    CGSCCAnalysisManager &, LazyCallGraph &, CGSCCUpdateResult &> {
+  MOCK_METHOD4(run,
+               PreservedAnalyses(LazyCallGraph::SCC &, CGSCCAnalysisManager &,
+                                 LazyCallGraph &, CGSCCUpdateResult &));
+
+};
+
+class OpenMPOptTest : public ::testing::Test {
+protected:
+  std::unique_ptr<LLVMContext> Ctx;
+
+  OpenMPOptTest() : Ctx(new LLVMContext) {}
+
+  std::unique_ptr<Module> parseModuleString(const char *ModuleString) {
+    SMDiagnostic Err;
+    auto M = parseAssemblyString(ModuleString, Err, *Ctx);
+    EXPECT_TRUE(M);
+    return M;
+  }
+};
+
+} // end anonymous namespace
+
+#endif // LLVM_UNITTESTS_TRANSFORMS_IPO_OPENMPOPT_H
 No newline at end of file