[mlir][ods] ODS ops get an `extraClassDefinition`

[Mogball]

Summary

Extra definitions are placed in the generated source file for each op class. The substitution $cppClass is replaced by the op's C++ class name.

This is useful when declaring but not defining methods in TableGen base classes:

class BaseOp<string mnemonic> 
    : Op<MyDialect, mnemonic, [DeclareOpInterfaceMethods<SomeInterface>] {
  let extraClassDeclaration = [{ 
    // ZOp is declared at at the bottom of the file and is incomplete here
    ZOp getParent();
  }];
  let extraClassDefinition = [{
    int $cppClass::someInterfaceMethod() {
      return someUtilityFunction(*this);
    }
    ZOp $cppClass::getParent() {
      return dyn_cast<ZOp>(this->getParentOp());
    }
  }];
}

Certain things may prevent defining these functions inline, in the declaration. In this example, ZOp in the same dialect is incomplete at the function declaration because ops classes are declared in alphabetical order. Alternatively, functions may be too big to be desired as inlined, or they may require dependencies that create cyclic includes, or they may be calling a templated utility function that one may not want to expose in a header. If the functions are not inlined, then inheriting from the base class N times means that each function will need to be defined N times. With extraClassDefinitions, they only need to be defined once.

[rriddle]

What types of use cases is this intended for? E.g. what examples do you have that would use this?

[Mogball]

The use cases I have in mind are both related to declaring but not defining methods in TableGen base classes.

class BaseOp<string mnemonic> : Op<MyDialect, mnemonic, [DeclareInterfaceMethods<SomeInterface>] {
  let extraClassDeclaration = [{ 
    // ZOp is declared at at the bottom of the file and is incomplete here
    ZOp getParent();
  }];
}

If I inherit from this class N times, I have to define N identical getParent methods and interface methods, and usually I have a templated function hiding in the source file I don't want to put in the header.

[mehdi_amini]

Can you capture such motivation in the commit message?

[jpienaar]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

563

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp:2125, isn't there already the convention of using cppClass ?

[Mogball]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

563

Oh, yeah. I didn't realize there wasn't an underscore.

[Mogball]

Rename the substitution _cppClass -> cppClass and update description

[Mogball]

@rriddle yay or nay?

[rriddle]

Is there a place in OpDefinitions.md to fit this?

The justification for base class stuff makes sense to me. I think we should discourage this for most other situations though.

mlir/include/mlir/IR/OpBase.td

2377–2378

Might want to add clarification here about where the generated code is placed w.r.t. the ops namespace.

[Mogball]

Update docs

mlir/include/mlir/IR/OpBase.td

class Op<Dialect dialect, string mnemonic, list<OpTrait> props = []> {
 
   // Op traits.
   // Note: The list of traits will be uniqued by ODS.
   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 = ?;
+
+  // Additional code that will be added to the generated source file.
+  // `$_cppClass` is replaced by the op's C++ class name.
+  code extraClassDefinition = ?;
 }
 
 // Base class for ops with static/dynamic offset, sizes and strides
 // attributes/arguments.
 class BaseOpWithOffsetSizesAndStrides<Dialect dialect, string mnemonic,
                                       list<OpTrait> traits = []> :
   Op<dialect, mnemonic, traits> {
 

mlir/include/mlir/TableGen/Class.h

public:
   /// Write the visibility declaration.
   void writeDeclTo(raw_indented_ostream &os) const override;
 
 private:
   /// The visibility of subsequent class declarations.
   Visibility visibility;
 };
 
-/// Unstructured extra class declarations, from TableGen definitions. The
-/// default visibility of extra class declarations is up to the owning class.
+/// Unstructured extra class declarations and definitions, from TableGen
+/// definitions. The default visibility of extra class declarations is up to the
+/// owning class.
 class ExtraClassDeclaration
     : public ClassDeclarationBase<ClassDeclaration::ExtraClassDeclaration> {
 public:
   /// Create an extra class declaration.
-  ExtraClassDeclaration(StringRef extraClassDeclaration)
-      : extraClassDeclaration(extraClassDeclaration) {}
+  ExtraClassDeclaration(StringRef extraClassDeclaration,
+                        StringRef extraClassDefinition = "")
+      : extraClassDeclaration(extraClassDeclaration),
+        extraClassDefinition(extraClassDefinition) {}
 
   /// Write the extra class declarations.
   void writeDeclTo(raw_indented_ostream &os) const override;
 
+  /// Write the extra class definitions.
+  void writeDefTo(raw_indented_ostream &os,
+                  StringRef namePrefix) const override;
+
 private:
   /// The string of the extra class declarations. It is re-indented before
   /// printed.
   StringRef extraClassDeclaration;
+  /// The string of the extra class definitions. It is re-indented before
+  /// printed.
+  StringRef extraClassDefinition;
 };
 
 /// A class used to emit C++ classes from Tablegen.  Contains a list of public
 /// methods and a list of private fields to be emitted.
 class Class {
 public:
   virtual ~Class() = default;
 

mlir/include/mlir/TableGen/Operator.h

public:
 
   // Query functions for the assembly format of the operator.
   bool hasAssemblyFormat() const;
   StringRef getAssemblyFormat() const;
 
   // Returns this op's extra class declaration code.
   StringRef getExtraClassDeclaration() const;
 
+  // Returns this op's extra class definition code.
+  StringRef getExtraClassDefinition() const;
+
   // Returns the Tablegen definition this operator was constructed from.
   // TODO: do not expose the TableGen record, this is a temporary solution to
   // OpEmitter requiring a Record because Operator does not provide enough
   // methods.
   const llvm::Record &getDef() const;
 
   // Returns the dialect of the op.
   const Dialect &getDialect() const { return dialect; }

mlir/lib/TableGen/Class.cpp

 //===----------------------------------------------------------------------===//
 // ExtraClassDeclaration definitions
 //===----------------------------------------------------------------------===//
 
 void ExtraClassDeclaration::writeDeclTo(raw_indented_ostream &os) const {
   os.printReindented(extraClassDeclaration);
 }
 
+void ExtraClassDeclaration::writeDefTo(raw_indented_ostream &os,
+                                       StringRef namePrefix) const {
+  os.printReindented(extraClassDefinition);
+}
+
 //===----------------------------------------------------------------------===//
 // Class definitions
 //===----------------------------------------------------------------------===//
 
 ParentClass &Class::addParent(ParentClass parent) {
   parents.push_back(std::move(parent));
   return parents.back();
 }

mlir/lib/TableGen/Operator.cpp

 
 StringRef Operator::getExtraClassDeclaration() const {
   constexpr auto attr = "extraClassDeclaration";
   if (def.isValueUnset(attr))
     return {};
   return def.getValueAsString(attr);
 }
 
+StringRef Operator::getExtraClassDefinition() const {
+  constexpr auto attr = "extraClassDefinition";
+  if (def.isValueUnset(attr))
+    return {};
+  return def.getValueAsString(attr);
+}
+
 const llvm::Record &Operator::getDef() const { return def; }
 
 bool Operator::skipDefaultBuilders() const {
   return def.getValueAsBit("skipDefaultBuilders");
 }
 
 auto Operator::result_begin() -> value_iterator { return results.begin(); }
 

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

 
 def ConversionCallOp : TEST_Op<"conversion_call_op",
     [CallOpInterface]> {
   let arguments = (ins Variadic<AnyType>:$arg_operands, SymbolRefAttr:$callee);
   let results = (outs Variadic<AnyType>);
 
   let extraClassDeclaration = [{
     /// Return the callee of this operation.
-    ::mlir::CallInterfaceCallable getCallableForCallee() {
+    ::mlir::CallInterfaceCallable getCallableForCallee();
+  }];
+  let extraClassDefinition = [{
+    ::mlir::CallInterfaceCallable $_cppClass::getCallableForCallee() {
       return (*this)->getAttrOfType<::mlir::SymbolRefAttr>("callee");
     }
   }];
 }
 
 def FunctionalRegionOp : TEST_Op<"functional_region_op",
     [CallableOpInterface]> {
   let regions = (region AnyRegion:$body);

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, StringRef 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;
+  /// 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, StringRef extraClassDeclaration,
+                 std::string extraClassDefinition)
     : Class(name.str()), extraClassDeclaration(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");
 }
 
 void OpClass::finalize() {
   Class::finalize();
   declare<VisibilityDeclaration>(Visibility::Public);
-  declare<ExtraClassDeclaration>(extraClassDeclaration);
+  declare<ExtraClassDeclaration>(extraClassDeclaration, extraClassDefinition);
 }

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

if (canEmitCondition) {
                         emitHelper.emitErrorPrefix(), attrName,
                         escapeString(attr.getSummary()));
       }
     }
     body << "  }\n";
   }
 }
 
+/// Op extra class definitions have a `$_cppClass` substitution that is to be
+/// replaced by the C++ class name.
+static std::string formatExtraDefinitions(const Operator &op) {
+  FmtContext ctx = FmtContext().addSubst("_cppClass", op.getCppClassName());
+  return tgfmt(op.getExtraClassDefinition(), &ctx).str();
+}
+
 OpEmitter::OpEmitter(const Operator &op,
                      const StaticVerifierFunctionEmitter &staticVerifierEmitter)
     : def(op.getDef()), op(op),
-      opClass(op.getCppClassName(), op.getExtraClassDeclaration()),
+      opClass(op.getCppClassName(), op.getExtraClassDeclaration(),
+              formatExtraDefinitions(op)),
       staticVerifierEmitter(staticVerifierEmitter) {
   verifyCtx.withOp("(*this->getOperation())");
   verifyCtx.addSubst("_ctxt", "this->getOperation()->getContext()");
 
   genTraits();
 
   // Generate C++ code for various op methods. The order here determines the
   // methods in the generated file.