Add hook for dialect specializing processing blocks post inlining calls

[jpienaar]

Summary

This allows for dialects to do different post-processing depending on operations with the inliner (my use case requires different attribute propagation rules depending on call op). This hook runs before the regular processInlinedBlocks method.

[rriddle]

Missing the test?

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

178

nit: Prefer early exit here.

[jpienaar]

Add forgotten test update & use early exit

mlir/include/mlir/Transforms/InliningUtils.h

public:
   ///   ... = foo.call @foo(%input : i32) -> i16
   ///
   /// 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;
   }
+
+  /// 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
 /// registered to each dialect within the given context.
 class InlinerInterface
     : public DialectInterfaceCollection<DialectInlinerInterface> {
 public:
public:
 
   //===--------------------------------------------------------------------===//
   // Transformation Hooks
   //===--------------------------------------------------------------------===//
 
   virtual void handleTerminator(Operation *op, Block *newDest) const;
   virtual void handleTerminator(Operation *op,
                                 ArrayRef<Value> valuesToRepl) const;
+  virtual void processInlinedCallBlocks(
+      Operation *call, iterator_range<Region::iterator> inlinedBlocks) const;
 };
 
 //===----------------------------------------------------------------------===//
 // Inline Methods.
 //===----------------------------------------------------------------------===//
 
 /// This function inlines a region, 'src', into another. This function returns
 /// failure if it is not possible to inline this function. If the function
LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
                            TypeRange regionResultTypes,
                            Optional<Location> inlineLoc = llvm::None,
                            bool shouldCloneInlinedRegion = true);
 
 /// This function is an overload of the above 'inlineRegion' that allows for
 /// providing the set of operands ('inlinedOperands') that should be used
 /// in-favor of the region arguments when inlining.
 LogicalResult inlineRegion(InlinerInterface &interface, Region *src,
-                           Operation *inlinePoint,
-                           ValueRange inlinedOperands,
+                           Operation *inlinePoint, ValueRange inlinedOperands,
                            ValueRange resultsToReplace,
                            Optional<Location> inlineLoc = llvm::None,
                            bool shouldCloneInlinedRegion = true);
 
 /// This function inlines a given region, 'src', of a callable operation,
 /// 'callable', into the location defined by the given call operation. This
 /// function returns failure if inlining is not possible, success otherwise. On
 /// failure, no changes are made to the module. 'shouldCloneInlinedRegion'

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);
 }
 
+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,
                             Region *insertRegion, bool shouldCloneInlinedRegion,
                             BlockAndValueMapping &valueMapping) {
   for (auto &block : *src) {
     for (auto &op : block) {
       // Check this operation.
       if (!interface.isLegalToInline(&op, insertRegion,
static bool isLegalToInline(InlinerInterface &interface, Region *src,
   }
   return true;
 }
 
 //===----------------------------------------------------------------------===//
 // Inline Methods
 //===----------------------------------------------------------------------===//
 
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Operation *inlinePoint,
-                                 BlockAndValueMapping &mapper,
-                                 ValueRange resultsToReplace,
-                                 TypeRange regionResultTypes,
-                                 Optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
+static LogicalResult
+inlineRegionImpl(InlinerInterface &interface, Region *src,
+                 Operation *inlinePoint, BlockAndValueMapping &mapper,
+                 ValueRange resultsToReplace, TypeRange regionResultTypes,
+                 Optional<Location> inlineLoc, bool shouldCloneInlinedRegion,
+                 Operation *call) {
   assert(resultsToReplace.size() == regionResultTypes.size());
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
 
   // Check that all of the region arguments have been mapped.
   auto *srcEntryBlock = &src->front();
   if (llvm::any_of(srcEntryBlock->getArguments(),
if (inlineLoc && !inlineLoc->isa<UnknownLoc>())
     remapInlinedLocations(newBlocks, *inlineLoc);
 
   // If the blocks were moved in-place, make sure to remap any necessary
   // operands.
   if (!shouldCloneInlinedRegion)
     remapInlinedOperands(newBlocks, mapper);
 
   // 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()) {
     // Have the interface handle the terminator of this block.
     auto *firstBlockTerminator = firstNewBlock->getTerminator();
     interface.handleTerminator(firstBlockTerminator,
                                llvm::to_vector<6>(resultsToReplace));
inlineRegionImpl(InlinerInterface &interface, Region *src,
 
   // Splice the instructions of the inlined entry block into the insert block.
   insertBlock->getOperations().splice(insertBlock->end(),
                                       firstNewBlock->getOperations());
   firstNewBlock->erase();
   return success();
 }
 
-/// This function is an overload of the above 'inlineRegion' that allows for
-/// providing the set of operands ('inlinedOperands') that should be used
-/// in-favor of the region arguments when inlining.
-LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
-                                 Operation *inlinePoint,
-                                 ValueRange inlinedOperands,
-                                 ValueRange resultsToReplace,
-                                 Optional<Location> inlineLoc,
-                                 bool shouldCloneInlinedRegion) {
+static LogicalResult
+inlineRegionImpl(InlinerInterface &interface, Region *src,
+                 Operation *inlinePoint, ValueRange inlinedOperands,
+                 ValueRange resultsToReplace, Optional<Location> inlineLoc,
+                 bool shouldCloneInlinedRegion, Operation *call) {
   // We expect the region to have at least one block.
   if (src->empty())
     return failure();
 
   auto *entryBlock = &src->front();
   if (inlinedOperands.size() != entryBlock->getNumArguments())
     return failure();
 
   // Map the provided call operands to the arguments of the region.
   BlockAndValueMapping mapper;
   for (unsigned i = 0, e = inlinedOperands.size(); i != e; ++i) {
     // Verify that the types of the provided values match the function argument
     // types.
     BlockArgument regionArg = entryBlock->getArgument(i);
     if (inlinedOperands[i].getType() != regionArg.getType())
       return failure();
     mapper.map(regionArg, inlinedOperands[i]);
   }
 
   // Call into the main region inliner function.
-  return inlineRegion(interface, src, inlinePoint, mapper, resultsToReplace,
+  return inlineRegionImpl(interface, src, inlinePoint, mapper, resultsToReplace,
                           resultsToReplace.getTypes(), inlineLoc,
-                      shouldCloneInlinedRegion);
+                          shouldCloneInlinedRegion, call);
+}
+
+LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
+                                 Operation *inlinePoint,
+                                 BlockAndValueMapping &mapper,
+                                 ValueRange resultsToReplace,
+                                 TypeRange regionResultTypes,
+                                 Optional<Location> inlineLoc,
+                                 bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(interface, src, inlinePoint, mapper, resultsToReplace,
+                          regionResultTypes, inlineLoc,
+                          shouldCloneInlinedRegion,
+                          /*call=*/nullptr);
+}
+
+LogicalResult mlir::inlineRegion(InlinerInterface &interface, Region *src,
+                                 Operation *inlinePoint,
+                                 ValueRange inlinedOperands,
+                                 ValueRange resultsToReplace,
+                                 Optional<Location> inlineLoc,
+                                 bool shouldCloneInlinedRegion) {
+  return inlineRegionImpl(interface, src, inlinePoint, inlinedOperands,
+                          resultsToReplace, inlineLoc, shouldCloneInlinedRegion,
+                          /*call=*/nullptr);
 }
 
 /// Utility function used to generate a cast operation from the given interface,
 /// or return nullptr if a cast could not be generated.
 static Value materializeConversion(const DialectInlinerInterface *interface,
                                    SmallVectorImpl<Operation *> &castOps,
                                    OpBuilder &castBuilder, Value arg, Type type,
                                    Location conversionLoc) {
for (unsigned i = 0, e = callResults.size(); i != e; ++i) {
     castResult.getDefiningOp()->replaceUsesOfWith(castResult, callResult);
   }
 
   // Check that it is legal to inline the callable into the call.
   if (!interface.isLegalToInline(call, callable, shouldCloneInlinedRegion))
     return cleanupState();
 
   // Attempt to inline the call.
-  if (failed(inlineRegion(interface, src, call, mapper, callResults,
+  if (failed(inlineRegionImpl(interface, src, call, mapper, callResults,
                               callableResultTypes, call.getLoc(),
-                          shouldCloneInlinedRegion)))
+                              shouldCloneInlinedRegion, call)))
     return cleanupState();
   return success();
 }

mlir/test/Transforms/inlining.mlir

func @convert_callee_fn_multiblock() -> i32 {
   br ^bb0
 ^bb0:
   %0 = constant 0 : i32
   return %0 : i32
 }
 
 // CHECK-LABEL: func @inline_convert_result_multiblock
 func @inline_convert_result_multiblock() -> i16 {
-// CHECK:   br ^bb1
+// CHECK:   br ^bb1 {inlined_conversion}
 // CHECK: ^bb1:
-// CHECK:   %[[C:.+]] = constant 0 : i32
+// CHECK:   %[[C:.+]] = constant {inlined_conversion} 0 : i32
 // CHECK:   br ^bb2(%[[C]] : i32)
 // CHECK: ^bb2(%[[BBARG:.+]]: i32):
 // CHECK:   %[[CAST_RESULT:.+]] = "test.cast"(%[[BBARG]]) : (i32) -> i16
 // CHECK:   return %[[CAST_RESULT]] : i16
 
   %res = "test.conversion_call_op"() { callee=@convert_callee_fn_multiblock } : () -> (i16)
   return %res : i16
 }

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

Operation *materializeCallConversion(OpBuilder &builder, Value input,
     // Only allow conversion for i16/i32 types.
     if (!(resultType.isSignlessInteger(16) ||
           resultType.isSignlessInteger(32)) ||
         !(input.getType().isSignlessInteger(16) ||
           input.getType().isSignlessInteger(32)))
       return nullptr;
     return builder.create<TestCastOp>(conversionLoc, resultType, input);
   }
+
+  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) {
+      block.walk([&](Operation *op) {
+        op->setAttr("inlined_conversion", UnitAttr::get(call->getContext()));
+      });
+    }
+  }
 };
 
 struct TestReductionPatternInterface : public DialectReductionPatternInterface {
 public:
   TestReductionPatternInterface(Dialect *dialect)
       : DialectReductionPatternInterface(dialect) {}
 
   virtual void