Generate builder for ops that use InferTypeOpInterface trait in ODS

For ops with infer type op interface defined, generate version that calls the inferal method on build. This is intermediate step to removing special casing of SameOperandsAndResultType & FirstAttrDereivedResultType. After that would be generating the inference code, with the initial focus on shaped container types. In between I plan to refactor these a bit to reuse generated paths. The intention would not be to add the type inference trait in multiple places, but rather to take advantage of the current modelling in ODS where possible to emit it instead.

Switch the `inferReturnTypes` method to be static.

Skipping ops with regions here as I don't like the Region vs unique_ptr<Region> difference at the moment, and I want the infer return type trait to be useful for verification too. So instead, just skip it for now to avoid churn.

PiperOrigin-RevId: 284217913

mlir/include/mlir/Analysis/InferTypeOpInterface.td

@@ -37,7 +37,7 @@ def InferTypeOpInterface : OpInterface<"InferTypeOpInterface"> {
   }];
 
   let methods = [
-    InterfaceMethod<
+    StaticInterfaceMethod<
       /*desc=*/[{Returns the return types that an op would generate.
 
       The method takes an optional location which, if set, will be used to

mlir/test/lib/TestDialect/TestPatterns.cpp

@@ -73,10 +73,7 @@ struct ReturnTypeOpMatch : public RewritePattern {
   PatternMatchResult matchAndRewrite(Operation *op,
                                      PatternRewriter &rewriter) const final {
     if (auto retTypeFn = dyn_cast<InferTypeOpInterface>(op)) {
-      SmallVector<Value *, 4> values;
-      values.reserve(op->getNumOperands());
-      for (auto &operand : op->getOpOperands())
-        values.push_back(operand.get());
+      SmallVector<Value *, 4> values(op->getOperands());
       auto res = retTypeFn.inferReturnTypes(op->getLoc(), values,
                                             op->getAttrs(), op->getRegions());
       SmallVector<Type, 1> result_types(op->getResultTypes());
@@ -84,6 +81,20 @@ struct ReturnTypeOpMatch : public RewritePattern {
         return op->emitOpError(
                    "inferred type incompatible with return type of operation"),
                matchFailure();
+
+      // TODO(jpienaar): Split this out to make the test more focused.
+      // Create new op with unknown location to verify building with
+      // InferTypeOpInterface is triggered.
+      auto fop = op->getParentOfType<FuncOp>();
+      if (values[0] == fop.getArgument(0)) {
+        // Use the 2nd function argument if the first function argument is used
+        // when constructing the new op so that a new return type is inferred.
+        values[0] = fop.getArgument(1);
+        values[1] = fop.getArgument(1);
+        // TODO(jpienaar): Expand to regions.
+        rewriter.create<OpWithInferTypeInterfaceOp>(
+            UnknownLoc::get(op->getContext()), values, op->getAttrs());
+      }
     }
     return matchFailure();
   }

mlir/test/mlir-tblgen/return-types.mlir

@@ -1,8 +1,19 @@
 // RUN: mlir-opt %s -test-return-type -split-input-file -verify-diagnostics | FileCheck %s --dump-input-on-failure
 
 // CHECK-LABEL: testReturnTypeOpInterface
-func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>) {
+func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>, %arg1 : tensor<20xi32>) {
   %good = "test.op_with_infer_type_if"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
+  // CHECK: test.op_with_infer_type_if
+  // CHECK-SAME: tensor<20xi32>
+  // CHECK: test.op_with_infer_type_if
+  // CHECK-SAME: tensor<10xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: testReturnTypeOpInterface
+func @testReturnTypeOpInterface(%arg0 : tensor<10xf32>) {
   // expected-error@+1 {{incompatible with return type}}
   %bad = "test.op_with_infer_type_if"(%arg0, %arg0) : (tensor<10xf32>, tensor<10xf32>) -> tensor<*xf32>
   return

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

@@ -541,6 +541,11 @@ private:
   // operand's type as all results' types.
   void genUseOperandAsResultTypeCollectiveParamBuilder();
 
+  // Generates the build() method that takes aggregate operands/attributes
+  // parameters. This build() method uses inferred types as result types.
+  // Requires: The type needs to be inferable via InferTypeOpInterface.
+  void genInferedTypeCollectiveParamBuilder();
+
   // Generates the build() method that takes each operand/attribute as a
   // stand-alone parameter. The generated build() method uses first attribute's
   // type as all result's types.
@@ -968,11 +973,6 @@ void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
   auto &m = opClass.newMethod("void", "build", params, OpMethod::MP_Static);
   auto &body = m.body();
 
-  // Result types
-  SmallVector<std::string, 2> resultTypes(numResults, "operands[0]->getType()");
-  body << "  " << builderOpState << ".addTypes({"
-       << llvm::join(resultTypes, ", ") << "});\n\n";
-
   // Operands
   body << "  " << builderOpState << ".addOperands(operands);\n\n";
 
@@ -984,6 +984,27 @@ void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
     for (int i = 0; i < numRegions; ++i)
       m.body() << "  (void)" << builderOpState << ".addRegion();\n";
   }
+
+  // Result types
+  SmallVector<std::string, 2> resultTypes(numResults, "operands[0]->getType()");
+  body << "  " << builderOpState << ".addTypes({"
+       << llvm::join(resultTypes, ", ") << "});\n\n";
+}
+
+void OpEmitter::genInferedTypeCollectiveParamBuilder() {
+  // TODO(jpienaar): Expand to support regions.
+  std::string params =
+      (Twine("Builder *, OperationState &") + builderOpState +
+       ", ArrayRef<Value *> operands, ArrayRef<NamedAttribute> attributes")
+          .str();
+  auto &m = opClass.newMethod("void", "build", params, OpMethod::MP_Static);
+  auto &body = m.body();
+
+  body << "  " << builderOpState << ".addOperands(operands);\n\n";
+  body << "  " << builderOpState << ".addAttributes(attributes);\n";
+  body << "  " << builderOpState << ".addTypes(" << opClass.getClassName()
+       << "::inferReturnTypes(" << builderOpState
+       << ".location, operands, attributes, /*regions=*/{}));\n";
 }
 
 void OpEmitter::genUseOperandAsResultTypeSeparateParamBuilder() {
@@ -1026,15 +1047,17 @@ void OpEmitter::genUseAttrAsResultTypeBuilder() {
   } else {
     resultType = "attr.second.getType()";
   }
-  SmallVector<std::string, 2> resultTypes(op.getNumResults(), resultType);
-  body << "    " << builderOpState << ".addTypes({"
-       << llvm::join(resultTypes, ", ") << "});\n";
-  body << "  }\n";
 
   // Operands
   body << "  " << builderOpState << ".addOperands(operands);\n\n";
   // Attributes
   body << "  " << builderOpState << ".addAttributes(attributes);\n";
+
+  // Result types
+  SmallVector<std::string, 2> resultTypes(op.getNumResults(), resultType);
+  body << "    " << builderOpState << ".addTypes({"
+       << llvm::join(resultTypes, ", ") << "});\n";
+  body << "  }\n";
 }
 
 void OpEmitter::genBuilder() {
@@ -1082,7 +1105,7 @@ void OpEmitter::genBuilder() {
   genCollectiveParamBuilder();
   // 4. one having a stand-alone parameter for each operand and attribute,
   //    use the first operand or attribute's type as all result types
-  // to facilitate different call patterns.
+  //    to facilitate different call patterns.
   if (op.getNumVariadicResults() == 0) {
     if (op.getTrait("OpTrait::SameOperandsAndResultType")) {
       genUseOperandAsResultTypeSeparateParamBuilder();
@@ -1091,6 +1114,11 @@ void OpEmitter::genBuilder() {
     if (op.getTrait("OpTrait::FirstAttrDerivedResultType"))
       genUseAttrAsResultTypeBuilder();
   }
+  // TODO(jpienaar): Subsume this with general checking if type can be infered
+  // automatically.
+  // TODO(jpienaar): Expand to handle regions.
+  if (op.getTrait("InferTypeOpInterface::Trait") && op.getNumRegions() == 0)
+    genInferedTypeCollectiveParamBuilder();
 }
 
 void OpEmitter::genCollectiveParamBuilder() {