[mlir] Prevent SubElementInterface from going into infinite recursion

[myhsu]

Summary

Since only mutable types and attributes can go into infinite recursion inside SubElementInterface::walkSubElement, and there are only a few of them (mutable types and attributes), we introduce a new flag in AbstractType/Attribute to indicate if a type or attribute is mutable. Then, inside SubElementInterface, we use a set to record visited mutable types and attributes that have been visited before.

NOTE: None of the code is using SubElementType/AttrInterface with potential recursion at this time, thus I don't attach any test case. However, we can merge this patch with D127539, which applies SubElementTypeInterface to LLVM dialect's types and provide chances to exercise the changes here, if that's the best way to move forward.

[Mogball]

Ah, I didn't think this was a possibility. It's unfortunate that walkSubElements has to be slowed down in general. It's particularly hot due to SymbolTable.

[Mogball]

I think we can optimize for the common case and have the first N recursions not track a set (1, 2?). Also, is it possible to have nested types stack overflow the walk function?

mlir/lib/IR/SubElementInterfaces.cpp

71

What about SmallPtrSet<Pointer union<Type,Attribute>>?

[rriddle]

Yeah, the set should be unnecessary in 99% of cases. This is something that is only necessary for mutable attributes/types, given that those are the only ones that could be recursive. Conceptually we could detect mutability on registration, add a bit for that to the AbstractAttribute/Type classes, and check that here.

[myhsu]

This is a major revision.

• Add an isMutable flag on both AbstractType/Attribute to indicate whether the underlying type/attribute (storage) is mutable. This is done by checking if there is a member function LogicalResult mutate(StorageAllocator &, ...) defined in the storage class via SFINAE facilities.
• SubElementTypeInterface now only does bookkeeping on mutable types or attributes.
• Use unit test to test the new mutable type flag. Which effectively moves part of the unit test infrastructure from D127539 to here.

[myhsu]

mlir/lib/IR/SubElementInterfaces.cpp

71

Given that these sets are not on the hot path anymore, do you think it's still better to change the container type? I feel like two separate sets are more readable.

[Mogball]

Could mutability not be a trait instead?

mlir/include/mlir/IR/StorageUniquerSupport.h

79

LLVM has a utility for this. LLVM::is_detected. You can find an example in AsmParser

mlir/lib/IR/SubElementInterfaces.cpp

19

I would have two separate recursive functions. One with the set and one without

71

Yeah no preference then

80–83

Drop llvm::

[myhsu]

In D127537#3582698, @Mogball wrote:

Could mutability not be a trait instead?

It is possible, for instance an opaque LLVMStruct cannot be mutated. But to have such granularity, we need to rely on non-static information, which is something AbstractType/AbstractAttribute cannot carry. Of course, we can have isMutable exposed on Type and Attribute, but I'm not sure about that. What do you think?

[Mogball]

I'm not sure I understand what you're getting at. Mutability is tied to the abstract type/attribute implementation. Instead of abstractType->isMutable(), could we not have abstractType->hasTrait<TypeTrait::Mutable>()?

[myhsu]

In D127537#3583128, @Mogball wrote:

I'm not sure I understand what you're getting at. Mutability is tied to the abstract type/attribute implementation. Instead of abstractType->isMutable(), could we not have abstractType->hasTrait<TypeTrait::Mutable>()?

I see, I thought you were referring to "trait" as a general C++ concept rather than MLIR's Trait.

IIRC trait needs to be manually annotated on a type or attribute, while my current approach detects mutability out of the box. The reason I chose the current way was because TypeStorage/AttributeStorage use the same way to see if a type/attribute can be mutated. And I'm a little worry that using trait might make such property out-of-sync with TypeStorage/AttributeStorage, if, let's say, a dialect forgets to annotate their mutable types or attributes with trait.

[Mogball]

I think @ftynse introduced mutable types and attributes and may be able to weigh in on this.

MLIR doesn't really "detect" mutability, more so "mutate" fails to compile if the type doesn't implement the function. Now that we're trying to query mutability, storing the results of a static check in a bit doesn't seem like the way to go. Traits make sense in this regard. To keep the two in sync, mutate can static assert that the concrete type implements the trait and the trait's verifier can static assert that the concrete type implements the function.

[rriddle]

In D127537#3585984, @myhsu wrote:

In D127537#3583128, @Mogball wrote:

I'm not sure I understand what you're getting at. Mutability is tied to the abstract type/attribute implementation. Instead of abstractType->isMutable(), could we not have abstractType->hasTrait<TypeTrait::Mutable>()?

I see, I thought you were referring to "trait" as a general C++ concept rather than MLIR's Trait.

IIRC trait needs to be manually annotated on a type or attribute, while my current approach detects mutability out of the box. The reason I chose the current way was because TypeStorage/AttributeStorage use the same way to see if a type/attribute can be mutated. And I'm a little worry that using trait might make such property out-of-sync with TypeStorage/AttributeStorage, if, let's say, a dialect forgets to annotate their mutable types or attributes with trait.

Having a MutableTypeInterface/MutableAttrInterface for specifying this seems like a good step forward. Manual annotation isn't necessarily a bad thing here, ODS can always sugar this better (though ODS doesn't support mutable attribute/types right now). We should consider the desired end state of specifying mutability for attributes and types.

[myhsu]

In D127537#3587874, @rriddle wrote:

In D127537#3585984, @myhsu wrote:

In D127537#3583128, @Mogball wrote:

I'm not sure I understand what you're getting at. Mutability is tied to the abstract type/attribute implementation. Instead of abstractType->isMutable(), could we not have abstractType->hasTrait<TypeTrait::Mutable>()?

I see, I thought you were referring to "trait" as a general C++ concept rather than MLIR's Trait.

IIRC trait needs to be manually annotated on a type or attribute, while my current approach detects mutability out of the box. The reason I chose the current way was because TypeStorage/AttributeStorage use the same way to see if a type/attribute can be mutated. And I'm a little worry that using trait might make such property out-of-sync with TypeStorage/AttributeStorage, if, let's say, a dialect forgets to annotate their mutable types or attributes with trait.

Having a MutableTypeInterface/MutableAttrInterface for specifying this seems like a good step forward. Manual annotation isn't necessarily a bad thing here, ODS can always sugar this better (though ODS doesn't support mutable attribute/types right now). We should consider the desired end state of specifying mutability for attributes and types.

Yes, ideally we can have MutableTypeInterface/MutableAttrInterface with a mutate function. But mutate functions in existing TypeStorage/AttrStorage relies on template-parameterized function arguments, which I don't think it can fit into the facilities for interface (An interface's Concept declares virtual functions but virtual functions can't have template).

Right now I'm planning to use a simple TypeTrait/AttrTrait to specify mutability to go forward on this patch.

[myhsu]

Use a TypeTrait/AttributeTrait to indicate mutability instead of a flag in AbstractType/AbstractAttribute. I also added static asserts to make sure it syncs up with mutate function defined in the Type/Attribute storage

[myhsu]

mlir/include/mlir/IR/StorageUniquerSupport.h

81

FYI: For variadic functions, we can't use the normal way (like decltype(declval<T>().foo(...))) to match function signature.

mlir/lib/IR/SubElementInterfaces.cpp

19

If we split into two versions, the function without set can only call the same variant* (i.e. the one without set), which effectively makes an assumption that all sub-types of an immutable type are immutable -- I'm afraid this is not the case, a good counterexample is !llvm.ptr<struct<...>> (LLVM pointer is immutable while LLVM struct is mutable).

*: Even if we use a "local" set to record visited types like

void walkImmutableImpl(...) {
  set<Type> localVisited;
  if (mutable)
    walkMutableImpl(..., localVisited);
  else
    walkImmutableImpl(...);
}

It still can't prevent infinite recursion.

[myhsu]

mlir/include/mlir/IR/AttributeSupport.h

266–267 ↗

(On Diff #437737)

I forgot to change this, just a second.

[myhsu]

Fix minor issues.

[Mogball]

mlir/include/mlir/IR/AttributeSupport.h

265 ↗

(On Diff #437737)

This should be somewhere else... In one of the mutate function wrappers?

[myhsu]

mlir/include/mlir/IR/AttributeSupport.h

265 ↗

(On Diff #437737)

Doing checks in one of the mutate wrappers allows up to only check for IsMutableType trait (because the wrapper will only be compiled / compilable if mutate function is defined) but we can't check the other direction: If IsMutableType is present, whether mutate function is defined. Thus, the check I wrote here checks for both at the same time (it's an "inverse XOR" check so it will only complain if only either mutate function or trait is present), in that sense I feel like putting it in the registration function is better

[Mogball]

mlir/include/mlir/IR/AttributeSupport.h

265 ↗

(On Diff #437737)

Can't you static_assert from the Trait::verify? Some other traits do that.

[rriddle]

mlir/include/mlir/IR/AttributeSupport.h

265 ↗

(On Diff #437737)

Can't you static_assert from the Trait::verify? Some other traits do that.

Attribute and Type traits don't have verifiers.

265 ↗

(On Diff #437743)

Can we pass the type of the trait to StorageUserBase instead? and have the static assert inside of StorageUserBase::mutate?

mlir/include/mlir/IR/Attributes.h

243

The Attr is redundant and doesn't add much.

mlir/include/mlir/IR/StorageUniquerSupport.h

79

We should be able to avoid this complexity. We really just need a static assert inside of StorageUserBase::mutate that it's only called when the parent is mutable. If necessary, we could do something similar to [1] that makes the static assert dependent on the template types of the mutate function (but ignores them)

[1]: https://github.com/llvm/llvm-project/blob/32bd0c1714b4733af7abb9f570e961ffc889c793/mlir/include/mlir/IR/PatternMatch.h#L823

mlir/include/mlir/IR/Types.h

234

Same here, I'd drop the Type.

[myhsu]

• Rename mutable attributes.
• Move static checks into StorageUserBase::mutate.

[Mogball]

mlir/include/mlir/IR/Attributes.h

30

Doesn't this mark all attributes with IsMutable?

[Mogball]

mlir/include/mlir/IR/Attributes.h

30

Ah, nevermind.

[Mogball]

mlir/include/mlir/IR/StorageUniquerSupport.h

186

This seems a little invasive to me -- that you have to pass in MutableTrait from the implementing classes. Can we not have a generic StorageUserTrait::IsMutable? That way every user of StorageUserBase doesn't need to also define its own mutability trait that's identical to all the other ones.

[myhsu]

Use a single StorageUserTrait::IsMutable to replace TypeTrait::IsMutable and AttributeTrait::IsMutable. The latter two, however, remain type aliases to StorageUserTrait::IsMutable for ease of use.

mlir/include/mlir/IR/StorageUniquerSupport.h

186

I agree that a single IsMutable trait is easier and less intrusive. But I feel like many dialect authors might find StorageUserTrait a little confusing since they don't necessarily know what storage is in this context, as it's an implementation detail. So I'm turning the Attribute/TypeTrait::IsMutable here into aliases to StorageUserTrait::IsMutable

[Mogball]

LGTM

[myhsu]

Hi @rriddle do you have any further comments on this?

[rriddle]

Can we have a test that checks that we no longer go into infinite recursion? Otherwise, LG!

(Sorry, coming back from OOO)

[myhsu]

In D127537#3613127, @rriddle wrote:

Can we have a test that checks that we no longer go into infinite recursion? Otherwise, LG!

(Sorry, coming back from OOO)

right, this patch is effectively preparing for D127539, which will go into infinite recursion without it. Otherwise, I don't think any of the type or attribute (as of the current HEAD) with this interface will go into infinite recursion. Thus, I think it's more difficult to write such test in this patch (D127539 has tests to make sure we don't go into infinite recursion though). What do you think?

[rriddle]

I would prefer if we could test this outside of the LLVM dialect if possible, given that it should just be a matter of using the Recursive type in the test dialect right?

[myhsu]

In D127537#3613422, @rriddle wrote:

I would prefer if we could test this outside of the LLVM dialect if possible, given that it should just be a matter of using the Recursive type in the test dialect right?

Yes I think this is possible, I forgot there is a test dialect where we can exercise this part of the code by making type alias, which calls SubElementInterface, on its recursive type.
I will update this diff accordingly later.

[myhsu]

Add generic tests to check if SubElemenTypeInterface goes into infinite recursion. Please take another look.

Note that since SubElementAttributeInterface shares most of the logics with SubElemenTypeInterface (including the newly introduced IsMutable trait), plus we don't have recursive attribute in TestDialect right now, I don't add a similar test for attribute.

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

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

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVTypes.h

 ///
 /// would be represented in MLIR as:
 ///
 /// !spv.struct<A, (!spv.ptr<!spv.struct<A>, Generic>)>
 ///
 /// In the above, expressing recursive struct types is accomplished by giving a
 /// recursive struct a unique identified and using that identifier in the struct
 /// definition for recursive references.
-class StructType : public Type::TypeBase<StructType, CompositeType,
-                                         detail::StructTypeStorage> {
+class StructType
+    : public Type::TypeBase<StructType, CompositeType,
+                            detail::StructTypeStorage, TypeTrait::IsMutable> {
 public:
   using Base::Base;
 
   // Type for specifying the offset of the struct members
   using OffsetInfo = uint32_t;
 
   // Type for specifying the decoration(s) on struct members
   struct MemberDecorationInfo {

mlir/include/mlir/IR/Attributes.h

 /// a thin wrapper around an underlying storage pointer. Attributes are usually
 /// passed by value.
 class Attribute {
 public:
   /// Utility class for implementing attributes.
   template <typename ConcreteType, typename BaseType, typename StorageType,
             template <typename T> class... Traits>
   using AttrBase = detail::StorageUserBase<ConcreteType, BaseType, StorageType,
                                            detail::AttributeUniquer, Traits...>;
 
   using ImplType = AttributeStorage;
   using ValueType = void;
   using AbstractTy = AbstractAttribute;
 
   constexpr Attribute() {}
   /* implicit */ Attribute(const ImplType *impl)
       : impl(const_cast<ImplType *>(impl)) {}
private:
   static typename InterfaceBase::Concept *getInterfaceFor(Attribute attr) {
     return attr.getAbstractAttribute().getInterface<ConcreteType>();
   }
 
   /// Allow access to 'getInterfaceFor'.
   friend InterfaceBase;
 };
 
+//===----------------------------------------------------------------------===//
+// Core AttributeTrait
+//===----------------------------------------------------------------------===//
+
+/// This trait is used to determine if an attribute is mutable or not. It is
+/// attached on an attribute if the corresponding ImplType defines a `mutate`
+/// function with proper signature.
+namespace AttributeTrait {
+template <typename ConcreteType>
+using IsMutable = detail::StorageUserTrait::IsMutable<ConcreteType>;
+} // namespace AttributeTrait
+
 } // namespace mlir.
 
 namespace llvm {
 
 // Attribute hash just like pointers.
 template <> struct DenseMapInfo<mlir::Attribute> {
   static mlir::Attribute getEmptyKey() {
     auto *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();

mlir/include/mlir/IR/StorageUniquerSupport.h

ConcreteType getInstance() const {
     // the concrete type will multiply derive from the (content free) TraitBase
     // class, and we need to be able to disambiguate the path for the C++
     // compiler.
     auto *trait = static_cast<const TraitType<ConcreteType> *>(this);
     return *static_cast<const ConcreteType *>(trait);
   }
 };
 
+namespace StorageUserTrait {
+/// This trait is used to determine if a storage user, like Type, is mutable
+/// or not. A storage user is mutable if ImplType of the derived class defines
+/// a `mutate` function with a proper signature. Note that this trait is not
+/// supposed to be used publicly. Users should use alias names like
+/// `TypeTrait::IsMutable` instead.
+template <typename ConcreteType>
+struct IsMutable : public StorageUserTraitBase<ConcreteType, IsMutable> {};
+} // namespace StorageUserTrait
+
 //===----------------------------------------------------------------------===//
 // StorageUserBase
 //===----------------------------------------------------------------------===//
 
 namespace storage_user_base_impl {
 /// Returns true if this given Trait ID matches the IDs of any of the provided
 /// trait types `Traits`.
 template <template <typename T> class... Traits>
 bool hasTrait(TypeID traitID) {
   TypeID traitIDs[] = {TypeID::get<Traits>()...};
   for (unsigned i = 0, e = sizeof...(Traits); i != e; ++i)
     if (traitIDs[i] == traitID)
       return true;
   return false;
 }
 
 // We specialize for the empty case to not define an empty array.
 template <>
 inline bool hasTrait(TypeID traitID) {
   return false;
 }
 } // namespace storage_user_base_impl
 
 /// Utility class for implementing users of storage classes uniqued by a
public:
   static ConcreteT getFromOpaquePointer(const void *ptr) {
     return ConcreteT((const typename BaseT::ImplType *)ptr);
   }
 
 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();
   }

mlir/include/mlir/IR/Types.h

static typename InterfaceBase::Concept *getInterfaceFor(Type type) {
     return type.getAbstractType().getInterface<ConcreteType>();
   }
 
   /// Allow access to 'getInterfaceFor'.
   friend InterfaceBase;
 };
 
 //===----------------------------------------------------------------------===//
+// Core TypeTrait
+//===----------------------------------------------------------------------===//
+
+/// This trait is used to determine if a type is mutable or not. It is attached
+/// on a type if the corresponding ImplType defines a `mutate` function with
+/// a proper signature.
+namespace TypeTrait {
+template <typename ConcreteType>
+using IsMutable = detail::StorageUserTrait::IsMutable<ConcreteType>;
+} // namespace TypeTrait
+
+//===----------------------------------------------------------------------===//
 // Type Utils
 //===----------------------------------------------------------------------===//
 
 // Make Type hashable.
 inline ::llvm::hash_code hash_value(Type arg) {
   return DenseMapInfo<const Type::ImplType *>::getHashValue(arg.impl);
 }
 

mlir/lib/IR/SubElementInterfaces.cpp

 //===- SubElementInterfaces.cpp - Attr and Type SubElement Interfaces -----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SubElementInterfaces.h"
 
+#include "llvm/ADT/DenseSet.h"
+
 using namespace mlir;
 
 template <typename InterfaceT>
 static void walkSubElementsImpl(InterfaceT interface,
                                 function_ref<void(Attribute)> walkAttrsFn,
-                                function_ref<void(Type)> walkTypesFn) {
+                                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);
+          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.
         if (auto interface = type.dyn_cast<SubElementTypeInterface>())
-          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn);
+          walkSubElementsImpl(interface, walkAttrsFn, walkTypesFn, visitedAttrs,
+                              visitedTypes);
 
         // 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);
+  DenseSet<Attribute> visitedAttrs;
+  DenseSet<Type> visitedTypes;
+  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
+                      visitedTypes);
 }
 
 void SubElementTypeInterface::walkSubElements(
     function_ref<void(Attribute)> walkAttrsFn,
     function_ref<void(Type)> walkTypesFn) {
   assert(walkAttrsFn && walkTypesFn && "expected valid walk functions");
-  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn);
+  DenseSet<Attribute> visitedAttrs;
+  DenseSet<Type> visitedTypes;
+  walkSubElementsImpl(*this, walkAttrsFn, walkTypesFn, visitedAttrs,
+                      visitedTypes);
 }
 
 //===----------------------------------------------------------------------===//
 // SubElementInterface Tablegen definitions
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/SubElementAttrInterfaces.cpp.inc"
 #include "mlir/IR/SubElementTypeInterfaces.cpp.inc"

mlir/test/IR/recursive-type.mlir

 // RUN: mlir-opt %s -test-recursive-types | FileCheck %s
 
+// CHECK: !testrec = !test.test_rec<type_to_alias, test_rec<type_to_alias>>
+
 // CHECK-LABEL: @roundtrip
 func.func @roundtrip() {
   // CHECK: !test.test_rec<a, test_rec<b, test_type>>
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<a, test_rec<b, test_type>>
   // CHECK: !test.test_rec<c, test_rec<c>>
   "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<c, test_rec<c>>
+  // Make sure walkSubElementType, which is used to generate aliases, doesn't go
+  // into inifinite recursion.
+  // CHECK: !testrec
+  "test.dummy_op_for_roundtrip"() : () -> !test.test_rec<type_to_alias, test_rec<type_to_alias>>
   return
 }
 
 // CHECK-LABEL: @create
 func.func @create() {
   // CHECK: !test.test_rec<some_long_and_unique_name, test_rec<some_long_and_unique_name>>
   return
 }

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

AliasResult getAlias(Type type, raw_ostream &os) const final {
     if (auto intType = type.dyn_cast<TestIntegerType>()) {
       if (intType.getSignedness() ==
               TestIntegerType::SignednessSemantics::Unsigned &&
           intType.getWidth() == 8) {
         os << "test_ui8";
         return AliasResult::FinalAlias;
       }
     }
+    if (auto recType = type.dyn_cast<TestRecursiveType>()) {
+      if (recType.getName() == "type_to_alias") {
+        // We only make alias for a specific recursive type.
+        os << "testrec";
+        return AliasResult::FinalAlias;
+      }
+    }
     return AliasResult::NoAlias;
   }
 
   //===------------------------------------------------------------------===//
   // Resources
   //===------------------------------------------------------------------===//
 
   std::string

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

 
 #include <tuple>
 
 #include "TestTraits.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/IR/SubElementInterfaces.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace test {
 class TestAttrWithFormatAttr;
 
 /// FieldInfo represents a field in the StructType data type. It is used as a
 /// parameter in TestTypeDefs.td.
struct TestRecursiveTypeStorage : public ::mlir::TypeStorage {
   ::mlir::Type body;
 };
 
 /// Simple recursive type identified by its name and pointing to another named
 /// type, potentially itself. This requires the body to be mutated separately
 /// from type creation.
 class TestRecursiveType
     : public ::mlir::Type::TypeBase<TestRecursiveType, ::mlir::Type,
-                                    TestRecursiveTypeStorage> {
+                                    TestRecursiveTypeStorage,
+                                    ::mlir::SubElementTypeInterface::Trait,
+                                    ::mlir::TypeTrait::IsMutable> {
 public:
   using Base::Base;
 
   static TestRecursiveType get(::mlir::MLIRContext *ctx,
                                ::llvm::StringRef name) {
     return Base::get(ctx, name);
   }
 
   /// Body getter and setter.
   ::mlir::LogicalResult setBody(Type body) { return Base::mutate(body); }
-  ::mlir::Type getBody() { return getImpl()->body; }
+  ::mlir::Type getBody() const { return getImpl()->body; }
 
   /// Name/key getter.
   ::llvm::StringRef getName() { return getImpl()->name; }
+
+  void walkImmediateSubElements(
+      ::llvm::function_ref<void(::mlir::Attribute)> walkAttrsFn,
+      ::llvm::function_ref<void(::mlir::Type)> walkTypesFn) const {
+    walkTypesFn(getBody());
+  }
 };
 
 } // namespace test
 
 #endif // MLIR_TESTTYPES_H

mlir/unittests/Dialect/CMakeLists.txt

 add_mlir_unittest(MLIRDialectTests
   BroadcastShapeTest.cpp
 )
 target_link_libraries(MLIRDialectTests
   PRIVATE
   MLIRIR
   MLIRDialect)
 
 add_subdirectory(Affine)
+add_subdirectory(LLVMIR)
 add_subdirectory(MemRef)
-
 add_subdirectory(Quant)
 add_subdirectory(SparseTensor)
 add_subdirectory(SPIRV)
 add_subdirectory(Utils)

mlir/unittests/Dialect/LLVMIR/CMakeLists.txt

+add_mlir_unittest(MLIRLLVMIRTests
+  LLVMTypeTest.cpp
+)
+target_link_libraries(MLIRLLVMIRTests
+  PRIVATE
+  MLIRLLVMDialect
+  )

mlir/unittests/Dialect/LLVMIR/LLVMTestBase.h

+//===- LLVMTestBase.h - Test fixure for LLVM dialect tests ------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Test fixure for LLVM dialect tests.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_UNITTEST_DIALECT_LLVMIR_LLVMTESTBASE_H
+#define MLIR_UNITTEST_DIALECT_LLVMIR_LLVMTESTBASE_H
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "gtest/gtest.h"
+
+class LLVMIRTest : public ::testing::Test {
+protected:
+  LLVMIRTest() { context.getOrLoadDialect<mlir::LLVM::LLVMDialect>(); }
+
+  mlir::MLIRContext context;
+};
+
+#endif

mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp

+//===- LLVMTypeTest.cpp - Tests for LLVM types ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "LLVMTestBase.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/IR/SubElementInterfaces.h"
+
+using namespace mlir;
+using namespace mlir::LLVM;
+
+TEST_F(LLVMIRTest, IsStructTypeMutable) {
+  auto structTy = LLVMStructType::getIdentified(&context, "foo");
+  ASSERT_TRUE(bool(structTy));
+  ASSERT_TRUE(structTy.hasTrait<TypeTrait::IsMutable>());
+}