[mlir][bufferization] Add bufferization.copy_tensor op

[springerm]

Summary

This operation is a "copy" operation on tensors. It is guaranteed to bufferize to a memcpy. This is different from "tensor.insert_slice", which may fold away.

Note: There is a symmetry between certain tensor, bufferization and memref ops:

• tensor.empty, bufferization.alloc_tensor, memref.alloc
• (none), bufferization.dealloc_tensor, memref.dealloc
• tensor.insert_slice, bufferization.copy_tensor, memref.copy

Tensor ops can generally canonicalize/fold away, while bufferization dialect ops can be used when a certain side effect is expected to materialize; so they do not fold away.

Depends On: D153514

[springerm]

rebase

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

%0 = linalg.matmul
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
     return %0 : tensor<?x?xf32, #SparseMatrix>
     ```
 
     ```mlir
     %c = bufferization.alloc_tensor(%d1, %d2) size_hint = %noe
       : tensor<?x?xf32, #SparseMatrix>
     ```
+
+    Note: An `alloc_tensor` with a `copy` should also be expressed as an
+    `alloc_tensor` without `copy`, followed by a `copy_tensor`.
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
                        Optional<AnyTensor>:$copy,
                        Optional<Index>:$size_hint,
                        OptionalAttr<AnyAttr>:$memory_space);
 
   let results = (outs AnyTensor:$result);
def Bufferization_CloneOp : Bufferization_Op<"clone", [
 
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// CopyTensorOp
+//===----------------------------------------------------------------------===//
+
+def Bufferization_CopyTensorOp : Bufferization_Op<"copy_tensor",
+    [BufferizableOpInterface, SameOperandsAndResultType,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+  let summary = "copy a tensor";
+
+  let description = [{
+    Copy the contents of the source tensor into the destination tensor. This
+    operation is guaranteed to bufferize to a memory copy.
+  }];
+
+  let arguments = (ins AnyTensor:$source,
+                       AnyTensor:$dest);
+
+  let results = (outs AnyTensor:$result);
+
+  let extraClassDeclaration = [{
+    LogicalResult bufferize(RewriterBase &rewriter,
+                            const BufferizationOptions &options);
+
+    bool bufferizesToMemoryRead(OpOperand &opOperand,
+                                const AnalysisState &state);
+
+    bool bufferizesToMemoryWrite(OpOperand &opOperand,
+                                 const AnalysisState &state);
+
+    AliasingOpResultList getAliasingOpResults(
+        OpOperand &opOperand, const AnalysisState &state);
+
+    RankedTensorType getType() {
+      return ::llvm::cast<RankedTensorType>(getResult().getType());
+    }
+  }];
+
+  let assemblyFormat = "$source `,` $dest attr-dict `:` type($source)";
+}
+
+//===----------------------------------------------------------------------===//
 // DeallocTensorOp
 //===----------------------------------------------------------------------===//
 
 def Bufferization_DeallocTensorOp : Bufferization_Op<"dealloc_tensor",
     [BufferizableOpInterface]> {
   string summary = "release underlying storage format of given tensor";
   string description = [{
     `bufferization.dealloc_tensor` is a buffer deallocation in tensor land. This

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

 Value AllocTensorOp::getDynamicSize(OpBuilder &b, unsigned idx) {
   assert(isDynamicDim(idx) && "expected dynamic dim");
   if (getCopy())
     return b.create<tensor::DimOp>(getLoc(), getCopy(), idx);
   return getOperand(getIndexOfDynamicSize(idx));
 }
 
 //===----------------------------------------------------------------------===//
+// CopyTensorOp
+//===----------------------------------------------------------------------===//
+
+bool CopyTensorOp::bufferizesToMemoryRead(OpOperand &opOperand,
+                                          const AnalysisState &state) {
+  if (&opOperand == &getOperation()->getOpOperand(0) /*source*/)
+    return true;
+  return false;
+}
+
+bool CopyTensorOp::bufferizesToMemoryWrite(OpOperand &opOperand,
+                                           const AnalysisState &state) {
+  if (&opOperand == &getOperation()->getOpOperand(1) /*dest*/)
+    return true;
+  return false;
+}
+
+AliasingOpResultList
+CopyTensorOp::getAliasingOpResults(OpOperand &opOperand,
+                                   const AnalysisState &state) {
+  if (&opOperand == &getOperation()->getOpOperand(1) /*dest*/)
+    return {{getOperation()->getResult(0), BufferRelation::Equivalent}};
+  return {};
+}
+
+LogicalResult CopyTensorOp::bufferize(RewriterBase &rewriter,
+                                      const BufferizationOptions &options) {
+  FailureOr<Value> buffer = getBuffer(rewriter, getDest(), options);
+  if (failed(buffer))
+    return failure();
+  rewriter.create<memref::TensorStoreOp>(getLoc(), getSource(), *buffer);
+  replaceOpWithBufferizedValues(rewriter, getOperation(), *buffer);
+  return success();
+}
+
+LogicalResult CopyTensorOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
+  reifiedReturnShapes[0] = tensor::getMixedSizes(builder, getLoc(), getDest());
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
 void CloneOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), getInput(),
                        SideEffects::DefaultResource::get());

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

 // -----
 
 // CHECK-LABEL: func @from_unranked_to_unranked
 func.func @from_unranked_to_unranked(%arg0: tensor<*xi32>) -> tensor<*xi32> {
   // CHECK: return %arg{{.*}} : tensor<*xi32>
   %0 = tensor.cast %arg0 : tensor<*xi32> to tensor<*xi32>
   return %0 : tensor<*xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @tensor_copy(
+//  CHECK-SAME:     %[[arg0:.*]]: tensor<5xf32>)
+func.func @tensor_copy(%arg0: tensor<5xf32>) -> tensor<5xf32> {
+  // CHECK: %[[m:.*]] = bufferization.to_memref %[[arg0]]
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+  // CHECK: memref.copy %[[m]], %[[alloc]]
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK: memref.dealloc %[[alloc]]
+  // CHECK: return %[[r]]
+  %dest = bufferization.alloc_tensor() : tensor<5xf32>
+  %0 = bufferization.copy_tensor %arg0, %dest : tensor<5xf32>
+  return %0 : tensor<5xf32>
+}

mlir/test/Dialect/Bufferization/invalid.mlir

 func.func private @invalid_writable_attribute(tensor<*xf32> {bufferization.writable = false})
 
 // -----
 
 func.func @invalid_writable_on_op() {
   // expected-error @+1{{attribute '"bufferization.writable"' not supported as an op attribute by the bufferization dialect}}
   arith.constant {bufferization.writable = true} 0  : index
 }
+
+// -----
+
+// expected-note @below{{prior use here}}
+func.func @invalid_tensor_copy(%arg0: tensor<?xf32>, %arg1: tensor<5xf32>) {
+  // expected-error @below{{expects different type than prior uses: 'tensor<?xf32>' vs 'tensor<5xf32>'}}
+  bufferization.copy_tensor %arg0, %arg1 : tensor<?xf32>
+}

mlir/test/Dialect/Bufferization/ops.mlir

 }
 
 // CHECK-LABEL: func @test_dealloc_tensor_op
 func.func @test_dealloc_tensor_op(%arg0: tensor<4xi32>) {
   // CHECK: bufferization.dealloc_tensor {{.*}} : tensor<4xi32>
   bufferization.dealloc_tensor %arg0 : tensor<4xi32>
   return
 }
+
+// CHECK-LABEL: func @test_copy_tensor_op
+func.func @test_copy_tensor_op(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>)
+    -> tensor<?xf32> {
+  // CHECK: bufferization.copy_tensor {{.*}} : tensor<?xf32>
+  %1 = bufferization.copy_tensor %arg0, %arg1 : tensor<?xf32>
+  return %1 : tensor<?xf32>
+}