[mlir][ods] Fix OpFormatGen calling inferReturnTypes before region/segment resolution

[Mogball]

Summary

The generated parser for ops with type inference calls inferReturnTypes before region resolution and segment attribute resolution, i.e. regions and the segment attributes are not passed to the inferReturnTypes even though it may need that information.

In particular, an op that has sized operand segments which queries those operands in its inferReturnTypes function will crash because the segment attributes hadn't been added yet.

[rriddle]

Your description is very light on details, please extrapolate what this is fixing/what is broken.

[Mogball]

Yeah sorry >.<

I added some details

[rriddle]

LG, thanks.

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

 
 // Test inferReturnTypes is called when allOperands is true and operand types
 // are separately specified.
 def FormatInferTypeAllTypesOp
     : FormatInferAllTypesBaseOp<"format_infer_type_all_types"> {
   let assemblyFormat = "`(` operands `)` attr-dict `:` type($args)";
 }
 
+// Test inferReturnTypes coupled with regions.
+def FormatInferTypeRegionsOp
+    : TEST_Op<"format_infer_type_regions", [InferTypeOpInterface]> {
+  let results = (outs Variadic<AnyType>:$outs);
+  let regions = (region AnyRegion:$region);
+  let assemblyFormat = "$region attr-dict";
+  let extraClassDeclaration = [{
+    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
+          ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
+          ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
+          ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+      if (regions.empty())
+        return ::mlir::failure();
+      auto types = regions.front()->getArgumentTypes();
+      inferredReturnTypes.assign(types.begin(), types.end());
+      return ::mlir::success();
+    }
+  }];
+}
+
+// Test inferReturnTypes coupled with variadic operands (operand_segment_sizes).
+def FormatInferTypeVariadicOperandsOp
+    : TEST_Op<"format_infer_type_variadic_operands",
+              [InferTypeOpInterface, AttrSizedOperandSegments]> {
+  let arguments = (ins Variadic<I32>:$a, Variadic<I64>:$b);
+  let results = (outs Variadic<AnyType>:$outs);
+  let assemblyFormat = "`(` $a `:` type($a) `)` `(` $b `:` type($b) `)` attr-dict";
+  let extraClassDeclaration = [{
+    static ::mlir::LogicalResult inferReturnTypes(::mlir::MLIRContext *context,
+          ::llvm::Optional<::mlir::Location> location, ::mlir::ValueRange operands,
+          ::mlir::DictionaryAttr attributes, ::mlir::RegionRange regions,
+          ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+      FormatInferTypeVariadicOperandsOpAdaptor adaptor(operands, attributes);
+      auto aTypes = adaptor.getA().getTypes();
+      auto bTypes = adaptor.getB().getTypes();
+      inferredReturnTypes.append(aTypes.begin(), aTypes.end());
+      inferredReturnTypes.append(bTypes.begin(), bTypes.end());
+      return ::mlir::success();
+    }
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Test SideEffects
 //===----------------------------------------------------------------------===//
 
 def SideEffectOp : TEST_Op<"side_effect_op",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<TestEffectOpInterface>]> {
   let results = (outs AnyType:$result);

mlir/test/mlir-tblgen/op-format.mlir

 %ignored_res9:2 = test.format_infer_type_all_types_one_operand(%i64, %i32) : i64, i32
 
 // CHECK: test.format_infer_type_all_types_two_operands(%[[I64]], %[[I32]]) (%[[I64]], %[[I32]]) : i64, i32, i64, i32
 %ignored_res10:4 = test.format_infer_type_all_types_two_operands(%i64, %i32) (%i64, %i32) : i64, i32, i64, i32
 
 // CHECK: test.format_infer_type_all_types(%[[I64]], %[[I32]]) : i64, i32
 %ignored_res11:2 = test.format_infer_type_all_types(%i64, %i32) : i64, i32
 
+// CHECK: test.format_infer_type_regions
+// CHECK-NEXT: ^bb0(%{{.*}}: {{.*}}, %{{.*}}: {{.*}}):
+%ignored_res12:2 = test.format_infer_type_regions {
+^bb0(%arg0: i32, %arg1: f32):
+  "test.terminator"() : () -> ()
+}
+
+// CHECK: test.format_infer_type_variadic_operands(%[[I32]], %[[I32]] : i32, i32) (%[[I64]], %[[I64]] : i64, i64)
+%ignored_res13:4 = test.format_infer_type_variadic_operands(%i32, %i32 : i32, i32) (%i64, %i64 : i64, i64)
+
 //===----------------------------------------------------------------------===//
 // Check DefaultValuedStrAttr
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.has_str_value
 test.has_str_value {}

mlir/tools/mlir-tblgen/OpFormatGen.cpp

void OperationFormat::genParser(Operator &op, OpClass &opClass) {
   attrTypeCtx.withBuilder("parser.getBuilder()");
 
   // Generate parsers for each of the elements.
   for (auto &element : elements)
     genElementParser(element.get(), body, attrTypeCtx);
 
   // Generate the code to resolve the operand/result types and successors now
   // that they have been parsed.
-  genParserTypeResolution(op, body);
   genParserRegionResolution(op, body);
   genParserSuccessorResolution(op, body);
   genParserVariadicSegmentResolution(op, body);
+  genParserTypeResolution(op, body);
 
   body << "  return ::mlir::success();\n";
 }
 
 void OperationFormat::genElementParser(Element *element, MethodBody &body,
                                        FmtContext &attrTypeCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {