[mlir] Refactor SubElementInterface replace support

[rriddle]

Summary

The current support was essentially the amount necessary
to support replacing SymbolRefAttrs, but suffers from various
deficiencies (both ergonomic and functional):

• Replace crashes if unsupported This makes it really hard to use safely, given that you don't know if you are going to crash or not when using it.

• Types aren't supported

This seems like a simple missed addition when the attribute replacement
support was originally added.

• The ergonomics are weird

It currently uses an index based replacement, which makes the implementations
quite clunky.

This commit refactors support to be a bit more ergonomic, and also
adds support for types in the process. This was also a great oppurtunity
to greatly simplify how replacement is done in the symbol table.

[lattner]

Oh my, this is delightful, thank you for tackling this!

As a follow-on, it would also be amazing for aggregate types like FunctionType etc to implement these types. I'm also curious whether terminal types/attributes should implement these interfaces as well. It is convenient for clients if IntegerAttr knows it has no subattrs for example.

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

419

nice catch

mlir/test/lib/Dialect/Test/TestTypes.h

161

I'm not sure what this means? You can't mutate a type of course, you can only return a new one, just like attributes. For example, if I want to programmatically turn all index types into i64, then you'd replace { f32, index, f32} by returning {f32, i64, f32}. Does that make sense?

[rriddle]

mlir/test/lib/Dialect/Test/TestTypes.h

161

You can mutate them. There is a very niche feature that allows for having "mutable" components of an attribute or type: https://mlir.llvm.org/docs/AttributesAndTypes/#mutable-attributes-and-types

It was added many moons ago for LLVM struct type, but it hasn't been evolved much since.

[lattner]

mlir/test/lib/Dialect/Test/TestTypes.h

161

Aha, got it, it's logically immutable but needs to be mutated during set up. Makes sense!

[rriddle]

In D130589#3680398, @lattner wrote:

Oh my, this is delightful, thank you for tackling this!

As a follow-on, it would also be amazing for aggregate types like FunctionType etc to implement these types. I'm also curious whether terminal types/attributes should implement these interfaces as well. It is convenient for clients if IntegerAttr knows it has no subattrs for example.

At first it was simply intended for opt-in behavior, but if we have cases that are failing conservatively without it, it might make sense to shift to a similar model to side effects (i.e. NoSubElements or something).

[lattner]

Yeah good point, a trait like that would also nicely de-boilerplate a noop implementation.

Of course right now all attributes have types :-), so every attribute theoretically needs an impl of this. I think we can overlook that though.

[csigg]

Hi River, I noticed a behavioral change that might or might not be intended:

gpu.module @foo {
  llvm.func @foo() { ... }
}
func.func @main() {
  gpu.launch_func @foo::@foo ...
}

SymbolTable::replaceAllSymbolUses(gpu_module_op, "bar", main_func_op) before this change updated the gpu.launch_func to @bar::@foo. After this change, it is updated to @bar::@bar.

It seems to me that SymbolTable::replaceAllSymbolUses(Operation* op, ...) should only update references to op, not all references that match op's symbol name. WDYT?

[rriddle]

In D130589#3684380, @csigg wrote:

Hi River, I noticed a behavioral change that might or might not be intended:

gpu.module @foo {
  llvm.func @foo() { ... }
}
func.func @main() {
  gpu.launch_func @foo::@foo ...
}

SymbolTable::replaceAllSymbolUses(gpu_module_op, "bar", main_func_op) before this change updated the gpu.launch_func to @bar::@foo. After this change, it is updated to @bar::@bar.

It seems to me that SymbolTable::replaceAllSymbolUses(Operation* op, ...) should only update references to op, not all references that match op's symbol name. WDYT?

Yep, that's definitely a bug. Thanks for the repro: https://reviews.llvm.org/D130693

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

public:
   unsigned getABIAlignment(const DataLayout &dataLayout,
                            DataLayoutEntryListRef params) const;
 
   unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMFunctionType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect function type. It consists of a single return type (unlike MLIR
 /// which can have multiple), a list of parameter types and can optionally be
public:
   ArrayRef<Type> params() { return getParams(); }
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type result, ArrayRef<Type> arguments, bool);
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMPointerType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect pointer type. This type typically represents a reference to an
 /// object in memory. Pointers may be opaque or parameterized by the element
unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
   bool areCompatible(DataLayoutEntryListRef oldLayout,
                      DataLayoutEntryListRef newLayout) const;
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMStructType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect structure type representing a collection of different-typed
 /// elements manipulated together. Structured can optionally be packed, meaning
public:
   bool areCompatible(DataLayoutEntryListRef oldLayout,
                      DataLayoutEntryListRef newLayout) const;
 
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMFixedVectorType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect fixed vector type, represents a sequence of elements of known
 /// length that can be processed as one.
public:
   /// only a subset of LLVM dialect types that don't have a built-in
   /// counter-part, e.g., pointers.
   static bool isValidElementType(Type type);
 
   /// Returns the element type of the vector.
   Type getElementType() const;
 
   /// Returns the number of elements in the fixed vector.
-  unsigned getNumElements();
+  unsigned getNumElements() const;
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type elementType, unsigned numElements);
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect scalable vector type, represents a sequence of elements of
 /// unknown length that is known to be divisible by some constant. These
public:
   static bool isValidElementType(Type type);
 
   /// Returns the element type of the vector.
   Type getElementType() const;
 
   /// Returns the scaling factor of the number of elements in the vector. The
   /// vector contains at least the resulting number of elements, or any non-zero
   /// multiple of this number.
-  unsigned getMinNumElements();
+  unsigned getMinNumElements() const;
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type elementType, unsigned minNumElements);
 
   void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn) const;
+  Type replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
 };
 
 //===----------------------------------------------------------------------===//
 // Printing and parsing.
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Parses an LLVM dialect type.

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap", [
   let typeBuilder = "IndexType::get($_value.getContext())";
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_ArrayAttr : Builtin_Attr<"Array", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface,
-        ["replaceImmediateSubAttribute"]>
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
   ]> {
   let summary = "A collection of other Attribute values";
   let description = [{
     Syntax:
 
     ```
     array-attribute ::= `[` (attribute-value (`,` attribute-value)*)? `]`
     ```
def Builtin_DenseStringElementsAttr : Builtin_Attr<
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface,
-        ["replaceImmediateSubAttribute"]>
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
   ]> {
   let summary = "An dictionary of named Attribute values";
   let description = [{
     Syntax:
 
     ```
     dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? `}`
     ```
def Builtin_StringAttr : Builtin_Attr<"String"> {
   let genStorageClass = 0;
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolRefAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef"> {
+def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "An Attribute containing a symbolic reference to an Operation";
   let description = [{
     Syntax:
 
     ```
     symbol-ref-attribute ::= symbol-ref-id (`::` symbol-ref-id)*
     ```
 

mlir/include/mlir/IR/SubElementInterfaces.td

 #define MLIR_IR_SUBELEMENTINTERFACES_TD_
 
 include "mlir/IR/OpBase.td"
 
 //===----------------------------------------------------------------------===//
 // SubElementInterfaceBase
 //===----------------------------------------------------------------------===//
 
-class SubElementInterfaceBase<string interfaceName, string derivedValue> {
+class SubElementInterfaceBase<string interfaceName, string attrOrType,
+                              string derivedValue> {
   string cppNamespace = "::mlir";
 
   list<InterfaceMethod> methods = [
     InterfaceMethod<
       /*desc=*/[{
         Walk all of the immediately nested sub-attributes and sub-types. This
         method does not recurse into sub elements.
       }], "void", "walkImmediateSubElements",
       (ins "llvm::function_ref<void(mlir::Attribute)>":$walkAttrsFn,
            "llvm::function_ref<void(mlir::Type)>":$walkTypesFn)
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Replace the attributes identified by the indices with the corresponding
-        value. The index is derived from the order of the attributes returned by
-        the attribute callback of `walkImmediateSubElements`. An index of 0 would
-        replace the very first attribute given by `walkImmediateSubElements`.
-        The new instance with the values replaced is returned.
-      }], cppNamespace # "::" # interfaceName, "replaceImmediateSubAttribute",
-      (ins "::llvm::ArrayRef<std::pair<size_t, ::mlir::Attribute>>":$replacements),
-      [{}],
-      /*defaultImplementation=*/[{
-        llvm_unreachable("Attribute or Type does not support replacing attributes");
-      }]
-    >,
+        Replace the immediately nested sub-attributes and sub-types with those provided.
+        The order of the provided elements is derived from the order of the elements
+        returned by the callbacks of `walkImmediateSubElements`. The element at index 0
+        would replace the very first attribute given by `walkImmediateSubElements`.
+        On success, the new instance with the values replaced is returned. If replacement
+        fails, nullptr is returned.
+      }], attrOrType, "replaceImmediateSubElements", (ins
+        "::llvm::ArrayRef<::mlir::Attribute>":$replAttrs,
+        "::llvm::ArrayRef<::mlir::Type>":$replTypes
+      )>,
   ];
 
   code extraClassDeclaration = [{
-    /// Walk all of the held sub-attributes.
-    void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
-      walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
-    }
-
-    /// Walk all of the held sub-types.
-    void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
-      walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
-    }
-
     /// Walk all of the held sub-attributes and sub-types.
     void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
                          llvm::function_ref<void(mlir::Type)> walkTypesFn);
-  }];
 
+    /// Recursively replace all of the nested sub-attributes and sub-types using the
+    /// provided map functions. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceSubElements(
+      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn,
+      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn
+    );
+  }];
   code extraTraitClassDeclaration = [{
+    /// Walk all of the held sub-attributes and sub-types.
+    void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
+                         llvm::function_ref<void(mlir::Type)> walkTypesFn) {
+      }] # interfaceName # " interface(" # derivedValue # [{);
+      interface.walkSubElements(walkAttrsFn, walkTypesFn);
+    }
+
+    /// Recursively replace all of the nested sub-attributes and sub-types using the
+    /// provided map functions. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceSubElements(
+      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn,
+      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn) {
+      }] # interfaceName # " interface(" # derivedValue # [{);
+      return interface.replaceSubElements(replaceAttrFn, replaceTypeFn);
+    }
+
+    /// Recursively replace all of the nested sub-attributes and sub-types using the
+    /// provided map functions. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceImmediateSubElements(
+      llvm::ArrayRef<mlir::Attribute> replAttrs,
+      llvm::function_ref<mlir::Type(mlir::Type)> replTypes) {
+      return nullptr;
+    }
+  }];
+  code extraSharedClassDeclaration = [{
     /// Walk all of the held sub-attributes.
     void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
       walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
     }
-
     /// Walk all of the held sub-types.
     void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
       walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
     }
     
-    /// Walk all of the held sub-attributes and sub-types.
-    void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
-                         llvm::function_ref<void(mlir::Type)> walkTypesFn) {
-      }] # interfaceName # " interface(" # derivedValue # [{);
-      interface.walkSubElements(walkAttrsFn, walkTypesFn);
+    /// Recursively replace all of the nested sub-attributes using the provided
+    /// map function. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceSubElements(
+      llvm::function_ref<mlir::Attribute(mlir::Attribute)> replaceAttrFn) {
+      return replaceSubElements(
+        replaceAttrFn, [](mlir::Type type) { return type; });
+    }
+    /// Recursively replace all of the nested sub-types using the provided map
+    /// function. Returns nullptr in the case of failure.
+    }] # attrOrType # [{ replaceSubElements(
+      llvm::function_ref<mlir::Type(mlir::Type)> replaceTypeFn) {
+      return replaceSubElements(
+        [](mlir::Attribute attr) { return attr; }, replaceTypeFn);
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SubElementAttrInterface
 //===----------------------------------------------------------------------===//
 
 def SubElementAttrInterface
     : AttrInterface<"SubElementAttrInterface">,
-      SubElementInterfaceBase<"SubElementAttrInterface", "$_attr"> {
+      SubElementInterfaceBase<"SubElementAttrInterface", "::mlir::Attribute",
+                              "$_attr"> {
   let description = [{
     An interface used to query and manipulate sub-elements, such as sub-types
     and sub-attributes of a composite attribute.
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // SubElementTypeInterface
 //===----------------------------------------------------------------------===//
 
 def SubElementTypeInterface
     : TypeInterface<"SubElementTypeInterface">,
-      SubElementInterfaceBase<"SubElementTypeInterface", "$_type"> {
+      SubElementInterfaceBase<"SubElementTypeInterface", "::mlir::Type",
+                              "$_type"> {
   let description = [{
     An interface used to query and manipulate sub-elements, such as sub-types
     and sub-attributes of a composite type.
   }];
 }
 
 #endif // MLIR_IR_SUBELEMENTINTERFACES_TD_

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

 }
 
 void LLVMArrayType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type LLVMArrayType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes.front(), getNumElements());
+}
+
 //===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMFunctionType::isValidArgumentType(Type type) {
   return !type.isa<LLVMVoidType, LLVMFunctionType>();
 }
 

 
 void LLVMFunctionType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Type type : llvm::concat<const Type>(getReturnTypes(), getParams()))
     walkTypesFn(type);
 }
 
+Type LLVMFunctionType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes.front(), replTypes.drop_front(), isVarArg());
+}
+
 //===----------------------------------------------------------------------===//
 // Pointer type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMPointerType::isValidElementType(Type type) {
   if (!type)
     return true;
   return isCompatibleOuterType(type)

 }
 
 void LLVMPointerType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type LLVMPointerType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes.front(), getAddressSpace());
+}
+
 //===----------------------------------------------------------------------===//
 // Struct type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMStructType::isValidElementType(Type type) {
   return !type.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                    LLVMFunctionType, LLVMTokenType, LLVMScalableVectorType>();
 }

 
 void LLVMStructType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Type type : getBody())
     walkTypesFn(type);
 }
 
+Type LLVMStructType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  // TODO: It's not clear how we support replacing sub-elements of mutable
+  // types.
+  return nullptr;
+}
+
 //===----------------------------------------------------------------------===//
 // Vector types.
 //===----------------------------------------------------------------------===//
 
 /// Verifies that the type about to be constructed is well-formed.
 template <typename VecTy>
 static LogicalResult
 verifyVectorConstructionInvariants(function_ref<InFlightDiagnostic()> emitError,
LLVMFixedVectorType::getChecked(function_ref<InFlightDiagnostic()> emitError,
   return Base::getChecked(emitError, elementType.getContext(), elementType,
                           numElements);
 }
 
 Type LLVMFixedVectorType::getElementType() const {
   return static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->elementType;
 }
 
-unsigned LLVMFixedVectorType::getNumElements() {
+unsigned LLVMFixedVectorType::getNumElements() const {
   return getImpl()->numElements;
 }
 
 bool LLVMFixedVectorType::isValidElementType(Type type) {
   return type.isa<LLVMPointerType, LLVMPPCFP128Type>();
 }
 
 LogicalResult
 LLVMFixedVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                             Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMFixedVectorType>(
       emitError, elementType, numElements);
 }
 
 void LLVMFixedVectorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type LLVMFixedVectorType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes[0], getNumElements());
+}
+
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType.
 //===----------------------------------------------------------------------===//
 
 LLVMScalableVectorType LLVMScalableVectorType::get(Type elementType,
                                                    unsigned minNumElements) {
   assert(elementType && "expected non-null subtype");
   return Base::get(elementType.getContext(), elementType, minNumElements);
 }
 
 LLVMScalableVectorType
 LLVMScalableVectorType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                                    Type elementType, unsigned minNumElements) {
   assert(elementType && "expected non-null subtype");
   return Base::getChecked(emitError, elementType.getContext(), elementType,
                           minNumElements);
 }
 
 Type LLVMScalableVectorType::getElementType() const {
   return static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->elementType;
 }
 
-unsigned LLVMScalableVectorType::getMinNumElements() {
+unsigned LLVMScalableVectorType::getMinNumElements() const {
   return getImpl()->numElements;
 }
 
 bool LLVMScalableVectorType::isValidElementType(Type type) {
   if (auto intType = type.dyn_cast<IntegerType>())
     return intType.isSignless();
 
   return isCompatibleFloatingPointType(type) || type.isa<LLVMPointerType>();
 }
 
 LogicalResult
 LLVMScalableVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                                Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMScalableVectorType>(
       emitError, elementType, numElements);
 }
 
 void LLVMScalableVectorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type LLVMScalableVectorType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes[0], getMinNumElements());
+}
+
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
 bool mlir::LLVM::isCompatibleOuterType(Type type) {
   // clang-format off
   if (type.isa<
       BFloat16Type,

mlir/lib/IR/BuiltinAttributes.cpp

 
 void ArrayAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Attribute attr : getValue())
     walkAttrsFn(attr);
 }
 
-SubElementAttrInterface ArrayAttr::replaceImmediateSubAttribute(
-    ArrayRef<std::pair<size_t, Attribute>> replacements) const {
-  std::vector<Attribute> vector = getValue().vec();
-  for (auto &it : replacements) {
-    vector[it.first] = it.second;
-  }
-  return get(getContext(), vector);
+Attribute
+ArrayAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                       ArrayRef<Type> replTypes) const {
+  return get(getContext(), replAttrs);
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// Helper function that does either an in place sort or sorts from source array
 /// into destination. If inPlace then storage is both the source and the

 
 void DictionaryAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (const NamedAttribute &attr : getValue())
     walkAttrsFn(attr.getValue());
 }
 
-SubElementAttrInterface DictionaryAttr::replaceImmediateSubAttribute(
-    ArrayRef<std::pair<size_t, Attribute>> replacements) const {
+Attribute
+DictionaryAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                            ArrayRef<Type> replTypes) const {
   std::vector<NamedAttribute> vec = getValue().vec();
-  for (auto &it : replacements)
-    vec[it.first].setValue(it.second);
+  for (auto &it : llvm::enumerate(replAttrs))
+    vec[it.index()].setValue(it.value());
 
   // The above only modifies the mapped value, but not the key, and therefore
   // not the order of the elements. It remains sorted
   return getWithSorted(getContext(), vec);
 }
 
 //===----------------------------------------------------------------------===//
 // StringAttr
FlatSymbolRefAttr SymbolRefAttr::get(Operation *symbol) {
   return SymbolRefAttr::get(symName);
 }
 
 StringAttr SymbolRefAttr::getLeafReference() const {
   ArrayRef<FlatSymbolRefAttr> nestedRefs = getNestedReferences();
   return nestedRefs.empty() ? getRootReference() : nestedRefs.back().getAttr();
 }
 
+void SymbolRefAttr::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkAttrsFn(getRootReference());
+  for (FlatSymbolRefAttr ref : getNestedReferences())
+    walkAttrsFn(ref);
+}
+
+Attribute
+SymbolRefAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                           ArrayRef<Type> replTypes) const {
+  ArrayRef<Attribute> rawNestedRefs = replAttrs.drop_front();
+  ArrayRef<FlatSymbolRefAttr> nestedRefs(
+      static_cast<const FlatSymbolRefAttr *>(rawNestedRefs.data()),
+      rawNestedRefs.size());
+  return get(replAttrs[0].cast<StringAttr>(), nestedRefs);
+}
+
 //===----------------------------------------------------------------------===//
 // IntegerAttr
 //===----------------------------------------------------------------------===//
 
 int64_t IntegerAttr::getInt() const {
   assert((getType().isIndex() || getType().isSignlessInteger()) &&
          "must be signless integer");
   return getValue().getSExtValue();

 // TypeAttr
 //===----------------------------------------------------------------------===//
 
 void TypeAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getValue());
 }
+
+Attribute
+TypeAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                      ArrayRef<Type> replTypes) const {
+  return get(replTypes[0]);
+}

mlir/lib/IR/BuiltinTypes.cpp

 
 void FunctionType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Type type : llvm::concat<const Type>(getInputs(), getResults()))
     walkTypesFn(type);
 }
 
+Type FunctionType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                               ArrayRef<Type> replTypes) const {
+  unsigned numInputs = getNumInputs();
+  return get(getContext(), replTypes.take_front(numInputs),
+             replTypes.drop_front(numInputs));
+}
+
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 /// Verify the construction of an opaque type.
 LogicalResult OpaqueType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  StringAttr dialect, StringRef typeData) {
   if (!Dialect::isValidNamespace(dialect.strref()))

 }
 
 void VectorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type VectorType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                             ArrayRef<Type> replTypes) const {
+  return get(getShape(), replTypes.front(), getNumScalableDims());
+}
+
 VectorType VectorType::cloneWith(Optional<ArrayRef<int64_t>> shape,
                                  Type elementType) const {
   return VectorType::get(shape.value_or(getShape()), elementType,
                          getNumScalableDims());
 }
 
 //===----------------------------------------------------------------------===//
 // TensorType

 void RankedTensorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
   if (Attribute encoding = getEncoding())
     walkAttrsFn(encoding);
 }
 
+Type RankedTensorType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(getShape(), replTypes.front(),
+             replAttrs.empty() ? Attribute() : replAttrs.back());
+}
+
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 UnrankedTensorType::verify(function_ref<InFlightDiagnostic()> emitError,
                            Type elementType) {
   return checkTensorElementType(emitError, elementType);
 }
 
 void UnrankedTensorType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
 }
 
+Type UnrankedTensorType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes.front());
+}
+
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 Type BaseMemRefType::getElementType() const {
   return llvm::TypeSwitch<BaseMemRefType, Type>(*this)
       .Case<MemRefType, UnrankedMemRefType>(
           [](auto type) { return type.getElementType(); });
void MemRefType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
   if (!getLayout().isIdentity())
     walkAttrsFn(getLayout());
   walkAttrsFn(getMemorySpace());
 }
 
+Type MemRefType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                             ArrayRef<Type> replTypes) const {
+  bool hasLayout = replAttrs.size() > 1;
+  return get(getShape(), replTypes[0],
+             hasLayout ? replAttrs[0].dyn_cast<MemRefLayoutAttrInterface>()
+                       : MemRefLayoutAttrInterface(),
+             hasLayout ? replAttrs[1] : replAttrs[0]);
+}
+
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 unsigned UnrankedMemRefType::getMemorySpaceAsInt() const {
   return detail::getMemorySpaceAsInt(getMemorySpace());
 }
 

 
 void UnrankedMemRefType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   walkTypesFn(getElementType());
   walkAttrsFn(getMemorySpace());
 }
 
+Type UnrankedMemRefType::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  return get(replTypes.front(), replAttrs.front());
+}
+
 //===----------------------------------------------------------------------===//
 /// TupleType
 //===----------------------------------------------------------------------===//
 
 /// Return the elements types for this tuple.
 ArrayRef<Type> TupleType::getTypes() const { return getImpl()->getTypes(); }
 
 /// Accumulate the types contained in this tuple and tuples nested within it.

 
 void TupleType::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) const {
   for (Type type : getTypes())
     walkTypesFn(type);
 }
 
+Type TupleType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                            ArrayRef<Type> replTypes) const {
+  return get(getContext(), replTypes);
+}
+
 //===----------------------------------------------------------------------===//
 // Type Utilities
 //===----------------------------------------------------------------------===//
 
 AffineMap mlir::makeStridedLinearLayoutMap(ArrayRef<int64_t> strides,
                                            int64_t offset,
                                            MLIRContext *context) {
   AffineExpr expr;

mlir/lib/IR/SubElementInterfaces.cpp

 //===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SubElementInterfaces.h"
 
 #include "llvm/ADT/DenseSet.h"
 
 using namespace mlir;
 
+//===----------------------------------------------------------------------===//
+// SubElementInterface
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// WalkSubElements
+
 template <typename InterfaceT>
 static void walkSubElementsImpl(InterfaceT interface,
                                 function_ref<void(Attribute)> walkAttrsFn,
                                 function_ref<void(Type)> walkTypesFn,
                                 DenseSet<Attribute> &visitedAttrs,
                                 DenseSet<Type> &visitedTypes) {
   interface.walkImmediateSubElements(
       [&](Attribute attr) {
void SubElementTypeInterface::walkSubElements(
   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
   DenseSet<Attribute> visitedAttrs;
   DenseSet<Type> visitedTypes;
   walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
                       visitedTypes);
 }
 
 //===----------------------------------------------------------------------===//
+// ReplaceSubElements
+
+/// Return if the given element is mutable.
+static bool isMutable(Attribute attr) {
+  return attr.hasTrait<AttributeTrait::IsMutable>();
+}
+static bool isMutable(Type type) {
+  return type.hasTrait<TypeTrait::IsMutable>();
+}
+
+template <typename InterfaceT, typename T, typename ReplaceSubElementFnT>
+static void updateSubElementImpl(T element, function_ref<T(T)> walkFn,
+                                 DenseMap<T, T> &visited,
+                                 SmallVectorImpl<T> &newElements,
+                                 FailureOr<bool> &changed,
+                                 ReplaceSubElementFnT &&replaceSubElementFn) {
+  // Bail early if we failed at any point.
+  if (failed(changed))
+    return;
+  newElements.push_back(element);
+
+  // Guard against potentially null inputs. We always map null to null.
+  if (!element)
+    return;
+
+  // Check for an existing mapping for this element, and walk it if we haven't
+  // yet.
+  T &mappedElement = visited[element];
+  if (!mappedElement) {
+    // Try walking this element.
+    if (!(mappedElement = walkFn(element))) {
+      changed = failure();
+      return;
+    }
+
+    // Handle replacing sub-elements if this element is also a container.
+    if (auto interface = mappedElement.template dyn_cast<InterfaceT>()) {
+      if (!(mappedElement = replaceSubElementFn(interface))) {
+        changed = failure();
+        return;
+      }
+    }
+  }
+
+  // Update to the mapped element.
+  if (mappedElement != element) {
+    newElements.back() = mappedElement;
+    changed = true;
+  }
+}
+
+template <typename InterfaceT>
+static typename InterfaceT::ValueType
+replaceSubElementsImpl(InterfaceT interface,
+                       function_ref<Attribute(Attribute)> walkAttrsFn,
+                       function_ref<Type(Type)> walkTypesFn,
+                       DenseMap<Attribute, Attribute> &visitedAttrs,
+                       DenseMap<Type, Type> &visitedTypes) {
+  // Walk the current sub-elements, replacing them as necessary.
+  SmallVector<Attribute, 16> newAttrs;
+  SmallVector<Type, 16> newTypes;
+  FailureOr<bool> changed = false;
+  auto replaceSubElementFn = [&](auto subInterface) {
+    return replaceSubElementsImpl(subInterface, walkAttrsFn, walkTypesFn,
+                                  visitedAttrs, visitedTypes);
+  };
+  interface.walkImmediateSubElements(
+      [&](Attribute element) {
+        updateSubElementImpl<SubElementAttrInterface>(
+            element, walkAttrsFn, visitedAttrs, newAttrs, changed,
+            replaceSubElementFn);
+      },
+      [&](Type element) {
+        updateSubElementImpl<SubElementTypeInterface>(
+            element, walkTypesFn, visitedTypes, newTypes, changed,
+            replaceSubElementFn);
+      });
+  if (failed(changed))
+    return {};
+
+  // If the sub-elements didn't change, just return the original value.
+  if (!*changed)
+    return interface;
+
+  // If this element is mutable, we don't support changing its sub elements, the
+  // sub element walk doesn't give us a valid ordering for what we need here. If
+  // we want to support mutable elements, we'll need something more.
+  if (isMutable(interface))
+    return {};
+
+  // Use the new elements during the replacement.
+  return interface.replaceImmediateSubElements(newAttrs, newTypes);
+}
+
+Attribute SubElementAttrInterface::replaceSubElements(
+    function_ref<Attribute(Attribute)> replaceAttrFn,
+    function_ref<Type(Type)> replaceTypeFn) {
+  assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
+  DenseMap<Attribute, Attribute> visitedAttrs;
+  DenseMap<Type, Type> visitedTypes;
+  return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
+                                visitedAttrs, visitedTypes);
+}
+
+Type SubElementTypeInterface::replaceSubElements(
+    function_ref<Attribute(Attribute)> replaceAttrFn,
+    function_ref<Type(Type)> replaceTypeFn) {
+  assert(replaceAttrFn && replaceTypeFn && "expected valid replace functions");
+  DenseMap<Attribute, Attribute> visitedAttrs;
+  DenseMap<Type, Type> visitedTypes;
+  return replaceSubElementsImpl(*this, replaceAttrFn, replaceTypeFn,
+                                visitedAttrs, visitedTypes);
+}
+
+//===----------------------------------------------------------------------===//
 // SubElementInterface Tablegen definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
 #include "mlir/IR/SubElementTypeInterfaces.cpp.inc"

mlir/lib/IR/SymbolTable.cpp

for (Operation &op : worklist.pop_back_val()->getOps()) {
         for (Region &region : op.getRegions())
           worklist.push_back(&region);
       }
     }
   }
   return WalkResult::advance();
 }
 
+/// Walk all of the operations nested under, and including, the given operation,
+/// without traversing into any nested symbol tables. Stops walking if the
+/// result of the callback is anything other than `WalkResult::advance`.
+static Optional<WalkResult>
+walkSymbolTable(Operation *op,
+                function_ref<Optional<WalkResult>(Operation *)> callback) {
+  Optional<WalkResult> result = callback(op);
+  if (result != WalkResult::advance() || op->hasTrait<OpTrait::SymbolTable>())
+    return result;
+  return walkSymbolTable(op->getRegions(), callback);
+}
+
 //===----------------------------------------------------------------------===//
 // SymbolTable
 //===----------------------------------------------------------------------===//
 
 /// Build a symbol table with the symbols within the given operation.
 SymbolTable::SymbolTable(Operation *symbolTableOp)
     : symbolTableOp(symbolTableOp) {
   assert(symbolTableOp->hasTrait<OpTrait::SymbolTable>() &&
LogicalResult detail::verifySymbol(Operation *op) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Symbol Use Lists
 //===----------------------------------------------------------------------===//
 
 /// Walk all of the symbol references within the given operation, invoking the
-/// provided callback for each found use. The callbacks takes as arguments: the
-/// use of the symbol, and the nested access chain to the attribute within the
-/// operation dictionary. An access chain is a set of indices into nested
-/// container attributes. For example, a symbol use in an attribute dictionary
-/// that looks like the following:
-///
-///    {use = [{other_attr, @symbol}]}
-///
-/// May have the following access chain:
-///
-///     [0, 0, 1]
-///
-static WalkResult walkSymbolRefs(
-    Operation *op,
-    function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
+/// provided callback for each found use. The callbacks takes the use of the
+/// symbol.
+static WalkResult
+walkSymbolRefs(Operation *op,
+               function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {
   // Check to see if the operation has any attributes.
   DictionaryAttr attrDict = op->getAttrDictionary();
   if (attrDict.empty())
     return WalkResult::advance();
 
   // A worklist of a container attribute and the current index into the held
   // attribute list.
   struct WorklistItem {
walkSymbolRefs(Operation *op,
   SmallVector<WorklistItem, 1> attrWorklist(1, WorklistItem(attrDict));
   SmallVector<int, 1> curAccessChain(1, /*Value=*/-1);
 
   // Process the symbol references within the given nested attribute range.
   auto processAttrs = [&](int &index,
                           WorklistItem &worklistItem) -> WalkResult {
     for (Attribute attr :
          llvm::drop_begin(worklistItem.immediateSubElements, index)) {
+      // Invoke the provided callback if we find a symbol use and check for a
+      // requested interrupt.
+      if (auto symbolRef = attr.dyn_cast<SymbolRefAttr>()) {
+        if (callback({op, symbolRef}).wasInterrupted())
+          return WalkResult::interrupt();
+
         /// Check for a nested container attribute, these will also need to be
         /// walked.
-      if (auto interface = attr.dyn_cast<SubElementAttrInterface>()) {
+      } else if (auto interface = attr.dyn_cast<SubElementAttrInterface>()) {
         attrWorklist.emplace_back(interface);
         curAccessChain.push_back(-1);
         return WalkResult::advance();
       }
-
-      // Invoke the provided callback if we find a symbol use and check for a
-      // requested interrupt.
-      if (auto symbolRef = attr.dyn_cast<SymbolRefAttr>())
-        if (callback({op, symbolRef}, curAccessChain).wasInterrupted())
-          return WalkResult::interrupt();
-
       // Make sure to keep the index counter in sync.
       ++index;
     }
 
     // Pop this container attribute from the worklist.
     attrWorklist.pop_back();
     curAccessChain.pop_back();
     return WalkResult::advance();
do {
     result = processAttrs(index, item);
   } while (!attrWorklist.empty() && !result.wasInterrupted());
   return result;
 }
 
 /// Walk all of the uses, for any symbol, that are nested within the given
 /// regions, invoking the provided callback for each. This does not traverse
 /// into any nested symbol tables.
-static Optional<WalkResult> walkSymbolUses(
-    MutableArrayRef<Region> regions,
-    function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
+static Optional<WalkResult>
+walkSymbolUses(MutableArrayRef<Region> regions,
+               function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {
   return walkSymbolTable(regions, [&](Operation *op) -> Optional<WalkResult> {
     // Check that this isn't a potentially unknown symbol table.
     if (isPotentiallyUnknownSymbolTable(op))
       return llvm::None;
 
     return walkSymbolRefs(op, callback);
   });
 }
 /// Walk all of the uses, for any symbol, that are nested within the given
 /// operation 'from', invoking the provided callback for each. This does not
 /// traverse into any nested symbol tables.
-static Optional<WalkResult> walkSymbolUses(
-    Operation *from,
-    function_ref<WalkResult(SymbolTable::SymbolUse, ArrayRef<int>)> callback) {
+static Optional<WalkResult>
+walkSymbolUses(Operation *from,
+               function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {
   // If this operation has regions, and it, as well as its dialect, isn't
   // registered then conservatively fail. The operation may define a
   // symbol table, so we can't opaquely know if we should traverse to find
   // nested uses.
   if (isPotentiallyUnknownSymbolTable(from))
     return llvm::None;
 
   // Walk the uses on this operation.
struct SymbolScope {
   }
   /// This variant is used when the callback type matches a stripped down type:
   /// void(SymbolTable::SymbolUse use)
   template <typename CallbackT,
             typename std::enable_if_t<std::is_same<
                 typename llvm::function_traits<CallbackT>::result_t,
                 void>::value> * = nullptr>
   Optional<WalkResult> walk(CallbackT cback) {
-    return walk([=](SymbolTable::SymbolUse use, ArrayRef<int>) {
+    return walk([=](SymbolTable::SymbolUse use) {
       return cback(use), WalkResult::advance();
     });
   }
 
+  /// Walk all of the operations nested under the current scope without
+  /// traversing into any nested symbol tables.
+  template <typename CallbackT>
+  Optional<WalkResult> walkSymbolTable(CallbackT &&cback) {
+    if (Region *region = limit.dyn_cast<Region *>())
+      return ::walkSymbolTable(*region, cback);
+    return ::walkSymbolTable(limit.get<Operation *>(), cback);
+  }
+
   /// The representation of the symbol within this scope.
   SymbolRefAttr symbol;
 
   /// The IR unit representing this scope.
   llvm::PointerUnion<Operation *, Region *> limit;
 };
 } // namespace
 

 
 //===----------------------------------------------------------------------===//
 // SymbolTable::getSymbolUses
 
 /// The implementation of SymbolTable::getSymbolUses below.
 template <typename FromT>
 static Optional<SymbolTable::UseRange> getSymbolUsesImpl(FromT from) {
   std::vector<SymbolTable::SymbolUse> uses;
-  auto walkFn = [&](SymbolTable::SymbolUse symbolUse, ArrayRef<int>) {
+  auto walkFn = [&](SymbolTable::SymbolUse symbolUse) {
     uses.push_back(symbolUse);
     return WalkResult::advance();
   };
   auto result = walkSymbolUses(from, walkFn);
   return result ? Optional<SymbolTable::UseRange>(std::move(uses)) : llvm::None;
 }
 
 /// Get an iterator range for all of the uses, for any symbol, that are nested

 //===----------------------------------------------------------------------===//
 // SymbolTable::symbolKnownUseEmpty
 
 /// The implementation of SymbolTable::symbolKnownUseEmpty below.
 template <typename SymbolT, typename IRUnitT>
 static bool symbolKnownUseEmptyImpl(SymbolT symbol, IRUnitT *limit) {
   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
     // Walk all of the symbol uses looking for a reference to 'symbol'.
-    if (scope.walk([&](SymbolTable::SymbolUse symbolUse, ArrayRef<int>) {
+    if (scope.walk([&](SymbolTable::SymbolUse symbolUse) {
           return isReferencePrefixOf(scope.symbol, symbolUse.getSymbolRef())
                      ? WalkResult::interrupt()
                      : WalkResult::advance();
         }) != WalkResult::advance())
       return false;
   }
   return true;
 }

 }
 bool SymbolTable::symbolKnownUseEmpty(Operation *symbol, Region *from) {
   return symbolKnownUseEmptyImpl(symbol, from);
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolTable::replaceAllSymbolUses
 
-/// Rebuild the given attribute container after replacing all references to a
-/// symbol with the updated attribute in 'accesses'.
-static SubElementAttrInterface rebuildAttrAfterRAUW(
-    SubElementAttrInterface container,
-    ArrayRef<std::pair<SmallVector<int, 1>, SymbolRefAttr>> accesses,
-    unsigned depth) {
-  // Given a range of Attributes, update the ones referred to by the given
-  // access chains to point to the new symbol attribute.
-
-  SmallVector<std::pair<size_t, Attribute>> replacements;
-
-  SmallVector<Attribute> subElements;
-  container.walkImmediateSubElements(
-      [&](Attribute attribute) { subElements.push_back(attribute); },
-      [](Type) {});
-  for (unsigned i = 0, e = accesses.size(); i != e;) {
-    ArrayRef<int> access = accesses[i].first;
-
-    // Check to see if this is a leaf access, i.e. a SymbolRef.
-    if (access.size() == depth + 1) {
-      replacements.emplace_back(access.back(), accesses[i].second);
-      ++i;
-      continue;
-    }
-
-    // Otherwise, this is a container. Collect all of the accesses for this
-    // index and recurse. The recursion here is bounded by the size of the
-    // largest access array.
-    auto nestedAccesses = accesses.drop_front(i).take_while([&](auto &it) {
-      ArrayRef<int> nextAccess = it.first;
-      return nextAccess.size() > depth + 1 &&
-             nextAccess[depth] == access[depth];
-    });
-    auto result = rebuildAttrAfterRAUW(subElements[access[depth]],
-                                       nestedAccesses, depth + 1);
-    replacements.emplace_back(access[depth], result);
-
-    // Skip over all of the accesses that refer to the nested container.
-    i += nestedAccesses.size();
-  }
-
-  return container.replaceImmediateSubAttribute(replacements);
-}
-
 /// Generates a new symbol reference attribute with a new leaf reference.
 static SymbolRefAttr generateNewRefAttr(SymbolRefAttr oldAttr,
                                         FlatSymbolRefAttr newLeafAttr) {
   if (oldAttr.isa<FlatSymbolRefAttr>())
     return newLeafAttr;
   auto nestedRefs = llvm::to_vector<2>(oldAttr.getNestedReferences());
   nestedRefs.back() = newLeafAttr;
   return SymbolRefAttr::get(oldAttr.getRootReference(), nestedRefs);
 }
 
 /// The implementation of SymbolTable::replaceAllSymbolUses below.
 template <typename SymbolT, typename IRUnitT>
 static LogicalResult
 replaceAllSymbolUsesImpl(SymbolT symbol, StringAttr newSymbol, IRUnitT *limit) {
-  // A collection of operations along with their new attribute dictionary.
-  std::vector<std::pair<Operation *, DictionaryAttr>> updatedAttrDicts;
-
-  // The current operation being processed.
-  Operation *curOp = nullptr;
-
-  // The set of access chains into the attribute dictionary of the current
-  // operation, as well as the replacement attribute to use.
-  SmallVector<std::pair<SmallVector<int, 1>, SymbolRefAttr>, 1> accessChains;
-
-  // Generate a new attribute dictionary for the current operation by replacing
-  // references to the old symbol.
-  auto generateNewAttrDict = [&] {
-    auto oldDict = curOp->getAttrDictionary();
-    auto newDict = rebuildAttrAfterRAUW(oldDict, accessChains, /*depth=*/0);
-    return newDict.cast<DictionaryAttr>();
-  };
-
   // Generate a new attribute to replace the given attribute.
   FlatSymbolRefAttr newLeafAttr = FlatSymbolRefAttr::get(newSymbol);
   for (SymbolScope &scope : collectSymbolScopes(symbol, limit)) {
+    SymbolRefAttr oldAttr = scope.symbol;
     SymbolRefAttr newAttr = generateNewRefAttr(scope.symbol, newLeafAttr);
-    auto walkFn = [&](SymbolTable::SymbolUse symbolUse,
-                      ArrayRef<int> accessChain) {
-      SymbolRefAttr useRef = symbolUse.getSymbolRef();
-      if (!isReferencePrefixOf(scope.symbol, useRef))
-        return WalkResult::advance();
 
-      // If we have a valid match, check to see if this is a proper
-      // subreference. If it is, then we will need to generate a different new
-      // attribute specifically for this use.
-      SymbolRefAttr replacementRef = newAttr;
-      if (useRef != scope.symbol) {
-        if (scope.symbol.isa<FlatSymbolRefAttr>()) {
-          replacementRef =
-              SymbolRefAttr::get(newSymbol, useRef.getNestedReferences());
-        } else {
-          auto nestedRefs = llvm::to_vector<4>(useRef.getNestedReferences());
-          nestedRefs[scope.symbol.getNestedReferences().size() - 1] =
-              newLeafAttr;
-          replacementRef =
-              SymbolRefAttr::get(useRef.getRootReference(), nestedRefs);
-        }
+    auto walkFn = [&](Operation *op) -> Optional<WalkResult> {
+      auto remapAttrFn = [&](Attribute attr) -> Attribute {
+        if (attr == oldAttr)
+          return newAttr;
+        // Handle prefix matches.
+        if (SymbolRefAttr symRef = attr.dyn_cast<SymbolRefAttr>()) {
+          if (isReferencePrefixOf(oldAttr, symRef)) {
+            auto oldNestedRefs = oldAttr.getNestedReferences();
+            auto nestedRefs = symRef.getNestedReferences();
+            if (oldNestedRefs.empty())
+              return SymbolRefAttr::get(newSymbol, nestedRefs);
+
+            auto newNestedRefs = llvm::to_vector<4>(nestedRefs);
+            newNestedRefs[oldNestedRefs.size() - 1] = newLeafAttr;
+            return SymbolRefAttr::get(symRef.getRootReference(), newNestedRefs);
           }
-
-      // If there was a previous operation, generate a new attribute dict
-      // for it. This means that we've finished processing the current
-      // operation, so generate a new dictionary for it.
-      if (curOp && symbolUse.getUser() != curOp) {
-        updatedAttrDicts.push_back({curOp, generateNewAttrDict()});
-        accessChains.clear();
         }
+        return attr;
+      };
+      // Generate a new attribute dictionary by replacing references to the old
+      // symbol.
+      auto newDict = op->getAttrDictionary().replaceSubElements(remapAttrFn);
+      if (!newDict)
+        return WalkResult::interrupt();
 
-      // Record this access.
-      curOp = symbolUse.getUser();
-      accessChains.push_back({llvm::to_vector<1>(accessChain), replacementRef});
+      op->setAttrs(newDict.template cast<DictionaryAttr>());
       return WalkResult::advance();
     };
-    if (!scope.walk(walkFn))
+    if (!scope.walkSymbolTable(walkFn))
       return failure();
-
-    // Check to see if we have a dangling op that needs to be processed.
-    if (curOp) {
-      updatedAttrDicts.push_back({curOp, generateNewAttrDict()});
-      curOp = nullptr;
-    }
-  }
-
-  // Update the attribute dictionaries as necessary.
-  for (auto &it : updatedAttrDicts)
-    it.first->setAttrs(it.second);
   return success();
 }
 
 /// Attempt to replace all uses of the given symbol 'oldSymbol' with the
 /// provided symbol 'newSymbol' that are nested within the given operation
 /// 'from'. This does not traverse into any nested symbol tables. If there are
 /// any unknown operations that may potentially be symbol tables, no uses are
 /// replaced and failure is returned.

mlir/test/lib/Dialect/Test/TestAttrDefs.td

let extraClassDeclaration = [{
     }
   }];
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
     DeclareAttrInterfaceMethods<SubElementAttrInterface,
-        ["replaceImmediateSubAttribute"]>
+        ["replaceImmediateSubElements"]>
   ]> {
-
   let mnemonic = "sub_elements_access";
 
   let parameters = (ins
     "::mlir::Attribute":$first,
     "::mlir::Attribute":$second,
     "::mlir::Attribute":$third
   );
   let hasCustomAssemblyFormat = 1;

mlir/test/lib/Dialect/Test/TestAttributes.cpp

 void TestSubElementsAccessAttr::walkImmediateSubElements(
     llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
     llvm::function_ref<void(mlir::Type)> walkTypesFn) const {
   walkAttrsFn(getFirst());
   walkAttrsFn(getSecond());
   walkAttrsFn(getThird());
 }
 
-SubElementAttrInterface TestSubElementsAccessAttr::replaceImmediateSubAttribute(
-    ArrayRef<std::pair<size_t, Attribute>> replacements) const {
-  Attribute first = getFirst();
-  Attribute second = getSecond();
-  Attribute third = getThird();
-  for (auto &it : replacements) {
-    switch (it.first) {
-    case 0:
-      first = it.second;
-      break;
-    case 1:
-      second = it.second;
-      break;
-    case 2:
-      third = it.second;
-      break;
-    }
-  }
-  return get(getContext(), first, second, third);
+Attribute TestSubElementsAccessAttr::replaceImmediateSubElements(
+    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
+  assert(replAttrs.size() == 3 && "invalid number of replacement attributes");
+  return get(getContext(), replAttrs[0], replAttrs[1], replAttrs[2]);
 }
 
 //===----------------------------------------------------------------------===//
 // TestExtern1DI64ElementsAttr
 //===----------------------------------------------------------------------===//
 
 ArrayRef<uint64_t> TestExtern1DI64ElementsAttr::getElements() const {
   return getHandle().getData()->getData();

mlir/test/lib/Dialect/Test/TestTypes.h

public:
   /// Name/key getter.
   ::llvm::StringRef getName() { return getImpl()->name; }
 
   void walkImmediateSubElements(
       ::llvm::function_ref<void(::mlir::Attribute)> walkAttrsFn,
       ::llvm::function_ref<void(::mlir::Type)> walkTypesFn) const {
     walkTypesFn(getBody());
   }
+  Type replaceImmediateSubElements(llvm::ArrayRef<mlir::Attribute> replAttrs,
+                                   llvm::ArrayRef<mlir::Type> replTypes) const {
+    // TODO: It's not clear how we support replacing sub-elements of mutable
+    // types.
+    return nullptr;
+  }
 };
 
 } // namespace test
 
 #endif // MLIR_TESTTYPES_H