diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp --- a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp +++ b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp @@ -15,7 +15,14 @@ using namespace mlir; +/// If the given block has the same successor with different arguments, +/// introduce dummy successor blocks so that all successors of the given block +/// are different. static void ensureDistinctSuccessors(Block &bb) { + // Early exit if the block cannot have successors. + if (bb.empty() || !bb.back().mightHaveTrait()) + return; + auto *terminator = bb.getTerminator(); // Find repeated successors with arguments. @@ -49,9 +56,11 @@ } void mlir::LLVM::ensureDistinctSuccessors(Operation *op) { - op->walk([](LLVMFuncOp f) { - for (auto &bb : f) { - ::ensureDistinctSuccessors(bb); + op->walk([](Operation *nested) { + for (Region ®ion : llvm::make_early_inc_range(nested->getRegions())) { + for (Block &block : llvm::make_early_inc_range(region)) { + ::ensureDistinctSuccessors(block); + } } }); } diff --git a/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir b/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir --- a/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir +++ b/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir @@ -29,3 +29,17 @@ // CHECK-NOT: ^{{.*}}: } + +// CHECK: @repeated_successor_openmp +llvm.func @repeated_successor_openmp(%arg0: i64, %arg1: i64, %arg2: i64, %arg3: i1) { + omp.wsloop (%arg4) : i64 = (%arg0) to (%arg1) step (%arg2) { + // CHECK: llvm.cond_br %{{.*}}, ^[[BB1:.*]]({{.*}}), ^[[BB2:.*]]({{.*}}) + llvm.cond_br %arg3, ^bb1(%arg0 : i64), ^bb1(%arg1 : i64) + // CHECK: ^[[BB1]] + ^bb1(%0: i64): // 2 preds: ^bb0, ^bb0 + omp.yield + // CHECK: ^[[BB2]](%[[ARG:.*]]: i64): + // CHECK: llvm.br ^[[BB1]](%[[ARG]] : i64) + } + llvm.return +} diff --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir --- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir +++ b/mlir/test/Target/LLVMIR/openmp-llvm.mlir @@ -892,3 +892,19 @@ } llvm.return } + +// ----- + +// Check that translation doesn't crash in presence of repeated successor +// blocks with different arguments within OpenMP operations: LLVM cannot +// represent this and a dummy block will be introduced for forwarding. The +// introduction mechanism itself is tested elsewhere. +// CHECK-LABEL: @repeated_successor +llvm.func @repeated_successor(%arg0: i64, %arg1: i64, %arg2: i64, %arg3: i1) { + omp.wsloop (%arg4) : i64 = (%arg0) to (%arg1) step (%arg2) { + llvm.cond_br %arg3, ^bb1(%arg0 : i64), ^bb1(%arg1 : i64) + ^bb1(%0: i64): // 2 preds: ^bb0, ^bb0 + omp.yield + } + llvm.return +}