[MLIR][OpenMP] Refactoring createTargetData in OMPIRBuilder

Key changes:
  - Refactor the createTargetData function to make use of the emitOffloadingArrays and emitOffloadingArraysArgument functions to generate code.
  - Added a new emitIfClause helper function to allow handling if clauses in a similar fashion to Clang.
  - Updated the MLIR side of code to account for changes to createTargetData.

Depends on D149872

Reviewed By: jdoerfert

Differential Revision: https://reviews.llvm.org/D146557

Author: TIFitis

llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h

@@ -1478,6 +1478,25 @@ class OpenMPIRBuilder {
   /// Computes the size of type in bytes.
   Value *getSizeInBytes(Value *BasePtr);
 
+  // Emit a branch from the current block to the Target block only if
+  // the current block has a terminator.
+  void emitBranch(BasicBlock *Target);
+
+  // If BB has no use then delete it and return. Else place BB after the current
+  // block, if possible, or else at the end of the function. Also add a branch
+  // from current block to BB if current block does not have a terminator.
+  void emitBlock(BasicBlock *BB, Function *CurFn, bool IsFinished = false);
+
+  /// Emits code for OpenMP 'if' clause using specified \a BodyGenCallbackTy
+  /// Here is the logic:
+  /// if (Cond) {
+  ///   ThenGen();
+  /// } else {
+  ///   ElseGen();
+  /// }
+  void emitIfClause(Value *Cond, BodyGenCallbackTy ThenGen,
+                    BodyGenCallbackTy ElseGen, InsertPointTy AllocaIP = {});
+
   /// Create the global variable holding the offload mappings information.
   GlobalVariable *createOffloadMaptypes(SmallVectorImpl<uint64_t> &Mappings,
                                         std::string VarName);
@@ -1987,29 +2006,41 @@ class OpenMPIRBuilder {
                                          StringRef EntryFnName,
                                          StringRef EntryFnIDName,
                                          int32_t NumTeams, int32_t NumThreads);
+  /// Type of BodyGen to use for region codegen
+  ///
+  /// Priv: If device pointer privatization is required, emit the body of the
+  /// region here. It will have to be duplicated: with and without
+  /// privatization.
+  /// DupNoPriv: If we need device pointer privatization, we need
+  /// to emit the body of the region with no privatization in the 'else' branch
+  /// of the conditional.
+  /// NoPriv: If we don't require privatization of device
+  /// pointers, we emit the body in between the runtime calls. This avoids
+  /// duplicating the body code.
+  enum BodyGenTy { Priv, DupNoPriv, NoPriv };
 
   /// Generator for '#omp target data'
   ///
   /// \param Loc The location where the target data construct was encountered.
+  /// \param AllocaIP The insertion points to be used for alloca instructions.
   /// \param CodeGenIP The insertion point at which the target directive code
   /// should be placed.
-  /// \param MapTypeFlags BitVector storing the mapType flags for the
-  /// mapOperands.
-  /// \param MapNames Names for the mapOperands.
-  /// \param MapperAllocas Pointers to the AllocInsts for the map clause.
   /// \param IsBegin If true then emits begin mapper call otherwise emits
   /// end mapper call.
   /// \param DeviceID Stores the DeviceID from the device clause.
   /// \param IfCond Value which corresponds to the if clause condition.
-  /// \param ProcessMapOpCB Callback that generates code for the map clause.
-  /// \param BodyGenCB Callback that will generate the region code.
+  /// \param Info Stores all information realted to the Target Data directive.
+  /// \param GenMapInfoCB Callback that populates the MapInfos and returns.
+  /// \param BodyGenCB Optional Callback to generate the region code.
   OpenMPIRBuilder::InsertPointTy createTargetData(
-      const LocationDescription &Loc, OpenMPIRBuilder::InsertPointTy CodeGenIP,
-      SmallVectorImpl<uint64_t> &MapTypeFlags,
-      SmallVectorImpl<Constant *> &MapNames,
-      struct MapperAllocas &MapperAllocas, bool IsBegin, int64_t DeviceID,
-      Value *IfCond, BodyGenCallbackTy ProcessMapOpCB,
-      BodyGenCallbackTy BodyGenCB = {});
+      const LocationDescription &Loc, InsertPointTy AllocaIP,
+      InsertPointTy CodeGenIP, Value *DeviceID, Value *IfCond,
+      TargetDataInfo &Info,
+      function_ref<MapInfosTy &(InsertPointTy CodeGenIP)> GenMapInfoCB,
+      omp::RuntimeFunction *MapperFunc = nullptr,
+      function_ref<InsertPointTy(InsertPointTy CodeGenIP,
+                                 BodyGenTy BodyGenType)>
+          BodyGenCB = nullptr);
 
   using TargetBodyGenCallbackTy = function_ref<InsertPointTy(
       InsertPointTy AllocaIP, InsertPointTy CodeGenIP)>;

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -4078,73 +4078,121 @@ Constant *OpenMPIRBuilder::registerTargetRegionFunction(
 }
 
 OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::createTargetData(
-    const LocationDescription &Loc, OpenMPIRBuilder::InsertPointTy CodeGenIP,
-    SmallVectorImpl<uint64_t> &MapTypeFlags,
-    SmallVectorImpl<Constant *> &MapNames, struct MapperAllocas &MapperAllocas,
-    bool IsBegin, int64_t DeviceID, Value *IfCond,
-    BodyGenCallbackTy ProcessMapOpCB, BodyGenCallbackTy BodyGenCB) {
+    const LocationDescription &Loc, InsertPointTy AllocaIP,
+    InsertPointTy CodeGenIP, Value *DeviceID, Value *IfCond,
+    TargetDataInfo &Info,
+    function_ref<MapInfosTy &(InsertPointTy CodeGenIP)> GenMapInfoCB,
+    omp::RuntimeFunction *MapperFunc,
+    function_ref<InsertPointTy(InsertPointTy CodeGenIP, BodyGenTy BodyGenType)>
+        BodyGenCB) {
   if (!updateToLocation(Loc))
     return InsertPointTy();
 
   Builder.restoreIP(CodeGenIP);
+  bool IsStandAlone = !BodyGenCB;
 
-  // LLVM utilities like blocks with terminators.
-  // The UI acts as a resume point for code insertion after the BodyGen
-  auto *UI = Builder.CreateUnreachable();
-  if (IfCond) {
-    auto *ThenTI =
-        SplitBlockAndInsertIfThen(IfCond, UI, /* Unreachable */ false);
-    ThenTI->getParent()->setName("omp_if.then");
-    Builder.SetInsertPoint(ThenTI);
-  } else {
-    Builder.SetInsertPoint(UI);
-  }
+  // Generate the code for the opening of the data environment. Capture all the
+  // arguments of the runtime call by reference because they are used in the
+  // closing of the region.
+  auto BeginThenGen = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+    emitOffloadingArrays(AllocaIP, Builder.saveIP(),
+                         GenMapInfoCB(Builder.saveIP()), Info,
+                         /*IsNonContiguous=*/true);
 
-  ProcessMapOpCB(Builder.saveIP(), Builder.saveIP());
+    TargetDataRTArgs RTArgs;
+    emitOffloadingArraysArgument(Builder, RTArgs, Info);
 
-  uint32_t SrcLocStrSize;
-  Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
-  Value *srcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
+    // Emit the number of elements in the offloading arrays.
+    Value *PointerNum = Builder.getInt32(Info.NumberOfPtrs);
 
-  GlobalVariable *MapTypesGV =
-      createOffloadMaptypes(MapTypeFlags, ".offload_maptypes");
-  Value *MapTypesArg = Builder.CreateConstInBoundsGEP2_32(
-      ArrayType::get(Builder.getInt64Ty(), MapTypeFlags.size()), MapTypesGV,
-      /*Idx0=*/0, /*Idx1=*/0);
+    // Source location for the ident struct
+    uint32_t SrcLocStrSize;
+    Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
+    Value *SrcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
+
+    Value *OffloadingArgs[] = {SrcLocInfo,           DeviceID,
+                               PointerNum,           RTArgs.BasePointersArray,
+                               RTArgs.PointersArray, RTArgs.SizesArray,
+                               RTArgs.MapTypesArray, RTArgs.MapNamesArray,
+                               RTArgs.MappersArray};
+
+    if (IsStandAlone) {
+      assert(MapperFunc && "MapperFunc missing for standalone target data");
+      Builder.CreateCall(getOrCreateRuntimeFunctionPtr(*MapperFunc),
+                         OffloadingArgs);
+    } else {
+      Function *BeginMapperFunc = getOrCreateRuntimeFunctionPtr(
+          omp::OMPRTL___tgt_target_data_begin_mapper);
 
-  GlobalVariable *MapNamesGV =
-      createOffloadMapnames(MapNames, ".offload_mapnames");
-  Value *MapNamesArg = Builder.CreateConstInBoundsGEP2_32(
-      ArrayType::get(Builder.getInt8PtrTy(), MapNames.size()), MapNamesGV,
-      /*Idx0=*/0, /*Idx1=*/0);
+      Builder.CreateCall(BeginMapperFunc, OffloadingArgs);
+
+      // If device pointer privatization is required, emit the body of the
+      // region here. It will have to be duplicated: with and without
+      // privatization.
+      Builder.restoreIP(BodyGenCB(Builder.saveIP(), BodyGenTy::Priv));
+    }
+  };
+
+  // If we need device pointer privatization, we need to emit the body of the
+  // region with no privatization in the 'else' branch of the conditional.
+  // Otherwise, we don't have to do anything.
+  auto BeginElseGen = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+    Builder.restoreIP(BodyGenCB(Builder.saveIP(), BodyGenTy::DupNoPriv));
+  };
+
+  // Generate code for the closing of the data region.
+  auto EndThenGen = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {
+    TargetDataRTArgs RTArgs;
+    emitOffloadingArraysArgument(Builder, RTArgs, Info, /*EmitDebug=*/false,
+                                 /*ForEndCall=*/true);
+
+    // Emit the number of elements in the offloading arrays.
+    Value *PointerNum = Builder.getInt32(Info.NumberOfPtrs);
+
+    // Source location for the ident struct
+    uint32_t SrcLocStrSize;
+    Constant *SrcLocStr = getOrCreateSrcLocStr(Loc, SrcLocStrSize);
+    Value *SrcLocInfo = getOrCreateIdent(SrcLocStr, SrcLocStrSize);
+
+    Value *OffloadingArgs[] = {SrcLocInfo,           DeviceID,
+                               PointerNum,           RTArgs.BasePointersArray,
+                               RTArgs.PointersArray, RTArgs.SizesArray,
+                               RTArgs.MapTypesArray, RTArgs.MapNamesArray,
+                               RTArgs.MappersArray};
+    Function *EndMapperFunc =
+        getOrCreateRuntimeFunctionPtr(omp::OMPRTL___tgt_target_data_end_mapper);
+
+    Builder.CreateCall(EndMapperFunc, OffloadingArgs);
+  };
 
-  Function *beginMapperFunc =
-      getOrCreateRuntimeFunctionPtr(omp::OMPRTL___tgt_target_data_begin_mapper);
-  Function *endMapperFunc =
-      getOrCreateRuntimeFunctionPtr(omp::OMPRTL___tgt_target_data_end_mapper);
+  // We don't have to do anything to close the region if the if clause evaluates
+  // to false.
+  auto EndElseGen = [&](InsertPointTy AllocaIP, InsertPointTy CodeGenIP) {};
 
   if (BodyGenCB) {
-    // Create call to start the data region.
-    emitMapperCall(Builder.saveIP(), beginMapperFunc, srcLocInfo, MapTypesArg,
-                   MapNamesArg, MapperAllocas, DeviceID, MapTypeFlags.size());
+    if (IfCond) {
+      emitIfClause(IfCond, BeginThenGen, BeginElseGen, AllocaIP);
+    } else {
+      BeginThenGen(AllocaIP, Builder.saveIP());
+    }
 
-    BodyGenCB(Builder.saveIP(), Builder.saveIP());
+    // If we don't require privatization of device pointers, we emit the body in
+    // between the runtime calls. This avoids duplicating the body code.
+    Builder.restoreIP(BodyGenCB(Builder.saveIP(), BodyGenTy::NoPriv));
 
-    Builder.SetInsertPoint(UI->getParent());
-    // Create call to end the data region.
-    emitMapperCall(Builder.saveIP(), endMapperFunc, srcLocInfo, MapTypesArg,
-                   MapNamesArg, MapperAllocas, DeviceID, MapTypeFlags.size());
+    if (IfCond) {
+      emitIfClause(IfCond, EndThenGen, EndElseGen, AllocaIP);
+    } else {
+      EndThenGen(AllocaIP, Builder.saveIP());
+    }
   } else {
-    emitMapperCall(Builder.saveIP(), IsBegin ? beginMapperFunc : endMapperFunc,
-                   srcLocInfo, MapTypesArg, MapNamesArg, MapperAllocas,
-                   DeviceID, MapTypeFlags.size());
+    if (IfCond) {
+      emitIfClause(IfCond, BeginThenGen, EndElseGen, AllocaIP);
+    } else {
+      BeginThenGen(AllocaIP, Builder.saveIP());
+    }
   }
 
-  // Update the insertion point and remove the terminator we introduced.
-  Builder.SetInsertPoint(UI->getParent());
-  if (IfCond)
-    UI->getParent()->setName("omp_if.end");
-  UI->eraseFromParent();
   return Builder.saveIP();
 }
 
@@ -4668,6 +4716,77 @@ void OpenMPIRBuilder::emitOffloadingArrays(
   emitNonContiguousDescriptor(AllocaIP, CodeGenIP, CombinedInfo, Info);
 }
 
+void OpenMPIRBuilder::emitBranch(BasicBlock *Target) {
+  BasicBlock *CurBB = Builder.GetInsertBlock();
+
+  if (!CurBB || CurBB->getTerminator()) {
+    // If there is no insert point or the previous block is already
+    // terminated, don't touch it.
+  } else {
+    // Otherwise, create a fall-through branch.
+    Builder.CreateBr(Target);
+  }
+
+  Builder.ClearInsertionPoint();
+}
+
+void OpenMPIRBuilder::emitBlock(BasicBlock *BB, Function *CurFn,
+                                bool IsFinished) {
+  BasicBlock *CurBB = Builder.GetInsertBlock();
+
+  // Fall out of the current block (if necessary).
+  emitBranch(BB);
+
+  if (IsFinished && BB->use_empty()) {
+    BB->eraseFromParent();
+    return;
+  }
+
+  // Place the block after the current block, if possible, or else at
+  // the end of the function.
+  if (CurBB && CurBB->getParent())
+    CurFn->insert(std::next(CurBB->getIterator()), BB);
+  else
+    CurFn->insert(CurFn->end(), BB);
+  Builder.SetInsertPoint(BB);
+}
+
+void OpenMPIRBuilder::emitIfClause(Value *Cond, BodyGenCallbackTy ThenGen,
+                                   BodyGenCallbackTy ElseGen,
+                                   InsertPointTy AllocaIP) {
+  // If the condition constant folds and can be elided, try to avoid emitting
+  // the condition and the dead arm of the if/else.
+  if (auto *CI = dyn_cast<ConstantInt>(Cond)) {
+    auto CondConstant = CI->getSExtValue();
+    if (CondConstant)
+      ThenGen(AllocaIP, Builder.saveIP());
+    else
+      ElseGen(AllocaIP, Builder.saveIP());
+    return;
+  }
+
+  Function *CurFn = Builder.GetInsertBlock()->getParent();
+
+  // Otherwise, the condition did not fold, or we couldn't elide it.  Just
+  // emit the conditional branch.
+  BasicBlock *ThenBlock = BasicBlock::Create(M.getContext(), "omp_if.then");
+  BasicBlock *ElseBlock = BasicBlock::Create(M.getContext(), "omp_if.else");
+  BasicBlock *ContBlock = BasicBlock::Create(M.getContext(), "omp_if.end");
+  Builder.CreateCondBr(Cond, ThenBlock, ElseBlock);
+  // Emit the 'then' code.
+  emitBlock(ThenBlock, CurFn);
+  ThenGen(AllocaIP, Builder.saveIP());
+  emitBranch(ContBlock);
+  // Emit the 'else' code if present.
+  // There is no need to emit line number for unconditional branch.
+  emitBlock(ElseBlock, CurFn);
+  ElseGen(AllocaIP, Builder.saveIP());
+  // There is no need to emit line number for unconditional branch.
+  emitBranch(ContBlock);
+  // Emit the continuation block for code after the if.
+  emitBlock(ContBlock, CurFn, /*IsFinished=*/true);
+}
+
 bool OpenMPIRBuilder::checkAndEmitFlushAfterAtomic(
     const LocationDescription &Loc, llvm::AtomicOrdering AO, AtomicKind AK) {
   assert(!(AO == AtomicOrdering::NotAtomic ||