[mlir][CF] Split cf-to-llvm from func-to-llvm

Author: matthias-springer

mlir/include/mlir/Conversion/Passes.td

@@ -460,10 +460,6 @@ def ConvertFuncToLLVMPass : Pass<"convert-func-to-llvm", "ModuleOp"> {
     1 value is returned, packed into an LLVM IR struct type. Function calls and
     returns are updated accordingly. Block argument types are updated to use
     LLVM IR types.
-
-    Note that until https://github.com/llvm/llvm-project/issues/70982 is resolved,
-    this pass includes patterns that lower `arith` and `cf` to LLVM. This is legacy
-    code due to when they were all converted in the same pass.
   }];
   let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [

mlir/lib/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.cpp

@@ -94,106 +94,111 @@ struct AssertOpLowering : public ConvertOpToLLVMPattern<cf::AssertOp> {
   bool abortOnFailedAssert = true;
 };
 
-/// The cf->LLVM lowerings for branching ops require that the blocks they jump
-/// to first have updated types which should be handled by a pattern operating
-/// on the parent op.
-static LogicalResult verifyMatchingValues(ConversionPatternRewriter &rewriter,
-                                          ValueRange operands,
-                                          ValueRange blockArgs, Location loc,
-                                          llvm::StringRef messagePrefix) {
-  for (const auto &idxAndTypes :
-       llvm::enumerate(llvm::zip(blockArgs, operands))) {
-    int64_t i = idxAndTypes.index();
-    Value argValue =
-        rewriter.getRemappedValue(std::get<0>(idxAndTypes.value()));
-    Type operandType = std::get<1>(idxAndTypes.value()).getType();
-    // In the case of an invalid jump, the block argument will have been
-    // remapped to an UnrealizedConversionCast. In the case of a valid jump,
-    // there might still be a no-op conversion cast with both types being equal.
-    // Consider both of these details to see if the jump would be invalid.
-    if (auto op = dyn_cast_or_null<UnrealizedConversionCastOp>(
-            argValue.getDefiningOp())) {
-      if (op.getOperandTypes().front() != operandType) {
-        return rewriter.notifyMatchFailure(loc, [&](Diagnostic &diag) {
-          diag << messagePrefix;
-          diag << "mismatched types from operand # " << i << " ";
-          diag << operandType;
-          diag << " not compatible with destination block argument type ";
-          diag << op.getOperandTypes().front();
-          diag << " which should be converted with the parent op.";
-        });
-      }
-    }
-  }
-  return success();
+/// Helper function for converting branch ops. This function converts the
+/// signature of the given block. If the new block signature is different from
+/// `expectedTypes`, returns "failure".
+static FailureOr<Block *> getConvertedBlock(ConversionPatternRewriter &rewriter,
+                                            const TypeConverter *converter,
+                                            Operation *branchOp, Block *block,
+                                            TypeRange expectedTypes) {
+  assert(converter && "expected non-null type converter");
+  assert(!block->isEntryBlock() && "entry blocks have no predecessors");
+
+  // There is nothing to do if the types already match.
+  if (block->getArgumentTypes() == expectedTypes)
+    return block;
+
+  // Compute the new block argument types and convert the block.
+  std::optional<TypeConverter::SignatureConversion> conversion =
+      converter->convertBlockSignature(block);
+  if (!conversion)
+    return rewriter.notifyMatchFailure(branchOp,
+                                       "could not compute block signature");
+  if (expectedTypes != conversion->getConvertedTypes())
+    return rewriter.notifyMatchFailure(
+        branchOp,
+        "mismatch between adaptor operand types and computed block signature");
+  return rewriter.applySignatureConversion(block, *conversion, converter);
 }
 
-/// Ensure that all block types were updated and then create an LLVM::BrOp
+/// Convert the destination block signature (if necessary) and lower the branch
+/// op to llvm.br.
 struct BranchOpLowering : public ConvertOpToLLVMPattern<cf::BranchOp> {
   using ConvertOpToLLVMPattern<cf::BranchOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(cf::BranchOp op, typename cf::BranchOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (failed(verifyMatchingValues(rewriter, adaptor.getDestOperands(),
-                                    op.getSuccessor()->getArguments(),
-                                    op.getLoc(),
-                                    /*messagePrefix=*/"")))
+    FailureOr<Block *> convertedBlock =
+        getConvertedBlock(rewriter, getTypeConverter(), op, op.getSuccessor(),
+                          TypeRange(adaptor.getOperands()));
+    if (failed(convertedBlock))
       return failure();
-
-    rewriter.replaceOpWithNewOp<LLVM::BrOp>(
-        op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs());
+    rewriter.replaceOpWithNewOp<LLVM::BrOp>(op, adaptor.getOperands(),
+                                            *convertedBlock);
     return success();
   }
 };
 
-/// Ensure that all block types were updated and then create an LLVM::CondBrOp
+/// Convert the destination block signatures (if necessary) and lower the
+/// branch op to llvm.cond_br.
 struct CondBranchOpLowering : public ConvertOpToLLVMPattern<cf::CondBranchOp> {
   using ConvertOpToLLVMPattern<cf::CondBranchOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(cf::CondBranchOp op,
                   typename cf::CondBranchOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (failed(verifyMatchingValues(rewriter, adaptor.getFalseDestOperands(),
-                                    op.getFalseDest()->getArguments(),
-                                    op.getLoc(), "in false case branch ")))
+    FailureOr<Block *> convertedTrueBlock =
+        getConvertedBlock(rewriter, getTypeConverter(), op, op.getTrueDest(),
+                          TypeRange(adaptor.getTrueDestOperands()));
+    if (failed(convertedTrueBlock))
       return failure();
-    if (failed(verifyMatchingValues(rewriter, adaptor.getTrueDestOperands(),
-                                    op.getTrueDest()->getArguments(),
-                                    op.getLoc(), "in true case branch ")))
+    FailureOr<Block *> convertedFalseBlock =
+        getConvertedBlock(rewriter, getTypeConverter(), op, op.getFalseDest(),
+                          TypeRange(adaptor.getFalseDestOperands()));
+    if (failed(convertedFalseBlock))
       return failure();
-
     rewriter.replaceOpWithNewOp<LLVM::CondBrOp>(
-        op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs());
+        op, adaptor.getCondition(), *convertedTrueBlock,
+        adaptor.getTrueDestOperands(), *convertedFalseBlock,
+        adaptor.getFalseDestOperands());
     return success();
   }
 };
 
-/// Ensure that all block types were updated and then create an LLVM::SwitchOp
+/// Convert the destination block signatures (if necessary) and lower the
+/// switch op to llvm.switch.
 struct SwitchOpLowering : public ConvertOpToLLVMPattern<cf::SwitchOp> {
   using ConvertOpToLLVMPattern<cf::SwitchOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(cf::SwitchOp op, typename cf::SwitchOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    if (failed(verifyMatchingValues(rewriter, adaptor.getDefaultOperands(),
-                                    op.getDefaultDestination()->getArguments(),
-                                    op.getLoc(), "in switch default case ")))
+    // Get or convert default block.
+    FailureOr<Block *> convertedDefaultBlock = getConvertedBlock(
+        rewriter, getTypeConverter(), op, op.getDefaultDestination(),
+        TypeRange(adaptor.getDefaultOperands()));
+    if (failed(convertedDefaultBlock))
       return failure();
 
-    for (const auto &i : llvm::enumerate(
-             llvm::zip(adaptor.getCaseOperands(), op.getCaseDestinations()))) {
-      if (failed(verifyMatchingValues(
-              rewriter, std::get<0>(i.value()),
-              std::get<1>(i.value())->getArguments(), op.getLoc(),
-              "in switch case " + std::to_string(i.index()) + " "))) {
+    // Get or convert all case blocks.
+    SmallVector<Block *> caseDestinations;
+    SmallVector<ValueRange> caseOperands = adaptor.getCaseOperands();
+    for (auto it : llvm::enumerate(op.getCaseDestinations())) {
+      Block *b = it.value();
+      FailureOr<Block *> convertedBlock =
+          getConvertedBlock(rewriter, getTypeConverter(), op, b,
+                            TypeRange(caseOperands[it.index()]));
+      if (failed(convertedBlock))
         return failure();
-      }
+      caseDestinations.push_back(*convertedBlock);
     }
 
     rewriter.replaceOpWithNewOp<LLVM::SwitchOp>(
-        op, adaptor.getOperands(), op->getSuccessors(), op->getAttrs());
+        op, adaptor.getFlag(), *convertedDefaultBlock,
+        adaptor.getDefaultOperands(), adaptor.getCaseValuesAttr(),
+        caseDestinations, caseOperands);
     return success();
   }
 };
@@ -230,14 +235,22 @@ struct ConvertControlFlowToLLVM
 
   /// Run the dialect converter on the module.
   void runOnOperation() override {
-    LLVMConversionTarget target(getContext());
-    RewritePatternSet patterns(&getContext());
-
-    LowerToLLVMOptions options(&getContext());
+    MLIRContext *ctx = &getContext();
+    LLVMConversionTarget target(*ctx);
+    // This pass lowers only CF dialect ops, but it also modifies block
+    // signatures inside other ops. These ops should be treated as legal. They
+    // are lowered by other passes.
+    target.markUnknownOpDynamicallyLegal([&](Operation *op) {
+      return op->getDialect() !=
+             ctx->getLoadedDialect<cf::ControlFlowDialect>();
+    });
+
+    LowerToLLVMOptions options(ctx);
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
 
-    LLVMTypeConverter converter(&getContext(), options);
+    LLVMTypeConverter converter(ctx, options);
+    RewritePatternSet patterns(ctx);
     mlir::cf::populateControlFlowToLLVMConversionPatterns(converter, patterns);
 
     if (failed(applyPartialConversion(getOperation(), target,

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

@@ -432,11 +432,11 @@ mlir::convertFuncOpToLLVMFuncOp(FunctionOpInterface funcOp,
 
   rewriter.inlineRegionBefore(funcOp.getFunctionBody(), newFuncOp.getBody(),
                               newFuncOp.end());
-  if (failed(rewriter.convertRegionTypes(&newFuncOp.getBody(), converter,
-                                         &result))) {
-    return rewriter.notifyMatchFailure(funcOp,
-                                       "region types conversion failed");
-  }
+  // Convert just the entry block. The remaining unstructured control flow is
+  // converted by ControlFlowToLLVM.
+  if (!newFuncOp.getBody().empty())
+    rewriter.applySignatureConversion(&newFuncOp.getBody().front(), result,
+                                      &converter);
 
   // Fix the type mismatch between the materialized `llvm.ptr` and the expected
   // pointee type in the function body when converting `llvm.byval`/`llvm.byref`
@@ -785,10 +785,6 @@ struct ConvertFuncToLLVMPass
     RewritePatternSet patterns(&getContext());
     populateFuncToLLVMConversionPatterns(typeConverter, patterns, symbolTable);
 
-    // TODO(https://github.com/llvm/llvm-project/issues/70982): Remove these in
-    // favor of their dedicated conversion passes.
-    cf::populateControlFlowToLLVMConversionPatterns(typeConverter, patterns);
-
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, std::move(patterns))))
       signalPassFailure();

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

@@ -9,6 +9,7 @@
 #include "mlir/Dialect/SparseTensor/Pipelines/Passes.h"
 
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
+#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/Passes.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
@@ -91,6 +92,7 @@ void mlir::sparse_tensor::buildSparsifier(OpPassManager &pm,
       createConvertVectorToLLVMPass(options.convertVectorToLLVMOptions()));
   pm.addPass(createConvertFuncToLLVMPass());
   pm.addPass(createArithToLLVMConversionPass());
+  pm.addPass(createConvertControlFlowToLLVMPass());
 
   // Finalize GPU code generation.
   if (gpuCodegen) {

mlir/test/Conversion/ControlFlowToLLVM/branch.mlir

@@ -0,0 +1,69 @@
+// RUN: mlir-opt %s -convert-cf-to-llvm -split-input-file | FileCheck %s
+
+// Unstructured control flow is converted, but the enclosing op is not
+// converted.
+
+// CHECK-LABEL: func.func @cf_br(
+//  CHECK-SAME:     %[[arg0:.*]]: index) -> index {
+//       CHECK:   %[[cast0:.*]] = builtin.unrealized_conversion_cast %[[arg0]] : index to i64
+//       CHECK:   llvm.br ^[[bb1:.*]](%[[cast0]] : i64)
+//       CHECK: ^[[bb1]](%[[arg1:.*]]: i64):
+//       CHECK:   %[[cast1:.*]] = builtin.unrealized_conversion_cast %[[arg1]] : i64 to index
+//       CHECK:   return %[[cast1]] : index
+//       CHECK: }
+func.func @cf_br(%arg0: index) -> index {
+  cf.br ^bb1(%arg0 : index)
+^bb1(%arg1: index):
+  return %arg1 : index
+}
+
+// -----
+
+// func.func and func.return types match. No unrealized_conversion_cast is
+// needed.
+
+// CHECK-LABEL: func.func @cf_br_type_match(
+//  CHECK-SAME:     %[[arg0:.*]]: i64) -> i64 {
+//       CHECK:   llvm.br ^[[bb1:.*]](%[[arg0:.*]] : i64)
+//       CHECK: ^[[bb1]](%[[arg1:.*]]: i64):
+//       CHECK:   return %[[arg1]] : i64
+//       CHECK: }
+func.func @cf_br_type_match(%arg0: i64) -> i64 {
+  cf.br ^bb1(%arg0 : i64)
+^bb1(%arg1: i64):
+  return %arg1 : i64
+}
+
+// -----
+
+// Test case for cf.cond_br.
+
+//   CHECK-LABEL: func.func @cf_cond_br
+// CHECK-COUNT-2:   unrealized_conversion_cast {{.*}} : index to i64
+//         CHECK:   llvm.cond_br %{{.*}}, ^{{.*}}(%{{.*}} : i64), ^{{.*}}(%{{.*}} : i64)
+//         CHECK: ^{{.*}}(%{{.*}}: i64):
+//         CHECK:   unrealized_conversion_cast {{.*}} : i64 to index
+//         CHECK: ^{{.*}}(%{{.*}}: i64):
+//         CHECK:   unrealized_conversion_cast {{.*}} : i64 to index
+func.func @cf_cond_br(%cond: i1, %a: index, %b: index) -> index {
+  cf.cond_br %cond, ^bb1(%a : index), ^bb2(%b : index)
+^bb1(%arg1: index):
+  return %arg1 : index
+^bb2(%arg2: index):
+  return %arg2 : index
+}
+
+// -----
+
+// Unreachable block (and IR in general) is not converted during a dialect
+// conversion.
+
+// CHECK-LABEL: func.func @unreachable_block()
+//       CHECK:   return
+//       CHECK: ^[[bb1:.*]](%[[arg0:.*]]: index):
+//       CHECK:   cf.br ^[[bb1]](%[[arg0]] : index)
+func.func @unreachable_block() {
+  return
+^bb1(%arg0: index):
+  cf.br ^bb1(%arg0 : index)
+}