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

[ftynse]

Summary

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.

mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp

 #include "PassDetail.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 
 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.
   llvm::SmallDenseMap<Block *, SmallVector<int, 4>> successorPositions;
   for (int i = 0, e = terminator->getNumSuccessors(); i < e; ++i) {
     Block *successor = terminator->getSuccessor(i);
     // Blocks with no arguments are safe even if they appear multiple times
     // because they don't need PHI nodes.
for (int position : llvm::drop_begin(successor.second, 1)) {
       dummyBlock->addArguments(successor.first->getArgumentTypes());
       builder.create<LLVM::BrOp>(terminator->getLoc(),
                                  dummyBlock->getArguments(), successor.first);
     }
   }
 }
 
 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);
+      }
     }
   });
 }
 
 namespace {
 struct LegalizeForExportPass
     : public LLVMLegalizeForExportBase<LegalizeForExportPass> {
   void runOnOperation() override {
     LLVM::ensureDistinctSuccessors(getOperation());
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> LLVM::createLegalizeForExportPass() {
   return std::make_unique<LegalizeForExportPass>();
 }

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

llvm.func @repeated_successor_no_args(%arg0: i1) {
   llvm.cond_br %arg0, ^bb1, ^bb1
 
 // CHECK: ^{{.*}}:
 ^bb1:
   llvm.return
 
 // 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

omp.sections nowait {
     omp.section {
       llvm.call @foo() : () -> ()
       omp.terminator
     }
     omp.terminator
   }
   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
+}