[mlir][linalg] Return newly created ops from bufferize_to_allocation

[springerm]

Summary

Return all ops that were generated as part of the bufferization, so that user do not have to match them in the enclosing op.

Depends On: D154956

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

def BufferizeToAllocationOp : Op<Transform_Dialect,
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      ReportTrackingListenerFailuresOpTrait]> {
   let description = [{
     This transform bufferizes the targeted operation and materializes the
     result in a new allocation. It replaces all original uses of the target
     result with the newly allocated buffer, wrapped in a
     `bufferization.to_tensor` op. It returns a handle to the newly allocated
-    buffer. Furthermore, it returns a handle to the result of the `to_tensor`
-    op.
+    buffer. Furthermore, it returns a handle that is mapped to all newly created
+    ops.
 
     Only bufferizable ops are that bufferize to a memory write or have an
     aliasing OpOperand (and do not themselves bufferize to an allocation) are
     supported. They are bufferized using their BufferizableOpInterface
     implementation. E.g.:
 
     ```
     %0 = tensor.insert %f into %dest[%pos] : tensor<10xf32>
- `vector.mask` is bufferized together with its region. The allocation is
       placed in front of the `vector.mask` op.
 
     An optional memory space attribute can be specified for the materialized
     buffer allocation.
 
     #### Return modes
 
     This operation consumes the `target` handle and produces the
-    `allocated_buffer` handle. It always succeeds.
+    `allocated_buffer` and `new_ops` handles. It always succeeds.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$target,
                        OptionalAttr<AnyAttr>:$memory_space);
-  let results = (outs Transform_AnyValue:$allocated_buffer);
+  let results = (outs Transform_AnyValue:$allocated_buffer,
+                      Transform_AnyOpType:$new_ops);
   let assemblyFormat = "$target attr-dict `:` type($target)";
   
   let builders = [
     OpBuilder<(ins "Value":$target, "Attribute":$memorySpace)>,
     OpBuilder<(ins "Value":$target, "int64_t":$memorySpace)>
   ];
 }
 

mlir/include/mlir/IR/PatternMatch.h

struct Listener : public OpBuilder::Listener {
 
     static bool classof(const OpBuilder::Listener *base);
   };
 
   /// A listener that forwards all notifications to another listener. This
   /// struct can be used as a base to create listener chains, so that multiple
   /// listeners can be notified of IR changes.
   struct ForwardingListener : public RewriterBase::Listener {
-    ForwardingListener(Listener *listener) : listener(listener) {}
+    ForwardingListener(OpBuilder::Listener *listener) : listener(listener) {}
 
     void notifyOperationInserted(Operation *op) override {
       listener->notifyOperationInserted(op);
     }
     void notifyBlockCreated(Block *block) override {
       listener->notifyBlockCreated(block);
     }
     void notifyOperationModified(Operation *op) override {
-      listener->notifyOperationModified(op);
+      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+        rewriteListener->notifyOperationModified(op);
     }
     void notifyOperationReplaced(Operation *op, Operation *newOp) override {
-      listener->notifyOperationReplaced(op, newOp);
+      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+        rewriteListener->notifyOperationReplaced(op, newOp);
     }
     void notifyOperationReplaced(Operation *op,
                                  ValueRange replacement) override {
-      listener->notifyOperationReplaced(op, replacement);
+      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+        rewriteListener->notifyOperationReplaced(op, replacement);
     }
     void notifyOperationRemoved(Operation *op) override {
-      listener->notifyOperationRemoved(op);
+      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+        rewriteListener->notifyOperationRemoved(op);
     }
     LogicalResult notifyMatchFailure(
         Location loc,
         function_ref<void(Diagnostic &)> reasonCallback) override {
-      return listener->notifyMatchFailure(loc, reasonCallback);
+      if (auto *rewriteListener = dyn_cast<RewriterBase::Listener>(listener))
+        return rewriteListener->notifyMatchFailure(loc, reasonCallback);
+      return failure();
     }
 
   private:
-    Listener *listener;
+    OpBuilder::Listener *listener;
   };
 
   /// Move the blocks that belong to "region" before the given position in
   /// another region "parent". The two regions must be different. The caller
   /// is responsible for creating or updating the operation transferring flow
   /// of control to the region and passing it the correct block arguments.
   virtual void inlineRegionBefore(Region &region, Region &parent,
                                   Region::iterator before);

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

 //===----------------------------------------------------------------------===//
 // BufferizeToAllocationOp
 //===----------------------------------------------------------------------===//
 
 void transform::BufferizeToAllocationOp::build(OpBuilder &b,
                                                OperationState &result,
                                                Value target,
                                                Attribute memorySpace) {
+  SmallVector<Type> resultTypes;
+  resultTypes.push_back(b.getType<transform::AnyValueType>());
+  resultTypes.push_back(b.getType<transform::AnyOpType>());
   return build(b, result,
-               /*resultTypes=*/b.getType<transform::AnyValueType>(),
+               /*resultTypes=*/resultTypes,
                /*target=*/target,
                /*memorySpace=*/memorySpace);
 }
 
 void transform::BufferizeToAllocationOp::build(OpBuilder &b,
                                                OperationState &result,
                                                Value target,
                                                int64_t memorySpace) {
+  SmallVector<Type> resultTypes;
+  resultTypes.push_back(b.getType<transform::AnyValueType>());
+  resultTypes.push_back(b.getType<transform::AnyOpType>());
   return build(b, result,
-               /*resultTypes=*/b.getType<transform::AnyValueType>(),
+               /*resultTypes=*/resultTypes,
                /*target=*/target,
                /*memorySpace=*/b.getI64IntegerAttr(memorySpace));
 }
 
+namespace {
+class NewOpsListener : public RewriterBase::ForwardingListener {
+public:
+  using RewriterBase::ForwardingListener::ForwardingListener;
+
+  SmallVector<Operation *> getNewOps() const {
+    return SmallVector<Operation *>(newOps.begin(), newOps.end());
+  }
+
+private:
+  void notifyOperationInserted(Operation *op) override {
+    ForwardingListener::notifyOperationInserted(op);
+    auto inserted = newOps.insert(op);
+    (void)inserted;
+    assert(inserted.second && "expected newly created op");
+  }
+
+  void notifyOperationRemoved(Operation *op) override {
+    ForwardingListener::notifyOperationRemoved(op);
+    op->walk([&](Operation *op) { newOps.erase(op); });
+  }
+
+  DenseSet<Operation *> newOps;
+};
+} // namespace
+
 DiagnosedSilenceableFailure transform::BufferizeToAllocationOp::apply(
     transform::TransformRewriter &rewriter,
     transform::TransformResults &results, transform::TransformState &state) {
+  // Attach listener to keep track of newly created ops.
+  OpBuilder::Listener *previousListener = rewriter.getListener();
+  auto resetListener =
+      llvm::make_scope_exit([&]() { rewriter.setListener(previousListener); });
+  NewOpsListener newOpsListener(previousListener);
+  rewriter.setListener(&newOpsListener);
+
+  // Bufferize ops.
   Attribute memorySpace =
       getMemorySpace().has_value() ? getMemorySpace().value() : Attribute();
   SmallVector<Value> allocatedBuffers;
   for (Operation *op : state.getPayloadOps(getTarget())) {
     Value buffer = linalg::bufferizeToAllocation(rewriter, op, memorySpace);
     if (!buffer) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
                                          << "failed to bufferize operation";
       diag.attachNote(op->getLoc()) << "target payload op";
       return diag;
     }
     allocatedBuffers.push_back(buffer);
   }
+
+  // Set results.
   results.setValues(cast<OpResult>(getAllocatedBuffer()), allocatedBuffers);
+  results.set(cast<OpResult>(getNewOps()), newOpsListener.getNewOps());
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::BufferizeToAllocationOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   producesHandle(getAllocatedBuffer(), effects);
+  producesHandle(getNewOps(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/test/Dialect/Linalg/pad-to-specific-memory-space.mlir

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
   } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
-  %buffer = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op
+  %buffer, %new_ops = transform.structured.bufferize_to_allocation %pad {memory_space = 3} : !transform.any_op
   %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 
 }
 
 // -----
 
 #map = affine_map<()[s0] -> (-s0 + 12, 7)>
 
^bb1(%arg1: !transform.any_op):
   %padded, %pad = transform.structured.pad %0 {
     padding_values=[0.0 : f32, 0.0 : f32, 0.0 : f32],
     padding_dimensions=[0, 1, 2],
     pack_paddings=[1, 1, 1]
   } : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
   transform.structured.masked_vectorize %pad vector_sizes [10, 12] : !transform.any_op
   %vector_write = transform.structured.match ops{["vector.transfer_write"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %mask_op = transform.get_parent_op %vector_write {op_name = "vector.mask"} : (!transform.any_op) -> !transform.any_op
-  %buffer = transform.structured.bufferize_to_allocation %mask_op {memory_space = 3} : !transform.any_op
+  %buffer, %new_ops = transform.structured.bufferize_to_allocation %mask_op {memory_space = 3} : !transform.any_op
   %2 = transform.bufferization.one_shot_bufferize %arg1 {bufferize_function_boundaries=true} : (!transform.any_op) -> !transform.any_op
 }

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

^bb0(%arg0: index, %arg1: index):
     tensor.yield %c : index
   } : tensor<?x10xindex> to tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.bufferize_to_allocation %0 : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 : !transform.any_op
+
+  // Ensure that one linalg.fill was generated.
+  %fill_op = transform.select "linalg.fill" in %new : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{1}}
+  test_print_number_of_associated_payload_ir_ops %fill_op : !transform.any_op
+
+  // Ensure that one memref.tensor_store was generated.
+  %tensor_store = transform.select "memref.tensor_store" in %new : (!transform.any_op) -> !transform.any_op
+  // expected-remark @below{{1}}
+  test_print_number_of_associated_payload_ir_ops %tensor_store : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_pad_constant(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[src:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc
 //       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]
 //       CHECK:   memref.copy %[[src]], %[[subview]]
 //       CHECK:   bufferization.to_tensor %[[alloc]] restrict writable
 func.func @tensor_pad_constant(%t: tensor<?x10xindex>, %l2: index, %h1: index,
                                %h2: index) -> tensor<?x?xindex> {
   %0 = tensor.pad %t low[5, %l2] high[%h1, %h2] {
   ^bb0(%arg0: index, %arg1: index):
     %c = arith.constant 50 : index
     tensor.yield %c : index
   } : tensor<?x10xindex> to tensor<?x?xindex>
   return %0 : tensor<?x?xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.pad"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.bufferize_to_allocation %0 : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 : !transform.any_op
   // Make sure that One-Shot Bufferize can bufferize the rest.
   %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_insert(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
 //       CHECK:   %[[m:.*]] = bufferization.to_memref %[[t]]
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
 //       CHECK:   memref.copy %[[m]], %[[alloc]]
 //       CHECK:   memref.store %{{.*}}, %[[alloc]]
 //       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @tensor_insert(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> {
   %r = tensor.insert %v into %t[%idx, %idx] : tensor<?x10xindex>
   return %r : tensor<?x10xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
   // Make sure that One-Shot Bufferize can bufferize the rest.
   %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @tensor_insert_into_empty(
 //       CHECK:   %[[alloc:.*]] = memref.alloc() : memref<10xindex, 4>
 //   CHECK-NOT:   memref.copy
 //       CHECK:   memref.store %{{.*}}, %[[alloc]]
 //       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @tensor_insert_into_empty(%idx: index, %v: index) -> tensor<10xindex> {
   %e = tensor.empty() : tensor<10xindex>
   %r = tensor.insert %v into %e[%idx] : tensor<10xindex>
   return %r : tensor<10xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
   // Make sure that One-Shot Bufferize can bufferize the rest.
   %4 = transform.bufferization.one_shot_bufferize %arg1 : (!transform.any_op) -> !transform.any_op
 }
 
 // -----
 
 func.func @tensor_extract(%t: tensor<?x10xindex>, %idx: index) -> index {
   // expected-note @below{{target payload op}}
   %r = tensor.extract %t[%idx, %idx] : tensor<?x10xindex>
   return %r : index
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.extract"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below{{failed to bufferize operation}}
-  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @vector_mask(
 //  CHECK-SAME:     %[[t:.*]]: tensor<?xf32>,
 //       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?xf32, 4>
 //       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
 //       CHECK:   vector.mask %{{.*}} { vector.transfer_write %{{.*}}, %[[alloc]]
 //       CHECK:   %[[r:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return %[[r]]
 func.func @vector_mask(%t: tensor<?xf32>, %val: vector<16xf32>, %idx: index, %m0: vector<16xi1>) -> tensor<?xf32> {
   %r = vector.mask %m0 { vector.transfer_write %val, %t[%idx] : vector<16xf32>, tensor<?xf32> } : vector<16xi1> -> tensor<?xf32>
   return %r : tensor<?xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["vector.mask"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-  %2 = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4} : !transform.any_op
 }