Add extraClassDeclaration field for ops.

Simple mechanism to allow specifying arbitrary function declarations. The modelling will not cover all cases so allow a means for users to declare a method function that they will define in their C++ files. The goal is to allow full C++ flexibility as the goal is to cover cases not modelled.

--

PiperOrigin-RevId: 245889819

mlir/include/mlir/IR/OpBase.td

@@ -892,6 +892,10 @@ class Op<Dialect dialect, string mnemonic, list<OpTrait> props = []> {
 
   // Op traits.
   list<OpTrait> traits = props;
+
+  // Additional code that will be added to the public part of the generated
+  // C++ code of the op declaration.
+  code extraClassDeclaration = ?;
 }
 
 // The arguments of an op.

mlir/include/mlir/TableGen/Operator.h

@@ -59,6 +59,9 @@ public:
   // Returns this op's C++ class name.
   StringRef getCppClassName() const;
 
+  // Returns this op's extra class declaration code.
+  StringRef getExtraClassDeclaration() const;
+
   // Returns the qualified C++ class name for the given TableGen def `name`.
   // The first `_` in `name` is treated as separating the dialect namespace
   // and the op class name if the dialect namespace is not empty. Otherwise,

mlir/lib/TableGen/Operator.cpp

@@ -76,6 +76,13 @@ int tblgen::Operator::getNumResults() const {
   return results->getNumArgs();
 }
 
+StringRef tblgen::Operator::getExtraClassDeclaration() const {
+  constexpr auto attr = "extraClassDeclaration";
+  if (def.isValueUnset(attr))
+    return {};
+  return def.getValueAsString(attr);
+}
+
 tblgen::TypeConstraint
 tblgen::Operator::getResultTypeConstraint(int index) const {
   DagInit *results = def.getValueAsDag("results");

mlir/test/mlir-tblgen/op-decl.td

@@ -30,6 +30,11 @@ def NS_AOp : NS_Op<"a_op", [NoSideEffect]> {
   let hasCanonicalizer = 1;
   let hasConstantFolder = 1;
   let hasFolder = 1;
+
+  let extraClassDeclaration = [{
+    // Display a graph for debugging purposes.
+    void displayGraph();
+  }];
 }
 
 // CHECK-LABEL: NS::AOp declarations
@@ -52,6 +57,8 @@ def NS_AOp : NS_Op<"a_op", [NoSideEffect]> {
 // CHECK:   static void getCanonicalizationPatterns(OwningRewritePatternList &results, MLIRContext *context);
 // CHECK:   LogicalResult constantFold(ArrayRef<Attribute> operands, SmallVectorImpl<Attribute> &results, MLIRContext *context);
 // CHECK:   bool fold(SmallVectorImpl<Value *> &results);
+// CHECK:   // Display a graph for debugging purposes.
+// CHECK:   void displayGraph();
 // CHECK: };
 
 // Check op trait for different number of operands

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

@@ -172,7 +172,7 @@ private:
 // Class for holding an op for C++ code emission
 class OpClass {
 public:
-  explicit OpClass(StringRef name);
+  explicit OpClass(StringRef name, StringRef extraClassDeclaration = "");
 
   // Adds an op trait.
   void addTrait(Twine trait);
@@ -188,7 +188,8 @@ public:
   void writeDefTo(raw_ostream &os) const;
 
 private:
-  std::string className;
+  StringRef className;
+  StringRef extraClassDeclaration;
   SmallVector<std::string, 4> traits;
   SmallVector<OpMethod, 8> methods;
 };
@@ -286,7 +287,8 @@ void OpMethod::writeDefTo(raw_ostream &os, StringRef namePrefix) const {
   os << "}";
 }
 
-OpClass::OpClass(StringRef name) : className(name) {}
+OpClass::OpClass(StringRef name, StringRef extraClassDeclaration)
+    : className(name), extraClassDeclaration(extraClassDeclaration) {}
 
 // Adds the given trait to this op. Prefixes "OpTrait::" to `trait` implicitly.
 void OpClass::addTrait(Twine trait) {
@@ -310,6 +312,8 @@ void OpClass::writeDeclTo(raw_ostream &os) const {
     method.writeDeclTo(os);
     os << "\n";
   }
+  // TODO: Add line control markers to make errors easier to debug.
+  os << extraClassDeclaration << "\n";
   os << "};";
 }
 
@@ -388,7 +392,8 @@ private:
 } // end anonymous namespace
 
 OpEmitter::OpEmitter(const Record &def)
-    : def(def), op(def), opClass(op.getCppClassName()) {
+    : def(def), op(def),
+      opClass(op.getCppClassName(), op.getExtraClassDeclaration()) {
   genTraits();
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.