From df296e79df1e01565995d138c9dcd779af33c336 Mon Sep 17 00:00:00 2001 From: Matthias Springer 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 | 80 +++++- .../llvm-project-overlay/mlir/BUILD.bazel | 33 +-- .../mlir/python/BUILD.bazel | 2 +- 27 files changed, 654 insertions(+), 290 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 c035190f43e39..e98b5728b38ef 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 72a4aa712f49c..e6b6d052df96a 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 traits = []> @@ -220,6 +220,7 @@ def Bufferization_MaterializeInDestinationOp AllElementTypesMatch<["source", "dest"]>, BufferizableOpInterface, DestinationStyleOpInterface, DeclareOpInterfaceMethods, + DeclareOpInterfaceMethods, DeclareOpInterfaceMethods, diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h index 023a46df26201..94b0fb25b5066 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 e21b07d8a2705..019da189a8c99 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 00f400aab5d50..7c2ffb7408d9a 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 36a04ff0eaeaf..d81298bb4daf0 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 3a6dfceadcce7..0000000000000 --- 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 ef94a8ae9a60e..0000000000000 --- 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 - %1 = tensor.insert_slice %x into %0[0][5][1] - : tensor<5xf32> into tensor - %2 = tensor.extract_slice %t[0][5][1] : tensor 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 - %2 = tensor.extract_slice %t[0][5][1] : tensor to tensor<5xf32> - ``` - }], - /*retType=*/"bool", - /*methodName=*/"isEquivalentSubset", - /*args=*/(ins - "::mlir::Value":$candidate, - "::llvm::function_ref":$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 - // Subset (built by this function): - %1 = tensor.extract_slice %t1[%pos][5][1] - : tensor 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 0000000000000..049cf2456a9c8 --- /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 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 0000000000000..9ebed2c94818d --- /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":$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":$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 + %1 = tensor.insert_slice %x into %0[0][5][1] + : tensor<5xf32> into tensor + %2 = tensor.extract_slice %t[0][5][1] : tensor 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 + %2 = tensor.extract_slice %t[0][5][1] : tensor to tensor<5xf32> + ``` + + Note: By default, this function calls + `SubsetOpInterface::operatesOnEquivalentSubset`. + }], + /*retType=*/"bool", + /*methodName=*/"isEquivalentSubset", + /*args=*/(ins + "::mlir::Value":$candidate, + "::llvm::function_ref":$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 + // Subset (built by this function): + %1 = tensor.extract_slice %t1[%pos][5][1] + : tensor 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 8f11c563e0cbd..8e2986a2d1f05 100644 --- a/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h +++ b/mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h @@ -19,6 +19,7 @@ #include namespace mlir { +class OffsetSizeAndStrideOpInterface; using ValueDimList = SmallVector>>; @@ -134,11 +135,11 @@ class ValueBoundsConstraintSet { std::optional 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 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 computeConstantBound( presburger::BoundType type, AffineMap map, ValueDimList mapOperands, StopConditionFn stopCondition = nullptr, bool closedUB = false); + static FailureOr computeConstantBound( + presburger::BoundType type, AffineMap map, ArrayRef 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 dim1 = std::nullopt, std::optional 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 + 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 52ff6ceeee85b..8f19245efdba6 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 equivalenceFn) { + return false; +} + +bool MaterializeInDestinationOp::operatesOnDisjointSubset( + SubsetOpInterface subsetOp, + function_ref equivalenceFn) { + return false; +} + LogicalResult MaterializeInDestinationOp::verify() { if (!isa(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 385d8dc9364e3..9895db9d93ce0 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 a6876c7c824e0..8617c17e7a5e5 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 6622cfefa76a2..4a418a05e6ff5 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 0bbfdba2b6e6e..c48402c3742a7 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 e0a43a29c32d8..7af1148bb93d5 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 e0819082102ef..6fcfa05468eea 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 { + bool operatesOnEquivalentSubset( + Operation *op, SubsetOpInterface candidate, + function_ref equivalenceFn) const { + // linalg.copy operates on the entire destination tensor. + if (auto otherCopyOp = dyn_cast(candidate.getOperation())) + return equivalenceFn(cast(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 equivalenceFn) const { + // In the absence of an analysis, "false" is a conservative way to implement + // this interface. + return false; + } +}; + struct LinalgCopyOpInterface : public SubsetInsertionOpInterface::ExternalModel { @@ -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(*ctx); linalg::CopyOp::attachInterface(*ctx); }); } diff --git a/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp index a95443db88b50..2cd57e7324b4d 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 a0c172ac52e4b..c5fd4e65bbf70 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 dbda9953684f4..7a1bafd409eea 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(op.getOperation())) + return extractionOp.getSourceOperand().get(); + if (auto insertionOp = + dyn_cast(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 +bool operateOnEquivalentSubsets( + OpTy op1, SubsetOpInterface op2, + function_ref equivalenceFn) { + auto offsetsSizesAndStrides2 = + dyn_cast(op2.getOperation()); + if (!offsetsSizesAndStrides2) + return false; + if (!sameOffsetsSizesAndStrides(op1, offsetsSizesAndStrides2, + isEqualConstantIntOrValue)) + return false; + return equivalenceFn( + getContainerOperand(cast(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 +bool operateOnDisjointSubsets(OpTy op1, SubsetOpInterface op2, + function_ref equivalenceFn) { + auto offsetsSizesAndStrides2 = + dyn_cast(op2.getOperation()); + if (!offsetsSizesAndStrides2) + return false; + FailureOr overlappingSlices = + ValueBoundsConstraintSet::areOverlappingSlices(op1, + offsetsSizesAndStrides2); + if (failed(overlappingSlices) || *overlappingSlices) + return false; + return equivalenceFn( + getContainerOperand(cast(op1.getOperation())), + getContainerOperand(op2)); +} + +struct ExtractSliceOpSubsetOpInterface + : public SubsetOpInterface::ExternalModel { + bool operatesOnEquivalentSubset( + Operation *op, SubsetOpInterface candidate, + function_ref equivalenceFn) const { + auto extractSliceOp = cast(op); + return operateOnEquivalentSubsets(extractSliceOp, candidate, equivalenceFn); + } + + bool operatesOnDisjointSubset( + Operation *op, SubsetOpInterface candidate, + function_ref equivalenceFn) const { + auto extractSliceOp = cast(op); + return operateOnDisjointSubsets(extractSliceOp, candidate, equivalenceFn); + } +}; + +struct ExtractSliceOpSubsetExtractionOpInterface + : public SubsetExtractionOpInterface::ExternalModel< + ExtractSliceOpSubsetExtractionOpInterface, tensor::ExtractSliceOp> { + OpOperand &getSourceOperand(Operation *op) const { + return cast(op).getSourceMutable(); + } +}; + template -struct InsertSliceLikeOpInterface +struct InsertSliceLikeOpSubsetOpInterface + : public SubsetOpInterface::ExternalModel< + InsertSliceLikeOpSubsetOpInterface, OpTy> { + bool operatesOnEquivalentSubset( + Operation *op, SubsetOpInterface candidate, + function_ref equivalenceFn) const { + auto insertSliceOp = cast(op); + return operateOnEquivalentSubsets(insertSliceOp, candidate, equivalenceFn); + } + + bool operatesOnDisjointSubset( + Operation *op, SubsetOpInterface candidate, + function_ref equivalenceFn) const { + auto insertSliceOp = cast(op); + return operateOnDisjointSubsets(insertSliceOp, candidate, equivalenceFn); + } +}; + +template +struct InsertSliceLikeOpSubsetInsertionOpInterface : public SubsetInsertionOpInterface::ExternalModel< - InsertSliceLikeOpInterface, OpTy> { + InsertSliceLikeOpSubsetInsertionOpInterface, OpTy> { OpOperand &getSourceOperand(Operation *op) const { return cast(op).getSourceMutable(); } @@ -28,23 +126,6 @@ struct InsertSliceLikeOpInterface return cast(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 equivalenceFn) const { - auto insertSliceOp = cast(op); - // Look for a matching tensor.extract_slice op. - auto extractSliceOp = candidate.getDefiningOp(); - 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(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>( + // Note: `SubsetExtractionOpInterface` and `SubsetInsertionOpInterface` + // require `SubsetOpInterface`. + ExtractSliceOp::attachInterface(*ctx); + ExtractSliceOp::attachInterface( *ctx); + InsertSliceOp::attachInterface< + InsertSliceLikeOpSubsetOpInterface>(*ctx); + InsertSliceOp::attachInterface< + InsertSliceLikeOpSubsetInsertionOpInterface>(*ctx); + ParallelInsertSliceOp::attachInterface< + InsertSliceLikeOpSubsetOpInterface>(*ctx); ParallelInsertSliceOp::attachInterface< - InsertSliceLikeOpInterface>(*ctx); + InsertSliceLikeOpSubsetInsertionOpInterface>( + *ctx); }); } diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt index f74306206d63f..2652d261f480b 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 b2b092287f96b..0000000000000 --- 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(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 0000000000000..7245ab20c499e --- /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(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(op); + assert(dstOp && "getUpdatedDestination must be implemented for non-DPS ops"); + auto insertionOp = cast(op); + return dstOp.getTiedOpResult(&insertionOp.getDestinationOperand()); +} + +bool detail::defaultIsEquivalentSubset( + Operation *op, Value candidate, + function_ref equivalenceFn) { + assert(isa(op) && + "expected SubsetInsertionOpInterface"); + if (!candidate.getDefiningOp()) + return false; + return cast(op).operatesOnEquivalentSubset( + candidate.getDefiningOp(), equivalenceFn); +} + +LogicalResult detail::verifySubsetOpInterface(SubsetOpInterface op) { + if (!(isa(op.getOperation()) ^ + isa(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 ff941115219f6..f0c37c872e6d3 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(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 ValueBoundsConstraintSet::computeConstantBound( return failure(); } +FailureOr ValueBoundsConstraintSet::computeConstantBound( + presburger::BoundType type, AffineMap map, ArrayRef 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 ValueBoundsConstraintSet::computeConstantDelta(Value value1, Value value2, std::optional dim1, @@ -512,6 +524,70 @@ ValueBoundsConstraintSet::areEqual(Value value1, Value value2, return *delta == 0; } +FailureOr 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 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 valueOperands; + AffineMap foldedMap = + foldAttributesIntoMap(b, map, ofrOperands, valueOperands); + FailureOr 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 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 valueOperands; + AffineMap foldedMap = + foldAttributesIntoMap(b, map, ofrOperands, valueOperands); + FailureOr 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 0a534cf1f1cda..0448ecbef655f 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 c83e4cc7ada23..348ee2beabeb0 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", ], )