[irdl] Add support for attribute wrappers

Author: math-fehr

mlir/include/mlir/Dialect/IRDL/AttributeWrapper.h

@@ -0,0 +1,144 @@
+//===- AttributeWrapper.h - IRDL type wrapper definition --------*- 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 declares wrappers around attribute and type definitions to
+// manipulate them in a unified way, using their names and a list of parameters
+// encoded as a list of attributes. These wrappers are necessary for IRDL, since
+// attributes and types don't have names, nor a way to interact with them in a
+// generic way.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_IRDL_ATTRIBUTEWRAPPER_H_
+#define MLIR_DIALECT_IRDL_ATTRIBUTEWRAPPER_H_
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/ExtensibleDialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "llvm/ADT/SmallString.h"
+
+namespace mlir {
+namespace irdl {
+
+/// A wrapper around an attribute definition to manipulate its attribute
+/// instances in a unified way, using a name and a list of parameters encoded as
+/// attributes.
+///
+/// If the attribute is defined in C++, CppAttributeWrapper should be preferred.
+class AttributeWrapper {
+public:
+  virtual ~AttributeWrapper(){};
+
+  /// Return the unique identifier of the attribute definition.
+  virtual TypeID getTypeID() = 0;
+
+  /// Check if the given attribute is an instance of the one wrapped.
+  virtual bool isCorrectAttribute(Attribute attr) = 0;
+
+  /// Get the parameters of an attribute. The attribute is expected to be an
+  /// instance of the wrapped attribute, which is checked with
+  /// `isCorrectAttribute`.
+  virtual SmallVector<Attribute> getAttributeParameters(Attribute attr) = 0;
+
+  /// Return the attribute definition name, including the dialect prefix.
+  virtual StringRef getName() = 0;
+
+  /// Instantiate the attribute from parameters.
+  /// It is expected that the amount of parameters is correct, which is checked
+  /// with `getParameterAmount`.
+  virtual Attribute
+  instantiateAttribute(function_ref<InFlightDiagnostic()> emitError,
+                       ArrayRef<Attribute> parameters) = 0;
+
+  /// Return the amount of parameters the attribute expects.
+  virtual size_t getParameterAmount() = 0;
+};
+
+/// A wrapper around a type definition to manipulate its type instances
+/// in an unified way, using a name and a list of parameters encoded as
+/// attributes. The wrappers also acts as an attribute wrapper, and expects
+/// the type to be nested in a `TypeAttr`.
+///
+/// If the type is defined as a C++ class, CppTypeWrapper should be preferred.
+class TypeWrapper : public AttributeWrapper {
+public:
+  virtual ~TypeWrapper(){};
+
+  /// Check if the given type is an instance of the one wrapped.
+  virtual bool isCorrectType(Type t) = 0;
+
+  /// Get the parameters of a type. The type is expected to be an instance of
+  /// the wrapped type, which is checked with `isCorrectType`.
+  virtual SmallVector<Attribute> getTypeParameters(Type t) = 0;
+
+  /// Instantiate the type from parameters.
+  /// It is expected that the amount of parameters is correct, which is checked
+  /// with `getParameterAmount`.
+  virtual Type instantiateType(function_ref<InFlightDiagnostic()> emitError,
+                               ArrayRef<Attribute> parameters) = 0;
+
+  /// Check if the given attribute is a `TypeAttr`, and that it contains an
+  /// instance of the type wrapped.
+  bool isCorrectAttribute(Attribute attr) override;
+
+  /// Get the parameters of a type. The type is expected to be nested in a
+  /// `TypeAttr`, and to be an instance of the wrapped type, which is checked
+  /// with `isCorrectType`.
+  SmallVector<Attribute> getAttributeParameters(Attribute attr) override;
+
+  /// Instantiate the type from parameters, and wrap it in a `TypeAttr`.
+  /// It is expected that the amount of parameters is correct, which is checked
+  /// with `getParameterAmount`.
+  Attribute instantiateAttribute(function_ref<InFlightDiagnostic()> emitError,
+                                 ArrayRef<Attribute> parameters) override;
+};
+
+using AttributeWrapperPtr = AttributeWrapper *;
+using TypeWrapperPtr = TypeWrapper *;
+
+/// A wrapper around an attribute definition defined with a C++ class to
+/// manipulate its attribute instances in a unified way, using a name and a list
+/// of parameters encoded as attributes.
+template <typename A>
+class CppAttributeWrapper : public AttributeWrapper {
+public:
+  TypeID getTypeID() override { return TypeID::get<A>(); }
+
+  /// Get the parameters of an attribute.
+  virtual SmallVector<Attribute> getAttributeParameters(A attr) = 0;
+
+  SmallVector<Attribute> getAttributeParameters(Attribute attr) override {
+    return getAttributeParameters(cast<A>(attr));
+  };
+
+  bool isCorrectAttribute(Attribute attr) override { return isa<A>(attr); }
+};
+
+/// A wrapper around a type definition defined with a C++ class to
+/// manipulate its type instances in a unified way, using a name and a list
+/// of parameters encoded as attributes.
+template <typename T>
+class CppTypeWrapper : public TypeWrapper {
+public:
+  TypeID getTypeID() override { return TypeID::get<T>(); }
+
+  /// Get the parameters of a type.
+  virtual SmallVector<Attribute> getTypeParameters(T t) = 0;
+
+  SmallVector<Attribute> getTypeParameters(Type type) override {
+    return getTypeParameters(cast<T>(type));
+  };
+
+  bool isCorrectType(Type type) override { return isa<T>(type); }
+};
+
+} // namespace irdl
+} // namespace mlir
+
+#endif // MLIR_DIALECT_IRDL_ATTRIBUTEWRAPPER_H_

mlir/include/mlir/Dialect/IRDL/IR/CMakeLists.txt

@@ -20,3 +20,10 @@ mlir_tablegen(IRDLTypesGen.h.inc -gen-typedef-decls)
 mlir_tablegen(IRDLTypesGen.cpp.inc -gen-typedef-defs)
 add_public_tablegen_target(MLIRIRDLTypesIncGen)
 add_dependencies(mlir-generic-headers MLIRIRDLTypesIncGen)
+
+# Add IRDL attributes
+set(LLVM_TARGET_DEFINITIONS IRDLAttributes.td)
+mlir_tablegen(IRDLAttributes.h.inc -gen-attrdef-decls)
+mlir_tablegen(IRDLAttributes.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRIRDLAttributesIncGen)
+add_dependencies(mlir-generic-headers MLIRIRDLAttributesIncGen)

mlir/include/mlir/Dialect/IRDL/IR/IRDL.h

@@ -13,8 +13,10 @@
 #ifndef MLIR_DIALECT_IRDL_IR_IRDL_H_
 #define MLIR_DIALECT_IRDL_IR_IRDL_H_
 
+#include "mlir/Dialect/IRDL/AttributeWrapper.h"
 #include "mlir/Dialect/IRDL/IR/IRDLInterfaces.h"
 #include "mlir/Dialect/IRDL/IR/IRDLTraits.h"
+#include "mlir/Dialect/IRDL/IRDLContext.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
@@ -37,6 +39,9 @@ class OpDefAttr;
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/IRDL/IR/IRDLTypesGen.h.inc"
 
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/IRDL/IR/IRDLAttributes.h.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/IRDL/IR/IRDLOps.h.inc"
 

mlir/include/mlir/Dialect/IRDL/IR/IRDL.td

@@ -70,6 +70,13 @@ def IRDL_Dialect : Dialect {
   }];
 
   let useDefaultTypePrinterParser = 1;
+  let useDefaultAttributePrinterParser = 1;
+
+  let extraClassDeclaration = [{
+  public:
+    /// Contains registered attribute and type wrappers.
+    ::mlir::irdl::IRDLContext irdlContext;
+  }];
 
   let name = "irdl";
   let cppNamespace = "::mlir::irdl";

mlir/include/mlir/Dialect/IRDL/IR/IRDLAttributes.td

@@ -0,0 +1,54 @@
+//===- IRDLAttributes.td - IR Definition Language Dialect --*- 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 declares the attributes used in IRDL.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_IRDL_IR_IRDLATTRIBUTES
+#define MLIR_DIALECT_IRDL_IR_IRDLATTRIBUTES
+
+include "IRDL.td"
+include "mlir/IR/AttrTypeBase.td"
+
+class IRDL_AttrDef<string name, string attrMnemonic, list<Trait> traits = []>
+  : AttrDef<IRDL_Dialect, name, traits, "::mlir::Attribute"> {
+  let mnemonic = attrMnemonic;
+}
+
+def TypeOrAttrDefRefAttr : IRDL_AttrDef<"TypeOrAttrDefRef",
+                                        "type_or_attr_def_ref"> {
+  let summary = "reference to a type or attribute definition";
+  let description = [{
+    An `irdl.type_or_attr_def_ref` references a type or attribute definition.
+    A type or attribute definition can either be a `TypeWrapper *` or
+    an `AttrWrapper *`, which refers to a C++-defined type or attribute,
+    or a `SymbolRefAttr` which refers to a type or attribute defined with IRDL.
+  }];
+  let parameters =
+    (ins OptionalParameter<"mlir::irdl::TypeWrapperPtr">:$typeWrapper,
+         OptionalParameter<"mlir::irdl::AttributeWrapperPtr">:$attrWrapper,
+         OptionalParameter<"mlir::SymbolRefAttr">:$symRef);
+
+  let builders = [
+    AttrBuilder<(ins "mlir::irdl::TypeWrapper*":$typeWrapper), [{
+      return $_get($_ctxt, typeWrapper, nullptr, mlir::SymbolRefAttr());
+    }]>,
+    AttrBuilder<(ins "mlir::irdl::AttributeWrapper*":$attrWrapper), [{
+      return $_get($_ctxt, nullptr, attrWrapper, mlir::SymbolRefAttr());
+    }]>,
+    AttrBuilderWithInferredContext<(ins "mlir::SymbolRefAttr":$symRef), [{
+      return $_get(symRef.getContext(), nullptr, nullptr, symRef);
+    }]>
+  ];
+
+  let genVerifyDecl = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+#endif // MLIR_DIALECT_IRDL_IR_IRDLATTRIBUTES

mlir/include/mlir/Dialect/IRDL/IR/IRDLOps.td

@@ -14,6 +14,7 @@
 #define MLIR_DIALECT_IRDL_IR_IRDLOPS
 
 include "IRDL.td"
+include "IRDLAttributes.td"
 include "IRDLTypes.td"
 include "IRDLInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
@@ -302,8 +303,8 @@ def IRDL_Parametric : IRDL_ConstraintOp<"parametric",
   let description = [{
     `irdl.parametric` defines a constraint that accepts only a single type
     or attribute base. The attribute base is defined by a symbolic reference
-    to the corresponding definition. It will additionally constraint the
-    parameters of the type/attribute.
+    to the corresponding definition, or a type/attribute wrapper. It will
+    additionally constraint the parameters of the type/attribute.
 
     Example:
 
@@ -325,7 +326,7 @@ def IRDL_Parametric : IRDL_ConstraintOp<"parametric",
     for any `T` takes a `cmath.complex` with parameter `T` and returns a `T`.
   }];
 
-  let arguments = (ins SymbolRefAttr:$base_type,
+  let arguments = (ins TypeOrAttrDefRefAttr:$base_type,
                        Variadic<IRDL_AttributeType>:$args);
   let results = (outs IRDL_AttributeType:$output);
   let assemblyFormat = " $base_type `<` $args `>` ` ` attr-dict ";

mlir/include/mlir/Dialect/IRDL/IRDLContext.h

@@ -0,0 +1,75 @@
+//===- IRDLContext.h - IRDL context -----------------------------*- C++ -*-===//
+//
+// This file is licensed 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Manages the registration context of IRDL dialects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_IRDL_IRDLCONTEXT_H_
+#define MLIR_DIALECT_IRDL_IRDLCONTEXT_H_
+
+#include "mlir/Dialect/IRDL/AttributeWrapper.h"
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/StringMap.h"
+
+namespace mlir {
+namespace irdl {
+
+/// Context for the runtime registration of IRDL dialect definitions.
+/// This class keeps track of all the attribute and types defined in C++
+/// that can be used in IRDL.
+class IRDLContext {
+  llvm::StringMap<std::unique_ptr<AttributeWrapper>> attributes;
+  llvm::StringMap<std::unique_ptr<TypeWrapper>> types;
+  DenseMap<TypeID, AttributeWrapper *> typeIDToAttributeWrapper;
+  DenseMap<TypeID, TypeWrapper *> typeIDToTypeWrapper;
+
+public:
+  IRDLContext();
+
+  /// Add a concrete attribute wrapper to IRDL.
+  /// The attribute definition wrapped can then be used in IRDL with its name.
+  template <typename A>
+  void addAttributeWrapper() {
+    addAttributeWrapper(std::make_unique<A>());
+  }
+
+  /// Add an attribute wrapper to IRDL.
+  /// The attribute definition wrapped can then be used in IRDL with its name.
+  void addAttributeWrapper(std::unique_ptr<AttributeWrapper> wrapper);
+
+  AttributeWrapper *getAttributeWrapper(StringRef typeName);
+  AttributeWrapper *getAttributeWrapper(TypeID typeID);
+
+  /// Add a concrete type wrapper to IRDL.
+  /// The type definition wrapped can then be used in IRDL with its name.
+  template <typename T>
+  void addTypeWrapper() {
+    addTypeWrapper(std::make_unique<T>());
+  }
+
+  /// Add a type wrapper to IRDL.
+  /// The type definition wrapped can then be used in IRDL with its name.
+  void addTypeWrapper(std::unique_ptr<TypeWrapper> wrapper);
+
+  TypeWrapper *getTypeWrapper(StringRef typeName);
+  TypeWrapper *getTypeWrapper(TypeID typeID);
+
+  llvm::StringMap<std::unique_ptr<AttributeWrapper>> const &getAllAttributes() {
+    return attributes;
+  }
+
+  llvm::StringMap<std::unique_ptr<TypeWrapper>> const &getAllTypes() {
+    return types;
+  }
+};
+
+} // namespace irdl
+} // namespace mlir
+
+#endif // MLIR_DIALECT_IRDL_IRDLCONTEXT_H_

mlir/include/mlir/Dialect/IRDL/IRDLVerifiers.h

@@ -13,6 +13,7 @@
 #ifndef MLIR_DIALECT_IRDL_IRDLVERIFIERS_H
 #define MLIR_DIALECT_IRDL_IRDLVERIFIERS_H
 
+#include "mlir/Dialect/IRDL/AttributeWrapper.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
@@ -97,6 +98,28 @@ class IsConstraint : public Constraint {
   Attribute expectedAttribute;
 };
 
+/// A constraint that checks that an attribute that has an attribute wrapper is
+/// of a specific attribute definition, and that all of its parameters satisfy
+/// the given constraints.
+/// This class also takes care of types, by expecting a `TypeAttr` attribute.
+class ParametricConstraint : public Constraint {
+public:
+  ParametricConstraint(::mlir::irdl::AttributeWrapper *expectedAttribute,
+                       SmallVector<unsigned> constraints)
+      : expectedAttribute(expectedAttribute),
+        constraints(std::move(constraints)) {}
+
+  virtual ~ParametricConstraint() = default;
+
+  LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
+                       Attribute attr,
+                       ConstraintVerifier &context) const override;
+
+private:
+  ::mlir::irdl::AttributeWrapper *expectedAttribute;
+  SmallVector<unsigned> constraints;
+};
+
 /// A constraint that checks that an attribute is of a
 /// specific dynamic attribute definition, and that all of its parameters
 /// satisfy the given constraints.