[mlir][linalg][bufferize] Fix bug in InitTensor elimination

After replacing then init_tensor with a new value, the new value must be inserted into the corresponding union/equivalence sets.

Differential Revision: https://reviews.llvm.org/D113374
This commit is contained in:
Matthias Springer 2021-11-11 10:26:14 +09:00
parent 476ab0f809
commit 996d4ffe30
3 changed files with 15 additions and 7 deletions

View File

@ -62,13 +62,20 @@ public:
/// Return true if `v1` and `v2` bufferize to equivalent buffers.
bool areEquivalentBufferizedValues(Value v1, Value v2) const {
// Return `false` if we have no information about `v1` or `v2`.
if (equivalentInfo.findValue(v1) == equivalentInfo.end() ||
equivalentInfo.findValue(v2) == equivalentInfo.end())
return false;
return equivalentInfo.isEquivalent(v1, v2);
}
return equivalentInfo.getLeaderValue(v1) ==
equivalentInfo.getLeaderValue(v2);
/// Return true if `v1` and `v2` bufferize to aliasing buffers.
bool areAliasingBufferizedValues(Value v1, Value v2) const {
return aliasInfo.isEquivalent(v1, v2);
}
/// Union the alias sets of `v1` and `v2`.
void unionAliasSets(Value v1, Value v2) { aliasInfo.unionSets(v1, v2); }
/// Union the equivalence classes of `v1` and `v2`.
void unionEquivalenceClasses(Value v1, Value v2) {
equivalentInfo.unionSets(v1, v2);
}
/// Apply `fun` to all the members of the equivalence class of `v`.

View File

@ -2079,6 +2079,8 @@ LogicalResult mlir::linalg::comprehensive_bufferize::initTensorElimination(
// InitTensorOps without uses are ignored by the bufferization.
initTensor.replaceAllUsesWith(replacement);
aliasInfo.createAliasInfoEntry(replacement);
aliasInfo.unionAliasSets(initTensor, replacement);
aliasInfo.unionEquivalenceClasses(initTensor, replacement);
// Run analysis on the newly created op.
if (auto opResult = replacement.dyn_cast<OpResult>()) {

View File

@ -860,7 +860,6 @@ func @buffer_forwarding_no_conflict(
%f = linalg.fill(%f0, %a) : f32, tensor<?xf32> -> tensor<?xf32>
// Self-copy canonicalizes away later.
// CHECK: linalg.copy(%[[T_SUBVIEW]], %[[T_SUBVIEW]])
%r1 = tensor.insert_slice %f into %t[42][%sz][1]: tensor<?xf32> into tensor<?xf32>
return %r1: tensor<?xf32>