[mlir] Infer SubElementInterface implementations using the storage KeyTy

[rriddle]

Summary

The KeyTy of attribute/type storage classes provide enough information for
automatically implementing the necessary sub element interface methods. This
removes the need for derived classes to do it themselves, which is both much
nicer and easier to handle certain invariants (e.g. null handling). In cases where
explicitly handling for parameter types is necessary, they can provide an implementation
of AttrTypeSubElementHandler to opt-in to support.

This tickles a few things alias wise, which annoyingly messes with tests that hard
code specific affine map numbers.

[Mogball]

Amazing

mlir/include/mlir/IR/SubElementInterfaces.h

94

I'm not sure this is better than just indexing directly into the array and casting. This is begging for an msvc problem where

get(repl.extract<IntegerAttr>(), reply.extract<FloatAttr>());

Crashes because the argument expressions are evaluated in reverse.

mlir/include/mlir/IR/SubElementInterfaces.td

59

QQ where is derviedValue coming from?

[rriddle]

mlir/include/mlir/IR/SubElementInterfaces.h

94

Not sure what you mean here. This is exactly what we do in a lot of cases already, given that at least for attributes they are guaranteed to be the same size. If this would break, it would have broken a long time ago. It would only break iff they could have different things in them depending on the platform, but that isn't the case here.

mlir/include/mlir/IR/SubElementInterfaces.td

59

It's a template parameter in the class (SubElementInterfaceBase), which expands to either $_attr or $_type depending on the interface.

[rriddle]

mlir/include/mlir/IR/SubElementInterfaces.h

94

Good point re: order of evaluation (I hate C++). Will fix that.

[rriddle]

mlir/include/mlir/IR/SubElementInterfaces.h

94

Updated to check size as well (for MSVC bad times). It's really just an optimization, so we can just skip it in the cases where MSVC is dumb.

flang/test/Fir/affine-promotion.fir

func.func @loop_with_load_and_store(%a1: !arr_d1, %a2: !arr_d1, %a3: !arr_d1) {
   }
   return
 }
 
 // CHECK:  func @loop_with_load_and_store(%[[VAL_0:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_1:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_2:.*]]: !fir.ref<!fir.array<?xf32>>) {
 // CHECK:    %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK:    %[[VAL_4:.*]] = arith.constant 100 : index
 // CHECK:    %[[VAL_5:.*]] = fir.shape %[[VAL_4]] : (index) -> !fir.shape<1>
-// CHECK:    %[[VAL_6:.*]] = affine.apply #map(){{\[}}%[[VAL_3]], %[[VAL_4]]]
+// CHECK:    %[[VAL_6:.*]] = affine.apply #{{.*}}(){{\[}}%[[VAL_3]], %[[VAL_4]]]
 // CHECK:    %[[VAL_7:.*]] = fir.alloca !fir.array<?xf32>, %[[VAL_6]]
 // CHECK:    %[[VAL_8:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
 // CHECK:    %[[VAL_9:.*]] = fir.convert %[[VAL_1]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
 // CHECK:    %[[VAL_10:.*]] = fir.convert %[[VAL_7]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
-// CHECK:    affine.for %[[VAL_11:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] {
-// CHECK:      %[[VAL_12:.*]] = affine.apply #map2(%[[VAL_11]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
+// CHECK:    affine.for %[[VAL_11:.*]] = %[[VAL_3]] to #{{.*}}(){{\[}}%[[VAL_4]]] {
+// CHECK:      %[[VAL_12:.*]] = affine.apply #{{.*}}(%[[VAL_11]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
 // CHECK:      %[[VAL_13:.*]] = affine.load %[[VAL_8]]{{\[}}%[[VAL_12]]] : memref<?xf32>
 // CHECK:      %[[VAL_14:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<?xf32>
 // CHECK:      %[[VAL_15:.*]] = arith.addf %[[VAL_13]], %[[VAL_14]] : f32
 // CHECK:      affine.store %[[VAL_15]], %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref<?xf32>
 // CHECK:    }
 // CHECK:    %[[VAL_16:.*]] = fir.convert %[[VAL_2]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
-// CHECK:    affine.for %[[VAL_17:.*]] = %[[VAL_3]] to #map1(){{\[}}%[[VAL_4]]] {
-// CHECK:      %[[VAL_18:.*]] = affine.apply #map2(%[[VAL_17]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
+// CHECK:    affine.for %[[VAL_17:.*]] = %[[VAL_3]] to #{{.*}}(){{\[}}%[[VAL_4]]] {
+// CHECK:      %[[VAL_18:.*]] = affine.apply #{{.*}}(%[[VAL_17]]){{\[}}%[[VAL_3]], %[[VAL_4]], %[[VAL_3]]]
 // CHECK:      %[[VAL_19:.*]] = affine.load %[[VAL_10]]{{\[}}%[[VAL_18]]] : memref<?xf32>
 // CHECK:      %[[VAL_20:.*]] = affine.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xf32>
 // CHECK:      %[[VAL_21:.*]] = arith.mulf %[[VAL_19]], %[[VAL_20]] : f32
 // CHECK:      affine.store %[[VAL_21]], %[[VAL_16]]{{\[}}%[[VAL_18]]] : memref<?xf32>
 // CHECK:    }
 // CHECK:    return
 // CHECK:  }
 

 
 // CHECK: func @loop_with_if(%[[VAL_0:.*]]: !fir.ref<!fir.array<?xf32>>, %[[VAL_1:.*]]: f32) {
 // CHECK:   %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK:   %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK:   %[[VAL_4:.*]] = arith.constant 2 : index
 // CHECK:   %[[VAL_5:.*]] = arith.constant 100 : index
 // CHECK:   %[[VAL_6:.*]] = fir.shape %[[VAL_5]] : (index) -> !fir.shape<1>
 // CHECK:   %[[VAL_7:.*]] = fir.convert %[[VAL_0]] : (!fir.ref<!fir.array<?xf32>>) -> memref<?xf32>
-// CHECK:   affine.for %[[VAL_8:.*]] = %[[VAL_3]] to #map(){{\[}}%[[VAL_5]]] {
-// CHECK:     %[[VAL_9:.*]] = affine.apply #map1(%[[VAL_8]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:   affine.for %[[VAL_8:.*]] = %[[VAL_3]] to #{{.*}}(){{\[}}%[[VAL_5]]] {
+// CHECK:     %[[VAL_9:.*]] = affine.apply #{{.*}}(%[[VAL_8]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
 // CHECK:     affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_9]]] : memref<?xf32>
 // CHECK:   }
-// CHECK:   affine.for %[[VAL_10:.*]] = %[[VAL_3]] to #map(){{\[}}%[[VAL_5]]] {
-// CHECK:     %[[VAL_11:.*]] = affine.apply #map1(%[[VAL_10]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:   affine.for %[[VAL_10:.*]] = %[[VAL_3]] to #{{.*}}(){{\[}}%[[VAL_5]]] {
+// CHECK:     %[[VAL_11:.*]] = affine.apply #{{.*}}(%[[VAL_10]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
 // CHECK:     affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_11]]] : memref<?xf32>
 // CHECK:   }
-// CHECK:   affine.for %[[VAL_12:.*]] = %[[VAL_3]] to #map(){{\[}}%[[VAL_5]]] {
+// CHECK:   affine.for %[[VAL_12:.*]] = %[[VAL_3]] to #{{.*}}(){{\[}}%[[VAL_5]]] {
 // CHECK:     %[[VAL_13:.*]] = arith.subi %[[VAL_12]], %[[VAL_4]] : index
 // CHECK:     affine.if #set(%[[VAL_12]]) {
-// CHECK:       %[[VAL_14:.*]] = affine.apply #map1(%[[VAL_12]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
+// CHECK:       %[[VAL_14:.*]] = affine.apply #{{.*}}(%[[VAL_12]]){{\[}}%[[VAL_3]], %[[VAL_5]], %[[VAL_3]]]
 // CHECK:       affine.store %[[VAL_1]], %[[VAL_7]]{{\[}}%[[VAL_14]]] : memref<?xf32>
 // CHECK:     }
 // CHECK:   }
 // CHECK:   return
 // CHECK: }

mlir/docs/AttributesAndTypes.md

- Provide a comparison method between an instance of the storage and the
   - `bool operator==(const KeyTy &) const`
 - Provide a method to generate the `KeyTy` from a list of arguments passed to
   the uniquer when building an Attribute or Type. (Note: This is only necessary
   if the `KeyTy` cannot be default constructed from these arguments).
   - `static KeyTy getKey(Args...&& args)`
 - Provide a method to hash an instance of the `KeyTy`. (Note: This is not
   necessary if an `llvm::DenseMapInfo<KeyTy>` specialization exists)
   - `static llvm::hash_code hashKey(const KeyTy &)`
+- Provide a method to generate the `KeyTy` from an instance of the storage class.
+  - `static KeyTy getAsKey()`
 
 Let's look at an example:
 
 ```c++
 /// Here we define a storage class for a ComplexType, that holds a non-zero
 /// integer and an integer type.
 struct ComplexTypeStorage : public TypeStorage {
   ComplexTypeStorage(unsigned nonZeroParam, Type integerType)
struct ComplexTypeStorage : public TypeStorage {
   }
 
   /// Define a construction method for creating a new instance of this storage.
   static ComplexTypeStorage *construct(StorageAllocator &allocator, const KeyTy &key) {
     return new (allocator.allocate<ComplexTypeStorage>())
         ComplexTypeStorage(key.first, key.second);
   }
 
+  /// Construct an instance of the key from this storage class.
+  KeyTy getAsKey() const {
+    return KeyTy(nonZeroParam, integerType);
+  }
+
   /// The parametric data held by the storage class.
   unsigned nonZeroParam;
   Type integerType;
 };
 ```
 
 ### Mutable attributes and types
 

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

def LLVM_DIBasicTypeAttr : LLVM_Attr<"DIBasicType", "di_basic_type",
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DICompileUnitAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_DICompileUnitAttr : LLVM_Attr<"DICompileUnit", "di_compile_unit", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DITypeAttr"> {
   let parameters = (ins
     LLVM_DITagParameter:$tag,
     "StringAttr":$name,
     OptionalParameter<"DIFileAttr">:$file,
     "uint32_t":$line,
     OptionalParameter<"DIScopeAttr">:$scope,
     "uint64_t":$sizeInBits,
     "uint64_t":$alignInBits,
     OptionalArrayRefParameter<"DINodeAttr">:$elements
   );
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DIDerivedTypeAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_DIDerivedTypeAttr : LLVM_Attr<"DIDerivedType", "di_derived_type", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DITypeAttr"> {
   let parameters = (ins
     LLVM_DITagParameter:$tag,
     "StringAttr":$name,
     "DITypeAttr":$baseType,
     "uint64_t":$sizeInBits,
     "uint32_t":$alignInBits,
     "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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DIScopeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     "DIFileAttr":$file,
     "unsigned":$line,
     "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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DIScopeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     "DIFileAttr":$file,
     "unsigned":$descriminator
   );
   let builders = [
     AttrBuilderWithInferredContext<(ins
       "DIScopeAttr":$scope, "DIFileAttr":$file, "unsigned":$descriminator
     ), [{
       return $_get(file.getContext(), scope, file, descriminator);
     }]>
   ];
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DILocalVariableAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_DILocalVariableAttr : LLVM_Attr<"DILocalVariable", "di_local_variable", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DINodeAttr"> {
   let parameters = (ins
     "DIScopeAttr":$scope,
     "StringAttr":$name,
     "DIFileAttr":$file,
     "unsigned":$line,
     "unsigned":$arg,
     "unsigned":$alignInBits,
def LLVM_DILocalVariableAttr : LLVM_Attr<"DILocalVariable", "di_local_variable", [
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DISubprogramAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_DISubprogramAttr : LLVM_Attr<"DISubprogram", "di_subprogram", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "DIScopeAttr"> {
   let parameters = (ins
     "DICompileUnitAttr":$compileUnit,
     "DIScopeAttr":$scope,
     "StringAttr":$name,
     "StringAttr":$linkageName,
     "DIFileAttr":$file,
     "unsigned":$line,
def LLVM_DISubrangeAttr : LLVM_Attr<"DISubrange", "di_subrange", /*traits=*/[],
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 //===----------------------------------------------------------------------===//
 // DISubroutineTypeAttr
 //===----------------------------------------------------------------------===//
 
 def LLVM_DISubroutineTypeAttr : LLVM_Attr<"DISubroutineType", "di_subroutine_type", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ], "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) `>`";
 }
 
 #endif // LLVMIR_ATTRDEFS

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap", [
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_ArrayAttr : Builtin_Attr<"Array", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   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 skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolRefAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   let summary = "An Attribute containing a symbolic reference to an Operation";
   let description = [{
     Syntax:
 
     ```
     symbol-ref-attribute ::= symbol-ref-id (`::` symbol-ref-id)*
     ```
def Builtin_SymbolRefAttr : Builtin_Attr<"SymbolRef", [
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TypeAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_TypeAttr : Builtin_Attr<"Type", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   let summary = "An Attribute containing a Type";
   let description = [{
     Syntax:
 
     ```
     type-attribute ::= type
     ```

mlir/include/mlir/IR/BuiltinLocationAttributes.td

class Builtin_LocationAttr<string name, list<Trait> traits = []>
   let mnemonic = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
 def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   let summary = "A tuple of other source locations";
   let description = [{
     Syntax:
 
     ```
     fused-location ::= `fused` fusion-metadata? `[` location (location `,`)* `]`
     fusion-metadata ::= `<` attribute-value `>`
def FusedLoc : Builtin_LocationAttr<"FusedLoc", [
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // NameLoc
 //===----------------------------------------------------------------------===//
 
 def NameLoc : Builtin_LocationAttr<"NameLoc", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   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", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+    SubElementAttrInterface
   ]> {
   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.
   }];

mlir/include/mlir/IR/Location.h

public:
   friend ::llvm::hash_code hash_value(Location arg);
 
   /// Methods for supporting PointerLikeTypeTraits.
   const void *getAsOpaquePointer() const { return impl.getAsOpaquePointer(); }
   static Location getFromOpaquePointer(const void *pointer) {
     return LocationAttr(reinterpret_cast<const AttributeStorage *>(pointer));
   }
 
+  /// Support llvm style casting.
+  static bool classof(Attribute attr) { return llvm::isa<LocationAttr>(attr); }
+
 protected:
   /// The internal backing location attribute.
   LocationAttr impl;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, const Location &loc) {
   loc.print(os);
   return os;

 
 /// Returns an instance of opaque location which contains a given pointer to
 /// 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
+//===----------------------------------------------------------------------===//
+
+/// Enable locations to be introspected as sub-elements.
+template <>
+struct AttrTypeSubElementHandler<Location> {
+  static void walk(Location param, AttrTypeSubElementWalker &walker) {
+    walker.walk(param);
+  }
+  static Location replace(Location param, AttrSubElementReplacements &attrRepls,
+                          TypeSubElementReplacements &typeRepls) {
+    return cast<LocationAttr>(attrRepls.take_front(1)[0]);
+  }
+};
+
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // LLVM Utilities
 //===----------------------------------------------------------------------===//
 
 namespace llvm {
 

mlir/include/mlir/IR/StorageUniquerSupport.h

static ConcreteT getChecked(function_ref<InFlightDiagnostic()> emitErrorFn,
     return UniquerT::template get<ConcreteT>(ctx, args...);
   }
 
   /// Get an instance of the concrete type from a void pointer.
   static ConcreteT getFromOpaquePointer(const void *ptr) {
     return ConcreteT((const typename BaseT::ImplType *)ptr);
   }
 
+  /// Utility for easy access to the storage instance.
+  ImplType *getImpl() const { return static_cast<ImplType *>(this->impl); }
+
 protected:
   /// Mutate the current storage instance. This will not change the unique key.
   /// The arguments are forwarded to 'ConcreteT::mutate'.
   template <typename... Args>
   LogicalResult mutate(Args &&...args) {
     static_assert(std::is_base_of<StorageUserTrait::IsMutable<ConcreteT>,
                                   ConcreteT>::value,
                   "The `mutate` function expects mutable trait "
                   "(e.g. TypeTrait::IsMutable) to be attached on parent.");
     return UniquerT::template mutate<ConcreteT>(this->getContext(), getImpl(),
                                                 std::forward<Args>(args)...);
   }
 
   /// Default implementation that just returns success.
   template <typename... Args>
   static LogicalResult verify(Args... args) {
     return success();
   }
 
-  /// Utility for easy access to the storage instance.
-  ImplType *getImpl() const { return static_cast<ImplType *>(this->impl); }
-
 private:
   /// Trait to check if T provides a 'ConcreteEntity' type alias.
   template <typename T>
   using has_concrete_entity_t = typename T::ConcreteEntity;
 
   /// A struct-wrapped type alias to T::ConcreteEntity if provided and to
   /// ConcreteT otherwise. This is akin to std::conditional but doesn't fail on
   /// the missing typedef. Useful for checking if the interface is targeting the

mlir/include/mlir/IR/SubElementInterfaces.h

 #include "mlir/IR/Types.h"
 #include "mlir/IR/Visitors.h"
 
 namespace mlir {
 template <typename T>
 using SubElementReplFn = function_ref<T(T)>;
 template <typename T>
 using SubElementResultReplFn = function_ref<std::pair<T, WalkResult>(T)>;
+
+//===----------------------------------------------------------------------===//
+/// AttrTypeSubElementHandler
+//===----------------------------------------------------------------------===//
+
+/// This class is used by AttrTypeSubElementHandler instances to walking sub
+/// attributes and types.
+class AttrTypeSubElementWalker {
+public:
+  AttrTypeSubElementWalker(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);
+  }
+  /// Walk a type.
+  void walk(Type element) {
+    if (element)
+      walkTypesFn(element);
+  }
+  /// Walk a range of attributes or types.
+  template <typename RangeT>
+  void walkRange(RangeT &&elements) {
+    for (auto element : elements)
+      walk(element);
+  }
+
+private:
+  function_ref<void(Attribute)> walkAttrsFn;
+  function_ref<void(Type)> walkTypesFn;
+};
+
+/// This class is used by AttrTypeSubElementHandler instances to process sub
+/// element replacements.
+template <typename T>
+class AttrTypeSubElementReplacements {
+public:
+  AttrTypeSubElementReplacements(ArrayRef<T> repls) : repls(repls) {}
+
+  /// Take the first N replacements as an ArrayRef, dropping them from
+  /// this replacement list.
+  ArrayRef<T> take_front(unsigned n) {
+    ArrayRef<T> elements = repls.take_front(n);
+    repls = repls.drop_front(n);
+    return elements;
+  }
+
+private:
+  /// The current set of replacements.
+  ArrayRef<T> repls;
+};
+using AttrSubElementReplacements = AttrTypeSubElementReplacements<Attribute>;
+using TypeSubElementReplacements = AttrTypeSubElementReplacements<Type>;
+
+/// This class provides support for interacting with the
+/// SubElementInterfaces for different types of parameters. An
+/// implementation of this class should be provided for any parameter class
+/// that may contain an attribute or type. There are two main methods of
+/// this class that need to be implemented:
+///
+///  - walk
+///
+///   This method should traverse into any sub elements of the parameter
+///   using the provided walker, or by invoking handlers for sub-types.
+///
+///  - replace
+///
+///   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) {}
+
+  /// 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);
+  }
+
+  /// Tag indicating that this handler does not support sub-elements.
+  using DefaultHandlerTag = void;
+};
+
+/// Detect if any of the given parameter types has a sub-element handler.
+namespace detail {
+template <typename T>
+using has_default_sub_element_handler_t = decltype(T::DefaultHandlerTag);
+} // namespace detail
+template <typename... Ts>
+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) {
+    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 if constexpr (!detail::IsInterface<T>::value &&
+                         std::is_base_of_v<Type, T>) {
+      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) {
+    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)) {
+      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)) {
+      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) {
+    std::apply(
+        [&](auto &&...params) {
+          (AttrTypeSubElementHandler<Ts>::walk(params, walker), ...);
+        },
+        param);
+  }
+  static auto replace(const std::tuple<Ts...> &param,
+                      AttrSubElementReplacements &attrRepls,
+                      TypeSubElementReplacements &typeRepls) {
+    return std::apply(
+        [&](const Ts &...params)
+            -> std::tuple<decltype(AttrTypeSubElementHandler<Ts>::replace(
+                params, attrRepls, typeRepls))...> {
+          return {AttrTypeSubElementHandler<Ts>::replace(params, attrRepls,
+                                                         typeRepls)...};
+        },
+        param);
+  }
+};
+
+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 {};
+
+/// 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();
+
+  // 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);
+    AttrTypeSubElementHandler<decltype(key)>::walk(key, walker);
+  }
+}
+
+/// 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,
+                                  ArrayRef<Type> &replTypes) {
+  auto key = static_cast<typename T::ImplType *>(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 T::Base::get(derived.getContext(),
+                                std::forward<decltype(params)>(params)...);
+          },
+          newKey);
+    } else {
+      return T::Base::get(derived.getContext(), newKey);
+    }
+  }
+}
+} // namespace detail
 } // namespace mlir
 
 /// Include the definitions of the sub elemnt interfaces.
 #include "mlir/IR/SubElementAttrInterfaces.h.inc"
 #include "mlir/IR/SubElementTypeInterfaces.h.inc"
 
 #endif // MLIR_IR_SUBELEMENTINTERFACES_H

mlir/include/mlir/IR/SubElementInterfaces.td

class SubElementInterfaceBase<string interfaceName, string attrOrType,
 
   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)
+           "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
-      )>,
+      }], 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

 
 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
+//===----------------------------------------------------------------------===//
+
+/// Enable TypeRange to be introspected for sub-elements.
+template <>
+struct AttrTypeSubElementHandler<TypeRange> {
+  static void walk(TypeRange param, AttrTypeSubElementWalker &walker) {
+    walker.walkRange(param);
+  }
+  static TypeRange replace(TypeRange param,
+                           AttrSubElementReplacements &attrRepls,
+                           TypeSubElementReplacements &typeRepls) {
+    return typeRepls.take_front(param.size());
+  }
+};
+
 } // namespace mlir
 
 namespace llvm {
 
 // Provide DenseMapInfo for TypeRange.
 template <>
 struct DenseMapInfo<mlir::TypeRange> {
   static mlir::TypeRange getEmptyKey() {

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

 // DITypeAttr
 //===----------------------------------------------------------------------===//
 
 bool DITypeAttr::classof(Attribute attr) {
   return llvm::isa<DIBasicTypeAttr, DISubroutineTypeAttr>(attr);
 }
 
 //===----------------------------------------------------------------------===//
-// DICompileUnitAttr
-//===----------------------------------------------------------------------===//
-
-void DICompileUnitAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getFile());
-  walkAttrsFn(getProducer());
-}
-
-Attribute
-DICompileUnitAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                               ArrayRef<Type> replTypes) const {
-  return get(getContext(), getSourceLanguage(), replAttrs[0].cast<DIFileAttr>(),
-             replAttrs[1].cast<StringAttr>(), getIsOptimized(),
-             getEmissionKind());
-}
-
-//===----------------------------------------------------------------------===//
-// DICompositeTypeAttr
-//===----------------------------------------------------------------------===//
-
-void DICompositeTypeAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getName());
-  walkAttrsFn(getFile());
-  walkAttrsFn(getScope());
-  for (DINodeAttr element : getElements())
-    walkAttrsFn(element);
-}
-
-Attribute DICompositeTypeAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  ArrayRef<Attribute> elements = replAttrs.drop_front(3);
-  return get(
-      getContext(), getTag(), replAttrs[0].cast<StringAttr>(),
-      cast_or_null<DIFileAttr>(replAttrs[1]), getLine(),
-      cast_or_null<DIScopeAttr>(replAttrs[2]), getSizeInBits(),
-      getAlignInBits(),
-      ArrayRef<DINodeAttr>(static_cast<const DINodeAttr *>(elements.data()),
-                           elements.size()));
-}
-
-//===----------------------------------------------------------------------===//
-// DIDerivedTypeAttr
-//===----------------------------------------------------------------------===//
-
-void DIDerivedTypeAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getName());
-  walkAttrsFn(getBaseType());
-}
-
-Attribute
-DIDerivedTypeAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                               ArrayRef<Type> replTypes) const {
-  return get(getContext(), getTag(), replAttrs[0].cast<StringAttr>(),
-             replAttrs[1].cast<DITypeAttr>(), getSizeInBits(), getAlignInBits(),
-             getOffsetInBits());
-}
-
-//===----------------------------------------------------------------------===//
-// DILexicalBlockAttr
-//===----------------------------------------------------------------------===//
-
-void DILexicalBlockAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getScope());
-  walkAttrsFn(getFile());
-}
-
-Attribute DILexicalBlockAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replAttrs[0].cast<DIScopeAttr>(), replAttrs[1].cast<DIFileAttr>(),
-             getLine(), getColumn());
-}
-
-//===----------------------------------------------------------------------===//
-// DILexicalBlockFileAttr
-//===----------------------------------------------------------------------===//
-
-void DILexicalBlockFileAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getScope());
-  walkAttrsFn(getFile());
-}
-
-Attribute DILexicalBlockFileAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replAttrs[0].cast<DIScopeAttr>(), replAttrs[1].cast<DIFileAttr>(),
-             getDescriminator());
-}
-
-//===----------------------------------------------------------------------===//
-// DILocalVariableAttr
-//===----------------------------------------------------------------------===//
-
-void DILocalVariableAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getScope());
-  walkAttrsFn(getName());
-  walkAttrsFn(getFile());
-  walkAttrsFn(getType());
-}
-
-Attribute DILocalVariableAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(getContext(), replAttrs[0].cast<DIScopeAttr>(),
-             replAttrs[1].cast<StringAttr>(), replAttrs[2].cast<DIFileAttr>(),
-             getLine(), getArg(), getAlignInBits(),
-             replAttrs[3].cast<DITypeAttr>());
-}
-
-//===----------------------------------------------------------------------===//
-// DISubprogramAttr
-//===----------------------------------------------------------------------===//
-
-void DISubprogramAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getCompileUnit());
-  walkAttrsFn(getScope());
-  walkAttrsFn(getName());
-  walkAttrsFn(getLinkageName());
-  walkAttrsFn(getFile());
-  walkAttrsFn(getType());
-}
-
-Attribute
-DISubprogramAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                              ArrayRef<Type> replTypes) const {
-  return get(getContext(), replAttrs[0].cast<DICompileUnitAttr>(),
-             replAttrs[1].cast<DIScopeAttr>(), replAttrs[2].cast<StringAttr>(),
-             replAttrs[3].cast<StringAttr>(), replAttrs[4].cast<DIFileAttr>(),
-             getLine(), getScopeLine(), getSubprogramFlags(),
-             replAttrs[5].cast<DISubroutineTypeAttr>());
-}
-
-//===----------------------------------------------------------------------===//
-// DISubroutineTypeAttr
-//===----------------------------------------------------------------------===//
-
-void DISubroutineTypeAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (DITypeAttr type : getTypes())
-    walkAttrsFn(type);
-}
-
-Attribute DISubroutineTypeAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(
-      getContext(), getCallingConvention(),
-      ArrayRef<DITypeAttr>(static_cast<const DITypeAttr *>(replAttrs.data()),
-                           replAttrs.size()));
-}
-
-//===----------------------------------------------------------------------===//
 // LoopOptionsAttrBuilder
 //===----------------------------------------------------------------------===//
 
 LoopOptionsAttrBuilder::LoopOptionsAttrBuilder(LoopOptionsAttr attr)
     : options(attr.getOptions().begin(), attr.getOptions().end()) {}
 
 template <typename T>
 LoopOptionsAttrBuilder &LoopOptionsAttrBuilder::setOption(LoopOptionCase tag,

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

 
 unsigned
 LLVMArrayType::getPreferredAlignment(const DataLayout &dataLayout,
                                      DataLayoutEntryListRef params) const {
   return dataLayout.getTypePreferredAlignment(getElementType());
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementTypeInterface
-
-void LLVMArrayType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type LLVMArrayType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes.front(), getNumElements());
-}
-
-//===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMFunctionType::isValidArgumentType(Type type) {
   return !type.isa<LLVMVoidType, LLVMFunctionType>();
 }
 
 bool LLVMFunctionType::isValidResultType(Type type) {
LLVMFunctionType::verify(function_ref<InFlightDiagnostic()> emitError,
   for (Type arg : arguments)
     if (!isValidArgumentType(arg))
       return emitError() << "invalid function argument type: " << arg;
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementTypeInterface
-
-void LLVMFunctionType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getReturnType());
-  for (Type type : getParams())
-    walkTypesFn(type);
-}
-
-Type LLVMFunctionType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes.front(), replTypes.drop_front(), isVarArg());
-}
-
-//===----------------------------------------------------------------------===//
 // LLVMPointerType
 //===----------------------------------------------------------------------===//
 
 bool LLVMPointerType::isValidElementType(Type type) {
   if (!type)
     return true;
   return isCompatibleOuterType(type)
              ? !type.isa<LLVMVoidType, LLVMTokenType, LLVMMetadataType,
if (extractPointerSpecValue(values, PtrDLEntryPos::Abi) >
       return emitError(loc) << "preferred alignment is expected to be at least "
                                "as large as ABI alignment";
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// SubElementTypeInterface
-
-void LLVMPointerType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type LLVMPointerType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(getContext(), replTypes.front(), getAddressSpace());
-}
-
-//===----------------------------------------------------------------------===//
 // Struct type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMStructType::isValidElementType(Type type) {
   return !type.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                    LLVMFunctionType, LLVMTokenType, LLVMScalableVectorType>();
 }
 

 
 LogicalResult
 LLVMFixedVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                             Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMFixedVectorType>(
       emitError, elementType, numElements);
 }
 
-void LLVMFixedVectorType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type LLVMFixedVectorType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes[0], getNumElements());
-}
-
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType.
 //===----------------------------------------------------------------------===//
 
 LLVMScalableVectorType LLVMScalableVectorType::get(Type elementType,
                                                    unsigned minNumElements) {
   assert(elementType && "expected non-null subtype");
   return Base::get(elementType.getContext(), elementType, minNumElements);

 
 LogicalResult
 LLVMScalableVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                                Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMScalableVectorType>(
       emitError, elementType, numElements);
 }
 
-void LLVMScalableVectorType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type LLVMScalableVectorType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes[0], getMinNumElements());
-}
-
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
 bool mlir::LLVM::isCompatibleOuterType(Type type) {
   // clang-format off
   if (type.isa<
       BFloat16Type,

mlir/lib/IR/BuiltinAttributes.cpp

 void BuiltinDialect::registerAttributes() {
   addAttributes<
 #define GET_ATTRDEF_LIST
 #include "mlir/IR/BuiltinAttributes.cpp.inc"
       >();
 }
 
 //===----------------------------------------------------------------------===//
-// ArrayAttr
-//===----------------------------------------------------------------------===//
-
-void ArrayAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (Attribute attr : getValue())
-    walkAttrsFn(attr);
-}
-
-Attribute
-ArrayAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                       ArrayRef<Type> replTypes) const {
-  return get(getContext(), replAttrs);
-}
-
-//===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 /// Helper function that does either an in place sort or sorts from source array
 /// into destination. If inPlace then storage is both the source and the
 /// destination, else value is the source and storage destination. Returns
 /// whether source was sorted.
 template <bool inPlace>
DictionaryAttr::iterator DictionaryAttr::end() const {
   return getValue().end();
 }
 size_t DictionaryAttr::size() const { return getValue().size(); }
 
 DictionaryAttr DictionaryAttr::getEmptyUnchecked(MLIRContext *context) {
   return Base::get(context, ArrayRef<NamedAttribute>());
 }
 
-void DictionaryAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (const NamedAttribute &attr : getValue())
-    walkAttrsFn(attr.getValue());
-}
-
-Attribute
-DictionaryAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                            ArrayRef<Type> replTypes) const {
-  std::vector<NamedAttribute> vec = getValue().vec();
-  for (auto &it : llvm::enumerate(replAttrs))
-    vec[it.index()].setValue(it.value());
-
-  // The above only modifies the mapped value, but not the key, and therefore
-  // not the order of the elements. It remains sorted
-  return getWithSorted(getContext(), vec);
-}
-
 //===----------------------------------------------------------------------===//
 // StridedLayoutAttr
 //===----------------------------------------------------------------------===//
 
 /// Prints a strided layout attribute.
 void StridedLayoutAttr::print(llvm::raw_ostream &os) const {
   auto printIntOrQuestion = [&](int64_t value) {
     if (value == ShapedType::kDynamicStrideOrOffset)
FlatSymbolRefAttr SymbolRefAttr::get(Operation *symbol) {
   return SymbolRefAttr::get(symName);
 }
 
 StringAttr SymbolRefAttr::getLeafReference() const {
   ArrayRef<FlatSymbolRefAttr> nestedRefs = getNestedReferences();
   return nestedRefs.empty() ? getRootReference() : nestedRefs.back().getAttr();
 }
 
-void SymbolRefAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getRootReference());
-  for (FlatSymbolRefAttr ref : getNestedReferences())
-    walkAttrsFn(ref);
-}
-
-Attribute
-SymbolRefAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                           ArrayRef<Type> replTypes) const {
-  ArrayRef<Attribute> rawNestedRefs = replAttrs.drop_front();
-  ArrayRef<FlatSymbolRefAttr> nestedRefs(
-      static_cast<const FlatSymbolRefAttr *>(rawNestedRefs.data()),
-      rawNestedRefs.size());
-  return get(replAttrs[0].cast<StringAttr>(), nestedRefs);
-}
-
 //===----------------------------------------------------------------------===//
 // IntegerAttr
 //===----------------------------------------------------------------------===//
 
 int64_t IntegerAttr::getInt() const {
   assert((getType().isIndex() || getType().isSignlessInteger()) &&
          "must be signless integer");
   return getValue().getSExtValue();
for (size_t i = 0, e = numSparseIndices; i != e; ++i) {
     if (!ElementsAttr::isValidIndex(type, index))
       return emitIndexError(i, index);
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
-// TypeAttr
-//===----------------------------------------------------------------------===//
-
-void TypeAttr::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getValue());
-}
-
-Attribute
-TypeAttr::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                      ArrayRef<Type> replTypes) const {
-  return get(replTypes[0]);
-}
-
-//===----------------------------------------------------------------------===//
 // Attribute Utilities
 //===----------------------------------------------------------------------===//
 
 AffineMap mlir::makeStridedLinearLayoutMap(ArrayRef<int64_t> strides,
                                            int64_t offset,
                                            MLIRContext *context) {
   AffineExpr expr;
   unsigned nSymbols = 0;

mlir/lib/IR/BuiltinTypes.cpp

FunctionType::getWithoutArgsAndResults(const BitVector &argIndices,
   SmallVector<Type> argStorage, resultStorage;
   TypeRange newArgTypes = function_interface_impl::filterTypesOut(
       getInputs(), argIndices, argStorage);
   TypeRange newResultTypes = function_interface_impl::filterTypesOut(
       getResults(), resultIndices, resultStorage);
   return clone(newArgTypes, newResultTypes);
 }
 
-void FunctionType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (Type type : llvm::concat<const Type>(getInputs(), getResults()))
-    walkTypesFn(type);
-}
-
-Type FunctionType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                               ArrayRef<Type> replTypes) const {
-  unsigned numInputs = getNumInputs();
-  return get(getContext(), replTypes.take_front(numInputs),
-             replTypes.drop_front(numInputs));
-}
-
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 /// Verify the construction of an opaque type.
 LogicalResult OpaqueType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  StringAttr dialect, StringRef typeData) {
   if (!Dialect::isValidNamespace(dialect.strref()))
if (auto et = getElementType().dyn_cast<IntegerType>())
     if (auto scaledEt = et.scaleElementBitwidth(scale))
       return VectorType::get(getShape(), scaledEt, getNumScalableDims());
   if (auto et = getElementType().dyn_cast<FloatType>())
     if (auto scaledEt = et.scaleElementBitwidth(scale))
       return VectorType::get(getShape(), scaledEt, getNumScalableDims());
   return VectorType();
 }
 
-void VectorType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type VectorType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                             ArrayRef<Type> replTypes) const {
-  return get(getShape(), replTypes.front(), getNumScalableDims());
-}
-
 VectorType VectorType::cloneWith(Optional<ArrayRef<int64_t>> shape,
                                  Type elementType) const {
   return VectorType::get(shape.value_or(getShape()), elementType,
                          getNumScalableDims());
 }
 
 //===----------------------------------------------------------------------===//
 // TensorType
for (int64_t s : shape)
     if (s < 0 && !ShapedType::isDynamic(s))
       return emitError() << "invalid tensor dimension size";
   if (auto v = encoding.dyn_cast_or_null<VerifiableTensorEncoding>())
     if (failed(v.verifyEncoding(shape, elementType, emitError)))
       return failure();
   return checkTensorElementType(emitError, elementType);
 }
 
-void RankedTensorType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-  if (Attribute encoding = getEncoding())
-    walkAttrsFn(encoding);
-}
-
-Type RankedTensorType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(getShape(), replTypes.front(),
-             replAttrs.empty() ? Attribute() : replAttrs.back());
-}
-
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 UnrankedTensorType::verify(function_ref<InFlightDiagnostic()> emitError,
                            Type elementType) {
   return checkTensorElementType(emitError, elementType);
 }
 
-void UnrankedTensorType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-}
-
-Type UnrankedTensorType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes.front());
-}
-
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 Type BaseMemRefType::getElementType() const {
   return llvm::TypeSwitch<BaseMemRefType, Type>(*this)
       .Case<MemRefType, UnrankedMemRefType>(
           [](auto type) { return type.getElementType(); });
if (failed(layout.verifyLayout(shape, emitError)))
     return failure();
 
   if (!isSupportedMemorySpace(memorySpace))
     return emitError() << "unsupported memory space Attribute";
 
   return success();
 }
 
-void MemRefType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-  if (!getLayout().isIdentity())
-    walkAttrsFn(getLayout());
-  walkAttrsFn(getMemorySpace());
-}
-
-Type MemRefType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                             ArrayRef<Type> replTypes) const {
-  bool hasLayout = replAttrs.size() > 1;
-  return get(getShape(), replTypes[0],
-             hasLayout ? replAttrs[0].dyn_cast<MemRefLayoutAttrInterface>()
-                       : MemRefLayoutAttrInterface(),
-             hasLayout ? replAttrs[1] : replAttrs[0]);
-}
-
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 unsigned UnrankedMemRefType::getMemorySpaceAsInt() const {
   return detail::getMemorySpaceAsInt(getMemorySpace());
 }
 
for (auto e : strideExprs) {
     if (auto c = e.dyn_cast<AffineConstantExpr>())
       strides.push_back(c.getValue());
     else
       strides.push_back(ShapedType::kDynamicStrideOrOffset);
   }
   return success();
 }
 
-void UnrankedMemRefType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkTypesFn(getElementType());
-  walkAttrsFn(getMemorySpace());
-}
-
-Type UnrankedMemRefType::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  return get(replTypes.front(), replAttrs.front());
-}
-
 //===----------------------------------------------------------------------===//
 /// TupleType
 //===----------------------------------------------------------------------===//
 
 /// Return the elements types for this tuple.
 ArrayRef<Type> TupleType::getTypes() const { return getImpl()->getTypes(); }
 
 /// Accumulate the types contained in this tuple and tuples nested within it.
 /// Note that this only flattens nested tuples, not any other container type,
 /// e.g. a tuple<i32, tensor<i32>, tuple<f32, tuple<i64>>> is flattened to
 /// (i32, tensor<i32>, f32, i64)
 void TupleType::getFlattenedTypes(SmallVectorImpl<Type> &types) {
   for (Type type : getTypes()) {
     if (auto nestedTuple = type.dyn_cast<TupleType>())
       nestedTuple.getFlattenedTypes(types);
     else
       types.push_back(type);
   }
 }
 
 /// Return the number of element types.
 size_t TupleType::size() const { return getImpl()->size(); }
 
-void TupleType::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (Type type : getTypes())
-    walkTypesFn(type);
-}
-
-Type TupleType::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                            ArrayRef<Type> replTypes) const {
-  return get(getContext(), replTypes);
-}
-
 //===----------------------------------------------------------------------===//
 // Type Utilities
 //===----------------------------------------------------------------------===//
 
 /// Return a version of `t` with identity layout if it can be determined
 /// statically that the layout is the canonical contiguous strided layout.
 /// Otherwise pass `t`'s layout into `simplifyAffineMap` and return a copy of
 /// `t` with simplified layout.

mlir/lib/IR/Location.cpp

 CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames) {
   assert(!frames.empty() && "required at least 1 call frame");
   Location caller = frames.back();
   for (auto frame : llvm::reverse(frames.drop_back()))
     caller = CallSiteLoc::get(frame, caller);
   return CallSiteLoc::get(name, caller);
 }
 
-void CallSiteLoc::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getCallee());
-  walkAttrsFn(getCaller());
-}
-
-Attribute
-CallSiteLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                         ArrayRef<Type> replTypes) const {
-  return get(replAttrs[0].cast<LocationAttr>(),
-             replAttrs[1].cast<LocationAttr>());
-}
-
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
 Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
                        MLIRContext *context) {
   // Unique the set of locations to be fused.
   llvm::SmallSetVector<Location, 4> decomposedLocs;
if (locs.empty()) {
     return Base::get(context, ArrayRef<Location>{UnknownLoc::get(context)},
                      metadata);
   }
   if (locs.size() == 1 && !metadata)
     return locs.front();
 
   return Base::get(context, locs, metadata);
 }
-
-void FusedLoc::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  for (Attribute attr : getLocations())
-    walkAttrsFn(attr);
-  walkAttrsFn(getMetadata());
-}
-
-Attribute
-FusedLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                      ArrayRef<Type> replTypes) const {
-  SmallVector<Location> newLocs;
-  newLocs.reserve(replAttrs.size() - 1);
-  for (Attribute attr : replAttrs.drop_back())
-    newLocs.push_back(attr.cast<LocationAttr>());
-  return get(getContext(), newLocs, replAttrs.back());
-}
-
-//===----------------------------------------------------------------------===//
-// NameLoc
-//===----------------------------------------------------------------------===//
-
-void NameLoc::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getName());
-  walkAttrsFn(getChildLoc());
-}
-
-Attribute NameLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                               ArrayRef<Type> replTypes) const {
-  return get(replAttrs[0].cast<StringAttr>(),
-             replAttrs[1].cast<LocationAttr>());
-}
-
-//===----------------------------------------------------------------------===//
-// OpaqueLoc
-//===----------------------------------------------------------------------===//
-
-void OpaqueLoc::walkImmediateSubElements(
-    function_ref<void(Attribute)> walkAttrsFn,
-    function_ref<void(Type)> walkTypesFn) const {
-  walkAttrsFn(getFallbackLocation());
-}
-
-Attribute
-OpaqueLoc::replaceImmediateSubElements(ArrayRef<Attribute> replAttrs,
-                                       ArrayRef<Type> replTypes) const {
-  return get(getUnderlyingLocation(), getUnderlyingTypeID(),
-             replAttrs[0].cast<LocationAttr>());
-}

mlir/lib/IR/SubElementInterfaces.cpp

 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;
         }
 
         // 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;
-
         // 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;
         }
 
         // Walk any sub elements first.
static void updateSubElementImpl(
     T element, function_ref<std::pair<T, WalkResult>(T)> walkFn,
     DenseMap<T, T> &visited, SmallVectorImpl<T> &newElements,
     FailureOr<bool> &changed, ReplaceSubElementFnT &&replaceSubElementFn) {
   // Bail early if we failed at any point.
   if (failed(changed))
     return;
   newElements.push_back(element);
 
-  // Guard against potentially null inputs. We always map null to null.
-  if (!element)
-    return;
-
   // Check for an existing mapping for this element, and walk it if we haven't
   // yet.
   T *mappedElement = &visited[element];
   if (!*mappedElement) {
     WalkResult result = WalkResult::advance();
     std::tie(*mappedElement, result) = walkFn(element);
 
     // Try walking this element.

mlir/lib/IR/TypeDetail.h

struct IntegerTypeStorage : public TypeStorage {
   }
 
   static IntegerTypeStorage *construct(TypeStorageAllocator &allocator,
                                        KeyTy key) {
     return new (allocator.allocate<IntegerTypeStorage>())
         IntegerTypeStorage(key.first, key.second);
   }
 
+  KeyTy getAsKey() const { return KeyTy(width, signedness); }
+
   unsigned width : 30;
   IntegerType::SignednessSemantics signedness : 2;
 };
 
 /// Function Type Storage and Uniquing.
 struct FunctionTypeStorage : public TypeStorage {
   FunctionTypeStorage(unsigned numInputs, unsigned numResults,
                       Type const *inputsAndResults)
       : numInputs(numInputs), numResults(numResults),
         inputsAndResults(inputsAndResults) {}
 
   /// The hash key used for uniquing.
-  using KeyTy = std::pair<TypeRange, TypeRange>;
+  using KeyTy = std::tuple<TypeRange, TypeRange>;
   bool operator==(const KeyTy &key) const {
     if (std::get<0>(key) == getInputs())
       return std::get<1>(key) == getResults();
     return false;
   }
 
   /// Construction.
   static FunctionTypeStorage *construct(TypeStorageAllocator &allocator,
                                         const KeyTy &key) {
-    TypeRange inputs = key.first, results = key.second;
+    auto [inputs, results] = key;
 
     // Copy the inputs and results into the bump pointer.
     SmallVector<Type, 16> types;
     types.reserve(inputs.size() + results.size());
     types.append(inputs.begin(), inputs.end());
     types.append(results.begin(), results.end());
     auto typesList = allocator.copyInto(ArrayRef<Type>(types));
 
     // Initialize the memory using placement new.
     return new (allocator.allocate<FunctionTypeStorage>())
         FunctionTypeStorage(inputs.size(), results.size(), typesList.data());
   }
 
   ArrayRef<Type> getInputs() const {
     return ArrayRef<Type>(inputsAndResults, numInputs);
   }
   ArrayRef<Type> getResults() const {
     return ArrayRef<Type>(inputsAndResults + numInputs, numResults);
   }
 
+  KeyTy getAsKey() const { return KeyTy(getInputs(), getResults()); }
+
   unsigned numInputs;
   unsigned numResults;
   Type const *inputsAndResults;
 };
 
 /// A type representing a collection of other types.
 struct TupleTypeStorage final
     : public TypeStorage,
struct TupleTypeStorage final
   /// Return the number of held types.
   unsigned size() const { return numElements; }
 
   /// Return the held types.
   ArrayRef<Type> getTypes() const {
     return {getTrailingObjects<Type>(), size()};
   }
 
+  KeyTy getAsKey() const { return getTypes(); }
+
   /// The number of tuple elements.
   unsigned numElements;
 };
 
 /// Checks if the memorySpace has supported Attribute type.
 bool isSupportedMemorySpace(Attribute memorySpace);
 
 /// Wraps deprecated integer memory space to the new Attribute form.

mlir/test/Dialect/Affine/loop-tiling.mlir

affine.for %i0 = 0 to %0 {
     }
   }
   return
 }
 
 // CHECK:       memref.dim %{{.*}}, %c0 : memref<?x?xf32>
 // CHECK-NEXT:  affine.for %{{.*}} = 0 to %{{.*}} step 32 {
 // CHECK-NEXT:    affine.for %{{.*}} = 0 to %{{.*}} step 32 {
-// CHECK-NEXT:      affine.for %{{.*}} = #map(%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
-// CHECK-NEXT:        affine.for %{{.*}} = #map(%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
+// CHECK-NEXT:      affine.for %{{.*}} = #[[$MAP:.*]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
+// CHECK-NEXT:        affine.for %{{.*}} = #[[$MAP]](%{{.*}}) to min [[$UBMAP]](%{{.*}})[%{{.*}}] {
 // CHECK-NEXT:          affine.store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
 // CHECK-NEXT:          affine.for %{{.*}} = 0 to %{{.*}} {
 // CHECK-NEXT:            affine.load
 // CHECK-NEXT:            affine.load
 // CHECK-NEXT:            arith.mulf
 // CHECK-NEXT:            affine.load
 // CHECK-NEXT:            arith.addf
 // CHECK-NEXT:            affine.store

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

linalg.generic #trait
      ins(%arg0, %arg1 : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32)
     outs(%shape : memref<?x1x?x1x?xf32>) {
        ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :
          linalg.yield %arg3 : f32
        }
   return %shape : memref<?x1x?x1x?xf32>
 }
 
-// CHECK:     #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
-// CHECK:     #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
-// CHECK:     #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+// CHECK-DAG:     #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>
+// CHECK-DAG:     #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>
+// CHECK-DAG:     #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
 // CHECK:     func @input_stays_same(
 // CHECK-SAME:  %[[ARG0:.*]]: memref<?x1x?xf32, strided<[?, 1, 1]>>,
 // CHECK-SAME:  %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<?x1x?x1x?xf32>)
 // CHECK-SAME   -> memref<?x1x?x1x?xf32> {
 // CHECK:      %[[OUT:.*]] = memref.collapse_shape %[[ARG2]] {{\[}}[0, 1], [2, 3], [4]]
 // CHECK-SAME:   : memref<?x1x?x1x?xf32> into memref<?x?x?xf32>
 // CHECK:      linalg.generic
 // CHECK-SAME:   {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]],

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

func.func @fold_static_stride_subview_with_affine_load_store_expand_shape_when_access_index_is_an_expression(%arg0: memref<1024x1024xf32>, %arg1: memref<1xf32>, %arg2: index) -> f32 {
   return %2 : f32
 }
 // CHECK-NEXT: affine.for %[[ARG3:.*]] = 0 to 1 {
 // CHECK-NEXT:  affine.for %[[ARG4:.*]] = 0 to 1024 {
 // CHECK-NEXT:   affine.for %[[ARG5:.*]] = 0 to 1020 {
 // CHECK-NEXT:    affine.for %[[ARG6:.*]] = 0 to 1 {
 // CHECK-NEXT:      %[[TMP1:.*]] = affine.apply #[[$MAP0]](%[[ARG3]], %[[ARG4]], %[[ARG5]], %[[ARG6]])
 // CHECK-NEXT:      %[[TMP2:.*]] = affine.apply #[[$MAP1]](%[[ARG3]], %[[TMP1]])
-// CHECK-NEXT:      %[[TMP3:.*]] = affine.apply #map2(%[[ARG5]], %[[ARG6]])
+// CHECK-NEXT:      %[[TMP3:.*]] = affine.apply #{{.*}}(%[[ARG5]], %[[ARG6]])
 // CHECK-NEXT:      affine.load %[[ARG0]][%[[TMP2]], %[[TMP3]]] : memref<1024x1024xf32>
 
 // -----
 
 // CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0 * 1024 + d1)>
 // CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0 + d1)>
 // CHECK-LABEL: fold_static_stride_subview_with_affine_load_store_expand_shape_with_constant_access_index
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1xf32>, %[[ARG2:.*]]: index)

mlir/test/Dialect/SCF/for-loop-specialization.mlir

func.func @for(%outer: index, %A: memref<?xf32>, %B: memref<?xf32>,
   return
 }
 
 // CHECK-LABEL:   func @for(
 // CHECK-SAME:              [[ARG0:%.*]]: index, [[ARG1:%.*]]: memref<?xf32>, [[ARG2:%.*]]: memref<?xf32>, [[ARG3:%.*]]: memref<?xf32>, [[ARG4:%.*]]: memref<?xf32>) {
 // CHECK:           [[CST_0:%.*]] = arith.constant 0 : index
 // CHECK:           [[CST_1:%.*]] = arith.constant 1 : index
 // CHECK:           [[DIM_0:%.*]] = memref.dim [[ARG1]], [[CST_0]] : memref<?xf32>
-// CHECK:           [[MIN:%.*]] = affine.min #map(){{\[}}[[DIM_0]], [[ARG0]]]
+// CHECK:           [[MIN:%.*]] = affine.min #{{.*}}(){{\[}}[[DIM_0]], [[ARG0]]]
 // CHECK:           [[CST_1024:%.*]] = arith.constant 1024 : index
 // CHECK:           [[PRED:%.*]] = arith.cmpi eq, [[MIN]], [[CST_1024]] : index
 // CHECK:           scf.if [[PRED]] {
 // CHECK:             scf.for [[IDX0:%.*]] = [[CST_0]] to [[CST_1024]] step [[CST_1]] {
 // CHECK:               memref.store
 // CHECK:             }
 // CHECK:           } else {
 // CHECK:             scf.for [[IDX0:%.*]] = [[CST_0]] to [[MIN]] step [[CST_1]] {
 // CHECK:               memref.store
 // CHECK:             }
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }

mlir/test/Dialect/SCF/parallel-loop-specialization.mlir

 }
 
 // CHECK-LABEL:   func @parallel_loop(
 // CHECK-SAME:                        [[VAL_0:%.*]]: index, [[VAL_1:%.*]]: index, [[VAL_2:%.*]]: memref<?x?xf32>, [[VAL_3:%.*]]: memref<?x?xf32>, [[VAL_4:%.*]]: memref<?x?xf32>, [[VAL_5:%.*]]: memref<?x?xf32>) {
 // CHECK:           [[VAL_6:%.*]] = arith.constant 0 : index
 // CHECK:           [[VAL_7:%.*]] = arith.constant 1 : index
 // CHECK:           [[VAL_8:%.*]] = memref.dim [[VAL_2]], [[VAL_6]] : memref<?x?xf32>
 // CHECK:           [[VAL_9:%.*]] = memref.dim [[VAL_2]], [[VAL_7]] : memref<?x?xf32>
-// CHECK:           [[VAL_10:%.*]] = affine.min #map(){{\[}}[[VAL_8]], [[VAL_0]]]
-// CHECK:           [[VAL_11:%.*]] = affine.min #map1(){{\[}}[[VAL_9]], [[VAL_1]]]
+// CHECK:           [[VAL_10:%.*]] = affine.min #{{.*}}(){{\[}}[[VAL_8]], [[VAL_0]]]
+// CHECK:           [[VAL_11:%.*]] = affine.min #{{.*}}(){{\[}}[[VAL_9]], [[VAL_1]]]
 // CHECK:           [[VAL_12:%.*]] = arith.constant 1024 : index
 // CHECK:           [[VAL_13:%.*]] = arith.cmpi eq, [[VAL_10]], [[VAL_12]] : index
 // CHECK:           [[VAL_14:%.*]] = arith.constant 64 : index
 // CHECK:           [[VAL_15:%.*]] = arith.cmpi eq, [[VAL_11]], [[VAL_14]] : index
 // CHECK:           [[VAL_16:%.*]] = arith.andi [[VAL_13]], [[VAL_15]] : i1
 // CHECK:           scf.if [[VAL_16]] {
 // CHECK:             scf.parallel ([[VAL_17:%.*]], [[VAL_18:%.*]]) = ([[VAL_6]], [[VAL_6]]) to ([[VAL_12]], [[VAL_14]]) step ([[VAL_7]], [[VAL_7]]) {
 // CHECK:               memref.store
 // CHECK:             }
 // CHECK:           } else {
 // CHECK:             scf.parallel ([[VAL_22:%.*]], [[VAL_23:%.*]]) = ([[VAL_6]], [[VAL_6]]) to ([[VAL_10]], [[VAL_11]]) step ([[VAL_7]], [[VAL_7]]) {
 // CHECK:               memref.store
 // CHECK:             }
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }

mlir/test/Dialect/SCF/parallel-loop-tiling.mlir

 // RUN: mlir-opt %s -pass-pipeline='builtin.module(func.func(scf-parallel-loop-tiling{parallel-loop-tile-sizes=1,4}))' -split-input-file | FileCheck %s
 
 func.func @parallel_loop(%arg0 : index, %arg1 : index, %arg2 : index,
                     %arg3 : index, %arg4 : index, %arg5 : index,
         %A: memref<?x?xf32>, %B: memref<?x?xf32>,
                     %C: memref<?x?xf32>, %result: memref<?x?xf32>) {
   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
     %B_elem = memref.load %B[%i0, %i1] : memref<?x?xf32>
     %C_elem = memref.load %C[%i0, %i1] : memref<?x?xf32>
     %sum_elem = arith.addf %B_elem, %C_elem : f32
     memref.store %sum_elem, %result[%i0, %i1] : memref<?x?xf32>
   }
   return
 }
 
-// CHECK:       #map = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
+// CHECK:       #[[$MAP:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
 // CHECK-LABEL:   func @parallel_loop(
 // CHECK-SAME:                        [[ARG1:%.*]]: index, [[ARG2:%.*]]: index, [[ARG3:%.*]]: index, [[ARG4:%.*]]: index, [[ARG5:%.*]]: index, [[ARG6:%.*]]: index, [[ARG7:%.*]]: memref<?x?xf32>, [[ARG8:%.*]]: memref<?x?xf32>, [[ARG9:%.*]]: memref<?x?xf32>, [[ARG10:%.*]]: memref<?x?xf32>) {
 // CHECK:           [[C0:%.*]] = arith.constant 0 : index
 // CHECK:           [[C1:%.*]] = arith.constant 1 : index
 // CHECK:           [[C4:%.*]] = arith.constant 4 : index
 // CHECK:           [[V1:%.*]] = arith.muli [[ARG5]], [[C1]] : index
 // CHECK:           [[V2:%.*]] = arith.muli [[ARG6]], [[C4]] : index
 // CHECK:           scf.parallel ([[V3:%.*]], [[V4:%.*]]) = ([[ARG1]], [[ARG2]]) to ([[ARG3]], [[ARG4]]) step ([[V1]], [[V2]]) {
-// CHECK:             [[V5:%.*]] = affine.min #map([[V1]], [[ARG3]], [[V3]])
-// CHECK:             [[V6:%.*]] = affine.min #map([[V2]], [[ARG4]], [[V4]])
+// CHECK:             [[V5:%.*]] = affine.min #[[$MAP]]([[V1]], [[ARG3]], [[V3]])
+// CHECK:             [[V6:%.*]] = affine.min #[[$MAP]]([[V2]], [[ARG4]], [[V4]])
 // CHECK:             scf.parallel ([[V7:%.*]], [[V8:%.*]]) = ([[C0]], [[C0]]) to ([[V5]], [[V6]]) step ([[ARG5]], [[ARG6]]) {
 // CHECK:               [[V9:%.*]] = arith.addi [[V7]], [[V3]] : index
 // CHECK:               [[V10:%.*]] = arith.addi [[V8]], [[V4]] : index
 // CHECK:               [[V11:%.*]] = memref.load [[ARG8]]{{\[}}[[V9]], [[V10]]] : memref<?x?xf32>
 // CHECK:               [[V12:%.*]] = memref.load [[ARG9]]{{\[}}[[V9]], [[V10]]] : memref<?x?xf32>
 // CHECK:               [[V13:%.*]] = arith.addf [[V11]], [[V12]] : f32
 // CHECK:               memref.store [[V13]], [[ARG10]]{{\[}}[[V9]], [[V10]]] : memref<?x?xf32>
 // CHECK:             }

 // CHECK:           [[C1:%.*]] = arith.constant 1 : index
 // CHECK:           scf.parallel ([[V1:%.*]], [[V2:%.*]]) = ([[C0]], [[C0]]) to ([[C2]], [[C2]]) step ([[C1]], [[C1]]) {
 // CHECK:             [[C0_1:%.*]] = arith.constant 0 : index
 // CHECK:             [[C1_1:%.*]] = arith.constant 1 : index
 // CHECK:             [[C4:%.*]] = arith.constant 4 : index
 // CHECK:             [[V3:%.*]] = arith.muli [[C1]], [[C1_1]] : index
 // CHECK:             [[V4:%.*]] = arith.muli [[C1]], [[C4]] : index
 // CHECK:             scf.parallel ([[V5:%.*]], [[V6:%.*]]) = ([[C0]], [[C0]]) to ([[C2]], [[C2]]) step ([[V3]], [[V4]]) {
-// CHECK:               [[V7:%.*]] = affine.min #map([[V4]], [[C2]], [[V6]])
+// CHECK:               [[V7:%.*]] = affine.min #{{.*}}([[V4]], [[C2]], [[V6]])
 // CHECK:               scf.parallel ([[V8:%.*]], [[V9:%.*]]) = ([[C0_1]], [[C0_1]]) to ([[V3]], [[V7]]) step ([[C1]], [[C1]]) {
 // CHECK:                 = arith.addi [[V8]], [[V5]] : index
 // CHECK:                 = arith.addi [[V9]], [[V6]] : index
 // CHECK:               }
 // CHECK:             }
 // CHECK:           }
 // CHECK:           [[C0_2:%.*]] = arith.constant 0 : index
 // CHECK:           [[C1_2:%.*]] = arith.constant 1 : index
 // CHECK:           [[C4_1:%.*]] = arith.constant 4 : index
 // CHECK:           [[V10:%.*]] = arith.muli [[C1]], [[C1_2]] : index
 // CHECK:           [[V11:%.*]] = arith.muli [[C1]], [[C4_1]] : index
 // CHECK:           scf.parallel ([[V12:%.*]], [[V13:%.*]]) = ([[C0]], [[C0]]) to ([[C2]], [[C2]]) step ([[V10]], [[V11]]) {
-// CHECK:             [[V14:%.*]] = affine.min #map([[V11]], [[C2]], [[V13]])
+// CHECK:             [[V14:%.*]] = affine.min #{{.*}}([[V11]], [[C2]], [[V13]])
 // CHECK:             scf.parallel ([[V15:%.*]], [[V16:%.*]]) = ([[C0_2]], [[C0_2]]) to ([[V10]], [[V14]]) step ([[C1]], [[C1]]) {
 // CHECK:               = arith.addi [[V15]], [[V12]] : index
 // CHECK:               = arith.addi [[V16]], [[V13]] : index
 // CHECK:             }
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 

mlir/test/Dialect/SparseTensor/roundtrip.mlir

-// RUN: mlir-opt %s -split-input-file | mlir-opt | FileCheck %s
+// RUN: mlir-opt %s -split-input-file | mlir-opt -split-input-file | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
 
 // CHECK-LABEL: func @sparse_new(
 // CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[T:.*]] = sparse_tensor.new %[[A]] : !llvm.ptr<i8> to tensor<128xf64, #{{.*}}>
 //       CHECK: return %[[T]] : tensor<128xf64, #{{.*}}>
 func.func @sparse_new(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {

mlir/test/IR/affine-map.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
 
 // Identity maps used in trivial compositions in MemRefs are optimized away.
-// CHECK-NOT: #map{{[0-9]*}} = affine_map<(d0, d1) -> (d0, d1)>
 #map0 = affine_map<(i, j) -> (i, j)>
-
-// CHECK-NOT: #map{{[0-9]*}} = affine_map<(d0, d1)[s0] -> (d0, d1)>
 #map1 = affine_map<(i, j)[s0] -> (i, j)>
 
 // CHECK: #map{{[0-9]*}} = affine_map<() -> (0)>
 // A map may have 0 inputs.
 // However, an affine.apply always takes at least one input.
 #map2 = affine_map<() -> (0)>
 
 // All the maps in the following block are equivalent and are unique'd as one

 // CHECK: #map{{[0-9]*}} = affine_map<(d0, d1) -> (d0 * 3, (d0 + d1) * 2, d0 mod 2)>
 #map58 = affine_map<(d0, d1) -> (4*d0 - 2*d0 + d0, (d0 + d1) + (d0 + d1), 2 * (d0 mod 2) - d0 mod 2)>
 
 // CHECK: #map{{[0-9]*}} = affine_map<(d0, d1) -> (d0 mod 5, (d1 mod 35) mod 4)>
 #map59 = affine_map<(d0, d1) -> ((d0 mod 35) mod 5, (d1 mod 35) mod 4)>
 
 // Check if parser can parse affine_map with identifiers that collide with
 // integer types.
-// CHECK: #map{{[0-9]*}} = affine_map<(d0, d1) -> (d0, d1)>
 #map60 = affine_map<(i0, i1) -> (i0, i1)>
 
 // Check if parser can parse affine_map with identifiers that collide with
 // reserved keywords.
 // CHECK: #map{{[0-9]*}} = affine_map<(d0, d1)[s0, s1] -> (-d0 + s0, -d1 + s1)>
 #map61 = affine_map<(d0, d1)[step, loc] -> (step - d0, loc - d1)>
 
 // CHECK: #map{{[0-9]*}} = affine_map<(d0, d1)[s0, s1] -> (-d0 + s0 floordiv 2, -d1 + s1 mod 3)>

mlir/test/IR/memory-ops.mlir

 // RUN: mlir-opt %s | FileCheck %s
 
-// CHECK: #map = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>
+// CHECK: #[[$MAP:.*]] = affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>
 
 // CHECK-LABEL: func @alloc() {
 func.func @alloc() {
 ^bb0:
   // Test simple alloc.
   // CHECK: %{{.*}} = memref.alloc() : memref<1024x64xf32, 1>
   %0 = memref.alloc() : memref<1024x64xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
 
   %c0 = "arith.constant"() {value = 0: index} : () -> index
   %c1 = "arith.constant"() {value = 1: index} : () -> index
 
   // Test alloc with dynamic dimensions.
   // CHECK: %{{.*}} = memref.alloc(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
   %1 = memref.alloc(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
 
   // Test alloc with no dynamic dimensions and one symbol.
-  // CHECK: %{{.*}} = memref.alloc()[%{{.*}}] : memref<2x4xf32, #map, 1>
+  // CHECK: %{{.*}} = memref.alloc()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
   %2 = memref.alloc()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
 
   // Test alloc with dynamic dimensions and one symbol.
-  // CHECK: %{{.*}} = memref.alloc(%{{.*}})[%{{.*}}] : memref<2x?xf32, #map, 1>
+  // CHECK: %{{.*}} = memref.alloc(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
   %3 = memref.alloc(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
 
   // Alloc with no mappings.
   // b/116054838 Parser crash while parsing ill-formed AllocOp
   // CHECK: %{{.*}} = memref.alloc() : memref<2xi32>
   %4 = memref.alloc() : memref<2 x i32>
 
   // CHECK:   return
^bb0:
   %c0 = "arith.constant"() {value = 0: index} : () -> index
   %c1 = "arith.constant"() {value = 1: index} : () -> index
 
   // Test alloca with dynamic dimensions.
   // CHECK: %{{.*}} = memref.alloca(%{{.*}}, %{{.*}}) : memref<?x?xf32, 1>
   %1 = memref.alloca(%c0, %c1) : memref<?x?xf32, affine_map<(d0, d1) -> (d0, d1)>, 1>
 
   // Test alloca with no dynamic dimensions and one symbol.
-  // CHECK: %{{.*}} = memref.alloca()[%{{.*}}] : memref<2x4xf32, #map, 1>
+  // CHECK: %{{.*}} = memref.alloca()[%{{.*}}] : memref<2x4xf32, #[[$MAP]], 1>
   %2 = memref.alloca()[%c0] : memref<2x4xf32, affine_map<(d0, d1)[s0] -> ((d0 + s0), d1)>, 1>
 
   // Test alloca with dynamic dimensions and one symbol.
-  // CHECK: %{{.*}} = memref.alloca(%{{.*}})[%{{.*}}] : memref<2x?xf32, #map, 1>
+  // CHECK: %{{.*}} = memref.alloca(%{{.*}})[%{{.*}}] : memref<2x?xf32, #[[$MAP]], 1>
   %3 = memref.alloca(%c1)[%c0] : memref<2x?xf32, affine_map<(d0, d1)[s0] -> (d0 + s0, d1)>, 1>
 
   // Alloca with no mappings, but with alignment.
   // CHECK: %{{.*}} = memref.alloca() {alignment = 64 : i64} : memref<2xi32>
   %4 = memref.alloca() {alignment = 64} : memref<2 x i32>
 
   return
 }

mlir/test/Transforms/loop-fusion-2.mlir

affine.for %i7 = 0 to 9 {
           %10 = affine.apply #map(%i6, %i10)
           %11 = affine.load %1[%10, %i8] : memref<144x4xf32>
         }
       }
     }
   }
   return
 }
-// MAXIMAL:      #map = affine_map<(d0, d1) -> (d0 * 16 + d1)>
+// MAXIMAL:      #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0 * 16 + d1)>
 // MAXIMAL-LABEL: func @fuse_across_dim_mismatch
 // MAXIMAL:        memref.alloc() : memref<1x1xf32>
 // MAXIMAL:        affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:    affine.for %{{.*}} = 0 to 9 {
 // MAXIMAL-NEXT:      affine.for %{{.*}} = 0 to 4 {
 // MAXIMAL-NEXT:        affine.for %{{.*}} = 0 to 16 {
-// MAXIMAL-NEXT:          affine.apply #map(%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply #[[$MAP]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.store %{{.*}}, %{{.*}}[0, 0] : memref<1x1xf32>
-// MAXIMAL-NEXT:          affine.apply #map(%{{.*}}, %{{.*}})
+// MAXIMAL-NEXT:          affine.apply #[[$MAP]](%{{.*}}, %{{.*}})
 // MAXIMAL-NEXT:          affine.load %{{.*}}[0, 0] : memref<1x1xf32>
 // MAXIMAL-NEXT:        }
 // MAXIMAL-NEXT:      }
 // MAXIMAL-NEXT:    }
 // MAXIMAL-NEXT:  }
 
 // -----
 

mlir/test/Transforms/normalize-memrefs-ops.mlir

func.func @test_norm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     // CHECK: "test.op_norm"(%[[ARG0]], %[[v0]]) : (memref<1x16x1x1x32x64xf32>, memref<1x16x1x1x32x64xf32>) -> ()
     // CHECK: memref.dealloc %[[v0]] : memref<1x16x1x1x32x64xf32>
     return
 }
 
 // Same test with op_nonnorm, with maps in the arguments and the operations in the function.
 
 // CHECK-LABEL: test_nonnorm
-// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #map>)
+// CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x14x14xf32, #[[MAP:.*]]>)
 func.func @test_nonnorm(%arg0 : memref<1x16x14x14xf32, #map0>) -> () {
     %0 = memref.alloc() : memref<1x16x14x14xf32, #map0>
     "test.op_nonnorm"(%arg0, %0) : (memref<1x16x14x14xf32, #map0>, memref<1x16x14x14xf32, #map0>) -> ()
     memref.dealloc %0 :  memref<1x16x14x14xf32, #map0>
 
-    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32, #map>
-    // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #map>, memref<1x16x14x14xf32, #map>) -> ()
-    // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32, #map>
+    // CHECK: %[[v0:.*]] = memref.alloc() : memref<1x16x14x14xf32, #[[MAP]]>
+    // CHECK: "test.op_nonnorm"(%[[ARG0]], %[[v0]]) : (memref<1x16x14x14xf32, #[[MAP]]>, memref<1x16x14x14xf32, #[[MAP]]>) -> ()
+    // CHECK: memref.dealloc %[[v0]] : memref<1x16x14x14xf32, #[[MAP]]>
     return
 }
 
 // Test with op_nonnorm whose memref map layouts are identity. This op_nonnorm
 // does not block the normalization of other operations.
 
 // CHECK-LABEL: test_nonnorm_identity_layout
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x16x1x1x32x64xf32>)

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

auto try_value_begin_impl(OverloadToken<llvm::APInt>) const {
       }).begin());
     }
   }];
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
-    DeclareAttrInterfaceMethods<SubElementAttrInterface,
-        ["replaceImmediateSubElements"]>
+    SubElementAttrInterface
   ]> {
   let mnemonic = "sub_elements_access";
 
   let parameters = (ins
     "::mlir::Attribute":$first,
     "::mlir::Attribute":$second,
     "::mlir::Attribute":$third
   );

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

Attribute TestSubElementsAccessAttr::parse(::mlir::AsmParser &parser,
   return get(parser.getContext(), first, second, third);
 }
 
 void TestSubElementsAccessAttr::print(::mlir::AsmPrinter &printer) const {
   printer << "<" << getFirst() << ", " << getSecond() << ", " << getThird()
           << ">";
 }
 
-void TestSubElementsAccessAttr::walkImmediateSubElements(
-    llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
-    llvm::function_ref<void(mlir::Type)> walkTypesFn) const {
-  walkAttrsFn(getFirst());
-  walkAttrsFn(getSecond());
-  walkAttrsFn(getThird());
-}
-
-Attribute TestSubElementsAccessAttr::replaceImmediateSubElements(
-    ArrayRef<Attribute> replAttrs, ArrayRef<Type> replTypes) const {
-  assert(replAttrs.size() == 3 && "invalid number of replacement attributes");
-  return get(getContext(), replAttrs[0], replAttrs[1], replAttrs[2]);
-}
-
 //===----------------------------------------------------------------------===//
 // TestExtern1DI64ElementsAttr
 //===----------------------------------------------------------------------===//
 
 ArrayRef<uint64_t> TestExtern1DI64ElementsAttr::getElements() const {
   if (auto *blob = getHandle().getBlob())
     return blob->getDataAs<uint64_t>();
   return llvm::None;

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

 
 void DefGen::emitKeyType() {
   std::string keyType("std::tuple<");
   llvm::raw_string_ostream os(keyType);
   llvm::interleaveComma(params, os,
                         [&](auto &param) { os << param.getCppType(); });
   os << '>';
   storageCls->declare<UsingDeclaration>("KeyTy", std::move(os.str()));
+
+  // Add a method to construct the key type from the storage.
+  Method *m = storageCls->addConstMethod<Method::Inline>("KeyTy", "getAsKey");
+  m->body().indent() << "return KeyTy(";
+  llvm::interleaveComma(params, m->body().indent(),
+                        [&](auto &param) { m->body() << param.getName(); });
+  m->body() << ");";
 }
 
 void DefGen::emitEquals() {
   Method *eq = storageCls->addConstMethod<Method::Inline>(
       "bool", "operator==", MethodParameter("const KeyTy &", "tblgenKey"));
   auto &body = eq->body().indent();
   auto scope = body.scope("return (", ");");
   const auto eachFn = [&](auto it) {

mlir/unittests/IR/SubElementInterfaceTest.cpp

 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("array", boolArrayAttr));
+      builder.getDictionaryAttr(builder.getNamedAttr(strAttr, boolArrayAttr));
 
   SmallVector<Attribute> subAttrs;
   dictAttr.walkSubAttrs([&](Attribute attr) { subAttrs.push_back(attr); });
   EXPECT_EQ(llvm::makeArrayRef(subAttrs),
-            ArrayRef<Attribute>({trueAttr, falseAttr, boolArrayAttr}));
+            ArrayRef<Attribute>({strAttr, trueAttr, falseAttr, boolArrayAttr}));
 }
 
 } // namespace