[mlir][LLVMIR] Apply SubElementTypeInterface on suitable types

[myhsu]

Summary

This feature is tested by unit test since not many places in the codebase use SubElementTypeInterface. This is also the preparation for changes about using type aliases for LLVM structs as we discussed on Discourse (Because SubElementTypeInterface is required to generate proper type aliases).

[myhsu]

Reflect changes to D127537. NFC

[myhsu]

Reflect minor updates from D127537. NFC

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

 //===- LLVMDialect.h - MLIR LLVM dialect types ------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the types for the LLVM dialect in MLIR. These MLIR types
 // correspond to the LLVM IR type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 #define MLIR_DIALECT_LLVMIR_LLVMTYPES_H_
 
+#include "mlir/IR/SubElementInterfaces.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace llvm {
 class ElementCount;
 class TypeSize;
 } // namespace llvm
 

 // LLVMArrayType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect array type. It is an aggregate type representing consecutive
 /// elements in memory, parameterized by the number of elements and the element
 /// type.
 class LLVMArrayType
     : public Type::TypeBase<LLVMArrayType, Type, detail::LLVMTypeAndSizeStorage,
-                            DataLayoutTypeInterface::Trait> {
+                            DataLayoutTypeInterface::Trait,
+                            SubElementTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
   using Base::getChecked;
 
   /// Checks if the given type can be used inside an array type.
   static bool isValidElementType(Type type);
 
public:
   unsigned getTypeSize(const DataLayout &dataLayout,
                        DataLayoutEntryListRef params) const;
 
   unsigned getABIAlignment(const DataLayout &dataLayout,
                            DataLayoutEntryListRef params) const;
 
   unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMFunctionType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect function type. It consists of a single return type (unlike MLIR
 /// which can have multiple), a list of parameter types and can optionally be
 /// variadic.
-class LLVMFunctionType
-    : public Type::TypeBase<LLVMFunctionType, Type,
-                            detail::LLVMFunctionTypeStorage> {
+class LLVMFunctionType : public Type::TypeBase<LLVMFunctionType, Type,
+                                               detail::LLVMFunctionTypeStorage,
+                                               SubElementTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
   using Base::getChecked;
 
   /// Checks if the given type can be used an argument in a function type.
   static bool isValidArgumentType(Type type);
 
public:
   getChecked(function_ref<InFlightDiagnostic()> emitError, Type result,
              ArrayRef<Type> arguments, bool isVarArg = false);
 
   /// Returns a clone of this function type with the given argument
   /// and result types.
   LLVMFunctionType clone(TypeRange inputs, TypeRange results) const;
 
   /// Returns the result type of the function.
-  Type getReturnType();
+  Type getReturnType() const;
 
   /// Returns the result type of the function as an ArrayRef, enabling better
   /// integration with generic MLIR utilities.
-  ArrayRef<Type> getReturnTypes();
+  ArrayRef<Type> getReturnTypes() const;
 
   /// Returns the number of arguments to the function.
   unsigned getNumParams();
 
   /// Returns `i`-th argument of the function. Asserts on out-of-bounds.
   Type getParamType(unsigned i);
 
   /// Returns a list of argument types of the function.
-  ArrayRef<Type> getParams();
+  ArrayRef<Type> getParams() const;
   ArrayRef<Type> params() { return getParams(); }
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type result, ArrayRef<Type> arguments, bool);
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMPointerType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect pointer type. This type typically represents a reference to an
 /// object in memory. Pointers may be opaque or parameterized by the element
 /// type. Both opaque and non-opaque pointers are additionally parameterized by
 /// the address space.
-class LLVMPointerType : public Type::TypeBase<LLVMPointerType, Type,
-                                              detail::LLVMPointerTypeStorage,
-                                              DataLayoutTypeInterface::Trait> {
+class LLVMPointerType
+    : public Type::TypeBase<
+          LLVMPointerType, Type, detail::LLVMPointerTypeStorage,
+          DataLayoutTypeInterface::Trait, SubElementTypeInterface::Trait> {
 public:
   /// Inherit base constructors.
   using Base::Base;
 
   /// Checks if the given type can have a pointer type pointing to it.
   static bool isValidElementType(Type type);
 
   /// Gets or creates an instance of LLVM dialect pointer type pointing to an
public:
   unsigned getABIAlignment(const DataLayout &dataLayout,
                            DataLayoutEntryListRef params) const;
   unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
   bool areCompatible(DataLayoutEntryListRef oldLayout,
                      DataLayoutEntryListRef newLayout) const;
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMStructType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect structure type representing a collection of different-typed
 /// elements manipulated together. Structured can optionally be packed, meaning

 /// body, one is not allowed to set another body (however, one can set exactly
 /// the same body).
 ///
 /// Note that the packedness of the struct takes place in uniquing of literal
 /// structs, but does not in uniquing of identified structs.
 class LLVMStructType
     : public Type::TypeBase<LLVMStructType, Type, detail::LLVMStructTypeStorage,
                             DataLayoutTypeInterface::Trait,
+                            SubElementTypeInterface::Trait,
                             TypeTrait::IsMutable> {
 public:
   /// Inherit base constructors.
   using Base::Base;
 
   /// Checks if the given type can be contained in a structure type.
   static bool isValidElementType(Type type);
 
public:
   unsigned getPreferredAlignment(const DataLayout &dataLayout,
                                  DataLayoutEntryListRef params) const;
 
   bool areCompatible(DataLayoutEntryListRef oldLayout,
                      DataLayoutEntryListRef newLayout) const;
 
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMFixedVectorType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect fixed vector type, represents a sequence of elements of known
 /// length that can be processed as one.
 class LLVMFixedVectorType
     : public Type::TypeBase<LLVMFixedVectorType, Type,
-                            detail::LLVMTypeAndSizeStorage> {
+                            detail::LLVMTypeAndSizeStorage,
+                            SubElementTypeInterface::Trait> {
 public:
   /// Inherit base constructor.
   using Base::Base;
   using Base::getChecked;
 
   /// Gets or creates a fixed vector type containing `numElements` of
   /// `elementType` in the same context as `elementType`.
   static LLVMFixedVectorType get(Type elementType, unsigned numElements);
   static LLVMFixedVectorType
   getChecked(function_ref<InFlightDiagnostic()> emitError, Type elementType,
              unsigned numElements);
 
   /// Checks if the given type can be used in a vector type. This type supports
   /// only a subset of LLVM dialect types that don't have a built-in
   /// counter-part, e.g., pointers.
   static bool isValidElementType(Type type);
 
   /// Returns the element type of the vector.
-  Type getElementType();
+  Type getElementType() const;
 
   /// Returns the number of elements in the fixed vector.
   unsigned getNumElements();
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type elementType, unsigned numElements);
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType.
 //===----------------------------------------------------------------------===//
 
 /// LLVM dialect scalable vector type, represents a sequence of elements of
 /// unknown length that is known to be divisible by some constant. These
 /// elements can be processed as one in SIMD context.
 class LLVMScalableVectorType
     : public Type::TypeBase<LLVMScalableVectorType, Type,
-                            detail::LLVMTypeAndSizeStorage> {
+                            detail::LLVMTypeAndSizeStorage,
+                            SubElementTypeInterface::Trait> {
 public:
   /// Inherit base constructor.
   using Base::Base;
   using Base::getChecked;
 
   /// Gets or creates a scalable vector type containing a non-zero multiple of
   /// `minNumElements` of `elementType` in the same context as `elementType`.
   static LLVMScalableVectorType get(Type elementType, unsigned minNumElements);
   static LLVMScalableVectorType
   getChecked(function_ref<InFlightDiagnostic()> emitError, Type elementType,
              unsigned minNumElements);
 
   /// Checks if the given type can be used in a vector type.
   static bool isValidElementType(Type type);
 
   /// Returns the element type of the vector.
-  Type getElementType();
+  Type getElementType() const;
 
   /// Returns the scaling factor of the number of elements in the vector. The
   /// vector contains at least the resulting number of elements, or any non-zero
   /// multiple of this number.
   unsigned getMinNumElements();
 
   /// Verifies that the type about to be constructed is well-formed.
   static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
                               Type elementType, unsigned minNumElements);
+
+  void walkImmediateSubElements(function_ref<void(Attribute)> walkAttrsFn,
+                                function_ref<void(Type)> walkTypesFn) const;
 };
 
 //===----------------------------------------------------------------------===//
 // Printing and parsing.
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 /// Parses an LLVM dialect type.

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

 }
 
 unsigned
 LLVMArrayType::getPreferredAlignment(const DataLayout &dataLayout,
                                      DataLayoutEntryListRef params) const {
   return dataLayout.getTypePreferredAlignment(getElementType());
 }
 
+void LLVMArrayType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMFunctionType::isValidArgumentType(Type type) {
   return !type.isa<LLVMVoidType, LLVMFunctionType>();
 }
 

 }
 
 LLVMFunctionType LLVMFunctionType::clone(TypeRange inputs,
                                          TypeRange results) const {
   assert(results.size() == 1 && "expected a single result type");
   return get(results[0], llvm::to_vector(inputs), isVarArg());
 }
 
-Type LLVMFunctionType::getReturnType() { return getImpl()->getReturnType(); }
-ArrayRef<Type> LLVMFunctionType::getReturnTypes() {
+Type LLVMFunctionType::getReturnType() const {
+  return getImpl()->getReturnType();
+}
+ArrayRef<Type> LLVMFunctionType::getReturnTypes() const {
   return getImpl()->getReturnType();
 }
 
 unsigned LLVMFunctionType::getNumParams() {
   return getImpl()->getArgumentTypes().size();
 }
 
 Type LLVMFunctionType::getParamType(unsigned i) {
   return getImpl()->getArgumentTypes()[i];
 }
 
 bool LLVMFunctionType::isVarArg() const { return getImpl()->isVariadic(); }
 
-ArrayRef<Type> LLVMFunctionType::getParams() {
+ArrayRef<Type> LLVMFunctionType::getParams() const {
   return getImpl()->getArgumentTypes();
 }
 
 LogicalResult
 LLVMFunctionType::verify(function_ref<InFlightDiagnostic()> emitError,
                          Type result, ArrayRef<Type> arguments, bool) {
   if (!isValidResultType(result))
     return emitError() << "invalid function result type: " << result;
 
   for (Type arg : arguments)
     if (!isValidArgumentType(arg))
       return emitError() << "invalid function argument type: " << arg;
 
   return success();
 }
 
+void LLVMFunctionType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  for (Type type : llvm::concat<const Type>(getReturnTypes(), getParams()))
+    walkTypesFn(type);
+}
+
 //===----------------------------------------------------------------------===//
 // Pointer type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMPointerType::isValidElementType(Type type) {
   if (!type)
     return true;
   return isCompatibleOuterType(type)
if (extractPointerSpecValue(values, DLEntryPos::Abi) >
         extractPointerSpecValue(values, DLEntryPos::Preferred)) {
       return emitError(loc) << "preferred alignment is expected to be at least "
                                "as large as ABI alignment";
     }
   }
   return success();
 }
 
+void LLVMPointerType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Struct type.
 //===----------------------------------------------------------------------===//
 
 bool LLVMStructType::isValidElementType(Type type) {
   return !type.isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
                    LLVMFunctionType, LLVMTokenType, LLVMScalableVectorType>();
 }
if (extractStructSpecValue(values, StructDLEntryPos::Abi) >
         extractStructSpecValue(values, StructDLEntryPos::Preferred)) {
       return emitError(loc) << "preferred alignment is expected to be at least "
                                "as large as ABI alignment";
     }
   }
   return mlir::success();
 }
 
+void LLVMStructType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  for (Type type : getBody())
+    walkTypesFn(type);
+}
+
 //===----------------------------------------------------------------------===//
 // Vector types.
 //===----------------------------------------------------------------------===//
 
 /// Verifies that the type about to be constructed is well-formed.
 template <typename VecTy>
 static LogicalResult
 verifyVectorConstructionInvariants(function_ref<InFlightDiagnostic()> emitError,

 LLVMFixedVectorType
 LLVMFixedVectorType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                                 Type elementType, unsigned numElements) {
   assert(elementType && "expected non-null subtype");
   return Base::getChecked(emitError, elementType.getContext(), elementType,
                           numElements);
 }
 
-Type LLVMFixedVectorType::getElementType() {
+Type LLVMFixedVectorType::getElementType() const {
   return static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->elementType;
 }
 
 unsigned LLVMFixedVectorType::getNumElements() {
   return getImpl()->numElements;
 }
 
 bool LLVMFixedVectorType::isValidElementType(Type type) {
   return type.isa<LLVMPointerType, LLVMPPCFP128Type>();
 }
 
 LogicalResult
 LLVMFixedVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                             Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMFixedVectorType>(
       emitError, elementType, numElements);
 }
 
+void LLVMFixedVectorType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // LLVMScalableVectorType.
 //===----------------------------------------------------------------------===//
 
 LLVMScalableVectorType LLVMScalableVectorType::get(Type elementType,
                                                    unsigned minNumElements) {
   assert(elementType && "expected non-null subtype");
   return Base::get(elementType.getContext(), elementType, minNumElements);
 }
 
 LLVMScalableVectorType
 LLVMScalableVectorType::getChecked(function_ref<InFlightDiagnostic()> emitError,
                                    Type elementType, unsigned minNumElements) {
   assert(elementType && "expected non-null subtype");
   return Base::getChecked(emitError, elementType.getContext(), elementType,
                           minNumElements);
 }
 
-Type LLVMScalableVectorType::getElementType() {
+Type LLVMScalableVectorType::getElementType() const {
   return static_cast<detail::LLVMTypeAndSizeStorage *>(impl)->elementType;
 }
 
 unsigned LLVMScalableVectorType::getMinNumElements() {
   return getImpl()->numElements;
 }
 
 bool LLVMScalableVectorType::isValidElementType(Type type) {
   if (auto intType = type.dyn_cast<IntegerType>())
     return intType.isSignless();
 
   return isCompatibleFloatingPointType(type) || type.isa<LLVMPointerType>();
 }
 
 LogicalResult
 LLVMScalableVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
                                Type elementType, unsigned numElements) {
   return verifyVectorConstructionInvariants<LLVMScalableVectorType>(
       emitError, elementType, numElements);
 }
 
+void LLVMScalableVectorType::walkImmediateSubElements(
+    function_ref<void(Attribute)> walkAttrsFn,
+    function_ref<void(Type)> walkTypesFn) const {
+  walkTypesFn(getElementType());
+}
+
 //===----------------------------------------------------------------------===//
 // Utility functions.
 //===----------------------------------------------------------------------===//
 
 bool mlir::LLVM::isCompatibleOuterType(Type type) {
   // clang-format off
   if (type.isa<
       BFloat16Type,

mlir/unittests/Dialect/LLVMIR/LLVMTypeTest.cpp

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