[mlir][AttrTypeReplacer] Make attribute dictionary replacement optional

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

Author: River707

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

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 different,
   // 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)