[llvm] [mlir][bufferization] Empty tensor elimination for materialize_in_destination (PR #65468)
Matthias Springer via llvm-commits
llvm-commits at lists.llvm.org
Thu Sep 14 01:09:15 PDT 2023
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/65468:
>From e3cabda8fb426fa9658c0c4c0b252ee6126df1f7 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Thu, 14 Sep 2023 10:06:56 +0200
Subject: [PATCH] [mlir][bufferization] Empty tensor elimination for
materialize_in_destination
This revision adds support for empty tensor elimination to "bufferization.materialize_in_destination" by implementing the `InferDestinationOpInterface`.
Furthermore, the One-Shot Bufferize conflict detection is improved for "bufferization.materialize_in_destination".
---
.../Dialect/Bufferization/IR/Bufferization.h | 2 ++
.../Bufferization/IR/BufferizationOps.td | 15 +++++++++-
.../Bufferization/IR/BufferizationOps.cpp | 27 ++++++++++++++++++
.../one-shot-bufferize-analysis.mlir | 28 +++++++++++++++++++
...ot-bufferize-empty-tensor-elimination.mlir | 14 ++++++++++
.../llvm-project-overlay/mlir/BUILD.bazel | 2 ++
6 files changed, 87 insertions(+), 1 deletion(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h b/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
index 450dfb37ddb2e18..16bb0d87fb6fe79 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
@@ -12,7 +12,9 @@
#include "mlir/Bytecode/BytecodeOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
+#include "mlir/Dialect/Bufferization/IR/SubsetInsertionOpInterface.h"
#include "mlir/Interfaces/CopyOpInterface.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index da60dbfbf884a3e..6fcba25a0f29752 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -12,6 +12,8 @@
include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.td"
include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td"
include "mlir/Dialect/Bufferization/IR/BufferizationBase.td"
+include "mlir/Dialect/Bufferization/IR/SubsetInsertionOpInterface.td"
+include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/Interfaces/CopyOpInterface.td"
@@ -215,7 +217,11 @@ def Bufferization_CloneOp : Bufferization_Op<"clone", [
def Bufferization_MaterializeInDestinationOp
: Bufferization_Op<"materialize_in_destination",
[BufferizableOpInterface, SameOperandsAndResultType,
- DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+ DestinationStyleOpInterface,
+ DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<SubsetInsertionOpInterface,
+ ["getSourceOperand", "getValuesNeededToBuildSubsetExtraction",
+ "buildSubsetExtraction", "isEquivalentSubset"]>]> {
let summary = "copy a tensor";
let description = [{
@@ -248,12 +254,19 @@ def Bufferization_MaterializeInDestinationOp
bool bufferizesToMemoryWrite(OpOperand &opOperand,
const AnalysisState &state);
+ bool bufferizesToElementwiseAccess(const AnalysisState &state,
+ ArrayRef<OpOperand *> opOperands);
+
AliasingValueList getAliasingValues(
OpOperand &opOperand, const AnalysisState &state);
RankedTensorType getType() {
return ::llvm::cast<RankedTensorType>(getResult().getType());
}
+
+ std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
+ return {1, 2}; // `dest` operand
+ }
}];
let assemblyFormat = "$source `in` $dest attr-dict `:` type($source)";
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index e5016c956804688..abfad44abaeffeb 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -580,6 +580,13 @@ MaterializeInDestinationOp::bufferize(RewriterBase &rewriter,
return success();
}
+bool MaterializeInDestinationOp::bufferizesToElementwiseAccess(
+ const AnalysisState &state, ArrayRef<OpOperand *> opOperands) {
+ // As elements are copied from the "source" buffer to the "dest" buffer,
+ // already copied elements are not read a second time.
+ return true;
+}
+
LogicalResult MaterializeInDestinationOp::reifyResultShapes(
OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
reifiedReturnShapes.resize(1, SmallVector<OpFoldResult>(getType().getRank()));
@@ -587,6 +594,26 @@ LogicalResult MaterializeInDestinationOp::reifyResultShapes(
return success();
}
+Value MaterializeInDestinationOp::buildSubsetExtraction(OpBuilder &builder,
+ Location loc) {
+ // The subset is the entire destination tensor.
+ return getDest();
+}
+
+bool MaterializeInDestinationOp::isEquivalentSubset(
+ Value candidate, function_ref<bool(Value, Value)> equivalenceFn) {
+ return equivalenceFn(getDest(), candidate);
+}
+
+SmallVector<Value>
+MaterializeInDestinationOp::getValuesNeededToBuildSubsetExtraction() {
+ return {getDest()};
+}
+
+OpOperand &MaterializeInDestinationOp::getSourceOperand() {
+ return getOperation()->getOpOperand(0) /*source*/;
+}
+
//===----------------------------------------------------------------------===//
// ToTensorOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
index 5a505c66892f113..a2fbb06d179ebda 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -158,3 +158,31 @@ func.func @bbarg_of_unknown_op_2(%f: f32) {
// CHECK: {__inplace_operands_attr__ = ["false"]} : (tensor<10xf32>) -> ()
return
}
+
+// -----
+
+// CHECK: func @materialize_in_destination_aliasing(
+func.func @materialize_in_destination_aliasing(%t: tensor<?xf32>, %p1: index, %p2: index, %sz: index) -> tensor<5xf32> {
+ %buffer = tensor.empty(%sz) : tensor<?xf32>
+ // CHECK: tensor.extract_slice
+ // CHECK-SAME: {__inplace_operands_attr__ = ["true", "none"]}
+ %src = tensor.extract_slice %t[%p1][5][1] : tensor<?xf32> to tensor<5xf32>
+ // CHECK: tensor.extract_slice
+ // CHECK-SAME: {__inplace_operands_attr__ = ["false", "none"]}
+ %dest = tensor.extract_slice %t[%p2][5][1] : tensor<?xf32> to tensor<5xf32>
+ // CHECK: bufferization.materialize_in_destination
+ // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
+ %r = bufferization.materialize_in_destination %src in %dest : tensor<5xf32>
+ return %r : tensor<5xf32>
+}
+
+// -----
+
+// CHECK: func @materialize_in_destination(
+func.func @materialize_in_destination(%t: tensor<?xf32>, %sz: index) -> tensor<?xf32> {
+ %buffer = tensor.empty(%sz) : tensor<?xf32>
+ // CHECK: bufferization.materialize_in_destination
+ // CHECK-SAME: {__inplace_operands_attr__ = ["true", "true"]}
+ %r = bufferization.materialize_in_destination %buffer in %buffer : tensor<?xf32>
+ return %r : tensor<?xf32>
+}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
index 3d15599915f0cfc..063ed07467f90d0 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
@@ -291,3 +291,17 @@ func.func @regression_multiple_insertion_points(%t1: tensor<?x?xf32>) -> tensor<
%2 = tensor.insert_slice %filled into %t1 [%0, %1] [2, 5] [1, 1] : tensor<2x5xf32> into tensor<?x?xf32>
return %2 : tensor<?x?xf32>
}
+
+// -----
+
+// CHECK-LABEL: func @materialize_in_destination(
+// CHECK-SAME: %[[m:.*]]: memref<5xf32, strided<[?], offset: ?>>,
+// CHECK: linalg.fill {{.*}} outs(%[[m]]
+// CHECK: return %[[m]]
+func.func @materialize_in_destination(%t: tensor<5xf32>, %f: f32) -> tensor<5xf32> {
+ %0 = tensor.empty() : tensor<5xf32>
+ %filled = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
+ %1 = bufferization.materialize_in_destination %filled in %t : tensor<5xf32>
+ //%1 = tensor.insert_slice %filled into %t[0][5][1] : tensor<5xf32> into tensor<5xf32>
+ return %1 : tensor<5xf32>
+}
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 2263414388b0811..2425f0ebcbac974 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -12060,6 +12060,8 @@ gentbl_cc_library(
deps = [
":BufferizableOpInterfaceTdFiles",
":BufferizationOpsTdFiles",
+ ":DestinationStyleOpInterfaceTdFiles",
+ ":SubsetInsertionOpInterfaceTdFiles",
],
)
More information about the llvm-commits
mailing list