[mlir][ODS] Automatically create `result_segment_sizes` in builder

[zero9178]

Summary

When using multiple variadic results of differing sizes, using AttrSizedResultSegments is currently a requirement. Unlike AttrSizedOperandSegments however, it is not created within the default builders created by tablegen. Instead, one has to explicitly add DenseI32ArrayAttr:$result_segments_sizes as argument and then also explicitly specify all the sizes when using the builder from C++.

This patch fixes that redundancy, by making the builder generate the attribute in similar fashion as it already does for AttrSizedOperandSegments. The sizes required are simply gathered from the result type arguments of the builder.

[mehdi_amini]

LG, but please wait for at least one more approval from any of the other reviewers. Thanks!

[Mogball]

LGTM Thanks!

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

1529

if {} else if {} else {} seems cleaner to me here since they're all one liners

1542–1543

[zero9178]

Address review comments

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

let arguments = (ins
     I32:$c,
     Variadic<I32>:$d,
     DenseI32ArrayAttr:$operand_segment_sizes
   );
 }
 
 def AttrSizedResultOp : TEST_Op<"attr_sized_results",
                                 [AttrSizedResultSegments]> {
-  let arguments = (ins
-    DenseI32ArrayAttr:$result_segment_sizes
-  );
   let results = (outs
     Variadic<I32>:$a,
     Variadic<I32>:$b,
     I32:$c,
     Variadic<I32>:$d
   );
 }
 
+def AttrSizedResultCompileTestOp : TEST_Op<"attr_sized_results_compile_test",
+                                           [AttrSizedResultSegments]> {
+  let results = (outs Variadic<I32>:$a, I32:$b, Optional<I32>:$c);
+}
+
 // This is used to test that the fallback for a custom op's parser and printer
 // is the dialect parser and printer hooks.
 def CustomFormatFallbackOp : TEST_Op<"dialect_custom_format_fallback">;
 
 // Ops related to OIList primitive
 def OIListTrivial : TEST_Op<"oilist_with_keywords_only"> {
   let arguments = (ins UnitAttr:$keyword, UnitAttr:$otherKeyword,
                        UnitAttr:$diffNameUnitAttrKeyword);

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

def OpL3 : NS_Op<"op_with_all_types_constraint",
   let arguments = (ins I32Attr:$a);
   let results = (outs AnyType:$b);
 }
 
 // CHECK-LABEL: LogicalResult OpL3::inferReturnTypes
 // CHECK-NOT: }
 // CHECK: ::mlir::Type odsInferredType0 = attributes.get("a").cast<::mlir::TypedAttr>().getType();
 // CHECK: inferredReturnTypes[0] = odsInferredType0;
+
+def OpM : NS_Op<"mix_diff_size_variadic_and_normal_results_op", [AttrSizedResultSegments]> {
+  let results = (outs Variadic<AnyTensor>:$output1, AnyTensor:$output2, Optional<AnyTensor>:$output3);
+}
+
+// CHECK-LABEL: OpM::build(::mlir::OpBuilder &odsBuilder, ::mlir::OperationState &odsState, ::mlir::TypeRange output1, ::mlir::Type output2, /*optional*/::mlir::Type output3)
+// CHECK: odsState.addAttribute(result_segment_sizesAttrName(odsState.name), odsBuilder.getDenseI32ArrayAttr({static_cast<int32_t>(output1.size()), 1, (output3 ? 1 : 0)}));

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

case TypeParamKind::None:
       return;
     case TypeParamKind::Separate:
       for (int i = 0, e = op.getNumResults(); i < e; ++i) {
         if (op.getResult(i).isOptional())
           body << "  if (" << resultNames[i] << ")\n  ";
         body << "  " << builderOpState << ".addTypes(" << resultNames[i]
              << ");\n";
       }
+
+      // Automatically create the 'result_segment_sizes' attribute using
+      // the length of the type ranges.
+      if (op.getTrait("::mlir::OpTrait::AttrSizedResultSegments")) {
+        std::string getterName = op.getGetterName(resultSegmentAttrName);
+        body << " " << builderOpState << ".addAttribute(" << getterName
+             << "AttrName(" << builderOpState << ".name), "
+             << "odsBuilder.getDenseI32ArrayAttr({";
+
+        interleaveComma(
+            llvm::seq<int>(0, op.getNumResults()), body, [&](int i) {
+              const NamedTypeConstraint &result = op.getResult(i);
+              if (!result.isVariableLength()) {
+                body << "1";
+              } else if (result.isOptional()) {
+                body << "(" << resultNames[i] << " ? 1 : 0)";
+              } else {
+                // VariadicOfVariadic of results are currently unsupported in
+                // MLIR, hence it can only be a simple variadic.
+                // TODO: Add implementation for VariadicOfVariadic results here
+                //       once supported.
+                assert(result.isVariadic());
+                body << "static_cast<int32_t>(" << resultNames[i] << ".size())";
+              }
+            });
+        body << "}));\n";
+      }
+
       return;
     case TypeParamKind::Collective: {
       int numResults = op.getNumResults();
       int numVariadicResults = op.getNumVariableLengthResults();
       int numNonVariadicResults = numResults - numVariadicResults;
       bool hasVariadicResult = numVariadicResults != 0;
 
       // Avoid emitting "resultTypes.size() >= 0u" which is always true.