[mlir][DeclareOpInterfaceMethods] Allow specifying a set of methods to force declaration generation for.

[rriddle]

Summary

Currently a declaration won't be generated if the method has a default implementation. Meaning that operations that wan't to override the default have to explicitly declare the method in the extraClassDeclarations. This revision adds an optional list parameter to DeclareOpInterfaceMethods to allow for specifying a set of methods that should always have the declarations generated, even if there is a default.

Depends On D79029

[jpienaar]

Thanks, so we have

• no default -> generate declaration
• default -> if ("op wants to specialize this") then generate declaration -> don't generate def

I was trying to think of what is the common case (wrt white list or black list approach), but probably if there is a default then it is the common case and so this should be least typing :)

[rriddle]

Add missing include

mlir/docs/OpDefinitions.md

let methods = [
     InterfaceMethod<"/*insert doc here*/",
       "unsigned", "getNumInputsAndOutputs", (ins), [{
         return op.getNumInputs() + op.getNumOutputs();
     }]>,
 
     // Provide only a default definition of the method.
     // Note: `ConcreteOp` corresponds to the derived operation typename.
     InterfaceMethod<"/*insert doc here*/",
-      "unsigned", "getNumInputsAndOutputs", (ins), /*methodBody=*/[{}], [{
+      "unsigned", "getNumWithDefault", (ins), /*methodBody=*/[{}], [{
         ConcreteOp op = cast<ConcreteOp>(getOperation());
         return op.getNumInputs() + op.getNumOutputs();
     }]>,
   ];
 }
 
 // Interfaces can optionally be wrapped inside DeclareOpInterfaceMethods. This
 // would result in autogenerating declarations for members `foo`, `bar` and
 // `fooStatic`. Methods with bodies are not declared inside the op
 // declaration but instead handled by the op interface trait directly.
 def OpWithInferTypeInterfaceOp : Op<...
     [DeclareOpInterfaceMethods<MyInterface>]> { ... }
+
+// Methods that have a default implementation do not have declarations
+// generated. If an operation wishes to override the default behavior, it can
+// explicitly specify the method that it wishes to override. This will force
+// the generation of a declaration for those methods.
+def OpWithOverrideInferTypeInterfaceOp : Op<...
+    [DeclareOpInterfaceMethods<MyInterface, ["getNumWithDefault"]>]> { ... }
 ```
 
 A verification method can also be specified on the `OpInterface` by setting
 `verify`. Setting `verify` results in the generated trait having a `verifyTrait`
 method that is applied to all operations implementing the trait.
 
 ### Builder methods
 

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def AtomicYieldOp : Std_Op<"atomic_yield", [
   let assemblyFormat = "$result attr-dict `:` type($result)";
 }
 
 //===----------------------------------------------------------------------===//
 // BranchOp
 //===----------------------------------------------------------------------===//
 
 def BranchOp : Std_Op<"br",
-    [DeclareOpInterfaceMethods<BranchOpInterface>, NoSideEffect, Terminator]> {
+    [DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
+     NoSideEffect, Terminator]> {
   let summary = "branch operation";
   let description = [{
     The `br` operation represents a branch operation in a function.
     The operation takes variable number of operands and produces no results.
     The operand number and types for each successor must match the arguments of
     the block successor.
 
     Example:
def BranchOp : Std_Op<"br",
   let verifier = ?;
 
   let extraClassDeclaration = [{
     Block *getDest();
     void setDest(Block *block);
 
     /// Erase the operand at 'index' from the operand list.
     void eraseOperand(unsigned index);
-
-    /// Returns the successor that would be chosen with the given constant
-    /// operands. Returns nullptr if a single successor could not be chosen.
-    Block *getSuccessorForOperands(ArrayRef<Attribute>);
   }];
 
   let hasCanonicalizer = 1;
   let assemblyFormat = [{
     $dest (`(` $destOperands^ `:` type($destOperands) `)`)? attr-dict
   }];
 }
 
def CmpIOp : Std_Op<"cmpi",
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
 //===----------------------------------------------------------------------===//
 // CondBranchOp
 //===----------------------------------------------------------------------===//
 
 def CondBranchOp : Std_Op<"cond_br",
-    [AttrSizedOperandSegments, DeclareOpInterfaceMethods<BranchOpInterface>,
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>,
      NoSideEffect, Terminator]> {
   let summary = "conditional branch operation";
   let description = [{
     The `cond_br` terminator operation represents a conditional branch on a
     boolean (1-bit integer) value. If the bit is set, then the first destination
     is jumped to; if it is false, the second destination is chosen. The count
     and types of operands must align with the arguments in the corresponding
     target blocks.
let extraClassDeclaration = [{
 
     unsigned getNumFalseOperands() { return getFalseOperands().size(); }
 
     /// Erase the operand at 'index' from the false operand list.
     void eraseFalseOperand(unsigned index) {
       eraseSuccessorOperand(falseIndex, index);
     }
 
-    /// Returns the successor that would be chosen with the given constant
-    /// operands. Returns nullptr if a single successor could not be chosen.
-    Block *getSuccessorForOperands(ArrayRef<Attribute> operands);
-
   private:
     /// Get the index of the first true destination operand.
     unsigned getTrueDestOperandIndex() { return 1; }
 
     /// Get the index of the first false destination operand.
     unsigned getFalseDestOperandIndex() {
       return getTrueDestOperandIndex() + getNumTrueOperands();
     }

mlir/include/mlir/IR/OpBase.td

class OpInterface<string name> : OpInterfaceTrait<name> {
 
   // An optional code block containing extra declarations to place in the
   // interface declaration.
   code extraClassDeclaration = "";
 }
 
 // Whether to declare the op interface methods in the op's header. This class
 // simply wraps an OpInterface but is used to indicate that the method
-// declarations should be generated.
-class DeclareOpInterfaceMethods<OpInterface interface> :
-  OpInterface<interface.cppClassName> {
+// declarations should be generated. This class takes an optional set of methods
+// that should have declarations generated even if the method has a default
+// implementation.
+class DeclareOpInterfaceMethods<OpInterface interface,
+                                list<string> overridenMethods = []>
+      : OpInterface<interface.cppClassName> {
     let description = interface.description;
     let cppClassName = interface.cppClassName;
     let methods = interface.methods;
+
+    // This field contains a set of method names that should always have their
+    // declarations generated. This allows for generating declarations for
+    // methods with default implementations that need to be overridden.
+    list<string> alwaysOverriddenMethods = overridenMethods;
 }
 
 //===----------------------------------------------------------------------===//
 // Op definitions
 //===----------------------------------------------------------------------===//
 
 // Marker used to identify the result list for an op.
 def outs;

mlir/include/mlir/TableGen/OpTrait.h

public:
   StringRef getTrait() const;
 
   static bool classof(const OpTrait *t) {
     return t->getKind() == Kind::Interface;
   }
 
   // Whether the declaration of methods for this trait should be emitted.
   bool shouldDeclareMethods() const;
+
+  // Returns the methods that should always be declared if this interface is
+  // emitting declarations.
+  std::vector<StringRef> getAlwaysDeclaredMethods() const;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_OPTRAIT_H_

mlir/lib/TableGen/OpTrait.cpp

 
 llvm::StringRef InterfaceOpTrait::getTrait() const {
   return def->getValueAsString("trait");
 }
 
 bool InterfaceOpTrait::shouldDeclareMethods() const {
   return def->isSubClassOf("DeclareOpInterfaceMethods");
 }
+
+std::vector<StringRef> InterfaceOpTrait::getAlwaysDeclaredMethods() const {
+  return def->getValueAsListOfStrings("alwaysOverriddenMethods");
+}

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

 // RUN: mlir-tblgen -gen-op-interface-decls -I %S/../../include %s | FileCheck %s --check-prefix=DECL --dump-input-on-failure
+// RUN: mlir-tblgen -gen-op-decls -I %S/../../include %s | FileCheck %s --check-prefix=OP_DECL --dump-input-on-failure
 
 include "mlir/IR/OpBase.td"
 
 def TestOpInterface : OpInterface<"TestOpInterface"> {
   let description = [{some op interface description}];
 
   let methods = [
     InterfaceMethod<
       /*desc=*/[{some function comment}],
       /*retTy=*/"int",
       /*methodName=*/"foo",
       /*args=*/(ins "int":$input)
     >,
+    InterfaceMethod<
+      /*desc=*/[{some function comment}],
+      /*retTy=*/"int",
+      /*methodName=*/"default_foo",
+      /*args=*/(ins "int":$input),
+      /*body=*/[{}],
+      /*defaultBody=*/[{ return 0; }]
+    >,
   ];
 }
 
 // Define Ops with TestOpInterface and
 // DeclareOpInterfaceMethods<TestOpInterface> traits to check that there
 // are not duplicated C++ classes generated.
 def TestDialect : Dialect {
   let name = "test";
 }
 
 def OpInterfaceOp : Op<TestDialect, "op_interface_op", [TestOpInterface]>;
 
 def DeclareMethodsOp : Op<TestDialect, "declare_methods_op",
                           [DeclareOpInterfaceMethods<TestOpInterface>]>;
 
+def DeclareMethodsWithDefaultOp : Op<TestDialect, "declare_methods_op",
+      [DeclareOpInterfaceMethods<TestOpInterface, ["default_foo"]>]>;
+
 // DECL-LABEL: TestOpInterfaceInterfaceTraits
 // DECL: class TestOpInterface : public OpInterface<TestOpInterface, detail::TestOpInterfaceInterfaceTraits>
 // DECL: int foo(int input);
 
 // DECL-NOT: TestOpInterface
+
+// OP_DECL-LABEL: class DeclareMethodsOp : public
+// OP_DECL: int foo(int input);
+// OP_DECL-NOT: int default_foo(int input);
+
+// OP_DECL-LABEL: class DeclareMethodsWithDefaultOp : public
+// OP_DECL: int foo(int input);
+// OP_DECL: int default_foo(int input);

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 }
 
 void OpEmitter::genOpInterfaceMethods() {
   for (const auto &trait : op.getTraits()) {
     auto opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait);
     if (!opTrait || !opTrait->shouldDeclareMethods())
       continue;
     auto interface = opTrait->getOpInterface();
-    for (auto method : interface.getMethods()) {
-      // Don't declare if the method has a body or a default implementation.
-      if (method.getBody() || method.getDefaultImplementation())
+
+    // Get the set of methods that should always be declared.
+    auto alwaysDeclaredMethodsVec = opTrait->getAlwaysDeclaredMethods();
+    llvm::StringSet<> alwaysDeclaredMethods;
+    alwaysDeclaredMethods.insert(alwaysDeclaredMethodsVec.begin(),
+                                 alwaysDeclaredMethodsVec.end());
+
+    for (const OpInterfaceMethod &method : interface.getMethods()) {
+      // Don't declare if the method has a body.
+      if (method.getBody())
+        continue;
+      // Don't declare if the method has a default implementation and the op
+      // didn't request that it always be declared.
+      if (method.getDefaultImplementation() &&
+          !alwaysDeclaredMethods.count(method.getName()))
         continue;
+
       std::string args;
       llvm::raw_string_ostream os(args);
       interleaveComma(method.getArguments(), os,
                       [&](const OpInterfaceMethod::Argument &arg) {
                         os << arg.type << " " << arg.name;
                       });
       opClass.newMethod(method.getReturnType(), method.getName(), os.str(),
                         method.isStatic() ? OpMethod::MP_Static