From 1ad48d6de27089d0223543c0b5c33961c08de877 Mon Sep 17 00:00:00 2001 From: Alex Zinenko Date: Tue, 18 Jan 2022 10:16:19 +0100 Subject: [PATCH] [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 ++++++++++++--- mlir/test/Dialect/LLVMIR/legalize-for-export.mlir | 14 ++++++++++++++ mlir/test/Target/LLVMIR/openmp-llvm.mlir | 16 ++++++++++++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/LegalizeForExport.cpp index 3bf306d..f4fa100 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 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 @@ 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 ®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 index 28cd980..b3f9e41 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 2c8bf20..91493b7 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 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 +} -- 2.7.4