[Mlir-commits] [mlir] 996d4ff - [mlir][linalg][bufferize] Fix bug in InitTensor elimination

[Matthias Springer]

Author: Matthias Springer
Date: 2021-11-11T10:28:17+09:00
New Revision: 996d4ffe30b5afb1abe51833ab7495c10675e5af

URL: https://github.com/llvm/llvm-project/commit/996d4ffe30b5afb1abe51833ab7495c10675e5af
DIFF: https://github.com/llvm/llvm-project/commit/996d4ffe30b5afb1abe51833ab7495c10675e5af.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
    mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
    mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
index 8ed623ac3694..b9021fac3606 100644
--- a/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
+++ b/mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h
@@ -62,13 +62,20 @@ class BufferizationAliasInfo {
 
   /// 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 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); }
 
-    return equivalentInfo.getLeaderValue(v1) ==
-           equivalentInfo.getLeaderValue(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`.

diff  --git a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 1d584011420d..89f6782389af 100644
--- a/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -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>()) {

diff  --git a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
index 564078ee0859..dcd409da32d0 100644
--- a/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir
@@ -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>