[MLIR][LLVMToControlFlow] Minimal LLVM to CF lifting passusers/dinistro/support-llvm-funcs-in-loop-lifting

This commit introduces a minimal example of a LLVM to CF lifting pass.

7 files changed, 149 insertions, 0 deletions

diff --git a/mlir/include/mlir/Conversion/LLVMToControlFlow/LLVMToControlFlow.h b/mlir/include/mlir/Conversion/LLVMToControlFlow/LLVMToControlFlow.h
new file mode 100644
index 0000000..deb3e41
--- /dev/null
+++ b/mlir/include/mlir/Conversion/LLVMToControlFlow/LLVMToControlFlow.h
@@ -0,0 +1,24 @@
+//===- LLVMToControlFlow.h - Convert LLVM to CF dialect ---------*- 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_CONVERSION_LLVMTOCONTROLFLOW_LLVMTOCONTROLFLOW_H
+#define MLIR_CONVERSION_LLVMTOCONTROLFLOW_LLVMTOCONTROLFLOW_H
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+class MLIRContext;
+class Pass;
+class RewritePatternSet;
+
+#define GEN_PASS_DECL_CONVERTLLVMTOCONTROLFLOW
+#include "mlir/Conversion/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_LLVMTOCONTROLFLOW_LLVMTOCONTROLFLOW_H
diff --git a/mlir/include/mlir/Conversion/Passes.h b/mlir/include/mlir/Conversion/Passes.h
index 2ab3283..5088011 100644
--- a/mlir/include/mlir/Conversion/Passes.h
+++ b/mlir/include/mlir/Conversion/Passes.h
@@ -42,6 +42,7 @@
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"
 #include "mlir/Conversion/IndexToSPIRV/IndexToSPIRV.h"
+#include "mlir/Conversion/LLVMToControlFlow/LLVMToControlFlow.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/MathToFuncs/MathToFuncs.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 4d272ba..ce0a816 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -1435,4 +1435,16 @@ def ConvertVectorToXeGPU : Pass<"convert-vector-to-xegpu"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// LLVMToControlFlow
+//===----------------------------------------------------------------------===//
+
+def ConvertLLVMToControlFlow : Pass<"convert-llvm-to-control-flow"> {
+  let summary = "Convert LLVM dialect to the control flow dialect";
+  let dependentDialects = [
+    "cf::ControlFlowDialect",
+    "LLVM::LLVMDialect"
+  ];
+}
+
 #endif // MLIR_CONVERSION_PASSES
diff --git a/mlir/lib/Conversion/CMakeLists.txt b/mlir/lib/Conversion/CMakeLists.txt
index 6651d87..be82c5e 100644
--- a/mlir/lib/Conversion/CMakeLists.txt
+++ b/mlir/lib/Conversion/CMakeLists.txt
@@ -33,6 +33,7 @@ add_subdirectory(IndexToLLVM)
 add_subdirectory(IndexToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(LLVMCommon)
+add_subdirectory(LLVMToControlFlow)
 add_subdirectory(MathToFuncs)
 add_subdirectory(MathToLibm)
 add_subdirectory(MathToLLVM)
diff --git a/mlir/lib/Conversion/LLVMToControlFlow/CMakeLists.txt b/mlir/lib/Conversion/LLVMToControlFlow/CMakeLists.txt
new file mode 100644
index 0000000..3fad4bd
--- /dev/null
+++ b/mlir/lib/Conversion/LLVMToControlFlow/CMakeLists.txt
@@ -0,0 +1,15 @@
+add_mlir_conversion_library(MLIRLLVMToControlFlow
+  LLVMToControlFlow.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/LLVMToControlFlow
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRControlFlowDialect
+  MLIRLLVMDialect
+  MLIRTransforms
+  )
+
diff --git a/mlir/lib/Conversion/LLVMToControlFlow/LLVMToControlFlow.cpp b/mlir/lib/Conversion/LLVMToControlFlow/LLVMToControlFlow.cpp
new file mode 100644
index 0000000..3fcb603
--- /dev/null
+++ b/mlir/lib/Conversion/LLVMToControlFlow/LLVMToControlFlow.cpp
@@ -0,0 +1,65 @@
+//===- LLVMToControlFlow.cpp - LLVM to CF conversion ----------------------===//
+//
+// 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/Conversion/LLVMToControlFlow/LLVMToControlFlow.h"
+
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace mlir {
+#define GEN_PASS_DEF_CONVERTLLVMTOCONTROLFLOW
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+namespace {
+
+struct BranchOpPattern : public OpConversionPattern<LLVM::BrOp> {
+  using OpConversionPattern<LLVM::BrOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(LLVM::BrOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<cf::BranchOp>(op, adaptor.getOperands(),
+                                              op.getSuccessor());
+    return success();
+  }
+};
+
+struct CondBranchOpPattern : public OpConversionPattern<LLVM::CondBrOp> {
+  using OpConversionPattern<LLVM::CondBrOp>::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(LLVM::CondBrOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<cf::CondBranchOp>(
+        op, adaptor.getCondition(), op.getTrueDest(),
+        adaptor.getTrueDestOperands(), op.getFalseDest(),
+        adaptor.getFalseDestOperands());
+    return success();
+  }
+};
+
+struct ConvertLLVMToControlFlowPass
+    : public impl::ConvertLLVMToControlFlowBase<ConvertLLVMToControlFlowPass> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void ConvertLLVMToControlFlowPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  patterns.add<BranchOpPattern, CondBranchOpPattern>(&getContext());
+  // Configure conversion to lower out SCF operations.
+  ConversionTarget target(getContext());
+  target.addIllegalOp<LLVM::BrOp, LLVM::CondBrOp>();
+  target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+  if (failed(
+          applyPartialConversion(getOperation(), target, std::move(patterns))))
+    signalPassFailure();
+}
diff --git a/mlir/test/Conversion/LLVMToControlFlow/convert-llvm-to-control-flow.mlir b/mlir/test/Conversion/LLVMToControlFlow/convert-llvm-to-control-flow.mlir
new file mode 100644
index 0000000..a471663
--- /dev/null
+++ b/mlir/test/Conversion/LLVMToControlFlow/convert-llvm-to-control-flow.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt -convert-llvm-to-control-flow -split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @branch
+// CHECK-SAME: %[[VAL:[[:alnum:]]+]]: i32
+llvm.func @branch(%val: i32) -> i32 {
+  // CHECK: cf.br ^[[BB1:.*]](%[[VAL]] : i32)
+  llvm.br ^bb1(%val : i32)
+// CHECK: ^[[BB1]](%[[BB_ARG:.*]]: i32):
+^bb1(%x: i32):
+  // CHECK: llvm.return %[[BB_ARG]] : i32
+  llvm.return %x : i32
+}
+
+// -----
+
+// CHECK-LABEL: @cond_branch
+// CHECK-SAME: %[[COND:[[:alnum:]]+]]: i1
+// CHECK-SAME: %[[VAL1:[[:alnum:]]+]]: i32
+// CHECK-SAME: %[[VAL2:[[:alnum:]]+]]: i32
+llvm.func @cond_branch(%cond: i1, %val1: i32, %val2: i32) -> i32 {
+  // CHECK: cf.cond_br %[[COND]], ^[[BB1:.*]](%[[VAL1]] : i32), ^[[BB2:.*]](%[[VAL2]] : i32)
+  llvm.cond_br %cond, ^bb1(%val1 : i32), ^bb2(%val2 : i32)
+// CHECK: ^[[BB1]](%[[BB1_ARG:.*]]: i32):
+^bb1(%x: i32):
+  // CHECK: llvm.return %[[BB1_ARG]] : i32
+  llvm.return %x : i32
+// CHECK: ^[[BB2]](%[[BB2_ARG:.*]]: i32):
+^bb2(%y: i32):
+  // CHECK: llvm.return %[[BB2_ARG]] : i32
+  llvm.return %y : i32
+}