[mlir] Support FuncOpSignatureConversion for more FunctionLike ops.

[mikeurbach]

Summary

This extracts the implementation of getType, setType, and getBody from
FunctionSupport.h into the mlir::impl namespace and defines them
generically in FunctionSupport.cpp. This allows them to be used
elsewhere for any FunctionLike ops that use FunctionType for their
type signature.

Using the new helpers, FuncOpSignatureConversion is generalized to
work with all such FunctionLike ops. Convenience helpers are added to
configure the pattern for a given concrete FunctionLike op type.

[mikeurbach]

Remove unrelated change.

[mikeurbach]

This is the patch discussed with @silvas in https://llvm.discourse.group/t/bufferization-framework-replacement/2446/7. I proceeded here along the same lines as when I wanted to generalize the eraseArguments and eraseResults helpers from FuncOp to FunctionLike, but perhaps there is a better way to achieve this. CC'ing @rriddle and @ftynse who had helpful feedback when I was working on FunctionSupport.h before.

[rriddle]

mlir/lib/Transforms/Utils/DialectConversion.cpp

2524

This doesn't generalize to all FunctionLike operations though, just the ones that are similar to FuncOp (e.g. this would crash LLVMFuncOp). I'm +1 on having a simple default implementation that other FunctionLike operations can hook into, but I don't think it needs to match all of them. Could we instead just make it as simple as having the user provide the op name of the function like operation they want to use here?

[ftynse]

mlir/lib/Transforms/Utils/DialectConversion.cpp

2524

I would just make it a template <OpTy> struct FuncLikeSignatureConversion : public ConversionPattern<OpTy>. The class will probably move to the header, but we can keep the actual implementation in the impl namespace.

[mikeurbach]

Switch to parameterizing the pattern by op type.

[mikeurbach]

Update revision description to match commit message.

[mikeurbach]

Remove unused include.

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

832

I don't think you need to expose the pattern. Just the operation name should be enough, which is all that OpConversionPattern does for you.

[mikeurbach]

mlir/lib/Transforms/Utils/DialectConversion.cpp

2524

Thanks for the feedback. I think I have updated the implementation along the lines Alex mentioned. Please let me know what you think of this approach.

[mikeurbach]

mlir/include/mlir/Transforms/DialectConversion.h

832

Ah ok I wasn't sure about this. I will try it out.

[mikeurbach]

mlir/include/mlir/Transforms/DialectConversion.h

832

@rriddle maybe I am misunderstanding, but I tried re-working this, and I am not sure how it should work. When you say don't need to expose the pattern, do you mean the pattern can be declared in the cpp file instead of this header?

I tried putting the pattern back in the cpp file, and then I had to update the populateFunctionLikeTypeConversionPattern helper that I call from outside MLIR. This templated helper needs to be defined in this header, correct?

I tried just forward declaring the pattern so I could use it in populateFunctionLikeTypeConversionPattern, but there are some enable_ifs that need to know about the base class, and other compile time assumptions about the pattern's constructor. So then I tried declaring the pattern here in the header, but defining it in the cpp file, and that failed with an error like so:

ld.lld: error: undefined symbol: mlir::FunctionLikeSignatureConversion<circt::handshake::FuncOp>::matchAndRewrite(circt::handshake::FuncOp, llvm::ArrayRef<mlir::Value>, mlir::ConversionPatternRewriter&) const

It seems like I need to define the templated matchAndRewrite here in the header as well. If I do that, and shove as much as I can into the impl namespace in the cpp file, that is basically what is in the current revision.

Does this sound reasonable or am I missing something? I think this is what Alex was getting at previously with the comment "The class will probably move to the header". Maybe I am missing the point here, and there is a better way to do this... Thanks for bearing with me.

[rriddle]

mlir/include/mlir/Transforms/DialectConversion.h

832

Apologies for the terse response, kept getting distracted by other things. What I mean is that OpConversionPattern is just a utility wrapper around ConversionPattern, that only serves to provide the name of the operation. What I was suggesting is that we expose the minimal amount necessary, which is really just what operation we are matching given that the things needed for the transformation already have generic utilities available(e.g. getting/setting the function's type). For example, could we do something like:

// Header file

void populateFuncOpTypeConversionPattern(StringRef functionLikeOpName,
                                         OwningRewritePatternList &patterns,
                                         MLIRContext *ctx,
                                         TypeConverter &converter);

template <typename FuncOpT>
void populateFuncOpTypeConversionPattern(OwningRewritePatternList &patterns,
                                         MLIRContext *ctx,
                                         TypeConverter &converter) {
  populateFuncOpTypeConversionPattern(FuncOpT::getOperationName(), patterns, ctx, converter);
}

// Source File

namespace {
struct FuncOpSignatureConversion : public ConversionPattern {
  FuncOpSignatureConversion(StringRef funcOpName, MLIRContext *ctx, TypeConverter &converter)
      : ConversionPattern(funcOpName, /*benefit=*/1, converter, ctx) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    // Implementation
  }
};
} // end anonymous namespace

void mlir::populateFuncOpTypeConversionPattern(
    StringRef functionLikeOpName, OwningRewritePatternList &patterns, MLIRContext *ctx,
    TypeConverter &converter) {
  patterns.insert<FuncOpSignatureConversion>(functionLikeOpName, ctx, converter);
}

?

[mikeurbach]

mlir/include/mlir/Transforms/DialectConversion.h

832

Ah I see, thank you that is very helpful. I do think we can go that route.

[mikeurbach]

Updates after most recent feedback.

[mikeurbach]

Clean up the diff a bit.

[mikeurbach]

Update title and description to match commit message.

[mikeurbach]

Fix typo in commit message.

[rriddle]

LGTM. Thanks!

mlir/include/mlir/Transforms/DialectConversion.h

425

I would add a note that this only works for FunctionLike operations that use FunctionType to represent their type.

[mikeurbach]

Add note that the new functionality only supports FunctionLike ops that use FunctionType for their signature.

[mikeurbach]

Thanks @rriddle for the iterations on this. Looking forward to using this in CIRCT with our FunctionLike ops.

mlir/include/mlir/IR/FunctionSupport.h

 /// Erase the specified arguments and update the function type attribute.
 void eraseFunctionArguments(Operation *op, ArrayRef<unsigned> argIndices,
                             unsigned originalNumArgs, Type newType);
 
 /// Erase the specified results and update the function type attribute.
 void eraseFunctionResults(Operation *op, ArrayRef<unsigned> resultIndices,
                           unsigned originalNumResults, Type newType);
 
+/// Get and set a FunctionLike operation's type signature.
+FunctionType getFunctionType(Operation *op);
+void setFunctionType(Operation *op, FunctionType newType);
+
+/// Get a FunctionLike operation's body.
+Region &getFunctionBody(Operation *op);
+
 } // namespace impl
 
 namespace OpTrait {
 
 /// This trait provides APIs for Ops that behave like functions.  In particular:
 /// - Ops must be symbols, i.e. also have the `Symbol` trait;
 /// - Ops have a single region with multiple blocks that corresponds to the body
 ///   of the function;
public:
 
   //===--------------------------------------------------------------------===//
   // Body Handling
   //===--------------------------------------------------------------------===//
 
   /// Returns true if this function is external, i.e. it has no body.
   bool isExternal() { return empty(); }
 
-  Region &getBody() { return this->getOperation()->getRegion(0); }
+  Region &getBody() {
+    return ::mlir::impl::getFunctionBody(this->getOperation());
+  }
 
   /// Delete all blocks from this function.
   void eraseBody() {
     getBody().dropAllReferences();
     getBody().getBlocks().clear();
   }
 
   /// This is the list of blocks in the function.
public:
   }
 
   /// Return the type of this function.
   ///
   /// Note that the concrete class must define a method with the same name to
   /// hide this one if the concrete class does not use FunctionType for the
   /// function type under the hood.
   FunctionType getType() {
-    return getTypeAttr().getValue().template cast<FunctionType>();
+    return ::mlir::impl::getFunctionType(this->getOperation());
   }
 
   /// Return the type of this function without the specified arguments and
   /// results. This is used to update the function's signature in the
   /// `eraseArguments` and `eraseResults` methods. The arrays of indices are
   /// allowed to have duplicates and can be in any order.
   ///
   /// Note that the concrete class must define a method with the same name to

 }
 
 //===----------------------------------------------------------------------===//
 // Function Type Attribute.
 //===----------------------------------------------------------------------===//
 
 template <typename ConcreteType>
 void FunctionLike<ConcreteType>::setType(FunctionType newType) {
-  SmallVector<char, 16> nameBuf;
-  auto oldType = getType();
-  auto *concreteOp = static_cast<ConcreteType *>(this);
-
-  for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
-    concreteOp->removeAttr(getArgAttrName(i, nameBuf));
-  for (int i = newType.getNumResults(), e = oldType.getNumResults(); i < e; i++)
-    concreteOp->removeAttr(getResultAttrName(i, nameBuf));
-  (*concreteOp)->setAttr(getTypeAttrName(), TypeAttr::get(newType));
+  ::mlir::impl::setFunctionType(this->getOperation(), newType);
 }
 
 //===----------------------------------------------------------------------===//
 // Function Argument Attribute.
 //===----------------------------------------------------------------------===//
 
 /// Set the attributes held by the argument at 'index'.
 template <typename ConcreteType>

mlir/include/mlir/Transforms/DialectConversion.h

matchAndRewrite(SourceOp op, ArrayRef<Value> operands,
     rewrite(op, operands, rewriter);
     return success();
   }
 
 private:
   using ConversionPattern::matchAndRewrite;
 };
 
+/// Add a pattern to the given pattern list to convert the signature of a
+/// FunctionLike op with the given type converter. This only supports
+/// FunctionLike ops which use FunctionType to represent their type.
+void populateFunctionLikeTypeConversionPattern(
+    StringRef functionLikeOpName, OwningRewritePatternList &patterns,
+    MLIRContext *ctx, TypeConverter &converter);
+
+template <typename FuncOpT>
+void populateFunctionLikeTypeConversionPattern(
+    OwningRewritePatternList &patterns, MLIRContext *ctx,
+    TypeConverter &converter) {
+  populateFunctionLikeTypeConversionPattern(FuncOpT::getOperationName(),
+                                            patterns, ctx, converter);
+}
+
 /// Add a pattern to the given pattern list to convert the signature of a FuncOp
 /// with the given type converter.
 void populateFuncOpTypeConversionPattern(OwningRewritePatternList &patterns,
                                          MLIRContext *ctx,
                                          TypeConverter &converter);
 
 //===----------------------------------------------------------------------===//
 // Conversion PatternRewriter

 /// Apply a partial conversion on the given operations and all nested
 /// operations. This method converts as many operations to the target as
 /// possible, ignoring operations that failed to legalize. This method only
 /// returns failure if there ops explicitly marked as illegal. If an
 /// `unconvertedOps` set is provided, all operations that are found not to be
 /// legalizable to the given `target` are placed within that set. (Note that if
 /// there is an op explicitly marked as illegal, the conversion terminates and
 /// the `unconvertedOps` set will not necessarily be complete.)
 LLVM_NODISCARD LogicalResult
 applyPartialConversion(ArrayRef<Operation *> ops, ConversionTarget &target,
                        const FrozenRewritePatternList &patterns,
                        DenseSet<Operation *> *unconvertedOps = nullptr);
 LLVM_NODISCARD LogicalResult
 applyPartialConversion(Operation *op, ConversionTarget &target,
                        const FrozenRewritePatternList &patterns,
                        DenseSet<Operation *> *unconvertedOps = nullptr);
 

mlir/lib/IR/FunctionSupport.cpp

void mlir::impl::eraseFunctionResults(Operation *op,
   for (unsigned i = 0, e = newResultAttrs.size(); i != e; ++i) {
     auto nameAttr = getResultAttrName(i, nameBuf);
     if (newResultAttrs[i] && !newResultAttrs[i].empty())
       op->setAttr(nameAttr, newResultAttrs[i]);
     else
       op->removeAttr(nameAttr);
   }
 }
+
+//===----------------------------------------------------------------------===//
+// Function type signature.
+//===----------------------------------------------------------------------===//
+
+FunctionType mlir::impl::getFunctionType(Operation *op) {
+  assert(op->hasTrait<OpTrait::FunctionLike>());
+  return op->getAttrOfType<TypeAttr>(mlir::impl::getTypeAttrName())
+      .getValue()
+      .cast<FunctionType>();
+}
+
+void mlir::impl::setFunctionType(Operation *op, FunctionType newType) {
+  assert(op->hasTrait<OpTrait::FunctionLike>());
+  SmallVector<char, 16> nameBuf;
+  FunctionType oldType = getFunctionType(op);
+
+  for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
+    op->removeAttr(getArgAttrName(i, nameBuf));
+  for (int i = newType.getNumResults(), e = oldType.getNumResults(); i < e; i++)
+    op->removeAttr(getResultAttrName(i, nameBuf));
+  op->setAttr(getTypeAttrName(), TypeAttr::get(newType));
+}
+
+//===----------------------------------------------------------------------===//
+// Function body.
+//===----------------------------------------------------------------------===//
+
+Region &mlir::impl::getFunctionBody(Operation *op) {
+  assert(op->hasTrait<OpTrait::FunctionLike>());
+  return op->getRegion(0);
+}

mlir/lib/Transforms/Utils/DialectConversion.cpp

 //===- DialectConversion.cpp - MLIR dialect conversion generic pass -------===//
 //
 // 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/Transforms/DialectConversion.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/FunctionSupport.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 #include "mlir/Transforms/Utils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/SaveAndRestore.h"
 #include "llvm/Support/ScopedPrinter.h"
auto TypeConverter::convertBlockSignature(Block *block)
     -> Optional<SignatureConversion> {
   SignatureConversion conversion(block->getNumArguments());
   if (failed(convertSignatureArgs(block->getArgumentTypes(), conversion)))
     return llvm::None;
   return conversion;
 }
 
 /// Create a default conversion pattern that rewrites the type signature of a
-/// FuncOp.
+/// FunctionLike op. This only supports FunctionLike ops which use FunctionType
+/// to represent their type.
 namespace {
-struct FuncOpSignatureConversion : public OpConversionPattern<FuncOp> {
-  FuncOpSignatureConversion(MLIRContext *ctx, TypeConverter &converter)
-      : OpConversionPattern(converter, ctx) {}
+struct FunctionLikeSignatureConversion : public ConversionPattern {
+  FunctionLikeSignatureConversion(StringRef functionLikeOpName,
+                                  MLIRContext *ctx, TypeConverter &converter)
+      : ConversionPattern(functionLikeOpName, /*benefit=*/1, converter, ctx) {}
 
-  /// Hook for derived classes to implement combined matching and rewriting.
+  /// Hook to implement combined matching and rewriting for FunctionLike ops.
   LogicalResult
-  matchAndRewrite(FuncOp funcOp, ArrayRef<Value> operands,
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
-    FunctionType type = funcOp.getType();
+    FunctionType type = mlir::impl::getFunctionType(op);
 
     // 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.getBody(), *typeConverter,
-                                           &result)))
+        failed(rewriter.convertRegionTypes(&mlir::impl::getFunctionBody(op),
+                                           *typeConverter, &result)))
       return failure();
 
     // Update the function signature in-place.
-    rewriter.updateRootInPlace(funcOp, [&] {
-      funcOp.setType(FunctionType::get(funcOp.getContext(),
-                                       result.getConvertedTypes(), newResults));
-    });
+    auto newType = FunctionType::get(rewriter.getContext(),
+                                     result.getConvertedTypes(), newResults);
+
+    rewriter.updateRootInPlace(
+        op, [&] { mlir::impl::setFunctionType(op, newType); });
+
     return success();
   }
 };
 } // end anonymous namespace
 
+void mlir::populateFunctionLikeTypeConversionPattern(
+    StringRef functionLikeOpName, OwningRewritePatternList &patterns,
+    MLIRContext *ctx, TypeConverter &converter) {
+  patterns.insert<FunctionLikeSignatureConversion>(functionLikeOpName, ctx,
+                                                   converter);
+}
+
 void mlir::populateFuncOpTypeConversionPattern(
     OwningRewritePatternList &patterns, MLIRContext *ctx,
     TypeConverter &converter) {
-  patterns.insert<FuncOpSignatureConversion>(ctx, converter);
+  populateFunctionLikeTypeConversionPattern<FuncOp>(patterns, ctx, converter);
 }
 
 //===----------------------------------------------------------------------===//
 // ConversionTarget
 //===----------------------------------------------------------------------===//
 
 /// Register a legality action for the given operation.
 void ConversionTarget::setOpAction(OperationName op,