[mlir] Promote the SubElementInterfaces to a core Attribute/Type construct

[rriddle]

Summary

This commit restructures the sub element infrastructure to be a core part
of attributes and types, instead of being relegated to an interface. This
establishes sub element walking/replacement as something "always there",
which makes it easier to rely on for correctness/etc (which various bits of
infrastructure want, such as Symbols).

Attribute/Type now have walk and replace methods directly
accessible, which provide power API for interacting with sub elements. As
part of this, a new AttrTypeWalker class is introduced that supports caching
walked attributes/types, and a friendlier API (see the simplification of symbol
walking in SymbolTable.cpp).

[Mogball]

Beautiful

I assume AttrTypeHandler is now the only way to override the walk behaviour for an attribute or type? (And provide it for unusual attributes and types that might be written in C++..)

mlir/include/mlir/IR/AttrTypeSubElements.h

112

Is this for cycles?

mlir/lib/IR/CMakeLists.txt

30

Alphabetical pls

mlir/lib/IR/SymbolTable.cpp

494

Nice

mlir/unittests/IR/InterfaceTest.cpp

50

I think this test was making just making sure pointer traits for interfaces work. Do we still need the test?

[jpienaar]

Nice graduation (gets us even a step towards introspection ... Well more pieces needed :-)). Have you thought about exposing these C side?

[rriddle]

In D142272#4070783, @Mogball wrote:

Beautiful

I assume AttrTypeHandler is now the only way to override the walk behaviour for an attribute or type? (And provide it for unusual attributes and types that might be written in C++..)

You can just define the immediate walk methods like before. I didn't include anything on it, because I don't think anyone should be doing that in 99.99% of cases.

mlir/include/mlir/IR/AttrTypeSubElements.h

112

Yeah, and for general caching. It removes the need for the user to cache visited attributes/types themselves, and only worry about the actual walk code.

mlir/unittests/IR/InterfaceTest.cpp

50

I just dropped it because it tests the exact same code path as types (given that they use the same infra, and that will likely never change), not really useful by itself IMO.

[rriddle]

In D142272#4071176, @jpienaar wrote:

Have you thought about exposing these C side?

Not particularly. I'm not that involved with the C API, and given that has been need driven (and I don't have a need), I generally leave that to interested parties.

[akuegel]

mlir/lib/IR/AttrTypeSubElements.cpp

31

This seems to prevent iterating through all elements of an ArrayAttr that contains the same attribute more than once.
Example:

xla_entry_computation_result_layout = [

  dense<> : tensor<0xindex>,
  dense<> : tensor<0xindex>
]

This is an attribute that captures the layouts of the results of an XLA computation. Previously, we used walkSubAttrs to iterate over all elements of the array, now it seems there is no comparable walk() method for this.

mlir/unittests/IR/AttributeTest.cpp

435

It would be good if this test could add trueAttr again, to show that it would be skipped during walk().

[rriddle]

mlir/lib/IR/AttrTypeSubElements.cpp

31

You'll want to use the walkImmediateSubElements functions (which might need a slight up lift in API), which guarantee they walk all of the immediately nested things. walkSubElements always walked recursively, which didn't really have any reason to re-walk the same Attribute/Type multiple times.

[akuegel]

mlir/lib/IR/AttrTypeSubElements.cpp

31

I tried (very shortly) to use walkImmediateSubElements, but couldn't make it to work. It also seems easier to instead do a for-loop through the array. I guess for the array use case, we are fine with the for loop. Having a walk() function that not skips elements would make our code slightly easier though, as we either have a single DenseElementsAttr or an array of DenseElementsAttr, and a single walk call would cover both cases.

mlir/unittests/IR/AttributeTest.cpp

435

I have created https://reviews.llvm.org/D142956 for this.

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

 //===-- LLVMAttrDefs.td - LLVM Attributes definition file --*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMIR_ATTRDEFS
 #define LLVMIR_ATTRDEFS
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/Dialect/LLVMIR/LLVMEnums.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
-include "mlir/IR/SubElementInterfaces.td"
 
 // All of the attributes will extend this class.
 class LLVM_Attr<string name, string attrMnemonic,
                 list<Trait> traits = [],
                 string baseCppClass = "::mlir::Attribute">
     : AttrDef<LLVM_Dialect, name, traits, baseCppClass> {
   let mnemonic = attrMnemonic;
 }
def LLVM_DIBasicTypeAttr : LLVM_Attr<"DIBasicType", "di_basic_type",
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DICompileUnitAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DICompileUnitAttr : LLVM_Attr<"DICompileUnit", "di_compile_unit", [
-    SubElementAttrInterface
-  ], "DIScopeAttr"> {
+def LLVM_DICompileUnitAttr : LLVM_Attr<"DICompileUnit", "di_compile_unit",
+                                       /*traits=*/[], "DIScopeAttr"> {
   let parameters = (ins
     LLVM_DILanguageParameter:$sourceLanguage,
     "DIFileAttr":$file,
     "StringAttr":$producer,
     "bool":$isOptimized,
     "DIEmissionKind":$emissionKind
   );
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DICompositeTypeAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DICompositeTypeAttr : LLVM_Attr<"DICompositeType", "di_composite_type", [
-    SubElementAttrInterface
-  ], "DITypeAttr"> {
+def LLVM_DICompositeTypeAttr : LLVM_Attr<"DICompositeType", "di_composite_type",
+                                         /*traits=*/[], "DITypeAttr"> {
   let parameters = (ins
     LLVM_DITagParameter:$tag,
     "StringAttr":$name,
     OptionalParameter<"DIFileAttr">:$file,
     OptionalParameter<"uint32_t">:$line,
     OptionalParameter<"DIScopeAttr">:$scope,
     OptionalParameter<"DITypeAttr">:$baseType,
     OptionalParameter<"DIFlags", "DIFlags::Zero">:$flags,
     OptionalParameter<"uint64_t">:$sizeInBits,
     OptionalParameter<"uint64_t">:$alignInBits,
     OptionalArrayRefParameter<"DINodeAttr">:$elements
   );
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DIDerivedTypeAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DIDerivedTypeAttr : LLVM_Attr<"DIDerivedType", "di_derived_type", [
-    SubElementAttrInterface
-  ], "DITypeAttr"> {
+def LLVM_DIDerivedTypeAttr : LLVM_Attr<"DIDerivedType", "di_derived_type",
+                                       /*traits=*/[], "DITypeAttr"> {
   let parameters = (ins
     LLVM_DITagParameter:$tag,
     OptionalParameter<"StringAttr">:$name,
     "DITypeAttr":$baseType,
     OptionalParameter<"uint64_t">:$sizeInBits,
     OptionalParameter<"uint32_t">:$alignInBits,
     OptionalParameter<"uint64_t">:$offsetInBits
   );
def LLVM_DIFileAttr : LLVM_Attr<"DIFile", "di_file", /*traits=*/[], "DIScopeAttr"> {
   ];
   let assemblyFormat = "`<` $name `in` $directory `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DILexicalBlockAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DILexicalBlockAttr : LLVM_Attr<"DILexicalBlock", "di_lexical_block", [
-    SubElementAttrInterface
-  ], "DIScopeAttr"> {
+def LLVM_DILexicalBlockAttr : LLVM_Attr<"DILexicalBlock", "di_lexical_block",
+                                        /*traits=*/[], "DIScopeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     OptionalParameter<"DIFileAttr">:$file,
     OptionalParameter<"unsigned">:$line,
     OptionalParameter<"unsigned">:$column
   );
   let builders = [
     AttrBuilderWithInferredContext<(ins
       "DIScopeAttr":$scope, "DIFileAttr":$file, "unsigned":$line,
       "unsigned":$column
     ), [{
       return $_get(file.getContext(), scope, file, line, column);
     }]>
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DILexicalBlockFileAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DILexicalBlockFile : LLVM_Attr<"DILexicalBlockFile", "di_lexical_block_file", [
-    SubElementAttrInterface
-  ], "DIScopeAttr"> {
+def LLVM_DILexicalBlockFile : LLVM_Attr<"DILexicalBlockFile", "di_lexical_block_file",
+                                        /*traits=*/[], "DIScopeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     OptionalParameter<"DIFileAttr">:$file,
     "unsigned":$discriminator
   );
   let builders = [
     AttrBuilderWithInferredContext<(ins
       "DIScopeAttr":$scope, "DIFileAttr":$file, "unsigned":$discriminator
     ), [{
       return $_get(file.getContext(), scope, file, discriminator);
     }]>
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DILocalVariableAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DILocalVariableAttr : LLVM_Attr<"DILocalVariable", "di_local_variable", [
-    SubElementAttrInterface
-  ], "DINodeAttr"> {
+def LLVM_DILocalVariableAttr : LLVM_Attr<"DILocalVariable", "di_local_variable",
+                                         /*traits=*/[], "DINodeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     "StringAttr":$name,
     OptionalParameter<"DIFileAttr">:$file,
     OptionalParameter<"unsigned">:$line,
     OptionalParameter<"unsigned">:$arg,
     OptionalParameter<"unsigned">:$alignInBits,
     OptionalParameter<"DITypeAttr">:$type
def LLVM_DILocalVariableAttr : LLVM_Attr<"DILocalVariable", "di_local_variable",
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DISubprogramAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DISubprogramAttr : LLVM_Attr<"DISubprogram", "di_subprogram", [
-    SubElementAttrInterface
-  ], "DIScopeAttr"> {
+def LLVM_DISubprogramAttr : LLVM_Attr<"DISubprogram", "di_subprogram",
+                                      /*traits=*/[], "DIScopeAttr"> {
   let parameters = (ins
     "DICompileUnitAttr":$compileUnit,
     "DIScopeAttr":$scope,
     "StringAttr":$name,
     OptionalParameter<"StringAttr">:$linkageName,
     "DIFileAttr":$file,
     OptionalParameter<"unsigned">:$line,
     OptionalParameter<"unsigned">:$scopeLine,
def LLVM_DISubrangeAttr : LLVM_Attr<"DISubrange", "di_subrange", /*traits=*/[],
   );
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DISubroutineTypeAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_DISubroutineTypeAttr : LLVM_Attr<"DISubroutineType", "di_subroutine_type", [
-    SubElementAttrInterface
-  ], "DITypeAttr"> {
+def LLVM_DISubroutineTypeAttr : LLVM_Attr<"DISubroutineType", "di_subroutine_type",
+                                          /*traits=*/[], "DITypeAttr"> {
   let parameters = (ins
     LLVM_DICallingConventionParameter:$callingConvention,
     OptionalArrayRefParameter<"DITypeAttr">:$types
   );
   let builders = [
     TypeBuilder<(ins "ArrayRef<DITypeAttr>":$types), [{
       return $_get($_ctxt, /*callingConvention=*/0, types);
     }]>
   ];
   let assemblyFormat = "`<` struct(params) `>`";
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // MemoryEffectsAttr
 //===----------------------------------------------------------------------===//
 
-def LLVM_MemoryEffectsAttr : LLVM_Attr<"MemoryEffects", "memory_effects", [
-    SubElementAttrInterface
-  ]> {
+def LLVM_MemoryEffectsAttr : LLVM_Attr<"MemoryEffects", "memory_effects"> {
   let parameters = (ins
     "ModRefInfo":$other,
     "ModRefInfo":$argMem,
     "ModRefInfo":$inaccessibleMem
   );
   let extraClassDeclaration = [{
     bool isReadWrite();
   }];
   let builders = [
     TypeBuilder<(ins "ArrayRef<ModRefInfo>":$memInfoArgs)>
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 #endif // LLVMIR_ATTRDEFS

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

 //===- LLVMTypes.h - MLIR LLVM dialect types --------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the types for the LLVM dialect in MLIR. These MLIR types
 // correspond to the LLVM IR type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
-#include "mlir/IR/SubElementInterfaces.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include <optional>
 
 namespace llvm {
 class ElementCount;
 class TypeSize;
 } // namespace llvm

 /// body, one is not allowed to set another body (however, one can set exactly
 /// the same body).
 ///
 /// Note that the packedness of the struct takes place in uniquing of literal
 /// structs, but does not in uniquing of identified structs.
 class LLVMStructType
     : public Type::TypeBase<LLVMStructType, Type, detail::LLVMStructTypeStorage,
                             DataLayoutTypeInterface::Trait,
-                            SubElementTypeInterface::Trait,
                             TypeTrait::IsMutable> {
 public:
   /// Inherit base constructors.
   using Base::Base;
 
   /// Checks if the given type can be contained in a structure type.
   static bool isValidElementType(Type type);
 

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td

 //===-- LLVMOps.td - LLVM IR dialect op definition file ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVMTYPES_TD
 #define LLVMTYPES_TD
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/AttrTypeBase.td"
-include "mlir/IR/SubElementInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
 
 /// Base class for all LLVM dialect types.
 class LLVMType<string typeName, string typeMnemonic, list<Trait> traits = []>
     : TypeDef<LLVM_Dialect, typeName, traits> {
   let mnemonic = typeMnemonic;
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMArrayType
 //===----------------------------------------------------------------------===//
 
 def LLVMArrayType : LLVMType<"LLVMArray", "array", [
-    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>,
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>]> {
   let summary = "LLVM array type";
   let description = [{
     The `!llvm.array` type represents a fixed-size array of element types.
     It is an aggregate type representing consecutive elements in memory,
     parameterized by the number of elements and the element type.
 
     Example:
 
let extraClassDeclaration = [{
     static bool isValidElementType(Type type);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMFunctionType
 //===----------------------------------------------------------------------===//
 
-def LLVMFunctionType : LLVMType<"LLVMFunction", "func", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+def LLVMFunctionType : LLVMType<"LLVMFunction", "func"> {
   let summary = "LLVM function type";
   let description = [{
     The `!llvm.func` is a function type. It consists of a single return type
     (unlike MLIR which can have multiple), a list of parameter types and can
     optionally be variadic.
 
     Example:
 

 }
 
 //===----------------------------------------------------------------------===//
 // LLVMPointerType
 //===----------------------------------------------------------------------===//
 
 def LLVMPointerType : LLVMType<"LLVMPointer", "ptr", [
     DeclareTypeInterfaceMethods<DataLayoutTypeInterface, [
-      "areCompatible", "verifyEntries"]>,
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+      "areCompatible", "verifyEntries"]>]> {
   let summary = "LLVM pointer type";
   let description = [{
     The `!llvm.ptr` type is an LLVM pointer type. This type typically represents
     a reference to an object in memory. Pointers may be opaque or parameterized
     by the element type. Both opaque and non-opaque pointers are additionally
     parameterized by the address space.
 
     Example:
let extraClassDeclaration = [{
     static bool isValidElementType(Type type);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMFixedVectorType
 //===----------------------------------------------------------------------===//
 
-def LLVMFixedVectorType : LLVMType<"LLVMFixedVector", "vec", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+def LLVMFixedVectorType : LLVMType<"LLVMFixedVector", "vec"> {
   let summary = "LLVM fixed vector type";
   let description = [{
     LLVM dialect scalable vector type, represents a sequence of elements of
     unknown length that is known to be divisible by some constant. These
     elements can be processed as one in SIMD context.
   }];
 
   let parameters = (ins "Type":$elementType, "unsigned":$numElements);
let extraClassDeclaration = [{
     static bool isValidElementType(Type type);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType
 //===----------------------------------------------------------------------===//
 
-def LLVMScalableVectorType : LLVMType<"LLVMScalableVector", "vec", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>]> {
+def LLVMScalableVectorType : LLVMType<"LLVMScalableVector", "vec"> {
   let summary = "LLVM scalable vector type";
   let description = [{
     LLVM dialect scalable vector type, represents a sequence of elements of
     unknown length that is known to be divisible by some constant. These
     elements can be processed as one in SIMD context.
   }];
 
   let parameters = (ins "Type":$elementType, "unsigned":$minNumElements);

mlir/include/mlir/IR/AttrTypeSubElements.h

-//===- SubElementInterfaces.h - Attr and Type SubElements -------*- C++ -*-===//
+//===- AttrTypeSubElements.h - Attr and Type SubElements -------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
-// This file contains interfaces and utilities for querying the sub elements of
-// an attribute or type.
+// This file contains utilities for querying the sub elements of an attribute or
+// type.
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_IR_SUBELEMENTINTERFACES_H
-#define MLIR_IR_SUBELEMENTINTERFACES_H
+#ifndef MLIR_IR_ATTRTYPESUBELEMENTS_H
+#define MLIR_IR_ATTRTYPESUBELEMENTS_H
 
-#include "mlir/IR/Attributes.h"
-#include "mlir/IR/Types.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Visitors.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include <optional>
 
 namespace mlir {
+class Attribute;
+class Type;
+
+//===----------------------------------------------------------------------===//
+/// AttrTypeWalker
+//===----------------------------------------------------------------------===//
+
+/// This class provides a utility for walking attributes/types, and their sub
+/// elements. Multiple walk functions may be registered.
+class AttrTypeWalker {
+public:
+  //===--------------------------------------------------------------------===//
+  // Application
+  //===--------------------------------------------------------------------===//
+
+  /// Walk the given attribute/type, and recursively walk any sub elements.
+  template <WalkOrder Order, typename T>
+  WalkResult walk(T element) {
+    return walkImpl(element, Order);
+  }
+  template <typename T>
+  WalkResult walk(T element) {
+    return walk<WalkOrder::PostOrder, T>(element);
+  }
+
+  //===--------------------------------------------------------------------===//
+  // Registration
+  //===--------------------------------------------------------------------===//
+
+  template <typename T>
+  using WalkFn = std::function<WalkResult(T)>;
+
+  /// Register a walk function for a given attribute or type. A walk function
+  /// must be convertible to any of the following forms(where `T` is a class
+  /// derived from `Type` or `Attribute`:
+  ///
+  ///   * WalkResult(T)
+  ///     - Returns a walk result, which can be used to control the walk
+  ///
+  ///   * void(T)
+  ///     - Returns void, i.e. the walk always continues.
+  ///
+  /// Note: When walking, the mostly recently added walk functions will be
+  ///       invoked first.
+  void addWalk(WalkFn<Attribute> &&fn) {
+    attrWalkFns.emplace_back(std::move(fn));
+  }
+  void addWalk(WalkFn<Type> &&fn) { typeWalkFns.push_back(std::move(fn)); }
+
+  /// Register a replacement function that doesn't match the default signature,
+  /// either because it uses a derived parameter type, or it uses a simplified
+  /// result type.
+  template <typename FnT,
+            typename T = typename llvm::function_traits<
+                std::decay_t<FnT>>::template arg_t<0>,
+            typename BaseT = std::conditional_t<std::is_base_of_v<Attribute, T>,
+                                                Attribute, Type>,
+            typename ResultT = std::invoke_result_t<FnT, T>>
+  std::enable_if_t<!std::is_same_v<T, BaseT> || std::is_same_v<ResultT, void>>
+  addWalk(FnT &&callback) {
+    addWalk([callback = std::forward<FnT>(callback)](BaseT base) -> WalkResult {
+      if (auto derived = dyn_cast<T>(base)) {
+        if constexpr (std::is_convertible_v<ResultT, WalkResult>)
+          return callback(derived);
+        else
+          callback(derived);
+      }
+      return WalkResult::advance();
+    });
+  }
+
+private:
+  WalkResult walkImpl(Attribute attr, WalkOrder order);
+  WalkResult walkImpl(Type type, WalkOrder order);
+
+  /// Internal implementation of the `walk` methods above.
+  template <typename T, typename WalkFns>
+  WalkResult walkImpl(T element, WalkFns &walkFns, WalkOrder order);
+
+  /// Walk the sub elements of the given interface.
+  template <typename T>
+  WalkResult walkSubElements(T interface, WalkOrder order);
+
+  /// The set of walk functions that map sub elements.
+  std::vector<WalkFn<Attribute>> attrWalkFns;
+  std::vector<WalkFn<Type>> typeWalkFns;
+
+  /// The set of visited attributes/types.
+  DenseMap<std::pair<const void *, int>, WalkResult> visitedAttrTypes;
+};
+
 //===----------------------------------------------------------------------===//
 /// AttrTypeReplacer
 //===----------------------------------------------------------------------===//
 
 /// This class provides a utility for replacing attributes/types, and their sub
 /// elements. Multiple replacement functions may be registered.
 class AttrTypeReplacer {
 public:
public:
   ///   * std::optional<std::pair<BaseT, WalkResult>>(T)
   ///     - Similar to the above, but also allows specifying a WalkResult to
   ///       control the replacement of sub elements of a given attribute or
   ///       type. Returning a `skip` result, for example, will not recursively
   ///       process the resultant attribute or type value.
   ///
   /// Note: When replacing, the mostly recently added replacement functions will
   ///       be invoked first.
-  void addReplacement(ReplaceFn<Attribute> fn) {
-    attrReplacementFns.emplace_back(std::move(fn));
-  }
-  void addReplacement(ReplaceFn<Type> fn) {
-    typeReplacementFns.push_back(std::move(fn));
-  }
+  void addReplacement(ReplaceFn<Attribute> fn);
+  void addReplacement(ReplaceFn<Type> fn);
 
   /// Register a replacement function that doesn't match the default signature,
   /// either because it uses a derived parameter type, or it uses a simplified
   /// result type.
   template <typename FnT,
             typename T = typename llvm::function_traits<
                 std::decay_t<FnT>>::template arg_t<0>,
             typename BaseT = std::conditional_t<std::is_base_of_v<Attribute, T>,
addReplacement([callback = std::forward<FnT>(callback)](
         }
       }
       return ReplaceFnResult<BaseT>();
     });
   }
 
 private:
   /// Internal implementation of the `replace` methods above.
-  template <typename InterfaceT, typename ReplaceFns, typename T>
-  T replaceImpl(T element, ReplaceFns &replaceFns, DenseMap<T, T> &map);
+  template <typename T, typename ReplaceFns>
+  T replaceImpl(T element, ReplaceFns &replaceFns);
 
   /// Replace the sub elements of the given interface.
-  template <typename InterfaceT, typename T = typename InterfaceT::ValueType>
-  T replaceSubElements(InterfaceT interface, DenseMap<T, T> &interfaceMap);
+  template <typename T>
+  T replaceSubElements(T interface);
 
   /// The set of replacement functions that map sub elements.
   std::vector<ReplaceFn<Attribute>> attrReplacementFns;
   std::vector<ReplaceFn<Type>> typeReplacementFns;
 
   /// The set of cached mappings for attributes/types.
-  DenseMap<Attribute, Attribute> attrMap;
-  DenseMap<Type, Type> typeMap;
+  DenseMap<const void *, const void *> attrTypeMap;
 };
 
 //===----------------------------------------------------------------------===//
 /// AttrTypeSubElementHandler
 //===----------------------------------------------------------------------===//
 
 /// This class is used by AttrTypeSubElementHandler instances to walking sub
 /// attributes and types.
-class AttrTypeSubElementWalker {
+class AttrTypeImmediateSubElementWalker {
 public:
-  AttrTypeSubElementWalker(function_ref<void(Attribute)> walkAttrsFn,
+  AttrTypeImmediateSubElementWalker(function_ref<void(Attribute)> walkAttrsFn,
                                     function_ref<void(Type)> walkTypesFn)
       : walkAttrsFn(walkAttrsFn), walkTypesFn(walkTypesFn) {}
 
   /// Walk an attribute.
-  void walk(Attribute element) {
-    if (element)
-      walkAttrsFn(element);
-  }
+  void walk(Attribute element);
   /// Walk a type.
-  void walk(Type element) {
-    if (element)
-      walkTypesFn(element);
-  }
+  void walk(Type element);
   /// Walk a range of attributes or types.
   template <typename RangeT>
   void walkRange(RangeT &&elements) {
     for (auto element : elements)
       walk(element);
   }
 
 private:

 ///
 ///   This method should extract any necessary sub elements using the
 ///   provided replacer, or by invoking handlers for sub-types. The new
 ///   post-replacement parameter value should be returned.
 ///
 template <typename T, typename Enable = void>
 struct AttrTypeSubElementHandler {
   /// Default walk implementation that does nothing.
-  static inline void walk(const T &param, AttrTypeSubElementWalker &walker) {}
+  static inline void walk(const T &param,
+                          AttrTypeImmediateSubElementWalker &walker) {}
 
   /// Default replace implementation just forwards the parameter.
   template <typename ParamT>
   static inline decltype(auto) replace(ParamT &&param,
                                        AttrSubElementReplacements &attrRepls,
                                        TypeSubElementReplacements &typeRepls) {
     return std::forward<ParamT>(param);
   }
inline constexpr bool has_sub_attr_or_type_v =
     (!llvm::is_detected<detail::has_default_sub_element_handler_t, Ts>::value ||
      ...);
 
 /// Implementation for derived Attributes and Types.
 template <typename T>
 struct AttrTypeSubElementHandler<
     T, std::enable_if_t<std::is_base_of_v<Attribute, T> ||
                         std::is_base_of_v<Type, T>>> {
-  static void walk(T param, AttrTypeSubElementWalker &walker) {
+  static void walk(T param, AttrTypeImmediateSubElementWalker &walker) {
     walker.walk(param);
   }
   static T replace(T param, AttrSubElementReplacements &attrRepls,
                    TypeSubElementReplacements &typeRepls) {
     if (!param)
       return T();
     if constexpr (std::is_base_of_v<Attribute, T>) {
       return cast<T>(attrRepls.take_front(1)[0]);
     } else {
       return cast<T>(typeRepls.take_front(1)[0]);
     }
   }
 };
-template <>
-struct AttrTypeSubElementHandler<NamedAttribute> {
-  template <typename T>
-  static void walk(T param, AttrTypeSubElementWalker &walker) {
-    walker.walk(param.getName());
-    walker.walk(param.getValue());
-  }
-  template <typename T>
-  static T replace(T param, AttrSubElementReplacements &attrRepls,
-                   TypeSubElementReplacements &typeRepls) {
-    ArrayRef<Attribute> paramRepls = attrRepls.take_front(2);
-    return T(cast<decltype(param.getName())>(paramRepls[0]), paramRepls[1]);
-  }
-};
 /// Implementation for derived ArrayRef.
 template <typename T>
 struct AttrTypeSubElementHandler<ArrayRef<T>,
                                  std::enable_if_t<has_sub_attr_or_type_v<T>>> {
   using EltHandler = AttrTypeSubElementHandler<T>;
 
-  static void walk(ArrayRef<T> param, AttrTypeSubElementWalker &walker) {
+  static void walk(ArrayRef<T> param,
+                   AttrTypeImmediateSubElementWalker &walker) {
     for (const T &subElement : param)
       EltHandler::walk(subElement, walker);
   }
   static auto replace(ArrayRef<T> param, AttrSubElementReplacements &attrRepls,
                       TypeSubElementReplacements &typeRepls) {
     // Normal attributes/types can extract using the replacer directly.
     if constexpr (std::is_base_of_v<Attribute, T> &&
-                  sizeof(T) == sizeof(Attribute)) {
+                  sizeof(T) == sizeof(void *)) {
       ArrayRef<Attribute> attrs = attrRepls.take_front(param.size());
       return ArrayRef<T>((const T *)attrs.data(), attrs.size());
     } else if constexpr (std::is_base_of_v<Type, T> &&
-                         sizeof(T) == sizeof(Type)) {
+                         sizeof(T) == sizeof(void *)) {
       ArrayRef<Type> types = typeRepls.take_front(param.size());
       return ArrayRef<T>((const T *)types.data(), types.size());
     } else {
       // Otherwise, we need to allocate storage for the new elements.
       SmallVector<T> newElements;
       for (const T &element : param)
         newElements.emplace_back(
             EltHandler::replace(element, attrRepls, typeRepls));
       return newElements;
     }
   }
 };
 /// Implementation for Tuple.
 template <typename... Ts>
 struct AttrTypeSubElementHandler<
     std::tuple<Ts...>, std::enable_if_t<has_sub_attr_or_type_v<Ts...>>> {
   static void walk(const std::tuple<Ts...> &param,
-                   AttrTypeSubElementWalker &walker) {
+                   AttrTypeImmediateSubElementWalker &walker) {
     std::apply(
         [&](const Ts &...params) {
           (AttrTypeSubElementHandler<Ts>::walk(params, walker), ...);
         },
         param);
   }
   static auto replace(const std::tuple<Ts...> &param,
                       AttrSubElementReplacements &attrRepls,

 
 namespace detail {
 template <typename T>
 struct is_tuple : public std::false_type {};
 template <typename... Ts>
 struct is_tuple<std::tuple<Ts...>> : public std::true_type {};
 template <typename T, typename... Ts>
 using has_get_method = decltype(T::get(std::declval<Ts>()...));
+template <typename T, typename... Ts>
+using has_get_as_key = decltype(std::declval<T>().getAsKey());
 
 /// This function provides the underlying implementation for the
 /// SubElementInterface walk method, using the key type of the derived
 /// attribute/type to interact with the individual parameters.
 template <typename T>
 void walkImmediateSubElementsImpl(T derived,
                                   function_ref<void(Attribute)> walkAttrsFn,
                                   function_ref<void(Type)> walkTypesFn) {
-  auto key = static_cast<typename T::ImplType *>(derived.getImpl())->getAsKey();
+  using ImplT = typename T::ImplType;
+  if constexpr (llvm::is_detected<has_get_as_key, ImplT>::value) {
+    auto key = static_cast<ImplT *>(derived.getImpl())->getAsKey();
 
     // If we don't have any sub-elements, there is nothing to do.
     if constexpr (!has_sub_attr_or_type_v<decltype(key)>) {
       return;
     } else {
-    AttrTypeSubElementWalker walker(walkAttrsFn, walkTypesFn);
+      AttrTypeImmediateSubElementWalker walker(walkAttrsFn, walkTypesFn);
       AttrTypeSubElementHandler<decltype(key)>::walk(key, walker);
     }
   }
+}
 
 /// This function invokes the proper `get` method for  a type `T` with the given
 /// values.
 template <typename T, typename... Ts>
-T constructSubElementReplacement(MLIRContext *ctx, Ts &&...params) {
+auto constructSubElementReplacement(MLIRContext *ctx, Ts &&...params) {
   // Prefer a direct `get` method if one exists.
   if constexpr (llvm::is_detected<has_get_method, T, Ts...>::value) {
     (void)ctx;
     return T::get(std::forward<Ts>(params)...);
   } else if constexpr (llvm::is_detected<has_get_method, T, MLIRContext *,
                                          Ts...>::value) {
     return T::get(ctx, std::forward<Ts>(params)...);
   } else {
     // Otherwise, pass to the base get.
     return T::Base::get(ctx, std::forward<Ts>(params)...);
   }
 }
 
 /// This function provides the underlying implementation for the
 /// SubElementInterface replace method, using the key type of the derived
 /// attribute/type to interact with the individual parameters.
 template <typename T>
-T replaceImmediateSubElementsImpl(T derived, ArrayRef<Attribute> &replAttrs,
+auto replaceImmediateSubElementsImpl(T derived, ArrayRef<Attribute> &replAttrs,
                                      ArrayRef<Type> &replTypes) {
-  auto key = static_cast<typename T::ImplType *>(derived.getImpl())->getAsKey();
+  using ImplT = typename T::ImplType;
+  if constexpr (llvm::is_detected<has_get_as_key, ImplT>::value) {
+    auto key = static_cast<ImplT *>(derived.getImpl())->getAsKey();
 
     // If we don't have any sub-elements, we can just return the original.
     if constexpr (!has_sub_attr_or_type_v<decltype(key)>) {
       return derived;
 
       // Otherwise, we need to replace any necessary sub-elements.
     } else {
       AttrSubElementReplacements attrRepls(replAttrs);
       TypeSubElementReplacements typeRepls(replTypes);
       auto newKey = AttrTypeSubElementHandler<decltype(key)>::replace(
           key, attrRepls, typeRepls);
       if constexpr (is_tuple<decltype(key)>::value) {
         return std::apply(
             [&](auto &&...params) {
               return constructSubElementReplacement<T>(
                   derived.getContext(),
                   std::forward<decltype(params)>(params)...);
             },
             newKey);
       } else {
         return constructSubElementReplacement<T>(derived.getContext(), newKey);
       }
     }
+  } else {
+    return derived;
+  }
 }
 } // namespace detail
 } // namespace mlir
 
-/// Include the definitions of the sub element interfaces.
-#include "mlir/IR/SubElementAttrInterfaces.h.inc"
-#include "mlir/IR/SubElementTypeInterfaces.h.inc"
-
-#endif // MLIR_IR_SUBELEMENTINTERFACES_H
+#endif // MLIR_IR_ATTRTYPESUBELEMENTS_H

mlir/include/mlir/IR/AttributeSupport.h

 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/StorageUniquerSupport.h"
 #include "mlir/IR/Types.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/Twine.h"
 
 namespace mlir {
-class MLIRContext;
-class Type;
-
 //===----------------------------------------------------------------------===//
 // AbstractAttribute
 //===----------------------------------------------------------------------===//
 
 /// This class contains all of the static information common to all instances of
 /// a registered Attribute.
 class AbstractAttribute {
 public:
   using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
+  using WalkImmediateSubElementsFn = function_ref<void(
+      Attribute, function_ref<void(Attribute)>, function_ref<void(Type)>)>;
+  using ReplaceImmediateSubElementsFn =
+      function_ref<Attribute(Attribute, ArrayRef<Attribute>, ArrayRef<Type>)>;
 
   /// Look up the specified abstract attribute in the MLIRContext and return a
   /// reference to it.
   static const AbstractAttribute &lookup(TypeID typeID, MLIRContext *context);
 
   /// This method is used by Dialect objects when they register the list of
   /// attributes they contain.
   template <typename T>
   static AbstractAttribute get(Dialect &dialect) {
     return AbstractAttribute(dialect, T::getInterfaceMap(), T::getHasTraitFn(),
+                             T::getWalkImmediateSubElementsFn(),
+                             T::getReplaceImmediateSubElementsFn(),
                              T::getTypeID());
   }
 
   /// This method is used by Dialect objects to register attributes with
   /// custom TypeIDs.
   /// The use of this method is in general discouraged in favor of
   /// 'get<CustomAttribute>(dialect)'.
-  static AbstractAttribute get(Dialect &dialect,
-                               detail::InterfaceMap &&interfaceMap,
-                               HasTraitFn &&hasTrait, TypeID typeID) {
+  static AbstractAttribute
+  get(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
+      HasTraitFn &&hasTrait,
+      WalkImmediateSubElementsFn walkImmediateSubElementsFn,
+      ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn,
+      TypeID typeID) {
     return AbstractAttribute(dialect, std::move(interfaceMap),
-                             std::move(hasTrait), typeID);
+                             std::move(hasTrait), walkImmediateSubElementsFn,
+                             replaceImmediateSubElementsFn, typeID);
   }
 
   /// Return the dialect this attribute was registered to.
   Dialect &getDialect() const { return const_cast<Dialect &>(dialect); }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this attribute, null otherwise. This should not be used
   /// directly.
public:
   template <template <typename T> class Trait>
   bool hasTrait() const {
     return hasTraitFn(TypeID::get<Trait>());
   }
 
   /// Returns true if the attribute has a particular trait.
   bool hasTrait(TypeID traitID) const { return hasTraitFn(traitID); }
 
+  /// Walk the immediate sub-elements of this attribute.
+  void walkImmediateSubElements(Attribute attr,
+                                function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
+
+  /// Replace the immediate sub-elements of this attribute.
+  Attribute replaceImmediateSubElements(Attribute attr,
+                                        ArrayRef<Attribute> replAttrs,
+                                        ArrayRef<Type> replTypes) const;
+
   /// Return the unique identifier representing the concrete attribute class.
   TypeID getTypeID() const { return typeID; }
 
 private:
   AbstractAttribute(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
-                    HasTraitFn &&hasTrait, TypeID typeID)
+                    HasTraitFn &&hasTraitFn,
+                    WalkImmediateSubElementsFn walkImmediateSubElementsFn,
+                    ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn,
+                    TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
-        hasTraitFn(std::move(hasTrait)), typeID(typeID) {}
+        hasTraitFn(std::move(hasTraitFn)),
+        walkImmediateSubElementsFn(walkImmediateSubElementsFn),
+        replaceImmediateSubElementsFn(replaceImmediateSubElementsFn),
+        typeID(typeID) {}
 
   /// Give StorageUserBase access to the mutable lookup.
   template <typename ConcreteT, typename BaseT, typename StorageT,
             typename UniquerT, template <typename T> class... Traits>
   friend class detail::StorageUserBase;
 
   /// Look up the specified abstract attribute in the MLIRContext and return a
   /// (mutable) pointer to it. Return a null pointer if the attribute could not
   /// be found in the context.
   static AbstractAttribute *lookupMutable(TypeID typeID, MLIRContext *context);
 
   /// This is the dialect that this attribute was registered to.
   const Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this attribute.
   detail::InterfaceMap interfaceMap;
 
   /// Function to check if the attribute has a particular trait.
   HasTraitFn hasTraitFn;
 
+  /// Function to walk the immediate sub-elements of this attribute.
+  WalkImmediateSubElementsFn walkImmediateSubElementsFn;
+
+  /// Function to replace the immediate sub-elements of this attribute.
+  ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn;
+
   /// The unique identifier of the derived Attribute class.
   const TypeID typeID;
 };
 
 //===----------------------------------------------------------------------===//
 // AttributeStorage
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Attributes.h

bool hasTrait() {
     return getAbstractAttribute().hasTrait<Trait>();
   }
 
   /// Return the abstract descriptor for this attribute.
   const AbstractTy &getAbstractAttribute() const {
     return impl->getAbstractAttribute();
   }
 
+  /// Walk all of the immediately nested sub-attributes and sub-types. This
+  /// method does not recurse into sub elements.
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const {
+    getAbstractAttribute().walkImmediateSubElements(*this, walkAttrsFn,
+                                                    walkTypesFn);
+  }
+
+  /// 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.
+  auto replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const {
+    return getAbstractAttribute().replaceImmediateSubElements(*this, replAttrs,
+                                                              replTypes);
+  }
+
+  /// Walk this attribute and all attibutes/types nested within using the
+  /// provided walk functions. See `AttrTypeWalker` for information on the
+  /// supported walk function types.
+  template <WalkOrder Order = WalkOrder::PostOrder, typename... WalkFns>
+  auto walk(WalkFns &&...walkFns) {
+    AttrTypeWalker walker;
+    (walker.addWalk(std::forward<WalkFns>(walkFns)), ...);
+    return walker.walk<Order>(*this);
+  }
+
+  /// Recursively replace all of the nested sub-attributes and sub-types using
+  /// the provided map functions. Returns nullptr in the case of failure. See
+  /// `AttrTypeReplacer` for information on the support replacement function
+  /// types.
+  template <typename... ReplacementFns>
+  auto replace(ReplacementFns &&...replacementFns) {
+    AttrTypeReplacer replacer;
+    (replacer.addReplacement(std::forward<ReplacementFns>(replacementFns)),
+     ...);
+    return replacer.replace(*this);
+  }
+
   /// Return the internal Attribute implementation.
   ImplType *getImpl() const { return impl; }
 
 protected:
   ImplType *impl{nullptr};
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Attribute attr) {
private:
   Attribute value;
 };
 
 inline ::llvm::hash_code hash_value(const NamedAttribute &arg) {
   using AttrPairT = std::pair<Attribute, Attribute>;
   return DenseMapInfo<AttrPairT>::getHashValue(AttrPairT(arg.name, arg.value));
 }
 
+/// Allow walking and replacing the subelements of a NamedAttribute.
+template <>
+struct AttrTypeSubElementHandler<NamedAttribute> {
+  template <typename T>
+  static void walk(T param, AttrTypeImmediateSubElementWalker &walker) {
+    walker.walk(param.getName());
+    walker.walk(param.getValue());
+  }
+  template <typename T>
+  static T replace(T param, AttrSubElementReplacements &attrRepls,
+                   TypeSubElementReplacements &typeRepls) {
+    ArrayRef<Attribute> paramRepls = attrRepls.take_front(2);
+    return T(cast<decltype(param.getName())>(paramRepls[0]), paramRepls[1]);
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // AttributeTraitBase
 //===----------------------------------------------------------------------===//
 
 namespace AttributeTrait {
 /// This class represents the base of an attribute trait.
 template <typename ConcreteType, template <typename> class TraitType>
 using TraitBase = detail::StorageUserTraitBase<ConcreteType, TraitType>;

mlir/include/mlir/IR/BuiltinAttributes.h

 //===- BuiltinAttributes.h - MLIR Builtin Attribute Classes -----*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINATTRIBUTES_H
 #define MLIR_IR_BUILTINATTRIBUTES_H
 
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
-#include "mlir/IR/SubElementInterfaces.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Sequence.h"
 #include <complex>
 #include <optional>
 
 namespace mlir {
 class AffineMap;
 class AsmResourceBlob;

mlir/include/mlir/IR/BuiltinAttributes.td

 
 #ifndef BUILTIN_ATTRIBUTES
 #define BUILTIN_ATTRIBUTES
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinDialect.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
-include "mlir/IR/SubElementInterfaces.td"
 
 // TODO: Currently the attributes defined in this file are prefixed with
 // `Builtin_`.  This is to differentiate the attributes here with the ones in
 // OpBase.td. We should remove the definitions in OpBase.td, and repoint users
 // to this file instead.
 
 // Base class for Builtin dialect attributes.
 class Builtin_Attr<string name, list<Trait> traits = [],
def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap", [
   }];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_ArrayAttr : Builtin_Attr<"Array", [
-    SubElementAttrInterface
-  ]> {
+def Builtin_ArrayAttr : Builtin_Attr<"Array"> {
   let summary = "A collection of other Attribute values";
   let description = [{
     Syntax:
 
     ```
     array-attribute ::= `[` (attribute-value (`,` attribute-value)*)? `]`
     ```
 
def Builtin_DenseResourceElementsAttr : Builtin_Attr<"DenseResourceElements", [
   }];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary", [
-    SubElementAttrInterface
-  ]> {
+def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary"> {
   let summary = "An dictionary of named Attribute values";
   let description = [{
     Syntax:
 
     ```
     dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? `}`
     ```
 
def Builtin_StringAttr : Builtin_Attr<"String", [TypedAttrInterface]> {
   let genStorageClass = 0;
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolRefAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef", [
-    SubElementAttrInterface
-  ]> {
+def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef"> {
   let summary = "An Attribute containing a symbolic reference to an Operation";
   let description = [{
     Syntax:
 
     ```
     symbol-ref-attribute ::= symbol-ref-id (`::` symbol-ref-id)*
     ```
 
     A symbol reference attribute is a literal attribute that represents a named
     reference to an operation that is nested within an operation with the
     `OpTrait::SymbolTable` trait. As such, this reference is given meaning by
     the nearest parent operation containing the `OpTrait::SymbolTable` trait. It
     may optionally contain a set of nested references that further resolve to a
     symbol nested within a different symbol table.
 
-    This attribute can only be held internally by
-    [array attributes](#array-attribute),
-    [dictionary attributes](#dictionary-attribute)(including the top-level
-    operation attribute dictionary) as well as attributes exposing it via
-    the `SubElementAttrInterface` interface. Symbol reference attributes
-    nested in types are currently not supported.
-
     **Rationale:** Identifying accesses to global data is critical to
     enabling efficient multi-threaded compilation. Restricting global
     data access to occur through symbols and limiting the places that can
     legally hold a symbol reference simplifies reasoning about these data
     accesses.
 
     See [`Symbols And SymbolTables`](../SymbolsAndSymbolTables.md) for more
     information.
def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef"> {
   }];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TypeAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_TypeAttr : Builtin_Attr<"Type", [
-    SubElementAttrInterface
-  ]> {
+def Builtin_TypeAttr : Builtin_Attr<"Type"> {
   let summary = "An Attribute containing a Type";
   let description = [{
     Syntax:
 
     ```
     type-attribute ::= type
     ```
 

mlir/include/mlir/IR/BuiltinLocationAttributes.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUILTIN_LOCATION_ATTRIBUTES_TD
 #define BUILTIN_LOCATION_ATTRIBUTES_TD
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinDialect.td"
-include "mlir/IR/SubElementInterfaces.td"
 
 // Base class for Builtin dialect location attributes.
 class Builtin_LocationAttr<string name, list<Trait> traits = []>
     : AttrDef<Builtin_Dialect, name, traits, "::mlir::LocationAttr"> {
   let cppClassName = name;
   let mnemonic = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
-def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc", [
-    SubElementAttrInterface
-  ]> {
+def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc"> {
   let summary = "A callsite source location";
   let description = [{
     Syntax:
 
     ```
     callsite-location ::= `callsite` `(` location `at` location `)`
     ```
 
def FileLineColLoc : Builtin_LocationAttr<"FileLineColLoc"> {
   ];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
-def FusedLoc : Builtin_LocationAttr<"FusedLoc", [
-    SubElementAttrInterface
-  ]> {
+def FusedLoc : Builtin_LocationAttr<"FusedLoc"> {
   let summary = "A tuple of other source locations";
   let description = [{
     Syntax:
 
     ```
     fused-location ::= `fused` fusion-metadata? `[` location (location `,`)* `]`
     fusion-metadata ::= `<` attribute-value `>`
     ```
let extraClassDeclaration = [{
     }
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // NameLoc
 //===----------------------------------------------------------------------===//
 
-def NameLoc : Builtin_LocationAttr<"NameLoc", [
-    SubElementAttrInterface
-  ]> {
+def NameLoc : Builtin_LocationAttr<"NameLoc"> {
   let summary = "A named source location";
   let description = [{
     Syntax:
 
     ```
     name-location ::= string-literal (`(` location `)`)?
     ```
 
def NameLoc : Builtin_LocationAttr<"NameLoc"> {
   ];
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // OpaqueLoc
 //===----------------------------------------------------------------------===//
 
-def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc", [
-    SubElementAttrInterface
-  ]> {
+def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc"> {
   let summary = "An opaque source location";
   let description = [{
     An instance of this location essentially contains a pointer to some data
     structure that is external to MLIR and an optional location that can be used
     if the first one is not suitable. Since it contains an external structure,
     only the optional location is used during serialization.
   }];
   let parameters = (ins "uintptr_t":$underlyingLocation,

mlir/include/mlir/IR/BuiltinTypes.h

 //===- BuiltinTypes.h - MLIR Builtin Type Classes ---------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINTYPES_H
 #define MLIR_IR_BUILTINTYPES_H
 
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/IR/SubElementInterfaces.h"
 
 namespace llvm {
 class BitVector;
 struct fltSemantics;
 } // namespace llvm
 
 //===----------------------------------------------------------------------===//
 // Tablegen Interface Declarations

mlir/include/mlir/IR/BuiltinTypes.td

 //===----------------------------------------------------------------------===//
 
 #ifndef BUILTIN_TYPES
 #define BUILTIN_TYPES
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinDialect.td"
 include "mlir/IR/BuiltinTypeInterfaces.td"
-include "mlir/IR/SubElementInterfaces.td"
 
 // TODO: Currently the types defined in this file are prefixed with `Builtin_`.
 // This is to differentiate the types here with the ones in OpBase.td. We should
 // remove the definitions in OpBase.td, and repoint users to this file instead.
 
 // Base class for Builtin dialect types.
 class Builtin_Type<string name, list<Trait> traits = [],
                    string baseCppClass = "::mlir::Type">

 def Builtin_Float128 : Builtin_FloatType<"Float128"> {
   let summary = "128-bit floating-point type";
 }
 
 //===----------------------------------------------------------------------===//
 // FunctionType
 //===----------------------------------------------------------------------===//
 
-def Builtin_Function : Builtin_Type<"Function", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>
-  ]> {
+def Builtin_Function : Builtin_Type<"Function"> {
   let summary = "Map from a list of inputs to a list of results";
   let description = [{
     Syntax:
 
     ```
     // Function types may have multiple results.
     function-result-type ::= type-list-parens | non-function-type
     function-type ::= type-list-parens `->` function-result-type
def Builtin_Integer : Builtin_Type<"Integer"> {
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // MemRefType
 //===----------------------------------------------------------------------===//
 
 def Builtin_MemRef : Builtin_Type<"MemRef", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>, ShapedTypeInterface
+    ShapedTypeInterface
   ], "BaseMemRefType"> {
   let summary = "Shaped reference to a region of memory";
   let description = [{
     Syntax:
 
     ```
     memref-type ::= `memref` `<` dimension-list-ranked type
                     (`,` layout-specification)? (`,` memory-space)? `>`
def Builtin_Opaque : Builtin_Type<"Opaque"> {
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // RankedTensorType
 //===----------------------------------------------------------------------===//
 
 def Builtin_RankedTensor : Builtin_Type<"RankedTensor", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>, ShapedTypeInterface
+    ShapedTypeInterface
   ], "TensorType"> {
   let summary = "Multi-dimensional array with a fixed number of dimensions";
   let description = [{
     Syntax:
 
     ```
     tensor-type ::= `tensor` `<` dimension-list type (`,` encoding)? `>`
     dimension-list ::= (dimension `x`)*
def Builtin_RankedTensor : Builtin_Type<"RankedTensor", [
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TupleType
 //===----------------------------------------------------------------------===//
 
-def Builtin_Tuple : Builtin_Type<"Tuple", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>
-  ]> {
+def Builtin_Tuple : Builtin_Type<"Tuple"> {
   let summary = "Fixed-sized collection of other types";
   let description = [{
     Syntax:
 
     ```
     tuple-type ::= `tuple` `<` (type ( `,` type)*)? `>`
     ```
 
def Builtin_Tuple : Builtin_Type<"Tuple"> {
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>, ShapedTypeInterface
+    ShapedTypeInterface
   ], "BaseMemRefType"> {
   let summary = "Shaped reference, with unknown rank, to a region of memory";
   let description = [{
     Syntax:
 
     ```
     unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>`
     memory-space ::= attribute-value
def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", [
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//
 
 def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>, ShapedTypeInterface
+    ShapedTypeInterface
   ], "TensorType"> {
   let summary = "Multi-dimensional array with unknown dimensions";
   let description = [{
     Syntax:
 
     ```
     tensor-type ::= `tensor` `<` `*` `x` type `>`
     ```
def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", [
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // VectorType
 //===----------------------------------------------------------------------===//
 
-def Builtin_Vector : Builtin_Type<"Vector", [
-    DeclareTypeInterfaceMethods<SubElementTypeInterface>, ShapedTypeInterface
-  ], "Type"> {
+def Builtin_Vector : Builtin_Type<"Vector", [ShapedTypeInterface], "Type"> {
   let summary = "Multi-dimensional SIMD vector type";
   let description = [{
     Syntax:
 
     ```
     vector-type ::= `vector` `<` vector-dim-list vector-element-type `>`
     vector-element-type ::= float-type | integer-type | index-type
     vector-dim-list := (static-dim-list `x`)? (`[` static-dim-list `]` `x`)?

mlir/include/mlir/IR/CMakeLists.txt

 add_public_tablegen_target(MLIRBuiltinTypeInterfacesIncGen)
 
 set(LLVM_TARGET_DEFINITIONS FunctionInterfaces.td)
 mlir_tablegen(FunctionOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(FunctionOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRFunctionInterfacesIncGen)
 add_dependencies(mlir-generic-headers MLIRFunctionInterfacesIncGen)
 
-set(LLVM_TARGET_DEFINITIONS SubElementInterfaces.td)
-mlir_tablegen(SubElementAttrInterfaces.h.inc -gen-attr-interface-decls)
-mlir_tablegen(SubElementAttrInterfaces.cpp.inc -gen-attr-interface-defs)
-mlir_tablegen(SubElementTypeInterfaces.h.inc -gen-type-interface-decls)
-mlir_tablegen(SubElementTypeInterfaces.cpp.inc -gen-type-interface-defs)
-add_public_tablegen_target(MLIRSubElementInterfacesIncGen)
-
 set(LLVM_TARGET_DEFINITIONS TensorEncoding.td)
 mlir_tablegen(TensorEncInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(TensorEncInterfaces.cpp.inc -gen-attr-interface-defs)
 add_public_tablegen_target(MLIRTensorEncodingIncGen)
 
 add_mlir_doc(BuiltinAttributes BuiltinAttributes Dialects/ -gen-attrdef-doc)
 add_mlir_doc(BuiltinLocationAttributes BuiltinLocationAttributes Dialects/ -gen-attrdef-doc)
 add_mlir_doc(BuiltinOps BuiltinOps Dialects/ -gen-op-doc)
 add_mlir_doc(BuiltinTypes BuiltinTypes Dialects/ -gen-typedef-doc)
 add_mlir_doc(BuiltinTypes BuiltinTypeInterfaces Dialects/ -gen-type-interface-docs)

mlir/include/mlir/IR/Location.h

 // location position information.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_LOCATION_H
 #define MLIR_IR_LOCATION_H
 
 #include "mlir/IR/Attributes.h"
-#include "mlir/IR/SubElementInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 
 class Location;
 class WalkResult;
 
 //===----------------------------------------------------------------------===//

 /// an object. The corresponding MLIR location is set to UnknownLoc.
 template <typename T>
 inline OpaqueLoc OpaqueLoc::get(T underlyingLocation, MLIRContext *context) {
   return get(reinterpret_cast<uintptr_t>(underlyingLocation), TypeID::get<T>(),
              UnknownLoc::get(context));
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementInterfaces
+// SubElements
 //===----------------------------------------------------------------------===//
 
 /// Enable locations to be introspected as sub-elements.
 template <>
 struct AttrTypeSubElementHandler<Location> {
-  static void walk(Location param, AttrTypeSubElementWalker &walker) {
+  static void walk(Location param, AttrTypeImmediateSubElementWalker &walker) {
     walker.walk(param);
   }
   static Location replace(Location param, AttrSubElementReplacements &attrRepls,
                           TypeSubElementReplacements &typeRepls) {
     return cast<LocationAttr>(attrRepls.take_front(1)[0]);
   }
 };
 

mlir/include/mlir/IR/StorageUniquerSupport.h

 //===- StorageUniquerSupport.h - MLIR Storage Uniquer Utilities -*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines utility classes for interfacing with StorageUniquer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_STORAGEUNIQUERSUPPORT_H
 #define MLIR_IR_STORAGEUNIQUERSUPPORT_H
 
+#include "mlir/IR/AttrTypeSubElements.h"
 #include "mlir/Support/InterfaceSupport.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/StorageUniquer.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/FunctionExtras.h"
 
 namespace mlir {
 class InFlightDiagnostic;
public:
   /// Returns the function that returns true if the given Trait ID matches the
   /// IDs of any of the traits defined by the storage user.
   static HasTraitFn getHasTraitFn() {
     return [](TypeID id) {
       return storage_user_base_impl::hasTrait<Traits...>(id);
     };
   }
 
+  /// Walk all of the immediately nested sub-attributes and sub-types. This
+  /// method does not recurse into sub elements.
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const {
+    ::mlir::detail::walkImmediateSubElementsImpl(
+        *static_cast<const ConcreteT *>(this), walkAttrsFn, walkTypesFn);
+  }
+
+  /// 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.
+  ///
+  /// Note that replacing the sub-elements of mutable types or attributes is
+  /// not currently supported by the interface. If an implementing type or
+  /// attribute is mutable, it should return `nullptr` if it has no mechanism
+  /// for replacing sub elements.
+  auto replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const {
+    return ::mlir::detail::replaceImmediateSubElementsImpl(
+        *static_cast<const ConcreteT *>(this), replAttrs, replTypes);
+  }
+
+  /// Returns a function that walks immediate sub elements of a given instance
+  /// of the storage user.
+  static auto getWalkImmediateSubElementsFn() {
+    return [](auto instance, function_ref<void(Attribute)> walkAttrsFn,
+              function_ref<void(Type)> walkTypesFn) {
+      cast<ConcreteT>(instance).walkImmediateSubElements(walkAttrsFn,
+                                                         walkTypesFn);
+    };
+  }
+
+  /// Returns a function that replaces immediate sub elements of a given
+  /// instance of the storage user.
+  static auto getReplaceImmediateSubElementsFn() {
+    return [](auto instance, ArrayRef<Attribute> replAttrs,
+              ArrayRef<Type> replTypes) {
+      return cast<ConcreteT>(instance).replaceImmediateSubElements(replAttrs,
+                                                                   replTypes);
+    };
+  }
+
   /// Attach the given models as implementations of the corresponding interfaces
   /// for the concrete storage user class. The type must be registered with the
   /// context, i.e. the dialect to which the type belongs must be loaded. The
   /// call will abort otherwise.
   template <typename... IfaceModels>
   static void attachInterface(MLIRContext &context) {
     typename ConcreteT::AbstractTy *abstract =
         ConcreteT::AbstractTy::lookupMutable(TypeID::get<ConcreteT>(),

mlir/include/mlir/IR/SubElementInterfaces.td

-//===-- SubElementInterfaces.td - Sub-Element Interfaces ---*- tablegen -*-===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-//
-// This file contains a set of interfaces that can be used to interface with
-// sub-elements, e.g. held attributes and types, of a composite attribute or
-// type.
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_IR_SUBELEMENTINTERFACES_TD_
-#define MLIR_IR_SUBELEMENTINTERFACES_TD_
-
-include "mlir/IR/OpBase.td"
-
-//===----------------------------------------------------------------------===//
-// SubElementInterfaceBase
-//===----------------------------------------------------------------------===//
-
-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),
-      /*methodBody=*/[{}], /*defaultImplementation=*/[{
-        ::mlir::detail::walkImmediateSubElementsImpl(
-          }] # derivedValue # [{, walkAttrsFn, walkTypesFn);
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/[{
-        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.
-
-        Note that replacing the sub-elements of mutable types or attributes is
-        not currently supported by the interface. If an implementing type or
-        attribute is mutable, it should return `nullptr` if it has no mechanism
-        for replacing sub elements.
-      }], attrOrType, "replaceImmediateSubElements",
-      (ins "::llvm::ArrayRef<::mlir::Attribute>":$replAttrs,
-           "::llvm::ArrayRef<::mlir::Type>":$replTypes),
-      /*methodBody=*/[{}], /*defaultImplementation=*/[{
-        return ::mlir::detail::replaceImmediateSubElementsImpl(
-           }] # derivedValue # [{, replAttrs, replTypes);
-      }]>,
-  ];
-
-  code extraClassDeclaration = [{
-    /// 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. See
-    /// `AttrTypeReplacer` for information on the support replacement function types.
-    template <typename... ReplacementFns>
-    }] # attrOrType # [{ replaceSubElements(ReplacementFns &&... replacementFns) {
-      AttrTypeReplacer replacer;
-      (replacer.addReplacement(std::forward<ReplacementFns>(replacementFns)), ...);
-      return replacer.replace(*this);
-    }
-  }];
-  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. See
-    /// `AttrTypeReplacer` for information on the support replacement function types.
-    template <typename... ReplacementFns>
-    }] # attrOrType # [{ replaceSubElements(ReplacementFns &&... replacementFns) {
-      AttrTypeReplacer replacer;
-      (replacer.addReplacement(std::forward<ReplacementFns>(replacementFns)), ...);
-      return replacer.replace(}] # derivedValue # [{);
-    }
-  }];
-  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);
-    }
-  }];
-}
-
-//===----------------------------------------------------------------------===//
-// SubElementAttrInterface
-//===----------------------------------------------------------------------===//
-
-def SubElementAttrInterface
-    : AttrInterface<"SubElementAttrInterface">,
-      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.
-
-    To support the introspection of custom parameters that hold sub-elements,
-    a specialization of the `AttrTypeSubElementHandler` class must be provided.
-  }];
-}
-
-//===----------------------------------------------------------------------===//
-// SubElementTypeInterface
-//===----------------------------------------------------------------------===//
-
-def SubElementTypeInterface
-    : TypeInterface<"SubElementTypeInterface">,
-      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.
-
-    To support the introspection of custom parameters that hold sub-elements,
-    a specialization of the `AttrTypeSubElementHandler` class must be provided.
-  }];
-}
-
-#endif // MLIR_IR_SUBELEMENTINTERFACES_TD_

mlir/include/mlir/IR/TypeRange.h

 
 template <typename RangeT>
 inline bool operator==(ArrayRef<Type> lhs, const ValueTypeRange<RangeT> &rhs) {
   return lhs.size() == static_cast<size_t>(llvm::size(rhs)) &&
          std::equal(lhs.begin(), lhs.end(), rhs.begin());
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementInterfaces
+// SubElements
 //===----------------------------------------------------------------------===//
 
 /// Enable TypeRange to be introspected for sub-elements.
 template <>
 struct AttrTypeSubElementHandler<TypeRange> {
-  static void walk(TypeRange param, AttrTypeSubElementWalker &walker) {
+  static void walk(TypeRange param, AttrTypeImmediateSubElementWalker &walker) {
     walker.walkRange(param);
   }
   static TypeRange replace(TypeRange param,
                            AttrSubElementReplacements &attrRepls,
                            TypeSubElementReplacements &typeRepls) {
     return typeRepls.take_front(param.size());
   }
 };

mlir/include/mlir/IR/TypeSupport.h

 // AbstractType
 //===----------------------------------------------------------------------===//
 
 /// This class contains all of the static information common to all instances of
 /// a registered Type.
 class AbstractType {
 public:
   using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
+  using WalkImmediateSubElementsFn = function_ref<void(
+      Type, function_ref<void(Attribute)>, function_ref<void(Type)>)>;
+  using ReplaceImmediateSubElementsFn =
+      function_ref<Type(Type, ArrayRef<Attribute>, ArrayRef<Type>)>;
 
   /// Look up the specified abstract type in the MLIRContext and return a
   /// reference to it.
   static const AbstractType &lookup(TypeID typeID, MLIRContext *context);
 
   /// This method is used by Dialect objects when they register the list of
   /// types they contain.
   template <typename T>
   static AbstractType get(Dialect &dialect) {
     return AbstractType(dialect, T::getInterfaceMap(), T::getHasTraitFn(),
-                        T::getTypeID());
+                        T::getWalkImmediateSubElementsFn(),
+                        T::getReplaceImmediateSubElementsFn(), T::getTypeID());
   }
 
   /// This method is used by Dialect objects to register types with
   /// custom TypeIDs.
   /// The use of this method is in general discouraged in favor of
   /// 'get<CustomType>(dialect)';
-  static AbstractType get(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
-                          HasTraitFn &&hasTrait, TypeID typeID) {
+  static AbstractType
+  get(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
+      HasTraitFn &&hasTrait,
+      WalkImmediateSubElementsFn walkImmediateSubElementsFn,
+      ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn,
+      TypeID typeID) {
     return AbstractType(dialect, std::move(interfaceMap), std::move(hasTrait),
-                        typeID);
+                        walkImmediateSubElementsFn,
+                        replaceImmediateSubElementsFn, typeID);
   }
 
   /// Return the dialect this type was registered to.
   Dialect &getDialect() const { return const_cast<Dialect &>(dialect); }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this type, null otherwise. This should not be used
   /// directly.
public:
   template <template <typename T> class Trait>
   bool hasTrait() const {
     return hasTraitFn(TypeID::get<Trait>());
   }
 
   /// Returns true if the type has a particular trait.
   bool hasTrait(TypeID traitID) const { return hasTraitFn(traitID); }
 
+  /// Walk the immediate sub-elements of the given type.
+  void walkImmediateSubElements(Type type,
+                                function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
+
+  /// Replace the immediate sub-elements of the given type.
+  Type replaceImmediateSubElements(Type type, ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const;
+
   /// Return the unique identifier representing the concrete type class.
   TypeID getTypeID() const { return typeID; }
 
 private:
   AbstractType(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
-               HasTraitFn &&hasTrait, TypeID typeID)
+               HasTraitFn &&hasTrait,
+               WalkImmediateSubElementsFn walkImmediateSubElementsFn,
+               ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn,
+               TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
-        hasTraitFn(std::move(hasTrait)), typeID(typeID) {}
+        hasTraitFn(std::move(hasTrait)),
+        walkImmediateSubElementsFn(walkImmediateSubElementsFn),
+        replaceImmediateSubElementsFn(replaceImmediateSubElementsFn),
+        typeID(typeID) {}
 
   /// Give StorageUserBase access to the mutable lookup.
   template <typename ConcreteT, typename BaseT, typename StorageT,
             typename UniquerT, template <typename T> class... Traits>
   friend class detail::StorageUserBase;
 
   /// Look up the specified abstract type in the MLIRContext and return a
   /// (mutable) pointer to it. Return a null pointer if the type could not
   /// be found in the context.
   static AbstractType *lookupMutable(TypeID typeID, MLIRContext *context);
 
   /// This is the dialect that this type was registered to.
   const Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this type.
   detail::InterfaceMap interfaceMap;
 
   /// Function to check if the type has a particular trait.
   HasTraitFn hasTraitFn;
 
+  /// Function to walk the immediate sub-elements of this type.
+  WalkImmediateSubElementsFn walkImmediateSubElementsFn;
+
+  /// Function to replace the immediate sub-elements of this type.
+  ReplaceImmediateSubElementsFn replaceImmediateSubElementsFn;
+
   /// The unique identifier of the derived Type class.
   const TypeID typeID;
 };
 
 //===----------------------------------------------------------------------===//
 // TypeStorage
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Types.h

public:
 
   /// Returns true if the type was registered with a particular trait.
   template <template <typename T> class Trait>
   bool hasTrait() {
     return getAbstractType().hasTrait<Trait>();
   }
 
   /// Return the abstract type descriptor for this type.
-  const AbstractTy &getAbstractType() { return impl->getAbstractType(); }
+  const AbstractTy &getAbstractType() const { return impl->getAbstractType(); }
 
   /// Return the Type implementation.
   ImplType *getImpl() const { return impl; }
 
+  /// Walk all of the immediately nested sub-attributes and sub-types. This
+  /// method does not recurse into sub elements.
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const {
+    getAbstractType().walkImmediateSubElements(*this, walkAttrsFn, walkTypesFn);
+  }
+
+  /// 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.
+  auto replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
+                                   ArrayRef<Type> replTypes) const {
+    return getAbstractType().replaceImmediateSubElements(*this, replAttrs,
+                                                         replTypes);
+  }
+
+  /// Walk this type and all attibutes/types nested within using the
+  /// provided walk functions. See `AttrTypeWalker` for information on the
+  /// supported walk function types.
+  template <WalkOrder Order = WalkOrder::PostOrder, typename... WalkFns>
+  auto walk(WalkFns &&...walkFns) {
+    AttrTypeWalker walker;
+    (walker.addWalk(std::forward<WalkFns>(walkFns)), ...);
+    return walker.walk<Order>(*this);
+  }
+
+  /// Recursively replace all of the nested sub-attributes and sub-types using
+  /// the provided map functions. Returns nullptr in the case of failure. See
+  /// `AttrTypeReplacer` for information on the support replacement function
+  /// types.
+  template <typename... ReplacementFns>
+  auto replace(ReplacementFns &&...replacementFns) {
+    AttrTypeReplacer replacer;
+    (replacer.addReplacement(std::forward<ReplacementFns>(replacementFns)),
+     ...);
+    return replacer.replace(*this);
+  }
+
 protected:
   ImplType *impl{nullptr};
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Type type) {
   type.print(os);
   return os;
 }

mlir/include/mlir/IR/Visitors.h

 ///   * Advance: the walk will continue.
 ///   * Skip: the walk of the current operation, region or block and their
 ///   nested elements that haven't been visited already will be skipped and will
 ///   continue with the next operation, region or block.
 class WalkResult {
   enum ResultEnum { Interrupt, Advance, Skip } result;
 
 public:
-  WalkResult(ResultEnum result) : result(result) {}
+  WalkResult(ResultEnum result = Advance) : result(result) {}
 
   /// Allow LogicalResult to interrupt the walk on failure.
   WalkResult(LogicalResult result)
       : result(failed(result) ? Interrupt : Advance) {}
 
   /// Allow diagnostics to interrupt the walk.
   WalkResult(Diagnostic &&) : result(Interrupt) {}
   WalkResult(InFlightDiagnostic &&) : result(Interrupt) {}

mlir/lib/Dialect/GPU/Transforms/LowerMemorySpaceAttributes.cpp

 //===----------------------------------------------------------------------===//
 
 IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
   return IntegerAttr::get(IntegerType::get(ctx, 64), space);
 }
 
 void mlir::gpu::populateMemorySpaceAttributeTypeConversions(
     TypeConverter &typeConverter, const MemorySpaceMapping &mapping) {
-  typeConverter.addConversion([mapping](Type type) -> std::optional<Type> {
-    auto subElementType = type.dyn_cast_or_null<SubElementTypeInterface>();
-    if (!subElementType)
-      return type;
-    Type newType = subElementType.replaceSubElements(
-        [mapping](Attribute attr) -> std::optional<Attribute> {
-          auto memorySpaceAttr = attr.dyn_cast_or_null<gpu::AddressSpaceAttr>();
-          if (!memorySpaceAttr)
-            return std::nullopt;
-          auto newValue = wrapNumericMemorySpace(
-              attr.getContext(), mapping(memorySpaceAttr.getValue()));
-          return newValue;
-        });
-    return newType;
+  typeConverter.addConversion([mapping](Type type) {
+    return type.replace([mapping](Attribute attr) -> std::optional<Attribute> {
   });
 }
 
 namespace {
 
 /// Converts any op that has operands/results/attributes with numeric MemRef
 /// memory spaces.
 struct LowerMemRefAddressSpacePattern final : public ConversionPattern {

mlir/lib/IR/AsmPrinter.cpp

 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectResourceBlobManager.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
-#include "mlir/IR/SubElementInterfaces.h"
 #include "mlir/IR/Verifier.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/ADT/SetVector.h"
if (auto memrefTy = type.dyn_cast<MemRefType>()) {
       if (!layout.isa<AffineMapAttr>() || !layout.isIdentity())
         printAttribute(memrefTy.getLayout());
       if (memrefTy.getMemorySpace())
         printAttribute(memrefTy.getMemorySpace());
       return;
     }
 
     // For most builtin types, we can simply walk the sub elements.
-    if (auto subElementInterface = dyn_cast<SubElementTypeInterface>(type)) {
-      auto visitFn = [&](auto element) {
-        if (element)
-          (void)printAlias(element);
-      };
-      subElementInterface.walkImmediateSubElements(visitFn, visitFn);
-    }
+    type.walkImmediateSubElements(visitFn, visitFn);
   }
 
   /// Consider the given type to be printed for an alias.
   void printType(Type type) override {
     recordAliasResult(initializer.visit(type, canBeDeferred));
   }
 
   /// Consider the given attribute to be printed for an alias.

mlir/lib/IR/AttrTypeSubElements.cpp

-//===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
+//===- AttrTypeSubElements.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 "mlir/IR/Operation.h"
-
-#include "llvm/ADT/DenseSet.h"
 #include <optional>
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
-// SubElementInterface
-//===----------------------------------------------------------------------===//
-
+// AttrTypeWalker
 //===----------------------------------------------------------------------===//
-// 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) {
-        // Guard against potentially null inputs. This removes the need for the
-        // derived attribute/type to do it.
-        if (!attr)
-          return;
 
-        // Avoid infinite recursion when visiting sub attributes later, if this
-        // is a mutable attribute.
-        if (LLVM_UNLIKELY(attr.hasTrait<AttributeTrait::IsMutable>())) {
-          if (!visitedAttrs.insert(attr).second)
-            return;
-        }
+WalkResult AttrTypeWalker::walkImpl(Attribute attr, WalkOrder order) {
+  return walkImpl(attr, attrWalkFns, order);
+}
+WalkResult AttrTypeWalker::walkImpl(Type type, WalkOrder order) {
+  return walkImpl(type, typeWalkFns, order);
 
-        // Walk any sub elements first.
-        if (auto interface = attr.dyn_cast<SubElementAttrInterface>())
-          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
-                              visitedTypes);
-
-        // Walk this attribute.
-        walkAttrsFn(attr);
-      },
-      [&](Type type) {
-        // Guard against potentially null inputs. This removes the need for the
-        // derived attribute/type to do it.
-        if (!type)
-          return;
+template <typename T, typename WalkFns>
+WalkResult AttrTypeWalker::walkImpl(T element, WalkFns &walkFns,
+                                    WalkOrder order) {
+  // Check if we've already walk this element before.
+  auto key = std::make_pair(element.getAsOpaquePointer(), (int)order);
+  auto it = visitedAttrTypes.find(key);
+  if (it != visitedAttrTypes.end())
+    return it->second;
+  visitedAttrTypes.try_emplace(key, WalkResult::advance());
 
-        // Avoid infinite recursion when visiting sub types later, if this
-        // is a mutable type.
-        if (LLVM_UNLIKELY(type.hasTrait<TypeTrait::IsMutable>())) {
-          if (!visitedTypes.insert(type).second)
-            return;
-        }
+  // If we are walking in post order, walk the sub elements first.
+  if (order == WalkOrder::PostOrder) {
+    if (walkSubElements(element, order).wasInterrupted())
+      return visitedAttrTypes[key] = WalkResult::interrupt();
 
-        // Walk any sub elements first.
-        if (auto interface = type.dyn_cast<SubElementTypeInterface>())
-          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
-                              visitedTypes);
+  // Walk this element, bailing if skipped or interrupted.
+  for (auto &walkFn : llvm::reverse(walkFns)) {
+    WalkResult walkResult = walkFn(element);
+    if (walkResult.wasInterrupted())
+      return visitedAttrTypes[key] = WalkResult::interrupt();
+    if (walkResult.wasSkipped())
+      return WalkResult::advance();
+  }
 
-        // Walk this type.
-        walkTypesFn(type);
-      });
+  // If we are walking in pre-order, walk the sub elements last.
+  if (order == WalkOrder::PreOrder) {
+    if (walkSubElements(element, order).wasInterrupted())
+      return WalkResult::interrupt();
+  }
+  return WalkResult::advance();
 }
 
-void SubElementAttrInterface::walkSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) {
-  assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
-  DenseSet<Attribute> visitedAttrs;
-  DenseSet<Type> visitedTypes;
-  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
-                      visitedTypes);
-}
-
-void SubElementTypeInterface::walkSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) {
-  assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
-  DenseSet<Attribute> visitedAttrs;
-  DenseSet<Type> visitedTypes;
-  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
-                      visitedTypes);
+template <typename T>
+WalkResult AttrTypeWalker::walkSubElements(T interface, WalkOrder order) {
+  WalkResult result = WalkResult::advance();
+  auto walkFn = [&](auto element) {
+    if (element && !result.wasInterrupted())
+      result = walkImpl(element, order);
+  };
+  interface.walkImmediateSubElements(walkFn, walkFn);
+  return result.wasInterrupted() ? result : WalkResult::advance();
 }
 
 //===----------------------------------------------------------------------===//
 /// AttrTypeReplacer
 //===----------------------------------------------------------------------===//
 
+void AttrTypeReplacer::addReplacement(ReplaceFn<Attribute> fn) {
+  attrReplacementFns.emplace_back(std::move(fn));
+}
+void AttrTypeReplacer::addReplacement(ReplaceFn<Type> fn) {
+  typeReplacementFns.push_back(std::move(fn));
+}
+
 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;
   };
void AttrTypeReplacer::recursivelyReplaceElementsIn(Operation *op,
                                                     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,
                                  SmallVectorImpl<T> &newElements,
                                  FailureOr<bool> &changed) {
   // Bail early if we failed at any point.
   if (failed(changed))
     return;
 
   // Guard against potentially null inputs. We always map null to null.
   if (!element) {
     newElements.push_back(nullptr);
     return;
   }
 
   // Replace the element.
   if (T result = replacer.replace(element)) {
     newElements.push_back(result);
     if (result != element)
       changed = true;
   } else {
     changed = failure();
   }
 }
 
-template <typename InterfaceT, typename T>
-T AttrTypeReplacer::replaceSubElements(InterfaceT interface,
-                                       DenseMap<T, T> &interfaceMap) {
+template <typename T>
+T AttrTypeReplacer::replaceSubElements(T interface) {
   // Walk the current sub-elements, replacing them as necessary.
   SmallVector<Attribute, 16> newAttrs;
   SmallVector<Type, 16> newTypes;
   FailureOr<bool> changed = false;
   interface.walkImmediateSubElements(
       [&](Attribute element) {
-        updateSubElementImpl(element, *this, attrMap, newAttrs, changed);
+        updateSubElementImpl(element, *this, newAttrs, changed);
       },
       [&](Type element) {
-        updateSubElementImpl(element, *this, typeMap, newTypes, changed);
+        updateSubElementImpl(element, *this, newTypes, changed);
       });
   if (failed(changed))
     return nullptr;
 
   // If any sub-elements changed, use the new elements during the replacement.
   T result = interface;
   if (*changed)
     result = interface.replaceImmediateSubElements(newAttrs, newTypes);
   return result;
 }
 
 /// Shared implementation of replacing a given attribute or type element.
-template <typename InterfaceT, typename ReplaceFns, typename T>
-T AttrTypeReplacer::replaceImpl(T element, ReplaceFns &replaceFns,
-                                DenseMap<T, T> &map) {
-  auto [it, inserted] = map.try_emplace(element, element);
+template <typename T, typename ReplaceFns>
+T AttrTypeReplacer::replaceImpl(T element, ReplaceFns &replaceFns) {
+  const void *opaqueElement = element.getAsOpaquePointer();
+  auto [it, inserted] = attrTypeMap.try_emplace(opaqueElement, opaqueElement);
   if (!inserted)
-    return it->second;
+    return T::getFromOpaquePointer(it->second);
 
   T result = element;
   WalkResult walkResult = WalkResult::advance();
   for (auto &replaceFn : llvm::reverse(replaceFns)) {
     if (std::optional<std::pair<T, WalkResult>> newRes = replaceFn(element)) {
       std::tie(result, walkResult) = *newRes;
       break;
     }
   }
 
   // If an error occurred, return nullptr to indicate failure.
-  if (walkResult.wasInterrupted() || !result)
-    return map[element] = nullptr;
+  if (walkResult.wasInterrupted() || !result) {
+    attrTypeMap[opaqueElement] = nullptr;
+    return nullptr;
+  }
 
   // Handle replacing sub-elements if this element is also a container.
   if (!walkResult.wasSkipped()) {
-    if (auto interface = dyn_cast<InterfaceT>(result)) {
-      // Replace the sub elements of this element, bailing if we fail.
-      if (!(result = replaceSubElements(interface, map)))
-        return map[element] = nullptr;
+    if (!(result = replaceSubElements(result))) {
+      attrTypeMap[opaqueElement] = nullptr;
+      return nullptr;
     }
   }
 
-  return map[element] = result;
+  attrTypeMap[opaqueElement] = result.getAsOpaquePointer();
+  return result;
 }
 
 Attribute AttrTypeReplacer::replace(Attribute attr) {
-  return replaceImpl<SubElementAttrInterface>(attr, attrReplacementFns,
-                                              attrMap);
+  return replaceImpl(attr, attrReplacementFns);
 }
 
 Type AttrTypeReplacer::replace(Type type) {
-  return replaceImpl<SubElementTypeInterface>(type, typeReplacementFns,
-                                              typeMap);
+  return replaceImpl(type, typeReplacementFns);
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementInterface Tablegen definitions
+// AttrTypeImmediateSubElementWalker
 //===----------------------------------------------------------------------===//
 
-#include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
-#include "mlir/IR/SubElementTypeInterfaces.cpp.inc"
+void AttrTypeImmediateSubElementWalker::walk(Attribute element) {
+  if (element)
+    walkAttrsFn(element);
+}
+
+void AttrTypeImmediateSubElementWalker::walk(Type element) {
+  if (element)
+    walkTypesFn(element);
+}

mlir/lib/IR/Attributes.cpp

 //===- Attributes.cpp - MLIR Affine Expr Classes --------------------------===//
 //
 // 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/Attributes.h"
 #include "mlir/IR/Dialect.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
+// AbstractAttribute
+//===----------------------------------------------------------------------===//
+
+void AbstractAttribute::walkImmediateSubElements(
+    Attribute attr, function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkImmediateSubElementsFn(attr, walkAttrsFn, walkTypesFn);
+}
+
+Attribute
+AbstractAttribute::replaceImmediateSubElements(Attribute attr,
+                                               ArrayRef<Attribute> replAttrs,
+                                               ArrayRef<Type> replTypes) const {
+  return replaceImmediateSubElementsFn(attr, replAttrs, replTypes);
+}
+
+//===----------------------------------------------------------------------===//
 // Attribute
 //===----------------------------------------------------------------------===//
 
 /// Return the context this attribute belongs to.
 MLIRContext *Attribute::getContext() const { return getDialect().getContext(); }
 
 //===----------------------------------------------------------------------===//
 // NamedAttribute

mlir/lib/IR/CMakeLists.txt

 add_mlir_library(MLIRIR
   AffineExpr.cpp
   AffineMap.cpp
   AsmPrinter.cpp
   Attributes.cpp
+  AttrTypeSubElements.cpp
   Block.cpp
   Builders.cpp
   BuiltinAttributeInterfaces.cpp
   BuiltinAttributes.cpp
   BuiltinDialect.cpp
   BuiltinDialectBytecode.cpp
   BuiltinTypes.cpp
   BuiltinTypeInterfaces.cpp
   Diagnostics.cpp
   Dialect.cpp
   DialectResourceBlobManager.cpp
   Dominance.cpp
   ExtensibleDialect.cpp
   FunctionImplementation.cpp
   FunctionInterfaces.cpp
   IntegerSet.cpp
   Location.cpp
   MLIRContext.cpp
   Operation.cpp
   OperationSupport.cpp
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp
-  SubElementInterfaces.cpp
   SymbolTable.cpp
   TensorEncoding.cpp
   Types.cpp
   TypeRange.cpp
   TypeUtilities.cpp
   Value.cpp
   ValueRange.cpp
   Verifier.cpp
   Visitors.cpp
add_mlir_library(MLIRIR
   MLIRBuiltinTypesIncGen
   MLIRBuiltinTypeInterfacesIncGen
   MLIRCallInterfacesIncGen
   MLIRCastInterfacesIncGen
   MLIRDataLayoutInterfacesIncGen
   MLIROpAsmInterfaceIncGen
   MLIRRegionKindInterfaceIncGen
   MLIRSideEffectInterfacesIncGen
-  MLIRSubElementInterfacesIncGen
   MLIRSymbolInterfacesIncGen
   MLIRTensorEncodingIncGen
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/IR/ExtensibleDialect.cpp

void ExtensibleDialect::registerDynamicType(
   (void)registered;
   assert(registered && "type TypeID was not unique");
 
   registered = nameToDynTypes.insert({name, typePtr}).second;
   (void)registered;
   assert(registered &&
          "Trying to create a new dynamic type with an existing name");
 
-  auto abstractType =
-      AbstractType::get(*dialect, DynamicAttr::getInterfaceMap(),
-                        DynamicType::getHasTraitFn(), typeID);
+  auto abstractType = AbstractType::get(
+      *dialect, DynamicAttr::getInterfaceMap(), DynamicType::getHasTraitFn(),
+      DynamicType::getWalkImmediateSubElementsFn(),
+      DynamicType::getReplaceImmediateSubElementsFn(), typeID);
 
   /// Add the type to the dialect and the type uniquer.
   addType(typeID, std::move(abstractType));
   typePtr->registerInTypeUniquer();
 }
 
 void ExtensibleDialect::registerDynamicAttr(
     std::unique_ptr<DynamicAttrDefinition> &&attr) {
void ExtensibleDialect::registerDynamicAttr(
   (void)registered;
   assert(registered && "attribute TypeID was not unique");
 
   registered = nameToDynAttrs.insert({name, attrPtr}).second;
   (void)registered;
   assert(registered &&
          "Trying to create a new dynamic attribute with an existing name");
 
-  auto abstractAttr =
-      AbstractAttribute::get(*dialect, DynamicAttr::getInterfaceMap(),
-                             DynamicAttr::getHasTraitFn(), typeID);
+  auto abstractAttr = AbstractAttribute::get(
+      *dialect, DynamicAttr::getInterfaceMap(), DynamicAttr::getHasTraitFn(),
+      DynamicAttr::getWalkImmediateSubElementsFn(),
+      DynamicAttr::getReplaceImmediateSubElementsFn(), typeID);
 
   /// Add the type to the dialect and the type uniquer.
   addAttribute(typeID, std::move(abstractAttr));
   attrPtr->registerInAttrUniquer();
 }
 
 void ExtensibleDialect::registerDynamicOp(
     std::unique_ptr<DynamicOpDefinition> &&op) {

mlir/lib/IR/SymbolTable.cpp

 //===----------------------------------------------------------------------===//
 
 /// Walk all of the symbol references within the given operation, invoking the
 /// 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 {
-    SubElementAttrInterface container;
-    SmallVector<Attribute> immediateSubElements;
-
-    explicit WorklistItem(SubElementAttrInterface container) {
-      SmallVector<Attribute> subElements;
-      container.walkImmediateSubElements(
-          [&](Attribute attr) { subElements.push_back(attr); }, [](Type) {});
-      immediateSubElements = std::move(subElements);
-    }
-  };
-
-  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>()) {
+  return op->getAttrDictionary().walk<WalkOrder::PreOrder>(
+      [&](SymbolRefAttr symbolRef) {
         if (callback({op, symbolRef}).wasInterrupted())
           return WalkResult::interrupt();
 
-        /// Check for a nested container attribute, these will also need to be
-        /// walked.
-      } else if (auto interface = attr.dyn_cast<SubElementAttrInterface>()) {
-        attrWorklist.emplace_back(interface);
-        curAccessChain.push_back(-1);
-        return WalkResult::advance();
-      }
-      // 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();
-  };
-
-  WalkResult result = WalkResult::advance();
-  do {
-    WorklistItem &item = attrWorklist.back();
-    int &index = curAccessChain.back();
-    ++index;
-
-    // Process the given attribute, which is guaranteed to be a container.
-    result = processAttrs(index, item);
-  } while (!attrWorklist.empty() && !result.wasInterrupted());
-  return result;
+        // Don't walk nested references.
+        return WalkResult::skip();
+      });
 }
 
 /// 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 std::optional<WalkResult>
 walkSymbolUses(MutableArrayRef<Region> regions,
                function_ref<WalkResult(SymbolTable::SymbolUse)> callback) {

mlir/lib/IR/TypeDetail.h

 
 /// Integer Type Storage and Uniquing.
 struct IntegerTypeStorage : public TypeStorage {
   IntegerTypeStorage(unsigned width,
                      IntegerType::SignednessSemantics signedness)
       : width(width), signedness(signedness) {}
 
   /// The hash key used for uniquing.
-  using KeyTy = std::pair<unsigned, IntegerType::SignednessSemantics>;
+  using KeyTy = std::tuple<unsigned, IntegerType::SignednessSemantics>;
 
   static llvm::hash_code hashKey(const KeyTy &key) {
     return llvm::hash_value(key);
   }
 
   bool operator==(const KeyTy &key) const {
     return KeyTy(width, signedness) == key;
   }
 
   static IntegerTypeStorage *construct(TypeStorageAllocator &allocator,
                                        KeyTy key) {
     return new (allocator.allocate<IntegerTypeStorage>())
-        IntegerTypeStorage(key.first, key.second);
+        IntegerTypeStorage(std::get<0>(key), std::get<1>(key));
   }
 
   KeyTy getAsKey() const { return KeyTy(width, signedness); }
 
   unsigned width : 30;
   IntegerType::SignednessSemantics signedness : 2;
 };
 

mlir/lib/IR/Types.cpp

 //===- Types.cpp - MLIR Type Classes --------------------------------------===//
 //
 // 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/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
+// AbstractType
+//===----------------------------------------------------------------------===//
+
+void AbstractType::walkImmediateSubElements(
+    Type type, function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkImmediateSubElementsFn(type, walkAttrsFn, walkTypesFn);
+}
+
+Type AbstractType::replaceImmediateSubElements(Type type,
+                                               ArrayRef<Attribute> replAttrs,
+                                               ArrayRef<Type> replTypes) const {
+  return replaceImmediateSubElementsFn(type, replAttrs, replTypes);
+}
+
+//===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//
 
 MLIRContext *Type::getContext() const { return getDialect().getContext(); }
 
 bool Type::isFloat8E5M2() const { return isa<Float8E5M2Type>(); }
 bool Type::isFloat8E4M3FN() const { return isa<Float8E4M3FNType>(); }
 bool Type::isBF16() const { return isa<BFloat16Type>(); }

mlir/test/IR/test-symbol-rauw.mlir

module {
   func.func private @failed_repl() attributes {sym.new_name = "replaced_name" }
 
   "foo.possibly_unknown_symbol_table"() ({
   }) : () -> ()
 }
 
 // -----
 
-// Check that replacement works in any implementations of SubElementsAttrInterface
+// Check that replacement works in any implementations of SubElements.
 module {
     // CHECK: func private @replaced_foo
     func.func private @symbol_foo() attributes {sym.new_name = "replaced_foo" }
 
     // CHECK: func @symbol_bar
     func.func @symbol_bar() {
       // CHECK: foo.op
       // CHECK-SAME: non_symbol_attr,

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

 // To get the test dialect definition.
 include "TestDialect.td"
 include "TestEnumDefs.td"
 include "mlir/Dialect/Utils/StructuredOpsUtils.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpAsmInterface.td"
-include "mlir/IR/SubElementInterfaces.td"
 
 // All of the attributes will extend this class.
 class Test_Attr<string name, list<Trait> traits = []>
     : AttrDef<Test_Dialect, name, traits>;
 
 def SimpleAttrA : Test_Attr<"SimpleA"> {
   let mnemonic = "smpla";
 }
auto try_value_begin_impl(OverloadToken<llvm::APInt>) const {
         return llvm::APInt(/*numBits=*/64, value);
       }).begin());
     }
   }];
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
-def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
-    SubElementAttrInterface
-  ]> {
+def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess"> {
   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/TestTypes.h

 #include <tuple>
 #include <optional>
 
 #include "TestTraits.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Operation.h"
-#include "mlir/IR/SubElementInterfaces.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace test {
 class TestAttrWithFormatAttr;
 
 /// FieldInfo represents a field in the StructType data type. It is used as a
 /// parameter in TestTypeDefs.td.

 };
 
 /// Simple recursive type identified by its name and pointing to another named
 /// type, potentially itself. This requires the body to be mutated separately
 /// from type creation.
 class TestRecursiveType
     : public ::mlir::Type::TypeBase<TestRecursiveType, ::mlir::Type,
                                     TestRecursiveTypeStorage,
-                                    ::mlir::SubElementTypeInterface::Trait,
                                     ::mlir::TypeTrait::IsMutable> {
 public:
   using Base::Base;
 
   static TestRecursiveType get(::mlir::MLIRContext *ctx,
                                ::llvm::StringRef name) {
     return Base::get(ctx, name);
   }

mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp

 //===- LLVMTypeTest.cpp - Tests for LLVM types ----------------------------===//
 //
 // 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 "LLVMTestBase.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/IR/SubElementInterfaces.h"
 
 using namespace mlir;
 using namespace mlir::LLVM;
 
 TEST_F(LLVMIRTest, IsStructTypeMutable) {
   auto structTy = LLVMStructType::getIdentified(&context, "foo");
   ASSERT_TRUE(bool(structTy));
   ASSERT_TRUE(structTy.hasTrait<TypeTrait::IsMutable>());
 }
 
 TEST_F(LLVMIRTest, MutualReferencedSubElementTypes) {
   auto fooStructTy = LLVMStructType::getIdentified(&context, "foo");
   ASSERT_TRUE(bool(fooStructTy));
   auto barStructTy = LLVMStructType::getIdentified(&context, "bar");
   ASSERT_TRUE(bool(barStructTy));
 
   // Created two structs that are referencing each other.
   Type fooBody[] = {LLVMPointerType::get(barStructTy)};
   ASSERT_TRUE(succeeded(fooStructTy.setBody(fooBody, /*isPacked=*/false)));
   Type barBody[] = {LLVMPointerType::get(fooStructTy)};
   ASSERT_TRUE(succeeded(barStructTy.setBody(barBody, /*isPacked=*/false)));
 
-  auto subElementInterface = fooStructTy.dyn_cast<SubElementTypeInterface>();
-  ASSERT_TRUE(bool(subElementInterface));
   // Test if walkSubElements goes into infinite loops.
   SmallVector<Type, 4> subElementTypes;
-  subElementInterface.walkSubElements(
-      [](Attribute attr) {},
-      [&](Type type) { subElementTypes.push_back(type); });
-  // We don't record LLVMPointerType (because it's immutable), thus
-  // !llvm.ptr<struct<"bar",...>> will be visited twice.
-  ASSERT_EQ(subElementTypes.size(), 5U);
+  fooStructTy.walk([&](Type type) { subElementTypes.push_back(type); });
+  ASSERT_EQ(subElementTypes.size(), 4U);
 
-  // !llvm.ptr<struct<"bar",...>>
+  // !llvm.ptr<struct<"foo",...>>
   ASSERT_TRUE(subElementTypes[0].isa<LLVMPointerType>());
 
-  // !llvm.struct<"foo",...>
+  // !llvm.struct<"bar",...>
   auto structType = subElementTypes[1].dyn_cast<LLVMStructType>();
   ASSERT_TRUE(bool(structType));
-  ASSERT_TRUE(structType.getName().equals("foo"));
+  ASSERT_TRUE(structType.getName().equals("bar"));
 
-  // !llvm.ptr<struct<"foo",...>>
+  // !llvm.ptr<struct<"bar",...>>
   ASSERT_TRUE(subElementTypes[2].isa<LLVMPointerType>());
 
-  // !llvm.struct<"bar",...>
+  // !llvm.struct<"foo",...>
   structType = subElementTypes[3].dyn_cast<LLVMStructType>();
   ASSERT_TRUE(bool(structType));
-  ASSERT_TRUE(structType.getName().equals("bar"));
-
-  // !llvm.ptr<struct<"bar",...>>
-  ASSERT_TRUE(subElementTypes[4].isa<LLVMPointerType>());
+  ASSERT_TRUE(structType.getName().equals("foo"));
 }

mlir/unittests/IR/AttributeTest.cpp

TEST(SparseElementsAttrTest, GetZero) {
   EXPECT_TRUE(zeroFloatValue.getType() == floatTy);
 
   auto zeroStringValue =
       sparseString.getValues<Attribute>()[{1, 1}].cast<StringAttr>();
   EXPECT_TRUE(zeroStringValue.getValue().empty());
   EXPECT_TRUE(zeroStringValue.getType() == stringTy);
 }
 
+//===----------------------------------------------------------------------===//
+// SubElements
+//===----------------------------------------------------------------------===//
+
+TEST(SubElementTest, Nested) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  BoolAttr trueAttr = builder.getBoolAttr(true);
+  BoolAttr falseAttr = builder.getBoolAttr(false);
+  ArrayAttr boolArrayAttr = builder.getArrayAttr({trueAttr, falseAttr});
+  StringAttr strAttr = builder.getStringAttr("array");
+  DictionaryAttr dictAttr =
+      builder.getDictionaryAttr(builder.getNamedAttr(strAttr, boolArrayAttr));
+
+  SmallVector<Attribute> subAttrs;
+  dictAttr.walk([&](Attribute attr) { subAttrs.push_back(attr); });
+  EXPECT_EQ(llvm::ArrayRef(subAttrs),
+            ArrayRef<Attribute>(
+                {strAttr, trueAttr, falseAttr, boolArrayAttr, dictAttr}));
+}
 } // namespace

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AdaptorTest.cpp
   AttributeTest.cpp
   DialectTest.cpp
   InterfaceTest.cpp
   IRMapping.cpp
   InterfaceAttachmentTest.cpp
   OperationSupportTest.cpp
   PatternMatchTest.cpp
   ShapedTypeTest.cpp
-  SubElementInterfaceTest.cpp
   TypeTest.cpp
 
   DEPENDS
   MLIRTestInterfaceIncGen
 )
 target_include_directories(MLIRIRTests PRIVATE "${MLIR_BINARY_DIR}/test/lib/Dialect/Test")
 target_link_libraries(MLIRIRTests
   PRIVATE
   MLIRIR
   MLIRTestDialect)

mlir/unittests/IR/InterfaceTest.cpp

TEST(InterfaceTest, OpInterfaceDenseMapKey) {
   opSet.insert(op1);
   opSet.insert(op2);
   opSet.erase(op1);
   EXPECT_FALSE(opSet.contains(op1));
   EXPECT_TRUE(opSet.contains(op2));
   EXPECT_FALSE(opSet.contains(op3));
 }
 
-TEST(InterfaceTest, AttrInterfaceDenseMapKey) {
-  MLIRContext context;
-  context.loadDialect<test::TestDialect>();
-
-  OpBuilder builder(&context);
-
-  DenseSet<SubElementAttrInterface> attrSet;
-  auto attr1 = builder.getArrayAttr({});
-  auto attr2 = builder.getI32ArrayAttr({0});
-  auto attr3 = builder.getI32ArrayAttr({1});
-  attrSet.insert(attr1);
-  attrSet.insert(attr2);
-  attrSet.erase(attr1);
-  EXPECT_FALSE(attrSet.contains(attr1));
-  EXPECT_TRUE(attrSet.contains(attr2));
-  EXPECT_FALSE(attrSet.contains(attr3));
-}
-
 TEST(InterfaceTest, TypeInterfaceDenseMapKey) {
   MLIRContext context;
   context.loadDialect<test::TestDialect>();
 
   OpBuilder builder(&context);
   DenseSet<DataLayoutTypeInterface> typeSet;
   auto type1 = builder.getType<test::TestTypeWithLayoutType>(1);
   auto type2 = builder.getType<test::TestTypeWithLayoutType>(2);
   auto type3 = builder.getType<test::TestTypeWithLayoutType>(3);
   typeSet.insert(type1);
   typeSet.insert(type2);
   typeSet.erase(type1);
   EXPECT_FALSE(typeSet.contains(type1));
   EXPECT_TRUE(typeSet.contains(type2));
   EXPECT_FALSE(typeSet.contains(type3));
 }

mlir/unittests/IR/SubElementInterfaceTest.cpp

-//===- SubElementInterfaceTest.cpp - SubElementInterface unit tests -------===//
-//
-// 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/Builders.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/SubElementInterfaces.h"
-#include "gtest/gtest.h"
-#include <cstdint>
-
-using namespace mlir;
-using namespace mlir::detail;
-
-namespace {
-TEST(SubElementInterfaceTest, Nested) {
-  MLIRContext context;
-  Builder builder(&context);
-
-  BoolAttr trueAttr = builder.getBoolAttr(true);
-  BoolAttr falseAttr = builder.getBoolAttr(false);
-  ArrayAttr boolArrayAttr = builder.getArrayAttr({trueAttr, falseAttr});
-  StringAttr strAttr = builder.getStringAttr("array");
-  DictionaryAttr dictAttr =
-      builder.getDictionaryAttr(builder.getNamedAttr(strAttr, boolArrayAttr));
-
-  SmallVector<Attribute> subAttrs;
-  dictAttr.walkSubAttrs([&](Attribute attr) { subAttrs.push_back(attr); });
-  EXPECT_EQ(llvm::ArrayRef(subAttrs),
-            ArrayRef<Attribute>({strAttr, trueAttr, falseAttr, boolArrayAttr}));
-}
-
-} // namespace