[ODS] Extra Concrete Declarations and Definitions under Traits

[amandatang]

Summary

Support extra concrete class declarations and definitions under NativeTrait that get injected into the class that specifies the trait. Extra declarations and definitions can be passed in as template arguments for NativeOpTraitNativeAttrTrait and NativeTypeTrait.

Usage examples of this feature include:

• Creating a wrapper Trait for authoring inferReturnTypes with the OpAdaptor by specifying necessary Op specific declarations and definitions directly in the trait
• Refactoring the InferTensorType trait

[amandatang]

clang format

[amandatang]

Quick fix

[jpienaar]

mlir/include/mlir/IR/OpBase.td

1963

Could you document these here?

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

241

You can remove this TODO now as not relevant anymore here

246

ArrayRef?

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

223

Does this end up only inserting inferReturnTypes ?

mlir/include/mlir/TableGen/Trait.h

71

s/type/instance/ ? given you can have Type's with this too

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1363

This cast shouldn't be needed. The adaptor returns a TypedValue<MemRefType>, so that getType() is of type MemRefType.

Now what is true, is that nothing has verified it until here along the build path (but building with invalid inputs is not expected to work) while along verify path this happens post verifying type, so should be fine in either case.

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

858

This is going to constantly resize the string below, could that be avoided? (same in other spots)

[amandatang]

Address comments

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

223

Yes

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

1363

The getSource() method in the Op itself returns TypedValue<MemRefType>, but the equivalent method in the adaptor returns a ValueT, which requires a cast.

[amandatang]

Update documentation

[jpienaar]

Looks good overall, thanks

mlir/docs/Traits.md

107 ↗

(On Diff #538839)

s/properties/fields/ ? (at least seems from TableGen reference that's the more common name)

110 ↗

(On Diff #538839)

They will be formatted too with variables being replaced

mlir/include/mlir/IR/AttrTypeBase.td

45

I thought I have asked about this, but why not have the same template args on NativeTrait too?

That way you cold just do NativeTrait<name, "Type", extraTypeDeclaration, extraTypeDefinition> here.

[amandatang]

Address final comments

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

 // ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_ExtractStridedMetadataOp : MemRef_Op<"extract_strided_metadata", [
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
     Pure,
     SameVariadicResultSize,
     ViewLikeOpInterface,
-    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+    InferTypeOpInterfaceAdaptor]> {
   let summary = "Extracts a buffer base with offset and strides";
   let description = [{
     Extracts a base buffer, offset and strides. This op allows additional layers
     of transformations and foldings to be added as lowering progresses from
     higher-level dialect to lower-level dialects such as the LLVM dialect.
 
     The op requires a strided memref source operand. If the source operand is not
     a strided memref, then verification fails.

mlir/include/mlir/IR/AttrTypeBase.td

 
 include "mlir/IR/OpBase.td"
 
 //-------------------------------------------------------------------------===//
 // AttrTrait definitions
 //===----------------------------------------------------------------------===//
 
 // These classes are used to define attribute specific traits.
-class NativeAttrTrait<string name> : NativeTrait<name, "Attribute">;
+class NativeAttrTrait<string name,
+                      code extraAttrDeclaration = [{}],
+                      code extraAttrDefinition = [{}]>
+    : NativeTrait<name, "Attribute"> {
+  let extraConcreteClassDeclaration = extraAttrDeclaration;
+  let extraConcreteClassDefinition = extraAttrDefinition;
+}
+
 class ParamNativeAttrTrait<string prop, string params>
     : ParamNativeTrait<prop, params, "Attribute">;
 class GenInternalAttrTrait<string prop> : GenInternalTrait<prop, "Attribute">;
 class PredAttrTrait<string descr, Pred pred> : PredTrait<descr, pred>;
 
 //===----------------------------------------------------------------------===//
 // TypeTrait definitions
 //===----------------------------------------------------------------------===//
 
 // These classes are used to define type specific traits.
-class NativeTypeTrait<string name> : NativeTrait<name, "Type">;
+class NativeTypeTrait<string name,
+                      code extraTypeDeclaration = [{}],
+                      code extraTypeDefinition = [{}]>
+    : NativeTrait<name, "Type"> {
+  let extraConcreteClassDeclaration = extraTypeDeclaration;
+  let extraConcreteClassDefinition = extraTypeDefinition;
+}
+
 class ParamNativeTypeTrait<string prop, string params>
     : ParamNativeTrait<prop, params, "Type">;
 class GenInternalTypeTrait<string prop> : GenInternalTrait<prop, "Type">;
 class PredTypeTrait<string descr, Pred pred> : PredTrait<descr, pred>;
 
 //===----------------------------------------------------------------------===//
 // Builders
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/OpBase.td

 }
 
 // NativeTrait corresponds to the MLIR C++ trait mechanism. The purpose to wrap
 // around C++ symbol string with this class is to make traits specified for
 // entities in TableGen less alien and more integrated.
 class NativeTrait<string name, string entityType> : Trait {
   string trait = name;
   string cppNamespace = "::mlir::" # entityType # "Trait";
+
+  code extraConcreteClassDeclaration = "";
+  code extraConcreteClassDefinition = "";
 }
 
 // ParamNativeTrait corresponds to the template-parameterized traits in the C++
 // implementation. MLIR uses nested class templates to implement such traits
 // leading to constructs of the form "TraitName<Parameters>::Impl". Use the
 // value in `prop` as the trait name and the value in `params` as parameters to
 // construct the native trait class name.
 class ParamNativeTrait<string prop, string params, string entityType>

 // OpTrait definitions
 //===----------------------------------------------------------------------===//
 
 // A trait that describes the structure of operation will be marked with
 // `StructuralOpTrait` and they will be verified first.
 class StructuralOpTrait;
 
 // These classes are used to define operation specific traits.
-class NativeOpTrait<string name, list<Trait> traits = []>
+class NativeOpTrait<string name, list<Trait> traits = [],
+                    code extraOpDeclaration = [{}],
+                    code extraOpDefinition = [{}]>
     : NativeTrait<name, "Op"> {
   // Specify the list of traits that need to be verified before the verification
   // of this NativeOpTrait.
   list<Trait> dependentTraits = traits;
+
+  let extraConcreteClassDeclaration = extraOpDeclaration;
+  let extraConcreteClassDefinition = extraOpDefinition;
 }
 class ParamNativeOpTrait<string prop, string params,
                          list<Trait> traits = []>
     : ParamNativeTrait<prop, params, "Op"> {
   // Specify the list of traits that need to be verified before the verification
   // of this ParamNativeOpTrait.
   list<Trait> dependentTraits = traits;
 }

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

private:
   // if interpreted as shape.
   ValueShapeMapFn valueToShape;
 };
 
 namespace detail {
 // Helper function to infer return tensor returns types given element and
 // shape inference function.
 //
 // TODO: Consider generating typedefs for trait member functions if this usage
 // becomes more common.
 LogicalResult inferReturnTensorTypes(
-    function_ref<
-        LogicalResult(MLIRContext *, std::optional<Location> location,
-                      ValueShapeRange operands, DictionaryAttr attributes,
-                      OpaqueProperties properties, RegionRange regions,
-                      SmallVectorImpl<ShapedTypeComponents> &retComponents)>
-        componentTypeFn,
-    MLIRContext *context, std::optional<Location> location, ValueRange operands,
-    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+    SmallVector<ShapedTypeComponents, 2> retComponents,
     SmallVectorImpl<Type> &inferredReturnTypes);
 
 /// Verifies that the inferred result types match the actual result types for
 /// the op. Precondition: op implements InferTypeOpInterface.
 LogicalResult verifyInferredResultTypes(Operation *op);
 } // namespace detail
 
 namespace OpTrait {
 template <typename ConcreteType>
 class InferTensorType;
 } // namespace OpTrait
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/InferTypeOpInterface.h.inc"
 
 namespace mlir {
 namespace OpTrait {
 
+template <typename ConcreteType>
+class InferTypeOpInterfaceAdaptor :
+    public TraitBase<ConcreteType, InferTypeOpInterfaceAdaptor> {};
+
 /// Tensor type inference trait that constructs a tensor from the inferred
 /// shape and elemental types.
 /// Requires: Op implements InferShapedTypeOpInterface and InferTypeOpInterface.
 ///   Less strict is possible (e.g., implements inferReturnTypeComponents and
 ///   these always populates all element types and shapes or fails, but this\
 ///   trait is currently only used where the interfaces are, so keep it
 ///   restricted for now).
 template <typename ConcreteType>
-class InferTensorType : public TraitBase<ConcreteType, InferTensorType> {
-public:
-  static LogicalResult
-  inferReturnTypes(MLIRContext *context, std::optional<Location> location,
-                   ValueRange operands, DictionaryAttr attributes,
-                   OpaqueProperties properties, RegionRange regions,
-                   SmallVectorImpl<Type> &inferredReturnTypes) {
-    static_assert(
-        ConcreteType::template hasTrait<InferShapedTypeOpInterface::Trait>(),
-        "requires InferShapedTypeOpInterface to ensure succesful invocation");
-    static_assert(
-        ConcreteType::template hasTrait<InferTypeOpInterface::Trait>(),
-        "requires InferTypeOpInterface to ensure succesful invocation");
-    return ::mlir::detail::inferReturnTensorTypes(
-        ConcreteType::inferReturnTypeComponents, context, location, operands,
-        attributes, properties, regions, inferredReturnTypes);
-  }
-};
+class InferTensorType : public TraitBase<ConcreteType, InferTensorType> {};
 
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

InterfaceMethod<
           "::mlir::ValueRange":$operands,
           "::llvm::SmallVectorImpl<::mlir::Value> &":$reifiedReturnShapes),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >
   ];
 }
 
+// Convenient trait to define a wrapper to inferReturnTypes that passes in the
+// Op Adaptor directly
+def InferTypeOpInterfaceAdaptor : TraitList<
+  [
+    // Op implements infer type op interface.
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    NativeOpTrait<
+      /*name=*/"InferTypeOpInterfaceAdaptor",
+      /*traits=*/[],
+      /*extraOpDeclaration=*/[{
+        static LogicalResult
+        inferReturnTypesAdaptor(MLIRContext *context,
+                                std::optional<Location> location,
+                                Adaptor adaptor,
+                                SmallVectorImpl<Type> &inferredReturnTypes);
+      }],
+      /*extraOpDefinition=*/[{
+        LogicalResult
+        $cppClass::inferReturnTypes(MLIRContext *context,
+                          std::optional<Location> location,
+                          ValueRange operands, DictionaryAttr attributes,
+                          OpaqueProperties properties, RegionRange regions,
+                          SmallVectorImpl<Type> &inferredReturnTypes) {
+          $cppClass::Adaptor adaptor(operands, attributes, properties, regions);
+          return $cppClass::inferReturnTypesAdaptor(context,
+            location, adaptor, inferredReturnTypes);
+        }
+      }]
+    >
+  ]>;
+
 // Convenience class grouping together type and shaped type op interfaces for
 // ops that have tensor return types.
 class InferTensorTypeBase<list<string> overridenMethods = []> : TraitList<
   [
     // Op implements infer type op interface.
-    InferTypeOpInterface,
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
     // The op will have methods implementing the ShapedType type inference
     // interface.
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
     // The op produces tensors and will use the ShapedType type infer interface
     // along with knowledge that it is producing Tensors to infer the type.
-    NativeOpTrait<"InferTensorType">
+    NativeOpTrait<
+      /*name=*/"InferTensorType",
+      /*traits=*/[],
+      /*extraOpDeclaration=*/[{}],
+      /*extraOpDefinition=*/[{
+        LogicalResult
+        $cppClass::inferReturnTypes(MLIRContext *context,
+                          std::optional<Location> location,
+                          ValueRange operands, DictionaryAttr attributes,
+                          OpaqueProperties properties, RegionRange regions,
+                          SmallVectorImpl<Type> &inferredReturnTypes) {
+          SmallVector<ShapedTypeComponents, 2> retComponents;
+          if (failed($cppClass::inferReturnTypeComponents(context, location,
+                                    operands, attributes, properties, regions,
+                                    retComponents)))
+            return failure();
+          return ::mlir::detail::inferReturnTensorTypes(retComponents,
+                                    inferredReturnTypes);
+        }
+      }]
+    >
   ]>;
 
 def InferTensorType : InferTensorTypeBase<["inferReturnTypeComponents"]>;
 def InferTensorTypeWithReify: InferTensorTypeBase<[
     "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
 def ReifyRankedShapedTypeOpInterface :
     OpInterface<"ReifyRankedShapedTypeOpInterface"> {

mlir/include/mlir/TableGen/Trait.h

 class NativeTrait : public Trait {
 public:
   // Returns the trait corresponding to a C++ trait class.
   std::string getFullyQualifiedTraitName() const;
 
   // Returns if this is a structural op trait.
   bool isStructuralOpTrait() const;
 
+  // Returns extra class declaration code to be added to the concrete type
+  // when the trait is specified
+  StringRef getExtraConcreteClassDeclaration() const;
+
+  // Returns extra class definition code to be added to the concrete type
+  // when the trait is specified
+  StringRef getExtraConcreteClassDefinition() const;
+
   static bool classof(const Trait *t) { return t->getKind() == Kind::Native; }
 };
 
 // Trait corresponding to a predicate on the operation.
 class PredTrait : public Trait {
 public:
   // Returns the template for constructing the predicate.
   std::string getPredTemplate() const;

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // ExtractStridedMetadataOp
 //===----------------------------------------------------------------------===//
 
 /// The number and type of the results are inferred from the
 /// shape of the source.
-LogicalResult ExtractStridedMetadataOp::inferReturnTypes(
-    MLIRContext *context, std::optional<Location> location, ValueRange operands,
-    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+LogicalResult ExtractStridedMetadataOp::inferReturnTypesAdaptor(
+    MLIRContext *context, std::optional<Location> location,
+    ExtractStridedMetadataOp::Adaptor adaptor,
     SmallVectorImpl<Type> &inferredReturnTypes) {
-  ExtractStridedMetadataOpAdaptor extractAdaptor(operands, attributes,
-                                                 properties);
   auto sourceType =
-      llvm::dyn_cast<MemRefType>(extractAdaptor.getSource().getType());
+      llvm::dyn_cast<MemRefType>(adaptor.getSource().getType());
   if (!sourceType)
     return failure();
 
   unsigned sourceRank = sourceType.getRank();
   IndexType indexType = IndexType::get(context);
   auto memrefType =
       MemRefType::get({}, sourceType.getElementType(),
                       MemRefLayoutAttrInterface{}, sourceType.getMemorySpace());

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 
 ShapeAdaptor ValueShapeRange::getShape(int index) const {
   if (index < 0 || static_cast<size_t>(index) >= size())
     return nullptr;
   return getShape(operator[](index));
 }
 
 LogicalResult mlir::detail::inferReturnTensorTypes(
-    function_ref<
-        LogicalResult(MLIRContext *, std::optional<Location> location,
-                      ValueShapeRange operands, DictionaryAttr attributes,
-                      OpaqueProperties properties, RegionRange regions,
-                      SmallVectorImpl<ShapedTypeComponents> &retComponents)>
-        componentTypeFn,
-    MLIRContext *context, std::optional<Location> location, ValueRange operands,
-    DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
+    SmallVector<ShapedTypeComponents, 2> retComponents,
     SmallVectorImpl<Type> &inferredReturnTypes) {
-  SmallVector<ShapedTypeComponents, 2> retComponents;
-  if (failed(componentTypeFn(context, location, operands, attributes,
-                             properties, regions, retComponents)))
-    return failure();
   for (const auto &shapeAndType : retComponents) {
     Type elementTy = shapeAndType.getElementType();
     assert(elementTy && "element type required to construct tensor");
 
     Attribute attr = shapeAndType.getAttribute();
     if (shapeAndType.hasRank()) {
       inferredReturnTypes.push_back(
           RankedTensorType::get(shapeAndType.getDims(), elementTy, attr));

mlir/lib/TableGen/Trait.cpp

std::string NativeTrait::getFullyQualifiedTraitName() const {
   return cppNamespace.empty() ? trait.str()
                               : (cppNamespace + "::" + trait).str();
 }
 
 bool NativeTrait::isStructuralOpTrait() const {
   return def->isSubClassOf("StructuralOpTrait");
 }
 
+StringRef NativeTrait::getExtraConcreteClassDeclaration() const {
+  return def->getValueAsString("extraConcreteClassDeclaration");
+}
+
+StringRef NativeTrait::getExtraConcreteClassDefinition() const {
+  return def->getValueAsString("extraConcreteClassDefinition");
+}
+
 //===----------------------------------------------------------------------===//
 // InternalTrait
 //===----------------------------------------------------------------------===//
 
 llvm::StringRef InternalTrait::getFullyQualifiedTraitName() const {
   return def->getValueAsString("trait");
 }
 

mlir/tools/mlir-tblgen/AttrOrTypeDefGen.cpp

for (auto &trait : def.getTraits()) {
     defParent.addTemplateParam(
         isa<NativeTrait>(&trait)
             ? cast<NativeTrait>(&trait)->getFullyQualifiedTraitName()
             : cast<InterfaceTrait>(&trait)->getFullyQualifiedTraitName());
   }
   defCls.addParent(std::move(defParent));
 }
 
+/// Include declarations specified on NativeTrait
+static std::string formatExtraDeclarations(const AttrOrTypeDef &def) {
+  std::string extraDeclarations = "";
+  // Include extra class declarations from NativeTrait
+  for (const auto &trait : def.getTraits()) {
+    if (auto *attrOrTypeTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+      if (auto *opTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+        StringRef value = opTrait->getExtraConcreteClassDeclaration();
+        if (value.empty()) continue;
+        extraDeclarations += value.str() + "\n";
+      }
+    }
+  }
+  if (std::optional<StringRef> extraDecl = def.getExtraDecls()) {
+    extraDeclarations += extraDecl->str();
+  }
+  return extraDeclarations;
+}
+
 /// Extra class definitions have a `$cppClass` substitution that is to be
 /// replaced by the C++ class name.
 static std::string formatExtraDefinitions(const AttrOrTypeDef &def) {
+  std::string extraDefinitions = "";
+  // Include extra class definitions from NativeTrait
+  for (const auto &trait : def.getTraits()) {
+    if (auto *attrOrTypeTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+      if (auto *opTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+        StringRef value = opTrait->getExtraConcreteClassDefinition();
+        if (value.empty()) continue;
+        extraDefinitions += value.str() + "\n";
+      }
+    }
+  }
   if (std::optional<StringRef> extraDef = def.getExtraDefs()) {
-    FmtContext ctx = FmtContext().addSubst("cppClass", def.getCppClassName());
-    return tgfmt(*extraDef, &ctx).str();
+    extraDefinitions += extraDef->str();
   }
-  return "";
+  FmtContext ctx = FmtContext().addSubst("cppClass", def.getCppClassName());
+  return tgfmt(extraDefinitions, &ctx).str();
 }
 
 void DefGen::emitTopLevelDeclarations() {
   // Inherit constructors from the attribute or type class.
   defCls.declare<VisibilityDeclaration>(Visibility::Public);
   defCls.declare<UsingDeclaration>("Base::Base");
 
   // Emit the extra declarations first in case there's a definition in there.
-  std::optional<StringRef> extraDecl = def.getExtraDecls();
+  std::string extraDecl = formatExtraDeclarations(def);
   std::string extraDef = formatExtraDefinitions(def);
-  defCls.declare<ExtraClassDeclaration>(extraDecl ? *extraDecl : "",
+  defCls.declare<ExtraClassDeclaration>(std::move(extraDecl),
                                         std::move(extraDef));
 }
 
 void DefGen::emitBuilders() {
   if (!def.skipDefaultBuilders()) {
     emitDefaultBuilder();
     if (def.genVerifyDecl())
       emitCheckedBuilder();

mlir/tools/mlir-tblgen/OpClass.h

 public:
   /// Create an operation class with extra class declarations, whose default
   /// visibility is public. Also declares at the top of the class:
   ///
   /// - inheritance of constructors from `Op`
   /// - inheritance of `print`
   /// - a type alias for the associated adaptor class
   ///
-  OpClass(StringRef name, StringRef extraClassDeclaration,
+  OpClass(StringRef name, std::string extraClassDeclaration,
           std::string extraClassDefinition);
 
   /// Add an op trait.
   void addTrait(Twine trait) { parent.addTemplateParam(trait.str()); }
 
   /// The operation class is finalized by calling `Class::finalize` to delcare
   /// all pending private and public methods (ops don't have custom constructors
   /// or fields). Then, the extra class declarations are appended to the end of
   /// the class declaration.
   void finalize() override;
 
 private:
   /// Hand-written extra class declarations.
-  StringRef extraClassDeclaration;
+  std::string extraClassDeclaration;
   /// Hand-written extra class definitions.
   std::string extraClassDefinition;
   /// The parent class, which also contains the traits to be inherited.
   ParentClass &parent;
 };
 
 } // namespace tblgen
 } // namespace mlir
 
 #endif // MLIR_TOOLS_MLIRTBLGEN_OPCLASS_H_

mlir/tools/mlir-tblgen/OpClass.cpp

 
 using namespace mlir;
 using namespace mlir::tblgen;
 
 //===----------------------------------------------------------------------===//
 // OpClass definitions
 //===----------------------------------------------------------------------===//
 
-OpClass::OpClass(StringRef name, StringRef extraClassDeclaration,
+OpClass::OpClass(StringRef name, std::string extraClassDeclaration,
                  std::string extraClassDefinition)
-    : Class(name.str()), extraClassDeclaration(extraClassDeclaration),
+    : Class(name.str()),
+      extraClassDeclaration(std::move(extraClassDeclaration)),
       extraClassDefinition(std::move(extraClassDefinition)),
       parent(addParent("::mlir::Op")) {
   parent.addTemplateParam(getClassName().str());
   declare<VisibilityDeclaration>(Visibility::Public);
   /// Inherit functions from Op.
   declare<UsingDeclaration>("Op::Op");
   declare<UsingDeclaration>("Op::print");
   /// Type alias for the adaptor class.
   declare<UsingDeclaration>("Adaptor", className + "Adaptor");
   declare<UsingDeclaration>("GenericAdaptor",
                             className + "GenericAdaptor<RangeT>")
       ->addTemplateParam("RangeT");
   declare<UsingDeclaration>(
       "FoldAdaptor", "GenericAdaptor<::llvm::ArrayRef<::mlir::Attribute>>");
 }
 
 void OpClass::finalize() {
   Class::finalize();
   declare<VisibilityDeclaration>(Visibility::Public);
-  declare<ExtraClassDeclaration>(extraClassDeclaration.str(),
-                                 extraClassDefinition);
+  declare<ExtraClassDeclaration>(extraClassDeclaration, extraClassDefinition);
 }

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

while (true) {{
   genNativeTraitAttrVerifier(body, emitHelper);
 
   bool isEmittingForOp = emitHelper.isEmittingForOp();
   for (const auto &namedAttr : emitHelper.getOp().getAttributes())
     if (canEmitAttrVerifier(namedAttr.attr, isEmittingForOp))
       emitVerifier(namedAttr.attr, namedAttr.name, getVarName(namedAttr.name));
 }
 
+/// Include declarations specified on NativeTrait
+static std::string formatExtraDeclarations(const Operator &op) {
+  std::string extraDeclarations = "";
+  // Include extra class declarations from NativeTrait
+  for (const auto &trait : op.getTraits()) {
+    if (auto *opTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+      StringRef value = opTrait->getExtraConcreteClassDeclaration();
+      if (value.empty()) continue;
+      extraDeclarations += value.str() + "\n";
+    }
+  }
+  extraDeclarations += op.getExtraClassDeclaration().str();
+  return extraDeclarations;
+}
+
 /// Op extra class definitions have a `$cppClass` substitution that is to be
 /// replaced by the C++ class name.
+/// Include declarations specified on NativeTrait
 static std::string formatExtraDefinitions(const Operator &op) {
+  std::string extraDefinitions = "";
+  // Include extra class definitions from NativeTrait
+  for (const auto &trait : op.getTraits()) {
+    if (auto *opTrait = dyn_cast<tblgen::NativeTrait>(&trait)) {
+      StringRef value = opTrait->getExtraConcreteClassDefinition();
+      if (value.empty()) continue;
+      extraDefinitions += value.str() + "\n";
+    }
+  }
+  extraDefinitions += op.getExtraClassDefinition().str();
   FmtContext ctx = FmtContext().addSubst("cppClass", op.getCppClassName());
-  return tgfmt(op.getExtraClassDefinition(), &ctx).str();
+  return tgfmt(extraDefinitions, &ctx).str();
 }
 
 OpEmitter::OpEmitter(const Operator &op,
                      const StaticVerifierFunctionEmitter &staticVerifierEmitter)
     : def(op.getDef()), op(op),
-      opClass(op.getCppClassName(), op.getExtraClassDeclaration(),
+      opClass(op.getCppClassName(), formatExtraDeclarations(op),
               formatExtraDefinitions(op)),
       staticVerifierEmitter(staticVerifierEmitter),
       emitHelper(op, /*emitForOp=*/true) {
   verifyCtx.addSubst("_op", "(*this->getOperation())");
   verifyCtx.addSubst("_ctxt", "this->getOperation()->getContext()");
 
   genTraits();