[mlir] Print bbArgs and attrs of linalg.map/reduce/tranpose on the next line.

[pifon2a]

Summary

%mapped = linalg.map
  ins(%arg0 : tensor<64xf32>)
  outs(%arg1 : tensor<64xf32>)
  (%in: f32) {
    %0 = math.absf %in : f32
    linalg.yield %0 : f32
  }
%reduced = linalg.reduce
  ins(%arg0 : tensor<16x32x64xf32>)
  outs(%arg1 : tensor<16x64xf32>)
  dimensions = [1]
  (%in: f32, %init: f32) {
    %0 = arith.addf %in, %init : f32
    linalg.yield %0 : f32
  }
%transposed = linalg.transpose
  ins(%arg0 : tensor<16x32x64xf32>)
  outs(%arg1 : tensor<32x64x16xf32>)
  permutation = [1, 2, 0]

[pifon2a]

Reformat

mlir/include/mlir/IR/OpImplementation.h

public:
 
   /// Print a loc(...) specifier if printing debug info is enabled.
   virtual void printOptionalLocationSpecifier(Location loc) = 0;
 
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   virtual void printNewline() = 0;
 
+  /// Increase indentation.
+  virtual void increaseIndent() = 0;
+
+  /// Decrease indentation.
+  virtual void decreaseIndent() = 0;
+
   /// Print a block argument in the usual format of:
   ///   %ssaName : type {attr1=42} loc("here")
   /// where location printing is controlled by the standard internal option.
   /// You may pass omitType=true to not print a type, and pass an empty
   /// attribute list if you don't care for attributes.
   virtual void printRegionArgument(BlockArgument arg,
                                    ArrayRef<NamedAttribute> argAttrs = {},
                                    bool omitType = false) = 0;

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) {
+  p.printNewline();
+  if (!inputs.empty())
+    p << "ins(" << inputs << " : " << inputs.getTypes() << ")";
+  p.printNewline();
+  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) {
ParseResult MapOp::parse(OpAsmParser &parser, OperationState &result) {
   Region *body = result.addRegion();
   if (parser.parseRegion(*body, regionArgs))
     return failure();
 
   return success();
 }
 
 void MapOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpParts(p, SmallVector<Value>(getInputOperands()),
-                               SmallVector<Value>(getOutputOperands()));
+  p.increaseIndent();
+  printCommonStructuredOpPartsWithNewLine(
+      p, SmallVector<Value>(getInputOperands()),
   p.printOptionalAttrDict((*this)->getAttrs());
 
+  p.printNewline();
   p << "(";
   llvm::interleaveComma(getMapper().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())
ParseResult ReduceOp::parse(OpAsmParser &parser, OperationState &result) {
   if (parser.parseRegion(*body, regionArgs))
     return failure();
 
   return success();
 }
 
 static void printDenseI64ArrayAttr(OpAsmPrinter &p, StringRef attributeName,
                                    ArrayRef<int64_t> attributeValue) {
-  p << " " << attributeName << " = [" << attributeValue << "] ";
+  p << attributeName << " = [" << attributeValue << "] ";
 }
 
 void ReduceOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpParts(p, SmallVector<Value>(getInputOperands()),
-                               SmallVector<Value>(getOutputOperands()));
+  p.increaseIndent();
+  printCommonStructuredOpPartsWithNewLine(
+      p, SmallVector<Value>(getInputOperands()),
+  p.printNewline();
   printDenseI64ArrayAttr(p, getDimensionsAttrName(), getDimensions());
   p.printOptionalAttrDict((*this)->getAttrs(), {getDimensionsAttrName()});
 
+  p.printNewline();
   p << "(";
   llvm::interleaveComma(getCombiner().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 < getNumInputs(); ++i) {
     if (getInputs()[i].getType().cast<ShapedType>().getShape() !=
         getInputs()[0].getType().cast<ShapedType>().getShape()) {

 
 void TransposeOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   if (!getResults().empty())
     setNameFn(getResults().front(), "transposed");
 }
 
 void TransposeOp::print(OpAsmPrinter &p) {
-  printCommonStructuredOpParts(p, SmallVector<Value>(getInputOperands()),
-                               SmallVector<Value>(getOutputOperands()));
+  p.increaseIndent();
+  printCommonStructuredOpPartsWithNewLine(
+      p, SmallVector<Value>(getInputOperands()),
+  p.printNewline();
   printDenseI64ArrayAttr(p, getPermutationAttrName(), getPermutation());
   p.printOptionalAttrDict((*this)->getAttrs(), {getPermutationAttrName()});
+  p.decreaseIndent();
 }
 
 LogicalResult TransposeOp::verify() {
   ArrayRef<int64_t> permutationRef = getPermutation();
 
   if (!isPermutation(permutationRef))
     return emitOpError("permutation is not valid");
 

mlir/lib/IR/AsmPrinter.cpp

private:
   raw_ostream &getStream() const override { return os; }
 
   /// The following are hooks of `OpAsmPrinter` that are not necessary for
   /// determining potential aliases.
   void printFloat(const APFloat &value) override {}
   void printAffineMapOfSSAIds(AffineMapAttr, ValueRange) override {}
   void printAffineExprOfSSAIds(AffineExpr, ValueRange, ValueRange) override {}
   void printNewline() override {}
+  void increaseIndent() override {}
+  void decreaseIndent() override {}
   void printOperand(Value) override {}
   void printOperand(Value, raw_ostream &os) override {
     // Users expect the output string to have at least the prefixed % to signal
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
   void printKeywordOrString(StringRef) override {}
public:
 
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   void printNewline() override {
     os << newLine;
     os.indent(currentIndent);
   }
 
+  /// Increase indentation.
+  void increaseIndent() override { currentIndent += indentWidth; }
+
+  /// Decrease indentation.
+  void decreaseIndent() override { currentIndent -= indentWidth; }
+
   /// Print a block argument in the usual format of:
   ///   %ssaName : type {attr1=42} loc("here")
   /// where location printing is controlled by the standard internal option.
   /// You may pass omitType=true to not print a type, and pass an empty
   /// attribute list if you don't care for attributes.
   void printRegionArgument(BlockArgument arg,
                            ArrayRef<NamedAttribute> argAttrs = {},
                            bool omitType = false) override;

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
-   // CHECK-SAME: ins(%[[LHS]], %[[RHS]] : memref<64xf32
+   // CHECK-NEXT: 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
-  // CHECK-SAME: ins(%[[INPUT]] : memref<16x32x64xf32
+  // CHECK-NEXT: 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
       }
   func.return %reduce : tensor<16x64xf32>
 }
 
 // -----
 
 // CHECK-LABEL: func @transpose
 // CHECK-SAME:  %[[ARG0:.*]]: memref<16x32x64xf32
 func.func @transpose(%input: tensor<16x32x64xf32>,
                      %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
   // CHECK:      linalg.transpose
-  // CHECK-SAME: ins(%[[ARG0]] : memref<16x32x64xf32
+  // CHECK-NEXT: ins(%[[ARG0]] : memref<16x32x64xf32
   %transpose = linalg.transpose
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<32x64x16xf32>)
       permutation = [1, 2, 0]
   func.return %transpose : tensor<32x64x16xf32>
 }
 
 // -----

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
+//       CHECK:   linalg.map
+//  CHECK-NEXT:   ins
+//  CHECK-NEXT:   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
+//       CHECK:   linalg.reduce
+//  CHECK-NEXT:   ins
+//  CHECK-NEXT:   outs
+//  CHECK-NEXT:   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>)
func.func @transpose(%input: tensor<16x32x64xf32>,
                      %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {
   %transpose = linalg.transpose
       ins(%input:tensor<16x32x64xf32>)
       outs(%init:tensor<32x64x16xf32>)
       permutation = [1, 2, 0]
   func.return %transpose : tensor<32x64x16xf32>
 }
 // CHECK-LABEL: func @transpose
+//      CHECK:    linalg.transpose
+// CHECK-NEXT:    ins
+// CHECK-NEXT:    outs
+// CHECK-NEXT:    permutation
 
 // -----
 
 func.func @transpose_memref(%input: memref<16x32x64xf32>,
                             %init: memref<32x64x16xf32>) {
   linalg.transpose
       ins(%input:memref<16x32x64xf32>)
       outs(%init:memref<32x64x16xf32>)
       permutation = [1, 2, 0]
   func.return
 }
 // CHECK-LABEL: func @transpose_memref