[Mlir-commits] [mlir] 0393443 - [mlir][ODS] Automatically create `result_segment_sizes` in builder

[Markus Böck]

Author: Markus Böck
Date: 2022-08-25T19:09:18+02:00
New Revision: 0393443bf8069fbfaabcc64ee276ac7da5cf2814

URL: https://github.com/llvm/llvm-project/commit/0393443bf8069fbfaabcc64ee276ac7da5cf2814
DIFF: https://github.com/llvm/llvm-project/commit/0393443bf8069fbfaabcc64ee276ac7da5cf2814.diff

LOG: [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

Added: 
    

Modified: 
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/op-result.td
    mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 6732f3407b859..9567aa153a0a8 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/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">;

diff  --git a/mlir/test/mlir-tblgen/op-result.td b/mlir/test/mlir-tblgen/op-result.td
index e3af519662cae..cf7fb6c30870d 100644
--- a/mlir/test/mlir-tblgen/op-result.td
+++ b/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)}));

diff  --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index c3a06df700c0a..9d5c391224c0d 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/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();