[llvm] [mlir][Interfaces] Add `SubsetOpInterface` and `SubsetExtractionOpInterface` (PR #70617)
Matthias Springer via llvm-commits
llvm-commits at lists.llvm.org
Sun Oct 29 22:07:01 PDT 2023
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/70617
There is currently an op interface for subset insertion ops (`SubsetInsertionOpInterface`), but not for subset extraction ops. This commit adds `SubsetExtractionOpInterface` to `mlir/Interfaces`, as well as a common dependent op interface: `SubsetOpInterface`.
- `SubsetOpInterface` is for ops that operate on tensor subsets. It provides interface methods to check if two subset ops operate on equivalent or disjoint subsets. Ops that implement this interface must implement either `SubsetExtractionOpInterface` or `SubsetInsertionOpInterface`.
- `SubsetExtractionOpInterface` is for ops that extract from a tensor at a subset. E.g., `tensor.extract_slice`, `tensor.gather`, `vector.transfer_read`. Current implemented only on `tensor.extract_slice`.
- `SubsetInsertionOpInterface` is for ops that insert into a destination tensor at a subset. E.g., `tensor.insert_slice`, `tensor.parallel_insert_slice`, `tensor.scatter`, `vector.transfer_write`. Currently only implemented on `tensor.insert_slice`, `tensor.parallel_insert_slice`.
Other changes:
- Rename `SubsetInsertionOpInterface.td` to `SubsetOpInterface.td`.
- Add helper functions to `ValueBoundsOpInterface.cpp` for checking whether two slices are disjoint.
The new interfaces will be utilized by a new "loop-invariant subset hoisting"
transformation. (This new transform is roughly
what `Linalg/Transforms/SubsetHoisting.cpp` is doing, but in a generic
and interface-driven way.)
>From dca48435bb7b882715250cdcd731807d5e7f0ed0 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Mon, 30 Oct 2023 13:22:26 +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 ---
.../mlir/Interfaces/SubsetOpInterface.h | 45 +++++
...ionOpInterface.td => SubsetOpInterface.td} | 159 +++++++++++++++---
.../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 | 22 ++-
.../BufferizableOpInterfaceImpl.cpp | 2 +-
.../Dialect/Tensor/Transforms/CMakeLists.txt | 2 +-
.../SubsetInsertionOpInterfaceImpl.cpp | 129 +++++++++++---
mlir/lib/Interfaces/CMakeLists.txt | 9 +-
.../Interfaces/SubsetInsertionOpInterface.cpp | 23 ---
mlir/lib/Interfaces/SubsetOpInterface.cpp | 58 +++++++
.../lib/Interfaces/ValueBoundsOpInterface.cpp | 80 ++++++++-
.../llvm-project-overlay/mlir/BUILD.bazel | 33 ++--
.../mlir/python/BUILD.bazel | 2 +-
26 files changed, 504 insertions(+), 162 deletions(-)
delete mode 100644 mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h
create mode 100644 mlir/include/mlir/Interfaces/SubsetOpInterface.h
rename mlir/include/mlir/Interfaces/{SubsetInsertionOpInterface.td => SubsetOpInterface.td} (50%)
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/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/SubsetInsertionOpInterface.td b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
similarity index 50%
rename from mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td
rename to mlir/include/mlir/Interfaces/SubsetOpInterface.td
index ef94a8ae9a60efd..13e45a1f7dcd753 100644
--- a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
@@ -1,4 +1,4 @@
-//===-- SubsetInsertionOpInterface.td - Tensor Subsets -----*- tablegen -*-===//
+//===-- 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.
@@ -6,29 +6,108 @@
//
//===----------------------------------------------------------------------===//
-#ifndef SUBSET_INSERTION_OP_INTERFACE
-#define SUBSET_INSERTION_OP_INTERFACE
+#ifndef SUBSET_OP_INTERFACE
+#define SUBSET_OP_INTERFACE
include "mlir/IR/OpBase.td"
-def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
+def SubsetOpInterface : OpInterface<"SubsetOpInterface"> {
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 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
+ an equivalent subset. Return "false" is 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" is 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));
+ }];
+}
+
+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
+ updates destination tensor is returned.
This interface provides helper methods for efficient bufferization of
subset-based tensor IR. Tensor subsets can bufferize to buffer "views"/
@@ -38,15 +117,13 @@ def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
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.
+ Return the source operand. The source operand can have any type.
}],
/*retType=*/"::mlir::OpOperand &",
/*methodName=*/"getSourceOperand",
@@ -54,7 +131,11 @@ def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
>,
InterfaceMethod<
/*desc=*/[{
- Return the destination tensor operand.
+ Return the destination operand. The destination operand must be a
+ tensor.
+
+ This function does not have to be implemented for destination style
+ ops that exactly one "init".
}],
/*retType=*/"::mlir::OpOperand &",
/*methodName=*/"getDestinationOperand",
@@ -65,6 +146,22 @@ def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
$_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
@@ -92,12 +189,20 @@ def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
: 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)
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultIsEquivalentSubset(
+ $_op.getOperation(), candidate, equivalenceFn);
+ }]
>,
InterfaceMethod<
/*desc=*/[{
@@ -152,4 +257,4 @@ def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
}];
}
-#endif // SUBSET_INSERTION_OP_INTERFACE
+#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 bad246c262979b7..bb90e5ee546d113 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -66,7 +66,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..ffc18a8a6a4ff4b 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -9,12 +9,29 @@
#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 {
+ return false;
+ }
+
+ bool operatesOnDisjointSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn) const {
+ // TODO
+ return false;
+ }
+};
+
struct LinalgCopyOpInterface
: public SubsetInsertionOpInterface::ExternalModel<LinalgCopyOpInterface,
linalg::CopyOp> {
@@ -48,9 +65,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..0a05d792b5dd2cf 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.cpp
@@ -9,17 +9,109 @@
#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 {
+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 `insertSliceOp` inserts into a subset that is equivalent
+/// to the subset defined by `candidate`. `equivalenceFn` is used to determine
+/// equivalence of tensors.
+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));
+}
+
+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 +120,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 +148,20 @@ struct InsertSliceLikeOpInterface
} // namespace
-void mlir::tensor::registerSubsetInsertionOpInterfaceExternalModels(
+void mlir::tensor::registerSubsetOpInterfaceExternalModels(
DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
- InsertSliceOp::attachInterface<InsertSliceLikeOpInterface<InsertSliceOp>>(
+ 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..f0c37c872e6d31d 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"
@@ -85,7 +86,7 @@ AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
return builder.getAffineConstantExpr(shapedType.getDimSize(*dim));
} else {
// Constant index value: return directly.
- if (auto constInt = getConstantIntValue(value))
+ if (auto constInt = ::getConstantIntValue(value))
return builder.getAffineConstantExpr(*constInt);
}
@@ -102,7 +103,7 @@ AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
if (Value value = llvm::dyn_cast_if_present<Value>(ofr))
return getExpr(value, /*dim=*/std::nullopt);
- auto constInt = getConstantIntValue(ofr);
+ auto constInt = ::getConstantIntValue(ofr);
assert(constInt.has_value() && "expected Integer constant");
return builder.getAffineConstantExpr(*constInt);
}
@@ -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 b1bd543ae7fce63..add146bf084d321 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",
@@ -10189,9 +10189,9 @@ gentbl_cc_library(
)
td_library(
- name = "SubsetInsertionOpInterfaceTdFiles",
+ name = "SubsetOpInterfaceTdFiles",
srcs = [
- "include/mlir/Interfaces/SubsetInsertionOpInterface.td",
+ "include/mlir/Interfaces/SubsetOpInterface.td",
],
includes = ["include"],
deps = [
@@ -10200,33 +10200,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",
],
@@ -10470,7 +10470,7 @@ cc_library(
":SCFTransforms",
":SCFUtils",
":SparseTensorDialect",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":Support",
":TensorDialect",
":TensorTilingInterfaceImpl",
@@ -10530,6 +10530,7 @@ cc_library(
":IR",
":Support",
":ValueBoundsOpInterfaceIncGen",
+ ":ViewLikeInterface",
"//llvm:Support",
],
)
@@ -12579,7 +12580,7 @@ gentbl_cc_library(
":BufferizableOpInterfaceTdFiles",
":BufferizationOpsTdFiles",
":DestinationStyleOpInterfaceTdFiles",
- ":SubsetInsertionOpInterfaceTdFiles",
+ ":SubsetOpInterfaceTdFiles",
],
)
@@ -12619,7 +12620,7 @@ cc_library(
":InferTypeOpInterface",
":MemRefDialect",
":SparseTensorDialect",
- ":SubsetInsertionOpInterface",
+ ":SubsetOpInterface",
":Support",
":TensorDialect",
"//llvm:Support",
@@ -12669,7 +12670,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