[mlir][ODS] Automatically create result_segment_sizes in builder

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.

Differential Revision: https://reviews.llvm.org/D132656

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

@@ -722,9 +722,6 @@ def AttrSizedOperandOp : TEST_Op<"attr_sized_operands",
 
 def AttrSizedResultOp : TEST_Op<"attr_sized_results",
                                 [AttrSizedResultSegments]> {
-  let arguments = (ins
-    DenseI32ArrayAttr:$result_segment_sizes
-  );
   let results = (outs
     Variadic<I32>:$a,
     Variadic<I32>:$b,
@@ -733,6 +730,11 @@ def AttrSizedResultOp : TEST_Op<"attr_sized_results",
   );
 }
 
+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">;

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

@@ -157,3 +157,10 @@ def OpL3 : NS_Op<"op_with_all_types_constraint",
 // 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

@@ -1514,6 +1514,34 @@ void OpEmitter::genSeparateArgParamBuilder() {
         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();