diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td --- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td +++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.td @@ -182,17 +182,26 @@ >, InterfaceMethod< /*desc=*/[{ - Return `true` if the given OpOperand can be written to in-place. This - is the case for most ops, but some ops such as ConstantOp may - bufferize to non-writable (read-only) memory locations. This method - will never be called on OpResults that do not have a tensor type. + Return `true` if the given Value can be written to in-place. Value is + either an OpResult of this operation or a BlockArgument of a block of + this operation. + + Most OpResult buffers can be written to, but some ops such as + ConstantOp may bufferize to non-writable (read-only) memory locations. + Therefore, by default, this method returns `true` for OpResults. This + method will never be called on OpResults that do not have a tensor + type. + + Whether a BlockArgument can be written to or not depends on the + operation. This method conservatively returns `false`. This method + will never be called on BlockArguments that do not have a tensor type. }], /*retType=*/"bool", /*methodName=*/"isWritable", - /*args=*/(ins "OpResult":$opResult), + /*args=*/(ins "Value":$value), /*methodBody=*/"", /*defaultImplementation=*/[{ - return true; + return value.isa(); }] > ]; diff --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h --- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h +++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h @@ -91,6 +91,12 @@ /// Specify that the value is known to bufferize to writable memory. void setBufferizesToWritableMemory(Value v); + /// Mark a value as in-place bufferized. + void markInPlace(OpResult v) { inplaceBufferized.insert(v); } + + /// Return `true` if a value was marked as in-place bufferized. + bool isInPlace(OpResult opResult) const; + /// Print to `os`. void printAliases(raw_ostream &os) const; void printEquivalences(raw_ostream &os) const; @@ -117,6 +123,9 @@ /// Set of tensors that are known to bufferize to writable memory. llvm::DenseSet bufferizeToWritableMemory; + /// Set of all OpResults that were decided to bufferize in-place. + llvm::DenseSet inplaceBufferized; + /// Auxiliary structure to store all the values a given value may alias with. /// Alias information is "may be" conservative: In the presence of branches, a /// value may alias with one of multiple other values. The concrete aliasing diff --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp @@ -221,48 +221,20 @@ //===----------------------------------------------------------------------===// // Bufferization-specific attribute manipulation. -// These could be simplified with helper structs on the side, for now attributes -// allow simple embedding in the IR which simplifies testing. -// This could also be folded in BufferizationAliasInfo or a Bufferizer class -// that uses BufferizationAliasInfo. +// These are for testing and debugging only. Bufferization information is +// stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR +// is annotated with the results of the analysis (copied from +// BufferizationAliasInfo), so that they can be checked in tests. //===----------------------------------------------------------------------===// /// Attribute marker to specify op results that can be bufferized inPlace. constexpr StringLiteral kInPlaceResultsAttrName = "__inplace_results_attr__"; -// TODO: proper enum. -enum class InPlaceSpec { - False, - True, - None, -}; - -static StringRef stringify(InPlaceSpec val) { - switch (val) { - case InPlaceSpec::False: - return "false"; - case InPlaceSpec::True: - return "true"; - case InPlaceSpec::None: - return "none"; - } - return ""; -} - -static Optional symbolize(StringRef str) { - return StringSwitch>(str) - .Case("false", InPlaceSpec::False) - .Case("true", InPlaceSpec::True) - .Case("none", InPlaceSpec::None) - .Default(None); -} - /// Mark whether OpResult can actually be bufferized inplace. -/// If `inPlace` is `InPlaceSpec::True`, the use-def chain analysis has -/// guaranteed that no subsequent write would occur to the bufferized -/// tensor value (i.e. the result can be bufferized inPlace). -static void setInPlaceOpResult(OpResult opResult, - InPlaceSpec inPlace = InPlaceSpec::True) { +/// If `inPlace` is `true`, the use-def chain analysis has guaranteed that no +/// subsequent write would occur to the bufferized tensor value (i.e. the result +/// can be bufferized inplace). +static void setInPlaceOpResult(OpResult opResult, bool inPlace) { if (!opResult) return; @@ -272,69 +244,20 @@ SmallVector inPlaceVector = attr ? SmallVector( llvm::to_vector<4>(attr.getAsValueRange())) - : SmallVector(op->getNumResults(), - stringify(InPlaceSpec::None)); - LDBG("->set inPlace=" << stringify(inPlace) << " <- #" - << opResult.getResultNumber() << ": " - << printOperationInfo(op) << "\n"); - inPlaceVector[opResult.getResultNumber()] = stringify(inPlace); + : SmallVector(op->getNumResults(), "false"); + LDBG("->set inPlace=" << inPlace << " <- #" << opResult.getResultNumber() + << ": " << printOperationInfo(op) << "\n"); + inPlaceVector[opResult.getResultNumber()] = inPlace ? "true" : "false"; op->setAttr(kInPlaceResultsAttrName, OpBuilder(op).getStrArrayAttr(inPlaceVector)); } -/// Get the InPlaceSpec attribute entry `kInPlaceResultsAttrName` for -/// `opResult`. If the result is `InPlaceSpec::True`, the use-def chain analysis -/// has guaranteed that no subsequent read of the tensor value occurs and the -/// result can be buferized inPlace. -/// If no InPlaceSpec attribute has been set for `opResult`, return -/// InPlaceSpec::None. -LLVM_ATTRIBUTE_UNUSED static InPlaceSpec getInPlace(OpResult opResult) { - if (!opResult) - return InPlaceSpec::None; - - Operation *op = opResult.getOwner(); - auto attr = - op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null(); - if (!attr) - return InPlaceSpec::None; - - // Must return a proper value. - return *symbolize(*(attr.getAsValueRange().begin() + - opResult.getResultNumber())); -} - -/// Get inPlace information for `bbArg`. -/// FuncOp allow argument attributes, we use those to encode the information. -/// BlockArgument of other ops delegate to their owner's parent op. -static InPlaceSpec getInPlace(BlockArgument bbArg) { - if (auto funcOp = dyn_cast(bbArg.getOwner()->getParentOp())) { - BoolAttr inplaceAttr = funcOp.getArgAttrOfType( - bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName); - if (!inplaceAttr) - return InPlaceSpec::None; - return inplaceAttr.getValue() ? InPlaceSpec::True : InPlaceSpec::False; - } - // Interestingly, scf::ForOp's and TiledLoop's bbArg can **always** be viewed - // inplace from the perspective of ops nested under: - // 1. Either the matching iter operand is not bufferized inplace and an - // alloc + optional copy makes the bbArg itself inplaceable. - // 2. Or the matching iter operand is bufferized inplace and bbArg just - // bufferizes to that too. - if (isa(bbArg.getOwner()->getParentOp())) - return InPlaceSpec::True; - // Unknown cases. - return InPlaceSpec::None; -} - /// Set the attribute that triggers inplace bufferization on a FuncOp argument /// `bbArg`. -static void -setInPlaceFuncArgument(BlockArgument bbArg, - InPlaceSpec inPlaceSpec = InPlaceSpec::True) { +static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) { auto funcOp = cast(bbArg.getOwner()->getParentOp()); - funcOp.setArgAttr( - bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName, - BoolAttr::get(bbArg.getContext(), inPlaceSpec == InPlaceSpec::True)); + funcOp.setArgAttr(bbArg.getArgNumber(), LinalgDialect::kInplaceableAttrName, + BoolAttr::get(bbArg.getContext(), inPlace)); } /// Remove the attribute that triggers inplace bufferization on a FuncOp @@ -347,12 +270,6 @@ LinalgDialect::kInplaceableAttrName); } -static InPlaceSpec getInPlace(Value v) { - if (auto bbArg = v.dyn_cast()) - return getInPlace(bbArg); - return getInPlace(v.cast()); -} - //===----------------------------------------------------------------------===// // Printing helpers. //===----------------------------------------------------------------------===// @@ -365,9 +282,6 @@ os << prefix; value.printAsOperand(os, state); os << " : " << value.getType(); - if (getInPlace(value) == InPlaceSpec::None) - return; - os << " [InPlace=" << stringify(getInPlace(value)) << "]"; } /// Print the operation name and bufferization information. @@ -428,12 +342,13 @@ } /// Return true if opOperand has been decided to bufferize in-place. -static bool isInplaceMemoryWrite(OpOperand &opOperand) { +static bool isInplaceMemoryWrite(OpOperand &opOperand, + const BufferizationAliasInfo &aliasInfo) { // Ops that do not bufferize to a memory write, cannot be write in-place. if (!bufferizesToMemoryWrite(opOperand)) return false; OpResult opResult = getAliasingOpResult(opOperand); - return opResult && getInPlace(opResult) == InPlaceSpec::True; + return opResult && aliasInfo.isInPlace(opResult); } BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) { @@ -460,7 +375,7 @@ "expected that OpResult has aliasing OpOperand"); for (OpOperand *operand : operands) aliasInfo.unionSets(operand->get(), opResult); - setInPlaceOpResult(opResult, InPlaceSpec::True); + markInPlace(opResult); } } }); @@ -492,39 +407,33 @@ /// is not writable. static bool aliasesNonWritableBuffer(Value value, const BufferizationAliasInfo &aliasInfo) { - LDBG("----Start aliasesNonWritableBuffer\n"); + LDBG("WRITABILITY ANALYSIS FOR " << printValueInfo(value) << "\n"); bool foundNonWritableBuffer = false; aliasInfo.applyOnAliases(value, [&](Value v) { - LDBG("-----------examine: " << printValueInfo(v) << '\n'); - if (aliasInfo.bufferizesToWritableMemory(v)) { - LDBG("-----------Value is known to be writable -> skip: " - << printValueInfo(v) << '\n'); + // Some values are known to be writable. + if (aliasInfo.bufferizesToWritableMemory(v)) return; - } - if (auto bbArg = v.dyn_cast()) { - if (getInPlace(bbArg) == InPlaceSpec::True) { - LDBG("-----------bbArg is writable -> skip: " << printValueInfo(bbArg) - << '\n'); + // Query BufferizableOpInterface to see if the OpResult is writable. + // TODO: Out-of-place bufferized OpResult could be considered writable. + if (auto bufferizableOp = v.getDefiningOp()) + if (bufferizableOp && bufferizableOp.isWritable(v)) return; - } - LDBG("-----------notWritable bbArg\n"); - foundNonWritableBuffer = true; - return; - } - auto bufferizableOp = dyn_cast(v.getDefiningOp()); - if (!bufferizableOp || !bufferizableOp.isWritable(v.cast())) { - // Unknown ops are treated conservatively: Assume that it is illegal to - // write to their OpResults in-place. - LDBG("-----------notWritable op\n"); - foundNonWritableBuffer = true; - return; - } + // Query BufferizableOpInterface to see if the BlockArgument is writable. + if (auto bbArg = v.dyn_cast()) + if (auto bufferizableOp = dyn_cast( + bbArg.getOwner()->getParentOp())) + if (bufferizableOp.isWritable(bbArg)) + return; + + foundNonWritableBuffer = true; }); - if (!foundNonWritableBuffer) - LDBG("---->value is writable\n"); + if (foundNonWritableBuffer) + LDBG("--> NON WRITABLE\n"); + else + LDBG("--> WRITABLE\n"); return foundNonWritableBuffer; } @@ -547,7 +456,7 @@ bool foundInplaceWrite = false; aliasInfo.applyOnAliases(value, [&](Value v) { for (auto &use : v.getUses()) { - if (isInplaceMemoryWrite(use)) { + if (isInplaceMemoryWrite(use, aliasInfo)) { LDBG("-----------wants to bufferize to inPlace write: " << printOperationInfo(use.getOwner()) << '\n'); foundInplaceWrite = true; @@ -562,10 +471,30 @@ return foundInplaceWrite; } +/// Return `true` if a value was marked as in-place bufferized. +bool BufferizationAliasInfo::isInPlace(OpResult opResult) const { + bool inplace = inplaceBufferized.contains(opResult); +#ifndef NDEBUG + if (inplace) { + auto bufferizableOp = + dyn_cast(opResult.getDefiningOp()); + assert(bufferizableOp && + "expected that in-place bufferized op is bufferizable"); + SmallVector operands = + bufferizableOp.getAliasingOpOperand(opResult); + for (OpOperand *operand : operands) + assert(areAliasingBufferizedValues(operand->get(), opResult) && + "expected that in-place bufferized OpResult aliases with " + "aliasing OpOperand"); + } +#endif // NDEBUG + return inplace; +} + /// Set the inPlace bufferization spec to true. void BufferizationAliasInfo::bufferizeInPlace(OpResult result, OpOperand &operand) { - setInPlaceOpResult(result, InPlaceSpec::True); + markInPlace(result); aliasInfo.unionSets(result, operand.get()); // Dump the updated alias analysis. LLVM_DEBUG(dumpAliases()); @@ -577,7 +506,8 @@ /// Set the inPlace bufferization spec to false. void BufferizationAliasInfo::bufferizeOutOfPlace(OpResult result) { - setInPlaceOpResult(result, InPlaceSpec::False); + if (inplaceBufferized.contains(result)) + inplaceBufferized.erase(result); } /// Starting from `value`, follow the use-def chain in reverse, always selecting @@ -862,7 +792,7 @@ aliasInfo.applyOnAliases(root, [&](Value alias) { for (auto &use : alias.getUses()) // Inplace write to a value that aliases root. - if (isInplaceMemoryWrite(use)) + if (isInplaceMemoryWrite(use, aliasInfo)) res.insert(&use); }); }; @@ -1234,8 +1164,7 @@ } // If bufferizing out-of-place, allocate a new buffer. - bool needCopy = getInPlace(result) != InPlaceSpec::True; - if (needCopy) { + if (!aliasInfo.isInPlace(result)) { // Ops with multiple aliasing operands can currently not bufferize // out-of-place. assert( @@ -1434,9 +1363,7 @@ // Bufferization analyses. //===----------------------------------------------------------------------===// -/// Determine if `operand` can be bufferized in-place with `result`. If so, set -/// InPlaceSpec::True on the result. Otherwise, set InPlaceSpec::False on the -/// result. +/// Determine if `operand` can be bufferized in-place with `result`. static LogicalResult bufferizableInPlaceAnalysisImpl(OpOperand &operand, OpResult result, BufferizationAliasInfo &aliasInfo, @@ -1470,8 +1397,7 @@ } /// Determine if `operand` can be bufferized in-place with one of the op's -/// results. If so, set InPlaceSpec::True on the result. Otherwise, set -/// InPlaceSpec::False on the result. +/// results. /// /// Even if an op does not read or write, it may still create an alias when /// bufferized in-place. An example of such ops is tensor.extract_slice. @@ -2014,12 +1940,12 @@ continue; SetVector maybeInitTensor = - findValueInReverseUseDefChain(operand.get(), [](Value val) { + findValueInReverseUseDefChain(operand.get(), [&](Value val) { // Continue traversal until this function returns true. OpResult opResult = val.dyn_cast(); if (!opResult) return true; - if (getInPlace(opResult) != InPlaceSpec::True) + if (!aliasInfo.isInPlace(opResult)) return true; // Only equivalent tensors are supported at the moment. // TODO: Support cases such as extract_slice(init_tensor). @@ -2103,12 +2029,12 @@ DominanceInfo &domInfo) { return initTensorElimination( funcOp, aliasInfo, domInfo, - [](OpOperand &operand) { + [&](OpOperand &operand) { auto insertSliceOp = dyn_cast(operand.getOwner()); if (!insertSliceOp) return false; // Only inplace bufferized InsertSliceOps are eligible. - if (getInPlace(insertSliceOp->getOpResult(0)) != InPlaceSpec::True) + if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0))) return false; return &operand == &insertSliceOp->getOpOperand(0) /*source*/; }, @@ -2121,6 +2047,23 @@ }); } +/// Annotate the IR with the result of the analysis. For testing/debugging only. +static void +annotateOpsWithBufferizationMarkers(Operation *op, + const BufferizationAliasInfo &aliasInfo) { + op->walk([&](Operation *op) { + for (OpResult opResult : op->getResults()) { + if (opResult.getType().isa()) + setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult)); + if (auto funcOp = dyn_cast(op)) + for (BlockArgument bbArg : funcOp.getArguments()) + if (bbArg.getType().isa()) + setInPlaceFuncArgument(bbArg, + aliasInfo.bufferizesToWritableMemory(bbArg)); + } + }); +} + LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize( ModuleOp moduleOp, const BufferizationOptions &options) { SmallVector orderedFuncOps; @@ -2149,7 +2092,7 @@ if (callerMap.find(funcOp) != callerMap.end()) for (BlockArgument bbArg : funcOp.getArguments()) if (bbArg.getType().isa()) - setInPlaceFuncArgument(bbArg); + aliasInfo.setBufferizesToWritableMemory(bbArg); // If the analysis fails, just return. if (failed(inPlaceAnalysisFuncOpBody(funcOp, aliasInfo, domInfo, @@ -2171,9 +2114,11 @@ return failure(); } } - // Don't drop the attributes if we only want to report the analysis. - if (options.testAnalysisOnly) + // Annotate operations if we only want to report the analysis. + if (options.testAnalysisOnly) { + annotateOpsWithBufferizationMarkers(moduleOp, aliasInfo); return success(); + } for (FuncOp funcOp : orderedFuncOps) { // Note: It would be good to apply cleanups here but we cannot as aliasInfo @@ -2196,8 +2141,6 @@ layoutPostProcessing(moduleOp); // Post-pass cleanup of inplaceable and buffer_layout attributes. - moduleOp.walk( - [&](Operation *op) { op->removeAttr(kInPlaceResultsAttrName); }); moduleOp.walk([&](FuncOp op) { for (BlockArgument bbArg : op.getArguments()) removeBufferizationFuncArguments(bbArg); @@ -2255,8 +2198,9 @@ return success(); } - bool isWritable(Operation *op, OpResult opResult) const { + bool isWritable(Operation *op, Value value) const { // Memory locations returned by memref::GetGlobalOp may not be written to. + assert(value.isa()); return false; } }; @@ -2460,6 +2404,16 @@ return BufferRelation::Equivalent; } + bool isWritable(Operation *op, Value value) const { + // Interestingly, linalg::TiledLoopOp's bbArg can **always** be viewed + // inplace from the perspective of ops nested under: + // 1. Either the matching iter operand is not bufferized inplace and an + // alloc + optional copy makes the bbArg itself inplaceable. + // 2. Or the matching iter operand is bufferized inplace and bbArg just + // bufferizes to that too. + return true; + } + LogicalResult bufferize(Operation *op, OpBuilder &b, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo, @@ -2749,6 +2703,16 @@ return BufferRelation::Equivalent; } + bool isWritable(Operation *op, Value value) const { + // Interestingly, scf::ForOp's bbArg can **always** be viewed + // inplace from the perspective of ops nested under: + // 1. Either the matching iter operand is not bufferized inplace and an + // alloc + optional copy makes the bbArg itself inplaceable. + // 2. Or the matching iter operand is bufferized inplace and bbArg just + // bufferizes to that too. + return true; + } + LogicalResult bufferize(Operation *op, OpBuilder &b, BlockAndValueMapping &bvm, BufferizationAliasInfo &aliasInfo, @@ -3106,8 +3070,7 @@ // If not inplaceable, alloc. Value alloc; - auto inPlace = getInPlace(extractSliceOp->getResult(0)); - if (inPlace != InPlaceSpec::True) + if (!aliasInfo.isInPlace(extractSliceOp->getResult(0))) alloc = createNewAllocDeallocPairForShapedValue( b, loc, extractSliceOp.result(), aliasInfo, allocationFn); @@ -3185,7 +3148,7 @@ if (extractSliceOp && areEquivalentExtractSliceOps(aliasInfo, extractSliceOp, insertSliceOp) && - getInPlace(extractSliceOp.result()) == InPlaceSpec::True) { + aliasInfo.isInPlace(extractSliceOp->getResult(0))) { LDBG("\tfound: " << extractSliceOp.getOperation() << '\n'); foundOp = true; } @@ -3264,11 +3227,10 @@ // slice is computed out of place into the inplace full tensor. // - The result is not inplace. This is the case where the whole tensor is // cloned and the clone needs to be updated. - auto inPlace = getInPlace(insertSliceOp->getResult(0)); // TODO: Is this necessary? if (!isSourceEquivalentToAMatchingInplaceExtractSliceOp(aliasInfo, insertSliceOp) || - inPlace != InPlaceSpec::True) { + !aliasInfo.isInPlace(insertSliceOp->getResult(0))) { LDBG("insert_slice needs extra source copy: " << insertSliceOp.source() << " -> copy\n"); // Take a subview of the dst. diff --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir --- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir +++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize-analysis.mlir @@ -978,7 +978,7 @@ %v2 = vector.transfer_read %e[%s], %cst : tensor, vector<5xf32> scf.yield %e, %v2 : tensor, vector<5xf32> } - // CHECK: __inplace_results_attr__ = ["true", "none"] + // CHECK: __inplace_results_attr__ = ["true", "false"] // Use %t3 in some way without reading it, so that it does not get DCE'd. // CHECK: linalg.generic