boilerplate for lowering func/cf/scf

j2kun · Sep 22, 2023 · 6e8a7aa · 6e8a7aa
1 parent 6779d11
commit 6e8a7aa
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/lib/Conversion/PolyToStandard/BUILD b/lib/Conversion/PolyToStandard/BUILD
@@ -31,6 +31,8 @@ cc_library(
         "pass_inc_gen",
         "//lib/Dialect/Poly",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncTransforms",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:TensorDialect",

diff --git a/lib/Conversion/PolyToStandard/PolyToStandard.cpp b/lib/Conversion/PolyToStandard/PolyToStandard.cpp
@@ -2,6 +2,8 @@
 
 #include "lib/Dialect/Poly/PolyOps.h"
 #include "lib/Dialect/Poly/PolyTypes.h"
+#include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project
+#include "mlir/include/mlir/Dialect/Func/Transforms/FuncConversions.h" // from @llvm-project
 #include "mlir/include/mlir/Transforms/DialectConversion.h" // from @llvm-project
 
 namespace mlir {
@@ -63,6 +65,29 @@ struct PolyToStandard : impl::PolyToStandardBase<PolyToStandard> {
     PolyToStandardTypeConverter typeConverter(context);
     patterns.add<ConvertAdd>(typeConverter, context);
 
+    populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
+        patterns, typeConverter);
+    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
+      return typeConverter.isSignatureLegal(op.getFunctionType()) &&
+             typeConverter.isLegal(&op.getBody());
+    });
+
+    populateReturnOpTypeConversionPattern(patterns, typeConverter);
+    target.addDynamicallyLegalOp<func::ReturnOp>(
+        [&](func::ReturnOp op) { return typeConverter.isLegal(op); });
+
+    populateCallOpTypeConversionPattern(patterns, typeConverter);
+    target.addDynamicallyLegalOp<func::CallOp>(
+        [&](func::CallOp op) { return typeConverter.isLegal(op); });
+
+    populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter);
+    target.markUnknownOpDynamicallyLegal([&](Operation *op) {
+      return isNotBranchOpInterfaceOrReturnLikeOp(op) ||
+             isLegalForBranchOpInterfaceTypeConversionPattern(op,
+                                                              typeConverter) ||
+             isLegalForReturnOpTypeConversionPattern(op, typeConverter);
+    });
+
     if (failed(applyPartialConversion(module, target, std::move(patterns)))) {
       signalPassFailure();
     }