[flang][OpenMP] Generalize fixing alloca IP pre-condition for private ops

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

@@ -1340,13 +1340,23 @@ allocatePrivateVars(llvm::IRBuilderBase &builder,
                     llvm::SmallVectorImpl<llvm::Value *> &llvmPrivateVars,
                     const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP,
                     llvm::DenseMap<Value, Value> *mappedPrivateVars = nullptr) {
-  llvm::IRBuilderBase::InsertPointGuard guard(builder);
   // Allocate private vars
   llvm::BranchInst *allocaTerminator =
       llvm::cast<llvm::BranchInst>(allocaIP.getBlock()->getTerminator());
+  if (allocaTerminator->getNumSuccessors() != 1) {
+    splitBB(llvm::OpenMPIRBuilder::InsertPointTy(
+                allocaIP.getBlock(), allocaTerminator->getIterator()),
+            true, "omp.region.after_alloca");
+  }
+
+  llvm::IRBuilderBase::InsertPointGuard guard(builder);
+  // Update the allocaTerminator in case the alloca block was split above.
+  allocaTerminator =
+      llvm::cast<llvm::BranchInst>(allocaIP.getBlock()->getTerminator());
   builder.SetInsertPoint(allocaTerminator);
   assert(allocaTerminator->getNumSuccessors() == 1 &&
          "This is an unconditional branch created by OpenMPIRBuilder");
+
   llvm::BasicBlock *afterAllocas = allocaTerminator->getSuccessor(0);
 
   // FIXME: Some of the allocation regions do more than just allocating.
@@ -1876,11 +1886,6 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
   SmallVector<llvm::Value *> privateReductionVariables(
       wsloopOp.getNumReductionVars());
 
-  splitBB(llvm::OpenMPIRBuilder::InsertPointTy(
-              allocaIP.getBlock(),
-              allocaIP.getBlock()->getTerminator()->getIterator()),
-          true, "omp.region.after_alloca");
-
   llvm::Expected<llvm::BasicBlock *> afterAllocas = allocatePrivateVars(
       builder, moduleTranslation, privateBlockArgs, privateDecls,
       mlirPrivateVars, llvmPrivateVars, allocaIP);

mlir/test/Target/LLVMIR/openmp-target-simd-on_device.mlir

@@ -0,0 +1,34 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {omp.is_target_device = true} {
+  omp.private {type = private} @simd_privatizer : !llvm.ptr alloc {
+  ^bb0(%arg0: !llvm.ptr):
+    omp.yield(%arg0 : !llvm.ptr)
+  }
+
+  llvm.func @test_target_simd() {
+    omp.target {
+      %5 = llvm.mlir.constant(1 : i32) : i32
+      %x = llvm.alloca %5 x i32 {bindc_name = "x"} : (i32) -> !llvm.ptr
+      omp.simd private(@simd_privatizer %x -> %arg1 : !llvm.ptr) {
+        omp.loop_nest (%arg2) : i32 = (%5) to (%5) inclusive step (%5) {
+          omp.yield
+        }
+      }
+      omp.terminator
+    }
+    llvm.return
+  }
+
+}
+
+// CHECK-LABEL: define {{.*}} @__omp_offloading_{{.*}}_test_target_simd_{{.*}}
+
+// CHECK:         %[[INT:.*]] = alloca i32, align 4
+// CHECK:         br label %[[LATE_ALLOC_BB:.*]]
+
+// CHECK:       [[LATE_ALLOC_BB]]:
+// CHECK:         br label %[[AFTER_ALLOC_BB:.*]]
+
+// CHECK:       [[AFTER_ALLOC_BB]]:
+// CHECK:         br i1 %{{.*}}, label %{{.*}}, label %{{.*}}