[mlir][linalg] Add nicer builders for `map` and `reduce`.

[olegshyshkov]

Summary

The new builders get a list of additional attrs, a lambda to build the region
body and infer return types from init.

[olegshyshkov]

Create region empty region in TransposeOp::parse to fix test breakage.

[olegshyshkov]

clang-format.

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

271

I personally find it very confusing to have something named init that is just ignored..
Init suggests it must be initialized properly, out does not have that baggage (but has other problems).

As an analogy, consider how one would write c = a + b vs c += a + b.
I am yet unaware the best way to proceed as it has implications on the dependence chains for bufferization.
The tradeoff goes relatively deep and also impacts how higher levels (or IR programmers) should target this, it is not just a trivial bikeshedding thing.

This is part of the bigger discussion in https://reviews.llvm.org/D136875.
This commit may need to be reverted depending on what the discussion above yields.

[pifon2a]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

271

Technically, it is init for the DPS interface, but we can rename to $out in the MapOp definition in TableGen. Would it be fine? I commented in the bigger discussion in the https://reviews.llvm.org/D136875.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

let arguments = (ins
     Variadic<TensorOrMemref>:$inputs,
 
     // Output arg
     TensorOrMemref:$init
   );
   let results = (outs Variadic<AnyTensor>:$result);
   let regions = (region SizedRegion<1>:$mapper);
 
+  let builders = [
+    OpBuilder<(ins "ValueRange":$inputs, "Value":$init,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>",
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
+  ];
+
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     // Implement functions necessary for LinalgStructuredInterface.
     SmallVector<StringRef> getIteratorTypesArray();
     ArrayAttr getIndexingMaps();
     std::string getLibraryCallName() {
       return "op_has_no_registered_library_name";
     }
 
let arguments = (ins
     Variadic<TensorOrMemref>:$inits,
 
     ConfinedAttr<DenseI64ArrayAttr,
                  [DenseArrayStrictlySorted<DenseI64ArrayAttr>]>:$dimensions
   );
   let results = (outs Variadic<AnyTensor>);
   let regions = (region SizedRegion<1>:$combiner);
 
+  let builders = [
+    OpBuilder<(ins "ValueRange":$inputs, "ValueRange":$inits,
+      "ArrayRef<int64_t>":$dimensions,
+      "function_ref<void(OpBuilder &, Location, ValueRange)>",
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
+  ];
+
   let extraClassDeclaration = structuredOpsBaseDecls # [{
     // Declare functions necessary for LinalgStructuredInterface.
     SmallVector<StringRef> getIteratorTypesArray();
     ArrayAttr getIndexingMaps();
     std::string getLibraryCallName() {
       return "op_has_no_registered_library_name";
     }
 

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

results
            FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
            FoldInsertPadIntoFill>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // GenericOp
 //===----------------------------------------------------------------------===//
 
+static void buildGenericRegion(
+    OpBuilder &builder, OperationState &result, ValueRange inputs,
+    ValueRange outputs,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild) {
+  SmallVector<Type, 4> blockArgTypes;
+  SmallVector<Location, 4> blockArgLocs;
+  for (ValueRange container : {inputs, outputs}) {
+    for (Value v : container) {
+      blockArgTypes.push_back(getElementTypeOrSelf(v));
+      blockArgLocs.push_back(v.getLoc());
+    }
+  }
+
+  OpBuilder::InsertionGuard guard(builder);
+  auto &region = *result.regions.front();
+  Block *bodyBlock =
+      builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
+  bodyBuild(builder, result.location, bodyBlock->getArguments());
+}
+
 void GenericOp::getAsmBlockArgumentNames(Region &region,
                                          OpAsmSetValueNameFn setNameFn) {
   for (Value v : getRegionInputArgs())
     setNameFn(v, "in");
   for (Value v : getRegionOutputArgs())
     setNameFn(v, "out");
 }
 
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayAttr indexingMaps,
     ArrayAttr iteratorTypes, StringAttr doc, StringAttr libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
     ArrayRef<NamedAttribute> attributes) {
   build(builder, result, resultTensorTypes, inputs, outputs, indexingMaps,
         iteratorTypes, doc, libraryCall);
   result.addAttributes(attributes);
-  if (!bodyBuild)
-    return;
-
-  SmallVector<Type, 4> blockArgTypes;
-  SmallVector<Location, 4> blockArgLocs;
-  for (ValueRange container : {inputs, outputs}) {
-    for (Value v : container) {
-      blockArgTypes.push_back(getElementTypeOrSelf(v));
-      blockArgLocs.push_back(v.getLoc());
-    }
-  }
-
-  OpBuilder::InsertionGuard guard(builder);
-  auto &region = *result.regions.front();
-  Block *bodyBlock =
-      builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
-  bodyBuild(builder, result.location, bodyBlock->getArguments());
+  if (bodyBuild)
+    buildGenericRegion(builder, result, inputs, outputs, bodyBuild);
 }
 
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
     ArrayRef<NamedAttribute> attributes) {
for (Value v : getRegionInputArgs())
     setNameFn(v, "in");
 }
 
 void MapOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "mapped");
 }
 
+void MapOp::build(
+    OpBuilder &builder, OperationState &result, ValueRange inputs, Value init,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, TypeRange{}, inputs, init);
+  result.addAttributes(attributes);
+
+  // Add output types for `RankedTensorType` output arguments.
+  Type initType = init.getType();
+  if (initType.isa<RankedTensorType>())
+    result.addTypes(initType);
+
+  if (bodyBuild)
+    buildGenericRegion(builder, result, inputs, /*outputs=*/{}, bodyBuild);
+}
+
 ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) {
   if (parseDstStyleOp(parser, result))
     return failure();
 
   SmallVector<OpAsmParser::Argument> regionArgs;
   if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
                                /*allowType=*/true, /*allowAttrs=*/true)) {
     return failure();

 }
 
 void ReduceOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "reduced");
 }
 
+void ReduceOp::build(
+    OpBuilder &builder, OperationState &result, ValueRange inputs,
+    ValueRange inits, ArrayRef<int64_t> dimensions,
+    function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,
+    ArrayRef<NamedAttribute> attributes) {
+  build(builder, result, TypeRange{}, inputs, inits, dimensions);
+  result.addAttributes(attributes);
+
+  // Add output types for `RankedTensorType` output arguments.
+  for (Value init : inits) {
+    Type initType = init.getType();
+    if (initType.isa<RankedTensorType>())
+      result.addTypes(initType);
+  }
+
+  if (bodyBuild)
+    buildGenericRegion(builder, result, inputs, inits, bodyBuild);
+}
+
 SmallVector<StringRef> ReduceOp::getIteratorTypesArray() {
   int64_t inputRank = getInputs()[0].getType().cast<ShapedType>().getRank();
   SmallVector<StringRef> iteratorTypes(inputRank,
                                        getParallelIteratorTypeName());
   for (int64_t reductionDim : getDimensions())
     iteratorTypes[reductionDim] = getReductionIteratorTypeName();
   return iteratorTypes;
 }
std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
                    mlir::ArrayRef<mlir::NamedAttribute>)>
 TransposeOp::getRegionBuilder() {
   return [](mlir::ImplicitLocOpBuilder &b, mlir::Block &block,
             mlir::ArrayRef<mlir::NamedAttribute>) {
     b.create<linalg::YieldOp>(block.getArguments().back());
   };
 }
 
-void TransposeOp::createRegion(::mlir::OpBuilder &opBuilder,
-                               ::mlir::OperationState &odsState) {
-  Region *region = odsState.addRegion();
-
-  SmallVector<Type> argTypes;
-  SmallVector<Location> argLocs;
-  for (auto t : odsState.operands) {
-    argTypes.push_back(getElementTypeOrSelf(t));
-    argLocs.push_back(opBuilder.getUnknownLoc());
-  }
-
-  // RAII.
-  OpBuilder::InsertionGuard guard(opBuilder);
-  Block *body =
-      opBuilder.createBlock(region, /*insertPt=*/{}, argTypes, argLocs);
-
-  ImplicitLocOpBuilder b(opBuilder.getUnknownLoc(), opBuilder);
-  getRegionBuilder()(b, *body, odsState.attributes.getAttrs());
-}
-
-void TransposeOp::build(::mlir::OpBuilder &odsBuilder,
-                        ::mlir::OperationState &odsState, Value input,
-                        Value init, DenseI64ArrayAttr permutation,
+void TransposeOp::build(::mlir::OpBuilder &builder,
+                        ::mlir::OperationState &result, Value input, Value init,
+                        DenseI64ArrayAttr permutation,
                         ArrayRef<NamedAttribute> attributes) {
-  odsState.addOperands(input);
-  odsState.addOperands(init);
-  odsState.addAttribute(getPermutationAttrName(odsState.name), permutation);
-  odsState.addAttributes(attributes);
-  odsState.addTypes(init.getType());
-
-  createRegion(odsBuilder, odsState);
+  result.addOperands(input);
+  result.addOperands(init);
+  result.addAttribute(getPermutationAttrName(result.name), permutation);
+  result.addAttributes(attributes);
+
+  // Add output types for `RankedTensorType` output arguments.
+  Type initType = init.getType();
+  if (initType.isa<RankedTensorType>())
+    result.addTypes(initType);
+
+  buildGenericRegion(builder, result, input, init,
+                     [&](OpBuilder &b, Location loc, ValueRange args) {
+                       b.create<linalg::YieldOp>(loc, args[0]);
+                     });
 }
 
-void TransposeOp::build(::mlir::OpBuilder &odsBuilder,
-                        ::mlir::OperationState &odsState, Value input,
-                        Value init, ArrayRef<int64_t> permutation,
+void TransposeOp::build(::mlir::OpBuilder &builder,
+                        ::mlir::OperationState &result, Value input, Value init,
+                        ArrayRef<int64_t> permutation,
                         ArrayRef<NamedAttribute> attributes) {
-  build(odsBuilder, odsState, input, init,
-        odsBuilder.getDenseI64ArrayAttr(permutation), attributes);
+  build(builder, result, input, init, builder.getDenseI64ArrayAttr(permutation),
+        attributes);
 }
 
 ParseResult TransposeOp::parse(OpAsmParser &parser, OperationState &result) {
   if (failed(parseDstStyleOp(
           parser, result, [&](OpAsmParser &parser, NamedAttrList &attributes) {
             return parseDenseI64ArrayAttr(parser, attributes, "permutation");
           })))
     return failure();
 
-  OpBuilder opBuilder(parser.getContext());
-  createRegion(opBuilder, result);
+  (void)result.addRegion();
+  OpBuilder builder(parser.getContext());
+  buildGenericRegion(builder, result, /*inputs=*/result.operands,
+                     /*outputs=*/{},
+                     [&](OpBuilder &b, Location loc, ValueRange args) {
+                       b.create<linalg::YieldOp>(loc, args[0]);
+                     });
   return success();
 }
 
 void TransposeOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "transposed");
 }