[mlir][bufferization] Add rename function to BufferViewFlowAnalysis

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).

Reviewed By: springerm

Differential Revision: https://reviews.llvm.org/D156665

2 files changed, 18 insertions, 0 deletions

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
index 5905204..24825db 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
@@ -57,6 +57,12 @@ public:
   /// 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.
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp
index f9c32d7..003fc62 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp
@@ -45,6 +45,18 @@ void BufferViewFlowAnalysis::remove(const SetVector<Value> &aliasValues) {
     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