diff --git a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp @@ -456,9 +456,12 @@ 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(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(op->getLoc(), operand->get(), + alloc); + } rewriter.updateRootInPlace(op, [&]() { operand->set(rewriter.create(op->getLoc(), alloc)); }); diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp --- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp @@ -267,6 +267,11 @@ 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 --- 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 @@ // ----- +// 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, %idx: index) -> index { // expected-note @below{{target payload op}} %r = tensor.extract %t[%idx, %idx] : tensor