[mlir] Allows to query traits from types and attributes

[math-fehr]

Summary

Types and attributes now have a hasTrait function that allow users to check
if a type defines a trait.
Also, AbstractType and AbstractAttribute has now a hasTraitFn field to carry
the implementation of the hasTrait function of the concrete type or attribute.
This patch also adds the remaining functions to access type and attribute traits
in TableGen.

[rriddle]

Looks great, thanks!

(Could I shamelessly ask you to do the same for attributes?)

[rriddle]

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

2

The filename here is off.

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

45

This include should be at the top of the file.

[math-fehr]

Address comments, and add support for traits in attributes

[math-fehr]

Hey @rriddle, gentle ping on this!

[ftynse]

Hi Matthieu,
River is out of office. This patch is approved, do you need help landing it?

[math-fehr]

Hi Alex!
Yes, would it be possible to land it for me?
Thanks!

mlir/include/mlir/IR/AttributeSupport.h

 //===----------------------------------------------------------------------===//
 // AbstractAttribute
 //===----------------------------------------------------------------------===//
 
 /// This class contains all of the static information common to all instances of
 /// a registered Attribute.
 class AbstractAttribute {
 public:
+  using HasTraitFn = llvm::unique_function<bool(TypeID) const>;
+
   /// Look up the specified abstract attribute in the MLIRContext and return a
   /// reference to it.
   static const AbstractAttribute &lookup(TypeID typeID, MLIRContext *context);
 
   /// This method is used by Dialect objects when they register the list of
   /// attributes they contain.
-  template <typename T> static AbstractAttribute get(Dialect &dialect) {
-    return AbstractAttribute(dialect, T::getInterfaceMap(), T::getTypeID());
+  template <typename T>
+  static AbstractAttribute get(Dialect &dialect) {
+    return AbstractAttribute(dialect, T::getInterfaceMap(), T::getHasTraitFn(),
+                             T::getTypeID());
   }
 
   /// Return the dialect this attribute was registered to.
   Dialect &getDialect() const { return const_cast<Dialect &>(dialect); }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this attribute, null otherwise. This should not be used
   /// directly.
-  template <typename T> typename T::Concept *getInterface() const {
+  template <typename T>
+  typename T::Concept *getInterface() const {
     return interfaceMap.lookup<T>();
   }
 
   /// Returns true if the attribute has the interface with the given ID
   /// registered.
   bool hasInterface(TypeID interfaceID) const {
     return interfaceMap.contains(interfaceID);
   }
 
+  /// Returns true if the attribute has a particular trait.
+  template <template <typename T> class Trait>
+  bool hasTrait() const {
+    return hasTraitFn(TypeID::get<Trait>());
+  }
+
+  /// Returns true if the attribute has a particular trait.
+  bool hasTrait(TypeID traitID) const { return hasTraitFn(traitID); }
+
   /// Return the unique identifier representing the concrete attribute class.
   TypeID getTypeID() const { return typeID; }
 
 private:
   AbstractAttribute(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
-                    TypeID typeID)
+                    HasTraitFn &&hasTrait, TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
-        typeID(typeID) {}
+        hasTraitFn(std::move(hasTrait)), typeID(typeID) {}
 
   /// Give StorageUserBase access to the mutable lookup.
   template <typename ConcreteT, typename BaseT, typename StorageT,
             typename UniquerT, template <typename T> class... Traits>
   friend class detail::StorageUserBase;
 
   /// Look up the specified abstract attribute in the MLIRContext and return a
   /// (mutable) pointer to it. Return a null pointer if the attribute could not
   /// be found in the context.
   static AbstractAttribute *lookupMutable(TypeID typeID, MLIRContext *context);
 
   /// This is the dialect that this attribute was registered to.
   const Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this attribute.
   detail::InterfaceMap interfaceMap;
 
+  /// Function to check if the attribute has a particular trait.
+  HasTraitFn hasTraitFn;
+
   /// The unique identifier of the derived Attribute class.
   const TypeID typeID;
 };
 
 //===----------------------------------------------------------------------===//
 // AttributeStorage
 //===----------------------------------------------------------------------===//
 

mlir/include/mlir/IR/Attributes.h

public:
   const void *getAsOpaquePointer() const { return impl; }
   /// Construct an attribute from the opaque pointer representation.
   static Attribute getFromOpaquePointer(const void *ptr) {
     return Attribute(reinterpret_cast<const ImplType *>(ptr));
   }
 
   friend ::llvm::hash_code hash_value(Attribute arg);
 
+  /// Returns true if the type was registered with a particular trait.
+  template <template <typename T> class Trait>
+  bool hasTrait() {
+    return getAbstractAttribute().hasTrait<Trait>();
+  }
+
   /// Return the abstract descriptor for this attribute.
   const AbstractTy &getAbstractAttribute() const {
     return impl->getAbstractAttribute();
   }
 
 protected:
   ImplType *impl;
 };

mlir/include/mlir/IR/OpBase.td

 
 // PredTrait is a trait implemented by way of a predicate on an entity.
 class PredTrait<string descr, Pred pred> : Trait {
   string summary = descr;
   Pred predicate = pred;
 }
 
 //===----------------------------------------------------------------------===//
+// TypeTrait definitions
+//===----------------------------------------------------------------------===//
+
+// TypeTrait represents a trait regarding a type.
+// TODO: Remove this class in favor of using Trait.
+class TypeTrait;
+
+// These classes are used to define type specific traits.
+class NativeTypeTrait<string name> : NativeTrait<name, "Type">, TypeTrait;
+class ParamNativeTypeTrait<string prop, string params>
+    : ParamNativeTrait<prop, params, "Type">, TypeTrait;
+class GenInternalTypeTrait<string prop>
+    : GenInternalTrait<prop, "Type">, TypeTrait;
+class PredTypeTrait<string descr, Pred pred>
+    : PredTrait<descr, pred>, TypeTrait;
+
+//===----------------------------------------------------------------------===//
+// AttrTrait definitions
+//===----------------------------------------------------------------------===//
+
+// AttrTrait represents a trait regarding an attribute.
+// TODO: Remove this class in favor of using Trait.
+class AttrTrait;
+
+// These classes are used to define attribute specific traits.
+class NativeAttrTrait<string name> : NativeTrait<name, "Attribute">, AttrTrait;
+class ParamNativeAttrTrait<string prop, string params>
+    : ParamNativeTrait<prop, params, "Attribute">, AttrTrait;
+class GenInternalAttrTrait<string prop>
+    : GenInternalTrait<prop, "Attribute">, AttrTrait;
+class PredAttrTrait<string descr, Pred pred>
+    : PredTrait<descr, pred>, AttrTrait;
+
+//===----------------------------------------------------------------------===//
 // OpTrait definitions
 //===----------------------------------------------------------------------===//
 
 // OpTrait represents a trait regarding an operation.
 // TODO: Remove this class in favor of using Trait.
 class OpTrait;
 
 // These classes are used to define operation specific traits.

mlir/include/mlir/IR/StorageUniquerSupport.h

ConcreteType getInstance() const {
     return *static_cast<const ConcreteType *>(trait);
   }
 };
 
 //===----------------------------------------------------------------------===//
 // 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
 /// StorageUniquer. Clients are not expected to interact with this class
 /// directly.
 template <typename ConcreteT, typename BaseT, typename StorageT,
           typename UniquerT, template <typename T> class... Traits>
 class StorageUserBase : public BaseT, public Traits<ConcreteT>... {
 public:
   using BaseT::BaseT;
 
   /// Utility declarations for the concrete attribute class.
   using Base = StorageUserBase<ConcreteT, BaseT, StorageT, UniquerT, Traits...>;
   using ImplType = StorageT;
+  using HasTraitFn = bool (*)(TypeID);
 
   /// Return a unique identifier for the concrete type.
   static TypeID getTypeID() { return TypeID::get<ConcreteT>(); }
 
   /// Provide an implementation of 'classof' that compares the type id of the
   /// provided value with that of the concrete type.
   template <typename T> static bool classof(T val) {
     static_assert(std::is_convertible<ConcreteT, T>::value,
                   "casting from a non-convertible type");
     return val.getTypeID() == getTypeID();
   }
 
   /// Returns an interface map for the interfaces registered to this storage
   /// user. This should not be used directly.
   static detail::InterfaceMap getInterfaceMap() {
     return detail::InterfaceMap::template get<Traits<ConcreteT>...>();
   }
 
+  /// Returns the function that returns true if the given Trait ID matches the
+  /// IDs of any of the traits defined by the storage user.
+  static HasTraitFn getHasTraitFn() {
+    return [](TypeID id) {
+      return storage_user_base_impl::hasTrait<Traits...>(id);
+    };
+  }
+
   /// Attach the given models as implementations of the corresponding interfaces
   /// for the concrete storage user class. The type must be registered with the
   /// context, i.e. the dialect to which the type belongs must be loaded. The
   /// call will abort otherwise.
   template <typename... IfaceModels>
   static void attachInterface(MLIRContext &context) {
     typename ConcreteT::AbstractTy *abstract =
         ConcreteT::AbstractTy::lookupMutable(TypeID::get<ConcreteT>(),

mlir/include/mlir/IR/TypeSupport.h

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

mlir/include/mlir/IR/Types.h

public:
   /// Methods for supporting PointerLikeTypeTraits.
   const void *getAsOpaquePointer() const {
     return static_cast<const void *>(impl);
   }
   static Type getFromOpaquePointer(const void *pointer) {
     return Type(reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
   }
 
+  /// Returns true if the type was registered with a particular trait.
+  template <template <typename T> class Trait>
+  bool hasTrait() {
+    return getAbstractType().hasTrait<Trait>();
+  }
+
   /// Return the abstract type descriptor for this type.
   const AbstractTy &getAbstractType() { return impl->getAbstractType(); }
 
 protected:
   ImplType *impl;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, Type type) {

mlir/test/IR/traits.mlir

func @graph_region_cant_have_blocks() {
   test.graph_region {
     // expected-error@-1 {{'test.graph_region' op expects graph region #0 to have 0 or 1 blocks}}
   ^bb42:
     br ^bb43
   ^bb43:
     "terminator"() : () -> ()
   }
 }
+
+// -----
+
+// Check that we can query traits in types
+func @succeeded_type_traits() {
+  // CHECK: "test.result_type_with_trait"() : () -> !test.test_type_with_trait
+  "test.result_type_with_trait"() : () -> !test.test_type_with_trait
+  return
+}
+
+// -----
+
+// Check that we can query traits in types
+func @failed_type_traits() {
+  // expected-error@+1 {{result type should have trait 'TestTypeTrait'}}
+  "test.result_type_with_trait"() : () -> i32
+  return
+}
+
+// -----
+
+// Check that we can query traits in attributes
+func @succeeded_attr_traits() {
+  // CHECK: "test.attr_with_trait"() {attr = #test.attr_with_trait} : () -> ()
+  "test.attr_with_trait"() {attr = #test.attr_with_trait} : () -> ()
+  return
+}
+
+// -----
+
+// Check that we can query traits in attributes
+func @failed_attr_traits() {
+  // expected-error@+1 {{'attr' attribute should have trait 'TestAttrTrait'}}
+  "test.attr_with_trait"() {attr = 42 : i32} : () -> ()
+  return
+}
 No newline at end of file

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

 
 #ifndef TEST_ATTRDEFS
 #define TEST_ATTRDEFS
 
 // To get the test dialect definition.
 include "TestOps.td"
 
 // All of the attributes will extend this class.
-class Test_Attr<string name> : AttrDef<Test_Dialect, name>;
+class Test_Attr<string name, list<Trait> traits = []>
+    : AttrDef<Test_Dialect, name, traits>;
 
 def SimpleAttrA : Test_Attr<"SimpleA"> {
   let mnemonic = "smpla";
 }
 
 // A more complex parameterized attribute.
 def CompoundAttrA : Test_Attr<"CompoundA"> {
   let mnemonic = "cmpnd_a";

 
 // An attribute testing AttributeSelfTypeParameter.
 def AttrWithTypeBuilder : Test_Attr<"AttrWithTypeBuilder"> {
   let mnemonic = "attr_with_type_builder";
   let parameters = (ins "::mlir::IntegerAttr":$attr);
   let typeBuilder = "$_attr.getType()";
 }
 
+def TestAttrTrait : NativeAttrTrait<"TestAttrTrait">;
+
+// The definition of a singleton attribute that has a trait.
+def AttrWithTrait : Test_Attr<"AttrWithTrait", [TestAttrTrait]> {
+  let mnemonic = "attr_with_trait";
+  let parameters = (ins );
+}
+
 #endif // TEST_ATTRDEFS

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTATTRIBUTES_H
 #define MLIR_TESTATTRIBUTES_H
 
 #include <tuple>
 
+#include "TestTraits.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 
 #include "TestAttrInterfaces.h.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.h.inc"
 
 #endif // MLIR_TESTATTRIBUTES_H

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

: TEST_Op<"string_attr_pretty_name",
            [DeclareOpInterfaceMethods<OpAsmOpInterface>]> {
   let arguments = (ins StrArrayAttr:$names);
   let results = (outs Variadic<I32>:$r);
 
   let printer = [{ return ::print(p, *this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
+// This operation requires its return type to have the trait 'TestTypeTrait'.
+def ResultTypeWithTraitOp : TEST_Op<"result_type_with_trait", []> {
+  let results = (outs AnyType);
+
+  let verifier = [{
+    if((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
+      return success();
+    return this->emitError("result type should have trait 'TestTypeTrait'");
+  }];
+}
+
+// This operation requires its "attr" attribute to have the
+// trait 'TestAttrTrait'.
+def AttrWithTraitOp : TEST_Op<"attr_with_trait", []> {
+  let arguments = (ins AnyAttr:$attr);
+
+  let verifier = [{
+    if (this->attr().hasTrait<AttributeTrait::TestAttrTrait>())
+      return success();
+    return this->emitError("'attr' attribute should have trait 'TestAttrTrait'");
+  }];
+}
+
+
 //===----------------------------------------------------------------------===//
 // Test Locations
 //===----------------------------------------------------------------------===//
 
 def TestLocationSrcOp : TEST_Op<"loc_src"> {
   let arguments = (ins I32:$input);
   let results = (outs I32:$output);
 }

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

+//===- TestTraits.h - MLIR Test Traits --------------------------*- 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 traits defined by the TestDialect.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TESTTRAITS_H
+#define MLIR_TESTTRAITS_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Types.h"
+
+namespace mlir {
+namespace TypeTrait {
+
+/// A trait defined on types for testing purposes.
+template <typename ConcreteType>
+class TestTypeTrait : public TypeTrait::TraitBase<ConcreteType, TestTypeTrait> {
+};
+
+} // namespace TypeTrait
+
+namespace AttributeTrait {
+
+/// A trait defined on attributes for testing purposes.
+template <typename ConcreteType>
+class TestAttrTrait
+    : public AttributeTrait::TraitBase<ConcreteType, TestAttrTrait> {};
+
+} // namespace AttributeTrait
+} // namespace mlir
+
+#endif // MLIR_TESTTRAITS_H

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

]> {
   }];
 }
 
 def TestMemRefElementType : Test_Type<"TestMemRefElementType",
                                       [MemRefElementTypeInterface]> {
   let mnemonic = "memref_element";
 }
 
+def TestTypeTrait : NativeTypeTrait<"TestTypeTrait">;
+
+// The definition of a singleton type that has a trait.
+def TestTypeWithTrait : Test_Type<"TestTypeWithTrait", [TestTypeTrait]> {
+  let mnemonic = "test_type_with_trait";
+}
+
 #endif // TEST_TYPEDEFS

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TESTTYPES_H
 #define MLIR_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/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 
 namespace mlir {
struct FieldInfo {
   FieldInfo allocateInto(TypeStorageAllocator &alloc) const {
     return FieldInfo{alloc.copyInto(name), type};
   }
 };
 
 } // namespace test
 } // namespace mlir
 
 #include "TestTypeInterfaces.h.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "TestTypeDefs.h.inc"
 
 namespace mlir {
 namespace test {
 
 /// Storage for simple named recursive types, where the type is identified by