[mlir][linalg] Reuploading: add a shortened printing/parsing form for linalg.map and linalg.reduce.

[khasanovaa]

Summary

Original review: https://reviews.llvm.org/D140406

Revert commit: https://reviews.llvm.org/rGc0b775a5b506408bcdd9ffe31a51400a99734f2c

[khasanovaa]

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

1098

Fixed stack-use-after-scope error here. I was taking a StringRef of rvalue before

[olegshyshkov]

Could you please add links to the original review and the revert commit?

[khasanovaa]

In D140535#4012697, @olegshyshkov wrote:

Could you please add links to the original review and the revert commit?

Added to description

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

 static void printCommonStructuredOpParts(OpAsmPrinter &p, ValueRange inputs,
                                          ValueRange outputs) {
   if (!inputs.empty())
     p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
   if (!outputs.empty())
     p << " outs(" << outputs << " : " << outputs.getTypes() << ")";
 }
 
-static void printCommonStructuredOpPartsWithNewLine(OpAsmPrinter &p,
-                                                    ValueRange inputs,
-                                                    ValueRange outputs) {
-  if (!inputs.empty()) {
-    p << " ins(" << inputs << " : " << inputs.getTypes() << ")";
-  }
-  if (!outputs.empty()) {
-    p << " outs(" << outputs << " : " << outputs.getTypes() << ")";
-  }
-}
 //===----------------------------------------------------------------------===//
 // Specific parsing and printing for named structured ops created by ods-gen.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseNamedStructuredOpRegion(
     OpAsmParser &parser, Region &region, unsigned numRegionArgs,
     TypeRange inputTypes, TypeRange outputTypes, ArrayRef<NamedAttribute> attrs,
     RegionBuilderFn regionBuilder) {
void MapOp::build(
   if (initType.isa<RankedTensorType>())
     result.addTypes(initType);
 
   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) {
+  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());
+  }
+
+  Operation *payloadOp = b.create(
+      result.location, b.getStringAttr(payloadOpName.getStringRef()),
+      block.getArguments(),
+      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 (failed(operationName))
+      return failure();
+    if (parser.parseOptionalAttrDict(payloadOpAttrs))
+      return failure();
+    payloadOpName = operationName.value();
+    if (parser.parseRBrace())
+      return failure();
+  }
+
   if (parseDstStyleOp(parser, result))
     return failure();
 
+  if (payloadOpName.has_value()) {
+    addBodyWithPayloadOp(parser, result, payloadOpName.value(), payloadOpAttrs,
+                         makeArrayRef(result.operands).drop_back());
+  } 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();
 }
 
+// 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 (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;
+  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()) {
+    auto fastAttr = attr.getValue().dyn_cast<mlir::arith::FastMathFlagsAttr>();
+    if (fastAttr && fastAttr.getValue() == mlir::arith::FastMathFlags::none) {
+      attrToElide = attr.getName().str();
+      elidedAttrs.push_back(attrToElide);
+      break;
+    }
+  }
+  p.printOptionalAttrDict(payloadOp->getAttrs(), elidedAttrs);
+  p << " }";
+}
+
 void MapOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpPartsWithNewLine(
-      p, SmallVector<Value>(getDpsInputOperands()),
+  Block *mapper = getBody();
+  Operation *payloadOp = findPayloadOp(mapper);
+  if (payloadOp) {
+    printShortForm(p, payloadOp);
+  }
+
+  printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
                                SmallVector<Value>(getDpsInitOperands()));
   p.printOptionalAttrDict((*this)->getAttrs());
 
+  if (!payloadOp) {
+    // Print region if the payload op was not detected.
     p.increaseIndent();
     p.printNewline();
     p << "(";
-  llvm::interleaveComma(getMapper().getArguments(), p,
+    llvm::interleaveComma(mapper->getArguments(), p,
                           [&](auto arg) { p.printRegionArgument(arg); });
     p << ") ";
 
     p.printRegion(getMapper(), /*printEntryBlockArgs=*/false);
     p.decreaseIndent();
   }
+}
 
 LogicalResult MapOp::verify() {
   auto *bodyBlock = getBody();
   auto blockArgs = bodyBlock->getArguments();
 
   // Checks if the number of `inputs` match the arity of the `mapper` region.
   if (getInputs().size() != blockArgs.size())
     return emitOpError() << "expects number of operands to match the arity of "
                             "mapper, but got: "
                          << getInputs().size() << " and " << blockArgs.size();
 
-  // The parameters of mapper should all match the element type // of inputs.
+  // The parameters of mapper should all match the element type of inputs.
   for (const auto &[bbArgType, inputArg] :
        llvm::zip(bodyBlock->getArgumentTypes(), getInputs())) {
     auto inputElemType = inputArg.getType().cast<ShapedType>().getElementType();
     if (bbArgType != inputElemType) {
       return emitOpError() << "expected element type of input " << inputElemType
                            << " to match bbArg type " << bbArgType;
     }
   }
static ParseResult parseDenseI64ArrayAttr(OpAsmParser &parser,
   if (parser.parseKeyword(attributeName) || parser.parseEqual())
     return failure();
 
   attributes.set(attributeName, DenseI64ArrayAttr::parse(parser, Type{}));
   return success();
 }
 
 ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
+  std::optional<OperationName> payloadOpName;
+  NamedAttrList payloadOpAttrs;
+  if (succeeded(parser.parseOptionalLBrace())) {
+    FailureOr<OperationName> operationName = parser.parseCustomOperationName();
+    if (failed(operationName))
+      return failure();
+    if (parser.parseOptionalAttrDict(payloadOpAttrs))
+      return failure();
+    payloadOpName = operationName.value();
+    if (parser.parseRBrace())
+      return failure();
+  }
+
   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));
+  } 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) {
-  printCommonStructuredOpPartsWithNewLine(
-      p, SmallVector<Value>(getDpsInputOperands()),
-      SmallVector<Value>(getDpsInitOperands()));
+  Block *mapper = getBody();
+  Operation *payloadOp = findPayloadOp(mapper);
+  if (payloadOp) {
+    printShortForm(p, payloadOp);
+  }
 
+  printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
+                               SmallVector<Value>(getDpsInitOperands()));
   printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
   p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
-
+  if (!payloadOp) {
+    // Print region if the payload op was not detected.
     p.increaseIndent();
     p.printNewline();
     p << "(";
-  llvm::interleaveComma(getCombiner().getArguments(), p,
+    llvm::interleaveComma(mapper->getArguments(), p,
                           [&](auto arg) { p.printRegionArgument(arg); });
     p << ") ";
 
     p.printRegion(getCombiner(), /*printEntryBlockArgs=*/false);
     p.decreaseIndent();
   }
+}
 
 LogicalResult ReduceOp::verify() {
   ArrayRef<int64_t> dimensionsRef = getDimensions();
 
   for (int64_t i = 1; i < getNumDpsInputs(); ++i) {
     if (getInputs()[i].getType().cast<ShapedType>().getShape() !=
         getInputs()[0].getType().cast<ShapedType>().getShape()) {
       return emitOpError() << "expects all inputs to have the same shapes. "

 
 void TransposeOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "transposed");
 }
 
 void TransposeOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpPartsWithNewLine(
-      p, SmallVector<Value>(getDpsInputOperands()),
+  printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
                                SmallVector<Value>(getDpsInitOperands()));
   printDenseI64ArrayAttr(p, getPermutationAttrName(), getPermutation());
   p.printOptionalAttrDict((*this)->getAttrs(), {getPermutationAttrName()});
 }
 
 LogicalResult TransposeOp::verify() {
   ArrayRef<int64_t> permutationRef = getPermutation();
 
   if (!isPermutationVector(permutationRef))

 
 void BroadcastOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "broadcasted");
 }
 
 void BroadcastOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpPartsWithNewLine(
-      p, SmallVector<Value>(getDpsInputOperands()),
+  printCommonStructuredOpParts(p, SmallVector<Value>(getDpsInputOperands()),
                                SmallVector<Value>(getDpsInitOperands()));
   printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
   p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
 }
 
 LogicalResult BroadcastOp::verify() {
   ArrayRef<int64_t> dimensionsRef = getDimensions();
 
   auto inputType = getInput().getType();

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

 
 // -----
 
 // CHECK-LABEL: func @map_binary
 // CHECK-SAME:  %[[LHS:[0-9a-zA-Z]*]]: memref<64xf32
 // CHECK-SAME:  %[[RHS:[0-9a-zA-Z]*]]: memref<64xf32
 func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
                       %init: tensor<64xf32>) -> tensor<64xf32> {
-   // CHECK:      linalg.map ins(%[[LHS]], %[[RHS]] : memref<64xf32
+   // CHECK:      linalg.map { arith.addf } ins(%[[LHS]], %[[RHS]] : memref<64xf32
    %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
           }
   func.return %add : tensor<64xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @reduce
 // CHECK-SAME:  %[[INPUT:.*]]: memref<16x32x64xf32
 func.func @reduce(%input: tensor<16x32x64xf32>,
                   %init: tensor<16x64xf32>) -> tensor<16x64xf32> {
-  // CHECK:     linalg.reduce ins(%[[INPUT]] : memref<16x32x64xf32
+  // 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
         linalg.yield %0: f32
       }

mlir/test/Dialect/Linalg/roundtrip.mlir

%add = linalg.map
           outs(%init:tensor<64xf32>)
           (%lhs_elem: f32, %rhs_elem: f32) {
             %0 = arith.addf %lhs_elem, %rhs_elem: f32
             linalg.yield %0: f32
           }
   func.return %add : tensor<64xf32>
 }
 // CHECK-LABEL: func @map_binary
-//       CHECK:   linalg.map ins
+//       CHECK:   linalg.map { arith.addf } ins
 //  CHECK-SAME:   outs
-//  CHECK-NEXT:   (%{{.*}}: f32, %{{.*}}: f32) {
-//  CHECK-NEXT:     arith.addf
-//  CHECK-NEXT:     linalg.yield
-//  CHECK-NEXT:   }
 
 // -----
 
 func.func @map_binary_memref(%lhs: memref<64xf32>, %rhs: memref<64xf32>,
                       %init: memref<64xf32>) {
    linalg.map
       ins(%lhs, %rhs: memref<64xf32>, memref<64xf32>)
       outs(%init:memref<64xf32>)
%reduce = linalg.reduce
       dimensions = [1]
       (%in: f32, %out: f32) {
         %0 = arith.addf %in, %out: f32
         linalg.yield %0: f32
       }
   func.return %reduce : tensor<16x64xf32>
 }
 // CHECK-LABEL: func @reduce
-//       CHECK:   linalg.reduce ins
+//       CHECK:   linalg.reduce { arith.addf } ins
 //  CHECK-SAME:   outs
 //  CHECK-SAME:   dimensions = [1]
-//  CHECK-NEXT:   (%{{.*}}: f32, %{{.*}}: f32) {
-//  CHECK-NEXT:     arith.addf
-//  CHECK-NEXT:     linalg.yield
-//  CHECK-NEXT:   }
 
 // -----
 
 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
         linalg.yield %0: f32
       }
   func.return
 }
-// CHECK-LABEL: func @reduce_memref
-//       CHECK:     linalg.reduce
+// CHECK-LABEL: func @reduce
+//       CHECK:   linalg.reduce { arith.addf } ins
+//  CHECK-SAME:   outs
+//  CHECK-SAME:   dimensions = [1]
 
 // -----
 
 func.func @variadic_reduce(%input1: tensor<16x32x64xf32>,
     %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xi64>,
     %init2: tensor<16x64xi64>)  -> (tensor<16x64xf32>, tensor<16x64xi64>) {
   %reduce, %reduce2 = linalg.reduce
       ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xi64>)
       outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xi64>)
       dimensions = [1]
       (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {
         %0 = arith.addf %in1, %out1: f32
         %1 = arith.addi %in2, %out2: i64
         linalg.yield %0, %1: f32, i64
       }
   func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xi64>
 }
 // CHECK-LABEL: func @variadic_reduce
 //       CHECK:     linalg.reduce
+//   CHECK-NOT:     { arith.addf
 
 // -----
 
 func.func @variadic_reduce_memref(%input1: memref<16x32x64xf32>,
     %init1: memref<16x64xf32>, %input2: memref<16x32x64xi64>,
     %init2: memref<16x64xi64>) {
   linalg.reduce
       ins(%input1, %input2 : memref<16x32x64xf32>, memref<16x32x64xi64>)
       outs(%init1, %init2 : memref<16x64xf32>, memref<16x64xi64>)
       dimensions = [1]
       (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {
         %0 = arith.addf %in1, %out1: f32
         %1 = arith.addi %in2, %out2: i64
         linalg.yield %0, %1: f32, i64
       }
   func.return
 }
-// CHECK-LABEL: func @variadic_reduce_memref
+//   CHECK-LABEL: func @variadic_reduce_memref
 //       CHECK:     linalg.reduce
+//   CHECK-NOT:     { arith.addf
 
 // -----
 
 func.func @transpose(%input: tensor<16x32x64xf32>,
                      %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
   %transpose = linalg.transpose
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<32x64x16xf32>)
linalg.broadcast
       dimensions = [1]
   func.return
 }
 
 // CHECK-LABEL: func @broadcast_memref
 //      CHECK:    linalg.broadcast ins
 // CHECK-SAME:    outs
 // CHECK-SAME:    dimensions
+
+// -----
+
+func.func @map_arith_with_attr(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,
+                      %init: tensor<64xf32>) -> tensor<64xf32> {
+  %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 fastmath<fast> : f32
+            linalg.yield %0: f32
+          }
+  func.return %add : tensor<64xf32>
+}
+
+// CHECK-LABEL: func @map_arith_with_attr
+// CHECK-NEXT:    %[[MAPPED:.*]] = linalg.map
+// CHECK-SAME:    { arith.addf {fastmath = #arith.fastmath<fast>} }
+// CHECK-SAME:    ins
+// CHECK-SAME:    outs
+// CHECK-NEXT:    return %[[MAPPED]] : tensor<64xf32>
+
+// -----
+
+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
+        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>