[Mlir-commits] [mlir] 913286b - [mlir][linalg] Add `SubsetInsertionOpInterface` to `linalg.copy` (#67524)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 27 01:04:41 PDT 2023
Author: Matthias Springer
Date: 2023-09-27T10:04:37+02:00
New Revision: 913286baed9bf8dd5bc1763184228c2e5e3286ea
URL: https://github.com/llvm/llvm-project/commit/913286baed9bf8dd5bc1763184228c2e5e3286ea
DIFF: https://github.com/llvm/llvm-project/commit/913286baed9bf8dd5bc1763184228c2e5e3286ea.diff
LOG: [mlir][linalg] Add `SubsetInsertionOpInterface` to `linalg.copy` (#67524)
This commit enables empty tensor elimination on `linalg.copy` ops.
Added:
mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
Modified:
mlir/include/mlir/InitAllDialects.h
mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-empty-tensor-elimination.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
new file mode 100644
index 000000000000000..023a46df2620109
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
@@ -0,0 +1,21 @@
+//===- SubsetInsertionOpInterfaceImpl.h - Tensor subsets --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_SUBSETINSERTIONOPINTERFACEIMPL_H
+#define MLIR_DIALECT_LINALG_SUBSETINSERTIONOPINTERFACEIMPL_H
+
+namespace mlir {
+class DialectRegistry;
+
+namespace linalg {
+void registerSubsetInsertionOpInterfaceExternalModels(
+ DialectRegistry ®istry);
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_SUBSETINSERTIONOPINTERFACEIMPL_H
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index 6c9cdd87274fe0e..5ec36a7f289e586 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -45,6 +45,7 @@
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
+#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
#include "mlir/Dialect/Math/IR/Math.h"
@@ -148,6 +149,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
gpu::registerBufferDeallocationOpInterfaceExternalModels(registry);
linalg::registerBufferizableOpInterfaceExternalModels(registry);
+ linalg::registerSubsetInsertionOpInterfaceExternalModels(registry);
linalg::registerTilingInterfaceExternalModels(registry);
linalg::registerValueBoundsOpInterfaceExternalModels(registry);
memref::registerAllocationOpInterfaceExternalModels(registry);
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
index 287949324b31d72..77ad13dacaa9838 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
@@ -110,6 +110,7 @@ LogicalResult mlir::bufferization::eliminateEmptyTensors(
// be replaced, but the transformation may not be beneficial.
if (!state.isInPlace(source))
return WalkResult::skip();
+
// All values that are needed to create the replacement op.
SmallVector<Value> neededValues =
op.getValuesNeededToBuildSubsetExtraction();
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index 5ae9b7f7b1efc3f..4e094609afa6a03 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -27,6 +27,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
Split.cpp
SplitReduction.cpp
SubsetHoisting.cpp
+ SubsetInsertionOpInterfaceImpl.cpp
SwapExtractSliceWithFillPatterns.cpp
Tiling.cpp
TilingInterfaceImpl.cpp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
new file mode 100644
index 000000000000000..ef7c1b76e0cd04a
--- /dev/null
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -0,0 +1,57 @@
+//===- SubsetInsertionOpInterfaceImpl.cpp - Tensor subsets ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
+
+#include "mlir/Dialect/Bufferization/IR/SubsetInsertionOpInterface.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+
+using namespace mlir;
+using namespace mlir::bufferization;
+using namespace mlir::linalg;
+
+namespace {
+struct LinalgCopyOpInterface
+ : public SubsetInsertionOpInterface::ExternalModel<LinalgCopyOpInterface,
+ linalg::CopyOp> {
+ OpOperand &getSourceOperand(Operation *op) const {
+ auto copyOp = cast<CopyOp>(op);
+ assert(copyOp.getInputs().size() == 1 && "expected single input");
+ return copyOp.getInputsMutable()[0];
+ }
+
+ bool
+ isEquivalentSubset(Operation *op, Value candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ auto copyOp = cast<CopyOp>(op);
+ assert(copyOp.getOutputs().size() == 1 && "expected single output");
+ return equivalenceFn(candidate, copyOp.getOutputs()[0]);
+ }
+
+ Value buildSubsetExtraction(Operation *op, OpBuilder &builder,
+ Location loc) const {
+ auto copyOp = cast<CopyOp>(op);
+ assert(copyOp.getOutputs().size() == 1 && "expected single output");
+ return copyOp.getOutputs()[0];
+ }
+
+ SmallVector<Value>
+ getValuesNeededToBuildSubsetExtraction(Operation *op) const {
+ auto copyOp = cast<CopyOp>(op);
+ assert(copyOp.getOutputs().size() == 1 && "expected single output");
+ return {copyOp.getOutputs()[0]};
+ }
+};
+} // namespace
+
+void mlir::linalg::registerSubsetInsertionOpInterfaceExternalModels(
+ DialectRegistry ®istry) {
+ registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
+ });
+}
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 e8beb272ccd5b68..41e43047657daff 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
@@ -305,3 +305,15 @@ func.func @materialize_in_destination(%t: tensor<5xf32>, %f: f32) -> tensor<5xf3
return %1 : tensor<5xf32>
}
+// -----
+
+// CHECK-LABEL: func @linalg_copy(
+// CHECK-SAME: %[[m:.*]]: memref<5xf32, strided<[?], offset: ?>>,
+// CHECK: linalg.fill {{.*}} outs(%[[m]]
+// CHECK: return %[[m]]
+func.func @linalg_copy(%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 = linalg.copy ins(%filled : tensor<5xf32>) outs(%t : tensor<5xf32>) -> tensor<5xf32>
+ return %1 : tensor<5xf32>
+}
More information about the Mlir-commits
mailing list