[Mlir-commits] [mlir] 1ccd8cd - [mlir][bufferization] Add bufferization.eliminate_empty_tensors transform op
Matthias Springer
llvmlistbot at llvm.org
Wed Apr 5 22:22:56 PDT 2023
Author: Matthias Springer
Date: 2023-04-06T14:22:47+09:00
New Revision: 1ccd8cd6e68c1edfdfc0dbc21f4afb75f3a338e0
URL: https://github.com/llvm/llvm-project/commit/1ccd8cd6e68c1edfdfc0dbc21f4afb75f3a338e0
DIFF: https://github.com/llvm/llvm-project/commit/1ccd8cd6e68c1edfdfc0dbc21f4afb75f3a338e0.diff
LOG: [mlir][bufferization] Add bufferization.eliminate_empty_tensors transform op
Differential Revision: https://reviews.llvm.org/D144401
Added:
Modified:
mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
index 880781b668693..692f0f0b68d66 100644
--- a/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td
@@ -17,6 +17,13 @@ include "mlir/Dialect/PDL/IR/PDLTypes.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/OpBase.td"
+def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
+def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
+
+//===----------------------------------------------------------------------===//
+// OneShotBufferizeOp
+//===----------------------------------------------------------------------===//
+
def OneShotBufferizeOp
: Op<Transform_Dialect, "bufferization.one_shot_bufferize",
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
@@ -63,12 +70,65 @@ def OneShotBufferizeOp
}
//===----------------------------------------------------------------------===//
-// EmptyTensorToAllocTensorOp
+// EliminateEmptyTensorsOp
//===----------------------------------------------------------------------===//
+def EliminateEmptyTensorsOp
+ : Op<Transform_Dialect, "bufferization.eliminate_empty_tensors",
+ [DeclareOpInterfaceMethods<TransformOpInterface>,
+ DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
+ let description = [{
+ Try to eliminate all `tensor.empty` ops within the targeted op by replacing
+ them with a destination tensor.
-def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
-def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
+ `tensor.empty` ops cannot be bufferizes. They can either be converted to
+ `bufferization.alloc_tensor` or replaced with another tensor (via this
+ transform). `tensor.empty` does not specify the contents of the returned
+ tensor so their results can be replaced with arbitrary tensor values as long
+ as the dimensions match.
+
+ This transform looks for `tensor.empty` ops where the SSA use-def chain of
+ the result ends in a supported "anchor op" (always following the aliasing
+ OpOperand/OpResult chain). Currently supported anchor ops are:
+ - `tensor.insert_slice`
+ - `bufferization.yield` (inside `bufferization.alloc_tensor`)
+
+ Example:
+
+ ```
+ %0 = tensor.empty() : tensor<5xf32>
+ %1 = linalg.fill ... outs(%0)
+ %2 = tensor.insert_slice %1 into %t[1][5][1]
+ ```
+
+ Is rewritten with:
+ ```
+ %0 = tensor.extract_slice %t[1][5][1]
+ %1 = linalg.fill ... outs(%0)
+ %2 = tensor.insert_slice %1 into %t[1][5][1]
+ ```
+
+ The above example can bufferize without an allocation (in the absence of
+ other conflicts) because there is no longer a `tensor.empty` op.
+
+ See `-eliminate-empty-tensors` for more details.
+
+ #### Return modes
+
+ This transform reads the target handle and modifies the payload. It does
+ not produce any handle.
+ }];
+
+ let arguments = (ins PDL_Operation:$target);
+
+ let results = (outs);
+
+ let assemblyFormat = "$target attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// EmptyTensorToAllocTensorOp
+//===----------------------------------------------------------------------===//
def EmptyTensorToAllocTensorOp
: Op<Transform_Dialect, "bufferization.empty_tensor_to_alloc_tensor",
diff --git a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
index 8f56357b8402f..58766e83f255b 100644
--- a/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/PDL/IR/PDL.h"
#include "mlir/Dialect/PDL/IR/PDLTypes.h"
@@ -62,6 +63,37 @@ transform::OneShotBufferizeOp::apply(TransformResults &transformResults,
return DiagnosedSilenceableFailure::success();
}
+//===----------------------------------------------------------------------===//
+// EliminateEmptyTensorsOp
+//===----------------------------------------------------------------------===//
+
+void transform::EliminateEmptyTensorsOp::getEffects(
+ SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+ onlyReadsHandle(getTarget(), effects);
+ modifiesPayload(effects);
+}
+
+DiagnosedSilenceableFailure
+transform::EliminateEmptyTensorsOp::apply(TransformResults &transformResults,
+ TransformState &state) {
+ IRRewriter rewriter(getContext());
+ OneShotBufferizationOptions options;
+ options.allowReturnAllocs = true;
+
+ ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+ for (Operation *target : payloadOps) {
+ OneShotAnalysisState state(target, options);
+ if (failed(analyzeOp(target, state)))
+ return mlir::emitSilenceableFailure(target->getLoc())
+ << "failed to analyze op";
+ if (failed(bufferization::insertSliceAnchoredEmptyTensorEliminationStep(
+ rewriter, target, state)))
+ return mlir::emitSilenceableFailure(target->getLoc())
+ << "failed to eliminate insert_slice anchored tensor.empty ops";
+ }
+ return DiagnosedSilenceableFailure::success();
+}
+
//===----------------------------------------------------------------------===//
// EmptyTensorToAllocTensorOp
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
index 7e6250098063c..05d6a1bbdd557 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir
@@ -130,3 +130,24 @@ func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
%0 = tensor.empty() : tensor<2x2xf32>
return %0 : tensor<2x2xf32>
}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+ %0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+ transform.bufferization.eliminate_empty_tensors %0
+}
+
+// CHECK-LABEL: func @empty_tensor_elimination(
+// CHECK: tensor.extract_slice
+// CHECK: linalg.fill
+// CHECK: tensor.insert_slice
+func.func @empty_tensor_elimination(
+ %t: tensor<10xf32>, %f: f32) -> tensor<10xf32> {
+ %0 = tensor.empty() : tensor<5xf32>
+ %1 = linalg.fill ins(%f : f32) outs(%0 : tensor<5xf32>) -> tensor<5xf32>
+ %2 = tensor.insert_slice %1 into %t [1][5][1]
+ : tensor<5xf32> into tensor<10xf32>
+ return %2 : tensor<10xf32>
+}
More information about the Mlir-commits
mailing list