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

[springerm]

Summary

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

Depends On D113373

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

public:
   /// Merge result's and operand's aliasing sets and iterate to a fixed point.
   void bufferizeInPlace(OpResult result, OpOperand &operand);
 
   /// Set the inPlace bufferization spec to false.
   void bufferizeOutOfPlace(OpResult result);
 
   /// 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`.
   void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
 
   /// Apply `fun` to all aliases of `v`.
   void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
 

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

for (OpOperand &operand : op->getOpOperands()) {
       Value replacement = rewriteFunc(b, initTensor.getLoc(), operand);
       if (!replacement)
         continue;
 
       // Uses of the InitTensorOp are replaced here, but the op is not deleted.
       // 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>()) {
         if (!skipAnalysis) {
           SmallVector<Operation *> ops(1, replacement.getDefiningOp());
           if (failed(inPlaceAnalysis(ops, aliasInfo, domInfo)))
             return WalkResult::interrupt();
         }

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @buffer_forwarding_no_conflict(
   // InitTensorOp replaces the init_tensor with an inplace extract_slice.
   // CHECK: %[[T_SUBVIEW:.*]] =  memref.subview %[[FUNC_ARG]][42] [%[[sz]]] [1]
   %a = linalg.init_tensor[%sz] : tensor<?xf32>
 
   // CHECK: linalg.fill({{.*}}, %[[T_SUBVIEW]]) : f32, memref<?xf32
   %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>
 }
 
 // -----
 
 // CHECK-LABEL: func @scf_if_inplace(