[mlir][bufferize] Add bufferization::WaitForBufferizationOp

[springerm]

Summary

This op canonicalizes away when the tensor operand has been bufferized.

This op is needed for partial bufferization of ops such as linalg.tiled_loop. Such ops can yield tensor values but not memref values. After bufferizing the loop (including the terminator) but not the loop body, ops in the loop body can DCE away because they no longer have any uses.

Example before bufferization (simplified):

linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
                  ins(%t0 : tensor<?xf32>)
		  outs(%t1 : tensor<?xf32>) {
  ...
  %0 = tensor.insert %f into %t0[...] : tensor<?xf32>
  linalg.yield %0
}

After bufferization of linalg.tiled_loop:

linalg.tiled_loop (%i) = (%c0) to (%c24) step (%c4)
                  ins(%m0 : memref<?xf32>)
		  outs(%m1 : memref<?xf32>) {
  ...
  %t0 = bufferization.to_tensor %m0
  %0 = tensor.insert %f into %t0[...] : tensor<?xf32>
  linalg.yield
}

Now the tensor.insert op can DCE away because it has no uses. This can be avoided by inserting the new op.

  ...
  %0 = tensor.insert %f into %t0[...] : tensor<?xf32>
  bufferization.wait_for_bufferization %0 : tensor<?xf32>
  linalg.yield
}

Note: WaitForBufferizationOp is also needed for a subsequent commit that switches the custom IR traversal of Comprehensive Bufferize to RewritePatterns (not dialect conversion but regular rewrite patterns). In that case, tensor.insert_slice ops that have a matching tensor.extract_slice op could DCE away.

Note: WaitForBufferizationOp is purpusely a new op in the bufferization dialect (as opposed to an anonymous/unnamed op) because it can survive partial bufferization. When used in Comprehensive Bufferize (One-Shot) bufferize, all WaitForBufferizationOps should have disappeared by the time bufferization is done (unless allow-unknown-ops). Other cases are considered a bufferization failure.

Depends On D116446

[nicolasvasilache]

putting a blocker until we have a deeper discussion, this does not pass my fishiness checks atm.

[springerm]

Putting this revision on hold. As we discussed, we likely won't need this and there's a better way to solve this issue.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

def Bufferization_ToMemrefOp : Bufferization_Op<"to_memref",
   let verifier = ?;
 
   let assemblyFormat = "$tensor attr-dict `:` type($memref)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// WaitForBufferizationOp
+//===----------------------------------------------------------------------===//
+
+def Bufferization_WaitForBufferizationOp
+    : Bufferization_Op<"wait_for_bufferization", []> {
+  let summary = "wait for the operand to be bufferized";
+  let description = [{
+    This op canonicalizes away once the tensor operand has been bufferized. The
+    tensor operand is bufferized if it is the result of a `to_tensor` op.
+
+    This op is useful for partial bufferization of certain ops. E.g., such ops
+    may have a region that yields a tensor value, but their corresponding memref
+    variant may not be yielding anything.
+
+    This op is used internally during bufferization. It should not be created
+    outside of `BufferizableOpInterface::bufferize` implementations.
+  }];
+
+  let arguments = (ins AnyTensor:$tensor);
+  let results = (outs);
+  // This op is fully verified by traits.
+  let verifier = ?;
+  let assemblyFormat = "$tensor attr-dict `:` type($tensor)";
+  let hasCanonicalizer = 1;
+}
+
 #endif // BUFFERIZATION_OPS

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

return builder.create<memref::DeallocOp>(alloc.getLoc(), alloc)
       .getOperation();
 }
 
 Optional<Value> CloneOp::buildClone(OpBuilder &builder, Value alloc) {
   return builder.create<CloneOp>(alloc.getLoc(), alloc).getResult();
 }
 
 //===----------------------------------------------------------------------===//
+// WaitForBufferizationOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Replace tensor.cast + to_memref by to_memref + memref.cast.
+struct FoldWaitForBufferizationOp
+    : public OpRewritePattern<WaitForBufferizationOp> {
+  using OpRewritePattern<WaitForBufferizationOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(WaitForBufferizationOp waitOp,
+                                PatternRewriter &rewriter) const final {
+    if (waitOp.tensor().getDefiningOp<ToTensorOp>()) {
+      rewriter.eraseOp(waitOp);
+      return success();
+    }
+
+    return failure();
+  }
+};
+} // namespace
+
+void WaitForBufferizationOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<FoldWaitForBufferizationOp>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/IR/BufferizationOps.cpp.inc"

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-partial.mlir

// CHECK-SCF: } else {
     // CHECK-SCF: scf.yield %[[insert_memref]]
     scf.yield %1, %pos : tensor<?xf32>, index
   }
 
   // CHECK-SCF: %[[r_tensor:.*]] = bufferization.to_tensor %[[r]]
   // CHECK-SCF: return %[[r_tensor]], %[[pos]]
   return %r1, %r2 : tensor<?xf32>, index
 }
+
+// -----
+
+// CHECK-LABEL: func @wait_for_bufferization_folds_away(
+//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+func @wait_for_bufferization_folds_away(
+    %t1: tensor<?xf32> {linalg.inplaceable = true}) -> f32 {
+  // CHECK-NOT: wait_for_bufferization
+  // CHECK-NOT: to_tensor
+  // CHECK-NOT: to_memref
+  bufferization.wait_for_bufferization %t1 : tensor<?xf32>
+  %c0 = arith.constant 0 : index
+  %1 = tensor.extract %t1[%c0] : tensor<?xf32>
+  return %1 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @wait_for_bufferization_does_not_fold_away(
+//  CHECK-SAME:     %[[m1:.*]]: memref<?xf32
+func @wait_for_bufferization_does_not_fold_away(
+    %t1: tensor<?xf32> {linalg.inplaceable = true}) -> f32 {
+  // CHECK: %[[m1_tensor:.*]] = bufferization.to_tensor %[[m1]]
+  // CHECK: %[[dummy:.*]] = "test.dummy_op"(%[[m1_tensor]])
+  %0 = "test.dummy_op"(%t1) : (tensor<?xf32>) -> tensor<?xf32>
+  // CHECK: bufferization.wait_for_bufferization %[[dummy]]
+  bufferization.wait_for_bufferization %0 : tensor<?xf32>
+  %c0 = arith.constant 0 : index
+  %1 = tensor.extract %0[%c0] : tensor<?xf32>
+  return %1 : f32
+}