Make convertOmpTarget use allocatePrivateVars.

Allows more code reuse by generalizing `allocatePrivateVars` to serve
`target` op conversion.

Author: ergawy

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

@@ -6807,8 +6807,11 @@ static Expected<Function *> createOutlinedFunction(
     OMPBuilder.ConstantAllocaRaiseCandidates.emplace_back(Func);
 
   // Insert target deinit call in the device compilation pass.
-  llvm::OpenMPIRBuilder::InsertPointOrErrorTy AfterIP =
-      CBFunc(Builder.saveIP(), Builder.saveIP());
+  BasicBlock *OutlinedBodyBB =
+      splitBB(Builder, /*CreateBranch=*/true, "outlined.body");
+  llvm::OpenMPIRBuilder::InsertPointOrErrorTy AfterIP = CBFunc(
+      Builder.saveIP(),
+      OpenMPIRBuilder::InsertPointTy(OutlinedBodyBB, OutlinedBodyBB->begin()));
   if (!AfterIP)
     return AfterIP.takeError();
   Builder.restoreIP(*AfterIP);

llvm/unittests/Frontend/OpenMPIRBuilderTest.cpp

@@ -6358,7 +6358,13 @@ TEST_F(OpenMPIRBuilderTest, TargetRegionDevice) {
   auto *Load2 = Load1->getNextNode();
   EXPECT_TRUE(isa<LoadInst>(Load2));
 
-  auto *Value1 = Load2->getNextNode();
+  auto *OutlinedBlockBr = Load2->getNextNode();
+  EXPECT_TRUE(isa<BranchInst>(OutlinedBlockBr));
+
+  auto *OutlinedBlock = OutlinedBlockBr->getSuccessor(0);
+  EXPECT_EQ(OutlinedBlock->getName(), "outlined.body");
+
+  auto *Value1 = OutlinedBlock->getFirstNonPHI();
   EXPECT_EQ(Value1, Value);
   EXPECT_EQ(Value1->getNextNode(), TargetStore);
   auto *Deinit = TargetStore->getNextNode();
@@ -6510,7 +6516,14 @@ TEST_F(OpenMPIRBuilderTest, ConstantAllocaRaise) {
   EXPECT_EQ(UserCodeBlock->getName(), "user_code.entry");
   auto *Load1 = UserCodeBlock->getFirstNonPHI();
   EXPECT_TRUE(isa<LoadInst>(Load1));
-  auto *Load2 = Load1->getNextNode();
+
+  auto *OutlinedBlockBr = Load1->getNextNode();
+  EXPECT_TRUE(isa<BranchInst>(OutlinedBlockBr));
+
+  auto *OutlinedBlock = OutlinedBlockBr->getSuccessor(0);
+  EXPECT_EQ(OutlinedBlock->getName(), "outlined.body");
+
+  auto *Load2 = OutlinedBlock->getFirstNonPHI();
   EXPECT_TRUE(isa<LoadInst>(Load2));
   EXPECT_EQ(Load2, Value);
   EXPECT_EQ(Load2->getNextNode(), TargetStore);

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -1286,6 +1286,44 @@ static LogicalResult allocAndInitializeReductionVars(
                            isByRef, deferredStores);
 }
 
+/// Return the llvm::Value * corresponding to the `privateVar` that
+/// is being privatized. It isn't always as simple as looking up
+/// moduleTranslation with privateVar. For instance, in case of
+/// an allocatable, the descriptor for the allocatable is privatized.
+/// This descriptor is mapped using an MapInfoOp. So, this function
+/// will return a pointer to the llvm::Value corresponding to the
+/// block argument for the mapped descriptor.
+static llvm::Value *
+findAssociatedValue(Value privateVar, llvm::IRBuilderBase &builder,
+                    LLVM::ModuleTranslation &moduleTranslation,
+                    omp::TargetOp targetOp = nullptr,
+                    llvm::DenseMap<Value, int> *mappedPrivateVars = nullptr) {
+  if (mappedPrivateVars != nullptr && mappedPrivateVars->contains(privateVar)) {
+    int blockArgIndex = (*mappedPrivateVars)[privateVar];
+    Value blockArg = targetOp.getRegion().getArgument(blockArgIndex);
+    Type privVarType = privateVar.getType();
+    Type blockArgType = blockArg.getType();
+    assert(isa<LLVM::LLVMPointerType>(blockArgType) &&
+           "A block argument corresponding to a mapped var should have "
+           "!llvm.ptr type");
+
+    if (privVarType == blockArg.getType())
+      return moduleTranslation.lookupValue(blockArg);
+
+    if (!isa<LLVM::LLVMPointerType>(privVarType)) {
+      // This typically happens when the privatized type is lowered from
+      // boxchar<KIND> and gets lowered to !llvm.struct<(ptr, i64)>. That is the
+      // struct/pair is passed by value. But, mapped values are passed only as
+      // pointers, so before we privatize, we must load the pointer.
+      llvm::Value *load =
+          builder.CreateLoad(moduleTranslation.convertType(privVarType),
+                             moduleTranslation.lookupValue(blockArg));
+      return load;
+    }
+  }
+  return moduleTranslation.lookupValue(privateVar);
+}
+
 /// Allocate delayed private variables. Returns the basic block which comes
 /// after all of these allocations. llvm::Value * for each of these private
 /// variables are populated in llvmPrivateVars.
@@ -1296,7 +1334,9 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
                     MutableArrayRef<omp::PrivateClauseOp> privateDecls,
                     MutableArrayRef<mlir::Value> mlirPrivateVars,
                     llvm::SmallVectorImpl<llvm::Value *> &llvmPrivateVars,
-                    const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP) {
+                    const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
+                    omp::TargetOp targetOp = nullptr,
+                    llvm::DenseMap<Value, int> *mappedPrivateVars = nullptr) {
   llvm::IRBuilderBase::InsertPointGuard guard(builder);
   // Allocate private vars
   llvm::BranchInst *allocaTerminator =
@@ -1326,7 +1366,8 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
     Region &allocRegion = privDecl.getAllocRegion();
 
     // map allocation region block argument
-    llvm::Value *nonPrivateVar = moduleTranslation.lookupValue(mlirPrivVar);
+    llvm::Value *nonPrivateVar = findAssociatedValue(
+        mlirPrivVar, builder, moduleTranslation, targetOp, mappedPrivateVars);
     assert(nonPrivateVar);
     moduleTranslation.mapValue(privDecl.getAllocMoldArg(), nonPrivateVar);
 
@@ -1341,6 +1382,7 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
     } else {
       builder.SetInsertPoint(privAllocBlock->getTerminator());
     }
+
     if (failed(inlineConvertOmpRegions(allocRegion, "omp.private.alloc",
                                        builder, moduleTranslation, &phis)))
       return llvm::createStringError(
@@ -3809,43 +3851,6 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
   return builder.saveIP();
 }
 
-/// Return the llvm::Value * corresponding to the `privateVar` that
-/// is being privatized. It isn't always as simple as looking up
-/// moduleTranslation with privateVar. For instance, in case of
-/// an allocatable, the descriptor for the allocatable is privatized.
-/// This descriptor is mapped using an MapInfoOp. So, this function
-/// will return a pointer to the llvm::Value corresponding to the
-/// block argument for the mapped descriptor.
-static llvm::Value *
-findHostAssociatedValue(Value privateVar, omp::TargetOp targetOp,
-                        llvm::DenseMap<Value, int> &mappedPrivateVars,
-                        llvm::IRBuilderBase &builder,
-                        LLVM::ModuleTranslation &moduleTranslation) {
-  if (mappedPrivateVars.contains(privateVar)) {
-    int blockArgIndex = mappedPrivateVars[privateVar];
-    Value blockArg = targetOp.getRegion().getArgument(blockArgIndex);
-    Type privVarType = privateVar.getType();
-    Type blockArgType = blockArg.getType();
-    assert(isa<LLVM::LLVMPointerType>(blockArgType) &&
-           "A block argument corresponding to a mapped var should have "
-           "!llvm.ptr type");
-
-    if (privVarType == blockArg.getType())
-      return moduleTranslation.lookupValue(blockArg);
-
-    if (!isa<LLVM::LLVMPointerType>(privVarType)) {
-      // This typically happens when the privatized type is lowered from
-      // boxchar<KIND> and gets lowered to !llvm.struct<(ptr, i64)>. That is the
-      // struct/pair is passed by value. But, mapped values are passed only as
-      // pointers, so before we privatize, we must load the pointer.
-      llvm::Value *load =
-          builder.CreateLoad(moduleTranslation.convertType(privVarType),
-                             moduleTranslation.lookupValue(blockArg));
-      return load;
-    }
-  }
-  return moduleTranslation.lookupValue(privateVar);
-}
 static LogicalResult
 convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
                  LLVM::ModuleTranslation &moduleTranslation) {
@@ -3949,68 +3954,49 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
         attr.isStringAttribute())
       llvmOutlinedFn->addFnAttr(attr);
 
-    builder.restoreIP(codeGenIP);
     for (auto [arg, mapOp] : llvm::zip_equal(mapBlockArgs, mapVars)) {
       auto mapInfoOp = cast<omp::MapInfoOp>(mapOp.getDefiningOp());
       llvm::Value *mapOpValue =
           moduleTranslation.lookupValue(mapInfoOp.getVarPtr());
       moduleTranslation.mapValue(arg, mapOpValue);
     }
+
     // Do privatization after moduleTranslation has already recorded
     // mapped values.
+    MutableArrayRef<BlockArgument> privateBlockArgs =
+        cast<omp::BlockArgOpenMPOpInterface>(opInst).getPrivateBlockArgs();
+    SmallVector<mlir::Value> mlirPrivateVars;
     SmallVector<llvm::Value *> llvmPrivateVars;
+    SmallVector<omp::PrivateClauseOp> privateDecls;
+    mlirPrivateVars.reserve(privateBlockArgs.size());
+    llvmPrivateVars.reserve(privateBlockArgs.size());
+    collectPrivatizationDecls(targetOp, privateDecls);
+    for (mlir::Value privateVar : targetOp.getPrivateVars())
+      mlirPrivateVars.push_back(privateVar);
+
+    llvm::Expected<llvm::BasicBlock *> afterAllocas =
+        allocatePrivateVars(builder, moduleTranslation, privateBlockArgs,
+                            privateDecls, mlirPrivateVars, llvmPrivateVars,
+                            allocaIP, targetOp, &mappedPrivateVars);
+
+    if (handleError(afterAllocas, *targetOp).failed())
+      return llvm::make_error<PreviouslyReportedError>();
+
     SmallVector<Region *> privateCleanupRegions;
-    if (!targetOp.getPrivateVars().empty()) {
-      builder.restoreIP(allocaIP);
-
-      OperandRange privateVars = targetOp.getPrivateVars();
-      std::optional<ArrayAttr> privateSyms = targetOp.getPrivateSyms();
-      MutableArrayRef<BlockArgument> privateBlockArgs =
-          cast<omp::BlockArgOpenMPOpInterface>(opInst).getPrivateBlockArgs();
-
-      for (auto [privVar, privatizerNameAttr, privBlockArg] :
-           llvm::zip_equal(privateVars, *privateSyms, privateBlockArgs)) {
-
-        SymbolRefAttr privSym = cast<SymbolRefAttr>(privatizerNameAttr);
-        omp::PrivateClauseOp privatizer = findPrivatizer(&opInst, privSym);
-        assert(privatizer.getDataSharingType() !=
-                   omp::DataSharingClauseType::FirstPrivate &&
-               "unsupported privatizer");
-        Region &allocRegion = privatizer.getAllocRegion();
-        BlockArgument allocRegionArg = allocRegion.getArgument(0);
-        moduleTranslation.mapValue(
-            allocRegionArg,
-            findHostAssociatedValue(privVar, targetOp, mappedPrivateVars,
-                                    builder, moduleTranslation));
-        SmallVector<llvm::Value *, 1> yieldedValues;
-        if (failed(inlineConvertOmpRegions(
-                allocRegion, "omp.targetop.privatizer", builder,
-                moduleTranslation, &yieldedValues))) {
-          return llvm::createStringError(
-              "failed to inline `alloc` region of `omp.private`");
-        }
-        assert(yieldedValues.size() == 1);
-        llvm::Value *llvmReplacementValue = yieldedValues.front();
-        moduleTranslation.mapValue(privBlockArg, llvmReplacementValue);
-        if (!privatizer.getDeallocRegion().empty()) {
-          llvmPrivateVars.push_back(llvmReplacementValue);
-          privateCleanupRegions.push_back(&privatizer.getDeallocRegion());
-        }
-        moduleTranslation.forgetMapping(allocRegion);
-        builder.restoreIP(builder.saveIP());
-      }
-    }
+    llvm::transform(privateDecls, std::back_inserter(privateCleanupRegions),
+                    [](omp::PrivateClauseOp privatizer) {
+                      return &privatizer.getDeallocRegion();
+                    });
 
+    builder.restoreIP(codeGenIP);
     llvm::Expected<llvm::BasicBlock *> exitBlock = convertOmpOpRegions(
         targetRegion, "omp.target", builder, moduleTranslation);
+
     if (!exitBlock)
       return exitBlock.takeError();
 
     builder.SetInsertPoint(*exitBlock);
-    if (!llvmPrivateVars.empty()) {
-      assert(llvmPrivateVars.size() == privateCleanupRegions.size() &&
-             "Number of private variables needing cleanup not equal to number"
-             "of privatizers with dealloc regions");
+    if (!privateCleanupRegions.empty()) {
       if (failed(inlineOmpRegionCleanup(
               privateCleanupRegions, llvmPrivateVars, moduleTranslation,
               builder, "omp.targetop.private.cleanup",
@@ -4020,10 +4006,10 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
             "op in the target region");
       }
     }
-    return builder.saveIP();
+
+    return InsertPointTy(exitBlock.get(), exitBlock.get()->end());
   };
 
-  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
   StringRef parentName = parentFn.getName();
 
   llvm::TargetRegionEntryInfo entryInfo;
@@ -4034,9 +4020,6 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   int32_t defaultValTeams = -1;
   int32_t defaultValThreads = 0;
 
-  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
-      findAllocaInsertPoint(builder, moduleTranslation);
-
   MapInfoData mapData;
   collectMapDataFromMapOperands(mapData, mapVars, moduleTranslation, dl,
                                 builder);
@@ -4084,6 +4067,10 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
   buildDependData(targetOp.getDependKinds(), targetOp.getDependVars(),
                   moduleTranslation, dds);
 
+  llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
+      findAllocaInsertPoint(builder, moduleTranslation);
+  llvm::OpenMPIRBuilder::LocationDescription ompLoc(builder);
+
   llvm::OpenMPIRBuilder::InsertPointOrErrorTy afterIP =
       moduleTranslation.getOpenMPBuilder()->createTarget(
           ompLoc, isOffloadEntry, allocaIP, builder.saveIP(), entryInfo,

mlir/test/Target/LLVMIR/omptarget-byref-bycopy-generation-device.mlir

@@ -33,9 +33,12 @@ module attributes {omp.is_target_device = true} {
 
 // CHECK: user_code.entry:                                  ; preds = %entry
 // CHECK: %[[LOAD_BYREF:.*]] = load ptr, ptr %[[ALLOCA_BYREF]], align 8
+// CHECK: br label %outlined.body
+
+// CHECK: outlined.body:
 // CHECK: br label %omp.target
 
-// CHECK: omp.target:                                       ; preds = %user_code.entry
+// CHECK: omp.target:
 // CHECK:  %[[VAL_LOAD_BYCOPY:.*]] = load i32, ptr %[[ALLOCA_BYCOPY]], align 4
 // CHECK:  store i32 %[[VAL_LOAD_BYCOPY]], ptr %[[LOAD_BYREF]], align 4
 // CHECK: br label %omp.region.cont

mlir/test/Target/LLVMIR/omptarget-declare-target-llvm-device.mlir

@@ -17,7 +17,7 @@ module attributes {omp.is_target_device = true} {
   llvm.func @_QQmain() attributes {} {
     %0 = llvm.mlir.addressof @_QMtest_0Esp : !llvm.ptr
 
-  // CHECK-DAG:   omp.target:                                       ; preds = %user_code.entry
+  // CHECK-DAG:   omp.target:                                       ; preds = %outlined.body
   // CHECK-DAG: %[[V:.*]] = load ptr, ptr @_QMtest_0Esp_decl_tgt_ref_ptr, align 8
   // CHECK-DAG: store i32 1, ptr %[[V]], align 4
   // CHECK-DAG: br label %omp.region.cont

mlir/test/Target/LLVMIR/openmp-target-use-device-nested.mlir

@@ -13,7 +13,11 @@
 // CHECK:            user_code.entry:                                  ; preds = %[[VAL_10:.*]]
 // CHECK-NEXT:         %[[VAL_11:.*]] = load ptr, ptr %[[VAL_3]], align 8
 // CHECK-NEXT:         br label %[[VAL_12:.*]]
-// CHECK:            omp.target:                                       ; preds = %[[VAL_8]]
+
+// CHECK:            [[VAL_12]]:
+// CHECK-NEXT:         br label %[[TARGET_REG_ENTRY:.*]]
+
+// CHECK:            [[TARGET_REG_ENTRY]]:                                       ; preds = %[[VAL_12]]
 // CHECK-NEXT:         %[[VAL_13:.*]] = load ptr, ptr %[[VAL_11]], align 8
 // CHECK-NEXT:         store i32 999, ptr %[[VAL_13]], align 4
 // CHECK-NEXT:         br label %[[VAL_14:.*]]