[mlir][ods] Fix OpDefinitionsGen infer return types builder with regions

[Mogball]

Summary

Despite handling regions and inferred return types, the builder was never generated for ops with both InferReturnTypeOpInterface and regions.

[jpienaar]

LG I'm not sure if TODO is addressed fully given condition move

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1224

We don't have any real uses where we utilize the types of the region to determine the type so this check was to conservatively skip until verified that it is useful. You move the check down, what does that change? (out of my cache)

[Mogball]

It's only partially addressed (region and successors -> regions).

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1224

This check was called in two places. In one place it stops a builder with inferredReturnTypes from being generated that does not handle regions (although there isn't any reason it shouldn't, since the extent of region handling is calling odsState.addRegion() N times). In another it does the same, except that builder has code for dealing with regions.

The builder gen in these edge cases could use a look-over.

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

 // Single variadic arg, non variadic results, with SameOperandsAndResultType.
 // Tests suppression of ambiguous build methods for operations with
 // SameOperandsAndResultType trait.
 def TableGenBuildOp4 : TEST_Op<"tblgen_build_4", [SameOperandsAndResultType]> {
   let arguments = (ins Variadic<AnyType>:$inputs);
   let results = (outs AnyType:$result);
 }
 
-// Single variadic arg with SameOperandsAndResultType and InferTypeOpInterface.
-// Tests suppression of ambiguous build methods for operations with
-// SameOperandsAndResultType and InferTypeOpInterface.
-def TableGenBuildOp5 : TEST_Op<"tblgen_build_5",
-      [SameOperandsAndResultType, InferTypeOpInterface]> {
+// Base class for testing `build` methods for ops with
+// InferReturnTypeOpInterface.
+class TableGenBuildInferReturnTypeBaseOp<string mnemonic,
+                                         list<OpTrait> traits = []>
+    : TEST_Op<mnemonic, [InferTypeOpInterface] # traits> {
   let arguments = (ins Variadic<AnyType>:$inputs);
   let results = (outs AnyType:$result);
 
   let extraClassDeclaration = [{
     static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *,
           ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
           ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
           ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
       inferredReturnTypes.assign({operands[0].getType()});
       return ::mlir::success();
     }
    }];
 }
 
+// Single variadic arg with SameOperandsAndResultType and InferTypeOpInterface.
+// Tests suppression of ambiguous build methods for operations with
+// SameOperandsAndResultType and InferTypeOpInterface.
+def TableGenBuildOp5 : TableGenBuildInferReturnTypeBaseOp<
+    "tblgen_build_5", [SameOperandsAndResultType]>;
+
+// Op with InferTypeOpInterface and regions.
+def TableGenBuildOp6 : TableGenBuildInferReturnTypeBaseOp<
+    "tblgen_build_6", [InferTypeOpInterface]> {
+  let regions = (region AnyRegion:$body);
+}
+
 //===----------------------------------------------------------------------===//
 // Test BufferPlacement
 //===----------------------------------------------------------------------===//
 
 def GetTupleElementOp: TEST_Op<"get_tuple_element"> {
   let description = [{
     Test op that returns a specified element of the tuple.
   }];

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

if (canUseUnwrappedRawValue(namedAttr.attr)) {
       canGenerate = true;
       break;
     }
   }
   return canGenerate;
 }
 
 static bool canInferType(Operator &op) {
-  return op.getTrait("::mlir::InferTypeOpInterface::Trait") &&
-         op.getNumRegions() == 0;
+  return op.getTrait("::mlir::InferTypeOpInterface::Trait");
 }
 
 void OpEmitter::genSeparateArgParamBuilder() {
   SmallVector<AttrParamKind, 2> attrBuilderType;
   attrBuilderType.push_back(AttrParamKind::WrappedAttr);
   if (canGenerateUnwrappedBuilder(op))
     attrBuilderType.push_back(AttrParamKind::UnwrappedValue);
 
auto emit = [&](AttrParamKind attrType, TypeParamKind paramKind,
     }
     llvm_unreachable("unhandled TypeParamKind");
   };
 
   // Some of the build methods generated here may be ambiguous, but TableGen's
   // ambiguous function detection will elide those ones.
   for (auto attrType : attrBuilderType) {
     emit(attrType, TypeParamKind::Separate, /*inferType=*/false);
-    if (canInferType(op))
+    if (canInferType(op) && op.getNumRegions() == 0)
       emit(attrType, TypeParamKind::None, /*inferType=*/true);
     emit(attrType, TypeParamKind::Collective, /*inferType=*/false);
   }
 }
 
 void OpEmitter::genUseOperandAsResultTypeCollectiveParamBuilder() {
   int numResults = op.getNumResults();
 
if (int numRegions = op.getNumRegions()) {
     body << llvm::formatv(
         "  for (unsigned i = 0; i != {0}; ++i)\n",
         (op.getNumVariadicRegions() ? "numRegions" : Twine(numRegions)));
     body << "    (void)" << builderOpState << ".addRegion();\n";
   }
 
   // Result types
   body << formatv(R"(
-    ::mlir::SmallVector<::mlir::Type, 2> inferredReturnTypes;
-    if (::mlir::succeeded({0}::inferReturnTypes(odsBuilder.getContext(),
-                  {1}.location, operands,
-                  {1}.attributes.getDictionary({1}.getContext()),
-                  /*regions=*/{{}, inferredReturnTypes))) {{)",
+          {1}.regions, inferredReturnTypes))) {{)",
                   opClass.getClassName(), builderOpState);
   if (numVariadicResults == 0 || numNonVariadicResults != 0)
-    body << "  assert(inferredReturnTypes.size()"
+    body << "\n    assert(inferredReturnTypes.size()"
          << (numVariadicResults != 0 ? " >= " : " == ") << numNonVariadicResults
-         << "u && \"mismatched number of return types\");\n";
-  body << "      " << builderOpState << ".addTypes(inferredReturnTypes);";
+         << "u && \"mismatched number of return types\");";
+  body << "\n    " << builderOpState << ".addTypes(inferredReturnTypes);";
 
   body << formatv(R"(
-    } else
-      ::llvm::report_fatal_error("Failed to infer result type(s).");)",
+  } else {{
+    ::llvm::report_fatal_error("Failed to infer result type(s).");
+  })",
                   opClass.getClassName(), builderOpState);
 }
 
 void OpEmitter::genUseOperandAsResultTypeSeparateParamBuilder() {
   SmallVector<MethodParameter> paramList;
   SmallVector<std::string, 4> resultNames;
   llvm::StringSet<> inferredAttributes;
   buildParamList(paramList, inferredAttributes, resultNames,
void OpEmitter::genCollectiveParamBuilder() {
   // Result types
   if (numVariadicResults == 0 || numNonVariadicResults != 0)
     body << "  assert(resultTypes.size()"
          << (numVariadicResults != 0 ? " >= " : " == ") << numNonVariadicResults
          << "u && \"mismatched number of return types\");\n";
   body << "  " << builderOpState << ".addTypes(resultTypes);\n";
 
   // Generate builder that infers type too.
-  // TODO: Expand to handle regions and successors.
+  // TODO: Expand to handle successors.
   if (canInferType(op) && op.getNumSuccessors() == 0)
     genInferredTypeCollectiveParamBuilder();
 }
 
 void OpEmitter::buildParamList(SmallVectorImpl<MethodParameter> &paramList,
                                llvm::StringSet<> &inferredAttributes,
                                SmallVectorImpl<std::string> &resultTypeNames,
                                TypeParamKind typeParamKind,

mlir/unittests/TableGen/OpBuildGen.cpp

 }
 
 TEST_F(
     OpBuildGenTest,
     BuildMethodsSameOperandsAndResultTypeAndInferOpTypeInterfaceSuppression) {
   testSingleVariadicInputInferredType<test::TableGenBuildOp5>();
 }
 
+TEST_F(OpBuildGenTest, BuildMethodsRegionsAndInferredType) {
+  auto op = builder.create<test::TableGenBuildOp6>(
+      loc, ValueRange{*cstI32, *cstF32}, /*attributes=*/noAttrs);
+  ASSERT_EQ(op->getNumRegions(), 1u);
+  verifyOp(std::move(op), {i32Ty}, {*cstI32, *cstF32}, noAttrs);
+}
+
 } // namespace mlir