[mlir][Inliner] Add a new hook for checking if it is legal to inline a callable into a call

[rriddle]

Summary

In certain situations it isn't legal to inline a call operation, but this isn't something that is possible(at least not easily) to prevent with the current hooks. This revision adds a new hook so that dialects with call operations that shouldn't be inlined can prevent it.

[jpienaar]

Nice, thanks

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

 
 ```c++
 /// This class defines the interface for handling inlining with Toy operations.
 /// We simplify inherit from the base interface class and override
 /// the necessary methods.
 struct ToyInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
+  /// This hook checks to see if the given callable operation is legal to inline
+  /// into the given call. For Toy this hook can simply return true, as the Toy
+  /// Call operation is always inlinable.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// This hook checks to see if the given operation is legal to inline into the
   /// given region. For Toy this hook can simply return true, as all Toy
   /// operations are inlinable.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 

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

 /// operations.
 struct ToyInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// All call operations within toy can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// All operations within toy can be inlined.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

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

 /// operations.
 struct ToyInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// All call operations within toy can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// All operations within toy can be inlined.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

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

 /// operations.
 struct ToyInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// All call operations within toy can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// All operations within toy can be inlined.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

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

 /// operations.
 struct ToyInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// All call operations within toy can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// All operations within toy can be inlined.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

mlir/include/mlir/IR/Operation.h

public:
 
   /// Return the specified attribute if present, null otherwise.
   Attribute getAttr(Identifier name) { return attrs.get(name); }
   Attribute getAttr(StringRef name) { return attrs.get(name); }
 
   template <typename AttrClass> AttrClass getAttrOfType(Identifier name) {
     return getAttr(name).dyn_cast_or_null<AttrClass>();
   }
-
   template <typename AttrClass> AttrClass getAttrOfType(StringRef name) {
     return getAttr(name).dyn_cast_or_null<AttrClass>();
   }
 
+  /// Return true if the operation has an attribute with the provided name,
+  /// false otherwise.
+  bool hasAttr(Identifier name) { return static_cast<bool>(getAttr(name)); }
+  bool hasAttr(StringRef name) { return static_cast<bool>(getAttr(name)); }
+  template <typename AttrClass, typename NameT>
+  bool hasAttrOfType(NameT &&name) {
+    return static_cast<bool>(
+        getAttrOfType<AttrClass>(std::forward<NameT>(name)));
+  }
+
   /// If the an attribute exists with the specified name, change it to the new
   /// value.  Otherwise, add a new attribute with the specified name/value.
   void setAttr(Identifier name, Attribute value) { attrs.set(name, value); }
   void setAttr(StringRef name, Attribute value) {
     setAttr(Identifier::get(name, getContext()), value);
   }
 
   /// Remove the attribute with the specified name if it exists.  The return

mlir/include/mlir/Transforms/InliningUtils.h

class DialectInlinerInterface
     : public DialectInterface::Base<DialectInlinerInterface> {
 public:
   DialectInlinerInterface(Dialect *dialect) : Base(dialect) {}
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// Returns true if the given operation 'callable', that implements the
+  /// 'CallableOpInterface', can be inlined into the position given call
+  /// operation 'call', that is registered to the current dialect and implements
+  /// the `CallOpInterface`.
+  virtual bool isLegalToInline(Operation *call, Operation *callable) const {
+    return false;
+  }
+
   /// Returns true if the given region 'src' can be inlined into the region
   /// 'dest' that is attached to an operation registered to the current dialect.
   /// 'valueMapping' contains any remapped values from within the 'src' region.
   /// This can be used to examine what values will replace entry arguments into
   /// the 'src' region for example.
   virtual bool isLegalToInline(Region *dest, Region *src,
                                BlockAndValueMapping &valueMapping) const {
     return false;
public:
   /// These hooks mirror the hooks for the DialectInlinerInterface, with default
   /// implementations that call the hook on the handler for the dialect 'op' is
   /// registered to.
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  virtual bool isLegalToInline(Operation *call, Operation *callable) const;
   virtual bool isLegalToInline(Region *dest, Region *src,
                                BlockAndValueMapping &valueMapping) const;
   virtual bool isLegalToInline(Operation *op, Region *dest,
                                BlockAndValueMapping &valueMapping) const;
   virtual bool shouldAnalyzeRecursively(Operation *op) const;
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

static inline bool containsReturn(Region &region) {
   });
 }
 
 namespace {
 /// This class defines the interface for inlining within the SPIR-V dialect.
 struct SPIRVInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
+  /// All call operations within SPIRV can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// Returns true if the given region 'src' can be inlined into the region
   /// 'dest' that is attached to an operation registered to the current dialect.
   bool isLegalToInline(Region *dest, Region *src,
                        BlockAndValueMapping &) const final {
     // Return true here when inlining into spv.func, spv.selection, and
     // spv.loop operations.
     auto *op = dest->getParentOp();
     return isa<spirv::FuncOp, spirv::SelectionOp, spirv::LoopOp>(op);

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 /// operations.
 struct StdInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  /// All call operations within standard ops can be inlined.
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    return true;
+  }
+
   /// All operations within standard ops can be inlined.
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks

mlir/lib/Transforms/Utils/InliningUtils.cpp

static void remapInlinedOperands(iterator_range<Region::iterator> inlinedBlocks,
   for (auto &block : inlinedBlocks)
     block.walk(remapOperands);
 }
 
 //===----------------------------------------------------------------------===//
 // InlinerInterface
 //===----------------------------------------------------------------------===//
 
+bool InlinerInterface::isLegalToInline(Operation *call,
+                                       Operation *callable) const {
+  auto *handler = getInterfaceFor(call);
+  return handler ? handler->isLegalToInline(call, callable) : false;
+}
+
 bool InlinerInterface::isLegalToInline(
     Region *dest, Region *src, BlockAndValueMapping &valueMapping) const {
   // Regions can always be inlined into functions.
   if (isa<FuncOp>(dest->getParentOp()))
     return true;
 
   auto *handler = getInterfaceFor(dest->getParentOp());
   return handler ? handler->isLegalToInline(dest, src, valueMapping) : false;
Value castResult =
         materializeConversion(callInterface, castOps, castBuilder, callResult,
                               callResult.getType(), castLoc);
     if (!castResult)
       return cleanupState();
     callResult.replaceAllUsesWith(castResult);
     castResult.getDefiningOp()->replaceUsesOfWith(castResult, callResult);
   }
 
+  // Check that it is legal to inline the callable into the call.
+  if (!interface.isLegalToInline(call, callable))
+    return cleanupState();
+
   // Attempt to inline the call.
   if (failed(inlineRegion(interface, src, call, mapper, callResults,
                           callableResultTypes, call.getLoc(),
                           shouldCloneInlinedRegion)))
     return cleanupState();
   return success();
 }

mlir/test/Transforms/inlining.mlir

 // INLINE_SIMPLIFY-LABEL: func @inline_simplify
 func @inline_simplify() -> i32 {
   // INLINE_SIMPLIFY-NEXT: %[[CST:.*]] = constant 0 : i32
   // INLINE_SIMPLIFY-NEXT: return %[[CST]]
   %fn = call @simplify_return_reference() : () -> (() -> i32)
   %res = call_indirect %fn() : () -> i32
   return %res : i32
 }
+
+// CHECK-LABEL: func @no_inline_invalid_call
+func @no_inline_invalid_call() -> i32 {
+  %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock, noinline } : () -> (i32)
+  return %res : i32
+}

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

 /// operations.
 struct TestInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
+  bool isLegalToInline(Operation *call, Operation *callable) const final {
+    // Don't allow inlining calls that are marked `noinline`.
+    return !call->hasAttr("noinline");
+  }
   bool isLegalToInline(Region *, Region *, BlockAndValueMapping &) const final {
     // Inlining into test dialect regions is legal.
     return true;
   }
   bool isLegalToInline(Operation *, Region *,
                        BlockAndValueMapping &) const final {
     return true;
   }