[mlir] Add `populateFunctionOpInterfaceTypeConversionPattern` version which operates on any `FunctionOpInterface`

[Hardcode84]

Summary

Exisitng version is always limited to some specific op.

[Hardcode84]

Do we still need a version operating on some specific op? If no, I can remove it completely.

[rriddle]

In D137469#3909397, @Hardcode84 wrote:

Do we still need a version operating on some specific op? If no, I can remove it completely.

I would say yes (at least for now). I would only expect people to use the interface variant if they might actually have multiple types of functions (where they want all to be converted).

mlir/lib/Transforms/Utils/DialectConversion.cpp

3101–3118

Can we share the implementation here? It's a bit sad to duplicate.

[Hardcode84]

remove copypaste

mlir/include/mlir/Transforms/DialectConversion.h

 
 template <typename FuncOpT>
 void populateFunctionOpInterfaceTypeConversionPattern(
     RewritePatternSet &patterns, TypeConverter &converter) {
   populateFunctionOpInterfaceTypeConversionPattern(FuncOpT::getOperationName(),
                                                    patterns, converter);
 }
 
+void populateAnyFunctionOpInterfaceTypeConversionPattern(
+    RewritePatternSet &patterns, TypeConverter &converter);
+
 //===----------------------------------------------------------------------===//
 // Conversion PatternRewriter
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 struct ConversionPatternRewriterImpl;
 } // namespace detail
 

mlir/lib/Transforms/Utils/DialectConversion.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto newType = FunctionType::get(rewriter.getContext(),
                                      result.getConvertedTypes(), newResults);
 
     rewriter.updateRootInPlace(op, [&] { funcOp.setType(newType); });
 
     return success();
   }
 };
+
+struct AnyFunctionOpInterfaceSignatureConversion
+    : public OpInterfaceConversionPattern<FunctionOpInterface> {
+  using OpInterfaceConversionPattern::OpInterfaceConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(FunctionOpInterface funcOp, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    FunctionType type = funcOp.getFunctionType().cast<FunctionType>();
+
+    // Convert the original function types.
+    TypeConverter::SignatureConversion result(type.getNumInputs());
+    SmallVector<Type, 1> newResults;
+    if (failed(typeConverter->convertSignatureArgs(type.getInputs(), result)) ||
+        failed(typeConverter->convertTypes(type.getResults(), newResults)) ||
+        failed(rewriter.convertRegionTypes(&funcOp.getFunctionBody(),
+                                           *typeConverter, &result)))
+      return failure();
+
+    // Update the function signature in-place.
+    auto newType = FunctionType::get(rewriter.getContext(),
+                                     result.getConvertedTypes(), newResults);
+
+    rewriter.updateRootInPlace(funcOp, [&] { funcOp.setType(newType); });
+
+    return success();
+  }
+};
 } // namespace
 
 void mlir::populateFunctionOpInterfaceTypeConversionPattern(
     StringRef functionLikeOpName, RewritePatternSet &patterns,
     TypeConverter &converter) {
   patterns.add<FunctionOpInterfaceSignatureConversion>(
       functionLikeOpName, patterns.getContext(), converter);
 }
 
+void mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(
+    RewritePatternSet &patterns, TypeConverter &converter) {
+  patterns.add<AnyFunctionOpInterfaceSignatureConversion>(
+      converter, patterns.getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // ConversionTarget
 //===----------------------------------------------------------------------===//
 
 void ConversionTarget::setOpAction(OperationName op,
                                    LegalizationAction action) {
   legalOperations[op].action = action;
 }

mlir/test/lib/Dialect/Test/TestPatterns.cpp

patterns
              TestCreateBlock, TestCreateIllegalBlock, TestUndoBlockArgReplace,
              TestUndoBlockErase, TestPassthroughInvalidOp, TestSplitReturnType,
              TestChangeProducerTypeI32ToF32, TestChangeProducerTypeF32ToF64,
              TestChangeProducerTypeF32ToInvalid, TestUpdateConsumerType,
              TestNonRootReplacement, TestBoundedRecursiveRewrite,
              TestNestedOpCreationUndoRewrite, TestReplaceEraseOp,
              TestCreateUnregisteredOp>(&getContext());
     patterns.add<TestDropOpSignatureConversion>(&getContext(), converter);
-    mlir::populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
-        patterns, converter);
+    mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns,
+                                                              converter);
     mlir::populateCallOpTypeConversionPattern(patterns, converter);
 
     // Define the conversion target used for the test.
     ConversionTarget target(getContext());
     target.addLegalOp<ModuleOp>();
     target.addLegalOp<LegalOpA, LegalOpB, LegalOpC, TestCastOp, TestValidOp,
                       TerminatorOp>();
     target
void runOnOperation() override {
 
     // Initialize the set of rewrite patterns.
     RewritePatternSet patterns(&getContext());
     patterns.add<TestTypeConsumerForward, TestTypeConversionProducer,
                  TestSignatureConversionUndo,
                  TestTestSignatureConversionNoConverter>(converter,
                                                          &getContext());
     patterns.add<TestTypeConversionAnotherProducer>(&getContext());
-    mlir::populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
-        patterns, converter);
+    mlir::populateAnyFunctionOpInterfaceTypeConversionPattern(patterns,
+                                                              converter);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace