1379 lines · cpp
1//===- OneShotAnalysis.cpp - One-Shot (Single Pass) Analysis --------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// One-Shot Analysis analyzes function bodies. By default, function boundaries10// (FuncOp bbArgs, CallOps, ReturnOps) are treated as "unknown" ops.11// OneShotModuleBufferization.cpp is an extension of One-Shot Analysis for12// simple call graphs without loops.13//14// One-Shot Bufferize consists of three phases.15//16// 1. Analyze ops to decide which OpOperands can bufferize inplace, i.e.,17// without inserting buffer copies. The analysis queries op bufferization18// semantics via `BufferizableOpInterface`.19// 2. Insert copies for OpOperands that were decided to bufferize out-of-place20// in tensor land during `TensorCopyInsertion`.21// 3. Bufferize ops by calling `BufferizableOpInterface::bufferize`.22//23// This file contains only the analysis. For convenience, this file also24// contains a helper function `runOneShotBufferize` that analyzes an op (and its25// nested ops) and then bufferizes it.26//27// Inplace bufferization decisions are passed from the analysis to the28// `TensorCopyInsertion` phase via `AnalysisState`. They can be printed for29// debugging purposes with `testAnalysisOnly`.30//31// Ops that do not implement `BufferizableOpInterface` can be analyzed but are32// treated conservatively. E.g., the analysis has to assume that their tensor33// OpOperands bufferize to memory writes. While such ops can be analyzed, they34// are not bufferized and remain in the IR. to_tensor and to_buffer ops are35// inserted at the bufferization boundary.36//37// This analysis caters to high-performance codegen where buffer reuse is deemed38// critical: the analysis should fail if the bufferized form of the function39// needs to return a buffer, unless `allowReturnAllocs` is enabled.40 41#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"42 43#include <random>44 45#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"46#include "mlir/Dialect/Bufferization/IR/Bufferization.h"47#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"48#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"49#include "mlir/Dialect/MemRef/IR/MemRef.h"50#include "mlir/IR/AsmState.h"51#include "mlir/IR/Dominance.h"52#include "mlir/IR/Iterators.h"53#include "mlir/IR/Operation.h"54#include "mlir/IR/TypeUtilities.h"55#include "mlir/Interfaces/ControlFlowInterfaces.h"56#include "mlir/Interfaces/SubsetOpInterface.h"57#include "llvm/ADT/DenseSet.h"58#include "llvm/ADT/SetVector.h"59#include "llvm/Support/DebugLog.h"60 61MLIR_DEFINE_EXPLICIT_TYPE_ID(mlir::bufferization::OneShotAnalysisState)62 63// Run mlir-opt with `-debug-only="one-shot-analysis"` for detailed debug64// output.65#define DEBUG_TYPE "one-shot-analysis"66 67using namespace mlir;68using namespace mlir::bufferization;69 70static bool isaTensor(Type t) { return isa<TensorType>(t); }71 72//===----------------------------------------------------------------------===//73// Bufferization-specific attribute manipulation.74// These are for testing and debugging only. Bufferization information is stored75// in OneShotBufferizationState. When run with `testAnalysisOnly`, the IR is76// annotated with the results of the analysis, so that they can be checked in77// tests.78//===----------------------------------------------------------------------===//79 80/// Attribute marker to specify op operands that bufferize in-place.81constexpr StringLiteral kInPlaceOperandsAttrName = "__inplace_operands_attr__";82 83constexpr StringLiteral kOpResultAliasSetAttrName =84 "__opresult_alias_set_attr__";85 86constexpr StringLiteral kBbArgAliasSetAttrName = "__bbarg_alias_set_attr__";87 88/// Mark whether OpOperand will be bufferized inplace.89static void setInPlaceOpOperand(OpOperand &opOperand, bool inPlace) {90 Operation *op = opOperand.getOwner();91 SmallVector<StringRef> inPlaceVector;92 if (auto attr = op->getAttr(kInPlaceOperandsAttrName)) {93 inPlaceVector = SmallVector<StringRef>(llvm::to_vector<4>(94 cast<ArrayAttr>(attr).getAsValueRange<StringAttr>()));95 } else {96 inPlaceVector = SmallVector<StringRef>(op->getNumOperands(), "none");97 for (OpOperand &opOperand : op->getOpOperands())98 if (isa<TensorType>(opOperand.get().getType()))99 inPlaceVector[opOperand.getOperandNumber()] = "false";100 }101 inPlaceVector[opOperand.getOperandNumber()] = inPlace ? "true" : "false";102 op->setAttr(kInPlaceOperandsAttrName,103 OpBuilder(op).getStrArrayAttr(inPlaceVector));104}105 106//===----------------------------------------------------------------------===//107// OneShotAnalysisState108//===----------------------------------------------------------------------===//109 110OneShotAnalysisState::OneShotAnalysisState(111 Operation *op, const OneShotBufferizationOptions &options)112 : AnalysisState(options, TypeID::get<OneShotAnalysisState>()) {113 // Set up alias sets.114 op->walk([&](Operation *op) {115 for (Value v : op->getResults())116 if (isa<TensorType>(v.getType()))117 createAliasInfoEntry(v);118 for (Region &r : op->getRegions())119 for (Block &b : r.getBlocks())120 for (auto bbArg : b.getArguments())121 if (isa<TensorType>(bbArg.getType()))122 createAliasInfoEntry(bbArg);123 });124 125 // Mark OpOperands in-place that must bufferize in-place.126 op->walk([&](BufferizableOpInterface bufferizableOp) {127 if (!options.isOpAllowed(bufferizableOp))128 return WalkResult::skip();129 for (OpOperand &opOperand : bufferizableOp->getOpOperands())130 if (isa<TensorType>(opOperand.get().getType()))131 if (bufferizableOp.mustBufferizeInPlace(opOperand, *this))132 bufferizeInPlace(opOperand);133 return WalkResult::advance();134 });135}136 137void OneShotAnalysisState::applyOnEquivalenceClass(138 Value v, function_ref<void(Value)> fun) const {139 auto leaderIt = equivalentInfo.findLeader(v);140 for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;141 ++mit) {142 fun(*mit);143 }144}145 146void OneShotAnalysisState::applyOnAliases(Value v,147 function_ref<void(Value)> fun) const {148 auto leaderIt = aliasInfo.findLeader(v);149 for (auto mit = leaderIt, meit = aliasInfo.member_end(); mit != meit; ++mit) {150 fun(*mit);151 }152}153 154bool OneShotAnalysisState::areEquivalentBufferizedValues(Value v1,155 Value v2) const {156 return equivalentInfo.isEquivalent(v1, v2);157}158 159bool OneShotAnalysisState::areAliasingBufferizedValues(Value v1,160 Value v2) const {161 return aliasInfo.isEquivalent(v1, v2);162}163 164void OneShotAnalysisState::bufferizeInPlace(OpOperand &operand) {165 if (inplaceBufferized.contains(&operand))166 return;167 inplaceBufferized.insert(&operand);168 for (AliasingValue alias : getAliasingValues(operand))169 aliasInfo.unionSets(alias.value, operand.get());170 ++statNumTensorInPlace;171}172 173void OneShotAnalysisState::bufferizeOutOfPlace(OpOperand &operand) {174 assert(!inplaceBufferized.contains(&operand) &&175 "OpOperand was already decided to bufferize inplace");176 ++statNumTensorOutOfPlace;177}178 179void OneShotAnalysisState::createAliasInfoEntry(Value v) {180 aliasInfo.insert(v);181 equivalentInfo.insert(v);182}183 184void OneShotAnalysisState::gatherUndefinedTensorUses(Operation *op) {185 op->walk([&](Operation *op) {186 // Skip unknown ops.187 auto bufferizableOp = getOptions().dynCastBufferizableOp(op);188 if (!bufferizableOp)189 return WalkResult::skip();190 191 // Check all tensor OpResults.192 for (OpResult opResult : op->getOpResults()) {193 if (!isa<TensorType>(opResult.getType()))194 continue;195 196 // If there is no preceding definition, the tensor contents are197 // undefined.198 if (opResult.getUses().empty())199 continue;200 // It does not really matter which use to take to search about201 // the value's definitions.202 OpOperand *opOperand = &(*opResult.getUses().begin());203 if (findDefinitionsCached(opOperand).empty())204 for (OpOperand &use : opResult.getUses())205 undefinedTensorUses.insert(&use);206 }207 208 return WalkResult::advance();209 });210}211 212bool OneShotAnalysisState::hasUndefinedContents(OpOperand *opOperand) const {213 return undefinedTensorUses.contains(opOperand);214}215 216bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {217 return inplaceBufferized.contains(&opOperand);218}219 220bool OneShotAnalysisState::isValueWritten(Value value) const {221 bool isWritten = false;222 applyOnAliases(value, [&](Value val) {223 for (OpOperand &use : val.getUses())224 if (isInPlace(use) && bufferizesToMemoryWrite(use))225 isWritten = true;226 });227 return isWritten;228}229 230bool OneShotAnalysisState::isWritable(Value value) const {231 // TODO: Out-of-place bufferized value could be considered writable.232 // Query BufferizableOpInterface to see if the BlockArgument is writable.233 if (auto bufferizableOp =234 getOptions().dynCastBufferizableOp(getOwnerOfValue(value)))235 return bufferizableOp.isWritable(value, *this);236 237 // Not a bufferizable op: The conservative answer is "not writable".238 return false;239}240 241void OneShotAnalysisState::unionAliasSets(Value v1, Value v2) {242 aliasInfo.unionSets(v1, v2);243}244 245void OneShotAnalysisState::unionEquivalenceClasses(Value v1, Value v2) {246 equivalentInfo.unionSets(v1, v2);247}248 249OneShotAnalysisState::Extension::~Extension() = default;250 251//===----------------------------------------------------------------------===//252// Bufferization-specific alias analysis.253//===----------------------------------------------------------------------===//254 255/// Return true if opOperand has been decided to bufferize in-place.256static bool isInplaceMemoryWrite(OpOperand &opOperand,257 const OneShotAnalysisState &state) {258 // OpOperands that do not bufferize to a memory write do not write in-place.259 if (!state.bufferizesToMemoryWrite(opOperand))260 return false;261 // Check current bufferization decisions.262 return state.isInPlace(opOperand);263}264 265/// Return true if `a` happens before `b`, i.e., `a` or one of its ancestors266/// properly dominates `b` and `b` is not inside `a`.267static bool happensBefore(Operation *a, Operation *b,268 const DominanceInfo &domInfo) {269 do {270 // TODO: Instead of isProperAncestor + properlyDominates, we should use271 // properlyDominatesImpl(a, b, /*enclosingOpOk=*/false)272 if (a->isProperAncestor(b))273 return false;274 if (domInfo.properlyDominates(a, b))275 return true;276 } while ((a = a->getParentOp()));277 return false;278}279 280/// Return `true` if op dominance can be used to rule out a read-after-write281/// conflicts based on the ordering of ops. Returns `false` if op dominance282/// cannot be used to due region-based loops.283///284/// Generalized op dominance can often be used to rule out potential conflicts285/// due to "read happens before write". E.g., the following IR is not a RaW286/// conflict because the read happens *before* the write.287///288/// Example 1:289/// %0 = ... : tensor<?xf32> // DEF290/// "reading_op"(%0) : tensor<?xf32> // READ291/// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE292///293/// This is no longer true inside loops (or repetitive regions). In such cases,294/// there may not be a meaningful `happensBefore` relationship because ops295/// could be executed multiple times. E.g.:296///297/// Example 2:298/// %0 = ... : tensor<?xf32> // DEF299/// scf.for ... {300/// "reading_op"(%0) : tensor<?xf32> // READ301/// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE302/// ...303/// }304///305/// In the above example, reading_op happens before writing_op according to306/// op dominance. However, both ops may happen multiple times; in307/// particular, the second execution of reading_op happens after the first308/// execution of writing_op. This is problematic because the tensor %0 they309/// operate on (i.e., the "definition") is defined outside of the loop.310///311/// On a high-level, there is a potential RaW in a program if there exists a312/// possible program execution such that there is a sequence of DEF, followed313/// by WRITE, followed by READ. Each additional DEF resets the sequence.314///315/// E.g.:316/// No conflict: DEF, WRITE, DEF, READ317/// Potential conflict: DEF, READ, WRITE, READ, WRITE318///319/// Example 1 has no conflict: DEF, READ, WRITE320/// Example 2 has a potential conflict: DEF, (READ, WRITE)*321//322/// Example 3:323/// scf.for ... {324/// %0 = ... : tensor<?xf32>325/// "reading_op"(%0) : tensor<?xf32>326/// %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32>327/// ...328/// }329/// This has no conflict: (DEF, READ, WRITE)*330///331/// Example 4:332/// %0 = ... : tensor<?xf32>333/// scf.for ... {334/// scf.for ... { "reading_op"(%0) }335/// %1 = "writing_op"(%0)336/// }337/// This has a potential conflict: DEF, ((READ)*, WRITE)*338///339/// Example 5:340/// %0 = ... : tensor<?xf32>341/// scf.for ... { %1 = "writing_op"(%0) }342/// scf.for ... { "reading_op"(%0) }343/// This has a potential conflict: DEF, WRITE*, READ*344///345/// The following rules are used to rule out RaW conflicts via ordering of ops:346///347/// 1. If the closest enclosing repetitive region of DEF is a proper ancestor of348/// a repetitive region that enclosing both READ and WRITE, we cannot rule349/// out RaW conflict due to the ordering of ops.350/// 2. Otherwise: There are no loops that interfere with our analysis; for351/// analysis purposes, we can assume that there are no loops/repetitive352/// regions. I.e., we can rule out a RaW conflict if READ happensBefore WRITE353/// or WRITE happensBefore DEF. (Checked in `hasReadAfterWriteInterference`.)354///355static bool canUseOpDominanceDueToRegions(OpOperand *uRead, OpOperand *uWrite,356 const SetVector<Value> &definitions,357 AnalysisState &state) {358 const BufferizationOptions &options = state.getOptions();359 for (Value def : definitions) {360 Region *rRead =361 state.getEnclosingRepetitiveRegion(uRead->getOwner(), options);362 Region *rDef = state.getEnclosingRepetitiveRegion(def, options);363 364 // READ and DEF are in the same repetitive region. `happensBefore` can be365 // used to rule out RaW conflicts due to op ordering.366 if (rRead == rDef)367 continue;368 369 // Find the enclosing repetitive region of READ that is closest to DEF but370 // not the repetitive region of DEF itself.371 while (true) {372 Region *nextRegion = getNextEnclosingRepetitiveRegion(rRead, options);373 if (nextRegion == rDef)374 break;375 assert(nextRegion && "expected to find another repetitive region");376 rRead = nextRegion;377 }378 379 // We cannot use op dominance if WRITE is inside the same repetitive region.380 if (rRead->getParentOp()->isAncestor(uWrite->getOwner()))381 return false;382 }383 384 return true;385}386 387/// Return `true` if op dominance can be used to rule out a read-after-write388/// conflicts based on the ordering of ops. Returns `false` if op dominance389/// cannot be used to due block-based loops within a region.390///391/// Refer to the `canUseOpDominanceDueToRegions` documentation for details on392/// how op domiance is used during RaW conflict detection.393///394/// On a high-level, there is a potential RaW in a program if there exists a395/// possible program execution such that there is a sequence of DEF, followed396/// by WRITE, followed by READ. Each additional DEF resets the sequence.397///398/// Op dominance cannot be used if there is a path from block(READ) to399/// block(WRITE) and a path from block(WRITE) to block(READ). block(DEF) should400/// not appear on that path.401static bool canUseOpDominanceDueToBlocks(OpOperand *uRead, OpOperand *uWrite,402 const SetVector<Value> &definitions,403 AnalysisState &state) {404 // Fast path: If READ and WRITE are in different regions, their block cannot405 // be reachable just via unstructured control flow. (Loops due to regions are406 // covered by `canUseOpDominanceDueToRegions`.)407 if (uRead->getOwner()->getParentRegion() !=408 uWrite->getOwner()->getParentRegion())409 return true;410 411 Block *readBlock = uRead->getOwner()->getBlock();412 Block *writeBlock = uWrite->getOwner()->getBlock();413 for (Value def : definitions) {414 Block *defBlock = def.getParentBlock();415 if (readBlock->isReachable(writeBlock, {defBlock}) &&416 writeBlock->isReachable(readBlock, {defBlock}))417 return false;418 }419 420 return true;421}422 423static bool canUseOpDominance(OpOperand *uRead, OpOperand *uWrite,424 const SetVector<Value> &definitions,425 AnalysisState &state) {426 return canUseOpDominanceDueToRegions(uRead, uWrite, definitions, state) &&427 canUseOpDominanceDueToBlocks(uRead, uWrite, definitions, state);428}429 430/// Annotate IR with details about the detected RaW conflict.431static void annotateConflict(OpOperand *uRead, OpOperand *uConflictingWrite,432 Value definition) {433 static uint64_t counter = 0;434 Operation *readingOp = uRead->getOwner();435 Operation *conflictingWritingOp = uConflictingWrite->getOwner();436 437 OpBuilder b(conflictingWritingOp->getContext());438 std::string id = "C_" + std::to_string(counter++);439 440 std::string conflictingWriteAttr =441 id +442 "[CONFL-WRITE: " + std::to_string(uConflictingWrite->getOperandNumber()) +443 "]";444 conflictingWritingOp->setAttr(conflictingWriteAttr, b.getUnitAttr());445 446 std::string readAttr =447 id + "[READ: " + std::to_string(uRead->getOperandNumber()) + "]";448 readingOp->setAttr(readAttr, b.getUnitAttr());449 450 if (auto opResult = dyn_cast<OpResult>(definition)) {451 std::string defAttr =452 id + "[DEF: result " + std::to_string(opResult.getResultNumber()) + "]";453 opResult.getDefiningOp()->setAttr(defAttr, b.getUnitAttr());454 } else {455 auto bbArg = cast<BlockArgument>(definition);456 std::string defAttr =457 id + "[DEF: bbArg " + std::to_string(bbArg.getArgNumber()) + "]";458 bbArg.getOwner()->getParentOp()->setAttr(defAttr, b.getUnitAttr());459 }460}461 462/// Return 'true' if a tensor that is equivalent to `other` can be found in the463/// reverse use-def chain of `start`. Note: If an OpOperand bufferizes out of464/// place along that use-def chain, the two tensors may not materialize as465/// equivalent buffers (but separate allocations).466///467/// Note: This function also requires that the two tensors have equivalent468/// indexing. I.e., the tensor types do not change along the use-def chain,469/// apart from static <-> dynamic dim casts.470static bool hasEquivalentValueInReverseUseDefChain(AnalysisState &state,471 OpOperand *start,472 Value other) {473 TraversalConfig config;474 config.followEquivalentOnly = true;475 config.alwaysIncludeLeaves = false;476 config.followSameTypeOrCastsOnly = true;477 return !state478 .findValueInReverseUseDefChain(479 start, [&](Value v) { return v == other; }, config)480 .empty();481}482 483/// Return "true" if the given operand's value is originating from a subset484/// that is equivalent to the subset that `subsetOp` inserts into.485static bool matchesInsertDestination(const AnalysisState &state,486 OpOperand *opOperand,487 SubsetInsertionOpInterface subsetOp) {488 auto matchingSubset = [&](Value val) {489 if (auto opResult = dyn_cast<OpResult>(val))490 if (subsetOp.isEquivalentSubset(opResult, [&](Value v1, Value v2) {491 return state.areEquivalentBufferizedValues(v1, v2);492 }))493 return true;494 return false;495 };496 // There may be multiple leaves at which the reverse SSA use-def chain lookup497 // terminates. All of them must be equivalent subsets.498 SetVector<Value> backwardSlice =499 state.findValueInReverseUseDefChain(opOperand, matchingSubset);500 return llvm::all_of(backwardSlice, matchingSubset);501}502 503/// Return "true" if the given "read" and potentially conflicting "write" are504/// not conflicting due to their subset relationship. The comments in this505/// function are expressed in terms of tensor.extract_slice/tensor.insert_slice506/// pairs, but apply to any subset ops that implement the507/// `SubsetInsertionOpInterface`.508static bool areNonConflictingSubsets(OpOperand *uRead,509 OpOperand *uConflictingWrite,510 const AnalysisState &state) {511 Operation *readingOp = uRead->getOwner();512 Operation *conflictingWritingOp = uConflictingWrite->getOwner();513 514 // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If515 // uRead is an InsertSliceOp...516 if (auto subsetOp = dyn_cast<SubsetInsertionOpInterface>(readingOp)) {517 // As an example, consider the following IR.518 //519 // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }520 // %1 = linalg.fill %cst, %0 {inplace= [true] }521 // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]522 // {inplace= [true] }523 524 if (uRead == &subsetOp.getDestinationOperand() &&525 matchesInsertDestination(state, uConflictingWrite, subsetOp))526 // Case 1: The main insight is that InsertSliceOp reads only part of527 // the destination tensor. The overwritten area is not read. If528 // uConflictingWrite writes into exactly the memory location that is529 // being read by uRead, this is not a conflict.530 //531 // In the above example:532 // uRead = OpOperand 1 (%t) of tensor.insert_slice533 // uConflictingWrite = OpOperand 1 (%0) of linalg.fill534 //535 // The read of %t does not conflict with the write of the FillOp536 // (same aliases!) because the area that the FillOp operates on is537 // exactly the one that is *not* read via %t.538 return true;539 540 if (uRead == &subsetOp.getSourceOperand() &&541 uConflictingWrite == &subsetOp.getDestinationOperand() &&542 matchesInsertDestination(state, uRead, subsetOp))543 // Case 2: The read of the source tensor and the write to the dest544 // tensor via an InsertSliceOp is not a conflict if the read is545 // reading exactly that part of an equivalent tensor that the546 // InsertSliceOp is writing.547 //548 // In the above example:549 // uRead = OpOperand 0 (%1) of tensor.insert_slice550 // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice551 return true;552 }553 554 // If uConflictingWrite is an InsertSliceOp...555 if (auto subsetOp =556 dyn_cast<SubsetInsertionOpInterface>(conflictingWritingOp))557 // As an example, consider the following IR.558 //559 // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }560 // %1 = linalg.fill %cst, %0 {inplace= [true] }561 // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]562 // {inplace= [true] }563 // %3 = vector.transfer_read %1, %cst564 //565 // In the above example:566 // uRead = OpOperand 0 (%1) of vector.transfer_read567 // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice568 // definition = %1569 //570 // This is not a conflict because the InsertSliceOp overwrites the571 // memory segment of %1 with the exact same data. (Effectively, there572 // is no memory write here.)573 if (uConflictingWrite == &subsetOp.getDestinationOperand() &&574 state.areEquivalentBufferizedValues(575 uRead->get(), subsetOp.getSourceOperand().get()) &&576 matchesInsertDestination(state, &subsetOp.getSourceOperand(), subsetOp))577 return true;578 579 return false;580}581 582/// Given sets of uses and writes, return true if there is a RaW conflict under583/// the assumption that all given reads/writes alias the same buffer and that584/// all given writes bufferize inplace.585///586/// A conflict is: According to SSA use-def chains, a read R is supposed to read587/// the result of a definition W1. But because of bufferization decisions, R588/// actually reads another definition W2.589static bool590hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,591 const DenseSet<OpOperand *> &usesWrite,592 const DominanceInfo &domInfo,593 OneShotAnalysisState &state) {594 const BufferizationOptions &options = state.getOptions();595 596 // Before going through the main RaW analysis, find cases where a buffer must597 // be privatized due to parallelism. If the result of a write is never read,598 // privatization is not necessary (and large parts of the IR are likely dead).599 if (options.checkParallelRegions && !usesRead.empty()) {600 for (OpOperand *uConflictingWrite : usesWrite) {601 // Find the allocation point or last write (definition) of the buffer.602 // Note: In contrast to `findDefinitions`, this also returns results of603 // ops that do not bufferize to memory write when no other definition604 // could be found. E.g., "bufferization.alloc_tensor" would be included,605 // even though that op just bufferizes to an allocation but does define606 // the contents of the buffer.607 SetVector<Value> definitionsOrLeaves =608 state.findValueInReverseUseDefChain(uConflictingWrite, [&](Value v) {609 return state.bufferizesToMemoryWrite(v);610 });611 assert(!definitionsOrLeaves.empty() &&612 "expected at least one definition or leaf");613 614 // The writing op must bufferize out-of-place if the definition is in a615 // different parallel region than this write.616 for (Value def : definitionsOrLeaves) {617 if (getParallelRegion(def.getParentRegion(), options) !=618 getParallelRegion(uConflictingWrite->getOwner()->getParentRegion(),619 options)) {620 LDBG() << "\n- bufferizes out-of-place due to parallel region:\n"621 << " unConflictingWrite = operand "622 << uConflictingWrite->getOperandNumber() << " of "623 << OpWithFlags(uConflictingWrite->getOwner(),624 OpPrintingFlags().skipRegions());625 return true;626 }627 }628 }629 }630 631 for (OpOperand *uRead : usesRead) {632 Operation *readingOp = uRead->getOwner();633 LDBG() << "\n- check conflict:\n"634 << " uRead = operand " << uRead->getOperandNumber() << " of "635 << OpWithFlags(readingOp, OpPrintingFlags().skipRegions());636 637 // Find the definition of uRead by following the SSA use-def chain.638 // E.g.:639 //640 // %0 = "writing_op"(%t) : tensor<?x32> -> tensor<?xf32>641 // %1 = "aliasing_op"(%0) : tensor<?x32> -> tensor<?xf32>642 // %2 = "reading_op"(%1) : : tensor<?x32> -> not_a_tensor_type643 //644 // In the above example, if uRead is the OpOperand of reading_op, the645 // definition is %0. Note that operations that create an alias but do not646 // bufferize to a memory write (such as ExtractSliceOp) are skipped.647 const SetVector<Value> &definitions = state.findDefinitionsCached(uRead);648 if (definitions.empty()) {649 // Fast path: No conflict if there are no definitions.650 LDBG() << " no conflict: read value has no definitions";651 continue;652 }653 654 // Look for conflicting memory writes. Potential conflicts are writes to an655 // alias that have been decided to bufferize inplace.656 for (OpOperand *uConflictingWrite : usesWrite) {657 LDBG() << " unConflictingWrite = operand "658 << uConflictingWrite->getOperandNumber() << " of "659 << OpWithFlags(uConflictingWrite->getOwner(),660 OpPrintingFlags().skipRegions());661 662 // Check if op dominance can be used to rule out read-after-write663 // conflicts.664 bool useDominance =665 canUseOpDominance(uRead, uConflictingWrite, definitions, state);666 LDBG() << "\n- useDominance = " << useDominance;667 668 // Throughout this loop, check for multiple requirements that have to be669 // met for uConflictingWrite to be an actual conflict.670 Operation *conflictingWritingOp = uConflictingWrite->getOwner();671 672 // Inside of repetitive regions, ops may be executed multiple times and op673 // dominance cannot be used to rule out conflicts.674 if (useDominance) {675 // No conflict if the readingOp dominates conflictingWritingOp, i.e.,676 // the write is not visible when reading.677 //678 // Note: If ops are executed multiple times (e.g., because they are679 // inside a loop), there may be no meaningful `happensBefore`680 // relationship.681 if (happensBefore(readingOp, conflictingWritingOp, domInfo)) {682 LDBG() << " no conflict: read happens before write";683 continue;684 }685 686 // No conflict if the reading use equals the use of the conflicting687 // write. A use cannot conflict with itself.688 //689 // Note: Just being the same op is not enough. It has to be the same690 // use.691 // Note: If the op is executed multiple times (e.g., because it is692 // inside a loop), it may be conflicting with itself.693 if (uConflictingWrite == uRead) {694 LDBG() << " no conflict: read and write are same use";695 continue;696 }697 698 // Ops are not conflicting if they are in mutually exclusive regions.699 //700 // Note: If ops are executed multiple times (e.g., because they are701 // inside a loop), mutually exclusive regions may be executed702 // multiple times.703 if (state.insideMutuallyExclusiveRegions(readingOp,704 conflictingWritingOp)) {705 LDBG() << " no conflict: read and write are in "706 "mutually exclusive regions";707 continue;708 }709 710 // Two equivalent operands of the same op are not conflicting if the op711 // bufferizes to element-wise access. I.e., all loads at a position712 // happen before all stores to the same position.713 if (conflictingWritingOp == readingOp) {714 if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp)) {715 if (bufferizableOp.bufferizesToElementwiseAccess(716 state, {uRead, uConflictingWrite})) {717 if (hasEquivalentValueInReverseUseDefChain(718 state, uRead, uConflictingWrite->get()) ||719 hasEquivalentValueInReverseUseDefChain(720 state, uConflictingWrite, uRead->get())) {721 LDBG() << " no conflict: op bufferizes to element-wise access";722 continue;723 }724 }725 }726 }727 }728 729 // No conflict if the operands are non-conflicting subsets.730 if (areNonConflictingSubsets(uRead, uConflictingWrite, state)) {731 LDBG() << " no conflict: non-conflicting subsets";732 continue;733 }734 735 // No conflict if the op interface says so.736 if (auto bufferizableOp = options.dynCastBufferizableOp(readingOp)) {737 if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite, state)) {738 LDBG() << " no conflict: op interace of reading op says 'no'";739 continue;740 }741 }742 743 if (conflictingWritingOp != readingOp) {744 if (auto bufferizableOp =745 options.dynCastBufferizableOp(conflictingWritingOp)) {746 if (bufferizableOp.isNotConflicting(uRead, uConflictingWrite,747 state)) {748 LDBG() << " no conflict: op interace of writing op says 'no'";749 continue;750 }751 }752 }753 754 // Check all possible definitions.755 for (Value definition : definitions) {756 LDBG() << " * definition = " << definition;757 758 // No conflict if the conflicting write happens before the definition.759 if (Operation *defOp = definition.getDefiningOp()) {760 if (happensBefore(conflictingWritingOp, defOp, domInfo)) {761 // conflictingWritingOp happens before defOp. No conflict.762 LDBG() << " no conflict: write happens before definition";763 continue;764 }765 // No conflict if conflictingWritingOp is contained in defOp.766 if (defOp->isProperAncestor(conflictingWritingOp)) {767 LDBG() << " no conflict: write is contained in definition";768 continue;769 }770 } else {771 auto bbArg = cast<BlockArgument>(definition);772 Block *block = bbArg.getOwner();773 if (!block->findAncestorOpInBlock(*conflictingWritingOp)) {774 LDBG() << " no conflict: definition is bbArg "775 "and write happens outside of block";776 // conflictingWritingOp happens outside of the block. No777 // conflict.778 continue;779 }780 }781 782 // No conflict if the conflicting write and the definition are the same783 // use.784 AliasingValueList aliases = state.getAliasingValues(*uConflictingWrite);785 if (aliases.getNumAliases() == 1 &&786 aliases.getAliases()[0].value == definition) {787 LDBG() << " no conflict: definition and write are same";788 continue;789 }790 791 // All requirements are met. Conflict found!792 793 if (options.printConflicts)794 annotateConflict(uRead, uConflictingWrite, definition);795 LDBG() << " => RaW CONFLICT FOUND";796 return true;797 }798 }799 }800 801 return false;802}803 804// Helper function to iterate on aliases of `root` and capture the writes.805static void getAliasingInplaceWrites(DenseSet<OpOperand *> &res, Value root,806 const OneShotAnalysisState &state) {807 state.applyOnAliases(root, [&](Value alias) {808 for (auto &use : alias.getUses())809 // Inplace write to a value that aliases root.810 if (isInplaceMemoryWrite(use, state))811 res.insert(&use);812 });813}814 815// Helper function to iterate on aliases of `root` and capture the reads.816static void getAliasingReads(DenseSet<OpOperand *> &res, Value root,817 const OneShotAnalysisState &state) {818 state.applyOnAliases(root, [&](Value alias) {819 for (auto &use : alias.getUses()) {820 // Read of a value that aliases root.821 if (state.bufferizesToMemoryRead(use)) {822 res.insert(&use);823 continue;824 }825 826 // Read of a dependent value in the SSA use-def chain. E.g.:827 //828 // %0 = ...829 // %1 = tensor.extract_slice %0 {not_analyzed_yet}830 // "read"(%1)831 //832 // In the above example, getAliasingReads(%0) includes the first OpOperand833 // of the tensor.extract_slice op. The extract_slice itself does not read834 // but its aliasing result is eventually fed into an op that does.835 //836 // Note: This is considered a "read" only if the use does not bufferize to837 // a memory write. (We already ruled out memory reads. In case of a memory838 // write, the buffer would be entirely overwritten; in the above example839 // there would then be no flow of data from the extract_slice operand to840 // its result's uses.)841 if (!state.bufferizesToMemoryWrite(use)) {842 AliasingValueList aliases = state.getAliasingValues(use);843 if (llvm::any_of(aliases, [&](AliasingValue a) {844 return state.isValueRead(a.value);845 }))846 res.insert(&use);847 }848 }849 });850}851 852/// Return true if bufferizing `operand` inplace would create a conflict. A read853/// R and a write W of the same alias set is a conflict if inplace bufferization854/// of W changes the value read by R to a value different from the one that855/// would be expected by tracing back R's origin through SSA use-def chains.856/// A conflict can only be introduced by a new alias and/or an inplace857/// bufferization decision.858///859/// Example:860/// %0 = tensor.extract_slice %t[...][...][1, 1] {inplace?}861/// %1 = vector.transfer_write %v1, %t {inplace} : vector<5xf32>, tensor<?xf32>862/// %e = tensor.extract_slice %1863/// %2 = vector.transfer_write %v2, %0 {inplace} : vector<6xf32>, tensor<?xf32>864/// %3 = vector.transfer_read %e, %cst : tensor<?xf32>, vector<7xf32>865///866/// In the above example, the two TransferWriteOps have already been decided to867/// bufferize inplace. Bufferizing the ExtractSliceOp inplace would create a868/// conflict because:869/// * According to SSA use-def chains, we expect to read the result of %1.870/// * However, adding an alias {%0, %t} would mean that the second871/// TransferWriteOp overwrites the result of the first one. Therefore, the872/// TransferReadOp would no longer be reading the result of %1.873///874/// If `checkConsistencyOnly` is true, this function checks if there is a875/// read-after-write conflict without bufferizing `operand` inplace. This would876/// indicate a problem with the current inplace bufferization decisions.877///878/// Note: If `checkConsistencyOnly`, this function may be called with a null879/// OpResult. In that case, only the consistency of bufferization decisions880/// involving aliases of the given OpOperand are checked.881static bool wouldCreateReadAfterWriteInterference(882 OpOperand &operand, const DominanceInfo &domInfo,883 OneShotAnalysisState &state, bool checkConsistencyOnly = false) {884 // Collect reads and writes of all aliases of OpOperand and OpResult.885 DenseSet<OpOperand *> usesRead, usesWrite;886 getAliasingReads(usesRead, operand.get(), state);887 getAliasingInplaceWrites(usesWrite, operand.get(), state);888 for (AliasingValue alias : state.getAliasingValues(operand)) {889 getAliasingReads(usesRead, alias.value, state);890 getAliasingInplaceWrites(usesWrite, alias.value, state);891 }892 if (!checkConsistencyOnly && state.bufferizesToMemoryWrite(operand))893 usesWrite.insert(&operand);894 895 return hasReadAfterWriteInterference(usesRead, usesWrite, domInfo, state);896}897 898/// Annotate IR with details about the detected non-writability conflict.899static void annotateNonWritableTensor(Value value) {900 static int64_t counter = 0;901 OpBuilder b(value.getContext());902 std::string id = "W_" + std::to_string(counter++);903 if (auto opResult = dyn_cast<OpResult>(value)) {904 std::string attr = id + "[NOT-WRITABLE: result " +905 std::to_string(opResult.getResultNumber()) + "]";906 opResult.getDefiningOp()->setAttr(attr, b.getUnitAttr());907 } else {908 auto bbArg = cast<BlockArgument>(value);909 std::string attr = id + "[NOT-WRITABLE: bbArg " +910 std::to_string(bbArg.getArgNumber()) + "]";911 bbArg.getOwner()->getParentOp()->setAttr(attr, b.getUnitAttr());912 }913}914 915/// Return true if bufferizing `operand` inplace would create a write to a916/// non-writable buffer.917static bool918wouldCreateWriteToNonWritableBuffer(OpOperand &operand,919 OneShotAnalysisState &state,920 bool checkConsistencyOnly = false) {921 bool foundWrite =922 !checkConsistencyOnly && state.bufferizesToMemoryWrite(operand);923 924 if (!foundWrite) {925 // Collect writes of all aliases of OpOperand and OpResult.926 DenseSet<OpOperand *> usesWrite;927 getAliasingInplaceWrites(usesWrite, operand.get(), state);928 for (AliasingValue alias : state.getAliasingValues(operand))929 getAliasingInplaceWrites(usesWrite, alias.value, state);930 foundWrite = !usesWrite.empty();931 }932 933 if (!foundWrite)934 return false;935 936 // Look for a read-only tensor among all aliases.937 bool foundReadOnly = false;938 auto checkReadOnly = [&](Value v) {939 if (!state.isWritable(v)) {940 foundReadOnly = true;941 if (state.getOptions().printConflicts)942 annotateNonWritableTensor(v);943 }944 };945 state.applyOnAliases(operand.get(), checkReadOnly);946 for (AliasingValue alias : state.getAliasingValues(operand))947 state.applyOnAliases(alias.value, checkReadOnly);948 if (foundReadOnly) {949 LDBG() << "=> NOT WRITABLE";950 return true;951 }952 953 return false;954}955 956//===----------------------------------------------------------------------===//957// Bufferization analyses.958//===----------------------------------------------------------------------===//959 960// Find the values that define the contents of the given operand's value.961const llvm::SetVector<Value> &962OneShotAnalysisState::findDefinitionsCached(OpOperand *opOperand) {963 Value value = opOperand->get();964 if (!cachedDefinitions.count(value))965 cachedDefinitions[value] = findDefinitions(opOperand);966 return cachedDefinitions[value];967}968 969void OneShotAnalysisState::resetCache() {970 AnalysisState::resetCache();971 cachedDefinitions.clear();972}973 974/// Determine if `operand` can be bufferized in-place.975static LogicalResult976bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,977 const DominanceInfo &domInfo) {978 LDBG() << "//===-------------------------------------------===//\n"979 << "Analyzing operand #" << operand.getOperandNumber() << " of "980 << OpWithFlags(operand.getOwner(), OpPrintingFlags().skipRegions());981 982 bool foundInterference =983 wouldCreateWriteToNonWritableBuffer(operand, state) ||984 wouldCreateReadAfterWriteInterference(operand, domInfo, state);985 986 if (foundInterference)987 state.bufferizeOutOfPlace(operand);988 else989 state.bufferizeInPlace(operand);990 991 LDBG() << "//===-------------------------------------------===//";992 return success();993}994 995LogicalResult996OneShotAnalysisState::analyzeSingleOp(Operation *op,997 const DominanceInfo &domInfo) {998 for (OpOperand &opOperand : op->getOpOperands())999 if (isa<TensorType>(opOperand.get().getType()))1000 if (failed(bufferizableInPlaceAnalysisImpl(opOperand, *this, domInfo)))1001 return failure();1002 return success();1003}1004 1005/// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.1006static void equivalenceAnalysis(SmallVector<Operation *> &ops,1007 OneShotAnalysisState &state) {1008 for (Operation *op : ops) {1009 if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op)) {1010 for (OpResult opResult : op->getOpResults()) {1011 if (!isa<TensorType>(opResult.getType()))1012 continue;1013 AliasingOpOperandList aliases = state.getAliasingOpOperands(opResult);1014 if (aliases.getNumAliases() == 0)1015 // Nothing to do if there are no aliasing OpOperands.1016 continue;1017 1018 Value firstOperand = aliases.begin()->opOperand->get();1019 bool allEquivalent = true;1020 for (AliasingOpOperand alias : aliases) {1021 bool isEquiv = alias.relation == BufferRelation::Equivalent;1022 bool isInPlace = state.isInPlace(*alias.opOperand);1023 Value operand = alias.opOperand->get();1024 if (isEquiv && isInPlace && alias.isDefinite) {1025 // Found a definite, equivalent alias. Merge equivalence sets.1026 // There can only be one definite alias, so we can stop here.1027 state.unionEquivalenceClasses(opResult, operand);1028 allEquivalent = false;1029 break;1030 }1031 if (!isEquiv || !isInPlace)1032 allEquivalent = false;1033 if (!state.areEquivalentBufferizedValues(operand, firstOperand))1034 allEquivalent = false;1035 }1036 1037 // If all "maybe" aliases are equivalent and the OpResult is not a new1038 // allocation, it is a definite, equivalent alias. E.g.:1039 //1040 // aliasingOpOperands(%r) = {(%t0, EQUIV, MAYBE), (%t1, EQUIV, MAYBE)}1041 // aliasingValues(%t0) = {(%r, EQUIV, MAYBE)}1042 // aliasingValues(%t1) = {(%r, EQUIV, MAYBE)}1043 // %r = arith.select %c, %t0, %t1 : tensor<?xf32>1044 //1045 // If %t0 and %t1 are equivalent, it is safe to union the equivalence1046 // classes of %r, %t0 and %t1.1047 if (allEquivalent && !bufferizableOp.bufferizesToAllocation(opResult))1048 state.unionEquivalenceClasses(opResult, firstOperand);1049 }1050 }1051 }1052}1053 1054/// Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained1055/// in `op`.1056static void equivalenceAnalysis(Operation *op, OneShotAnalysisState &state) {1057 // Traverse ops in PostOrder: Nested ops first, then enclosing ops.1058 SmallVector<Operation *> ops;1059 op->walk<WalkOrder::PostOrder>([&](Operation *op) {1060 // No tensors => no buffers.1061 if (none_of(op->getResultTypes(), isaTensor))1062 return;1063 ops.push_back(op);1064 });1065 1066 equivalenceAnalysis(ops, state);1067}1068 1069/// "Bottom-up from terminators" heuristic.1070static SmallVector<Operation *>1071bottomUpFromTerminatorsHeuristic(Operation *op,1072 const OneShotAnalysisState &state) {1073 SetVector<Operation *> traversedOps;1074 1075 // Find region terminators.1076 op->walk<WalkOrder::PostOrder>([&](RegionBranchTerminatorOpInterface term) {1077 if (!traversedOps.insert(term))1078 return;1079 // Follow the reverse SSA use-def chain from each yielded value as long as1080 // we stay within the same region.1081 SmallVector<OpResult> worklist;1082 for (Value v : term->getOperands()) {1083 if (!isa<TensorType>(v.getType()))1084 continue;1085 auto opResult = dyn_cast<OpResult>(v);1086 if (!opResult)1087 continue;1088 worklist.push_back(opResult);1089 }1090 while (!worklist.empty()) {1091 OpResult opResult = worklist.pop_back_val();1092 Operation *defOp = opResult.getDefiningOp();1093 if (!traversedOps.insert(defOp))1094 continue;1095 if (!term->getParentRegion()->findAncestorOpInRegion(*defOp))1096 continue;1097 AliasingOpOperandList aliases = state.getAliasingOpOperands(opResult);1098 for (auto alias : aliases) {1099 Value v = alias.opOperand->get();1100 if (!isa<TensorType>(v.getType()))1101 continue;1102 auto opResult = dyn_cast<OpResult>(v);1103 if (!opResult)1104 continue;1105 worklist.push_back(opResult);1106 }1107 }1108 });1109 1110 // Analyze traversed ops, then all remaining ops.1111 SmallVector<Operation *> result(traversedOps.begin(), traversedOps.end());1112 op->walk<WalkOrder::PostOrder, ReverseIterator>([&](Operation *op) {1113 if (!traversedOps.contains(op) && hasTensorSemantics(op))1114 result.push_back(op);1115 });1116 return result;1117}1118 1119LogicalResult OneShotAnalysisState::analyzeOp(Operation *op,1120 const DominanceInfo &domInfo) {1121 OneShotBufferizationOptions::AnalysisHeuristic heuristic =1122 getOptions().analysisHeuristic;1123 1124 SmallVector<Operation *> orderedOps;1125 if (heuristic ==1126 OneShotBufferizationOptions::AnalysisHeuristic::BottomUpFromTerminators) {1127 orderedOps = bottomUpFromTerminatorsHeuristic(op, *this);1128 } else {1129 op->walk([&](Operation *op) {1130 // No tensors => no buffers.1131 if (!hasTensorSemantics(op))1132 return;1133 orderedOps.push_back(op);1134 });1135 switch (heuristic) {1136 case OneShotBufferizationOptions::AnalysisHeuristic::BottomUp: {1137 // Default: Walk ops in reverse for better interference analysis.1138 std::reverse(orderedOps.begin(), orderedOps.end());1139 break;1140 }1141 case OneShotBufferizationOptions::AnalysisHeuristic::TopDown: {1142 // Ops are already sorted top-down in `orderedOps`.1143 break;1144 }1145 case OneShotBufferizationOptions::AnalysisHeuristic::Fuzzer: {1146 assert(getOptions().analysisFuzzerSeed &&1147 "expected that fuzzer seed it set");1148 // This is a fuzzer. For testing purposes only. Randomize the order in1149 // which operations are analyzed. The bufferization quality is likely1150 // worse, but we want to make sure that no assertions are triggered1151 // anywhere.1152 std::mt19937 g(getOptions().analysisFuzzerSeed);1153 llvm::shuffle(orderedOps.begin(), orderedOps.end(), g);1154 break;1155 }1156 default: {1157 llvm_unreachable("unsupported heuristic");1158 }1159 }1160 }1161 1162 // Analyze ops in the computed order.1163 for (Operation *op : orderedOps)1164 if (failed(analyzeSingleOp(op, domInfo)))1165 return failure();1166 1167 equivalenceAnalysis(op, *this);1168 return success();1169}1170 1171/// Perform various checks on the input IR to see if it contains IR constructs1172/// that are unsupported by One-Shot Bufferize.1173static LogicalResult1174checkPreBufferizationAssumptions(Operation *op, const DominanceInfo &domInfo,1175 OneShotAnalysisState &state) {1176 const BufferizationOptions &options = state.getOptions();1177 1178 // Note: This walk cannot be combined with the one below because interface1179 // methods of invalid/unsupported ops may be called during the second walk.1180 // (On ops different from `op`.)1181 WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {1182 // Skip ops that are not in the filter.1183 if (!options.isOpAllowed(op.getOperation()))1184 return WalkResult::advance();1185 1186 // Check for unsupported unstructured control flow.1187 if (!op.supportsUnstructuredControlFlow()) {1188 for (Region &r : op->getRegions()) {1189 if (r.getBlocks().size() > 1) {1190 op->emitOpError("op or BufferizableOpInterface implementation does "1191 "not support unstructured control flow, but at least "1192 "one region has multiple blocks");1193 return WalkResult::interrupt();1194 }1195 }1196 }1197 1198 return WalkResult::advance();1199 });1200 if (walkResult.wasInterrupted())1201 return failure();1202 1203 walkResult = op->walk([&](BufferizableOpInterface op) {1204 // Skip ops that are not in the filter.1205 if (!options.isOpAllowed(op.getOperation()))1206 return WalkResult::advance();1207 1208 // Input IR may not contain any ToTensorOps without the "restrict"1209 // attribute. Such tensors may alias any other tensor, which is currently1210 // not handled in the analysis.1211 if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {1212 if (!toTensorOp.getRestrict() && !toTensorOp->getUses().empty()) {1213 op->emitOpError("to_tensor ops without `restrict` are not supported by "1214 "One-Shot Analysis");1215 return WalkResult::interrupt();1216 }1217 }1218 1219 for (OpOperand &opOperand : op->getOpOperands()) {1220 if (isa<TensorType>(opOperand.get().getType())) {1221 if (wouldCreateReadAfterWriteInterference(1222 opOperand, domInfo, state,1223 /*checkConsistencyOnly=*/true)) {1224 // This error can happen if certain "mustBufferizeInPlace" interface1225 // methods are implemented incorrectly, such that the IR already has1226 // a RaW conflict before making any bufferization decisions. It can1227 // also happen if the bufferization.materialize_in_destination is used1228 // in such a way that a RaW conflict is not avoidable.1229 op->emitOpError("not bufferizable under the given constraints: "1230 "cannot avoid RaW conflict");1231 return WalkResult::interrupt();1232 }1233 1234 if (state.isInPlace(opOperand) &&1235 wouldCreateWriteToNonWritableBuffer(1236 opOperand, state, /*checkConsistencyOnly=*/true)) {1237 op->emitOpError("not bufferizable under the given constraints: would "1238 "write to read-only buffer");1239 return WalkResult::interrupt();1240 }1241 }1242 }1243 1244 return WalkResult::advance();1245 });1246 1247 return success(!walkResult.wasInterrupted());1248}1249 1250/// Annotate the IR with the result of the analysis. For testing/debugging only.1251static void1252annotateOpsWithBufferizationMarkers(Operation *op,1253 const OneShotAnalysisState &state) {1254 // Add __inplace_operands_attr__.1255 op->walk([&](Operation *op) {1256 for (OpOperand &opOperand : op->getOpOperands())1257 if (isa<TensorType>(opOperand.get().getType()))1258 setInPlaceOpOperand(opOperand, state.isInPlace(opOperand));1259 });1260}1261 1262static void annotateOpsWithAliasSets(Operation *op,1263 const OneShotAnalysisState &state) {1264 AsmState asmState(op);1265 Builder b(op->getContext());1266 // Helper function to build an array attribute of aliasing SSA value strings.1267 auto buildAliasesArray = [&](Value v) {1268 SmallVector<Attribute> aliases;1269 state.applyOnAliases(v, [&](Value alias) {1270 std::string buffer;1271 llvm::raw_string_ostream stream(buffer);1272 alias.printAsOperand(stream, asmState);1273 aliases.push_back(b.getStringAttr(buffer));1274 });1275 return b.getArrayAttr(aliases);1276 };1277 1278 op->walk([&](Operation *op) {1279 // Build alias set array for every OpResult.1280 SmallVector<Attribute> opResultAliasSets;1281 for (OpResult opResult : op->getOpResults()) {1282 if (llvm::isa<TensorType>(opResult.getType())) {1283 opResultAliasSets.push_back(buildAliasesArray(opResult));1284 }1285 }1286 if (!opResultAliasSets.empty())1287 op->setAttr(kOpResultAliasSetAttrName, b.getArrayAttr(opResultAliasSets));1288 1289 // Build alias set array for every BlockArgument.1290 SmallVector<Attribute> regionAliasSets;1291 bool hasTensorBbArg = false;1292 for (Region &r : op->getRegions()) {1293 SmallVector<Attribute> blockAliasSets;1294 for (Block &block : r.getBlocks()) {1295 SmallVector<Attribute> bbArgAliasSets;1296 for (BlockArgument bbArg : block.getArguments()) {1297 if (llvm::isa<TensorType>(bbArg.getType())) {1298 bbArgAliasSets.push_back(buildAliasesArray(bbArg));1299 hasTensorBbArg = true;1300 }1301 }1302 blockAliasSets.push_back(b.getArrayAttr(bbArgAliasSets));1303 }1304 regionAliasSets.push_back(b.getArrayAttr(blockAliasSets));1305 }1306 if (hasTensorBbArg)1307 op->setAttr(kBbArgAliasSetAttrName, b.getArrayAttr(regionAliasSets));1308 });1309}1310 1311LogicalResult bufferization::analyzeOp(Operation *op,1312 OneShotAnalysisState &state,1313 BufferizationStatistics *statistics) {1314 DominanceInfo domInfo(op);1315 const OneShotBufferizationOptions &options = state.getOptions();1316 1317 if (failed(checkPreBufferizationAssumptions(op, domInfo, state)))1318 return failure();1319 1320 // If the analysis fails, just return.1321 if (failed(state.analyzeOp(op, domInfo)))1322 return failure();1323 1324 if (statistics) {1325 statistics->numTensorInPlace = state.getStatNumTensorInPlace();1326 statistics->numTensorOutOfPlace = state.getStatNumTensorOutOfPlace();1327 }1328 1329 bool failedAnalysis = false;1330 1331 // Gather some extra analysis data.1332 state.gatherUndefinedTensorUses(op);1333 1334 // Analysis verification: After setting up alias/equivalence sets, each op1335 // can check for expected invariants/limitations and fail the analysis if1336 // necessary.1337 op->walk([&](Operation *op) {1338 if (BufferizableOpInterface bufferizableOp =1339 options.dynCastBufferizableOp(op))1340 failedAnalysis |= failed(bufferizableOp.verifyAnalysis(state));1341 });1342 1343 // Annotate operations if we only want to report the analysis.1344 if (options.testAnalysisOnly)1345 annotateOpsWithBufferizationMarkers(op, state);1346 if (options.dumpAliasSets)1347 annotateOpsWithAliasSets(op, state);1348 1349 return success(!failedAnalysis);1350}1351 1352LogicalResult bufferization::runOneShotBufferize(1353 Operation *op, const OneShotBufferizationOptions &options,1354 BufferizationState &state, BufferizationStatistics *statistics) {1355 // copy-before-write deactivates the analysis. It cannot be used together with1356 // test-analysis-only.1357 assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&1358 "invalid combination of bufferization flags");1359 1360 if (options.copyBeforeWrite) {1361 // Copy buffer before each write. No analysis is needed.1362 } else {1363 // Run One-Shot Analysis and insert buffer copies (on the tensor level)1364 // only where needed. This is the default and much more efficient than1365 // copy-before-write.1366 if (failed(insertTensorCopies(op, options, state, statistics)))1367 return failure();1368 1369 // If test-analysis-only is set, the IR was annotated with RaW conflict1370 // markers (attributes) during One-Shot Analysis.1371 if (options.testAnalysisOnly)1372 return success();1373 }1374 1375 // Bufferize the op and its nested ops. If options.copyBeforeWrite is set,1376 // a new buffer copy is allocated every time a buffer is written to.1377 return bufferizeOp(op, options, state, statistics);1378}1379