[Mlir-commits] [mlir] d473dac - [mlir][SubElementInterfaces] Add a recursivelyReplaceElementsIn helper to AttrTypeReplacer

[River Riddle]

Author: River Riddle
Date: 2022-11-18T02:09:57-08:00
New Revision: d473dac3d63f7b54dd29f9c3e8077267457a4254

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

LOG: [mlir][SubElementInterfaces] Add a recursivelyReplaceElementsIn helper to AttrTypeReplacer

This somewhat improves the ergonomics when replacing recursively within
a set of IR.

Added: 
    

Modified: 
    mlir/include/mlir/IR/SubElementInterfaces.h
    mlir/lib/IR/SubElementInterfaces.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/SubElementInterfaces.h b/mlir/include/mlir/IR/SubElementInterfaces.h
index 5c362c1d3d57..9b14b1790ee5 100644
--- a/mlir/include/mlir/IR/SubElementInterfaces.h
+++ b/mlir/include/mlir/IR/SubElementInterfaces.h
@@ -40,6 +40,12 @@ class AttrTypeReplacer {
   void replaceElementsIn(Operation *op, bool replaceAttrs = true,
                          bool replaceLocs = false, bool replaceTypes = false);
 
+  /// Replace the elements within the given operation, and all nested
+  /// operations.
+  void recursivelyReplaceElementsIn(Operation *op, bool replaceAttrs = true,
+                                    bool replaceLocs = false,
+                                    bool replaceTypes = false);
+
   /// Replace the given attribute/type, and recursively replace any sub
   /// elements. Returns either the new attribute/type, or nullptr in the case of
   /// failure.

diff  --git a/mlir/lib/IR/SubElementInterfaces.cpp b/mlir/lib/IR/SubElementInterfaces.cpp
index d6130a672845..606b3c6041fe 100644
--- a/mlir/lib/IR/SubElementInterfaces.cpp
+++ b/mlir/lib/IR/SubElementInterfaces.cpp
@@ -145,6 +145,15 @@ void AttrTypeReplacer::replaceElementsIn(Operation *op, bool replaceAttrs,
   }
 }
 
+void AttrTypeReplacer::recursivelyReplaceElementsIn(Operation *op,
+                                                    bool replaceAttrs,
+                                                    bool replaceLocs,
+                                                    bool replaceTypes) {
+  op->walk([&](Operation *nestedOp) {
+    replaceElementsIn(nestedOp, replaceAttrs, replaceLocs, replaceTypes);
+  });
+}
+
 template <typename T>
 static void updateSubElementImpl(T element, AttrTypeReplacer &replacer,
                                  DenseMap<T, T> &elementMap,