[mlir] Add new SubElementAttr/SubElementType Interfaces

[rriddle]

Summary

These interfaces allow for a composite attribute or type to opaquely provide access to any held attributes or types. There are several intended use cases for this interface. The first of which is to allow the printer to create aliases for non-builtin dialect attributes and types. In the future, this interface will also be extended to allow for SymbolRefAttr to be placed on other entities aside from just DictionaryAttr and ArrayAttr.

To limit potential test breakages, this revision only adds the new interfaces to the builtin attributes/types that are currently hardcoded during AsmPrinter alias generation. In a followup the remaining builtin attributes/types, and non-builtin attributes/types can be extended to support it.

[jpienaar]

Could you add a C++ unit test for this?

mlir/include/mlir/IR/SubElementInterfaces.td

33

Why the Impl suffix?

[rriddle]

update

[rriddle]

In D102945#2796118, @jpienaar wrote:

Could you add a C++ unit test for this?

Done.

mlir/include/mlir/IR/SubElementInterfaces.td

33

I wouldn't see this as something generally called by users, given that it only walks one level. Renamed to the more appropriate: walkImmediateSubElements.

[jpienaar]

mlir/include/mlir/IR/SubElementInterfaces.td

31

immediate ?

mlir/include/mlir/IR/BuiltinAttributes.h

 //===- BuiltinAttributes.h - MLIR Builtin Attribute Classes -----*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINATTRIBUTES_H
 #define MLIR_IR_BUILTINATTRIBUTES_H
 
-#include "mlir/IR/Attributes.h"
+#include "SubElementInterfaces.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Sequence.h"
 #include <complex>
 
 namespace mlir {
 class AffineMap;
 class BoolAttr;
 class DenseIntElementsAttr;

mlir/include/mlir/IR/BuiltinAttributes.td

 // validity of and defining the IR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUILTIN_ATTRIBUTES
 #define BUILTIN_ATTRIBUTES
 
 include "mlir/IR/BuiltinDialect.td"
+include "mlir/IR/SubElementInterfaces.td"
 
 // TODO: Currently the attributes defined in this file are prefixed with
 // `Builtin_`.  This is to differentiate the attributes here with the ones in
 // OpBase.td. We should remove the definitions in OpBase.td, and repoint users
 // to this file instead.
 
 // Base class for Builtin dialect attributes.
-class Builtin_Attr<string name, string baseCppClass = "::mlir::Attribute">
-    : AttrDef<Builtin_Dialect, name, [], baseCppClass> {
+class Builtin_Attr<string name, list<Trait> traits = [],
+                   string baseCppClass = "::mlir::Attribute">
+    : AttrDef<Builtin_Dialect, name, traits, baseCppClass> {
   let mnemonic = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // AffineMapAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap"> {
def Builtin_AffineMapAttr : Builtin_Attr<"AffineMap"> {
   let skipDefaultBuilders = 1;
   let typeBuilder = "IndexType::get($_value.getContext())";
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_ArrayAttr : Builtin_Attr<"Array"> {
+def Builtin_ArrayAttr : Builtin_Attr<"Array", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "A collection of other Attribute values";
   let description = [{
     Syntax:
 
     ```
     array-attribute ::= `[` (attribute-value (`,` attribute-value)*)? `]`
     ```
 
def Builtin_ArrayAttr : Builtin_Attr<"Array", [
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // DenseIntOrFPElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseIntOrFPElementsAttr
-    : Builtin_Attr<"DenseIntOrFPElements", "DenseElementsAttr"> {
+    : Builtin_Attr<"DenseIntOrFPElements", /*traits=*/[], "DenseElementsAttr"> {
   let summary = "An Attribute containing a dense multi-dimensional array of "
                 "integer or floating-point values";
   let description = [{
     Syntax:
 
     ```
     dense-intorfloat-elements-attribute ::= `dense` `<` attribute-value `>` `:`
                                             ( tensor-type | vector-type )
def Builtin_DenseIntOrFPElementsAttr
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DenseStringElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DenseStringElementsAttr
-    : Builtin_Attr<"DenseStringElements", "DenseElementsAttr"> {
+    : Builtin_Attr<"DenseStringElements", /*traits=*/[], "DenseElementsAttr"> {
   let summary = "An Attribute containing a dense multi-dimensional array of "
                 "strings";
   let description = [{
     Syntax:
 
     ```
     dense-string-elements-attribute ::= `dense` `<` attribute-value `>` `:`
                                         ( tensor-type | vector-type )
def Builtin_DenseStringElementsAttr
   let genStorageClass = 0;
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary"> {
+def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "An dictionary of named Attribute values";
   let description = [{
     Syntax:
 
     ```
     dictionary-attribute ::= `{` (attribute-entry (`,` attribute-entry)*)? `}`
     ```
 
def Builtin_OpaqueAttr : Builtin_Attr<"Opaque"> {
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // OpaqueElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_OpaqueElementsAttr
-    : Builtin_Attr<"OpaqueElements", "ElementsAttr"> {
+    : Builtin_Attr<"OpaqueElements", /*traits=*/[], "ElementsAttr"> {
   let summary = "An opaque representation of a multi-dimensional array";
   let description = [{
     Syntax:
 
     ```
     opaque-elements-attribute ::= `opaque` `<` dialect-namespace  `,`
                                   hex-string-literal `>` `:`
                                   ( tensor-type | vector-type )
def Builtin_OpaqueElementsAttr
   let skipDefaultBuilders = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SparseElementsAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_SparseElementsAttr
-    : Builtin_Attr<"SparseElements", "ElementsAttr"> {
+    : Builtin_Attr<"SparseElements", /*traits=*/[], "ElementsAttr"> {
   let summary = "An opaque representation of a multi-dimensional array";
   let description = [{
     Syntax:
 
     ```
     sparse-elements-attribute ::= `sparse` `<` attribute-value `,`
                                   attribute-value `>` `:`
                                   ( tensor-type | vector-type )
let extraClassDeclaration = [{
     StringRef getLeafReference() const;
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // TypeAttr
 //===----------------------------------------------------------------------===//
 
-def Builtin_TypeAttr : Builtin_Attr<"Type"> {
+def Builtin_TypeAttr : Builtin_Attr<"Type", [
+    DeclareAttrInterfaceMethods<SubElementAttrInterface>
+  ]> {
   let summary = "An Attribute containing a Type";
   let description = [{
     Syntax:
 
     ```
     type-attribute ::= type
     ```
 

mlir/include/mlir/IR/BuiltinTypes.h

 //===- BuiltinTypes.h - MLIR Builtin Type Classes ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_BUILTINTYPES_H
 #define MLIR_IR_BUILTINTYPES_H
 
-#include "mlir/IR/Attributes.h"
-#include "mlir/IR/Types.h"
+#include "SubElementInterfaces.h"
 
 namespace llvm {
 struct fltSemantics;
 } // namespace llvm
 
 namespace mlir {
 class AffineExpr;
 class AffineMap;

mlir/include/mlir/IR/BuiltinTypes.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUILTIN_TYPES
 #define BUILTIN_TYPES
 
 include "mlir/IR/BuiltinDialect.td"
 include "mlir/IR/BuiltinTypeInterfaces.td"
+include "mlir/IR/SubElementInterfaces.td"
 
 // TODO: Currently the types defined in this file are prefixed with `Builtin_`.
 // This is to differentiate the types here with the ones in OpBase.td. We should
 // remove the definitions in OpBase.td, and repoint users to this file instead.
 
 // Base class for Builtin dialect types.
-class Builtin_Type<string name, string baseCppClass = "::mlir::Type">
-    : TypeDef<Builtin_Dialect, name, [], baseCppClass> {
+class Builtin_Type<string name, list<Trait> traits = [],
+                   string baseCppClass = "::mlir::Type">
+    : TypeDef<Builtin_Dialect, name, traits, baseCppClass> {
   let mnemonic = ?;
 }
 
 //===----------------------------------------------------------------------===//
 // ComplexType
 //===----------------------------------------------------------------------===//
 
 def Builtin_Complex : Builtin_Type<"Complex"> {
def Builtin_Complex : Builtin_Type<"Complex"> {
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // FloatType
 //===----------------------------------------------------------------------===//
 
 // Base class for Builtin dialect float types.
-class Builtin_FloatType<string name> : Builtin_Type<name, "::mlir::FloatType"> {
+class Builtin_FloatType<string name>
+    : Builtin_Type<name, /*traits=*/[], "::mlir::FloatType"> {
   let extraClassDeclaration = [{
     static }] # name # [{Type get(MLIRContext *context);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // BFloat16Type
 

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

mlir/include/mlir/IR/CMakeLists.txt

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

mlir/include/mlir/IR/SubElementInterfaces.h

+//===- SubElementInterfaces.h - Attr and Type SubElements -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains interfaces and utilities for querying the sub elements of
+// an attribute or type.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_SUBELEMENTINTERFACES_H
+#define MLIR_INTERFACES_SUBELEMENTINTERFACES_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Types.h"
+
+/// Include the definitions of the sub elemnt interfaces.
+#include "mlir/IR/SubElementAttrInterfaces.h.inc"
+#include "mlir/IR/SubElementTypeInterfaces.h.inc"
+
+#endif // MLIR_INTERFACES_SUBELEMENTINTERFACES_H

mlir/include/mlir/IR/SubElementInterfaces.td

+//===-- SubElementInterfaces.td - Sub-Element Interfaces ---*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a set of interfaces that can be used to interface with
+// sub-elements, e.g. held attributes and types, of a composite attribute or
+// type.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_SUBELEMENTINTERFACES_TD_
+#define MLIR_IR_SUBELEMENTINTERFACES_TD_
+
+include "mlir/IR/OpBase.td"
+
+//===----------------------------------------------------------------------===//
+// SubElementInterfaceBase
+//===----------------------------------------------------------------------===//
+
+class SubElementInterfaceBase<string interfaceName, string derivedValue> {
+  string cppNamespace = "::mlir";
+
+  list<InterfaceMethod> methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Walk all of the immediately nested sub-attributes and sub-types. This
+        method does not recurse into sub elements.
+      }], "void", "walkImmediateSubElements",
+      (ins "llvm::function_ref<void(mlir::Attribute)>":$walkAttrsFn,
+           "llvm::function_ref<void(mlir::Type)>":$walkTypesFn)
+    >,
+  ];
+
+  code extraClassDeclaration = [{
+    /// Walk all of the held sub-attributes.
+    void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
+      walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
+    }
+
+    /// Walk all of the held sub-types.
+    void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
+      walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
+    }
+
+    /// Walk all of the held sub-attributes and sub-types.
+    void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
+                         llvm::function_ref<void(mlir::Type)> walkTypesFn);
+  }];
+
+  code extraTraitClassDeclaration = [{
+    /// Walk all of the held sub-attributes.
+    void walkSubAttrs(llvm::function_ref<void(mlir::Attribute)> walkFn) {
+      walkSubElements(walkFn, /*walkTypesFn=*/[](mlir::Type) {});
+    }
+
+    /// Walk all of the held sub-types.
+    void walkSubTypes(llvm::function_ref<void(mlir::Type)> walkFn) {
+      walkSubElements(/*walkAttrsFn=*/[](mlir::Attribute) {}, walkFn);
+    }
+
+    /// Walk all of the held sub-attributes and sub-types.
+    void walkSubElements(llvm::function_ref<void(mlir::Attribute)> walkAttrsFn,
+                         llvm::function_ref<void(mlir::Type)> walkTypesFn) {
+      }] # interfaceName # " interface(" # derivedValue # [{);
+      interface.walkSubElements(walkAttrsFn, walkTypesFn);
+    }
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// SubElementAttrInterface
+//===----------------------------------------------------------------------===//
+
+def SubElementAttrInterface
+    : AttrInterface<"SubElementAttrInterface">,
+      SubElementInterfaceBase<"SubElementAttrInterface", "$_attr"> {
+  let description = [{
+    An interface used to query and manipulate sub-elements, such as sub-types
+    and sub-attributes of a composite attribute.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// SubElementTypeInterface
+//===----------------------------------------------------------------------===//
+
+def SubElementTypeInterface
+    : TypeInterface<"SubElementTypeInterface">,
+      SubElementInterfaceBase<"SubElementTypeInterface", "$_type"> {
+  let description = [{
+    An interface used to query and manipulate sub-elements, such as sub-types
+    and sub-attributes of a composite type.
+  }];
+}
+
+#endif // MLIR_IR_SUBELEMENTINTERFACES_TD_

mlir/lib/IR/AsmPrinter.cpp

 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/SubElementInterfaces.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
void AliasInitializer::visit(Attribute attr, bool canBeDeferred) {
 
   // Try to generate an alias for this attribute.
   if (succeeded(generateAlias(attr, aliasToAttr))) {
     if (canBeDeferred)
       deferrableAttributes.insert(attr);
     return;
   }
 
-  if (auto arrayAttr = attr.dyn_cast<ArrayAttr>()) {
-    for (Attribute element : arrayAttr.getValue())
-      visit(element);
-  } else if (auto dictAttr = attr.dyn_cast<DictionaryAttr>()) {
-    for (const NamedAttribute &attr : dictAttr)
-      visit(attr.second);
-  } else if (auto typeAttr = attr.dyn_cast<TypeAttr>()) {
-    visit(typeAttr.getValue());
+  // Check for any sub elements.
+  if (auto subElementInterface = attr.dyn_cast<SubElementAttrInterface>()) {
+    subElementInterface.walkSubElements([&](Attribute attr) { visit(attr); },
+                                        [&](Type type) { visit(type); });
   }
 }
 
 void AliasInitializer::visit(Type type) {
   if (!visitedTypes.insert(type).second)
     return;
 
   // Try to generate an alias for this type.
   if (succeeded(generateAlias(type, aliasToType)))
     return;
 
-  // Visit several subtypes that contain types or attributes.
-  if (auto funcType = type.dyn_cast<FunctionType>()) {
-    // Visit input and result types for functions.
-    for (auto input : funcType.getInputs())
-      visit(input);
-    for (auto result : funcType.getResults())
-      visit(result);
-  } else if (auto shapedType = type.dyn_cast<ShapedType>()) {
-    visit(shapedType.getElementType());
-
-    // Visit affine maps in memref type.
-    if (auto memref = type.dyn_cast<MemRefType>())
-      for (auto map : memref.getAffineMaps())
-        visit(AffineMapAttr::get(map));
+  // Check for any sub elements.
+  if (auto subElementInterface = type.dyn_cast<SubElementTypeInterface>()) {
+    subElementInterface.walkSubElements([&](Attribute attr) { visit(attr); },
+                                        [&](Type type) { visit(type); });
   }
 }
 
 template <typename T>
 LogicalResult AliasInitializer::generateAlias(
     T symbol, llvm::MapVector<StringRef, std::vector<T>> &aliasToSymbol) {
   SmallString<16> tempBuffer;
   for (const auto &interface : interfaces) {

mlir/lib/IR/BuiltinAttributes.cpp

void BuiltinDialect::registerAttributes() {
   addAttributes<AffineMapAttr, ArrayAttr, DenseIntOrFPElementsAttr,
                 DenseStringElementsAttr, DictionaryAttr, FloatAttr,
                 SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr,
                 OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr,
                 UnitAttr>();
 }
 
 //===----------------------------------------------------------------------===//
+// ArrayAttr
+//===----------------------------------------------------------------------===//
+
+void ArrayAttr::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  for (Attribute attr : getValue())
+    walkAttrsFn(attr);
+}
+
+//===----------------------------------------------------------------------===//
 // 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 (Attribute attr : llvm::make_second_range(getValue()))
+    walkAttrsFn(attr);
+}
+
 //===----------------------------------------------------------------------===//
 // StringAttr
 //===----------------------------------------------------------------------===//
 
 StringAttr StringAttr::getEmptyStringAttrUnchecked(MLIRContext *context) {
   return Base::get(context, "", NoneType::get(context));
 }
 
std::vector<ptrdiff_t> SparseElementsAttr::getFlattenedSparseIndices() const {
   // Otherwise, reinterpret each index as an ArrayRef when flattening.
   auto numSparseIndices = sparseIndices.getType().getDimSize(0);
   size_t rank = getType().getRank();
   for (size_t i = 0, e = numSparseIndices; i != e; ++i)
     flatSparseIndices.push_back(getFlattenedIndex(
         {&*std::next(sparseIndexValues.begin(), i * rank), rank}));
   return flatSparseIndices;
 }
+
+//===----------------------------------------------------------------------===//
+// TypeAttr
+//===----------------------------------------------------------------------===//
+
+void TypeAttr::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getValue());
+}

mlir/lib/IR/BuiltinTypes.cpp

iterateIndicesExcept(originalNumResults, resultIndices, [&](unsigned i) {
       newResultTypesBuffer.emplace_back(getResult(i));
     });
     newResultTypes = newResultTypesBuffer;
   }
 
   return get(getContext(), newInputTypes, 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);
+}
+
 //===----------------------------------------------------------------------===//
 // OpaqueType
 //===----------------------------------------------------------------------===//
 
 /// Verify the construction of an opaque type.
 LogicalResult OpaqueType::verify(function_ref<InFlightDiagnostic()> emitError,
                                  Identifier 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);
   if (auto et = getElementType().dyn_cast<FloatType>())
     if (auto scaledEt = et.scaleElementBitwidth(scale))
       return VectorType::get(getShape(), scaledEt);
   return VectorType();
 }
 
+void VectorType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // TensorType
 //===----------------------------------------------------------------------===//
 
 // Check if "elementType" can be an element type of a tensor.
 static LogicalResult
 checkTensorElementType(function_ref<InFlightDiagnostic()> emitError,
                        Type elementType) {
for (int64_t s : shape)
     if (s < -1)
       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());
+}
+
 //===----------------------------------------------------------------------===//
 // 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());
+}
+
 //===----------------------------------------------------------------------===//
 // BaseMemRefType
 //===----------------------------------------------------------------------===//
 
 Attribute BaseMemRefType::getMemorySpace() const {
   if (auto rankedMemRefTy = dyn_cast<MemRefType>())
     return rankedMemRefTy.getMemorySpace();
   return cast<UnrankedMemRefType>().getMemorySpace();
LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
 
   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());
+  walkAttrsFn(getMemorySpace());
+  for (AffineMap map : getAffineMaps())
+    walkAttrsFn(AffineMapAttr::get(map));
+}
+
 //===----------------------------------------------------------------------===//
 // 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());
+}
+
 //===----------------------------------------------------------------------===//
 /// 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 Utilities
 //===----------------------------------------------------------------------===//
 
 AffineMap mlir::makeStridedLinearLayoutMap(ArrayRef<int64_t> strides,
                                            int64_t offset,
                                            MLIRContext *context) {
   AffineExpr expr;

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   IntegerSet.cpp
   Location.cpp
   MLIRContext.cpp
   Operation.cpp
   OperationSupport.cpp
   PatternMatch.cpp
   Region.cpp
   RegionKindInterface.cpp
+  SubElementInterfaces.cpp
   SymbolTable.cpp
   TensorEncoding.cpp
   Types.cpp
   TypeRange.cpp
   TypeUtilities.cpp
   Value.cpp
   Verifier.cpp
   Visitors.cpp
add_mlir_library(MLIRIR
   MLIRBuiltinTypesIncGen
   MLIRBuiltinTypeInterfacesIncGen
   MLIRCallInterfacesIncGen
   MLIRCastInterfacesIncGen
   MLIRDataLayoutInterfacesIncGen
   MLIROpAsmInterfaceIncGen
   MLIRRegionKindInterfaceIncGen
   MLIRSideEffectInterfacesIncGen
+  MLIRSubElementInterfacesIncGen
   MLIRSymbolInterfacesIncGen
   MLIRTensorEncodingIncGen
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/IR/SubElementInterfaces.cpp

+//===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/SubElementInterfaces.h"
+
+using namespace mlir;
+
+template <typename InterfaceT>
+static void walkSubElementsImpl(InterfaceT interface,
+                                function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) {
+  interface.walkImmediateSubElements(
+      [&](Attribute attr) {
+        // Guard against potentially null inputs. This removes the need for the
+        // derived attribute/type to do it.
+        if (!attr)
+          return;
+
+        // Walk any sub elements first.
+        if (auto interface = attr.dyn_cast<SubElementAttrInterface>())
+          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn);
+
+        // 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;
+
+        // Walk any sub elements first.
+        if (auto interface = type.dyn_cast<SubElementTypeInterface>())
+          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn);
+
+        // Walk this type.
+        walkTypesFn(type);
+      });
+}
+
+void SubElementAttrInterface::walkSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) {
+  assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
+  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn);
+}
+
+void SubElementTypeInterface::walkSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) {
+  assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
+  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn);
+}
+
+//===----------------------------------------------------------------------===//
+// SubElementInterface Tablegen definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
+#include "mlir/IR/SubElementTypeInterfaces.cpp.inc"

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
   DialectTest.cpp
   MemRefTypeTest.cpp
   OperationSupportTest.cpp
   ShapedTypeTest.cpp
+  SubElementInterfaceTest.cpp
 )
 target_link_libraries(MLIRIRTests
   PRIVATE
   MLIRIR)

mlir/unittests/IR/SubElementInterfaceTest.cpp

+//===- SubElementInterfaceTest.cpp - SubElementInterface unit tests -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/SubElementInterfaces.h"
+#include "gtest/gtest.h"
+#include <cstdint>
+
+using namespace mlir;
+using namespace mlir::detail;
+
+namespace {
+TEST(SubElementInterfaceTest, Nested) {
+  MLIRContext context;
+  Builder builder(&context);
+
+  BoolAttr trueAttr = builder.getBoolAttr(true);
+  BoolAttr falseAttr = builder.getBoolAttr(false);
+  ArrayAttr boolArrayAttr = builder.getArrayAttr({trueAttr, falseAttr});
+  DictionaryAttr dictAttr =
+      builder.getDictionaryAttr(builder.getNamedAttr("array", boolArrayAttr));
+
+  SmallVector<Attribute> subAttrs;
+  dictAttr.walkSubAttrs([&](Attribute attr) { subAttrs.push_back(attr); });
+  EXPECT_EQ(llvm::makeArrayRef(subAttrs),
+            ArrayRef<Attribute>({trueAttr, falseAttr, boolArrayAttr}));
+}
+
+} // end namespace