From f486f826390cb230ee75da0d99cb7bd3c6384218 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 3 Oct 2023 10:50:47 +0200 Subject: [PATCH 1/2] [mlir][tensor][bufferize] `tensor.empty` bufferizes to an allocation Make `tensor.empty` bufferizable, so that the `-empty-tensor-to-alloc-tensor` pass becomes optional. This makes the bufferization easier to use. `tensor.empty` used to be non-bufferizable, so that there two separate ops, one that can be optimized away (`tensor.empty`) and one that is guaranteed to bufferize to an allocation (`bufferization.alloc_tensor`). With the recent improvements of "empty tensor elimination" this is no longer needed and `bufferization.alloc_tensor` can be phased out. --- .../BufferizableOpInterfaceImpl.cpp | 21 ++++++++++--------- mlir/test/Dialect/Tensor/bufferize.mlir | 9 +++++--- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp index b08283f007078..5f34fd45679b7 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp @@ -253,10 +253,7 @@ struct DimOpInterface } }; -/// Bufferization of tensor.empty. This op does not bufferize, but we need an -/// interface implementation, so that the result of this op is considered -/// "writable" (default impl. of `isWritable`). Results of ops that do not -/// implement `BufferizableOpInterface` are not writable. +/// Bufferization of tensor.empty. Replace with "bufferization.alloc_tensor". struct EmptyOpInterface : public BufferizableOpInterface::ExternalModel { @@ -268,17 +265,21 @@ struct EmptyOpInterface LogicalResult bufferize(Operation *op, RewriterBase &rewriter, const BufferizationOptions &options) const { + auto emptyOp = cast(op); + + // Optimization: Fold away the op if it has no uses. 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"); + // Allocate a tensor. This emits a "bufferization.alloc_tensor" op. + FailureOr allocTensor = allocateTensorForShapedValue( + rewriter, op->getLoc(), emptyOp.getResult(), options, /*copy=*/false); + if (failed(allocTensor)) + return failure(); + rewriter.replaceOp(op, *allocTensor); + return success(); } }; diff --git a/mlir/test/Dialect/Tensor/bufferize.mlir b/mlir/test/Dialect/Tensor/bufferize.mlir index c7b16315bfed1..a8b3c6af9ae89 100644 --- a/mlir/test/Dialect/Tensor/bufferize.mlir +++ b/mlir/test/Dialect/Tensor/bufferize.mlir @@ -1,4 +1,4 @@ -// RUN: mlir-opt %s -tensor-bufferize -cse -split-input-file -verify-diagnostics | FileCheck %s +// RUN: mlir-opt %s -tensor-bufferize -cse -split-input-file | FileCheck %s // CHECK-LABEL: func @dim( // CHECK-SAME: %[[TENSOR:.*]]: tensor<*xf32>, @@ -62,9 +62,12 @@ func.func @tensor.cast_to_unranked(%arg0: tensor<2xf32>) -> tensor<*xf32> { } // ----- + +// CHECK-LABEL: func @tensor.empty( +// CHECK: %[[ALLOC:.*]] = memref.alloc() {{.*}} : memref<5xf32> +// CHECK: %[[RET:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<5xf32> +// CHECK: return %[[RET]] : tensor<5xf32> func.func @tensor.empty() -> tensor<5xf32> { - // expected-error@+2 {{failed to bufferize op}} - // expected-error@+1 {{cannot be bufferized, but can be converted to bufferization.alloc_tensor}} %0 = tensor.empty() : tensor<5xf32> return %0 : tensor<5xf32> } From 1b0808f384549e8a8a8c6dac1e797677c4039fb3 Mon Sep 17 00:00:00 2001 From: Matthias Springer Date: Tue, 3 Oct 2023 15:59:44 +0200 Subject: [PATCH 2/2] Update mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp Co-authored-by: Martin Erhart --- .../Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp index 5f34fd45679b7..e082e2c54ef36 100644 --- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp @@ -253,7 +253,7 @@ struct DimOpInterface } }; -/// Bufferization of tensor.empty. Replace with "bufferization.alloc_tensor". +/// Bufferization of "tensor.empty". Replace with "bufferization.alloc_tensor". struct EmptyOpInterface : public BufferizableOpInterface::ExternalModel {