[mlir] [llvm] [mlir][Interfaces] Add `SubsetOpInterface` and `SubsetExtractionOpInterface` (PR #70617)
Matthias Springer via llvm-commits
llvm-commits at lists.llvm.org
Tue Oct 31 18:13:22 PDT 2023
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/70617
>From feec2b8cce9e8371ef58e067150d895475613fb3 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 1 Nov 2023 10:11:45 +0900
Subject: [PATCH] [mlir] `SubsetOpInterface` and `SubsetExtractionOpInterface`
---
.../Dialect/Bufferization/IR/Bufferization.h | 2 +-
.../Bufferization/IR/BufferizationOps.td | 3 +-
.../SubsetInsertionOpInterfaceImpl.h | 3 +-
.../SubsetInsertionOpInterfaceImpl.h | 3 +-
mlir/include/mlir/InitAllDialects.h | 4 +-
mlir/include/mlir/Interfaces/CMakeLists.txt | 2 +-
.../Interfaces/SubsetInsertionOpInterface.h | 27 --
.../Interfaces/SubsetInsertionOpInterface.td | 155 ----------
.../mlir/Interfaces/SubsetOpInterface.h | 45 +++
.../mlir/Interfaces/SubsetOpInterface.td | 267 ++++++++++++++++++
.../mlir/Interfaces/ValueBoundsOpInterface.h | 36 +--
.../Bufferization/IR/BufferizationOps.cpp | 12 +
.../Dialect/Bufferization/IR/CMakeLists.txt | 2 +-
.../Bufferization/Transforms/CMakeLists.txt | 2 +-
.../Transforms/EmptyTensorElimination.cpp | 2 +-
.../Transforms/OneShotAnalysis.cpp | 2 +-
.../Dialect/Linalg/Transforms/CMakeLists.txt | 2 +-
.../SubsetInsertionOpInterfaceImpl.cpp | 29 +-
.../BufferizableOpInterfaceImpl.cpp | 2 +-
.../Dialect/Tensor/Transforms/CMakeLists.txt | 2 +-
.../SubsetInsertionOpInterfaceImpl.cpp | 137 +++++++--
mlir/lib/Interfaces/CMakeLists.txt | 9 +-
.../Interfaces/SubsetInsertionOpInterface.cpp | 23 --
mlir/lib/Interfaces/SubsetOpInterface.cpp | 58 ++++
.../lib/Interfaces/ValueBoundsOpInterface.cpp | 76 +++++
.../llvm-project-overlay/mlir/BUILD.bazel | 33 +--
.../mlir/python/BUILD.bazel | 2 +-
27 files changed, 652 insertions(+), 288 deletions(-)
delete mode 100644 mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h
delete mode 100644 mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td
create mode 100644 mlir/include/mlir/Interfaces/SubsetOpInterface.h
create mode 100644 mlir/include/mlir/Interfaces/SubsetOpInterface.td
delete mode 100644 mlir/lib/Interfaces/SubsetInsertionOpInterface.cpp
create mode 100644 mlir/lib/Interfaces/SubsetOpInterface.cpp
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h b/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
index c035190f43e3950..e98b5728b38ef81 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/Bufferization.h
@@ -15,7 +15,7 @@
#include "mlir/Interfaces/CopyOpInterface.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
//===----------------------------------------------------------------------===//
// Bufferization Dialect
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 72a4aa712f49c98..e6b6d052df96a8c 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -15,7 +15,7 @@ include "mlir/Dialect/Bufferization/IR/BufferizationBase.td"
include "mlir/Interfaces/DestinationStyleOpInterface.td"
include "mlir/Interfaces/InferTypeOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
-include "mlir/Interfaces/SubsetInsertionOpInterface.td"
+include "mlir/Interfaces/SubsetOpInterface.td"
include "mlir/Interfaces/CopyOpInterface.td"
class Bufferization_Op<string mnemonic, list<Trait> traits = []>
@@ -220,6 +220,7 @@ def Bufferization_MaterializeInDestinationOp
AllElementTypesMatch<["source", "dest"]>,
BufferizableOpInterface, DestinationStyleOpInterface,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<SubsetOpInterface>,
DeclareOpInterfaceMethods<SubsetInsertionOpInterface,
["getSourceOperand", "getValuesNeededToBuildSubsetExtraction",
"buildSubsetExtraction", "isEquivalentSubset"]>,
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
index 023a46df2620109..94b0fb25b506650 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h
@@ -13,8 +13,7 @@ namespace mlir {
class DialectRegistry;
namespace linalg {
-void registerSubsetInsertionOpInterfaceExternalModels(
- DialectRegistry ®istry);
+void registerSubsetOpInterfaceExternalModels(DialectRegistry ®istry);
} // namespace linalg
} // namespace mlir
diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h
index e21b07d8a2705a0..019da189a8c991b 100644
--- a/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h
@@ -13,8 +13,7 @@ namespace mlir {
class DialectRegistry;
namespace tensor {
-void registerSubsetInsertionOpInterfaceExternalModels(
- DialectRegistry ®istry);
+void registerSubsetOpInterfaceExternalModels(DialectRegistry ®istry);
} // namespace tensor
} // namespace mlir
diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h
index 00f400aab5d50a0..7c2ffb7408d9afd 100644
--- a/mlir/include/mlir/InitAllDialects.h
+++ b/mlir/include/mlir/InitAllDialects.h
@@ -151,7 +151,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
cf::registerBufferDeallocationOpInterfaceExternalModels(registry);
gpu::registerBufferDeallocationOpInterfaceExternalModels(registry);
linalg::registerBufferizableOpInterfaceExternalModels(registry);
- linalg::registerSubsetInsertionOpInterfaceExternalModels(registry);
+ linalg::registerSubsetOpInterfaceExternalModels(registry);
linalg::registerTilingInterfaceExternalModels(registry);
linalg::registerValueBoundsOpInterfaceExternalModels(registry);
memref::registerAllocationOpInterfaceExternalModels(registry);
@@ -167,7 +167,7 @@ inline void registerAllDialects(DialectRegistry ®istry) {
tensor::registerBufferizableOpInterfaceExternalModels(registry);
tensor::registerFindPayloadReplacementOpInterfaceExternalModels(registry);
tensor::registerInferTypeOpInterfaceExternalModels(registry);
- tensor::registerSubsetInsertionOpInterfaceExternalModels(registry);
+ tensor::registerSubsetOpInterfaceExternalModels(registry);
tensor::registerTilingInterfaceExternalModels(registry);
tensor::registerValueBoundsOpInterfaceExternalModels(registry);
vector::registerBufferizableOpInterfaceExternalModels(registry);
diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt
index 36a04ff0eaeaf4b..d81298bb4daf014 100644
--- a/mlir/include/mlir/Interfaces/CMakeLists.txt
+++ b/mlir/include/mlir/Interfaces/CMakeLists.txt
@@ -12,7 +12,7 @@ add_mlir_interface(ParallelCombiningOpInterface)
add_mlir_interface(RuntimeVerifiableOpInterface)
add_mlir_interface(ShapedOpInterfaces)
add_mlir_interface(SideEffectInterfaces)
-add_mlir_interface(SubsetInsertionOpInterface)
+add_mlir_interface(SubsetOpInterface)
add_mlir_interface(TilingInterface)
add_mlir_interface(ValueBoundsOpInterface)
add_mlir_interface(VectorInterfaces)
diff --git a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h b/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h
deleted file mode 100644
index 3a6dfceadcce7c0..000000000000000
--- a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h
+++ /dev/null
@@ -1,27 +0,0 @@
-//===- SubsetInsertionOpInterface.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_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_
-#define MLIR_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_
-
-#include "mlir/IR/OpDefinition.h"
-
-namespace mlir {
-namespace detail {
-
-/// Return the destination/"init" operand of the op if it implements the
-/// `DestinationStyleOpInterface` and has exactly one "init" operand. Asserts
-/// otherwise.
-OpOperand &defaultGetDestinationOperand(Operation *op);
-
-} // namespace detail
-} // namespace mlir
-
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h.inc"
-
-#endif // MLIR_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_
diff --git a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td b/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td
deleted file mode 100644
index ef94a8ae9a60efd..000000000000000
--- a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td
+++ /dev/null
@@ -1,155 +0,0 @@
-//===-- SubsetInsertionOpInterface.td - Tensor Subsets -----*- tablegen -*-===//
-//
-// 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 SUBSET_INSERTION_OP_INTERFACE
-#define SUBSET_INSERTION_OP_INTERFACE
-
-include "mlir/IR/OpBase.td"
-
-def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
- let description = [{
- This interface can be implemented by ops that insert a source tensor into
- a destination tensor.
-
- The elements in the destination tensor that are overwritten by this
- insertion are called the "subset". How the subset is defined is up to the
- op. E.g., "tensor.insert_slice" defines the subset via a hyperrectangular
- slice. A scatter operation could define the subset via a list of indices.
-
- Ops that deal with tensor subsets come in two flavours:
- - Insertion flavor: Ops that insert a source tensor into a destination
- tensor at the specified subset. Such ops usually return a new destination
- tensor and implement the `DestinationStyleOpInterface`. Insertion ops can
- implement the `SubsetInsertionOpInterface`. Example: "tensor.insert_slice"
- - Extraction flavor: Ops that define a tensor subset. They extract a
- specified subset from a tensor. There is currently no op interface for
- such ops. Example: "tensor.extract_slice"
-
- This interface provides helper methods for efficient bufferization of
- subset-based tensor IR. Tensor subsets can bufferize to buffer "views"/
- "aliases" (in contrast to one or multiple less efficient buffer allocation).
-
- This interface is queried by One-Shot Bufferize to detect cases where a
- seeming read-after-write is not actually a conflict because the respective
- ops are operating on equivalent subsets. More details can be found in the
- documentation of One-Shot Analysis (see `areNonConflictingSubsets`).
-
- Note: This interface currently assumes that a subset op inserts a single
- tensor (source) into a destination tensor at a single subset.
- }];
- let cppNamespace = "::mlir";
- let methods = [
- InterfaceMethod<
- /*desc=*/[{
- Return the source tensor operand.
- }],
- /*retType=*/"::mlir::OpOperand &",
- /*methodName=*/"getSourceOperand",
- /*args=*/(ins)
- >,
- InterfaceMethod<
- /*desc=*/[{
- Return the destination tensor operand.
- }],
- /*retType=*/"::mlir::OpOperand &",
- /*methodName=*/"getDestinationOperand",
- /*args=*/(ins),
- /*methodBody=*/"",
- /*defaultImplementation=*/[{
- return ::mlir::detail::defaultGetDestinationOperand(
- $_op.getOperation());
- }]
- >,
- InterfaceMethod<
- /*desc=*/[{
- Return "true" if this operation inserts into a subset that is
- equivalent to the subset defined by `candidate`.
-
- Two subsets are "equivalent" and "same" if they can bufferize to the
- same buffer views/aliases. If they are "equivalent", the tensor IR
- may be expressed in terms of different SSA values (but they could
- bufferize to MemRef SSA values that can CSE without breaking
- correctness). `equivalenceFn` should return "true" if the two given
- values are equivalent.
-
- Example:
- ```
- // The subset of the SubsetInsertionOpInterface op %1 is equivalent to
- // the subset defined by %2 (but not "same"):
- %0 = arith.select %c, %t, %t : tensor<?xf32>
- %1 = tensor.insert_slice %x into %0[0][5][1]
- : tensor<5xf32> into tensor<?xf32>
- %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
-
- // The subset of the SubsetInsertionOpInterface op %1 is equivalent to
- // and "same" as the subset defined by %2.
- %1 = tensor.insert_slice %x into %t[0][5][1]
- : tensor<5xf32> into tensor<?xf32>
- %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
- ```
- }],
- /*retType=*/"bool",
- /*methodName=*/"isEquivalentSubset",
- /*args=*/(ins
- "::mlir::Value":$candidate,
- "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn)
- >,
- InterfaceMethod<
- /*desc=*/[{
- Return the subset of the destination tensor that this operation
- inserts into.
-
- Example:
- ```
- // SubsetOpInterface op:
- %0 = tensor.insert_slice %t0 into %t1[%pos][5][1]
- : tensor<5xf32> into tensor<?xf32>
- // Subset (built by this function):
- %1 = tensor.extract_slice %t1[%pos][5][1]
- : tensor<?xf32> to tensor<5xf32>
- ```
-
- Note: Implementations do not necessarily have to build new IR. They
- may return existing SSA values.
- }],
- /*retType=*/"::mlir::Value",
- /*methodName=*/"buildSubsetExtraction",
- /*args=*/(ins "::mlir::OpBuilder &":$builder, "Location":$loc)
- >,
- InterfaceMethod<
- /*desc=*/[{
- Return all SSA values that are needed (i.e., must be in scope) at the
- insertion of the builder when calling `buildSubsetExtraction`. Users
- of `buildSubsetExtraction` can use this helper method to find a
- suitable insertion point.
-
- Example: The SSA values needed to build the subset in the example of
- `buildSubsetExtraction` are %t1 and %pos.
- }],
- /*retType=*/"::llvm::SmallVector<::mlir::Value>",
- /*methodName=*/"getValuesNeededToBuildSubsetExtraction",
- /*args=*/(ins)
- >,
- ];
-
- let extraClassDeclaration = [{
- /// Return "true" if this operation inserts into the same subset as defined
- /// by `candidate`.
- ///
- /// Note: This function is useful outside of bufferization, where no tensor
- /// equivalence information is available.
- bool isSameSubset(OpResult candidate) {
- auto subsetOp = cast<::mlir::SubsetInsertionOpInterface>(
- getOperation());
- return subsetOp.isEquivalentSubset(
- candidate, [](Value v1, Value v2) { return v1 == v2; });
- }
- }];
-}
-
-#endif // SUBSET_INSERTION_OP_INTERFACE
diff --git a/mlir/include/mlir/Interfaces/SubsetOpInterface.h b/mlir/include/mlir/Interfaces/SubsetOpInterface.h
new file mode 100644
index 000000000000000..049cf2456a9c842
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.h
@@ -0,0 +1,45 @@
+//===- SubsetOpInterface.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_INTERFACES_SUBSETOPINTERFACE_H_
+#define MLIR_INTERFACES_SUBSETOPINTERFACE_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+class SubsetOpInterface;
+class SubsetExtractionOpInterface;
+class SubsetInsertionOpInterface;
+
+namespace detail {
+
+/// Return the destination/"init" operand of the op if it implements the
+/// `DestinationStyleOpInterface` and has exactly one "init" operand. Asserts
+/// otherwise.
+OpOperand &defaultGetDestinationOperand(Operation *op);
+
+/// Return the updated destination result of the op if it implements the
+/// `DestinationStyleOpInterface`.
+OpResult defaultGetUpdatedDestination(Operation *op);
+
+/// Default implementation of `isEquivalentSubset`.
+bool defaultIsEquivalentSubset(Operation *op, Value candidate,
+ function_ref<bool(Value, Value)> equivalenceFn);
+
+/// Verify `SubsetOpInterface`.
+LogicalResult verifySubsetOpInterface(SubsetOpInterface op);
+
+/// Verify `SubsetExtractionOpInterface`.
+LogicalResult verifySubsetExtractionOpInterface(SubsetExtractionOpInterface op);
+
+} // namespace detail
+} // namespace mlir
+
+#include "mlir/Interfaces/SubsetOpInterface.h.inc"
+
+#endif // MLIR_INTERFACES_SUBSETOPINTERFACE_H_
diff --git a/mlir/include/mlir/Interfaces/SubsetOpInterface.td b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
new file mode 100644
index 000000000000000..9ebed2c94818d19
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
@@ -0,0 +1,267 @@
+//===-- SubsetOpInterface.td - Tensor Subsets --------------*- tablegen -*-===//
+//
+// 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 SUBSET_OP_INTERFACE
+#define SUBSET_OP_INTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def SubsetOpInterface : OpInterface<"SubsetOpInterface"> {
+ let description = [{
+ This interface can be implemented by ops that operate on tensor subsets. A
+ "subset" is a part of a tensor. This interface describes the subset that
+ an implementing op operates on. Only the specified subset may be accessed by
+ the op.
+
+ Subset ops come in two flavours and ops that implement the
+ `SubsetOpInterface` must also implement one of the respective interfaces.
+ - Insertion flavor: Ops that insert a source value into a destination
+ tensor at the specified subset. Such ops return an updated destination
+ tensor and usually implement the `DestinationStyleOpInterface`. Insertion
+ ops must implement the `SubsetInsertionOpInterface`.
+ - Extraction flavor: Ops that extract at a subset. Extraction ops must
+ implement the `SubsetExtractionOpInterface`.
+
+ How the subset is specified is up to the implementing op. E.g.:
+ - `tensor.extract_slice/insert_slice` describe the subset as a
+ hyperrectangular slice.
+ - `tensor.gather/scatter` describe the subset as list of indices. (Not
+ implemented yet.)
+
+ Note: This interface does not expose any interface methods to get a
+ description of the accessed subset. That is because there is currently no
+ efficient way to describe arbitrary subsets. This interface merely provides
+ interface methods to check if two subsets are equivalent or disjoint.
+ }];
+
+ let cppNamespace = "::mlir";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return "true" if this op and the given candidate subset op operate on
+ equivalent subsets. Return "false" if the two subsets are disjoint
+ or cannot be proven to be equivalent.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"operatesOnEquivalentSubset",
+ /*args=*/(ins
+ "::mlir::SubsetOpInterface":$candidate,
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return "true" if this op and the given candidate subset op operate on
+ disjoint subsets. Return "false" if the two subsets are equivalent,
+ overlapping or cannot be proven to be disjoint.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"operatesOnDisjointSubset",
+ /*args=*/(ins
+ "::mlir::SubsetOpInterface":$candidate,
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn)
+ >,
+ ];
+
+ let verify = [{
+ return ::mlir::detail::verifySubsetOpInterface(
+ ::mlir::cast<::mlir::SubsetOpInterface>($_op));
+ }];
+}
+
+def SubsetExtractionOpInterface
+ : OpInterface<"SubsetExtractionOpInterface", [SubsetOpInterface]> {
+ let description = [{
+ This interface can be implemented by ops that extract a value from
+ a source tensor at a specified subset. The elements in the source tensor
+ that are read by this extraction are called "subset".
+
+ Extraction ops must have a single result value.
+ }];
+
+ let cppNamespace = "::mlir";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the source tensor operand.
+ }],
+ /*retType=*/"::mlir::OpOperand &",
+ /*methodName=*/"getSourceOperand",
+ /*args=*/(ins)
+ >,
+ ];
+
+ let verify = [{
+ return ::mlir::detail::verifySubsetExtractionOpInterface(
+ ::mlir::cast<::mlir::SubsetExtractionOpInterface>($_op));
+ }];
+
+ let extraClassDeclaration = [{
+ /// Return the single result of this op.
+ ::mlir::Value getResult() {
+ return getOperation()->getResult(0);
+ }
+ }];
+}
+
+def SubsetInsertionOpInterface
+ : OpInterface<"SubsetInsertionOpInterface", [SubsetOpInterface]> {
+ let description = [{
+ This interface can be implemented by ops that insert a source value into
+ a destination tensor at a specified subset. The elements in the destination
+ tensor that are overwritten by this insertion are called "subset". The
+ updated destination tensor is returned.
+
+ This interface provides helper methods for efficient bufferization of
+ subset-based tensor IR. Tensor subsets can bufferize to buffer "views"/
+ "aliases" (in contrast to one or multiple less efficient buffer allocation).
+
+ This interface is queried by One-Shot Bufferize to detect cases where a
+ seeming read-after-write is not actually a conflict because the respective
+ ops are operating on equivalent subsets. More details can be found in the
+ documentation of One-Shot Analysis (see `areNonConflictingSubsets`).
+ }];
+
+ let cppNamespace = "::mlir";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the source operand. The source operand can have any type.
+ }],
+ /*retType=*/"::mlir::OpOperand &",
+ /*methodName=*/"getSourceOperand",
+ /*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the destination operand. The destination operand must be a
+ tensor.
+
+ This function does not have to be implemented for destination style
+ ops that have exactly one "init" operand.
+ }],
+ /*retType=*/"::mlir::OpOperand &",
+ /*methodName=*/"getDestinationOperand",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultGetDestinationOperand(
+ $_op.getOperation());
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the updated destination result.
+
+ This function does not have to be implemented for destination style
+ ops.
+ }],
+ /*retType=*/"::mlir::OpResult",
+ /*methodName=*/"getUpdatedDestination",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultGetUpdatedDestination(
+ $_op.getOperation());
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return "true" if this operation inserts into a subset that is
+ equivalent to the subset defined by `candidate`.
+
+ Two subsets are "equivalent" and "same" if they can bufferize to the
+ same buffer views/aliases. If they are "equivalent", the tensor IR
+ may be expressed in terms of different SSA values (but they could
+ bufferize to MemRef SSA values that can CSE without breaking
+ correctness). `equivalenceFn` should return "true" if the two given
+ values are equivalent.
+
+ Example:
+ ```
+ // The subset of the SubsetInsertionOpInterface op %1 is equivalent to
+ // the subset defined by %2 (but not "same"):
+ %0 = arith.select %c, %t, %t : tensor<?xf32>
+ %1 = tensor.insert_slice %x into %0[0][5][1]
+ : tensor<5xf32> into tensor<?xf32>
+ %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
+
+ // The subset of the SubsetInsertionOpInterface op %1 is equivalent to
+ // and "same" as the subset defined by %2.
+ %1 = tensor.insert_slice %x into %t[0][5][1]
+ : tensor<5xf32> into tensor<?xf32>
+ %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32>
+ ```
+
+ Note: By default, this function calls
+ `SubsetOpInterface::operatesOnEquivalentSubset`.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"isEquivalentSubset",
+ /*args=*/(ins
+ "::mlir::Value":$candidate,
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultIsEquivalentSubset(
+ $_op.getOperation(), candidate, equivalenceFn);
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return the subset of the destination tensor that this operation
+ inserts into.
+
+ Example:
+ ```
+ // SubsetOpInterface op:
+ %0 = tensor.insert_slice %t0 into %t1[%pos][5][1]
+ : tensor<5xf32> into tensor<?xf32>
+ // Subset (built by this function):
+ %1 = tensor.extract_slice %t1[%pos][5][1]
+ : tensor<?xf32> to tensor<5xf32>
+ ```
+
+ Note: Implementations do not necessarily have to build new IR. They
+ may return existing SSA values.
+ }],
+ /*retType=*/"::mlir::Value",
+ /*methodName=*/"buildSubsetExtraction",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder, "Location":$loc)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return all SSA values that are needed (i.e., must be in scope) at the
+ insertion of the builder when calling `buildSubsetExtraction`. Users
+ of `buildSubsetExtraction` can use this helper method to find a
+ suitable insertion point.
+
+ Example: The SSA values needed to build the subset in the example of
+ `buildSubsetExtraction` are %t1 and %pos.
+ }],
+ /*retType=*/"::llvm::SmallVector<::mlir::Value>",
+ /*methodName=*/"getValuesNeededToBuildSubsetExtraction",
+ /*args=*/(ins)
+ >,
+ ];
+
+ let extraClassDeclaration = [{
+ /// Return "true" if this operation inserts into the same subset as defined
+ /// by `candidate`.
+ ///
+ /// Note: This function is useful outside of bufferization, where no tensor
+ /// equivalence information is available.
+ bool isSameSubset(OpResult candidate) {
+ auto subsetOp = cast<::mlir::SubsetInsertionOpInterface>(
+ getOperation());
+ return subsetOp.isEquivalentSubset(
+ candidate, [](Value v1, Value v2) { return v1 == v2; });
+ }
+ }];
+}
+
+#endif // SUBSET_OP_INTERFACE
diff --git a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
index 8f11c563e0cbd91..8e2986a2d1f05f6 100644
--- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
+++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h
@@ -19,6 +19,7 @@
#include <queue>
namespace mlir {
+class OffsetSizeAndStrideOpInterface;
using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>>;
@@ -134,11 +135,11 @@ class ValueBoundsConstraintSet {
std::optional<int64_t> dim, ValueRange independencies,
bool closedUB = false);
- /// Compute a constant bound for the given index-typed value or shape
- /// dimension size.
+ /// Compute a constant bound for the given affine map, where dims and symbols
+ /// are bound to the given operands. The affine map must have exactly one
+ /// result.
///
- /// `dim` must be `nullopt` if and only if `value` is index-typed. This
- /// function traverses the backward slice of the given value in a
+ /// This function traverses the backward slice of the given operands in a
/// worklist-driven manner until `stopCondition` evaluates to "true". The
/// constraint set is populated according to `ValueBoundsOpInterface` for each
/// visited value. (No constraints are added for values for which the stop
@@ -155,26 +156,12 @@ class ValueBoundsConstraintSet {
std::optional<int64_t> dim = std::nullopt,
StopConditionFn stopCondition = nullptr,
bool closedUB = false);
-
- /// Compute a constant bound for the given affine map, where dims and symbols
- /// are bound to the given operands. The affine map must have exactly one
- /// result.
- ///
- /// This function traverses the backward slice of the given operands in a
- /// worklist-driven manner until `stopCondition` evaluates to "true". The
- /// constraint set is populated according to `ValueBoundsOpInterface` for each
- /// visited value. (No constraints are added for values for which the stop
- /// condition evaluates to "true".)
- ///
- /// The stop condition is optional: If none is specified, the backward slice
- /// is traversed in a breadth-first manner until a constant bound could be
- /// computed.
- ///
- /// By default, lower/equal bounds are closed and upper bounds are open. If
- /// `closedUB` is set to "true", upper bounds are also closed.
static FailureOr<int64_t> computeConstantBound(
presburger::BoundType type, AffineMap map, ValueDimList mapOperands,
StopConditionFn stopCondition = nullptr, bool closedUB = false);
+ static FailureOr<int64_t> computeConstantBound(
+ presburger::BoundType type, AffineMap map, ArrayRef<Value> mapOperands,
+ StopConditionFn stopCondition = nullptr, bool closedUB = false);
/// Compute a constant delta between the given two values. Return "failure"
/// if a constant delta could not be determined.
@@ -195,6 +182,13 @@ class ValueBoundsConstraintSet {
std::optional<int64_t> dim1 = std::nullopt,
std::optional<int64_t> dim2 = std::nullopt);
+ /// Return "true" if the given slices are guaranteed to be overlapping.
+ /// Return "false" if the given slices are guaranteed to be non-overlapping.
+ /// Return "failure" if unknown.
+ static FailureOr<bool>
+ areOverlappingSlices(OffsetSizeAndStrideOpInterface slice1,
+ OffsetSizeAndStrideOpInterface slice2);
+
/// Add a bound for the given index-typed value or shaped value. This function
/// returns a builder that adds the bound.
BoundBuilder bound(Value value) { return BoundBuilder(*this, value); }
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 52ff6ceeee85b03..8f19245efdba6c8 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -643,6 +643,18 @@ OpOperand &MaterializeInDestinationOp::getSourceOperand() {
return getOperation()->getOpOperand(0) /*source*/;
}
+bool MaterializeInDestinationOp::operatesOnEquivalentSubset(
+ SubsetOpInterface subsetOp,
+ function_ref<bool(Value, Value)> equivalenceFn) {
+ return false;
+}
+
+bool MaterializeInDestinationOp::operatesOnDisjointSubset(
+ SubsetOpInterface subsetOp,
+ function_ref<bool(Value, Value)> equivalenceFn) {
+ return false;
+}
+
LogicalResult MaterializeInDestinationOp::verify() {
if (!isa<TensorType, BaseMemRefType>(getDest().getType()))
return emitOpError("'dest' must be a tensor or a memref");
diff --git a/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
index 385d8dc9364e379..9895db9d93ce0bb 100644
--- a/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/IR/CMakeLists.txt
@@ -23,7 +23,7 @@ add_mlir_dialect_library(MLIRBufferizationDialect
MLIRFunctionInterfaces
MLIRIR
MLIRSparseTensorDialect
- MLIRSubsetInsertionOpInterface
+ MLIRSubsetOpInterface
MLIRTensorDialect
MLIRMemRefDialect
)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
index a6876c7c824e0ce..8617c17e7a5e5e5 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
@@ -36,7 +36,7 @@ add_mlir_dialect_library(MLIRBufferizationTransforms
MLIRTensorDialect
MLIRSCFDialect
MLIRSideEffectInterfaces
- MLIRSubsetInsertionOpInterface
+ MLIRSubsetOpInterface
MLIRTransforms
MLIRViewLikeInterface
MLIRSupport
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
index 6622cfefa76a26f..4a418a05e6ff565 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorElimination.cpp
@@ -15,7 +15,7 @@
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Dominance.h"
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
#include "mlir/Pass/Pass.h"
namespace mlir {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 0bbfdba2b6e6ef9..c48402c3742a77b 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -54,7 +54,7 @@
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index e0a43a29c32d88b..7af1148bb93d5a0 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -67,7 +67,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
MLIRSCFTransforms
MLIRSCFUtils
MLIRPass
- MLIRSubsetInsertionOpInterface
+ MLIRSubsetOpInterface
MLIRSparseTensorDialect
MLIRTensorDialect
MLIRTensorTilingInterfaceImpl
diff --git a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index e0819082102ef66..6fcfa05468eeab8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -9,12 +9,36 @@
#include "mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
using namespace mlir;
using namespace mlir::linalg;
namespace {
+struct LinalgCopyOpSubsetOpInterface
+ : public SubsetOpInterface::ExternalModel<LinalgCopyOpSubsetOpInterface,
+ linalg::CopyOp> {
+ bool operatesOnEquivalentSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ // linalg.copy operates on the entire destination tensor.
+ if (auto otherCopyOp = dyn_cast<linalg::CopyOp>(candidate.getOperation()))
+ return equivalenceFn(cast<linalg::CopyOp>(op).getOutputs()[0],
+ otherCopyOp.getOutputs()[0]);
+ // In the absence of an analysis, "false" is a conservative way to implement
+ // this interface.
+ return false;
+ }
+
+ bool operatesOnDisjointSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ // In the absence of an analysis, "false" is a conservative way to implement
+ // this interface.
+ return false;
+ }
+};
+
struct LinalgCopyOpInterface
: public SubsetInsertionOpInterface::ExternalModel<LinalgCopyOpInterface,
linalg::CopyOp> {
@@ -48,9 +72,10 @@ struct LinalgCopyOpInterface
};
} // namespace
-void mlir::linalg::registerSubsetInsertionOpInterfaceExternalModels(
+void mlir::linalg::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+ linalg::CopyOp::attachInterface<LinalgCopyOpSubsetOpInterface>(*ctx);
linalg::CopyOp::attachInterface<LinalgCopyOpInterface>(*ctx);
});
}
diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
index a95443db88b50b2..2cd57e7324b4dc5 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1066,5 +1066,5 @@ void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
// Bufferization requires SubsetInsertionOpInterface models. Make sure that
// they are registered.
- tensor::registerSubsetInsertionOpInterfaceExternalModels(registry);
+ tensor::registerSubsetOpInterfaceExternalModels(registry);
}
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index a0c172ac52e4be8..c5fd4e65bbf7028 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -30,7 +30,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRMemRefDialect
MLIRPass
MLIRSCFDialect
- MLIRSubsetInsertionOpInterface
+ MLIRSubsetOpInterface
MLIRTensorDialect
MLIRTensorUtils
MLIRTilingInterface
diff --git a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
index dbda9953684f41d..7a1bafd409eea60 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -9,17 +9,115 @@
#include "mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Interfaces/SubsetInsertionOpInterface.h"
+#include "mlir/Interfaces/SubsetOpInterface.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
using namespace mlir;
using namespace mlir::tensor;
namespace {
+/// Return the tensor that the given subset op operates on.
+Value getContainerOperand(SubsetOpInterface op) {
+ if (auto extractionOp =
+ dyn_cast<SubsetExtractionOpInterface>(op.getOperation()))
+ return extractionOp.getSourceOperand().get();
+ if (auto insertionOp =
+ dyn_cast<SubsetInsertionOpInterface>(op.getOperation()))
+ return insertionOp.getDestinationOperand().get();
+ llvm_unreachable("expected SubsetExtraction/InsertionOpInterface");
+}
+
+/// Return "true" if the two ops operate on an equivalent subset.
+/// `equivalenceFn` is used to determine equivalence of tensors. Return "false"
+/// if the two ops operate non-equivalent subsets, if equivalence cannot be
+/// determined or if `op1` is not a subset op.
+template <typename OpTy>
+bool operateOnEquivalentSubsets(
+ OpTy op1, SubsetOpInterface op2,
+ function_ref<bool(Value, Value)> equivalenceFn) {
+ auto offsetsSizesAndStrides2 =
+ dyn_cast<OffsetSizeAndStrideOpInterface>(op2.getOperation());
+ if (!offsetsSizesAndStrides2)
+ return false;
+ if (!sameOffsetsSizesAndStrides(op1, offsetsSizesAndStrides2,
+ isEqualConstantIntOrValue))
+ return false;
+ return equivalenceFn(
+ getContainerOperand(cast<SubsetOpInterface>(op1.getOperation())),
+ getContainerOperand(op2));
+}
+
+/// Return "true" if the two ops operate on a disjoint subsets.
+/// `equivalenceFn` is used to determine equivalence of tensors. Return "false"
+/// if the two ops operate non-disjoint subsets, if disjointness cannot be
+/// determined or if `op1` is not a subset op.
+template <typename OpTy>
+bool operateOnDisjointSubsets(OpTy op1, SubsetOpInterface op2,
+ function_ref<bool(Value, Value)> equivalenceFn) {
+ auto offsetsSizesAndStrides2 =
+ dyn_cast<OffsetSizeAndStrideOpInterface>(op2.getOperation());
+ if (!offsetsSizesAndStrides2)
+ return false;
+ FailureOr<bool> overlappingSlices =
+ ValueBoundsConstraintSet::areOverlappingSlices(op1,
+ offsetsSizesAndStrides2);
+ if (failed(overlappingSlices) || *overlappingSlices)
+ return false;
+ return equivalenceFn(
+ getContainerOperand(cast<SubsetOpInterface>(op1.getOperation())),
+ getContainerOperand(op2));
+}
+
+struct ExtractSliceOpSubsetOpInterface
+ : public SubsetOpInterface::ExternalModel<ExtractSliceOpSubsetOpInterface,
+ tensor::ExtractSliceOp> {
+ bool operatesOnEquivalentSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+ return operateOnEquivalentSubsets(extractSliceOp, candidate, equivalenceFn);
+ }
+
+ bool operatesOnDisjointSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ auto extractSliceOp = cast<tensor::ExtractSliceOp>(op);
+ return operateOnDisjointSubsets(extractSliceOp, candidate, equivalenceFn);
+ }
+};
+
+struct ExtractSliceOpSubsetExtractionOpInterface
+ : public SubsetExtractionOpInterface::ExternalModel<
+ ExtractSliceOpSubsetExtractionOpInterface, tensor::ExtractSliceOp> {
+ OpOperand &getSourceOperand(Operation *op) const {
+ return cast<tensor::ExtractSliceOp>(op).getSourceMutable();
+ }
+};
+
template <typename OpTy>
-struct InsertSliceLikeOpInterface
+struct InsertSliceLikeOpSubsetOpInterface
+ : public SubsetOpInterface::ExternalModel<
+ InsertSliceLikeOpSubsetOpInterface<OpTy>, OpTy> {
+ bool operatesOnEquivalentSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ auto insertSliceOp = cast<OpTy>(op);
+ return operateOnEquivalentSubsets(insertSliceOp, candidate, equivalenceFn);
+ }
+
+ bool operatesOnDisjointSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ auto insertSliceOp = cast<OpTy>(op);
+ return operateOnDisjointSubsets(insertSliceOp, candidate, equivalenceFn);
+ }
+};
+
+template <typename OpTy>
+struct InsertSliceLikeOpSubsetInsertionOpInterface
: public SubsetInsertionOpInterface::ExternalModel<
- InsertSliceLikeOpInterface<OpTy>, OpTy> {
+ InsertSliceLikeOpSubsetInsertionOpInterface<OpTy>, OpTy> {
OpOperand &getSourceOperand(Operation *op) const {
return cast<OpTy>(op).getSourceMutable();
}
@@ -28,23 +126,6 @@ struct InsertSliceLikeOpInterface
return cast<OpTy>(op).getDestMutable();
}
- /// Return "true" if `insertSliceOp` inserts into a subset that is equivalent
- /// to the subset defined by `candidate`. `equivalenceFn` is used to determine
- /// equivalence of tensors.
- bool
- isEquivalentSubset(Operation *op, Value candidate,
- function_ref<bool(Value, Value)> equivalenceFn) const {
- auto insertSliceOp = cast<OpTy>(op);
- // Look for a matching tensor.extract_slice op.
- auto extractSliceOp = candidate.getDefiningOp<tensor::ExtractSliceOp>();
- if (!extractSliceOp)
- return false;
- if (!equivalenceFn(extractSliceOp.getSource(), insertSliceOp.getDest()))
- return false;
- return sameOffsetsSizesAndStrides(extractSliceOp, insertSliceOp,
- isEqualConstantIntOrValue);
- }
-
Value buildSubsetExtraction(Operation *op, OpBuilder &builder,
Location loc) const {
auto insertSliceOp = cast<OpTy>(op);
@@ -73,12 +154,22 @@ struct InsertSliceLikeOpInterface
} // namespace
-void mlir::tensor::registerSubsetInsertionOpInterfaceExternalModels(
+void mlir::tensor::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- InsertSliceOp::attachInterface<InsertSliceLikeOpInterface<InsertSliceOp>>(
+ // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface`
+ // require `SubsetOpInterface`.
+ ExtractSliceOp::attachInterface<ExtractSliceOpSubsetOpInterface>(*ctx);
+ ExtractSliceOp::attachInterface<ExtractSliceOpSubsetExtractionOpInterface>(
*ctx);
+ InsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetOpInterface<InsertSliceOp>>(*ctx);
+ InsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetInsertionOpInterface<InsertSliceOp>>(*ctx);
+ ParallelInsertSliceOp::attachInterface<
+ InsertSliceLikeOpSubsetOpInterface<ParallelInsertSliceOp>>(*ctx);
ParallelInsertSliceOp::attachInterface<
- InsertSliceLikeOpInterface<ParallelInsertSliceOp>>(*ctx);
+ InsertSliceLikeOpSubsetInsertionOpInterface<ParallelInsertSliceOp>>(
+ *ctx);
});
}
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index f74306206d63f14..2652d261f480ba4 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -16,7 +16,7 @@ set(LLVM_OPTIONAL_SOURCES
RuntimeVerifiableOpInterface.cpp
ShapedOpInterfaces.cpp
SideEffectInterfaces.cpp
- SubsetInsertionOpInterface.cpp
+ SubsetOpInterface.cpp
TilingInterface.cpp
ValueBoundsOpInterface.cpp
VectorInterfaces.cpp
@@ -84,15 +84,15 @@ add_mlir_interface_library(RuntimeVerifiableOpInterface)
add_mlir_interface_library(ShapedOpInterfaces)
add_mlir_interface_library(SideEffectInterfaces)
-add_mlir_library(MLIRSubsetInsertionOpInterface
- SubsetInsertionOpInterface.cpp
+add_mlir_library(MLIRSubsetOpInterface
+ SubsetOpInterface.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
DEPENDS
MLIRDestinationStyleOpInterface
- MLIRSubsetInsertionOpInterfaceIncGen
+ MLIRSubsetOpInterfaceIncGen
LINK_LIBS PUBLIC
MLIRDestinationStyleOpInterface
@@ -112,6 +112,7 @@ add_mlir_library(MLIRValueBoundsOpInterface
DEPENDS
MLIRDestinationStyleOpInterface
MLIRValueBoundsOpInterfaceIncGen
+ MLIRViewLikeInterface
LINK_LIBS PUBLIC
MLIRAnalysis
diff --git a/mlir/lib/Interfaces/SubsetInsertionOpInterface.cpp b/mlir/lib/Interfaces/SubsetInsertionOpInterface.cpp
deleted file mode 100644
index b2b092287f96ba6..000000000000000
--- a/mlir/lib/Interfaces/SubsetInsertionOpInterface.cpp
+++ /dev/null
@@ -1,23 +0,0 @@
-//===- SubsetInsertionOpInterface.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/Interfaces/SubsetInsertionOpInterface.h"
-#include "mlir/Interfaces/DestinationStyleOpInterface.h"
-
-#include "mlir/Interfaces/SubsetInsertionOpInterface.cpp.inc"
-
-using namespace mlir;
-
-OpOperand &detail::defaultGetDestinationOperand(Operation *op) {
- auto dstOp = dyn_cast<DestinationStyleOpInterface>(op);
- assert(dstOp && "getDestination must be implemented for non-DPS ops");
- assert(
- dstOp.getNumDpsInits() == 1 &&
- "getDestination must be implemented for ops with 0 or more than 1 init");
- return *dstOp.getDpsInitOperand(0);
-}
diff --git a/mlir/lib/Interfaces/SubsetOpInterface.cpp b/mlir/lib/Interfaces/SubsetOpInterface.cpp
new file mode 100644
index 000000000000000..7245ab20c499e20
--- /dev/null
+++ b/mlir/lib/Interfaces/SubsetOpInterface.cpp
@@ -0,0 +1,58 @@
+//===- SubsetOpInterface.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/Interfaces/SubsetOpInterface.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+
+#include "mlir/Interfaces/SubsetOpInterface.cpp.inc"
+
+using namespace mlir;
+
+OpOperand &detail::defaultGetDestinationOperand(Operation *op) {
+ auto dstOp = dyn_cast<DestinationStyleOpInterface>(op);
+ assert(dstOp && "getDestination must be implemented for non-DPS ops");
+ assert(
+ dstOp.getNumDpsInits() == 1 &&
+ "getDestination must be implemented for ops with 0 or more than 1 init");
+ return *dstOp.getDpsInitOperand(0);
+}
+
+OpResult detail::defaultGetUpdatedDestination(Operation *op) {
+ auto dstOp = dyn_cast<DestinationStyleOpInterface>(op);
+ assert(dstOp && "getUpdatedDestination must be implemented for non-DPS ops");
+ auto insertionOp = cast<SubsetInsertionOpInterface>(op);
+ return dstOp.getTiedOpResult(&insertionOp.getDestinationOperand());
+}
+
+bool detail::defaultIsEquivalentSubset(
+ Operation *op, Value candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) {
+ assert(isa<SubsetInsertionOpInterface>(op) &&
+ "expected SubsetInsertionOpInterface");
+ if (!candidate.getDefiningOp<SubsetExtractionOpInterface>())
+ return false;
+ return cast<SubsetOpInterface>(op).operatesOnEquivalentSubset(
+ candidate.getDefiningOp<SubsetOpInterface>(), equivalenceFn);
+}
+
+LogicalResult detail::verifySubsetOpInterface(SubsetOpInterface op) {
+ if (!(isa<SubsetExtractionOpInterface>(op.getOperation()) ^
+ isa<SubsetInsertionOpInterface>(op.getOperation())))
+ return op->emitOpError(
+ "SubsetOpInterface ops must implement either "
+ "SubsetExtractionOpInterface or SubsetInsertionOpInterface");
+ return success();
+}
+
+LogicalResult
+detail::verifySubsetExtractionOpInterface(SubsetExtractionOpInterface op) {
+ if (op->getNumResults() != 1)
+ return op->emitOpError(
+ "SubsetExtractionOpInterface ops must have one result");
+ return success();
+}
diff --git a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
index ff941115219f68b..200f7f3b25b731b 100644
--- a/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
+++ b/mlir/lib/Interfaces/ValueBoundsOpInterface.cpp
@@ -11,6 +11,7 @@
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/APSInt.h"
#include "llvm/Support/Debug.h"
@@ -484,6 +485,17 @@ FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
return failure();
}
+FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
+ presburger::BoundType type, AffineMap map, ArrayRef<Value> operands,
+ StopConditionFn stopCondition, bool closedUB) {
+ ValueDimList valueDims;
+ for (Value v : operands) {
+ assert(v.getType().isIndex() && "expected index type");
+ valueDims.emplace_back(v, std::nullopt);
+ }
+ return computeConstantBound(type, map, valueDims, stopCondition, closedUB);
+}
+
FailureOr<int64_t>
ValueBoundsConstraintSet::computeConstantDelta(Value value1, Value value2,
std::optional<int64_t> dim1,
@@ -512,6 +524,70 @@ ValueBoundsConstraintSet::areEqual(Value value1, Value value2,
return *delta == 0;
}
+FailureOr<bool> ValueBoundsConstraintSet::areOverlappingSlices(
+ OffsetSizeAndStrideOpInterface slice1,
+ OffsetSizeAndStrideOpInterface slice2) {
+ assert(slice1.getStaticOffsets().size() == slice1.getStaticOffsets().size() &&
+ "expected slices of same rank");
+ assert(slice1.getStaticSizes().size() == slice1.getStaticSizes().size() &&
+ "expected slices of same rank");
+ assert(slice1.getStaticStrides().size() == slice1.getStaticStrides().size() &&
+ "expected slices of same rank");
+
+ Builder b(slice1.getContext());
+ bool foundUnknownBound = false;
+ for (int64_t i = 0, e = slice1.getStaticOffsets().size(); i < e; ++i) {
+ AffineMap map =
+ AffineMap::get(/*dimCount=*/0, /*symbolCount=*/4,
+ b.getAffineSymbolExpr(0) +
+ b.getAffineSymbolExpr(1) * b.getAffineSymbolExpr(2) -
+ b.getAffineSymbolExpr(3));
+ {
+ // Case 1: Slices are guaranteed to be non-overlapping if
+ // offset1 + size1 * stride1 <= offset2 (for at least one dimension).
+ SmallVector<OpFoldResult> ofrOperands;
+ ofrOperands.push_back(slice1.getMixedOffsets()[i]);
+ ofrOperands.push_back(slice1.getMixedSizes()[i]);
+ ofrOperands.push_back(slice1.getMixedStrides()[i]);
+ ofrOperands.push_back(slice2.getMixedOffsets()[i]);
+ SmallVector<Value> valueOperands;
+ AffineMap foldedMap =
+ foldAttributesIntoMap(b, map, ofrOperands, valueOperands);
+ FailureOr<int64_t> constBound = computeConstantBound(
+ presburger::BoundType::EQ, foldedMap, valueOperands);
+ foundUnknownBound |= failed(constBound);
+ if (succeeded(constBound) && *constBound <= 0)
+ return false;
+ }
+ {
+ // Case 2: Slices are guaranteed to be non-overlapping if
+ // offset2 + size2 * stride2 <= offset1 (for at least one dimension).
+ SmallVector<OpFoldResult> ofrOperands;
+ ofrOperands.push_back(slice2.getMixedOffsets()[i]);
+ ofrOperands.push_back(slice2.getMixedSizes()[i]);
+ ofrOperands.push_back(slice2.getMixedStrides()[i]);
+ ofrOperands.push_back(slice1.getMixedOffsets()[i]);
+ SmallVector<Value> valueOperands;
+ AffineMap foldedMap =
+ foldAttributesIntoMap(b, map, ofrOperands, valueOperands);
+ FailureOr<int64_t> constBound = computeConstantBound(
+ presburger::BoundType::EQ, foldedMap, valueOperands);
+ foundUnknownBound |= failed(constBound);
+ if (succeeded(constBound) && *constBound <= 0)
+ return false;
+ }
+ }
+
+ // If at least one bound could not be computed, we cannot be certain that the
+ // slices are really overlapping.
+ if (foundUnknownBound)
+ return failure();
+
+ // All bounds could be computed and none of the above cases applied.
+ // Therefore, the slices are guaranteed to overlap.
+ return true;
+}
+
ValueBoundsConstraintSet::BoundBuilder &
ValueBoundsConstraintSet::BoundBuilder::operator[](int64_t dim) {
assert(!this->dim.has_value() && "dim was already set");
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 0a534cf1f1cda0e..0448ecbef655fa5 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -6933,7 +6933,7 @@ cc_library(
":MemRefDialect",
":Pass",
":SCFDialect",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":TensorDialect",
":TensorPassIncGen",
":TensorUtils",
@@ -10190,9 +10190,9 @@ gentbl_cc_library(
)
td_library(
- name = "SubsetInsertionOpInterfaceTdFiles",
+ name = "SubsetOpInterfaceTdFiles",
srcs = [
- "include/mlir/Interfaces/SubsetInsertionOpInterface.td",
+ "include/mlir/Interfaces/SubsetOpInterface.td",
],
includes = ["include"],
deps = [
@@ -10201,33 +10201,33 @@ td_library(
)
gentbl_cc_library(
- name = "SubsetInsertionOpInterfaceIncGen",
+ name = "SubsetOpInterfaceIncGen",
tbl_outs = [
(
["-gen-op-interface-decls"],
- "include/mlir/Interfaces/SubsetInsertionOpInterface.h.inc",
+ "include/mlir/Interfaces/SubsetOpInterface.h.inc",
),
(
["-gen-op-interface-defs"],
- "include/mlir/Interfaces/SubsetInsertionOpInterface.cpp.inc",
+ "include/mlir/Interfaces/SubsetOpInterface.cpp.inc",
),
],
tblgen = ":mlir-tblgen",
- td_file = "include/mlir/Interfaces/SubsetInsertionOpInterface.td",
+ td_file = "include/mlir/Interfaces/SubsetOpInterface.td",
deps = [
- ":SubsetInsertionOpInterfaceTdFiles",
+ ":SubsetOpInterfaceTdFiles",
],
)
cc_library(
- name = "SubsetInsertionOpInterface",
- srcs = ["lib/Interfaces/SubsetInsertionOpInterface.cpp"],
- hdrs = ["include/mlir/Interfaces/SubsetInsertionOpInterface.h"],
+ name = "SubsetOpInterface",
+ srcs = ["lib/Interfaces/SubsetOpInterface.cpp"],
+ hdrs = ["include/mlir/Interfaces/SubsetOpInterface.h"],
includes = ["include"],
deps = [
":DestinationStyleOpInterface",
":IR",
- ":SubsetInsertionOpInterfaceIncGen",
+ ":SubsetOpInterfaceIncGen",
":Support",
"//llvm:Support",
],
@@ -10471,7 +10471,7 @@ cc_library(
":SCFTransforms",
":SCFUtils",
":SparseTensorDialect",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":Support",
":TensorDialect",
":TensorTilingInterfaceImpl",
@@ -10531,6 +10531,7 @@ cc_library(
":IR",
":Support",
":ValueBoundsOpInterfaceIncGen",
+ ":ViewLikeInterface",
"//llvm:Support",
],
)
@@ -12580,7 +12581,7 @@ gentbl_cc_library(
":BufferizableOpInterfaceTdFiles",
":BufferizationOpsTdFiles",
":DestinationStyleOpInterfaceTdFiles",
- ":SubsetInsertionOpInterfaceTdFiles",
+ ":SubsetOpInterfaceTdFiles",
],
)
@@ -12620,7 +12621,7 @@ cc_library(
":InferTypeOpInterface",
":MemRefDialect",
":SparseTensorDialect",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":Support",
":TensorDialect",
"//llvm:Support",
@@ -12670,7 +12671,7 @@ cc_library(
":Pass",
":SCFDialect",
":SideEffectInterfaces",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":Support",
":TensorDialect",
":Transforms",
diff --git a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
index c83e4cc7ada23e3..348ee2beabeb061 100644
--- a/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/python/BUILD.bazel
@@ -413,7 +413,7 @@ gentbl_filegroup(
deps = [
":BufferizationOpsPyTdFiles",
"//mlir:DestinationStyleOpInterfaceTdFiles",
- "//mlir:SubsetInsertionOpInterfaceTdFiles",
+ "//mlir:SubsetOpInterfaceTdFiles",
],
)
More information about the llvm-commits
mailing list