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

[hamax97]

Summary

getValuesInOffloadArrays goes through the offload arrays in __tgt_target_data_begin_mapper getting the values stored in them before the call is issued.

call void @__tgt_target_data_begin_mapper(arg0, arg1,
    i8** %offload_baseptrs, i8** %offload_ptrs, i64* %offload_sizes,
...)

PD: I tried hard to keep it small.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

761

Don't you need a reference here? We also need a test, maybe just one that prints the result?

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

761

If I pass an array it's passed as a pointer, isn't it?.

A unittest?

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

761

I see. Make it an ArrayRef<OffloadArrayPtr> instead. That will make it clear and you should not need to modify the call site.

A unit test or a lit test hat looks for some output that we print for debug purposes.

[hamax97]

• Adding regression test.

[sstefan1]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

488

I think it is preferred to have this on the new line. There's at least one more below.

Also I think it is preferred to use the actual type in the loop above.

[hamax97]

NIT: Refactor some variable declarations.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

451

Style: Avoid comments in the same line as code (always, except in calls like ,/* .. */ foo) ). Put them before the declaration please.

llvm/test/Transforms/OpenMP/values_in_offload_arrays.ll

4

you need to require assertions (IIRC) otherwise this tests fails (as non-assert builds do not have debug print)

[hamax97]

• Require comments
• Added REQUIRES: assertions to regression test.

[hamax97]

• NIT changes.
• Tried to split the patch, but then we couldn't test for the values we got.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

517

In line 499 and 500 you compute the Dst already and compare it to Array and here again. Maybe use GetPointerBaseWithConstantOffset instead (above) and pass in the offset. That way you don't traverse the use chain twice with different "methods" and there is no special case for the location 0. You take the array element type and ask for the index to the offset, that should be available via a helper in the struct type, IIRC.

737

Let's keep this simple for now. Just an array of 3 OffloadArray objects. Make the initilize a member function and no need for the extra pointers and stuff.

802

Why the global check here?

847

I somehow feel it would be better to print them by index, so that all information for one item is together, though it's not a strong feeling.

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

517

That function is pretty handy, thank you. What I don't find is the helper for returning the index given the offset.

737

What do you mean by make initialize a member function?. A member of OpenMPOpt?.

802

The %offload_sizes is sometimes a global array. If it is, the idea is to not analyze it.

847

Hahaha. I'm fine with both ways.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

517

I forgot half of it:

Take the datalayout, use it with the struct type to get a StructLayout (const StructLayout *getStructLayout(StructType *Ty) const;).
The struct layout has a method: getElementContainingOffset that does what u want, though you need to verify you are at the beginning of the element with getElementOffset .

737

remove the static keyword basically

802

be convservative now: pick one use case and go with that, a negative test for the other with a fixme is good too.

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

517

I see. But I don't have an StructType. I would have to create it with only one field, that is, the array being allocated. I also found out that the offset calculated by GetPointerBaseWithConstantOffset always increases x8, 0, 8, 16 ..., no matter the type of the array. If I divide that offset by the size of the pointer type (given by DataLayout) then I get the index being accessed.

802

I think that's what it is doing, isn't it?. I only analyze %offload_sizes if it is not a global value. When it is a global value, is because it's a constant array, so we don't have to bother tracking the instructions for setting it up. Also, we have test cases for both. But I don't completely understand what should I do.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

517

oh, it's an array, right. you can use the store size of the element type and divide the offset by it, it should be a multiple and that gives you the index.

802

I guess this is OK. maybe use an early exist though:

if (global)
  return true.

you also need to consider that OAs[2] is later maybe unitialized, but that is for follow up patches.

[hamax97]

• Removing factory function for OffloadArray, now using it as member function.
• Calculating the accessed index of the OffloadArray using the offset returned by GetPointerBaseWithConstantOffset.

[jdoerfert]

LG, I added a bunch of minor change requests, nothing major.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

450

initialize to nullptr

467

Nit: pass Array to getValues, and initialize the member only if that was successful.

505

bail if Idx * PointerSize != Offset.

742

Make this:
LLVM_DEBUG(...)

the test then requires the -debug-only=openmpopt flag

802

return true only if the global is a constant, false otherwise. so "return global->isconstant"

817

call this "printXXX" or "dumpXXX"

[hamax97]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

505

Why would that happen?. I tried getting the size of the type stored, the problem is that even for %offload_sizes which is an array of i64 the offset returned is 8 bytes for example, not 64 bytes. That's why I assumed it to be always divided by the size of a pointer instead of the size of the type stored.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

505

it should not happen but it can it the pointer is casted before the store.

[hamax97]

rG8931add61705 [OpenMPOpt][HideMemTransfersLatency] Get values stored in offload arrays.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

 #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"
 
 using namespace llvm;
 using namespace omp;
 
 #define DEBUG_TYPE "openmp-opt"
 
 static cl::opt<bool> DisableOpenMPOptimizations(
     "openmp-opt-disable", cl::ZeroOrMore,
#include "llvm/Frontend/OpenMP/OMPKinds.def"
 
     // TODO: We should attach the attributes defined in OMPKinds.def.
   }
 
   /// Collection of known kernels (\see Kernel) in the module.
   SmallPtrSetImpl<Kernel> &Kernels;
 };
 
+/// Used to map the values physically (in the IR) stored in an offload
+/// array, to a vector in memory.
+struct OffloadArray {
+  /// Physical array (in the IR).
+  AllocaInst *Array;
+  /// Mapped values.
+  SmallVector<Value *, 8> StoredValues;
+  /// Last stores made in the offload array.
+  SmallVector<StoreInst *, 8> LastAccesses;
+
+  OffloadArray() = default;
+
+  /// Initializes the OffloadArray with the values stored in \p Array before
+  /// instruction \p Before is reached. Returns false if the initialization
+  /// fails.
+  /// This MUST be used immediately after construction the object.
+  bool initialize(AllocaInst &Array, Instruction &Before) {
+    if (!Array.getAllocatedType()->isArrayTy())
+      return false;
+
+    this->Array = &Array;
+    if (!getValues(Before))
+      return false;
+
+    return true;
+  }
+
+private:
+  /// Traverses the BasicBlock where Array is, collecting the stores made to
+  /// Array, leaving StoredValues with the values stored before the instruction
+  /// \p Before is reached.
+  bool getValues(Instruction &Before) {
+    // Initialize container.
+    const uint64_t NumValues =
+        Array->getAllocatedType()->getArrayNumElements();
+    StoredValues.assign(NumValues, nullptr);
+    LastAccesses.assign(NumValues, nullptr);
+
+    // TODO: This assumes the instruction \p Before is in the same
+    //  BasicBlock as Array. Make it general, for any control flow graph.
+    BasicBlock *BB = Array->getParent();
+    if (BB != Before.getParent())
+      return false;
+
+    const DataLayout &DL = Array->getModule()->getDataLayout();
+    const unsigned int PointerSize = DL.getPointerSize();
+
+    for (Instruction &I : *BB) {
+      if (&I == &Before)
+        break;
+
+      if (!isa<StoreInst>(&I))
+        continue;
+
+      auto *S = cast<StoreInst>(&I);
+      int64_t Offset = -1;
+      auto *Dst = GetPointerBaseWithConstantOffset(S->getPointerOperand(),
+                                                   Offset, DL);
+      if (Dst == Array) {
+        int64_t Idx = Offset / PointerSize;
+        StoredValues[Idx] = getUnderlyingObject(S->getValueOperand());
+        LastAccesses[Idx] = S;
+      }
+    }
+
+    return isFilled();
+  }
+
+  /// Returns true if all values in StoredValues and
+  /// LastAccesses are not nullptrs.
+  bool isFilled() {
+    const unsigned NumValues = StoredValues.size();
+    for (unsigned I = 0; I < NumValues; ++I) {
+      if (!StoredValues[I] || !LastAccesses[I])
+        return false;
+    }
+
+    return true;
+  }
+};
+
 struct OpenMPOpt {
 
   using OptimizationRemarkGetter =
       function_ref<OptimizationRemarkEmitter &(Function *)>;
 
   OpenMPOpt(SmallVectorImpl<Function *> &SCC, CallGraphUpdater &CGUpdater,
             OptimizationRemarkGetter OREGetter,
             OMPInformationCache &OMPInfoCache, Attributor &A)
private:
   bool hideMemTransfersLatency() {
     auto &RFI = OMPInfoCache.RFIs[OMPRTL___tgt_target_data_begin_mapper];
     bool Changed = false;
     auto SplitMemTransfers = [&](Use &U, Function &Decl) {
       auto *RTCall = getCallIfRegularCall(U, &RFI);
       if (!RTCall)
         return false;
 
+      OffloadArray OffloadArrays[3];
+      if (!getValuesInOffloadArrays(*RTCall, OffloadArrays))
+        return false;
+
+#ifndef NDEBUG
+      debugValuesInOffloadArrays(OffloadArrays);
+#endif
+
       // 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;
   }
 
+  /// Maps the values stored in the offload arrays passed as arguments to
+  /// \p RuntimeCall into the offload arrays in \p OAs.
+  bool getValuesInOffloadArrays(CallInst &RuntimeCall,
+                                MutableArrayRef<OffloadArray> OAs) {
+    assert(OAs.size() == 3 && "Need space for three offload arrays!");
+
+    // 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);
+    // i8** %offload_ptrs.
+    const unsigned PtrsArgNum = 3;
+    Value *PtrsArg = RuntimeCall.getArgOperand(PtrsArgNum);
+    // i8** %offload_sizes.
+    const unsigned SizesArgNum = 4;
+    Value *SizesArg = RuntimeCall.getArgOperand(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 global array don't analyze it.
+    if (isa<GlobalValue>(V))
+      return true;
+    if (!isa<AllocaInst>(V))
+      return false;
+
+    auto *SizesArray = cast<AllocaInst>(V);
+    if (!OAs[2].initialize(*SizesArray, RuntimeCall))
+      return false;
+
+    return true;
+  }
+
+  /// Prints the values in the OffloadArrays \p OAs using LLVM_DEBUG.
+  /// For now this is a way to test that the function getValuesInOffloadArrays
+  /// is working properly.
+  /// TODO: Move this to a unittest when unittests are available for OpenMPOpt.
+  void debugValuesInOffloadArrays(ArrayRef<OffloadArray> OAs) {
+    assert(OAs.size() == 3 && "There are three offload arrays to debug!");
+
+    LLVM_DEBUG(dbgs() << TAG << " Successfully got offload values:\n");
+    std::string ValuesStr;
+    raw_string_ostream Printer(ValuesStr);
+    std::string Separator = " --- ";
+
+    for (auto *BP : OAs[0].StoredValues) {
+      BP->print(Printer);
+      Printer << Separator;
+    }
+    LLVM_DEBUG(dbgs() << "\t\toffload_baseptrs: " << Printer.str() << "\n");
+    ValuesStr.clear();
+
+    for (auto *P : OAs[1].StoredValues) {
+      P->print(Printer);
+      Printer << Separator;
+    }
+    LLVM_DEBUG(dbgs() << "\t\toffload_ptrs: " << Printer.str() << "\n");
+    ValuesStr.clear();
+
+    for (auto *S : OAs[2].StoredValues) {
+      S->print(Printer);
+      Printer << Separator;
+    }
+    LLVM_DEBUG(dbgs() << "\t\toffload_sizes: " << Printer.str() << "\n");
+  }
+
   /// 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())) {
 

llvm/test/Transforms/OpenMP/values_in_offload_arrays.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"
+
+@.__omp_offloading_heavyComputation.region_id = weak constant i8 0
+@.offload_maptypes. = private unnamed_addr constant [2 x i64] [i64 35, i64 35]
+
+; CHECK-LABEL: {{[^@]+}}Successfully got offload values:
+; CHECK-NEXT: offload_baseptrs: double* %a ---   %size.addr = alloca i32, align 4 ---
+; CHECK-NEXT: offload_ptrs: double* %a ---   %size.addr = alloca i32, align 4 ---
+; CHECK-NEXT: offload_sizes:   %0 = shl nuw nsw i64 %conv, 3 --- i64 4 ---
+
+;int heavyComputation(double* 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* %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
+
+  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)
+  %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
+}
+
+declare void @__tgt_target_data_begin_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
+declare void @__tgt_target_data_end_mapper(i64, i32, i8**, i8**, i64*, i64*, i8**)
+
+declare dso_local i32 @rand(...)