[mlir][linalg] BufferizeToAllocationOp: Do not copy uninitialized buffers

[springerm]

Summary

Tensors/buffers that do not have any defined contents (e.g., tensor.empty) are no longer copied.

Depends On: D154071

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
   // Allocate buffers.
   OpBuilder::InsertionGuard g(rewriter);
   rewriter.setInsertionPoint(insertionPoint ? insertionPoint : op);
   SmallVector<Value> allocs;
   for (OpOperand *operand : outOfPlaceOperands) {
     Value alloc = createAllocationForTensor(rewriter, op->getLoc(),
                                             operand->get(), memorySpace);
     allocs.push_back(alloc);
-    // Initialize buffer with a copy of the operand data.
-    // TODO: Do not copy uninitialized tensors such as tensor.empty.
-    rewriter.create<memref::TensorStoreOp>(op->getLoc(), operand->get(), alloc);
+    if (!state.findDefinitions(operand->get()).empty()) {
+      // Initialize buffer with a copy of the operand data. Not needed if the
+      // tensor is uninitialized.
+      rewriter.create<memref::TensorStoreOp>(op->getLoc(), operand->get(),
+                                             alloc);
+    }
     rewriter.updateRootInPlace(op, [&]() {
       operand->set(rewriter.create<ToTensorOp>(op->getLoc(), alloc));
     });
   }
 
   // Bufferize the op.
   rewriter.setInsertionPoint(op);
   if (failed(bufferizableOp.bufferize(rewriter, options)))

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

struct EmptyOpInterface
   bool resultBufferizesToMemoryWrite(Operation *op, OpResult opResult,
                                      const AnalysisState &state) const {
     // The returned tensor does not have specified contents.
     return false;
   }
 
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
+    if (op->getUses().empty()) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+
     // tensor.empty ops are used to indicate the shape of a tensor. They have
     // no defined contents and cannot be bufferized. However, they can be
     // converted to bufferization.alloc_tensor ops, which then bufferize to an
     // allocation (--empty-tensor-to-alloc-tensor).
     return op->emitOpError("cannot be bufferized, but can be converted to "
                            "bufferization.alloc_tensor");
   }
 };

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

^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
   // 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
+  // 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):