[mlir][Inliner] Add a `wouldBeCloned` flag to each of the `isLegalToInline` hooks.

[rriddle]

Summary

Often times the legality of inlining can change depending on if the callable is going to be inlined in-place, or cloned. For example, some operations are not allowed to be duplicated and can only be inlined if the original callable will cease to exist afterwards. The new wouldBeCloned flag allows for dialects to hook into this when determining legality.

Depends On D90359

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

 /// 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) 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 *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   /// This hook is called when a terminator operation has been inlined. The only
   /// terminator that we have in the Toy dialect is the return
   /// operation(toy.return). We handle the return by replacing the values
   /// previously returned by the call operation with the operands of the

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

 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     return true;
   }
 
   /// All operations within toy can be inlined.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 

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

 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     return true;
   }
 
   /// All operations within toy can be inlined.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 

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

 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     return true;
   }
 
   /// All operations within toy can be inlined.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 

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

 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     return true;
   }
 
   /// All operations within toy can be inlined.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 

mlir/include/mlir/Transforms/InliningUtils.h

public:
 
   //===--------------------------------------------------------------------===//
   // 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 {
+  /// the `CallOpInterface`. 'wouldBeCloned' is set to true if the region of the
+  /// given 'callable' is set to be cloned during the inlining process, or false
+  /// if the region is set to be moved in-place(i.e. no duplicates would be
+  /// created).
+  virtual bool isLegalToInline(Operation *call, Operation *callable,
+                               bool wouldBeCloned) 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,
+  /// 'wouldBeCloned' is set to true if the given 'src' region is set to be
+  /// cloned during the inlining process, or false if the region is set to be
+  /// moved in-place(i.e. no duplicates would be created). '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, bool wouldBeCloned,
                                BlockAndValueMapping &valueMapping) const {
     return false;
   }
 
   /// Returns true if the given operation 'op', that is registered to this
   /// dialect, can be inlined into the given region, false otherwise.
-  /// 'valueMapping' contains any remapped values from within the 'src' region.
-  /// This can be used to examine what values may potentially replace the
-  /// operands to 'op'.
-  virtual bool isLegalToInline(Operation *op, Region *dest,
+  /// 'wouldBeCloned' is set to true if the given 'op' is set to be cloned
+  /// during the inlining process, or false if the operation is set to be moved
+  /// in-place(i.e. no duplicates would be created). 'valueMapping' contains any
+  /// remapped values from within the 'src' region. This can be used to examine
+  /// what values may potentially replace the operands to 'op'.
+  virtual bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                                BlockAndValueMapping &valueMapping) const {
     return false;
   }
 
   /// This hook is invoked on an operation that contains regions. It should
   /// return true if the analyzer should recurse within the regions of this
   /// operation when computing legality and cost, false otherwise. The default
   /// implementation returns true.
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,
+  virtual bool isLegalToInline(Operation *call, Operation *callable,
+                               bool wouldBeCloned) const;
+  virtual bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                                BlockAndValueMapping &valueMapping) const;
-  virtual bool isLegalToInline(Operation *op, Region *dest,
+  virtual bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                                BlockAndValueMapping &valueMapping) const;
   virtual bool shouldAnalyzeRecursively(Operation *op) const;
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   virtual void handleTerminator(Operation *op, Block *newDest) const;

mlir/lib/Dialect/Affine/IR/AffineOps.cpp

struct AffineInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
   /// 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,
+  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
     // Conservatively don't allow inlining into affine structures.
     return false;
   }
 
   /// Returns true if the given operation 'op', that is registered to this
   /// dialect, can be inlined into the given region, false otherwise.
-  bool isLegalToInline(Operation *op, Region *region,
+  bool isLegalToInline(Operation *op, Region *region, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
     // Always allow inlining affine operations into the top-level region of a
     // function. There are some edge cases when inlining *into* affine
     // structures, but that is handled in the other 'isLegalToInline' hook
     // above.
     // TODO: We should be able to inline into other regions than functions.
     return isa<FuncOp>(region->getParentOp());
   }

mlir/lib/Dialect/Linalg/IR/LinalgTypes.cpp

 
 namespace {
 
 struct LinalgInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   // We don't have any special restrictions on what can be inlined into
   // destination regions (e.g. while/conditional bodies). Always allow it.
-  bool isLegalToInline(Region *dest, Region *src,
+  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
     return true;
   }
   // Operations in Linalg dialect are always legal to inline.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
   // Handle the given inlined terminator by replacing it with a new operation
   // as necessary. Required when the region has only one block.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {}
 };

mlir/lib/Dialect/SCF/SCF.cpp

 // SCFDialect Dialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct SCFInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
   // We don't have any special restrictions on what can be inlined into
   // destination regions (e.g. while/conditional bodies). Always allow it.
-  bool isLegalToInline(Region *dest, Region *src,
+  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        BlockAndValueMapping &valueMapping) const final {
     return true;
   }
   // Operations in scf dialect are always legal to inline since they are
   // pure.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
   // Handle the given inlined terminator by replacing it with a new operation
   // as necessary. Required when the region has only one block.
   void handleTerminator(Operation *op,
                         ArrayRef<Value> valuesToRepl) const final {
     auto retValOp = dyn_cast<scf::YieldOp>(op);

mlir/lib/Dialect/SPIRV/SPIRVDialect.cpp

 }
 
 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) 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,
+  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        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);
   }
 
   /// Returns true if the given operation 'op', that is registered to this
   /// dialect, can be inlined into the region 'dest' that is attached to an
   /// operation registered to the current dialect.
-  bool isLegalToInline(Operation *op, Region *dest,
+  bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                        BlockAndValueMapping &) const final {
     // TODO: Enable inlining structured control flows with return.
     if ((isa<spirv::SelectionOp, spirv::LoopOp>(op)) &&
         containsReturn(op->getRegion(0)))
       return false;
     // TODO: we need to filter OpKill here to avoid inlining it to
     // a loop continue construct:
     // https://github.com/KhronosGroup/SPIRV-Headers/issues/86

mlir/lib/Dialect/Shape/IR/Shape.cpp

 
 namespace {
 /// This class defines the interface for inlining shape dialect ops.
 struct ShapeInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   // 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,
+  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   // Returns true if the given operation 'op', that is registered to this
   // dialect, can be inlined into the region 'dest' that is attached to an
   // operation registered to the current dialect.
-  bool isLegalToInline(Operation *op, Region *dest,
+  bool isLegalToInline(Operation *op, Region *dest, bool wouldBeCloned,
                        BlockAndValueMapping &) const final {
     return true;
   }
 };
 } // namespace
 
 void ShapeDialect::initialize() {
   addOperations<

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

 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 {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     return true;
   }
 
   /// All operations within standard ops can be inlined.
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        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(Operation *call, Operation *callable,
+                                       bool wouldBeCloned) const {
+  if (auto *handler = getInterfaceFor(call))
+    return handler->isLegalToInline(call, callable, wouldBeCloned);
+  return false;
 }
 
 bool InlinerInterface::isLegalToInline(
-    Region *dest, Region *src, BlockAndValueMapping &valueMapping) const {
+    Region *dest, Region *src, bool wouldBeCloned,
+    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;
+  if (auto *handler = getInterfaceFor(dest->getParentOp()))
+    return handler->isLegalToInline(dest, src, wouldBeCloned, valueMapping);
+  return false;
 }
 
 bool InlinerInterface::isLegalToInline(
-    Operation *op, Region *dest, BlockAndValueMapping &valueMapping) const {
-  auto *handler = getInterfaceFor(op);
-  return handler ? handler->isLegalToInline(op, dest, valueMapping) : false;
+    Operation *op, Region *dest, bool wouldBeCloned,
+    BlockAndValueMapping &valueMapping) const {
+  if (auto *handler = getInterfaceFor(op))
+    return handler->isLegalToInline(op, dest, wouldBeCloned, valueMapping);
+  return false;
 }
 
 bool InlinerInterface::shouldAnalyzeRecursively(Operation *op) const {
   auto *handler = getInterfaceFor(op);
   return handler ? handler->shouldAnalyzeRecursively(op) : true;
 }
 
 /// Handle the given inlined terminator by replacing it with a new operation
void InlinerInterface::handleTerminator(Operation *op,
                                         ArrayRef<Value> valuesToRepl) const {
   auto *handler = getInterfaceFor(op);
   assert(handler && "expected valid dialect handler");
   handler->handleTerminator(op, valuesToRepl);
 }
 
 /// Utility to check that all of the operations within 'src' can be inlined.
 static bool isLegalToInline(InlinerInterface &interface, Region *src,
-                            Region *insertRegion,
+                            Region *insertRegion, bool shouldCloneInlinedRegion,
                             BlockAndValueMapping &valueMapping) {
   for (auto &block : *src) {
     for (auto &op : block) {
       // Check this operation.
-      if (!interface.isLegalToInline(&op, insertRegion, valueMapping)) {
+      if (!interface.isLegalToInline(&op, insertRegion,
+                                     shouldCloneInlinedRegion, valueMapping)) {
         LLVM_DEBUG({
           llvm::dbgs() << "* Illegal to inline because of op: ";
           op.dump();
         });
         return false;
       }
       // Check any nested regions.
       if (interface.shouldAnalyzeRecursively(&op) &&
           llvm::any_of(op.getRegions(), [&](Region &region) {
             return !isLegalToInline(interface, &region, insertRegion,
-                                    valueMapping);
+                                    shouldCloneInlinedRegion, valueMapping);
           }))
         return false;
     }
   }
   return true;
 }
 
 //===----------------------------------------------------------------------===//
LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
 
   // The insertion point must be within a block.
   Block *insertBlock = inlinePoint->getBlock();
   if (!insertBlock)
     return failure();
   Region *insertRegion = insertBlock->getParent();
 
   // Check that the operations within the source region are valid to inline.
-  if (!interface.isLegalToInline(insertRegion, src, mapper) ||
-      !isLegalToInline(interface, src, insertRegion, mapper))
+  if (!interface.isLegalToInline(insertRegion, src, shouldCloneInlinedRegion,
+                                 mapper) ||
+      !isLegalToInline(interface, src, insertRegion, shouldCloneInlinedRegion,
+                       mapper))
     return failure();
 
   // Split the insertion block.
   Block *postInsertBlock =
       insertBlock->splitBlock(++inlinePoint->getIterator());
 
   // 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
Value castResult =
                               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))
+  if (!interface.isLegalToInline(call, callable, shouldCloneInlinedRegion))
     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/lib/Dialect/Test/TestDialect.cpp

 /// operations.
 struct TestInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;
 
   //===--------------------------------------------------------------------===//
   // Analysis Hooks
   //===--------------------------------------------------------------------===//
 
-  bool isLegalToInline(Operation *call, Operation *callable) const final {
+  bool isLegalToInline(Operation *call, Operation *callable,
+                       bool wouldBeCloned) const final {
     // Don't allow inlining calls that are marked `noinline`.
     return !call->hasAttr("noinline");
   }
-  bool isLegalToInline(Region *, Region *, BlockAndValueMapping &) const final {
+  bool isLegalToInline(Region *, Region *, bool,
+                       BlockAndValueMapping &) const final {
     // Inlining into test dialect regions is legal.
     return true;
   }
-  bool isLegalToInline(Operation *, Region *,
+  bool isLegalToInline(Operation *, Region *, bool,
                        BlockAndValueMapping &) const final {
     return true;
   }
 
   bool shouldAnalyzeRecursively(Operation *op) const final {
     // Analyze recursively if this is not a functional region operation, it
     // froms a separate functional scope.
     return !isa<FunctionalRegionOp>(op);