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

[ftynse]

Summary

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.

[ftynse]

rebase

[ftynse]

more cmake

[rriddle]

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

29 ↗

(On Diff #248207)

Can we give this a better name? Legalize seems general, and doesn't really hint that this is just for exporting to LLVMIR. LegalizeForExport seems slightly better, but I'm sure there is a better suggestion.

[stephenneuendorffer]

I tend to agree with riddle that having a better passname. Maybe "LegalizeForLLVM" ?

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

8

PUBLIC here please.

mlir/lib/Target/CMakeLists.txt

11

PUBLIC here please.

[ftynse]

Rename Legalize -> LegalizeForExport and rebase.

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

 /// Creates a pass to convert the Standard dialect into the LLVMIR dialect.
 /// By default stdlib malloc/free are used for allocating MemRef payloads.
 /// Specifying `useAlloca-true` emits stack allocations instead. In the future
 /// this may become an enum when we have concrete uses for other options.
 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

+//===- 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

 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LoopsToGPU/LoopsToGPUPass.h"
 #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"
 #include "mlir/Dialect/SPIRV/Passes.h"
 #include "mlir/Quantizer/Transforms/Passes.h"
 #include "mlir/Transforms/LocationSnapshot.h"
 #include "mlir/Transforms/Passes.h"
 
#endif
   createLinalgTilingPass();
   createLinalgTilingToParallelLoopsPass();
   createLinalgPromotionPass(0);
   createConvertLinalgToLoopsPass();
   createConvertLinalgToParallelLoopsPass();
   createConvertLinalgToAffineLoopsPass();
   createConvertLinalgToLLVMPass();
 
+  // LLVM
+  LLVM::createLegalizeForExportPass();
+
   // LoopOps
   createParallelLoopFusionPass();
   createParallelLoopSpecializationPass();
   createParallelLoopTilingPass();
 
   // QuantOps
   quant::createConvertSimulatedQuantPass();
   quant::createConvertConstPass();

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

 // the corresponding LLVMIR module. It only handles core LLVM IR operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TARGET_LLVMIR_MODULETRANSLATION_H
 #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"
 
 #include "llvm/Frontend/OpenMP/OMPIRBuilder.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
static std::unique_ptr<llvm::Module> translateModule(Operation *m) {
     if (!satisfiesLLVMModule(m))
       return nullptr;
     if (failed(checkSupportedModuleOps(m)))
       return nullptr;
     auto llvmModule = prepareLLVMModule(m);
     if (!llvmModule)
       return nullptr;
 
+    LLVM::ensureDistinctSuccessors(m);
+
     T translator(m, std::move(llvmModule));
     if (failed(translator.convertGlobals()))
       return nullptr;
     if (failed(translator.convertFunctions()))
       return nullptr;
 
     return std::move(translator.llvmModule);
   }

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     rewriter.replaceOp(op, {newLoaded});
 
     return matchSuccess();
   }
 };
 
 } // 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) {
   // FIXME: this should be tablegen'ed
   // clang-format off
   patterns.insert<
       AbsFOpLowering,
       AddFOpLowering,
if (useBarePtrCallConv && emitCWrappers) {
       getModule().emitError()
           << "incompatible conversion options: bare-pointer calling convention "
              "and C wrapper emission";
       signalPassFailure();
       return;
     }
 
     ModuleOp m = getModule();
-    LLVM::ensureDistinctSuccessors(m);
 
     LLVMTypeConverterCustomization customs;
     customs.funcArgConverter = useBarePtrCallConv ? barePtrFuncArgTypeConverter
                                                   : structFuncArgTypeConverter;
     LLVMTypeConverter typeConverter(&getContext(), customs);
 
     OwningRewritePatternList patterns;
     if (useBarePtrCallConv)

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

+add_subdirectory(Transforms)
+
 add_mlir_dialect_library(MLIRLLVMIR
   IR/LLVMDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
 
   DEPENDS
   MLIRLLVMOpsIncGen

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

+add_mlir_dialect_library(MLIRLLVMIRTransforms
+  LegalizeForExport.cpp
+  )
+
+target_link_libraries(MLIRLLVMIRTransforms
+  PUBLIC
+  MLIRIR
+  MLIRLLVMIR
+  MLIRPass
+  )

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

+//===- 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

 add_mlir_library(MLIRTargetLLVMIRModuleTranslation
   LLVMIR/DebugTranslation.cpp
   LLVMIR/ModuleTranslation.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVMIR
   DEPENDS
   intrinsics_gen
   )
 target_link_libraries(MLIRTargetLLVMIRModuleTranslation
   PUBLIC
   MLIRLLVMIR
+  MLIRLLVMIRTransforms
   LLVMCore
   LLVMIRReader
   LLVMSupport
   LLVMTransformUtils
   MLIRTranslation
   )
 
 add_mlir_library(MLIRTargetLLVMIR

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

^bb1:
   %2 = addi %arg0, %0 : i32
   return %2 : i32
 
 // CHECK-NEXT: ^bb2:
 ^bb2:
 // 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) -> () {
 ^bb0(%arg0: f32, %arg1: f32):
   // CHECK:      llvm.fcmp "oeq" %arg0, %arg1 : !llvm.float
   // CHECK-NEXT: llvm.fcmp "ogt" %arg0, %arg1 : !llvm.float
   // CHECK-NEXT: llvm.fcmp "oge" %arg0, %arg1 : !llvm.float
   // CHECK-NEXT: llvm.fcmp "olt" %arg0, %arg1 : !llvm.float

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

+// 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: ^{{.*}}:
+}