[Mlir-commits] [mlir] ef4f535 - [mlir][linalg] BufferizeToAllocationOp: Do not copy uninitialized buffers

[Matthias Springer]

Author: Matthias Springer
Date: 2023-07-04T16:40:10+02:00
New Revision: ef4f5357e3d3b1071f90a4d72f27fa0e83440c00

URL: https://github.com/llvm/llvm-project/commit/ef4f5357e3d3b1071f90a4d72f27fa0e83440c00
DIFF: https://github.com/llvm/llvm-project/commit/ef4f5357e3d3b1071f90a4d72f27fa0e83440c00.diff

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

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

Differential Revision: https://reviews.llvm.org/D154081

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
    mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
index 570f6d255f097d..ea236a94fa81dd 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp
@@ -456,9 +456,12 @@ Value linalg::bufferizeToAllocation(RewriterBase &rewriter, Operation *op,
     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));
     });

diff  --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index e824c73a1f0797..a4f14ea1bf9e10 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -267,6 +267,11 @@ struct EmptyOpInterface
 
   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

diff  --git a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
index 9f98d6728ed36c..e51c20334e9554 100644
--- a/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
+++ b/mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir
@@ -88,6 +88,29 @@ transform.sequence failures(propagate) {
 
 // -----
 
+// 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>