[mlir][linalg][transform] Support dynamic tile sizes in TileToForeachThreadOp

[springerm]

Summary

TileToForeachThreadOp now accepts mixed SSA value operands / index attributes for tile_sizes and num_threads. (Reusing OperandsOrIntegersSizesList.) In case of an operand, a PDL_Operation must be specified that is mapped to a payload op that returns the tile size or number of threads.

[springerm]

fix bug

[springerm]

fix bug

[nicolasvasilache]

Generally LGTM once comments are addressed.

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

1027

Some comment for this plz?

1047

nit: add the index for better debugging: "num_threads #" << index << "..."

1055

nit: add the index for better debugging: "num_threads #" << index << "..."

1066

Seems like an opportunity to factor out all this duplicated code

1136

Hmm this is frustrating .. I think I saw @Mogball touch something related to these lists of index or attr recently ?

1213

There should already be code available to avoid duplicating all this, see dispatchIndexOpFoldResults
If that does not fit the bill, could you evolve into a better refactoring?.

1227

ditto

[springerm]

address comments

[springerm]

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

1136

oilist did the trick!

[nicolasvasilache]

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

1136

great!

mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td

let extraClassDeclaration = [{
     /// Returns the list of tile sizes, which may be static (Attribute) or
     /// dynamic (Value).
     SmallVector<OpFoldResult> getMixedSizes();
   }];
 }
 
 def TileToForeachThreadOp :
     Op<Transform_Dialect, "structured.tile_to_foreach_thread_op",
-      [FunctionalStyleTransformOpTrait,
-       MemoryEffectsOpInterface,
-       TransformEachOpTrait,
+      [AttrSizedOperandSegments,
+       DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
        TransformOpInterface]> {
   let description = [{
-    Tile a TilingInterface op to a tiled `scf.foreach_thread`. Tiling is
-    applied by either specifying `num_threads` or `tile_size`. If `num_threads`
-    is specified, then the tile size for each dimension `i` is calculated
-    dynamically via `ceilDiv(dimSize[i], num_threads[i])`.
-    If non-empty, the `thread_dim_mapping` is added as an attribute to the
-    resulting `scf.foreach_thread`.
-    Zero tile sizes indicate that the dimension is not tiled and can be
+    Tile a TilingInterface op to a tiled `scf.foreach_thread`.
+
+    Tiling is applied by either specifying `num_threads` or `tile_size`. If
+    `num_threads` is specified, then the tile size for each dimension `i` is
+    calculated dynamically via `ceilDiv(dimSize[i], num_threads[i])`.
+    `num_threads` and `tile_size` can be either static index attributes or SSA
+    values of PDL operation handle type (or a mix thereof). Operation handles
+    must be mapped to exactly one op that has exactly one result of index type.
+
+    Static zero tile sizes indicate that the dimension is not tiled and can be
     thought of as tiling by the full size of data.
+
     It is the user's responsibility to ensure that `num_threads/tile_sizes` is
     a valid tiling specification (i.e. that only tiles parallel dimensions, 
     e.g. in the Linalg case).
-    
+    If non-empty, the `thread_dim_mapping` is added as an attribute to the
+    resulting `scf.foreach_thread`.
+
     #### Return modes
-    
     This operation ignores ops that do not implement the TilingInterface and
     drops them in the return.
 
     If all the operations referred to by the `target` PDLOperation tile
     successfully, the transform succeeds.
     Otherwise the transform silently fails.
     
     The two returned handles point to only the subset of successfully produced
     tiled operations, which can all be empty.
     
     These two returned handles point to:
       - the new scf.foreach_thread op,
       - the tiled op that implements TilingInterface.
 
     ### Example using `num_threads`
 
     ```
-    %0 = pdl_match @match_matmul in %arg1    
+    %0 = pdl_match @match_matmul in %arg1
     %3:2 = transform.structured.tile_to_foreach_thread_op %0 num_threads [10, 20]
     ```
 
     ### Example using `tile_sizes`
-    
     ```
-    %0 = pdl_match @match_matmul in %arg1    
-    %3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 20, 0]
+    %0 = pdl_match @match_matmul in %arg1
+    %sz = pdl_match @match_size_op in %arg1
+    %3:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [0, %sz, 20]
     ```
   }];
 
   let arguments = (ins PDL_Operation:$target,
-                   // TODO: dynamic number of threads.
-                   OptionalAttr<DefaultValuedAttr<I64ArrayAttr, "{}">>:$num_threads,
-                   OptionalAttr<DefaultValuedAttr<I64ArrayAttr, "{}">>:$tile_sizes,
+                   Variadic<PDL_Operation>:$num_threads,
+                   Variadic<PDL_Operation>:$tile_sizes,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$static_num_threads,
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$static_tile_sizes,
                    OptionalAttr<I64ArrayAttr>:$thread_dim_mapping);
   let results = (outs PDL_Operation:$foreach_thread_op,
                       PDL_Operation:$tiled_op);
-
-  let assemblyFormat = [{
-    $target  (`num_threads` $num_threads^) : (`tile_sizes` $tile_sizes)?
-      (`(` `mapped` `to` `dims` $thread_dim_mapping^ `)`)? attr-dict
-  }];
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
 
   let extraClassDeclaration = [{
-    ::mlir::DiagnosedSilenceableFailure applyToOne(
-        ::mlir::TilingInterface target, 
-        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+    ::mlir::DiagnosedSilenceableFailure apply(
+        ::mlir::transform::TransformResults &transformResults,
         ::mlir::transform::TransformState &state);
+
+    ::llvm::SmallVector<::mlir::OpFoldResult> getMixedNumThreads();
+    ::llvm::SmallVector<::mlir::OpFoldResult> getMixedTileSizes();
   }];
 }
 
 def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let description = [{
     Indicates that the given `target` op all the ops it contains should be

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

void transform::TileOp::getEffects(
   producesHandle(getLoops(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
 // TileToForeachThreadOp
 //===----------------------------------------------------------------------===//
 
-DiagnosedSilenceableFailure transform::TileToForeachThreadOp::applyToOne(
-    TilingInterface target, SmallVectorImpl<Operation *> &results,
+DiagnosedSilenceableFailure transform::TileToForeachThreadOp::apply(
+    transform::TransformResults &transformResults,
     transform::TransformState &state) {
   IRRewriter rewriter(getContext());
-  rewriter.setInsertionPoint(target);
+  ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
+  SmallVector<Operation *> tileOps;
+  SmallVector<Operation *> tiledOps;
+
+  // getMixedNumThreads are OpFoldResults[index attributes or PDL operation].
+  // Convert to OpFoldResults[index attributes or payload op].
+  SmallVector<OpFoldResult> numThreads;
+  if (!getMixedNumThreads().empty()) {
+    for (OpFoldResult ofr : getMixedNumThreads()) {
+      if (ofr.is<Attribute>()) {
+        numThreads.push_back(ofr);
+        continue;
+      }
+      ArrayRef<Operation *> dynamicNumThreads =
+          state.getPayloadOps(ofr.get<Value>());
+      if (dynamicNumThreads.size() != 1) {
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << "num_threads handle must be mapped to exactly 1 payload op";
+        diag.attachNote(ofr.get<Value>().getLoc()) << "num_threads handle";
+        return diag;
+      }
+      Operation *op = dynamicNumThreads[0];
+      if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << "num_threads payload op must have exactly 1 index result";
+        diag.attachNote(op->getLoc()) << "num_threads payload";
+        return diag;
+      }
+      numThreads.push_back(op->getResult(0));
+    }
+  }
+
+  // getMixedTileSizes are OpFoldResults[index attributes or PDL operation].
+  // Convert to OpFoldResults[index attributes or payload op].
+  SmallVector<OpFoldResult> tileSizes;
+  if (!getMixedTileSizes().empty()) {
+    for (OpFoldResult ofr : getMixedTileSizes()) {
+      if (ofr.is<Attribute>()) {
+        tileSizes.push_back(ofr);
+        continue;
+      }
+      ArrayRef<Operation *> dynamicTileSizes =
+          state.getPayloadOps(ofr.get<Value>());
+      if (dynamicTileSizes.size() != 1) {
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << "tile_sizes handle must be mapped to exactly 1 payload op";
+        diag.attachNote(ofr.get<Value>().getLoc()) << "tile_sizes handle";
+        return diag;
+      }
+      Operation *op = dynamicTileSizes[0];
+      if (op->getNumResults() != 1 || !op->getResult(0).getType().isIndex()) {
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << "tile_sizes payload op must have exactly 1 index result";
+        diag.attachNote(op->getLoc()) << "tile_sizes payload";
+        return diag;
+      }
+      tileSizes.push_back(op->getResult(0));
+    }
+  }
+
+  for (Operation *target : targets) {
+    auto tilableOp = dyn_cast<TilingInterface>(target);
+    if (!tilableOp) {
+      DiagnosedSilenceableFailure diag =
+          emitSilenceableError() << "only TilingInterface ops are supported";
+      diag.attachNote(target->getLoc()) << "target op";
+      return diag;
+    }
+    rewriter.setInsertionPoint(tilableOp);
     auto maybeThreadDimMappingAttr = getThreadDimMapping();
-  auto dimMapping =
-      llvm::to_vector(maybeThreadDimMappingAttr
-                          ? extractFromI64ArrayAttr(*maybeThreadDimMappingAttr)
-                          : ArrayRef<int64_t>{});
+    auto dimMapping = llvm::to_vector(
+        maybeThreadDimMappingAttr
 
     FailureOr<ForeachThreadTilingResult> tilingResult = failure();
-  if (Optional<ArrayAttr> numThreads = getNumThreads())
-    tilingResult = linalg::tileToForeachThreadOp(
-        rewriter, target, getAsOpFoldResult(*numThreads), dimMapping);
-
-  if (Optional<ArrayAttr> tileSizes = getTileSizes())
+    if (!getMixedNumThreads().empty()) {
+      tilingResult = linalg::tileToForeachThreadOp(rewriter, tilableOp,
+                                                   numThreads, dimMapping);
+    } else {
       tilingResult = linalg::tileToForeachThreadOpUsingTileSizes(
-        rewriter, target, getAsOpFoldResult(*tileSizes), dimMapping);
+          rewriter, tilableOp, tileSizes, dimMapping);
+    }
 
     if (failed(tilingResult))
-    return emitDefaultSilenceableFailure(target);
-  rewriter.replaceOp(target, tilingResult->tileOp->getResults());
-  results.assign({tilingResult->tileOp, tilingResult->tiledOp});
+      return emitDefaultSilenceableFailure(tilableOp);
+    rewriter.replaceOp(tilableOp, tilingResult->tileOp->getResults());
+
+    tileOps.push_back(tilingResult->tileOp);
+    tiledOps.push_back(tilingResult->tiledOp);
+  }
+
+  transformResults.set(getForeachThreadOp().cast<OpResult>(), tileOps);
+  transformResults.set(getTiledOp().cast<OpResult>(), tiledOps);
+
   return DiagnosedSilenceableFailure(success());
 }
 
+void transform::TileToForeachThreadOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getTarget(), effects);
+  onlyReadsHandle(getTileSizes(), effects);
+  onlyReadsHandle(getNumThreads(), effects);
+  producesHandle(getResults(), effects);
+}
+
+ParseResult transform::TileToForeachThreadOp::parse(OpAsmParser &parser,
+                                                    OperationState &result) {
+  OpAsmParser::UnresolvedOperand target;
+  SmallVector<OpAsmParser::UnresolvedOperand> dynamicNumThreads,
+      dynamicTileSizes;
+  ArrayAttr staticNumThreads, staticTileSizes, threadDimMapping;
+  auto pdlOperationType = pdl::OperationType::get(parser.getContext());
+  if (parser.parseOperand(target) ||
+      parser.resolveOperand(target, pdlOperationType, result.operands))
+    return failure();
+  if (succeeded(parser.parseOptionalKeyword("num_threads"))) {
+    if (parseOperandsOrIntegersSizesList(parser, dynamicNumThreads,
+                                         staticNumThreads) ||
+        parser.resolveOperands(dynamicNumThreads, pdlOperationType,
+                               result.operands))
+      return failure();
+    result.addAttribute(getStaticNumThreadsAttrName(result.name),
+                        staticNumThreads);
+  } else if (succeeded(parser.parseOptionalKeyword("tile_sizes"))) {
+    if (parseOperandsOrIntegersSizesList(parser, dynamicTileSizes,
+                                         staticTileSizes) ||
+        parser.resolveOperands(dynamicTileSizes, pdlOperationType,
+                               result.operands))
+      return failure();
+    result.addAttribute(getStaticTileSizesAttrName(result.name),
+                        staticTileSizes);
+  }
+
+  if (succeeded(parser.parseOptionalLParen())) {
+    if (parser.parseKeyword("mapped") || parser.parseKeyword("to") ||
+        parser.parseKeyword("dims") ||
+        parser.parseCustomAttributeWithFallback(
+            threadDimMapping, parser.getBuilder().getType<NoneType>(),
+            getThreadDimMappingAttrName(result.name).getValue(),
+            result.attributes) ||
+        parser.parseOptionalRParen())
+      return failure();
+  }
+
+  if (parser.parseOptionalAttrDict(result.attributes))
+    return ParseResult::failure();
+
+  result.addAttribute("operand_segment_sizes",
+                      parser.getBuilder().getI32VectorAttr(
+                          {static_cast<int32_t>(1),
+                           static_cast<int32_t>(dynamicNumThreads.size()),
+                           static_cast<int32_t>(dynamicTileSizes.size())}));
+  result.addTypes({pdlOperationType, pdlOperationType});
+  return success();
+}
+
+void TileToForeachThreadOp::print(OpAsmPrinter &p) {
+  p << ' ' << getTarget();
+  if (!getStaticNumThreads().empty()) {
+    p << " num_threads ";
+    printOperandsOrIntegersSizesList(p, getOperation(), getNumThreads(),
+                                     getStaticNumThreads());
+  } else {
+    p << " tile_sizes ";
+    printOperandsOrIntegersSizesList(p, getOperation(), getTileSizes(),
+                                     getStaticTileSizes());
+  }
+  if (getThreadDimMapping()) {
+    p << " (mapped to dims ";
+    p.printAttributeWithoutType(getThreadDimMappingAttr());
+    p << ")";
+  }
+  p.printOptionalAttrDict((*this)->getAttrs(), {getStaticNumThreadsAttrName(),
+                                                getStaticTileSizesAttrName(),
+                                                getThreadDimMappingAttrName(),
+                                                "operand_segment_sizes"});
+}
+
+SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedNumThreads() {
+  ValueRange dynamic = getNumThreads();
+  SmallVector<int64_t> numThreads =
+      extractFromI64ArrayAttr(getStaticNumThreads());
+  SmallVector<OpFoldResult> results;
+  results.reserve(numThreads.size());
+  unsigned dynamicPos = 0;
+  Builder builder(getContext());
+  for (int64_t size : numThreads) {
+    if (size == ShapedType::kDynamicSize) {
+      results.push_back(dynamic[dynamicPos++]);
+    } else {
+      results.push_back(builder.getIndexAttr(size));
+    }
+  }
+  return results;
+}
+
+SmallVector<OpFoldResult> TileToForeachThreadOp::getMixedTileSizes() {
+  ValueRange dynamic = getTileSizes();
+  SmallVector<int64_t> tileSizes =
+      extractFromI64ArrayAttr(getStaticTileSizes());
+  SmallVector<OpFoldResult> results;
+  results.reserve(tileSizes.size());
+  unsigned dynamicPos = 0;
+  Builder builder(getContext());
+  for (int64_t size : tileSizes) {
+    if (size == ShapedType::kDynamicSize) {
+      results.push_back(dynamic[dynamicPos++]);
+    } else {
+      results.push_back(builder.getIndexAttr(size));
+    }
+  }
+  return results;
+}
+
+LogicalResult TileToForeachThreadOp::verify() {
+  if (getMixedNumThreads().empty() == getMixedTileSizes().empty())
+    return emitOpError("either num_threads or tile_sizes must be specified");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {

mlir/test/Dialect/Linalg/tile-to-foreach-thread.mlir

 transform.with_pdl_patterns {
 ^bb0(%arg0: !pdl.operation):
   transform.sequence %arg0 failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
     %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [10, 21]
   }
 }
+
+// -----
+
+// In this test case, matmul dims and tile size are dynamic.
+
+// CHECK-DAG: #[[$map0:.+]] = affine_map<()[s0, s1] -> (s0 ceildiv s1)>
+// CHECK-DAG: #[[$map1:.+]] = affine_map<()[s0] -> (s0 ceildiv 20)>
+// CHECK-DAG: #[[$map2:.+]] = affine_map<(d0)[s0, s1] -> (-(d0 * s1) + s0, s1)>
+// CHECK-DAG: #[[$map3:.+]] = affine_map<(d0)[s0] -> (d0 * -20 + s0, 20)>
+// CHECK-DAG: #[[$map4:.+]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$map5:.+]] = affine_map<(d0) -> (d0 * 20)>
+
+// CHECK-LABEL: matmul_tile_size_dynamic_dynamic(
+//  CHECK-SAME:   %[[A:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[B:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[C:[0-9a-z]+]]: tensor<?x?xf32>
+func.func @matmul_tile_size_dynamic_dynamic(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  //  CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
+  //  CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
+  //  CHECK-DAG: %[[tile_size:.*]] = "test.dummy"()
+  //  CHECK-DAG: %[[M:.+]] = tensor.dim %[[A]], %[[c0]] :
+  //  CHECK-DAG: %[[N:.+]] = tensor.dim %[[B]], %c1 :
+  //  CHECK-DAG: %[[NT0:.+]] = affine.apply #[[$map0]]()[%[[M]], %[[tile_size]]]
+  //  CHECK-DAG: %[[NT1:.+]] = affine.apply #[[$map1]]()[%[[N]]]
+  //      CHECK: scf.foreach_thread (%[[IV0:.+]], %[[IV1:.+]]) in (%[[NT0]], %[[NT1]])
+  //      CHECK    tensor.extract_slice %[[A]]
+  //      CHECK    tensor.extract_slice %[[B]]
+  //      CHECK    tensor.extract_slice %[[C]]
+  //      CHECK:   linalg.matmul
+  //      CHECK:   scf.foreach_thread.perform_concurrently
+  // CHECK-NEXT:    tensor.parallel_insert_slice
+  %tile_size = "test.dummy"() : () -> (index)
+  %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+                    outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
+  return %0 : tensor<?x?xf32>
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  transform.sequence %arg0 failures(propagate) {
+  ^bb1(%arg1: !pdl.operation):
+    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+    %sz = transform.structured.match ops{["test.dummy"]} in %arg1
+    %1:2 = transform.structured.tile_to_foreach_thread_op %0 tile_sizes [%sz, 20]
+  }
+}