[mlir][linalg] Add ods-gen helper to simplify the build methods.

[gysit]

Summary

Add a helper used to implement the build methods generated by ods-gen. The change reduces code size and compilation time since all structured op builders use the same build method. The change reduces the LinalgOps.cpp compilation time from 10.2s to 9.8s (debug build).

Depends On D123987

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

static void fillStructuredOpRegion(OpBuilder &opBuilder, Region &region,
   ImplicitLocOpBuilder b(opBuilder.getUnknownLoc(), opBuilder);
   regionBuilder(b, *body, attrs);
 
   // indexing_maps is an auto-generated method.
 
   // iterator_types is an auto-generated method.
 }
 
-/// Create the region and fill the block of a structured operation given
-/// `inputTypes` and `outputTypes` as well as a `regionBuilder`.
-void createAndFillStructuredOpRegion(OpBuilder &opBuilder,
-                                     OperationState &result,
-                                     TypeRange inputTypes,
-                                     TypeRange outputTypes,
-                                     RegionBuilderFn regionBuilder) {
-  Region &region = *result.addRegion();
-  fillStructuredOpRegion(opBuilder, region, inputTypes, outputTypes,
-                         result.attributes.getAttrs(), regionBuilder);
+/// Creates a structured operation given `inputs`, `outputs`, and `attributes`.
+/// The result types are derived automatically if `resultTensorTypes` is none.
+/// The body of the operation is filled using `regionBuilder`. All ods-gen
+/// created structured operations use the method to implement their builders.
+static void buildStructuredOp(OpBuilder &b, OperationState &state,
+                              llvm::Optional<TypeRange> resultTensorTypes,
+                              ValueRange inputs, ValueRange outputs,
+                              ArrayRef<NamedAttribute> attributes,
+  // Derive the result types if needed.
+  SmallVector<Type> derivedResultTypes =
+      resultTensorTypes.getValueOr(TypeRange());
+  if (!resultTensorTypes.hasValue())
+    copy_if(outputs.getTypes(), std::back_inserter(derivedResultTypes),
+            [](Type type) { return type.isa<RankedTensorType>(); });
+
+  state.addOperands(inputs);
+  state.addOperands(outputs);
+  state.addTypes(derivedResultTypes);
+  state.addAttributes(attributes);
+  state.addAttribute(
+      "operand_segment_sizes",
+      b.getI32VectorAttr({static_cast<int32_t>(inputs.size()),
+                          static_cast<int32_t>(outputs.size())}));
+
+  // Create and fill the region of the structured operation.
+  Region &region = *state.addRegion();
+  fillStructuredOpRegion(b, region, TypeRange(inputs), TypeRange(outputs),
+                         state.attributes.getAttrs(), regionBuilder);
 }
 
 /// Common parsing used for both named structured ops created by ods-gen and by
 /// manually defined C++ ops. Does not handle regions.
 static ParseResult
 parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
                              SmallVectorImpl<Type> &inputTypes,
                              SmallVectorImpl<Type> &outputTypes) {

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

 #       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 #  ODS-NEXT:       "ValueRange":$outputs,
 #  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
 
 #       ODS:  (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
 #  ODS-NEXT:       "ValueRange":$outputs, "Attribute":$cast,
 #  ODS-NEXT:       CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
 
-#       ODS:    $_state.addOperands(inputs);
-#  ODS-NEXT:    $_state.addOperands(outputs);
-#  ODS-NEXT:    $_state.addTypes(resultTensorTypes);
-#  ODS-NEXT:    $_state.addAttribute("cast", cast)
-#  ODS-NEXT:    $_state.addAttributes(attributes);
-#  ODS-NEXT:    $_state.addAttribute(
-#  ODS-NEXT:      "operand_segment_sizes",
-#  ODS-NEXT:      $_builder.getI32VectorAttr({
-#  ODS-NEXT:        static_cast<int32_t>(inputs.size()),
-#  ODS-NEXT:        static_cast<int32_t>(outputs.size())}));
-#  ODS-NEXT:    createAndFillStructuredOpRegion(
-#  ODS-NEXT:      $_builder,
-#  ODS-NEXT:      $_state,
-#  ODS-NEXT:      TypeRange(inputs),
-#  ODS-NEXT:      TypeRange(outputs),
-#  ODS-NEXT:      Test1Op::getRegionBuilder()
+#       ODS:    buildStructuredOp($_builder, $_state, resultTensorTypes,
+#  ODS-NEXT:      attributes, Test1Op::getRegionBuilder())
 
 
 # IMPL-LABEL:  void Test1Op::regionBuilder(ImplicitLocOpBuilder &b,
 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
 #       IMPL:  TypeFn castVal = TypeFn::cast_signed;
 #  IMPL-NEXT:  auto castIter = llvm::find_if(attrs, [&](const NamedAttribute &attr) {
 #  IMPL-NEXT:                                return attr.getName() == "cast"; });
 #  IMPL-NEXT:  if (castIter != attrs.end()) {

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

/*extraInterfaces=*/[{2}])> {
     let regions = (region AnyRegion:$region);
 
     let skipDefaultBuilders = 1;
     let builders = [
       OpBuilder<
       (ins "ValueRange":$inputs, "ValueRange":$outputs,
             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
       [{{
-        $_state.addOperands(inputs);
-        $_state.addOperands(outputs);
-        SmallVector<Type> resultTensorTypes;
-        copy_if(outputs.getTypes(),
-                std::back_inserter(resultTensorTypes),
-                [](Type type) {{ return type.isa<RankedTensorType>(); });
-        $_state.addTypes(resultTensorTypes);
-        $_state.addAttribute(
-          "operand_segment_sizes",
-          $_builder.getI32VectorAttr({{
-            static_cast<int32_t>(inputs.size()),
-            static_cast<int32_t>(outputs.size())}));
-        $_state.addAttributes(attributes);
-        createAndFillStructuredOpRegion(
-          $_builder,
-          $_state,
-          TypeRange(inputs),
-          TypeRange(outputs),
-          {0}::getRegionBuilder());
+        buildStructuredOp($_builder, $_state, llvm::None, inputs, outputs,
+          attributes, {0}::getRegionBuilder());
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
             "ValueRange":$outputs,
             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
       [{{
-        $_state.addOperands(inputs);
-        $_state.addOperands(outputs);
-        $_state.addTypes(resultTensorTypes);
-        $_state.addAttributes(attributes);
-        $_state.addAttribute(
-          "operand_segment_sizes",
-          $_builder.getI32VectorAttr({{
-            static_cast<int32_t>(inputs.size()),
-            static_cast<int32_t>(outputs.size())}));
-        createAndFillStructuredOpRegion(
-          $_builder,
-          $_state,
-          TypeRange(inputs),
-          TypeRange(outputs),
-          {0}::getRegionBuilder());
+        buildStructuredOp($_builder, $_state, resultTensorTypes,
+          inputs, outputs, attributes, {0}::getRegionBuilder());
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
       [{{
         $_state.addOperands(operands);
         $_state.addAttributes(attributes);
         $_state.addTypes(resultTensorTypes);

 // {1}: Comma interleaved attribute parameters
 // {2}: Attribute initialization
 static const char structuredOpBuilderFormat[] = R"FMT(
   , OpBuilder<
   (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
        "ValueRange":$outputs, {1},
        CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
   [{{
-    $_state.addOperands(inputs);
-    $_state.addOperands(outputs);
-    $_state.addTypes(resultTensorTypes);
     {2}
-    $_state.addAttributes(attributes);
-    $_state.addAttribute(
-      "operand_segment_sizes",
-      $_builder.getI32VectorAttr({{
-        static_cast<int32_t>(inputs.size()),
-        static_cast<int32_t>(outputs.size())}));
-    createAndFillStructuredOpRegion(
-      $_builder,
-      $_state,
-      TypeRange(inputs),
-      TypeRange(outputs),
-      {0}::getRegionBuilder());
+    buildStructuredOp($_builder, $_state, resultTensorTypes, inputs, outputs,
+      attributes, {0}::getRegionBuilder());
   }]>
 )FMT";
 
 // The iterator_types() method for structured ops. Parameters:
 // {0}: Class name
 // {1}: Comma interleaved iterator type names.
 static const char structuredOpIteratorTypesFormat[] =
     R"FMT(