[mlir] LLVM dialect: move ensureDistinctSuccessors out of std->LLVM conversion

MLIR supports terminators that have the same successor block with different
block operands, which cannot be expressed in the LLVM's phi-notation as the
block identifier is used to tell apart the predecessors. This limitation can be
worked around by branching to a new block instead, with this new block
unconditionally branching to the original successor and forwarding the
argument. Until now, this transformation was performed during the conversion
from the Standard to the LLVM dialect. This does not scale well to multiple
dialects targeting the LLVM dialect as all of them would have to be aware of
this limitation and perform the preparatory transformation. Instead, do it as a
separate pass and run it immediately before the translation.

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

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

@@ -61,15 +61,6 @@ std::unique_ptr<OpPassBase<ModuleOp>>
 createLowerToLLVMPass(bool useAlloca = false, bool useBarePtrCallConv = false,
                       bool emitCWrappers = false);
 
-namespace LLVM {
-/// Make argument-taking successors of each block distinct.  PHI nodes in LLVM
-/// IR use the predecessor ID to identify which value to take.  They do not
-/// support different values coming from the same predecessor.  If a block has
-/// another block as a successor more than once with different values, insert
-/// a new dummy block for LLVM PHI nodes to tell the sources apart.
-void ensureDistinctSuccessors(ModuleOp m);
-} // namespace LLVM
-
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_STANDARDTOLLVM_CONVERTSTANDARDTOLLVMPASS_H_

mlir/include/mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h

@@ -0,0 +1,34 @@
+//===- LegalizeForExport.h - Prepare for translation to LLVM IR -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LLVMIR_TRANSFORMS_LEGALIZE_FOR_EXPORT_H
+#define MLIR_DIALECT_LLVMIR_TRANSFORMS_LEGALIZE_FOR_EXPORT_H
+
+#include <memory>
+
+namespace mlir {
+class Operation;
+class Pass;
+
+namespace LLVM {
+
+/// Make argument-taking successors of each block distinct.  PHI nodes in LLVM
+/// IR use the predecessor ID to identify which value to take. They do not
+/// support different values coming from the same predecessor. If a block has
+/// another block as a successor more than once with different values, insert
+/// a new dummy block for LLVM PHI nodes to tell the sources apart.
+void ensureDistinctSuccessors(Operation *op);
+
+/// Creates a pass that legalizes the LLVM dialect operations so that they can
+/// be translated to LLVM IR.
+std::unique_ptr<Pass> createLegalizeForExportPass();
+
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_LEGALIZE_FOR_EXPORT_H

mlir/include/mlir/InitAllPasses.h

@@ -26,6 +26,7 @@
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Dialect/FxpMathOps/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/LoopOps/Passes.h"
 #include "mlir/Dialect/QuantOps/Passes.h"
@@ -107,6 +108,9 @@ inline void registerAllPasses() {
   createConvertLinalgToAffineLoopsPass();
   createConvertLinalgToLLVMPass();
 
+  // LLVM
+  LLVM::createLegalizeForExportPass();
+
   // LoopOps
   createParallelLoopFusionPass();
   createParallelLoopSpecializationPass();

mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h

@@ -15,6 +15,7 @@
 #define MLIR_TARGET_LLVMIR_MODULETRANSLATION_H
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Value.h"
@@ -57,6 +58,8 @@ public:
     if (!llvmModule)
       return nullptr;
 
+    LLVM::ensureDistinctSuccessors(m);
+
     T translator(m, std::move(llvmModule));
     if (failed(translator.convertGlobals()))
       return nullptr;

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

@@ -2785,49 +2785,6 @@ struct AtomicCmpXchgOpLowering : public LoadStoreOpLowering<AtomicRMWOp> {
 
 } // namespace
 
-static void ensureDistinctSuccessors(Block &bb) {
-  Operation *terminator = bb.getTerminator();
-  if (terminator->getNumSuccessors() < 2)
-    return;
-
-  // 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.
-    if (successor->getNumArguments() == 0)
-      continue;
-    successorPositions[successor].push_back(i);
-  }
-
-  // If a successor appears for the second or more time in the terminator,
-  // create a new dummy block that unconditionally branches to the original
-  // destination, and retarget the terminator to branch to this new block.
-  // There is no need to pass arguments to the dummy block because it will be
-  // dominated by the original block and can therefore use any values defined in
-  // the original block.
-  OpBuilder builder(terminator->getContext());
-  for (const auto &successor : successorPositions) {
-    // Start from the second occurrence of a block in the successor list.
-    for (int position : llvm::drop_begin(successor.second, 1)) {
-      Block *dummyBlock = builder.createBlock(bb.getParent());
-      terminator->setSuccessor(dummyBlock, position);
-      dummyBlock->addArguments(successor.first->getArgumentTypes());
-      builder.create<BranchOp>(terminator->getLoc(), successor.first,
-                               dummyBlock->getArguments());
-    }
-  }
-}
-
-void mlir::LLVM::ensureDistinctSuccessors(ModuleOp m) {
-  for (auto f : m.getOps<FuncOp>()) {
-    for (auto &bb : f.getBlocks()) {
-      ::ensureDistinctSuccessors(bb);
-    }
-  }
-}
-
 /// Collect a set of patterns to convert from the Standard dialect to LLVM.
 void mlir::populateStdToLLVMNonMemoryConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
@@ -3024,7 +2981,6 @@ struct LLVMLoweringPass : public ModulePass<LLVMLoweringPass> {
     }
 
     ModuleOp m = getModule();
-    LLVM::ensureDistinctSuccessors(m);
 
     LLVMTypeConverterCustomization customs;
     customs.funcArgConverter = useBarePtrCallConv ? barePtrFuncArgTypeConverter

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_subdirectory(Transforms)
+
 add_mlir_dialect_library(MLIRLLVMIR
   IR/LLVMDialect.cpp
 

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_mlir_dialect_library(MLIRLLVMIRTransforms
+  LegalizeForExport.cpp
+  )
+
+target_link_libraries(MLIRLLVMIRTransforms
+  PUBLIC
+  MLIRIR
+  MLIRLLVMIR
+  MLIRPass
+  )

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

@@ -0,0 +1,73 @@
+//===- LegalizeForExport.cpp - Prepare for translation to LLVM IR ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Module.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+static void ensureDistinctSuccessors(Block &bb) {
+  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.
+    if (successor->getNumArguments() == 0)
+      continue;
+    successorPositions[successor].push_back(i);
+  }
+
+  // If a successor appears for the second or more time in the terminator,
+  // create a new dummy block that unconditionally branches to the original
+  // destination, and retarget the terminator to branch to this new block.
+  // There is no need to pass arguments to the dummy block because it will be
+  // dominated by the original block and can therefore use any values defined in
+  // the original block.
+  OpBuilder builder(terminator->getContext());
+  for (const auto &successor : successorPositions) {
+    // Start from the second occurrence of a block in the successor list.
+    for (int position : llvm::drop_begin(successor.second, 1)) {
+      Block *dummyBlock = builder.createBlock(bb.getParent());
+      terminator->setSuccessor(dummyBlock, position);
+      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.getBlocks()) {
+      ::ensureDistinctSuccessors(bb);
+    }
+  });
+}
+
+namespace {
+struct LegalizeForExportPass : public OperationPass<LegalizeForExportPass> {
+  void runOnOperation() override {
+    LLVM::ensureDistinctSuccessors(getOperation());
+  }
+};
+} // end namespace
+
+std::unique_ptr<Pass> LLVM::createLegalizeForExportPass() {
+  return std::make_unique<LegalizeForExportPass>();
+}
+
+static PassRegistration<LegalizeForExportPass>
+    pass("llvm-legalize-for-export",
+         "Legalize LLVM dialect to be convertible to LLVM IR");

mlir/lib/Target/CMakeLists.txt

@@ -10,6 +10,7 @@ add_mlir_library(MLIRTargetLLVMIRModuleTranslation
 target_link_libraries(MLIRTargetLLVMIRModuleTranslation
   PUBLIC
   MLIRLLVMIR
+  MLIRLLVMIRTransforms
   LLVMCore
   LLVMIRReader
   LLVMSupport

mlir/test/Conversion/StandardToLLVM/convert-to-llvmir.mlir

@@ -554,19 +554,6 @@ func @dfs_block_order(%arg0: i32) -> (i32) {
 // CHECK-NEXT:  llvm.br ^bb1
   br ^bb1
 }
-// CHECK-LABEL: func @cond_br_same_target(%arg0: !llvm.i1, %arg1: !llvm.i32, %arg2: !llvm.i32)
-func @cond_br_same_target(%arg0: i1, %arg1: i32, %arg2 : i32) -> (i32) {
-// CHECK-NEXT:  llvm.cond_br %arg0, ^[[origBlock:bb[0-9]+]](%arg1 : !llvm.i32), ^[[dummyBlock:bb[0-9]+]](%arg2 : !llvm.i32)
-  cond_br %arg0, ^bb1(%arg1 : i32), ^bb1(%arg2 : i32)
-
-// CHECK:      ^[[origBlock]](%[[BLOCKARG1:.*]]: !llvm.i32):
-// CHECK-NEXT:  llvm.return %[[BLOCKARG1]] : !llvm.i32
-^bb1(%0 : i32):
-  return %0 : i32
-
-// CHECK:      ^[[dummyBlock]](%[[BLOCKARG2:.*]]: !llvm.i32):
-// CHECK-NEXT:  llvm.br ^[[origBlock]](%[[BLOCKARG2]] : !llvm.i32)
-}
 
 // CHECK-LABEL: func @fcmp(%arg0: !llvm.float, %arg1: !llvm.float) {
 func @fcmp(f32, f32) -> () {

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

@@ -0,0 +1,31 @@
+// RUN: mlir-opt -llvm-legalize-for-export %s | FileCheck %s
+
+// Verifies that duplicate successor with different arguments are deduplicated
+// by introducing a new block that forwards its arguments to the original
+// successor through an unconditional branch.
+// CHECK-LABEL: @repeated_successor_different_args
+llvm.func @repeated_successor_different_args(%arg0: !llvm.i1, %arg1: !llvm.i32, %arg2: !llvm.i32) {
+  // CHECK: llvm.cond_br %{{.*}}, ^[[BB1:.*]]({{.*}}), ^[[BB2:.*]]({{.*}})
+  llvm.cond_br %arg0, ^bb1(%arg1: !llvm.i32), ^bb1(%arg2: !llvm.i32)
+
+// CHECK: ^[[BB1]]({{.*}}):
+^bb1(%arg3: !llvm.i32):
+  llvm.return
+
+// CHECK: ^[[BB2]](%[[ARG:.*]]: !llvm.i32):
+// CHECK:  llvm.br ^[[BB1]](%[[ARG]] : !llvm.i32)
+}
+
+// Verifies that duplicate successors without arguments do not lead to the
+// introduction of new blocks during legalization.
+// CHECK-LABEL: @repeated_successor_no_args
+llvm.func @repeated_successor_no_args(%arg0: !llvm.i1) {
+  // CHECK: llvm.cond_br
+  llvm.cond_br %arg0, ^bb1, ^bb1
+
+// CHECK: ^{{.*}}:
+^bb1:
+  llvm.return
+
+// CHECK-NOT: ^{{.*}}:
+}