//===- OneShotAnalysis.cpp - One-Shot (Single Pass) Analysis --------------===// // // 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 // //===----------------------------------------------------------------------===// // // One-Shot Analysis analyzes function bodies. By default, function boundaries // (FuncOp bbArgs, CallOps, ReturnOps) are treated as "unknown" ops. // OneShotModuleBufferization.cpp is an extension of One-Shot Analysis for // simple call graphs without loops. // // One-Shot Bufferize consists of three phases. // // 1. Analyze ops to decide which OpOperands can bufferize inplace, i.e., // without inserting buffer copies. The analysis queries op bufferization // semantics via `BufferizableOpInterface`. // 2. Insert copies for OpOperands that were decided to bufferize out-of-place // in tensor land during `TensorCopyInsertion`. // 3. Bufferize ops by calling `BufferizableOpInterface::bufferize`. // // This file contains only the analysis. For convenience, this file also // contains a helper function `runOneShotBufferize` that analyzes an op (and its // nested ops) and then bufferizes it. // // Inplace bufferization decisions are passed from the analysis to the // `TensorCopyInsertion` phase via `AnalysisState`. They can be printed for // debugging purposes with `testAnalysisOnly`. // // Ops that do not implement `BufferizableOpInterface` can be analyzed but are // treated conservatively. E.g., the analysis has to assume that their tensor // OpOperands bufferize to memory writes. While such ops can be analyzed, they // are not bufferized and remain in the IR. to_tensor and to_memref ops are // inserted at the bufferization boundary. // // This analysis caters to high-performance codegen where buffer reuse is deemed // critical: the analysis should fail if the bufferized form of the function // needs to return a buffer, unless `allowReturnAllocs` is enabled. #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h" #include <optional> #include <random> #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h" #include "mlir/Dialect/Bufferization/Transforms/Transforms.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/AsmState.h" #include "mlir/IR/Dominance.h" #include "mlir/IR/Iterators.h" #include "mlir/IR/Operation.h" #include "mlir/IR/TypeUtilities.h" #include "mlir/Interfaces/ControlFlowInterfaces.h" #include "mlir/Interfaces/SubsetOpInterface.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SetVector.h" MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::bufferization::OneShotAnalysisState) // Run mlir-opt with `-debug-only="one-shot-analysis"` for detailed debug // output. #define DEBUG_TYPE … usingnamespacemlir; usingnamespacemlir::bufferization; static bool isaTensor(Type t) { … } //===----------------------------------------------------------------------===// // Bufferization-specific attribute manipulation. // These are for testing and debugging only. Bufferization information is stored // in OneShotBufferizationState. When run with `testAnalysisOnly`, the IR is // annotated with the results of the analysis, so that they can be checked in // tests. //===----------------------------------------------------------------------===// /// Attribute marker to specify op operands that bufferize in-place. constexpr StringLiteral kInPlaceOperandsAttrName = …; constexpr StringLiteral kOpResultAliasSetAttrName = …; constexpr StringLiteral kBbArgAliasSetAttrName = …; /// Mark whether OpOperand will be bufferized inplace. static void setInPlaceOpOperand(OpOperand &opOperand, bool inPlace) { … } //===----------------------------------------------------------------------===// // OneShotAnalysisState //===----------------------------------------------------------------------===// OneShotAnalysisState::OneShotAnalysisState( Operation *op, const OneShotBufferizationOptions &options) : … { … } void OneShotAnalysisState::applyOnEquivalenceClass( Value v, function_ref<void(Value)> fun) const { … } void OneShotAnalysisState::applyOnAliases(Value v, function_ref<void(Value)> fun) const { … } bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1, Value v2) const { … } bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1, Value v2) const { … } void OneShotAnalysisState::bufferizeInPlace(OpOperand &operand) { … } void OneShotAnalysisState::bufferizeOutOfPlace(OpOperand &operand) { … } void OneShotAnalysisState::createAliasInfoEntry(Value v) { … } void OneShotAnalysisState::gatherUndefinedTensorUses(Operation *op) { … } bool OneShotAnalysisState::hasUndefinedContents(OpOperand *opOperand) const { … } bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const { … } bool OneShotAnalysisState::isValueWritten(Value value) const { … } bool OneShotAnalysisState::isWritable(Value value) const { … } void OneShotAnalysisState::unionAliasSets(Value v1, Value v2) { … } void OneShotAnalysisState::unionEquivalenceClasses(Value v1, Value v2) { … } OneShotAnalysisState::Extension::~Extension() = default; //===----------------------------------------------------------------------===// // Bufferization-specific alias analysis. //===----------------------------------------------------------------------===// /// Return true if opOperand has been decided to bufferize in-place. static bool isInplaceMemoryWrite(OpOperand &opOperand, const OneShotAnalysisState &state) { … } /// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors /// properly dominates `b` and `b` is not inside `a`. static bool happensBefore(Operation *a, Operation *b, const DominanceInfo &domInfo) { … } static bool isReachable(Block *from, Block *to, ArrayRef<Block *> except) { … } /// Return `true` if op dominance can be used to rule out a read-after-write /// conflicts based on the ordering of ops. Returns `false` if op dominance /// cannot be used to due region-based loops. /// /// Generalized op dominance can often be used to rule out potential conflicts /// due to "read happens before write". E.g., the following IR is not a RaW /// conflict because the read happens *before* the write. /// /// Example 1: /// %0 = ... : tensor<?xf32> // DEF /// "reading_op"(%0) : tensor<?xf32> // READ /// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE /// /// This is no longer true inside loops (or repetitive regions). In such cases, /// there may not be a meaningful `happensBefore` relationship because ops /// could be executed multiple times. E.g.: /// /// Example 2: /// %0 = ... : tensor<?xf32> // DEF /// scf.for ... { /// "reading_op"(%0) : tensor<?xf32> // READ /// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE /// ... /// } /// /// In the above example, reading_op happens before writing_op according to /// op dominance. However, both ops may happen multiple times; in /// particular, the second execution of reading_op happens after the first /// execution of writing_op. This is problematic because the tensor %0 they /// operate on (i.e., the "definition") is defined outside of the loop. /// /// On a high-level, there is a potential RaW in a program if there exists a /// possible program execution such that there is a sequence of DEF, followed /// by WRITE, followed by READ. Each additional DEF resets the sequence. /// /// E.g.: /// No conflict: DEF, WRITE, DEF, READ /// Potential conflict: DEF, READ, WRITE, READ, WRITE /// /// Example 1 has no conflict: DEF, READ, WRITE /// Example 2 has a potential conflict: DEF, (READ, WRITE)* // /// Example 3: /// scf.for ... { /// %0 = ... : tensor<?xf32> /// "reading_op"(%0) : tensor<?xf32> /// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> /// ... /// } /// This has no conflict: (DEF, READ, WRITE)* /// /// Example 4: /// %0 = ... : tensor<?xf32> /// scf.for ... { /// scf.for ... { "reading_op"(%0) } /// %1 = "writing_op"(%0) /// } /// This has a potential conflict: DEF, ((READ)*, WRITE)* /// /// Example 5: /// %0 = ... : tensor<?xf32> /// scf.for ... { %1 = "writing_op"(%0) } /// scf.for ... { "reading_op"(%0) } /// This has a potential conflict: DEF, WRITE*, READ* /// /// The following rules are used to rule out RaW conflicts via ordering of ops: /// /// 1. If the closest enclosing repetitive region of DEF is a proper ancestor of /// a repetitive region that enclosing both READ and WRITE, we cannot rule /// out RaW conflict due to the ordering of ops. /// 2. Otherwise: There are no loops that interfere with our analysis; for /// analysis purposes, we can assume that there are no loops/repetitive /// regions. I.e., we can rule out a RaW conflict if READ happensBefore WRITE /// or WRITE happensBefore DEF. (Checked in `hasReadAfterWriteInterference`.) /// static bool canUseOpDominanceDueToRegions(OpOperand *uRead, OpOperand *uWrite, const SetVector<Value> &definitions, AnalysisState &state) { … } /// Return `true` if op dominance can be used to rule out a read-after-write /// conflicts based on the ordering of ops. Returns `false` if op dominance /// cannot be used to due block-based loops within a region. /// /// Refer to the `canUseOpDominanceDueToRegions` documentation for details on /// how op domiance is used during RaW conflict detection. /// /// On a high-level, there is a potential RaW in a program if there exists a /// possible program execution such that there is a sequence of DEF, followed /// by WRITE, followed by READ. Each additional DEF resets the sequence. /// /// Op dominance cannot be used if there is a path from block(READ) to /// block(WRITE) and a path from block(WRITE) to block(READ). block(DEF) should /// not appear on that path. static bool canUseOpDominanceDueToBlocks(OpOperand *uRead, OpOperand *uWrite, const SetVector<Value> &definitions, AnalysisState &state) { … } static bool canUseOpDominance(OpOperand *uRead, OpOperand *uWrite, const SetVector<Value> &definitions, AnalysisState &state) { … } /// Annotate IR with details about the detected RaW conflict. static void annotateConflict(OpOperand *uRead, OpOperand *uConflictingWrite, Value definition) { … } /// Return 'true' if a tensor that is equivalent to `other` can be found in the /// reverse use-def chain of `start`. Note: If an OpOperand bufferizes out of /// place along that use-def chain, the two tensors may not materialize as /// equivalent buffers (but separate allocations). /// /// Note: This function also requires that the two tensors have equivalent /// indexing. I.e., the tensor types do not change along the use-def chain, /// apart from static <-> dynamic dim casts. static bool hasEquivalentValueInReverseUseDefChain(AnalysisState &state, Value start, Value other) { … } /// Return "true" if `value` is originating from a subset that is equivalent to /// the subset that `subsetOp` inserts into. static bool matchesInsertDestination(const AnalysisState &state, Value value, SubsetInsertionOpInterface subsetOp) { … } /// Return "true" if the given "read" and potentially conflicting "write" are /// not conflicting due to their subset relationship. The comments in this /// function are expressed in terms of tensor.extract_slice/tensor.insert_slice /// pairs, but apply to any subset ops that implement the /// `SubsetInsertionOpInterface`. static bool areNonConflictingSubsets(OpOperand *uRead, OpOperand *uConflictingWrite, const AnalysisState &state) { … } /// Given sets of uses and writes, return true if there is a RaW conflict under /// the assumption that all given reads/writes alias the same buffer and that /// all given writes bufferize inplace. /// /// A conflict is: According to SSA use-def chains, a read R is supposed to read /// the result of a definition W1. But because of bufferization decisions, R /// actually reads another definition W2. static bool hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead, const DenseSet<OpOperand *> &usesWrite, const DominanceInfo &domInfo, OneShotAnalysisState &state) { … } // Helper function to iterate on aliases of `root` and capture the writes. static void getAliasingInplaceWrites(DenseSet<OpOperand *> &res, Value root, const OneShotAnalysisState &state) { … } // Helper function to iterate on aliases of `root` and capture the reads. static void getAliasingReads(DenseSet<OpOperand *> &res, Value root, const OneShotAnalysisState &state) { … } /// Return true if bufferizing `operand` inplace would create a conflict. A read /// R and a write W of the same alias set is a conflict if inplace bufferization /// of W changes the value read by R to a value different from the one that /// would be expected by tracing back R's origin through SSA use-def chains. /// A conflict can only be introduced by a new alias and/or an inplace /// bufferization decision. /// /// Example: /// %0 = tensor.extract_slice %t[...][...][1, 1] {inplace?} /// %1 = vector.transfer_write %v1, %t {inplace} : vector<5xf32>, tensor<?xf32> /// %e = tensor.extract_slice %1 /// %2 = vector.transfer_write %v2, %0 {inplace} : vector<6xf32>, tensor<?xf32> /// %3 = vector.transfer_read %e, %cst : tensor<?xf32>, vector<7xf32> /// /// In the above example, the two TransferWriteOps have already been decided to /// bufferize inplace. Bufferizing the ExtractSliceOp inplace would create a /// conflict because: /// * According to SSA use-def chains, we expect to read the result of %1. /// * However, adding an alias {%0, %t} would mean that the second /// TransferWriteOp overwrites the result of the first one. Therefore, the /// TransferReadOp would no longer be reading the result of %1. /// /// If `checkConsistencyOnly` is true, this function checks if there is a /// read-after-write conflict without bufferizing `operand` inplace. This would /// indicate a problem with the current inplace bufferization decisions. /// /// Note: If `checkConsistencyOnly`, this function may be called with a null /// OpResult. In that case, only the consistency of bufferization decisions /// involving aliases of the given OpOperand are checked. static bool wouldCreateReadAfterWriteInterference( OpOperand &operand, const DominanceInfo &domInfo, OneShotAnalysisState &state, bool checkConsistencyOnly = false) { … } /// Annotate IR with details about the detected non-writability conflict. static void annotateNonWritableTensor(Value value) { … } /// Return true if bufferizing `operand` inplace would create a write to a /// non-writable buffer. static bool wouldCreateWriteToNonWritableBuffer(OpOperand &operand, OneShotAnalysisState &state, bool checkConsistencyOnly = false) { … } //===----------------------------------------------------------------------===// // Bufferization analyses. //===----------------------------------------------------------------------===// // Find the values that define the contents of the given value. const llvm::SetVector<Value> & OneShotAnalysisState::findDefinitionsCached(Value value) { … } void OneShotAnalysisState::resetCache() { … } /// Determine if `operand` can be bufferized in-place. static LogicalResult bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state, const DominanceInfo &domInfo) { … } LogicalResult OneShotAnalysisState::analyzeSingleOp(Operation *op, const DominanceInfo &domInfo) { … } /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops. static void equivalenceAnalysis(SmallVector<Operation *> &ops, OneShotAnalysisState &state) { … } /// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained /// in `op`. static void equivalenceAnalysis(Operation *op, OneShotAnalysisState &state) { … } /// "Bottom-up from terminators" heuristic. static SmallVector<Operation *> bottomUpFromTerminatorsHeuristic(Operation *op, const OneShotAnalysisState &state) { … } LogicalResult OneShotAnalysisState::analyzeOp(Operation *op, const DominanceInfo &domInfo) { … } /// Perform various checks on the input IR to see if it contains IR constructs /// that are unsupported by One-Shot Bufferize. static LogicalResult checkPreBufferizationAssumptions(Operation *op, const DominanceInfo &domInfo, OneShotAnalysisState &state) { … } /// Annotate the IR with the result of the analysis. For testing/debugging only. static void annotateOpsWithBufferizationMarkers(Operation *op, const OneShotAnalysisState &state) { … } static void annotateOpsWithAliasSets(Operation *op, const OneShotAnalysisState &state) { … } LogicalResult bufferization::analyzeOp(Operation *op, OneShotAnalysisState &state, BufferizationStatistics *statistics) { … } LogicalResult bufferization::runOneShotBufferize(Operation *op, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics) { … }