[mlir] Allow for using interface class name in ODS interface definitions

[ftynse]

Summary

It may be necessary for interface methods to process or return variables with
the interface class type, in particular for attribute and type interfaces that
can return modified attributes and types that implement the same interface.
However, the code generated by ODS in this case would not compile because the
signature (and the body if provided) appear in the definition of the Model
class and before the interface class, which derives from the Model. Change the ODS
interface method generator to emit only method declarations in the Model class
itself, and emit method definitions after the interface class. Mark as "inline"
since their definitions are still emitted in the header and are no longer
implicitly inline. Add a forward declaration of the interface class before the
Concept+Model classes to make the class name usable in declarations.

[rriddle]

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

361

This deserves a comment.

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

InterfaceMethod<"Prints the type name.",
       "void", "printTypeB", (ins "Location":$loc),
       [{}], /*defaultImplementation=*/[{
         emitRemark(loc) << $_type << " - TestB";
       }]
     >,
     InterfaceMethod<"Prints the type name.",
       "void", "printTypeC", (ins "Location":$loc)
     >,
+    // It should be possible to use the interface type name as result type
+    // as well as in the implementation.
+    InterfaceMethod<"Prints the type name and returns the type as interface.",
+      "TestTypeInterface", "printTypeRet", (ins "Location":$loc),
+      [{}], /*defaultImplementation=*/[{
+        emitRemark(loc) << $_type << " - TestRet";
+        return $_type;
+      }]
+    >,
   ];
   let extraClassDeclaration = [{
     /// Prints the type name.
     void printTypeD(Location loc) const {
       emitRemark(loc) << *this << " - TestD";
     }
   }];
   let extraTraitClassDeclaration = [{
     /// Prints the type name.
     void printTypeE(Location loc) const {
       emitRemark(loc) << $_type << " - TestE";
     }
   }];
 }
 
 #endif // TEST_INTERFACES

mlir/test/lib/IR/TestInterfaces.cpp

struct TestTypeInterfaces
   void runOnOperation() override {
     getOperation().walk([](Operation *op) {
       for (Type type : op->getResultTypes()) {
         if (auto testInterface = type.dyn_cast<TestTypeInterface>()) {
           testInterface.printTypeA(op->getLoc());
           testInterface.printTypeB(op->getLoc());
           testInterface.printTypeC(op->getLoc());
           testInterface.printTypeD(op->getLoc());
+          // Just check that we can assign the result to a variable of interface
+          // type.
+          TestTypeInterface result = testInterface.printTypeRet(op->getLoc());
+          (void)result;
         }
         if (auto testType = type.dyn_cast<TestType>())
           testType.printTypeE(op->getLoc());
       }
     });
   }
 };
 } // end anonymous namespace

mlir/test/mlir-tblgen/interfaces.mlir

 // RUN: mlir-opt -test-type-interfaces -allow-unregistered-dialect -verify-diagnostics %s
 
 // expected-remark@below {{'!test.test_type' - TestA}}
 // expected-remark@below {{'!test.test_type' - TestB}}
 // expected-remark@below {{'!test.test_type' - TestC}}
 // expected-remark@below {{'!test.test_type' - TestD}}
+// expected-remark@below {{'!test.test_type' - TestRet}}
 // expected-remark@below {{'!test.test_type' - TestE}}
 %foo0 = "foo.test"() : () -> (!test.test_type)
 
 // Type without the test interface.
 %foo1 = "foo.test"() : () -> (i32)

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

 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 ::mlir::OpInterface<TestOpInterface, detail::TestOpInterfaceInterfaceTraits>
+
 // DECL: int foo(int input);
 
-// DECL-NOT: TestOpInterface
+// DECL: template<typename ConcreteOp>
+// DECL: int detail::TestOpInterfaceInterfaceTraits::Model<ConcreteOp>::foo
 
 // 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/OpInterfacesGen.cpp

public:
   bool emitInterfaceDocs();
 
 protected:
   InterfaceGenerator(std::vector<llvm::Record *> &&defs, raw_ostream &os)
       : defs(std::move(defs)), os(os) {}
 
   void emitConceptDecl(Interface &interface);
   void emitModelDecl(Interface &interface);
+  void emitModelMethodsDef(Interface &interface);
   void emitTraitDecl(Interface &interface, StringRef interfaceName,
                      StringRef interfaceTraitsName);
   void emitInterfaceDecl(Interface interface);
 
   /// The set of interface records to emit.
   std::vector<llvm::Record *> defs;
   // The stream to emit to.
   raw_ostream &os;
void InterfaceGenerator::emitModelDecl(Interface &interface) {
   os << "    Model() : Concept{";
   llvm::interleaveComma(
       interface.getMethods(), os,
       [&](const InterfaceMethod &method) { os << method.getName(); });
   os << "} {}\n\n";
 
   // Insert each of the virtual method overrides.
   for (auto &method : interface.getMethods()) {
-    emitCPPType(method.getReturnType(), os << "    static ");
+    emitCPPType(method.getReturnType(), os << "    static inline ");
+    emitMethodNameAndArgs(method, os, valueType,
+                          /*addThisArg=*/!method.isStatic(),
+                          /*addConst=*/false);
+    os << ";\n";
+  }
+  os << "  };\n";
+}
+
+void InterfaceGenerator::emitModelMethodsDef(Interface &interface) {
+  for (auto &method : interface.getMethods()) {
+    os << "template<typename " << valueTemplate << ">\n";
+    emitCPPType(method.getReturnType(), os);
+    os << "detail::" << interface.getName() << "InterfaceTraits::Model<"
+       << valueTemplate << ">::";
     emitMethodNameAndArgs(method, os, valueType,
                           /*addThisArg=*/!method.isStatic(),
                           /*addConst=*/false);
-    os << " {\n      ";
+    os << " {\n  ";
 
     // Check for a provided body to the function.
     if (Optional<StringRef> body = method.getBody()) {
       if (method.isStatic())
         os << body->trim();
       else
         os << tblgen::tgfmt(body->trim(), &nonStaticMethodFmt);
-      os << "\n    }\n";
+      os << "\n}\n";
       continue;
     }
 
     // Forward to the method on the concrete operation type.
     if (method.isStatic())
       os << "return " << valueTemplate << "::";
     else
       os << tblgen::tgfmt("return $_self.", &nonStaticMethodFmt);
 
     // Add the arguments to the call.
     os << method.getName() << '(';
     llvm::interleaveComma(
         method.getArguments(), os,
         [&](const InterfaceMethod::Argument &arg) { os << arg.name; });
-    os << ");\n    }\n";
+    os << ");\n}\n";
   }
-  os << "  };\n";
 }
 
 void InterfaceGenerator::emitTraitDecl(Interface &interface,
                                        StringRef interfaceName,
                                        StringRef interfaceTraitsName) {
   os << llvm::formatv("  template <typename {3}>\n"
                       "  struct {0}Trait : public ::mlir::{2}<{0},"
                       " detail::{1}>::Trait<{3}> {{\n",
void InterfaceGenerator::emitInterfaceDecl(Interface interface) {
   llvm::SmallVector<StringRef, 2> namespaces;
   llvm::SplitString(interface.getCppNamespace(), namespaces, "::");
   for (StringRef ns : namespaces)
     os << "namespace " << ns << " {\n";
 
   StringRef interfaceName = interface.getName();
   auto interfaceTraitsName = (interfaceName + "InterfaceTraits").str();
 
+  // Emit a forward declaration of the interface class so that it becomes usable
+  // in the signature of its methods.
+  os << "class " << interfaceName << ";\n";
+
   // Emit the traits struct containing the concept and model declarations.
   os << "namespace detail {\n"
      << "struct " << interfaceTraitsName << " {\n";
   emitConceptDecl(interface);
   emitModelDecl(interface);
   os << "};\n} // end namespace detail\n";
 
   // Emit the main interface class declaration.
void InterfaceGenerator::emitInterfaceDecl(Interface interface) {
   }
 
   // Emit any extra declarations.
   if (Optional<StringRef> extraDecls = interface.getExtraClassDeclaration())
     os << *extraDecls << "\n";
 
   os << "};\n";
 
+  emitModelMethodsDef(interface);
+
   for (StringRef ns : llvm::reverse(namespaces))
     os << "} // namespace " << ns << "\n";
 }
 
 bool InterfaceGenerator::emitInterfaceDecls() {
   llvm::emitSourceFileHeader("Interface Declarations", os);
 
   for (const auto *def : defs)