diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp index d6112fa9af118..6ffe169038c54 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp @@ -1340,13 +1340,23 @@ allocatePrivateVars(llvm::IRBuilderBase &builder, llvm::SmallVectorImpl &llvmPrivateVars, const llvm::OpenMPIRBuilder::InsertPointTy &allocaIP, llvm::DenseMap *mappedPrivateVars = nullptr) { - llvm::IRBuilderBase::InsertPointGuard guard(builder); // Allocate private vars llvm::BranchInst *allocaTerminator = llvm::cast(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(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 privateReductionVariables( wsloopOp.getNumReductionVars()); - splitBB(llvm::OpenMPIRBuilder::InsertPointTy( - allocaIP.getBlock(), - allocaIP.getBlock()->getTerminator()->getIterator()), - true, "omp.region.after_alloca"); - llvm::Expected afterAllocas = allocatePrivateVars( builder, moduleTranslation, privateBlockArgs, privateDecls, mlirPrivateVars, llvmPrivateVars, allocaIP); diff --git a/mlir/test/Target/LLVMIR/openmp-target-simd-on_device.mlir b/mlir/test/Target/LLVMIR/openmp-target-simd-on_device.mlir new file mode 100644 index 0000000000000..0ce90578ea9d6 --- /dev/null +++ b/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 %{{.*}}