[MLIR] Generate inferReturnTypes declaration using InferTypeOpInterface trait.

[jurahul]

Summary

• Instead of hardcoding the parameters and return types of 'inferReturnTypes', use the InferTypeOpInterface trait to generate the method declaration.
• Fix InferTypeOfInterface to use fully qualified type for inferReturnTypes results.

[jurahul]

Fix clang-tidy warning

[jpienaar]

Nice, thanks!

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

296

methods?

1746

How about hoisting the loop into a lambda function? I think it could read easier (just a suggestion)

1753

Nit: && "unable to find inferReturnTypes interface method" ?

[jurahul]

Address feedback

[jurahul]

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1746

Is this something you had in mind?

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

let methods = [
     StaticInterfaceMethod<
       /*desc=*/[{Infer the return types that an op would generate.
 
       The method takes an optional location which, if set, will be used to
       report errors on. The operands and attributes correspond to those with
       which an Operation would be created (e.g., as used in Operation::create)
       and the regions of the op.
       }],
-      /*retTy=*/"LogicalResult",
+      /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypes",
       /*args=*/(ins "::mlir::MLIRContext *":$context,
                     "::llvm::Optional<::mlir::Location>":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::llvm::SmallVectorImpl<::mlir::Type>&":$inferredReturnTypes)
     >,

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

private:
 
   // Generates verify statements for successors in the operation.
   // The generated code will be attached to `body`.
   void genSuccessorVerifier(OpMethodBody &body);
 
   // Generates the traits used by the object.
   void genTraits();
 
-  // Generate the OpInterface methods.
+  // Generate the OpInterface methods for all interfaces.
   void genOpInterfaceMethods();
 
-  // Generate op interface method.
-  void genOpInterfaceMethod(const tblgen::InterfaceOpTrait *trait);
+  // Generate op interface method for the given interface.
+  void genOpInterfaceMethods(const tblgen::InterfaceOpTrait *trait);
+
+  // Generate op interface method for the given interface method. If
+  // 'declaration' is true, generates a declaration, else a definition.
+  OpMethod *genOpInterfaceMethod(const tblgen::InterfaceMethod &method,
+                                 bool declaration = true);
 
   // Generate the side effect interface methods.
   void genSideEffectInterfaceMethods();
 
   // Generate the type inference interface methods.
   void genTypeInterfaceMethods();
 
 private:
if (hasSingleResult) {
       paramList.emplace_back("::llvm::SmallVectorImpl<::mlir::OpFoldResult> &",
                              "results");
       opClass.addMethodAndPrune("::mlir::LogicalResult", "fold",
                                 OpMethod::MP_Declaration, std::move(paramList));
     }
   }
 }
 
-void OpEmitter::genOpInterfaceMethod(const tblgen::InterfaceOpTrait *opTrait) {
+void OpEmitter::genOpInterfaceMethods(const tblgen::InterfaceOpTrait *opTrait) {
   auto interface = opTrait->getOpInterface();
 
   // 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 InterfaceMethod &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;
+    genOpInterfaceMethod(method);
+  }
+}
 
-    SmallVector<OpMethodParameter, 4> paramList;
-    for (const InterfaceMethod::Argument &arg : method.getArguments())
-      paramList.emplace_back(arg.type, arg.name);
+OpMethod *OpEmitter::genOpInterfaceMethod(const InterfaceMethod &method,
+                                          bool declaration) {
 
-    auto properties = method.isStatic() ? OpMethod::MP_StaticDeclaration
-                                        : OpMethod::MP_Declaration;
-    opClass.addMethodAndPrune(method.getReturnType(), method.getName(),
+  auto properties = method.isStatic() ? OpMethod::MP_Static : OpMethod::MP_None;
+  if (declaration)
+    properties =
+        static_cast<OpMethod::Property>(properties | OpMethod::MP_Declaration);
+  return opClass.addMethodAndPrune(method.getReturnType(), method.getName(),
                                    properties, std::move(paramList));
-  }
-}
 
 void OpEmitter::genOpInterfaceMethods() {
   for (const auto &trait : op.getTraits()) {
     if (const auto *opTrait = dyn_cast<tblgen::InterfaceOpTrait>(&trait))
       if (opTrait->shouldDeclareMethods())
-        genOpInterfaceMethod(opTrait);
+        genOpInterfaceMethods(opTrait);
   }
 }
 
 void OpEmitter::genSideEffectInterfaceMethods() {
   enum EffectKind { Operand, Result, Symbol, Static };
   struct EffectLocation {
     /// The effect applied.
     SideEffect effect;
for (auto &location : it.second) {
       }
     }
   }
 }
 
 void OpEmitter::genTypeInterfaceMethods() {
   if (!op.allResultTypesKnown())
     return;
-
-  SmallVector<OpMethodParameter, 4> paramList;
-  paramList.emplace_back("::mlir::MLIRContext *", "context");
-  paramList.emplace_back("::llvm::Optional<::mlir::Location>", "location");
-  paramList.emplace_back("::mlir::ValueRange", "operands");
-  paramList.emplace_back("::mlir::DictionaryAttr", "attributes");
-  paramList.emplace_back("::mlir::RegionRange", "regions");
-  paramList.emplace_back("::llvm::SmallVectorImpl<::mlir::Type>&",
-                         "inferredReturnTypes");
-  auto *method =
-      opClass.addMethodAndPrune("::mlir::LogicalResult", "inferReturnTypes",
-                                OpMethod::MP_Static, std::move(paramList));
+  // Generate 'inferReturnTypes' method declaration using the interface method
+  // declared in 'InferTypeOpInterface' op interface.
+  const auto *trait = dyn_cast<InterfaceOpTrait>(
+      op.getTrait("::mlir::InferTypeOpInterface::Trait"));
+  auto interface = trait->getOpInterface();
+  OpMethod *method = nullptr;
+  for (const InterfaceMethod &interfaceMethod : interface.getMethods()) {
+    if (interfaceMethod.getName() == "inferReturnTypes") {
+      method = genOpInterfaceMethod(interfaceMethod, /*declaration=*/false);
+      break;
+    }
+  }
+  assert(method != nullptr);
   auto &body = method->body();
   body << "  inferredReturnTypes.resize(" << op.getNumResults() << ");\n";
 
   FmtContext fctx;
   fctx.withBuilder("odsBuilder");
   body << "  ::mlir::Builder odsBuilder(context);\n";
 
   auto emitType =