[mlir][Interfaces] LISH: Add helpers for hyperrectangular subsets #70628
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
llvmbot
added
mlir:core
|
@llvm/pr-subscribers-mlir-bufferization @llvm/pr-subscribers-mlir-linalg Author: Matthias Springer (matthias-springer) ChangesThe majority of subset ops operate on hyperrectangular subsets. This commit adds a new optional interface method ( The comparison logic for hyperrectangular subsets (is disjoint/equivalent) is implemented with Depends on #70535, #70617, #70619 and #70623. Only review the top commit. Patch is 98.93 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/70628.diff 35 Files Affected:
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..9dc6afcaab31c86 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,8 @@ def Bufferization_MaterializeInDestinationOp
AllElementTypesMatch<["source", "dest"]>,
BufferizableOpInterface, DestinationStyleOpInterface,
DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
+ DeclareOpInterfaceMethods<SubsetOpInterface,
+ ["operatesOnEquivalentSubset", "operatesOnDisjointSubset"]>,
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/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index 8ab37c1d51d6b6c..bd68c27445744e3 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -268,6 +268,11 @@ class OpFoldResult : public PointerUnion<Attribute, Value> {
public:
void dump() const { llvm::errs() << *this << "\n"; }
+
+ MLIRContext *getContext() const {
+ return is<Attribute>() ? get<Attribute>().getContext()
+ : get<Value>().getContext();
+ }
};
// Temporarily exit the MLIR namespace to add casting support as later code in
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..98c33ec65012fca
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.h
@@ -0,0 +1,59 @@
+//===- 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"
+#include "mlir/Interfaces/ValueBoundsOpInterface.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 `SubsetInsertionOpInterface::isEquivalentSubset`.
+bool defaultIsEquivalentSubset(Operation *op, Value candidate,
+ function_ref<bool(Value, Value)> equivalenceFn);
+
+/// Default implementation of `SubsetOpInterface::operatesOnEquivalentSubset`.
+bool defaultOperatesOnEquivalentSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn);
+
+/// Default implementation of `SubsetOpInterface::operatesOnDisjointSubset`.
+bool defaultOperatesOnDisjointSubset(
+ Operation *op, SubsetOpInterface candidate,
+ function_ref<bool(Value, Value)> equivalenceFn);
+
+/// Return the container that the given subset op is operating on.
+Value getTensorContainer(Operation *op);
+
+/// 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..a00e398618a0118
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.td
@@ -0,0 +1,306 @@
+//===-- 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.)
+ }];
+
+ 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.
+
+ This interface method does not have to be implemented if
+ `getAccessedHyperrectangularSlice` is implemented.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"operatesOnEquivalentSubset",
+ /*args=*/(ins
+ "::mlir::SubsetOpInterface":$candidate,
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultOperatesOnEquivalentSubset(
+ $_op, candidate, 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.
+
+ This interface method does not have to be implemented if
+ `getAccessedHyperrectangularSlice` is implemented.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"operatesOnDisjointSubset",
+ /*args=*/(ins
+ "::mlir::SubsetOpInterface":$candidate,
+ "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::detail::defaultOperatesOnDisjointSubset(
+ $_op, candidate, equivalenceFn);
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ ...
[truncated]
|
This was referenced Oct 30, 2023
matthias-springer
force-pushed
the
hyperrectangular_slice
branch
from
387be39 to
7566fe1
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The majority of subset ops operate on hyperrectangular subsets. This commit adds a new optional interface method (
getAccessedHyperrectangularSlice) that can be implemented by such subset ops. If implemented, the otheroperatesOn...interface methods of theSubsetOpInterfacedo not have to be implemented anymore.The comparison logic for hyperrectangular subsets (is disjoint/equivalent) is implemented with
ValueBoundsOpInterface. This makes the subset hoisting more powerful: simple cases where two different SSA values always have the same runtime value can now be supported.