[Mlir-commits] [mlir] 1311306 - [mlir][AttrTypeReplacer] Make attribute dictionary replacement optional

[River Riddle]

Author: River Riddle
Date: 2022-11-14T18:02:18-08:00
New Revision: 1311306410e4c665ff93651bd64b69e14888596d

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

LOG: [mlir][AttrTypeReplacer] Make attribute dictionary replacement optional

This provides an optimization opportunity for clients that don't want/need
to recurse attribute dictionaries.

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 016269282c22e..5c362c1d3d57e 100644
--- a/mlir/include/mlir/IR/SubElementInterfaces.h
+++ b/mlir/include/mlir/IR/SubElementInterfaces.h
@@ -31,13 +31,14 @@ class AttrTypeReplacer {
   // Application
   //===--------------------------------------------------------------------===//
 
-  /// Replace the elements within the given operation. By default this includes
-  /// the attributes within the operation. If `replaceLocs` is true, this also
-  /// updates its location, the locations of any nested block arguments. If
-  /// `replaceTypes` is true, this also updates the result types of the
-  /// operation, and the types of any nested block arguments.
-  void replaceElementsIn(Operation *op, bool replaceLocs = false,
-                         bool replaceTypes = false);
+  /// Replace the elements within the given operation. If `replaceAttrs` is
+  /// true, this updates the attribute dictionary of the operation. If
+  /// `replaceLocs` is true, this also updates its location, and the locations
+  /// of any nested block arguments. If `replaceTypes` is true, this also
+  /// updates the result types of the operation, and the types of any nested
+  /// block arguments.
+  void replaceElementsIn(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

diff  --git a/mlir/lib/IR/SubElementInterfaces.cpp b/mlir/lib/IR/SubElementInterfaces.cpp
index 88aeaf191d90a..d6130a672845f 100644
--- a/mlir/lib/IR/SubElementInterfaces.cpp
+++ b/mlir/lib/IR/SubElementInterfaces.cpp
@@ -95,17 +95,20 @@ void SubElementTypeInterface::walkSubElements(
 /// AttrTypeReplacer
 //===----------------------------------------------------------------------===//
 
-void AttrTypeReplacer::replaceElementsIn(Operation *op, bool replaceLocs,
-                                         bool replaceTypes) {
+void AttrTypeReplacer::replaceElementsIn(Operation *op, bool replaceAttrs,
+                                         bool replaceLocs, bool replaceTypes) {
   // Functor that replaces the given element if the new value is 
diff erent,
   // otherwise returns nullptr.
   auto replaceIfDifferent = [&](auto element) {
     auto replacement = replace(element);
     return (replacement && replacement != element) ? replacement : nullptr;
   };
-  // Check the attribute dictionary for replacements.
-  if (auto newAttrs = replaceIfDifferent(op->getAttrDictionary()))
-    op->setAttrs(cast<DictionaryAttr>(newAttrs));
+
+  // Update the attribute dictionary.
+  if (replaceAttrs) {
+    if (auto newAttrs = replaceIfDifferent(op->getAttrDictionary()))
+      op->setAttrs(cast<DictionaryAttr>(newAttrs));
+  }
 
   // If we aren't updating locations or types, we're done.
   if (!replaceTypes && !replaceLocs)