[mlir] Update how scalable indices are printed

[awarzynski]

Summary

This patch makes sure that scalable indices (that would normally
represent scalable tile or vector sizes) are printed correctly, i.e.
with additional square brackets:

%1, %loop = transform.structured.tile %0 [2, 8, [4]]

This change complements https://reviews.llvm.org/D150944 and is a part
of a larger effort to enable scalable vectorisation in Linalg. See this
RFC for more context:

• https://discourse.llvm.org/t/rfc-scalable-vectorisation-in-linalg/

[ftynse]

mlir/include/mlir/Interfaces/ViewLikeInterface.h

55

Mlir code uses markdown-style backticks for inline code in comments rather than doxygen tags. If it had been using tags, it would have been \p rather than /p.

mlir/lib/Dialect/SCF/IR/SCF.cpp

1224

The function has the default value for isTrailingIdxScalable, no need to repeat it here and below.

[awarzynski]

Thank you for the review - I'll address your comments before landing this.

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 ///
 /// where `values` is of ODS type `Variadic<*>` and `integers` is of ODS
 /// type `I64ArrayAttr`. Prints a list with either (1) the static integer value
 /// in `integers` is `kDynamic` or (2) the next value otherwise. If `valueTypes`
 /// is non-empty, it is expected to contain as many elements as `values`
 /// indicating their types. This allows idiomatic printing of mixed value and
 /// integer attributes in a list. E.g.
 /// `[%arg0 : index, 7, 42, %arg42 : i32]`.
+///
+/// If /p isTrailingIdxScalable is true, then wrap the trailing index with
+/// square brackets, e.g. `[42]`, to denote scalability. This would normally be
+/// used for scalable tile or vector sizes.
 void printDynamicIndexList(
     OpAsmPrinter &printer, Operation *op, OperandRange values,
     ArrayRef<int64_t> integers, TypeRange valueTypes = TypeRange(),
-    AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square);
+    AsmParser::Delimiter delimiter = AsmParser::Delimiter::Square,
+    bool isTrailingIdxScalable = false);
 
 /// Parser hook for custom directive in assemblyFormat.
 ///
 ///   custom<DynamicIndexList>($values, $integers)
 ///   custom<DynamicIndexList>($values, $integers, type($values))
 ///
 /// where `values` is of ODS type `Variadic<*>` and `integers` is of ODS
 /// type `I64ArrayAttr`. Parse a mixed list with either (1) static integer

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

ParseResult transform::TileOp::parse(OpAsmParser &parser,
 
   result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
   result.addTypes(functionalType.getResults());
   return success();
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ' << getTarget();
-  printDynamicIndexList(p, getOperation(), getDynamicSizes(), getStaticSizes());
+  printDynamicIndexList(p, getOperation(), getDynamicSizes(), getStaticSizes(),
+                        /*valueTypes=*/{}, OpAsmParser::Delimiter::Square,
+                        getLastTileSizeScalable());
   printOptionalInterchange(p, getInterchange());
   p << " : ";
   p.printFunctionalType(getOperands().getTypes(), getResults().getTypes());
 }
 
 void transform::TileOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);

mlir/lib/Dialect/SCF/IR/SCF.cpp

 }
 
 void ForallOp::print(OpAsmPrinter &p) {
   Operation *op = getOperation();
   p << " (" << getInductionVars();
   if (isNormalized()) {
     p << ") in ";
     printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren,
+                          /*isTrailingIdxScalable=*/false);
   } else {
     p << ") = ";
     printDynamicIndexList(p, op, getDynamicLowerBound(), getStaticLowerBound(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren,
+                          /*isTrailingIdxScalable=*/false);
     p << " to ";
     printDynamicIndexList(p, op, getDynamicUpperBound(), getStaticUpperBound(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren,
+                          /*isTrailingIdxScalable=*/false);
     p << " step ";
     printDynamicIndexList(p, op, getDynamicStep(), getStaticStep(),
-                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren);
+                          /*valueTypes=*/{}, OpAsmParser::Delimiter::Paren,
+                          /*isTrailingIdxScalable=*/false);
   }
   printInitializationList(p, getRegionOutArgs(), getOutputs(), " shared_outs");
   p << " ";
   if (!getRegionOutArgs().empty())
     p << "-> (" << getResultTypes() << ") ";
   p.printRegion(getRegion(),
                 /*printEntryBlockArgs=*/false,
                 /*printBlockTerminators=*/getNumResults() > 0);
if (parser.parseArgumentList(ivs, OpAsmParser::Delimiter::Paren))
     return failure();
 
   DenseI64ArrayAttr staticLbs, staticUbs, staticSteps;
   SmallVector<OpAsmParser::UnresolvedOperand> dynamicLbs, dynamicUbs,
       dynamicSteps;
   if (succeeded(parser.parseOptionalKeyword("in"))) {
     // Parse upper bounds.
     if (parseDynamicIndexList(
-            parser, dynamicUbs, staticUbs, /*scalable=*/nullptr,
+            parser, dynamicUbs, staticUbs, /*isTrailingIdxScalable=*/nullptr,
             /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicUbs, indexType, result.operands))
       return failure();
 
     unsigned numLoops = ivs.size();
     staticLbs = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 0));
     staticSteps = b.getDenseI64ArrayAttr(SmallVector<int64_t>(numLoops, 1));
   } else {
     // Parse lower bounds.
     if (parser.parseEqual() ||
         parseDynamicIndexList(
-            parser, dynamicLbs, staticLbs, /*scalable=*/nullptr,
+            parser, dynamicLbs, staticLbs, /*isTrailingIdxScalable=*/nullptr,
             /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
 
         parser.resolveOperands(dynamicLbs, indexType, result.operands))
       return failure();
 
     // Parse upper bounds.
     if (parser.parseKeyword("to") ||
         parseDynamicIndexList(
-            parser, dynamicUbs, staticUbs, /*scalable=*/nullptr,
+            parser, dynamicUbs, staticUbs, /*isTrailingIdxScalable=*/nullptr,
             /*valueTypes=*/nullptr, OpAsmParser::Delimiter::Paren) ||
         parser.resolveOperands(dynamicUbs, indexType, result.operands))
       return failure();
 
     // Parse step values.
     if (parser.parseKeyword("step") ||
         parseDynamicIndexList(
             parser, dynamicSteps, staticSteps, /*scalable=*/nullptr,

mlir/lib/Dialect/Transform/Utils/Utils.cpp

if (parser.parseLParen().failed() ||
         parser.parseRParen().failed()) {
       return failure();
     }
     packed.emplace(packedOperand);
     integers = parser.getBuilder().getDenseI64ArrayAttr({});
     return success();
   }
 
-  return parseDynamicIndexList(parser, values, integers, /*scalable=*/nullptr,
-                               &valueTypes);
+  return parseDynamicIndexList(parser, values, integers,
+                               /*isTrailingIdxScalable=*/nullptr, &valueTypes);
 }

mlir/lib/Interfaces/ViewLikeInterface.cpp

default:
     llvm_unreachable("unsupported delimiter");
   }
 }
 
 void mlir::printDynamicIndexList(OpAsmPrinter &printer, Operation *op,
                                  OperandRange values,
                                  ArrayRef<int64_t> integers,
                                  TypeRange valueTypes,
-                                 AsmParser::Delimiter delimiter) {
+                                 AsmParser::Delimiter delimiter,
+                                 bool isTrailingIdxScalable) {
   char leftDelimiter = getLeftDelimiter(delimiter);
   char rightDelimiter = getRightDelimiter(delimiter);
   printer << leftDelimiter;
   if (integers.empty()) {
     printer << rightDelimiter;
     return;
   }
+
+  int64_t trailingScalableInteger;
+  if (isTrailingIdxScalable) {
+    // ATM only the trailing idx can be scalable
+    trailingScalableInteger = integers.back();
+    integers = integers.drop_back();
+  }
+
   unsigned idx = 0;
   llvm::interleaveComma(integers, printer, [&](int64_t integer) {
     if (ShapedType::isDynamic(integer)) {
       printer << values[idx];
       if (!valueTypes.empty())
         printer << " : " << valueTypes[idx];
       ++idx;
     } else {
       printer << integer;
     }
   });
+
+  // Print the trailing scalable index
+  if (isTrailingIdxScalable) {
+    printer << ", ";
+    printer << "[";
+    printer << trailingScalableInteger;
+    printer << "]";
+  }
+
   printer << rightDelimiter;
 }
 
 ParseResult mlir::parseDynamicIndexList(
     OpAsmParser &parser,
     SmallVectorImpl<OpAsmParser::UnresolvedOperand> &values,
     DenseI64ArrayAttr &integers, bool *isTrailingIdxScalable,
     SmallVectorImpl<Type> *valueTypes, AsmParser::Delimiter delimiter) {

mlir/test/Dialect/Transform/ops.mlir

 // CHECK: print
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !transform.any_op):
   transform.print %arg0 : !transform.any_op
   transform.print
   transform.print %arg0 {name = "test"} : !transform.any_op
   transform.print {name = "test"}
 }
+
+// CHECK: transform.sequence
+// CHECK: transform.structured.tile %0[4, 4, [4]]
+transform.sequence failures(propagate) {
+^bb0(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.structured.tile %0 [4, 4, [4]] : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
+}