diff --git a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp @@ -522,11 +522,6 @@ Operation *opToBufferize, const DominanceInfo &domInfo) const; - /// Return true if we find any read to opOperand.get() or any of its aliases, - /// that does not dominate opOperand.getOwner(). - bool existsNonDominatingRead(OpOperand &opOperand, - const DominanceInfo &domInfo) const; - /// Return true if `v1` and `v2` bufferize to equivalent buffers. bool areEquivalentBufferizedValues(Value v1, Value v2) const { return equivalentInfo.getLeaderValue(v1) == @@ -732,6 +727,10 @@ setInPlaceOpResult(result, InPlaceSpec::False); } +static bool isTheSSADef(OpOperand &write, OpOperand &read) { + return getInplaceableOpResult(write) == read.get(); +} + /// Return true if merging the alias sets of `rootWrite` and `rootRead` would /// result in a semantic change in the program (i.e. RAW violation). /// @@ -747,15 +746,22 @@ Value rootWrite, Value rootRead, Operation *opToBufferize, const DominanceInfo &domInfo) const { LDBG("----Start wouldCreateReadAfterWriteInterference\n"); + LDBG("--------rootWrite: " << rootWrite << "\n"); + LDBG("-------- rootRead: " << rootRead << "\n"); - // Collect all the inplace write uses of some alias of `rootWrite`. + // Collect: + // 1. all the inplace write uses of some alias of `rootWrite`. + // 2. all the write uses that belong to `opToBufferize`. DenseSet usesWrite; auto &aliasListWrite = getAliasInfoRef(rootWrite); for (Value vWrite : aliasListWrite) { + // rootWrite is not yet inplace, we want to determine if it can be inplace + // so we consider all its write uses, not just the inplace ones. for (auto &uWrite : vWrite.getUses()) { - if (!bufferizesToMemoryWrite(uWrite, InPlaceSpec::True)) - continue; - usesWrite.insert(&uWrite); + if ((uWrite.getOwner() == opToBufferize && + bufferizesToMemoryWrite(uWrite)) || + bufferizesToMemoryWrite(uWrite, InPlaceSpec::True)) + usesWrite.insert(&uWrite); } } @@ -775,16 +781,17 @@ LDBG("----++++aliasRead #" << uRead->getOperandNumber() << " in: " << *aliasingReadOp << '\n'); for (OpOperand *uWrite : usesWrite) { - // Don't consider self-use of the same operand. - // Uses within the same op is fine though. + // Don't consider self-use of the same operand for interference. + // Multiple different uses within the same op is fair game though. if (uWrite == uRead) continue; + Operation *aliasingWriteOp = uWrite->getOwner(); LDBG("---- aliasWrite #" << uWrite->getOperandNumber() << " in: " << *aliasingWriteOp << '\n'); - // If read and written value already alias, no interference would be added - // by bufferizing inplace. - if (getAliasInfoRef(uRead->get()).contains(uWrite->get())) + // If the candidate write is the one that produces the read value (in the + // SSA def-use sense), this is not considered an interference. + if (getInplaceableOpResult(*uWrite) == uRead->get()) continue; // If aliasingReadOp properly dominates aliasingWriteOp, the read cannot // be affected by the write: there is no interference. @@ -815,35 +822,6 @@ return false; } -/// Return true if we find any read to opOperand.get() or any of its aliases, -/// that does not dominate opOperand.getOwner(). -bool BufferizationAliasInfo::existsNonDominatingRead( - OpOperand &opOperand, const DominanceInfo &domInfo) const { - LDBG("----Start existsNonDominatingRead\n"); - Operation *op = opOperand.getOwner(); - for (Value alias : getAliasInfoRef(opOperand.get())) { - for (OpOperand &wantReadUse : alias.getUses()) { - LDBG("--------current operand #" << wantReadUse.getOperandNumber() << ": " - << *(wantReadUse.getOwner()) << '\n'); - if (!bufferizesToMemoryRead(wantReadUse)) { - LDBG("------------not a read -> skip\n"); - continue; - } - if (&wantReadUse == &opOperand) { - LDBG("------------self-read is not an interference -> skip\n"); - continue; - } - if (domInfo.properlyDominates(wantReadUse.getOwner(), op)) { - LDBG("------------read properly dominates -> skip\n"); - continue; - } - LDBG("----found interfering read of " << wantReadUse.get() << '\n'); - return true; - } - } - return false; -} - /// Return true if the source of a `insertSliceOp` bufferizes to an /// equivalent ExtractSliceOp. bool BufferizationAliasInfo::isSourceEquivalentToAMatchingExtractSliceOp( @@ -985,12 +963,6 @@ assert(!domInfo.properlyDominates(aliasingReadOp, aliasingWriteOp) && "Unexpected aliasingReadOp properly dominates aliasingWriteOp"); - assert(((rootRead.isa() && - rootRead.getDefiningOp() == opToBufferize) || - (rootWrite.isa() && - rootWrite.getDefiningOp() == opToBufferize)) && - "Expected rootRead or rootWrite to be produced by opToBufferize"); - // Bail if the write does not dominate the read: it may clobber but only on // a strict subset of paths, which is not enough for safety. if (!domInfo.dominates(aliasingWriteOp, aliasingReadOp)) { @@ -1616,8 +1588,10 @@ << result << '\n'); // `result` must bufferize to a writeable buffer to be a candidate. - // This means the use->def chain not backpropagate to a function that is - // not inplaceable or to a constant op to be considered. + // This means the operand must not alias either: + // 1. a function bbArg that is not inplaceable or + // 2. a constant op. + // to be considered for inplace bufferization bool wouldCreateAliasingWriteToNonWriteableBuffer = aliasInfo.aliasesNonWriteableBuffer(operand); if (wouldCreateAliasingWriteToNonWriteableBuffer) @@ -1625,15 +1599,14 @@ else LDBG("->bufferizes to writeable inplace buffer\n"); - Value s = operand.get(), r = result; + assert(result == getInplaceableOpResult(operand)); + Value s = operand.get(); bool foundInterference = wouldCreateAliasingWriteToNonWriteableBuffer || - aliasInfo.existsNonDominatingRead(operand, domInfo) || - // Do not consider (s, s) and (r, r) as all the aliasings already - // exist by construction; we are interested in new interfering aliases - // only. - aliasInfo.wouldCreateReadAfterWriteInterference(s, r, op, domInfo) || - aliasInfo.wouldCreateReadAfterWriteInterference(r, s, op, domInfo); + // To determine whether this can bufferize inplace, check whether an + // aliasing read of `s` and an aliasing write of `s` would create a RaW + // violation. + aliasInfo.wouldCreateReadAfterWriteInterference(s, s, op, domInfo); if (foundInterference) aliasInfo.bufferizeOutOfPlace(result);