[mlir][Transforms] LISH: Improve bypass analysis for loop-like ops #70623
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 
Member
| @llvm/pr-subscribers-mlir-bufferization @llvm/pr-subscribers-mlir-core Author: Matthias Springer (matthias-springer) ChangesImprove the bypass analysis for loop-like ops. Until now, loop-like ops were treated like any other non-subset ops: they prevent hoisting of any sort because the analysis does not know which parts of a tensor init operand are accessed by the loop-like op. With this change, the analysis can look into loop-like ops and analyze which subset they are operating on. Depends on #70535, #70617 and #70619. Only review the top commit. Patch is 90.58 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/70623.diff 34 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..e6b6d052df96a8c 100644 --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td @@ -15,7 +15,7 @@ include "mlir/Dialect/Bufferization/IR/BufferizationBase.td" include "mlir/Interfaces/DestinationStyleOpInterface.td" include "mlir/Interfaces/InferTypeOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" -include "mlir/Interfaces/SubsetInsertionOpInterface.td" +include "mlir/Interfaces/SubsetOpInterface.td" include "mlir/Interfaces/CopyOpInterface.td" class Bufferization_Op<string mnemonic, list<Trait> traits = []> @@ -220,6 +220,7 @@ def Bufferization_MaterializeInDestinationOp AllElementTypesMatch<["source", "dest"]>, BufferizableOpInterface, DestinationStyleOpInterface, DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>, + DeclareOpInterfaceMethods<SubsetOpInterface>, DeclareOpInterfaceMethods<SubsetInsertionOpInterface, ["getSourceOperand", "getValuesNeededToBuildSubsetExtraction", "buildSubsetExtraction", "isEquivalentSubset"]>, diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h index 023a46df2620109..94b0fb25b506650 100644 --- a/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/SubsetInsertionOpInterfaceImpl.h @@ -13,8 +13,7 @@ namespace mlir { class DialectRegistry; namespace linalg { -void registerSubsetInsertionOpInterfaceExternalModels( - DialectRegistry ®istry); +void registerSubsetOpInterfaceExternalModels(DialectRegistry ®istry); } // namespace linalg } // namespace mlir diff --git a/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h b/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h index e21b07d8a2705a0..019da189a8c991b 100644 --- a/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h +++ b/mlir/include/mlir/Dialect/Tensor/Transforms/SubsetInsertionOpInterfaceImpl.h @@ -13,8 +13,7 @@ namespace mlir { class DialectRegistry; namespace tensor { -void registerSubsetInsertionOpInterfaceExternalModels( - DialectRegistry ®istry); +void registerSubsetOpInterfaceExternalModels(DialectRegistry ®istry); } // namespace tensor } // namespace mlir diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h index 00f400aab5d50a0..7c2ffb7408d9afd 100644 --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -151,7 +151,7 @@ inline void registerAllDialects(DialectRegistry ®istry) { cf::registerBufferDeallocationOpInterfaceExternalModels(registry); gpu::registerBufferDeallocationOpInterfaceExternalModels(registry); linalg::registerBufferizableOpInterfaceExternalModels(registry); - linalg::registerSubsetInsertionOpInterfaceExternalModels(registry); + linalg::registerSubsetOpInterfaceExternalModels(registry); linalg::registerTilingInterfaceExternalModels(registry); linalg::registerValueBoundsOpInterfaceExternalModels(registry); memref::registerAllocationOpInterfaceExternalModels(registry); @@ -167,7 +167,7 @@ inline void registerAllDialects(DialectRegistry ®istry) { tensor::registerBufferizableOpInterfaceExternalModels(registry); tensor::registerFindPayloadReplacementOpInterfaceExternalModels(registry); tensor::registerInferTypeOpInterfaceExternalModels(registry); - tensor::registerSubsetInsertionOpInterfaceExternalModels(registry); + tensor::registerSubsetOpInterfaceExternalModels(registry); tensor::registerTilingInterfaceExternalModels(registry); tensor::registerValueBoundsOpInterfaceExternalModels(registry); vector::registerBufferizableOpInterfaceExternalModels(registry); diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt index 36a04ff0eaeaf4b..d81298bb4daf014 100644 --- a/mlir/include/mlir/Interfaces/CMakeLists.txt +++ b/mlir/include/mlir/Interfaces/CMakeLists.txt @@ -12,7 +12,7 @@ add_mlir_interface(ParallelCombiningOpInterface) add_mlir_interface(RuntimeVerifiableOpInterface) add_mlir_interface(ShapedOpInterfaces) add_mlir_interface(SideEffectInterfaces) -add_mlir_interface(SubsetInsertionOpInterface) +add_mlir_interface(SubsetOpInterface) add_mlir_interface(TilingInterface) add_mlir_interface(ValueBoundsOpInterface) add_mlir_interface(VectorInterfaces) diff --git a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h b/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h deleted file mode 100644 index 3a6dfceadcce7c0..000000000000000 --- a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.h +++ /dev/null @@ -1,27 +0,0 @@ -//===- SubsetInsertionOpInterface.h - Tensor Subsets ------------*- C++ -*-===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - -#ifndef MLIR_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_ -#define MLIR_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_ - -#include "mlir/IR/OpDefinition.h" - -namespace mlir { -namespace detail { - -/// Return the destination/"init" operand of the op if it implements the -/// `DestinationStyleOpInterface` and has exactly one "init" operand. Asserts -/// otherwise. -OpOperand &defaultGetDestinationOperand(Operation *op); - -} // namespace detail -} // namespace mlir - -#include "mlir/Interfaces/SubsetInsertionOpInterface.h.inc" - -#endif // MLIR_INTERFACES_SUBSETINSERTIONOPINTERFACE_H_ diff --git a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td b/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td deleted file mode 100644 index ef94a8ae9a60efd..000000000000000 --- a/mlir/include/mlir/Interfaces/SubsetInsertionOpInterface.td +++ /dev/null @@ -1,155 +0,0 @@ -//===-- SubsetInsertionOpInterface.td - Tensor Subsets -----*- tablegen -*-===// -// -// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. -// See https://llvm.org/LICENSE.txt for license information. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception -// -//===----------------------------------------------------------------------===// - -#ifndef SUBSET_INSERTION_OP_INTERFACE -#define SUBSET_INSERTION_OP_INTERFACE - -include "mlir/IR/OpBase.td" - -def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> { - let description = [{ - This interface can be implemented by ops that insert a source tensor into - a destination tensor. - - The elements in the destination tensor that are overwritten by this - insertion are called the "subset". How the subset is defined is up to the - op. E.g., "tensor.insert_slice" defines the subset via a hyperrectangular - slice. A scatter operation could define the subset via a list of indices. - - Ops that deal with tensor subsets come in two flavours: - - Insertion flavor: Ops that insert a source tensor into a destination - tensor at the specified subset. Such ops usually return a new destination - tensor and implement the `DestinationStyleOpInterface`. Insertion ops can - implement the `SubsetInsertionOpInterface`. Example: "tensor.insert_slice" - - Extraction flavor: Ops that define a tensor subset. They extract a - specified subset from a tensor. There is currently no op interface for - such ops. Example: "tensor.extract_slice" - - This interface provides helper methods for efficient bufferization of - subset-based tensor IR. Tensor subsets can bufferize to buffer "views"/ - "aliases" (in contrast to one or multiple less efficient buffer allocation). - - This interface is queried by One-Shot Bufferize to detect cases where a - seeming read-after-write is not actually a conflict because the respective - ops are operating on equivalent subsets. More details can be found in the - documentation of One-Shot Analysis (see `areNonConflictingSubsets`). - - Note: This interface currently assumes that a subset op inserts a single - tensor (source) into a destination tensor at a single subset. - }]; - let cppNamespace = "::mlir"; - let methods = [ - InterfaceMethod< - /*desc=*/[{ - Return the source tensor operand. - }], - /*retType=*/"::mlir::OpOperand &", - /*methodName=*/"getSourceOperand", - /*args=*/(ins) - >, - InterfaceMethod< - /*desc=*/[{ - Return the destination tensor operand. - }], - /*retType=*/"::mlir::OpOperand &", - /*methodName=*/"getDestinationOperand", - /*args=*/(ins), - /*methodBody=*/"", - /*defaultImplementation=*/[{ - return ::mlir::detail::defaultGetDestinationOperand( - $_op.getOperation()); - }] - >, - InterfaceMethod< - /*desc=*/[{ - Return "true" if this operation inserts into a subset that is - equivalent to the subset defined by `candidate`. - - Two subsets are "equivalent" and "same" if they can bufferize to the - same buffer views/aliases. If they are "equivalent", the tensor IR - may be expressed in terms of different SSA values (but they could - bufferize to MemRef SSA values that can CSE without breaking - correctness). `equivalenceFn` should return "true" if the two given - values are equivalent. - - Example: - ``` - // The subset of the SubsetInsertionOpInterface op %1 is equivalent to - // the subset defined by %2 (but not "same"): - %0 = arith.select %c, %t, %t : tensor<?xf32> - %1 = tensor.insert_slice %x into %0[0][5][1] - : tensor<5xf32> into tensor<?xf32> - %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32> - - // The subset of the SubsetInsertionOpInterface op %1 is equivalent to - // and "same" as the subset defined by %2. - %1 = tensor.insert_slice %x into %t[0][5][1] - : tensor<5xf32> into tensor<?xf32> - %2 = tensor.extract_slice %t[0][5][1] : tensor<?xf32> to tensor<5xf32> - ``` - }], - /*retType=*/"bool", - /*methodName=*/"isEquivalentSubset", - /*args=*/(ins - "::mlir::Value":$candidate, - "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn) - >, - InterfaceMethod< - /*desc=*/[{ - Return the subset of the destination tensor that this operation - inserts into. - - Example: - ``` - // SubsetOpInterface op: - %0 = tensor.insert_slice %t0 into %t1[%pos][5][1] - : tensor<5xf32> into tensor<?xf32> - // Subset (built by this function): - %1 = tensor.extract_slice %t1[%pos][5][1] - : tensor<?xf32> to tensor<5xf32> - ``` - - Note: Implementations do not necessarily have to build new IR. They - may return existing SSA values. - }], - /*retType=*/"::mlir::Value", - /*methodName=*/"buildSubsetExtraction", - /*args=*/(ins "::mlir::OpBuilder &":$builder, "Location":$loc) - >, - InterfaceMethod< - /*desc=*/[{ - Return all SSA values that are needed (i.e., must be in scope) at the - insertion of the builder when calling `buildSubsetExtraction`. Users - of `buildSubsetExtraction` can use this helper method to find a - suitable insertion point. - - Example: The SSA values needed to build the subset in the example of - `buildSubsetExtraction` are %t1 and %pos. - }], - /*retType=*/"::llvm::SmallVector<::mlir::Value>", - /*methodName=*/"getValuesNeededToBuildSubsetExtraction", - /*args=*/(ins) - >, - ]; - - let extraClassDeclaration = [{ - /// Return "true" if this operation inserts into the same subset as defined - /// by `candidate`. - /// - /// Note: This function is useful outside of bufferization, where no tensor - /// equivalence information is available. - bool isSameSubset(OpResult candidate) { - auto subsetOp = cast<::mlir::SubsetInsertionOpInterface>( - getOperation()); - return subsetOp.isEquivalentSubset( - candidate, [](Value v1, Value v2) { return v1 == v2; }); - } - }]; -} - -#endif // SUBSET_INSERTION_OP_INTERFACE diff --git a/mlir/include/mlir/Interfaces/SubsetOpInterface.h b/mlir/include/mlir/Interfaces/SubsetOpInterface.h new file mode 100644 index 000000000000000..049cf2456a9c842 --- /dev/null +++ b/mlir/include/mlir/Interfaces/SubsetOpInterface.h @@ -0,0 +1,45 @@ +//===- SubsetOpInterface.h - Tensor Subsets ---------------------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_SUBSETOPINTERFACE_H_ +#define MLIR_INTERFACES_SUBSETOPINTERFACE_H_ + +#include "mlir/IR/OpDefinition.h" + +namespace mlir { +class SubsetOpInterface; +class SubsetExtractionOpInterface; +class SubsetInsertionOpInterface; + +namespace detail { + +/// Return the destination/"init" operand of the op if it implements the +/// `DestinationStyleOpInterface` and has exactly one "init" operand. Asserts +/// otherwise. +OpOperand &defaultGetDestinationOperand(Operation *op); + +/// Return the updated destination result of the op if it implements the +/// `DestinationStyleOpInterface`. +OpResult defaultGetUpdatedDestination(Operation *op); + +/// Default implementation of `isEquivalentSubset`. +bool defaultIsEquivalentSubset(Operation *op, Value candidate, + function_ref<bool(Value, Value)> equivalenceFn); + +/// Verify `SubsetOpInterface`. +LogicalResult verifySubsetOpInterface(SubsetOpInterface op); + +/// Verify `SubsetExtractionOpInterface`. +LogicalResult verifySubsetExtractionOpInterface(SubsetExtractionOpInterface op); + +} // namespace detail +} // namespace mlir + +#include "mlir/Interfaces/SubsetOpInterface.h.inc" + +#endif // MLIR_INTERFACES_SUBSETOPINTERFACE_H_ diff --git a/mlir/include/mlir/Interfaces/SubsetOpInterface.td b/mlir/include/mlir/Interfaces/SubsetOpInterface.td new file mode 100644 index 000000000000000..07d62b8319c2961 --- /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 + 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)); + }]; + + 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 buffe... [truncated] |
matthias-springer changed the title matthias-springer changed the title This was referenced Oct 30, 2023
nicolasvasilache approved these changes Oct 31, 2023
matthias-springer force-pushed the bypass_analysis_for_loops branch from 959e925 to 5feaa44 Compare 
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
3 participants
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.
Improve the bypass analysis for loop-like ops. Until now, loop-like ops were treated like any other non-subset ops: they prevent hoisting of any sort because the analysis does not know which parts of a tensor init operand are accessed by the loop-like op. With this change, the analysis can look into loop-like ops and analyze which subset they are operating on.