Apply shortened printing/parsing form to linalg.reduce.

[olegshyshkov]

Awesome! Thank you!

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

1035

You can use a range-based loop over block.getArguments().drop_back() to avoid index computations. Or maybe even llvm::copy with std::back_inserter

1036

nit: you can use here and in other places below:

block.getArgument(i)

1103

nit: getOperand(i - 1)

1103

nit: getArgument(i)

[khasanovaa]

Resolve comments.

[khasanovaa]

Replace getArguments()[] with getArgument().

[stella.stamenova]

Looks like this added a failure to the already red mlir windows buildbot: https://lab.llvm.org/buildbot/#/builders/13/builds/30167

If you can't address this today, can you please revert it, so we don't go into the weekend with a red bot?

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

```
       %add = linalg.map
           ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
           outs(%init: tensor<64xf32>)
           (%lhs_elem: f32, %rhs_elem: f32) {
             %0 = arith.addf %lhs_elem, %rhs_elem: f32
             linalg.yield %0: f32
           }
     ```
+
+    Shortened print form is available. Applies to simple maps with one 
+    non-yield operation inside the body.
+
+    The example above will be printed as:
+    ```
+      %add = linalg.map { arith.addf }
+          ins(%lhs, %rhs : tensor<64xf32>, tensor<64xf32>)
+          outs(%init: tensor<64xf32>)
+    ```
   }];
 
   let arguments = (ins
     // Input args
     Variadic<TensorOrMemref>:$inputs,
 
     // Output arg
     TensorOrMemref:$init
let description = [{
 
     Example:
     ```
       %reduce = linalg.reduce
           ins(%input:tensor<16x32x64xf32>)
           outs(%init:tensor<16x64xf32>)
           dimensions = [1]
           (%in: f32, %out: f32) {
-            %0 = arith.addf %in, %out: f32
+            %0 = arith.addf %out, %in: f32
             linalg.yield %0: f32
           }
     ```
+
+    Shortened print form is available. Applies to simple (not variadic) reduces
+    with one non-yield operation inside the body. Applies only if the operation
+    takes `%out` as the first argument.
+
+    The example above will be printed as:
+    ```
+          %reduce = linalg.reduce { arith.addf }
+          ins(%input:tensor<16x32x64xf32>)
+          outs(%init:tensor<16x64xf32>)
+          dimensions = [1]
+    ```
   }];
 
   let arguments = (ins
     // Input arg
     Variadic<TensorOrMemref>:$inputs,
     // Output arg
     Variadic<TensorOrMemref>:$inits,
 

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

void MapOp::build(
   if (bodyBuild)
     buildGenericRegion(builder, result.location, *result.regions.front(),
                        inputs, /*outputs=*/{}, bodyBuild);
 }
 
 static void addBodyWithPayloadOp(OpAsmParser &parser, OperationState &result,
                                  const OperationName &payloadOpName,
                                  const NamedAttrList &payloadOpAttrs,
-                                 ArrayRef<Value> operands) {
+                                 ArrayRef<Value> operands,
+                                 bool initFirst = false) {
   OpBuilder b(parser.getContext());
   Region *body = result.addRegion();
   Block &block = body->emplaceBlock();
   b.setInsertionPointToStart(&block);
   SmallVector<Value> bbArgs;
   for (auto &operand : operands) {
     block.addArgument(operand.getType().cast<ShapedType>().getElementType(),
                       b.getUnknownLoc());
   }
+  SmallVector<Value> payloadOpOperands;
+  // If initFirst flag is enabled, we consider init as the first position of
+  // payload operands.
+  if (initFirst) {
+    payloadOpOperands.push_back(block.getArguments().back());
+    for (const auto& arg : block.getArguments().drop_back())
+      payloadOpOperands.push_back(arg);
+  } else {
+    payloadOpOperands = {block.getArguments().begin(),
+                         block.getArguments().end()};
+  }
 
   Operation *payloadOp = b.create(
       result.location, b.getStringAttr(payloadOpName.getStringRef()),
-      block.getArguments(),
+      payloadOpOperands,
       TypeRange{
           result.operands.back().getType().cast<ShapedType>().getElementType()},
       payloadOpAttrs);
-
   b.create<YieldOp>(result.location, payloadOp->getResults());
 }
 
 ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) {
   std::optional<OperationName> payloadOpName;
   NamedAttrList payloadOpAttrs;
   if (succeeded(parser.parseOptionalLBrace())) {
     FailureOr<OperationName> operationName = parser.parseCustomOperationName();
if (payloadOpName.has_value()) {
     if (parser.parseRegion(*body, regionArgs))
       return failure();
   }
   return success();
 }
 
 // Retrieve the operation from the body, if it is the only one (except
 // yield) and if it gets the same amount of arguments as the body does.
-static Operation *findPayloadOp(Block *body) {
+// If initFirst flag is enabled, we check that init takes the first position in
+// operands of payload.
+static Operation *findPayloadOp(Block *body, bool initFirst = false) {
   if (body->getOperations().size() != 2)
     return nullptr;
   Operation &payload = body->getOperations().front();
   assert(isa<YieldOp>(body->getOperations().back()));
 
   if (payload.getNumOperands() == 0 ||
       payload.getNumOperands() != body->getNumArguments())
     return nullptr;
+  if (initFirst) {
+    // check init
+    if (payload.getOperands().back() != body->getArgument(0))
+      return nullptr;
+    // check rest
+    for (int i = 1; i < body->getNumArguments(); ++i) {
+      if (payload.getOperand(i - 1) != body->getArgument(i)) {
+        return nullptr;
+      }
+    }
+  } else {
     for (const auto &[bbArg, operand] :
          llvm::zip(payload.getOperands(), body->getArguments())) {
       if (bbArg != operand)
         return nullptr;
     }
+  }
   return &payload;
 }
 
 void printShortForm(OpAsmPrinter &p, Operation *payloadOp) {
   SmallVector<StringRef> elidedAttrs;
   std::string attrToElide;
   p << " { " << payloadOp->getName().getStringRef();
   for (const auto &attr : payloadOp->getAttrs()) {
ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
   if (parseDstStyleOp(
           parser, result, [&](OpAsmParser &parser, NamedAttrList &attributes) {
             return parseDenseI64ArrayAttr(parser, attributes, "dimensions");
           }))
     return failure();
 
   if (payloadOpName.has_value()) {
     addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
-                         makeArrayRef(result.operands));
+                         makeArrayRef(result.operands), /*initFirst=*/true);
   } else {
     SmallVector<OpAsmParser::Argument> regionArgs;
     if (parser.parseArgumentList(regionArgs, OpAsmParser::Delimiter::Paren,
                                  /*allowType=*/true, /*allowAttrs=*/true)) {
       return failure();
     }
 
     Region *body = result.addRegion();
     if (parser.parseRegion(*body, regionArgs))
       return failure();
   }
 
   return success();
 }
 
 static void printDenseI64ArrayAttr(OpAsmPrinter &p, StringRef attributeName,
                                    ArrayRef<int64_t> attributeValue) {
   p << ' ' << attributeName << " = [" << attributeValue << "] ";
 }
 
 void ReduceOp::print(OpAsmPrinter &p) {
   Block *mapper = getBody();
-  Operation *payloadOp = findPayloadOp(mapper);
+  Operation *payloadOp = findPayloadOp(mapper, /*initFirst=*/true);
   if (payloadOp) {
     printShortForm(p, payloadOp);
   }
 
   printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
                                SmallVector<Value>(getDpsInitOperands()));
   printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
   p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});

mlir/test/Dialect/Linalg/one-shot-bufferize.mlir

 func.func @reduce(%input: tensor<16x32x64xf32>,
                   %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
   // CHECK:     linalg.reduce { arith.addf } ins(%[[INPUT]] : memref<16x32x64xf32
   %reduce = linalg.reduce
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<16x64xf32>)
       dimensions = [1]
       (%in: f32, %out: f32) {
-        %0 = arith.addf %in, %out: f32
+        %0 = arith.addf %out, %in: f32
         linalg.yield %0: f32
       }
   func.return %reduce : tensor<16x64xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose

mlir/test/Dialect/Linalg/roundtrip.mlir

 
 func.func @reduce(%input: tensor<16x32x64xf32>,
                   %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
   %reduce = linalg.reduce
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<16x64xf32>)
       dimensions = [1]
       (%in: f32, %out: f32) {
-        %0 = arith.addf %in, %out: f32
+        %0 = arith.addf %out, %in: f32
         linalg.yield %0: f32
       }
   func.return %reduce : tensor<16x64xf32>
 }
 // CHECK-LABEL: func @reduce
 //       CHECK:   linalg.reduce { arith.addf } ins
 //  CHECK-SAME:   outs
 //  CHECK-SAME:   dimensions = [1]
 
 // -----
 
 func.func @reduce_memref(%input: memref<16x32x64xf32>,
                          %init: memref<16x64xf32>) {
   linalg.reduce
       ins(%input:memref<16x32x64xf32>)
       outs(%init:memref<16x64xf32>)
       dimensions = [1]
       (%in: f32, %out: f32) {
-        %0 = arith.addf %in, %out: f32
+        %0 = arith.addf %out, %in: f32
         linalg.yield %0: f32
       }
   func.return
 }
 // CHECK-LABEL: func @reduce
 //       CHECK:   linalg.reduce { arith.addf } ins
 //  CHECK-SAME:   outs
 //  CHECK-SAME:   dimensions = [1]

 
 func.func @reduce_arith_with_attr(%input: tensor<16x32x64xf32>,
                   %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
   %reduce = linalg.reduce
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<16x64xf32>)
       dimensions = [1]
       (%in: f32, %out: f32) {
-        %0 = arith.addf %in, %out fastmath<fast> : f32
+        %0 = arith.addf %out, %in fastmath<fast> : f32
         linalg.yield %0: f32
       }
   func.return %reduce : tensor<16x64xf32>
 }
 // CHECK-LABEL: func @reduce_arith_with_attr
 // CHECK-NEXT:    %[[REDUCED:.*]] = linalg.reduce
 // CHECK-SAME:    { arith.addf {fastmath = #arith.fastmath<fast>} }
 // CHECK-SAME:    ins
 // CHECK-SAME:    outs
 // CHECK-SAME:    dimensions = [1]
 // CHECK-NEXT:    return %[[REDUCED]] : tensor<16x64xf32>