[mlir] handle nested regions in llvm-legalize-for-export

The translation from the MLIR LLVM dialect to LLVM IR includes a mechanism that
ensures the successors of a block to be different blocks in case block
arguments are passed to them since the opposite cannot be expressed in LLVM IR.
This mechanism previously only worked for functions because it was written
prior to the introduction of other region-carrying operations such as the
OpenMP dialect, which also translates directly to LLVM IR. Modify this
mechanism to handle all regions in the module and not only functions.

Reviewed By: wsmoses

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

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<OpTrait::IsTerminator>())
+    return;
+
   auto *terminator = bb.getTerminator();
 
   // Find repeated successors with arguments.
@@ -49,9 +56,11 @@ static void ensureDistinctSuccessors(Block &bb) {
 }
 
 void mlir::LLVM::ensureDistinctSuccessors(Operation *op) {
-  op->walk([](LLVMFuncOp f) {
-    for (auto &bb : f) {
-      ::ensureDistinctSuccessors(bb);
+  op->walk([](Operation *nested) {
+    for (Region &region : llvm::make_early_inc_range(nested->getRegions())) {
+      for (Block &block : llvm::make_early_inc_range(region)) {
+        ::ensureDistinctSuccessors(block);
+      }
     }
   });
 }

mlir/test/Dialect/LLVMIR/legalize-for-export.mlir

@@ -29,3 +29,17 @@ llvm.func @repeated_successor_no_args(%arg0: i1) {
 
 // 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
+}

mlir/test/Target/LLVMIR/openmp-llvm.mlir

@@ -892,3 +892,19 @@ llvm.func @omp_sections_with_clauses() -> () {
   }
   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
+}