[mlir] Argument and result attribute handling during inlining.

[gysit]

Summary

The revision adds the handleArgument and handleResult handlers that
allow users of the inlining interface to implement argument and result
conversions that take argument and result attributes into account. The
motivating use cases for this revision are taken from the LLVM dialect
inliner, which has to copy arguments that are marked as byval and that
also has to consider zeroext / signext when converting integers.

All type conversions are currently handled by the
materializeCallConversion hook. It runs before isLegalToInline and
supports only the introduction of a single cast operation since it may
have to rollback. The new handlers run shortly before and after
inlining and cannot fail. As a result, they can introduce more complex
ir such as copying a struct argument. At the moment, the new hooks
cannot be used to perform type conversions since all type conversions
have to be done using the materializeCallConversion. A follow up
revision will either relax this constraint or drop
materializeCallConversion in favor of the new and more flexible
handlers.

The revision also extends the CallableOpInterface to provide access
to the argument and result attributes if available.

[Dinistro]

Dropped some comments.

mlir/lib/Transforms/Utils/InliningUtils.cpp

178

A comment suggest that the dictionary attributes passed to the handleArgument are non-null. Thus, I would prefer to use cast here, as it will assert when this cast fails and not silently return null and then violate the pre-condition of handleArgument.

207

As above.

[gysit]

Address review comments and rebase.

[Dinistro]

LGTM!

[mehdi_amini]

Please get reviews from @rriddle here.

[rriddle]

Patch looks reasonable to me.

mlir/include/mlir/Interfaces/CallInterfaces.td

94–108

What's the benefit of a default implementation here? Why not force users to provide this? (I assume most will)

mlir/lib/Transforms/Utils/InliningUtils.cpp

110

What's the tradeoff of having this anchored on the call vs the callable?

168–169

I'd hoist these checks out, they make a bit more sense in the caller.

198–199

Same here

[gysit]

Address review comments.

[gysit]

Thanks for the review!

mlir/include/mlir/Interfaces/CallInterfaces.td

94–108

Yes makes sense!

mlir/lib/Transforms/Utils/InliningUtils.cpp

110

It is not completely clear to me which one makes more sense here. I used the call since the existing type conversion hook did so. Another argument in-favor of using the call is that the newly created IR is introduced in-place of the call operation. On the other hand, the callable carries the argument and result attributes. I think the later is a pretty strong argument for anchoring on the callable. I updated the revision accordingly. Let me know if you think it is not an improvement.

[gysit]

Fix toy dialect.

[rriddle]

Still, LGTM (feel free to land).

mlir/lib/Transforms/Utils/InliningUtils.cpp

110

LGTM. I was thinking the callable as well given it is the one that should have an indication of what the attributes are, though if we really need a different handling that can be added as necessary.

[gysit]

In D145582#4210671, @rriddle wrote:

Still, LGTM (feel free to land).

I was waiting due to the blocking review of @jpienaar. Is it ok to land nevertheless?

[rriddle]

In D145582#4210683, @gysit wrote:

In D145582#4210671, @rriddle wrote:

Still, LGTM (feel free to land).

I was waiting due to the blocking review of @jpienaar. Is it ok to land nevertheless?

Yeah, it should be fine in this case.

[gysit]

Rebase and implement new methods on a newly added callable op.

mlir/docs/Interfaces.md

 
 ##### CallInterfaces
 
 *   `CallOpInterface` - Used to represent operations like 'call'
     -   `CallInterfaceCallable getCallableForCallee()`
 *   `CallableOpInterface` - Used to represent the target callee of call.
     -   `Region * getCallableRegion()`
     -   `ArrayRef<Type> getCallableResults()`
+    -   `ArrayAttr getCallableArgAttrs()`
+    -   `ArrayAttr getCallableResAttrs()`
 
 ##### RegionKindInterfaces
 
 *   `RegionKindInterface` - Used to describe the abstract semantics of regions.
     -   `RegionKind getRegionKind(unsigned index)` - Return the kind of the
         region with the given index inside this operation.
         -   RegionKind::Graph - represents a graph region without control flow
             semantics

mlir/docs/Tutorials/Toy/Ch-4.md

 ```c++
 /// Returns the region on the function operation that is callable.
 Region *FuncOp::getCallableRegion() { return &getBody(); }
 
 /// Returns the results types that the callable region produces when
 /// executed.
 ArrayRef<Type> FuncOp::getCallableResults() { return getType().getResults(); }
 
+/// Returns the argument attributes for all callable region arguments or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+/// Returns the result attributes for all callable region results or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 // ....
 
 /// Return the callee of the generic call operation, this is required by the
 /// call interface.
 CallInterfaceCallable GenericCallOp::getCallableForCallee() {
   return getAttrOfType<SymbolRefAttr>("callee");
 }
 

mlir/examples/toy/Ch4/mlir/Dialect.cpp

 mlir::Region *FuncOp::getCallableRegion() { return &getBody(); }
 
 /// Returns the results types that the callable region produces when
 /// executed.
 llvm::ArrayRef<mlir::Type> FuncOp::getCallableResults() {
   return getFunctionType().getResults();
 }
 
+/// Returns the argument attributes for all callable region arguments or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+/// Returns the result attributes for all callable region results or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // GenericCallOp
 //===----------------------------------------------------------------------===//
 
 void GenericCallOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                           StringRef callee, ArrayRef<mlir::Value> arguments) {
   // Generic call always returns an unranked Tensor initially.
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));

mlir/examples/toy/Ch5/mlir/Dialect.cpp

 mlir::Region *FuncOp::getCallableRegion() { return &getBody(); }
 
 /// Returns the results types that the callable region produces when
 /// executed.
 llvm::ArrayRef<mlir::Type> FuncOp::getCallableResults() {
   return getFunctionType().getResults();
 }
 
+/// Returns the argument attributes for all callable region arguments or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+/// Returns the result attributes for all callable region results or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // GenericCallOp
 //===----------------------------------------------------------------------===//
 
 void GenericCallOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                           StringRef callee, ArrayRef<mlir::Value> arguments) {
   // Generic call always returns an unranked Tensor initially.
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));

mlir/examples/toy/Ch6/mlir/Dialect.cpp

 mlir::Region *FuncOp::getCallableRegion() { return &getBody(); }
 
 /// Returns the results types that the callable region produces when
 /// executed.
 llvm::ArrayRef<mlir::Type> FuncOp::getCallableResults() {
   return getFunctionType().getResults();
 }
 
+/// Returns the argument attributes for all callable region arguments or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+/// Returns the result attributes for all callable region results or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // GenericCallOp
 //===----------------------------------------------------------------------===//
 
 void GenericCallOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                           StringRef callee, ArrayRef<mlir::Value> arguments) {
   // Generic call always returns an unranked Tensor initially.
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));

mlir/examples/toy/Ch7/mlir/Dialect.cpp

 mlir::Region *FuncOp::getCallableRegion() { return &getBody(); }
 
 /// Returns the results types that the callable region produces when
 /// executed.
 llvm::ArrayRef<mlir::Type> FuncOp::getCallableResults() {
   return getFunctionType().getResults();
 }
 
+/// Returns the argument attributes for all callable region arguments or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+/// Returns the result attributes for all callable region results or
+/// null if there are none.
+ArrayAttr FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // GenericCallOp
 //===----------------------------------------------------------------------===//
 
 void GenericCallOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
                           StringRef callee, ArrayRef<mlir::Value> arguments) {
   // Generic call always returns an unranked Tensor initially.
   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

let extraClassDeclaration = [{
     ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr
                                                               : &getBody(); }
 
     /// Returns the results types that the callable region produces when
     /// executed.
     ArrayRef<Type> getCallableResults() { return getFunctionType()
                                                     .getResults(); }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
     //===------------------------------------------------------------------===//
     // FunctionOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// Returns the argument types of this async function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
 
     /// Returns the result types of this async function.

mlir/include/mlir/Dialect/Func/IR/FuncOps.td

let extraClassDeclaration = [{
     /// return null in the case of an external callable object, e.g. an external
     /// function.
     ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }
 
     /// Returns the results types that the callable region produces when
     /// executed.
     ArrayRef<Type> getCallableResults() { return getFunctionType().getResults(); }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
     //===------------------------------------------------------------------===//
     // FunctionOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
 
     /// Returns the result types of this function.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let extraClassDeclaration = [{
     bool isVarArg() { return getFunctionType().isVarArg(); }
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getParams(); }
 
     /// Returns the result types of this function.
     ArrayRef<Type> getResultTypes() { return getFunctionType().getReturnTypes(); }
 
+    //===------------------------------------------------------------------===//
+    // CallableOpInterface
+    //===------------------------------------------------------------------===//
+
     /// Returns the callable region, which is the function body. If the function
     /// is external, returns null.
     Region *getCallableRegion();
 
     /// Returns the callable result type, which is the single function return
     /// type if it is not void, or an empty array if the function's return type
     /// is void, as void is not assignable to a value.
     ArrayRef<Type> getCallableResults() {
       if (getFunctionType().getReturnType().isa<LLVM::LLVMVoidType>())
         return {};
       return getFunctionType().getReturnTypes();
     }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 
 def LLVM_NullOp

mlir/include/mlir/Dialect/MLProgram/IR/MLProgramOps.td

let extraClassDeclaration = [{
     ::mlir::Region *getCallableRegion() {
       return isExternal() ? nullptr : &getBody();
     }
 
     /// Returns the results types that the callable region produces when
     /// executed.
     ArrayRef<Type> getCallableResults() { return getFunctionType().getResults(); }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
     //===------------------------------------------------------------------===//
     // FunctionOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
 
     /// Returns the result types of this function.
let extraClassDeclaration = [{
     /// return null in the case of an external callable object, e.g. an external
     /// function.
     ::mlir::Region *getCallableRegion() { return isExternal() ? nullptr : &getBody(); }
 
     /// Returns the results types that the callable region produces when
     /// executed.
     ArrayRef<Type> getCallableResults() { return getFunctionType().getResults(); }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
     //===------------------------------------------------------------------===//
     // FunctionOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
 
     /// Returns the result types of this function.

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

let extraClassDeclaration = [{
     }
 
     /// Returns the results types that the callable region produces when
     /// executed.
     ArrayRef<Type> getCallableResults() {
       return getFunctionType().getResults();
     }
 
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
     //===------------------------------------------------------------------===//
     // FunctionOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
 
     /// Returns the result types of this function.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let extraClassDeclaration = [{
     }
 
     ::mlir::Region *getCallableRegion() {
       return &getBody();
     }
     ::llvm::ArrayRef<::mlir::Type> getCallableResults() {
       return getFunctionType().getResults();
     }
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
   }];
 }
 
 def SplitHandlesOp : TransformDialectOp<"split_handles",
     [FunctionalStyleTransformOpTrait,
      DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "Splits handles from a union of payload ops to a list";

mlir/include/mlir/Interfaces/CallInterfaces.td

InterfaceMethod<[{
       "::mlir::Region *", "getCallableRegion"
     >,
     InterfaceMethod<[{
         Returns the results types that the callable region produces when
         executed.
       }],
       "::llvm::ArrayRef<::mlir::Type>", "getCallableResults"
     >,
+    InterfaceMethod<[{
+        Returns the argument attributes for all callable region arguments or
+        null if there are none.
+      }],
+      "::mlir::ArrayAttr", "getCallableArgAttrs"
+    >,
+    InterfaceMethod<[{
+        Returns the result attributes for all callable region results or null
+        if there are none.
+      }],
+      "::mlir::ArrayAttr", "getCallableResAttrs"
+    >
   ];
 }
 
 #endif // MLIR_INTERFACES_CALLINTERFACES

mlir/include/mlir/Transforms/InliningUtils.h

 //===- InliningUtils.h - Inliner utilities ----------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines interfaces for various inlining utility methods.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TRANSFORMS_INLININGUTILS_H
 #define MLIR_TRANSFORMS_INLININGUTILS_H
 
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Region.h"
 #include <optional>
 
 namespace mlir {
 
 class Block;
public:
   ///
   /// NOTE: This hook may be invoked before the 'isLegal' checks above.
   virtual Operation *materializeCallConversion(OpBuilder &builder, Value input,
                                                Type resultType,
                                                Location conversionLoc) const {
     return nullptr;
   }
 
+  /// Hook to transform the call arguments before using them to replace the
+  /// callee arguments. It returns the transformation result or `argument`
+  /// itself if the hook did not change anything. The type of the returned value
+  /// has to match `targetType`, and the `argumentAttrs` dictionary is non-null
+  /// even if no attribute is present. The hook is called after converting the
+  /// callsite argument types using the materializeCallConversion callback, and
+  /// right before inlining the callee region. Any operations created using the
+  /// provided `builder` are inserted right before the inlined callee region.
+  /// Example use cases are the insertion of copies for by value arguments, or
+  /// integer conversions that require signedness information.
+  virtual Value handleArgument(OpBuilder &builder, Operation *call,
+                               Operation *callable, Value argument,
+                               Type targetType,
+                               DictionaryAttr argumentAttrs) const {
+    return argument;
+  }
+
+  /// Hook to transform the callee results before using them to replace the call
+  /// results. It returns the transformation result or the `result` itself if
+  /// the hook did not change anything. The type of the returned values has to
+  /// match `targetType`, and the `resultAttrs` dictionary is non-null even if
+  /// no attribute is present. The hook is called right before handling
+  /// terminators, and obtains the callee result before converting its type
+  /// using the `materializeCallConversion` callback. Any operations created
+  /// using the provided `builder` are inserted right after the inlined callee
+  /// region. Example use cases are the insertion of copies for by value results
+  /// or integer conversions that require signedness information.
+  /// NOTE: This hook is invoked after inlining the `callable` region.
+  virtual Value handleResult(OpBuilder &builder, Operation *call,
+                             Operation *callable, Value result, Type targetType,
+                             DictionaryAttr resultAttrs) const {
+    return result;
+  }
+
   /// Process a set of blocks that have been inlined for a call. This callback
   /// is invoked before inlined terminator operations have been processed.
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {}
 };
 
 /// This interface provides the hooks into the inlining interface.
 /// Note: this class automatically collects 'DialectInlinerInterface' objects
public:
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   virtual void handleTerminator(Operation *op, Block *newDest) const;
   virtual void handleTerminator(Operation *op,
                                 ArrayRef<Value> valuesToRepl) const;
+
+  virtual Value handleArgument(OpBuilder &builder, Operation *call,
+                               Operation *callable, Value argument,
+                               Type targetType,
+                               DictionaryAttr argumentAttrs) const;
+  virtual Value handleResult(OpBuilder &builder, Operation *call,
+                             Operation *callable, Value result, Type targetType,
+                             DictionaryAttr resultAttrs) const;
+
   virtual void processInlinedCallBlocks(
       Operation *call, iterator_range<Region::iterator> inlinedBlocks) const;
 };
 
 //===----------------------------------------------------------------------===//
 // Inline Methods.
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

Region *spirv::FuncOp::getCallableRegion() {
   return isExternal() ? nullptr : &getBody();
 }
 
 // CallableOpInterface
 ArrayRef<Type> spirv::FuncOp::getCallableResults() {
   return getFunctionType().getResults();
 }
 
+// CallableOpInterface
+::mlir::ArrayAttr spirv::FuncOp::getCallableArgAttrs() {
+  return getArgAttrs().value_or(nullptr);
+}
+
+// CallableOpInterface
+::mlir::ArrayAttr spirv::FuncOp::getCallableResAttrs() {
+  return getResAttrs().value_or(nullptr);
+}
+
 //===----------------------------------------------------------------------===//
 // spirv.FunctionCall
 //===----------------------------------------------------------------------===//
 
 LogicalResult spirv::FunctionCallOp::verify() {
   auto fnName = getCalleeAttr();
 
   auto funcOp = dyn_cast_or_null<spirv::FuncOp>(

mlir/lib/Transforms/Utils/InliningUtils.cpp

 /// as necessary.
 void InlinerInterface::handleTerminator(Operation *op,
                                         ArrayRef<Value> valuesToRepl) const {
   auto *handler = getInterfaceFor(op);
   assert(handler && "expected valid dialect handler");
   handler->handleTerminator(op, valuesToRepl);
 }
 
+Value InlinerInterface::handleArgument(OpBuilder &builder, Operation *call,
+                                       Operation *callable, Value argument,
+                                       Type targetType,
+                                       DictionaryAttr argumentAttrs) const {
+  auto *handler = getInterfaceFor(callable);
+  assert(handler && "expected valid dialect handler");
+  return handler->handleArgument(builder, call, callable, argument, targetType,
+                                 argumentAttrs);
+}
+
+Value InlinerInterface::handleResult(OpBuilder &builder, Operation *call,
+                                     Operation *callable, Value result,
+                                     Type targetType,
+                                     DictionaryAttr resultAttrs) const {
+  auto *handler = getInterfaceFor(callable);
+  assert(handler && "expected valid dialect handler");
+  return handler->handleResult(builder, call, callable, result, targetType,
+                               resultAttrs);
+}
+
 void InlinerInterface::processInlinedCallBlocks(
     Operation *call, iterator_range<Region::iterator> inlinedBlocks) const {
   auto *handler = getInterfaceFor(call);
   assert(handler && "expected valid dialect handler");
   handler->processInlinedCallBlocks(call, inlinedBlocks);
 }
 
 /// Utility to check that all of the operations within 'src' can be inlined.
static bool isLegalToInline(InlinerInterface &interface, Region *src,
   }
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // Inline Methods
 //===----------------------------------------------------------------------===//
 
+static void handleArgumentImpl(InlinerInterface &interface, OpBuilder &builder,
+                               CallOpInterface call,
+                               CallableOpInterface callable,
+                               IRMapping &mapper) {
+  // Unpack the argument attributes if there are any.
+  SmallVector<DictionaryAttr> argAttrs(
+      callable.getCallableRegion()->getNumArguments(),
+      builder.getDictionaryAttr({}));
+  if (ArrayAttr arrayAttr = callable.getCallableArgAttrs()) {
+    assert(arrayAttr.size() == argAttrs.size());
+    for (auto [idx, attr] : llvm::enumerate(arrayAttr))
+      argAttrs[idx] = cast<DictionaryAttr>(attr);
+  }
+
+  // Run the argument attribute handler for the given argument and attribute.
+  for (auto [blockArg, argAttr] :
+       llvm::zip(callable.getCallableRegion()->getArguments(), argAttrs)) {
+    Value newArgument = interface.handleArgument(builder, call, callable,
+                                                 mapper.lookup(blockArg),
+                                                 blockArg.getType(), argAttr);
+    assert(newArgument.getType() == blockArg.getType() &&
+           "expected the handled argument type to match the target type");
+
+    // Update the mapping to point the new argument returned by the handler.
+    mapper.map(blockArg, newArgument);
+  }
+}
+
+static void handleResultImpl(InlinerInterface &interface, OpBuilder &builder,
+                             CallOpInterface call, CallableOpInterface callable,
+                             ValueRange results) {
+  // Unpack the result attributes if there are any.
+  SmallVector<DictionaryAttr> resAttrs(results.size(),
+                                       builder.getDictionaryAttr({}));
+  if (ArrayAttr arrayAttr = callable.getCallableResAttrs()) {
+    assert(arrayAttr.size() == resAttrs.size());
+    for (auto [idx, attr] : llvm::enumerate(arrayAttr))
+      resAttrs[idx] = cast<DictionaryAttr>(attr);
+  }
+
+  // Run the result attribute handler for the given result and attribute.
+  SmallVector<DictionaryAttr> resultAttributes;
+  for (auto [result, resAttr] : llvm::zip(results, resAttrs)) {
+    // Store the original result users before running the handler.
+    DenseSet<Operation *> resultUsers;
+    for (Operation *user : result.getUsers())
+      resultUsers.insert(user);
+
+    // TODO: Use the type of the call result to replace once the hook can be
+    // used for type conversions. At the moment, all type conversions have to be
+    // done using materializeCallConversion.
+    Type targetType = result.getType();
+
+    Value newResult = interface.handleResult(builder, call, callable, result,
+                                             targetType, resAttr);
+    assert(newResult.getType() == targetType &&
+           "expected the handled result type to match the target type");
+
+    // Replace the result uses except for the ones introduce by the handler.
+    result.replaceUsesWithIf(newResult, [&](OpOperand &operand) {
+      return resultUsers.count(operand.getOwner());
+    });
+  }
+}
+
 static LogicalResult
 inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
                  Block::iterator inlinePoint, IRMapping &mapper,
                  ValueRange resultsToReplace, TypeRange regionResultTypes,
                  std::optional<Location> inlineLoc,
                  bool shouldCloneInlinedRegion, Operation *call = nullptr) {
   assert(resultsToReplace.size() == regionResultTypes.size());
   // We expect the region to have at least one block.
inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
   // Check that the operations within the source region are valid to inline.
   Region *insertRegion = inlineBlock->getParent();
   if (!interface.isLegalToInline(insertRegion, src, shouldCloneInlinedRegion,
                                  mapper) ||
       !isLegalToInline(interface, src, insertRegion, shouldCloneInlinedRegion,
                        mapper))
     return failure();
 
+  // Run the argument attribute handler before inlining the callable region.
+  OpBuilder builder(inlineBlock, inlinePoint);
+  auto callable = dyn_cast<CallableOpInterface>(src->getParentOp());
+  if (call && callable)
+    handleArgumentImpl(interface, builder, call, callable, mapper);
+
   // Check to see if the region is being cloned, or moved inline. In either
   // case, move the new blocks after the 'insertBlock' to improve IR
   // readability.
   Block *postInsertBlock = inlineBlock->splitBlock(inlinePoint);
   if (shouldCloneInlinedRegion)
     src->cloneInto(insertRegion, postInsertBlock->getIterator(), mapper);
   else
     insertRegion->getBlocks().splice(postInsertBlock->getIterator(),
inlineRegionImpl(InlinerInterface &interface, Region *src, Block *inlineBlock,
 
   // Process the newly inlined blocks.
   if (call)
     interface.processInlinedCallBlocks(call, newBlocks);
   interface.processInlinedBlocks(newBlocks);
 
   // Handle the case where only a single block was inlined.
   if (std::next(newBlocks.begin()) == newBlocks.end()) {
+    // Run the result attribute handler on the terminator operands.
+    Operation *firstBlockTerminator = firstNewBlock->getTerminator();
+    builder.setInsertionPoint(firstBlockTerminator);
+    if (call && callable)
+      handleResultImpl(interface, builder, call, callable,
+                       firstBlockTerminator->getOperands());
+
     // Have the interface handle the terminator of this block.
-    auto *firstBlockTerminator = firstNewBlock->getTerminator();
     interface.handleTerminator(firstBlockTerminator,
                                llvm::to_vector<6>(resultsToReplace));
     firstBlockTerminator->erase();
 
     // Merge the post insert block into the cloned entry block.
     firstNewBlock->getOperations().splice(firstNewBlock->end(),
                                           postInsertBlock->getOperations());
     postInsertBlock->erase();
   } else {
     // Otherwise, there were multiple blocks inlined. Add arguments to the post
     // insertion block to represent the results to replace.
     for (const auto &resultToRepl : llvm::enumerate(resultsToReplace)) {
       resultToRepl.value().replaceAllUsesWith(
           postInsertBlock->addArgument(regionResultTypes[resultToRepl.index()],
                                        resultToRepl.value().getLoc()));
     }
 
+    // Run the result attribute handler on the post insertion block arguments.
+    builder.setInsertionPointToStart(postInsertBlock);
+    if (call && callable)
+      handleResultImpl(interface, builder, call, callable,
+                       postInsertBlock->getArguments());
+
     /// Handle the terminators for each of the new blocks.
     for (auto &newBlock : newBlocks)
       interface.handleTerminator(newBlock.getTerminator(), postInsertBlock);
   }
 
   // Splice the instructions of the inlined entry block into the insert block.
   inlineBlock->getOperations().splice(inlineBlock->end(),
                                       firstNewBlock->getOperations());

mlir/test/Transforms/inlining.mlir

func.func @func_with_block_args_location_callee1(%arg0 : i32) {
   return
 }
 
 // CHECK-LABEL: func @func_with_block_args_location_callee2
 func.func @func_with_block_args_location_callee2(%arg0 : i32) {
   call @func_with_block_args_location(%arg0) : (i32) -> ()
   return
 }
+
+// Check that we can handle argument and result attributes.
+test.conversion_func_op @handle_attr_callee_fn_multi_arg(%arg0 : i16, %arg1 : i16 {"test.handle_argument"}) -> (i16 {"test.handle_result"}, i16) {
+  %0 = arith.addi %arg0, %arg1 : i16
+  %1 = arith.subi %arg0, %arg1 : i16
+  "test.return"(%0, %1) : (i16, i16) -> ()
+}
+test.conversion_func_op @handle_attr_callee_fn(%arg0 : i32 {"test.handle_argument"}) -> (i32 {"test.handle_result"}) {
+  "test.return"(%arg0) : (i32) -> ()
+}
+
+// CHECK-LABEL: func @inline_handle_attr_call
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]
+func.func @inline_handle_attr_call(%arg0 : i16, %arg1 : i16) -> (i16, i16) {
+
+  // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[ARG1]]) : (i16) -> i16
+  // CHECK: %[[SUM:.*]] = arith.addi %[[ARG0]], %[[CHANGE_INPUT]]
+  // CHECK: %[[DIFF:.*]] = arith.subi %[[ARG0]], %[[CHANGE_INPUT]]
+  // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[SUM]]) : (i16) -> i16
+  // CHECK-NEXT: return %[[CHANGE_RESULT]], %[[DIFF]]
+  %res0, %res1 = "test.conversion_call_op"(%arg0, %arg1) { callee=@handle_attr_callee_fn_multi_arg } : (i16, i16) -> (i16, i16)
+  return %res0, %res1 : i16, i16
+}
+
+// CHECK-LABEL: func @inline_convert_and_handle_attr_call
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]
+func.func @inline_convert_and_handle_attr_call(%arg0 : i16) -> (i16) {
+
+  // CHECK: %[[CAST_INPUT:.*]] = "test.cast"(%[[ARG0]]) : (i16) -> i32
+  // CHECK: %[[CHANGE_INPUT:.*]] = "test.type_changer"(%[[CAST_INPUT]]) : (i32) -> i32
+  // CHECK: %[[CHANGE_RESULT:.*]] = "test.type_changer"(%[[CHANGE_INPUT]]) : (i32) -> i32
+  // CHECK: %[[CAST_RESULT:.*]] = "test.cast"(%[[CHANGE_RESULT]]) : (i32) -> i16
+  // CHECK: return %[[CAST_RESULT]]
+  %res = "test.conversion_call_op"(%arg0) { callee=@handle_attr_callee_fn } : (i16) -> (i16)
+  return %res : i16
+}

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

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Reducer/ReductionPatternInterface.h"
 #include "mlir/Transforms/FoldUtils.h"
Operation *materializeCallConversion(OpBuilder &builder, Value input,
     if (!(resultType.isSignlessInteger(16) ||
           resultType.isSignlessInteger(32)) ||
         !(input.getType().isSignlessInteger(16) ||
           input.getType().isSignlessInteger(32)))
       return nullptr;
     return builder.create<TestCastOp>(conversionLoc, resultType, input);
   }
 
+  Value handleArgument(OpBuilder &builder, Operation *call, Operation *callable,
+                       Value argument, Type targetType,
+                       DictionaryAttr argumentAttrs) const final {
+    if (!argumentAttrs.contains("test.handle_argument"))
+      return argument;
+    return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
+                                             argument);
+  }
+
+  Value handleResult(OpBuilder &builder, Operation *call, Operation *callable,
+                     Value result, Type targetType,
+                     DictionaryAttr resultAttrs) const final {
+    if (!resultAttrs.contains("test.handle_result"))
+      return result;
+    return builder.create<TestTypeChangerOp>(call->getLoc(), targetType,
+                                             result);
+  }
+
   void processInlinedCallBlocks(
       Operation *call,
       iterator_range<Region::iterator> inlinedBlocks) const final {
     if (!isa<ConversionCallOp>(call))
       return;
 
     // Set attributed on all ops in the inlined blocks.
     for (Block &block : inlinedBlocks) {
if (!fnAttr)
     return emitOpError("requires a 'callee' symbol reference attribute");
   if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))
     return emitOpError() << "'" << fnAttr.getValue()
                          << "' does not reference a valid function";
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ConversionFuncOp
+//===----------------------------------------------------------------------===//
+
+ParseResult ConversionFuncOp::parse(OpAsmParser &parser,
+                                    OperationState &result) {
+  auto buildFuncType =
+      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
+         function_interface_impl::VariadicFlag,
+         std::string &) { return builder.getFunctionType(argTypes, results); };
+
+  return function_interface_impl::parseFunctionOp(
+      parser, result, /*allowVariadic=*/false,
+      getFunctionTypeAttrName(result.name), buildFuncType,
+      getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
+}
+
+void ConversionFuncOp::print(OpAsmPrinter &p) {
+  function_interface_impl::printFunctionOp(
+      p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),
+      getArgAttrsAttrName(), getResAttrsAttrName());
+}
+
+//===----------------------------------------------------------------------===//
 // TestFoldToCallOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
   using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(FoldToCallOp op,

mlir/test/lib/Dialect/Test/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // 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 TEST_OPS
 #define TEST_OPS
 
 include "TestDialect.td"
 include "TestInterfaces.td"
 include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/PatternBase.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
def ConversionCallOp : TEST_Op<"conversion_call_op",
   }];
   let extraClassDefinition = [{
     ::mlir::CallInterfaceCallable $cppClass::getCallableForCallee() {
       return (*this)->getAttrOfType<::mlir::SymbolRefAttr>("callee");
     }
   }];
 }
 
+def ConversionFuncOp : TEST_Op<"conversion_func_op", [CallableOpInterface,
+                                                      FunctionOpInterface]> {
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttrOf<FunctionType>:$function_type,
+                       OptionalAttr<DictArrayAttr>:$arg_attrs,
+                       OptionalAttr<DictArrayAttr>:$res_attrs,
+                       OptionalAttr<StrAttr>:$sym_visibility);
+  let regions = (region AnyRegion:$body);
+
+  let extraClassDeclaration = [{
+    //===------------------------------------------------------------------===//
+    // CallableOpInterface
+    //===------------------------------------------------------------------===//
+
+    /// Returns the region on the current operation that is callable. This may
+    /// return null in the case of an external callable object, e.g. an external
+    /// function.
+    ::mlir::Region *getCallableRegion() {
+      return isExternal() ? nullptr : &getBody();
+    }
+
+    /// Returns the results types that the callable region produces when
+    /// executed.
+    ::mlir::ArrayRef<::mlir::Type> getCallableResults() {
+      return getFunctionType().getResults();
+    }
+
+    /// Returns the argument attributes for all callable region arguments or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return getArgAttrs().value_or(nullptr);
+    }
+
+    /// Returns the result attributes for all callable region results or
+    /// null if there are none.
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return getResAttrs().value_or(nullptr);
+    }
+
+    //===------------------------------------------------------------------===//
+    // FunctionOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    /// Returns the argument types of this async function.
+    ::mlir::ArrayRef<::mlir::Type> getArgumentTypes() {
+      return getFunctionType().getInputs();
+    }
+
+    /// Returns the result types of this async function.
+    ::mlir::ArrayRef<::mlir::Type> getResultTypes() {
+      return getFunctionType().getResults();
+    }
+
+    /// Returns the number of results of this async function
+    unsigned getNumResults() {return getResultTypes().size();}
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+}
+
 def FunctionalRegionOp : TEST_Op<"functional_region_op",
     [CallableOpInterface]> {
   let regions = (region AnyRegion:$body);
   let results = (outs FunctionType);
 
   let extraClassDeclaration = [{
     ::mlir::Region *getCallableRegion() { return &getBody(); }
     ::llvm::ArrayRef<::mlir::Type> getCallableResults() {
       return getType().cast<::mlir::FunctionType>().getResults();
     }
+    ::mlir::ArrayAttr getCallableArgAttrs() {
+      return nullptr;
+    }
+    ::mlir::ArrayAttr getCallableResAttrs() {
+      return nullptr;
+    }
   }];
 }
 
 
 def FoldToCallOp : TEST_Op<"fold_to_call_op"> {
   let arguments = (ins FlatSymbolRefAttr:$callee);
   let hasCanonicalizer = 1;
 }