[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

30

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

12 ↗

(On Diff #248207)

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/Legalize.h

+//===- Legalize.h - Make the LLVM dialect convertible 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
+//
+//===----------------------------------------------------------------------===//
+
+#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> createLegalizePass();
+
+} // namespace LLVM
+} // namespace mlir

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/Legalize.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::createLegalizePass();
+
   // 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/Legalize.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Value.h"
 
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Value.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));
     translator.convertGlobals();
     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) {
-  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.
-  for (const auto &successor : successorPositions) {
-    const auto &positions = successor.second;
-    // Start from the second occurrence of a block in the successor list.
-    for (auto position = std::next(positions.begin()), end = positions.end();
-         position != end; ++position) {
-      auto *dummyBlock = new Block();
-      bb.getParent()->push_back(dummyBlock);
-      auto builder = OpBuilder(dummyBlock);
-      SmallVector<Value, 8> operands(
-          terminator->getSuccessorOperands(*position));
-      builder.create<BranchOp>(terminator->getLoc(), successor.first, operands);
-      terminator->setSuccessor(dummyBlock, *position);
-      for (int i = 0, e = terminator->getNumSuccessorOperands(*position); i < e;
-           ++i)
-        terminator->eraseSuccessorOperand(*position, i);
-    }
-  }
-}
-
-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
   )
 add_dependencies(MLIRLLVMIR MLIRLLVMOpsIncGen MLIRLLVMConversionsIncGen LLVMAsmParser LLVMCore LLVMSupport)
 target_link_libraries(MLIRLLVMIR LLVMAsmParser LLVMCore LLVMSupport MLIRIR)

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

+add_mlir_dialect_library(MLIRLLVMIRTransforms
+  Legalize.cpp
+  )
+
+target_link_libraries(MLIRLLVMIRTransforms
+  MLIRIR
+  MLIRLLVMIR
+  MLIRPass
+  )

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

+//===- Legalize.cpp - Make the LLVM dialect convertible 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/Legalize.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.
+  for (const auto &successor : successorPositions) {
+    const auto &positions = successor.second;
+    // Start from the second occurrence of a block in the successor list.
+    for (auto position = std::next(positions.begin()), end = positions.end();
+         position != end; ++position) {
+      auto *dummyBlock = new Block();
+      bb.getParent()->push_back(dummyBlock);
+      auto builder = OpBuilder(dummyBlock);
+      SmallVector<Value, 8> operands(
+          terminator->getSuccessorOperands(*position));
+      builder.create<LLVM::BrOp>(terminator->getLoc(), successor.first,
+                                 operands);
+      terminator->setSuccessor(dummyBlock, *position);
+      for (int i = 0, e = terminator->getNumSuccessorOperands(*position); i < e;
+           ++i)
+        terminator->eraseSuccessorOperand(*position, i);
+    }
+  }
+}
+
+void mlir::LLVM::ensureDistinctSuccessors(Operation *op) {
+  op->walk([](LLVMFuncOp f) {
+    for (auto &bb : f.getBlocks()) {
+      ::ensureDistinctSuccessors(bb);
+    }
+  });
+}
+
+namespace {
+struct LLVMLegalizePass : public OperationPass<LLVMLegalizePass> {
+  void runOnOperation() override {
+    LLVM::ensureDistinctSuccessors(getOperation());
+  }
+};
+} // end namespace
+
+std::unique_ptr<Pass> LLVM::createLegalizePass() {
+  return std::make_unique<LLVMLegalizePass>();
+}
+
+static PassRegistration<LLVMLegalizePass>
+    pass("llvm-legalize", "Legalize LLVM dialect to be convertible to LLVM IR");