[mlir][bufferization] Add rename function to BufferViewFlowAnalysis

[maerhart]

Summary

This new function to replace a Value with another Value saves us from re-running
the entire alias analysis when an operation has to be re-build because
additional result values have to be added (e.g., when adding more iter_args to
an scf.for).

[springerm]

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h

62

built

[maerhart]

Address comment

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h

public:
   /// C = B
   ///
   /// Results in resolve(B) returning {B, C}
   ValueSetT resolve(Value value) const;
 
   /// Removes the given values from all alias sets.
   void remove(const SetVector<Value> &aliasValues);
 
+  /// Replaces all occurrences of 'from' in the internal datastructures with
+  /// 'to'. This is useful when the defining operation of a value has to be
+  /// re-built because additional results have to be added or the types of
+  /// results have to be changed.
+  void rename(Value from, Value to);
+
 private:
   /// This function constructs a mapping from values to its immediate
   /// dependencies.
   void build(Operation *op);
 
   /// Maps values to all immediate dependencies this value can have.
   ValueMapT dependencies;
 };
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERVIEWFLOWANALYSIS_H

mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

 }
 
 /// Removes the given values from all alias sets.
 void BufferViewFlowAnalysis::remove(const SetVector<Value> &aliasValues) {
   for (auto &entry : dependencies)
     llvm::set_subtract(entry.second, aliasValues);
 }
 
+void BufferViewFlowAnalysis::rename(Value from, Value to) {
+  dependencies[to] = dependencies[from];
+  dependencies.erase(from);
+
+  for (auto &[key, value] : dependencies) {
+    if (value.contains(from)) {
+      value.insert(to);
+      value.erase(from);
+    }
+  }
+}
+
 /// This function constructs a mapping from values to its immediate
 /// dependencies. It iterates over all blocks, gets their predecessors,
 /// determines the values that will be passed to the corresponding block
 /// arguments and inserts them into the underlying map. Furthermore, it wires
 /// successor regions and branch-like return operations from nested regions.
 void BufferViewFlowAnalysis::build(Operation *op) {
   // Registers all dependencies of the given values.
   auto registerDependencies = [&](ValueRange values, ValueRange dependencies) {