[Mlir-commits] [mlir] 1ad48d6 - [mlir] handle nested regions in llvm-legalize-for-export

[Alex Zinenko]

Author: Alex Zinenko
Date: 2022-01-18T17:09:14+01:00
New Revision: 1ad48d6de27089d0223543c0b5c33961c08de877

URL: https://github.com/llvm/llvm-project/commit/1ad48d6de27089d0223543c0b5c33961c08de877
DIFF: https://github.com/llvm/llvm-project/commit/1ad48d6de27089d0223543c0b5c33961c08de877.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
    mlir/test/Dialect/LLVMIR/legalize-for-export.mlir
    mlir/test/Target/LLVMIR/openmp-llvm.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp
index 3bf306dba493a..f4fa100e6a78f 100644
--- 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 
diff erent arguments,
+/// introduce dummy successor blocks so that all successors of the given block
+/// are 
diff erent.
 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);
+      }
     }
   });
 }

diff  --git a/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir b/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir
index 28cd980df5048..b3f9e41180841 100644
--- a/mlir/test/Dialect/LLVMIR/legalize-for-export.mlir
+++ b/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
+}

diff  --git a/mlir/test/Target/LLVMIR/openmp-llvm.mlir b/mlir/test/Target/LLVMIR/openmp-llvm.mlir
index 2c8bf2098ca31..91493b7b9cdbd 100644
--- a/mlir/test/Target/LLVMIR/openmp-llvm.mlir
+++ b/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 
diff erent 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
+}