[Mlir-commits] [mlir] b9982b2 - [mlir][bufferization] Add rename function to BufferViewFlowAnalysis

[Martin Erhart]

Author: Martin Erhart
Date: 2023-08-02T10:59:19Z
New Revision: b9982b203b5c22152510e49fc05e505ecd970495

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

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

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
    mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
index 5905204fe9993d..24825db69f90c5 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h
@@ -57,6 +57,12 @@ class BufferViewFlowAnalysis {
   /// 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 f9c32d777a3bd8..003fc62f657857 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