[mlir] make multi-size tiling use transform parameters

[ftynse]

Summary

Use the recently introduced transform dialect parameter mechanism to
perform controllable multi-size tiling with sizes computed at the
transformation time rather than at runtime.

This requires to generalize tile and split structured transform
operations to work with any transform dialect handle types, which is
desirable in itself to avoid unchecked overuse of PDL OperationType.

[springerm]

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

23

This almost looks like a transform dialect equivalent of OpFoldResult. If we had Values instead of Operation *. Maybe worth adding as a comment.

25

Should this be TransformParamTypeInterface?

320–323

Result types are missing

334–337

Result types are missing

[ftynse]

Address review.

[ftynse]

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

25

No, I want specifically the Transform_ParamType that was checked to be associated with an integer attribute.

[ftynse]

Rebase

[ftynse]

Provide default types in the C++ and Python builders.

[shabalin]

LGTM, thanks for adding extra builders to smooth out breaking changes!

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

 //===- LinalgTransformOps.td - Linalg transform ops --------*- tablegen -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_TRANSFORM_OPS
 #define LINALG_TRANSFORM_OPS
 
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Dialect/SCF/IR/DeviceMappingInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/EnumAttr.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/RegionKindInterface.td"
 
+// This is roughly similar to OpFoldResult assuming the handle produces a single
+// value in the payload IR.
+def TransformParamTypeOrAnyHandle : Type<
+    Or<[TransformHandleTypeInterface.predicate,
+        Transform_ParamType.predicate]>,
+    "transform 'param' type or any handle type">;
+
 //===----------------------------------------------------------------------===//
 // DecomposeOp
 //===----------------------------------------------------------------------===//
 
 def DecomposeOp : Op<Transform_Dialect, "structured.decompose",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
let description = [{
     pointing to single-result index-typed operations (which may be arithmetic
     constant operations) defining the two respective tile sizes and the product
     of the first tile size with the number of tiles of that size (useful for
     splitting the iteration space).
 
     This operation composes with the regular tiling when applied per-dimension:
 
     ```mlir
     %sz1, %sz2, %split = structured.multitile_sizes %target
                          { target_size = 10, dimension = 1 }
+                       : !transform.any_op, !transform.param<i64>,
+                         !transform.param<i64>, !transform.param<i64>
     %low, %high = structured.split %target after %split { dimension = 1 }
-    %tiled_low = structured.tile %low [0, %sz1]
-    %tiled_high = structured.tile %high [0, %sz2]
-    %common = merge_handles %tiled_low, %tiled_high
+                : !transform.any_op, !transform.param<i64>
+    %tiled_low, %loop1 = structured.tile %low [0, %sz1]
+                       : (!transform.any_op, !transform.param<i64>)
+                      -> (!transform.any_op, !transform.any_op)
+    %tiled_high, %loop2 = structured.tile %high [0, %sz2]
+                        : (!transform.any_op, !transform.param<i64>)
+                       -> (!transform.any_op, !transform.any_op)
+    %common = merge_handles %tiled_low, %tiled_high : !transform.any_op
 
     %sz3, %sz4, %split = structured.multitile_size %target
                          { target_size = 42, dimension = 0 }
+                       : !transform.any_op, !transform.any_op,
+                         !transform.any_op, !transform.any_op
     %sz3r, %sz4r, %splitr = replicate num(%common) %sz3, %sz4, %splitr
+             : !transform.any_op, !transform.any_op, !transform.any_op
     structured.split %common after %splitr { dimension = 0 }
+             : !transform.any_op, !transform.any_op
     // ...
     ```
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
                        I64Attr:$dimension,
                        I64Attr:$target_size,
                        DefaultValuedAttr<I64Attr, "1">:$divisor);
-  let results = (outs PDL_Operation:$low_size,
-                      PDL_Operation:$high_size,
-                      PDL_Operation:$split_point);
-  let assemblyFormat = "$target attr-dict";
+  let results = (outs TransformParamTypeOrAnyHandle:$low_size,
+                      TransformParamTypeOrAnyHandle:$high_size,
+                      TransformParamTypeOrAnyHandle:$split_point);
+  let hasVerifier = 1;
+  let assemblyFormat =
+    "$target attr-dict `:` custom<MultitileSizesTypes>("
+    "type($target), type($low_size), type($high_size), type($split_point))";
 
   let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::linalg::LinalgOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         TransformState &state);
   }];
 }
let description = [{
 
     The operation consumes the target handle, but preserves the split point
     handle if provided. It produces two new handles pointing to the two parts
     of the structured op after splitting, in the same order as the target
     operand, with the first handle corresponding to the part with lower
     iteration space indices.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
+  let arguments = (ins TransformHandleTypeInterface:$target,
                        I64Attr:$dimension,
-                       Optional<PDL_Operation>:$dynamic_split_point,
+                       Optional<TransformParamTypeOrAnyHandle>:$dynamic_split_point,
                        I64Attr:$static_split_point);
-  let results = (outs PDL_Operation:$first, PDL_Operation:$second);
+  let results = (outs TransformHandleTypeInterface:$first,
+                      TransformHandleTypeInterface:$second);
   let hasVerifier = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // SplitReductionOp
 //===----------------------------------------------------------------------===//
 
let description = [{
     This operation produces a silenceable failure if the `dynamic_sizes` handles
     are associated with lists of payload operations of a size different than
     that of the list associated with the `target` handle.
 
     If the internal implementation of tiling for any of the operations fails,
     produces a definite failure.
   }];
 
-  let arguments = (ins PDL_Operation:$target,
-                   Variadic<PDL_Operation>:$dynamic_sizes,
+  let arguments = (ins TransformHandleTypeInterface:$target,
+                   Variadic<TransformParamTypeOrAnyHandle>:$dynamic_sizes,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_sizes,
                    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$interchange);
-  let results = (outs PDL_Operation:$tiled_linalg_op,
-                      Variadic<PDL_Operation>:$loops);
+  let results = (outs TransformHandleTypeInterface:$tiled_linalg_op,
+                      Variadic<TransformHandleTypeInterface>:$loops);
   let builders = [
+    OpBuilder<(ins "TypeRange":$loopTypes,
+                   "Value":$target,
+                   "ArrayRef<int64_t>":$staticTileSizes,
+                   CArg<"ArrayRef<int64_t>", "{}">:$interchange)>,
+    OpBuilder<(ins "TypeRange":$loopTypes,
+                   "Value":$target,
+                   "ArrayRef<OpFoldResult>":$mixedTileSizes,
+                   CArg<"ArrayRef<int64_t>", "{}">:$interchange)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<int64_t>":$staticTileSizes,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange)>,
     OpBuilder<(ins "Value":$target,
                    "ArrayRef<OpFoldResult>":$mixedTileSizes,
                    CArg<"ArrayRef<int64_t>", "{}">:$interchange)>
+
   ];
 
   let hasCustomAssemblyFormat = 1;
 
   let extraClassDeclaration = [{
     /// Returns the list of tile sizes, which may be static (Attribute) or
     /// dynamic (Value).
     SmallVector<OpFoldResult> getMixedSizes();

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 /// of the LinalgOp to the corresponding non-empty range indices of newly
 /// created loops.
 using LoopIndexToRangeIndexMap = DenseMap<int, int>;
 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
 makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map,
                     ArrayRef<OpFoldResult> allShapeSizes,
                     ArrayRef<OpFoldResult> allTileSizes);
 
-/// A description of a multi-size tiling comprising tile sizes and numbers of
-/// tiles, expressed as Values which may or may not be constant. Multi-size
-/// currently means two-size.
-struct MultiSizeSpecification {
+namespace detail {
+template <typename T>
+struct MultiSizeSpecificationBase {
   /// Tile sizes.
-  Value lowTileSize, highTileSize;
+  T lowTileSize, highTileSize;
   /// Number of tiles associated with each size.
-  Value lowTripCount, highTripCount;
+  T lowTripCount, highTripCount;
 };
+} // namespace detail
+
+/// A description of a multi-size tiling comprising tile sizes and numbers of
+/// tiles, expressed as Values which may or may not be constant. Multi-size
+/// currently means two-size.
+struct MultiSizeSpecification
+    : public detail::MultiSizeSpecificationBase<Value> {};
+struct StaticMultiSizeSpecification
+    : public detail::MultiSizeSpecificationBase<int64_t> {};
 
 /// Emits the IR computing the multi-sized tiling specification with two tile
 /// sizes not exceeding `targetSize`, each divisible by `sizeDivisor`, such
 /// that there exist numbers of tiles with these sizes that fully cover the
 /// given iteration space `dimension` of the structured `op`.
 ///
 /// The computation is as follows:
 ///

 /// assertion that size computation succeeded.
 ///
 /// Returns the specification consisting of both tile values and the number of
 /// tiles of each size.
 FailureOr<MultiSizeSpecification>
 computeMultiTileSizes(OpBuilder &builder, LinalgOp op, unsigned dimension,
                       OpFoldResult targetSize, OpFoldResult divisor,
                       bool emitAssertions = true);
+FailureOr<StaticMultiSizeSpecification>
+computeStaticMultiTileSizes(LinalgOp op, unsigned dimension, int64_t targetSize,
+                            int64_t divisor);
 
 /// Rewrite a TilingInterface `op` to a tiled `scf.foreach_thread`, applying
 /// tiling by `numThreads`.
 /// If non-empty, the `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
 /// thought of as tiling by the full size of data. It is the user's
 /// responsibility to ensure that `numThreads` is a valid tiling specification

mlir/include/mlir/Dialect/Transform/IR/TransformInterfaces.td

def TransformParamTypeInterface
     : TransformTypeInterfaceBase<"TransformParamTypeInterface",
                                  "::mlir::Attribute"> {
   let description = [{
     Types that can be used for the Transform dialect parameter values. Such types
     define the structure of the parameters associated with the value, e.g., their
     underlying type. A user of the value can assume that the parameter has been
     verified.
   }];
-
 }
 
+def Transform_AnyHandleOrParamType
+  : Type<Or<[TransformParamTypeInterface.predicate,
+             TransformHandleTypeInterface.predicate]>,
+         "any transform handle or parameter">;
+
 def FunctionalStyleTransformOpTrait
     : NativeOpTrait<"FunctionalStyleTransformOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }
 
 def TransformEachOpTrait : NativeOpTrait<"TransformEachOpTrait"> {
   let cppNamespace = "::mlir::transform";
 }

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

let description = [{
     that the associated payload IR ops are destroyed, so having the same op
     listed more than once will lead to double-free. Single-operand
     MergeHandlesOp may be used to deduplicate the associated list of payload IR
     ops when necessary. Furthermore, a combination of ReplicateOp and
     MergeHandlesOp can be used to construct arbitrary lists with repetitions.
   }];
 
   let arguments = (ins TransformHandleTypeInterface:$pattern,
-                       Variadic<TransformHandleTypeInterface>:$handles);
-  let results = (outs Variadic<TransformHandleTypeInterface>:$replicated);
+                       Variadic<Transform_AnyHandleOrParamType>:$handles);
+  let results = (outs Variadic<Transform_AnyHandleOrParamType>:$replicated);
   let assemblyFormat = "`num` `(` $pattern `)` $handles attr-dict `:` "
                        "type($pattern) `,` type($handles)";
 }
 
 def SequenceOp : TransformDialectOp<"sequence",
     [DeclareOpInterfaceMethods<RegionBranchOpInterface,
         ["getSuccessorEntryOperands", "getSuccessorRegions",
          "getRegionInvocationBounds"]>,

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

 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformUtils.h"
 #include "mlir/Dialect/Transform/Utils/Utils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
transform::MatchOp::apply(transform::TransformResults &results,
   results.set(getResult().cast<OpResult>(), res);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===---------------------------------------------------------------------===//
 // MultiTileSizesOp
 //===---------------------------------------------------------------------===//
 
+static void printMultitileSizesTypes(OpAsmPrinter &printer, Operation *op,
+                                     Type targetType, Type lowSizeType, Type,
+                                     Type) {
+  printer.printFunctionalType(TypeRange{targetType}, TypeRange{lowSizeType});
+}
+
+static ParseResult parseMultitileSizesTypes(OpAsmParser &parser,
+                                            Type &targetType, Type &lowSizeType,
+                                            Type &highSizeType,
+                                            Type &splitPointType) {
+  FunctionType funcType;
+  llvm::SMLoc typeLoc = parser.getCurrentLocation();
+  if (failed(parser.parseType<FunctionType>(funcType)))
+    return failure();
+
+  if (funcType.getNumInputs() != 1 || funcType.getNumResults() != 1) {
+    parser.emitError(typeLoc) << "expects a trailing functional type with one "
+                                 "argument and one result";
+  }
+  targetType = funcType.getInput(0);
+  lowSizeType = highSizeType = splitPointType = funcType.getResult(0);
+
+  return success();
+}
+
 DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
     LinalgOp target, transform::ApplyToEachResultList &results,
     TransformState &state) {
+  if (getLowSize().getType().isa<TransformParamTypeInterface>()) {
+    if (target.hasDynamicShape()) {
+      results.assign(
+          ArrayRef<Attribute>({Attribute(), Attribute(), Attribute()}));
+      auto diag = emitSilenceableError()
+                  << "cannot compute parametric tile sizes for dynamically "
+                     "shaped payload op";
+      diag.attachNote(target->getLoc()) << "payload op";
+      return diag;
+    }
+
+    FailureOr<StaticMultiSizeSpecification> spec = computeStaticMultiTileSizes(
+        target, getDimension(), getTargetSize(), getDivisor());
+    if (failed(spec)) {
+      results.assign(
+          ArrayRef<Attribute>({Attribute(), Attribute(), Attribute()}));
+      return emitSilenceableError()
+             << "failed to compute multi-size tiling sizes";
+    }
+
+    Builder builder(target.getContext());
+    results.assign(llvm::map_range(
+        ArrayRef<int64_t>({spec->lowTileSize, spec->highTileSize,
+                           spec->lowTileSize * spec->lowTripCount}),
+        [&builder, this](int64_t value) {
+          return builder.getIntegerAttr(
+              getLowSize().getType().cast<ParamType>().getType(), value);
+        }));
+    return DiagnosedSilenceableFailure::success();
+  }
+
   OpBuilder builder(target.getContext());
   builder.setInsertionPoint(target);
   OpFoldResult targetSize = builder.getIndexAttr(getTargetSize());
   OpFoldResult divisor = builder.getIndexAttr(getDivisor());
   FailureOr<MultiSizeSpecification> spec = computeMultiTileSizes(
       builder, target, getDimension(), targetSize, divisor);
   if (failed(spec)) {
     return emitSilenceableError() << "could not generate tile size computation";
DiagnosedSilenceableFailure transform::MultiTileSizesOp::applyToOne(
   results.push_back(splitPoint);
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::MultiTileSizesOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getTarget(), effects);
   producesHandle(getResults(), effects);
+  if (getLowSize().getType().isa<TransformParamTypeInterface>())
+    onlyReadsPayload(effects);
+  else
     modifiesPayload(effects);
 }
 
+LogicalResult transform::MultiTileSizesOp::verify() {
+  if (getLowSize().getType() != getHighSize().getType() ||
+      getLowSize().getType() != getSplitPoint().getType()) {
+    return emitOpError() << "expects all results type to be the same";
+  }
+  return success();
+}
+
 //===---------------------------------------------------------------------===//
 // PackOp
 //===---------------------------------------------------------------------===//
 
 SmallVector<OpFoldResult> transform::PackOp::getMixedPackedSizes() {
   Builder b(getContext());
   return getMixedValues(getStaticPackedSizes(), getPackedSizes(), b);
 }
DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
                                            TransformState &state) {
   // Collect the dynamic split points if provided.
   ArrayRef<Operation *> payload = state.getPayloadOps(getTarget());
   TrivialPatternRewriter rewriter(getContext());
   SmallVector<OpFoldResult> splitPoints;
   splitPoints.reserve(payload.size());
   if (getDynamicSplitPoint()) {
     auto diag = DiagnosedSilenceableFailure::success();
+    if (getDynamicSplitPoint().getType().isa<TransformHandleTypeInterface>()) {
       splitPoints = llvm::to_vector(llvm::map_range(
           state.getPayloadOps(getDynamicSplitPoint()), [&](Operation *op) {
             if (op->getNumResults() != 1 ||
                 !op->getResult(0).getType().isIndex()) {
               diag = emitSilenceableError()
                      << "expected dynamic split point handle to point to a "
                         "single-result index-typed op";
               diag.attachNote(op->getLoc()) << "dynamic split point";
             }
             return OpFoldResult(op->getResult(0));
           }));
+    } else {
+      splitPoints = llvm::to_vector(
+          llvm::map_range(state.getParams(getDynamicSplitPoint()),
+                          [](Attribute attr) { return OpFoldResult(attr); }));
+    }
     if (diag.isSilenceableFailure()) {
       results.set(getFirst().cast<OpResult>(), {});
       results.set(getSecond().cast<OpResult>(), {});
       return diag;
     }
 
     if (splitPoints.size() != payload.size()) {
       return emitDefiniteFailure()
if (getDynamicSplitPoint())
     onlyReadsHandle(getDynamicSplitPoint(), effects);
   producesHandle(getResults(), effects);
   modifiesPayload(effects);
 }
 
 ParseResult SplitOp::parse(OpAsmParser &parser, OperationState &result) {
   OpAsmParser::UnresolvedOperand target, dynamicSplitPoint;
   IntegerAttr staticSplitPoint;
-  auto pdlOperationType =
-      pdl::OperationType::get(parser.getBuilder().getContext());
-  if (parser.parseOperand(target) ||
-      parser.resolveOperand(target, pdlOperationType, result.operands) ||
-      parser.parseKeyword("after"))
+  if (parser.parseOperand(target) || parser.parseKeyword("after"))
     return failure();
 
   OptionalParseResult dynamicPointParseResult =
       parser.parseOptionalOperand(dynamicSplitPoint);
   if (!dynamicPointParseResult.has_value()) {
     int64_t staticSplitPointValue;
     if (failed(parser.parseInteger(staticSplitPointValue)))
       return failure();
 
     staticSplitPoint =
         parser.getBuilder().getI64IntegerAttr(staticSplitPointValue);
-  } else {
-    if (failed(*dynamicPointParseResult) ||
-        parser.resolveOperand(dynamicSplitPoint, pdlOperationType,
+  }
+
+  Type targetType;
+  if (parser.parseOptionalAttrDict(result.attributes) ||
+      parser.parseColonType(targetType) ||
+      parser.resolveOperand(target, targetType, result.operands)) {
+    return failure();
+  }
+  if (dynamicPointParseResult.has_value()) {
+    Type splitPointType;
+    if (failed(*dynamicPointParseResult) || parser.parseComma() ||
+        parser.parseType(splitPointType) ||
+        parser.resolveOperand(dynamicSplitPoint, splitPointType,
                               result.operands)) {
       return failure();
     }
 
     staticSplitPoint =
         parser.getBuilder().getI64IntegerAttr(ShapedType::kDynamic);
   }
 
   result.addAttribute(
       SplitOp::getStaticSplitPointAttrName(result.name).getValue(),
       staticSplitPoint);
-  if (failed(parser.parseOptionalAttrDict(result.attributes)))
-    return failure();
-
-  result.addTypes({pdlOperationType, pdlOperationType});
+  result.addTypes({targetType, targetType});
   return success();
 }
 
 void SplitOp::print(OpAsmPrinter &printer) {
   printer << " " << getTarget() << " after ";
   int64_t staticSplitSize = static_cast<int64_t>(getStaticSplitPoint());
   if (staticSplitSize != ShapedType::kDynamic)
     printer << staticSplitSize;
   else
     printer << getDynamicSplitPoint();
   printer << " ";
   printer.printOptionalAttrDict(getOperation()->getAttrs(),
                                 {getStaticSplitPointAttrName()});
+  printer << " : " << getTarget().getType();
+  if (staticSplitSize == ShapedType::kDynamic)
+    printer << ", " << getDynamicSplitPoint().getType();
 }
 
 LogicalResult SplitOp::verify() {
   if ((static_cast<int64_t>(getStaticSplitPoint()) != ShapedType::kDynamic) ^
       (getDynamicSplitPoint() == nullptr)) {
     return emitOpError() << "expects either a dynamic or a static split "
                             "point to be provided";
   }
transform::TileReductionUsingForeachThreadOp::applyToOne(
   results.push_back(result->parallelTiledOp);
   results.push_back(result->mergeOp);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
+
 void transform::TileOp::build(OpBuilder &builder, OperationState &result,
-                              Value target, ArrayRef<int64_t> staticTileSizes,
+                              TypeRange loopTypes, Value target,
+                              ArrayRef<int64_t> staticTileSizes,
                               ArrayRef<int64_t> interchange) {
-  return build(builder, result,
+  return build(builder, result, loopTypes,
                /*target=*/target,
                /*mixedTileSizes=*/
                getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
                interchange);
 }
 
 void transform::TileOp::build(OpBuilder &builder, OperationState &result,
+                              Value target, ArrayRef<int64_t> staticTileSizes,
+                              ArrayRef<int64_t> interchange) {
+  build(builder, result, target,
+        getAsOpFoldResult(builder.getI64ArrayAttr(staticTileSizes)),
+        interchange);
+}
+
+void transform::TileOp::build(OpBuilder &builder, OperationState &result,
                               Value target,
                               ArrayRef<OpFoldResult> mixedTileSizes,
                               ArrayRef<int64_t> interchange) {
+  // Loop types are automaticaly splat by the callee, setting up one is enough.
+  SmallVector<Type> loopTypes(1, builder.getType<transform::AnyOpType>());
+  build(builder, result, loopTypes, target, mixedTileSizes, interchange);
+}
+
+void transform::TileOp::build(OpBuilder &builder, OperationState &result,
+                              TypeRange loopTypes, Value target,
+                              ArrayRef<OpFoldResult> mixedTileSizes,
+                              ArrayRef<int64_t> interchange) {
   SmallVector<int64_t> staticTileSizes;
   SmallVector<Value> dynamicTileSizes;
   dispatchIndexOpFoldResults(mixedTileSizes, dynamicTileSizes, staticTileSizes);
   // Call the default builder which sets up the proper operands segment sizes
   // attributes for multiple variadic operands. In the absence of this, horrible
   // bugs ensue.
-  MLIRContext *ctx = builder.getContext();
-  auto operationType = pdl::OperationType::get(ctx);
   auto staticTileSizesAttr = builder.getDenseI64ArrayAttr(staticTileSizes);
-  build(builder, result,
-        /*resultTypes=*/TypeRange{operationType, operationType},
+  unsigned numExpectedLoops =
+      staticTileSizes.size() - llvm::count(staticTileSizes, 0);
+  SmallVector<Type> resultTypes;
+  resultTypes.reserve(numExpectedLoops);
+  assert((loopTypes.size() == 1 || loopTypes.size() == numExpectedLoops) &&
+         "expected one loop type or as many as loops");
+  if (loopTypes.size() == 1)
+    resultTypes.append(numExpectedLoops, loopTypes[0]);
+  else
+    llvm::append_range(resultTypes, loopTypes);
+  build(builder, result, /*tiled_linalg_op=*/target.getType(),
+        /*loops=*/resultTypes,
         /*target=*/target,
         /*dynamic_sizes=*/dynamicTileSizes,
         /*static_sizes=*/staticTileSizesAttr,
         /*interchange=*/builder.getDenseI64ArrayAttr(interchange));
 }
 
 DiagnosedSilenceableFailure
 transform::TileOp::apply(TransformResults &transformResults,
                          TransformState &state) {
   ArrayRef<int64_t> tileSizes = getStaticSizes();
 
   ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
   SmallVector<ArrayRef<Operation *>> dynamicSizeProducers;
+  SmallVector<SmallVector<int64_t>> paramSizes;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
-  for (Value dynamicSizeProducerHandle : getDynamicSizes()) {
-    dynamicSizeProducers.push_back(
-        state.getPayloadOps(dynamicSizeProducerHandle));
+  paramSizes.reserve(getDynamicSizes().size());
+  for (Value transformValue : getDynamicSizes()) {
+    if (transformValue.getType().isa<ParamType>()) {
+      dynamicSizeProducers.push_back({});
+      ArrayRef<Attribute> params = state.getParams(transformValue);
+      paramSizes.push_back(
+          llvm::to_vector(llvm::map_range(params, [](Attribute attr) {
+            return attr.cast<IntegerAttr>().getValue().getSExtValue();
+          })));
+
+      if (paramSizes.back().size() != targets.size()) {
+        for (OpResult r : getResults())
+          transformResults.set(r, {});
+        DiagnosedSilenceableFailure diag =
+            emitSilenceableError()
+            << "expected as many parameter values ("
+            << dynamicSizeProducers.back().size() << ") as target ops ("
+            << targets.size() << ")";
+        diag.attachNote(transformValue.getLoc()) << "for this parameter";
+        return diag;
+      }
+
+      continue;
+    }
+    paramSizes.push_back({});
+    dynamicSizeProducers.push_back(state.getPayloadOps(transformValue));
 
     if (dynamicSizeProducers.back().size() != targets.size()) {
+      for (OpResult r : getResults())
+        transformResults.set(r, {});
       DiagnosedSilenceableFailure diag =
           emitSilenceableError()
           << "expected as many dynamic size-producing operations ("
           << dynamicSizeProducers.back().size() << ") as target ops ("
           << targets.size() << ")";
-      diag.attachNote(dynamicSizeProducerHandle.getLoc()) << "for this handle";
+      diag.attachNote(transformValue.getLoc()) << "for this handle";
       return diag;
     }
 
     for (Operation *op : dynamicSizeProducers.back()) {
       if (op->getNumResults() == 1 &&
           op->getResult(0).getType().isa<IndexType>())
         continue;
+
+      for (OpResult r : getResults())
+        transformResults.set(r, {});
       DiagnosedSilenceableFailure diag =
           emitSilenceableError() << "expected sizes to be produced by ops "
                                     "with a single index-type result";
       diag.attachNote(op->getLoc()) << "size producer op";
-      diag.attachNote(dynamicSizeProducerHandle.getLoc()) << "for this handle";
+      diag.attachNote(transformValue.getLoc()) << "for this handle";
       return diag;
     }
   }
 
   SmallVector<Operation *> tiled;
   SmallVector<SmallVector<Operation *, 4>, 4> loops;
   loops.resize(getLoops().size());
   for (auto &en : llvm::enumerate(targets)) {
if (!tileSizes.empty()) {
           [&, index](OpBuilder &b, Operation *) {
             SmallVector<Value, 4> sizes;
             sizes.reserve(tileSizes.size());
             unsigned dynamicIdx = 0;
             for (OpFoldResult ofr : getMixedSizes()) {
               if (auto attr = ofr.dyn_cast<Attribute>()) {
                 sizes.push_back(b.create<arith::ConstantIndexOp>(
                     getLoc(), attr.cast<IntegerAttr>().getInt()));
-              } else {
+                continue;
+              }
+              ArrayRef<Operation *> dynamicSizes =
+                  dynamicSizeProducers[dynamicIdx];
+              ArrayRef<int64_t> params = paramSizes[dynamicIdx];
+              ++dynamicIdx;
+              assert((dynamicSizes.empty() ^ params.empty()) &&
+                     "expected either dynamic sizes or parameters");
+              if (!params.empty()) {
                 sizes.push_back(
-                    dynamicSizeProducers[dynamicIdx++][index]->getResult(0));
+                    b.create<arith::ConstantIndexOp>(getLoc(), params[index]));
+              } else {
+                sizes.push_back(dynamicSizes[index]->getResult(0));
               }
             }
             return sizes;
           });
     }
 
     tilingOptions.setInterchange(getInterchange());
     TrivialPatternRewriter rewriter(linalgOp.getContext());
void printOptionalInterchange(OpAsmPrinter &p,
   }
 }
 
 ParseResult transform::TileOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
   OpAsmParser::UnresolvedOperand target;
   SmallVector<OpAsmParser::UnresolvedOperand> dynamicSizes;
   DenseI64ArrayAttr staticSizes;
-  auto pdlOperationType = pdl::OperationType::get(parser.getContext());
-  if (parser.parseOperand(target) ||
-      parser.resolveOperand(target, pdlOperationType, result.operands) ||
+  FunctionType functionalType;
+  llvm::SMLoc operandLoc;
+  if (parser.parseOperand(target) || parser.getCurrentLocation(&operandLoc) ||
       parseDynamicIndexList(parser, dynamicSizes, staticSizes) ||
-      parser.resolveOperands(dynamicSizes, pdlOperationType, result.operands))
+      parseOptionalInterchange(parser, result) ||
+      parser.parseColonType(functionalType))
     return ParseResult::failure();
 
-  // Parse optional interchange.
-  if (failed(parseOptionalInterchange(parser, result)))
-    return ParseResult::failure();
-  result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
   size_t numExpectedLoops =
       staticSizes.size() - llvm::count(staticSizes.asArrayRef(), 0);
-  result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOperationType));
+  if (functionalType.getNumResults() != numExpectedLoops + 1) {
+    return parser.emitError(parser.getNameLoc())
+           << "expected " << (numExpectedLoops + 1) << " result type(s)";
+  }
+  if (functionalType.getNumInputs() != dynamicSizes.size() + 1) {
+    return parser.emitError(operandLoc)
+           << "expected " << dynamicSizes.size() + 1 << " operand type(s)";
+  }
+  if (parser.resolveOperand(target, functionalType.getInputs().front(),
+                            result.operands) ||
+      parser.resolveOperands(dynamicSizes,
+                             functionalType.getInputs().drop_front(),
+                             operandLoc, result.operands)) {
+    return failure();
+  }
+
+  result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
+  result.addTypes(functionalType.getResults());
   return success();
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ' << getTarget();
   printDynamicIndexList(p, getOperation(), getDynamicSizes(), getStaticSizes());
   printOptionalInterchange(p, getInterchange());
+  p << " : ";
+  p.printFunctionalType(getOperands().getTypes(), getResults().getTypes());
 }
 
 void transform::TileOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getDynamicSizes(), effects);
   producesHandle(getTiledLinalgOp(), effects);
   producesHandle(getLoops(), effects);

mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp

static void emitIsPositiveIndexAssertion(ImplicitLocOpBuilder &b,
   Value zero = b.create<arith::ConstantIndexOp>(0);
   Value condition = b.create<arith::CmpIOp>(arith::CmpIPredicate::sgt,
                                             value.get<Value>(), zero);
   b.create<cf::AssertOp>(
       condition,
       b.getStringAttr("expected strictly positive tile size and divisor"));
 }
 
+FailureOr<StaticMultiSizeSpecification>
+mlir::linalg::computeStaticMultiTileSizes(LinalgOp op, unsigned dimension,
+                                          int64_t targetSize, int64_t divisor) {
+  assert(!op.hasDynamicShape() &&
+         "cannot compute static multi-tile sizes for an op with dynamic shape");
+  assert(targetSize > 0 && "target size must be non-negative");
+  assert(divisor > 0 && "divisor must be non-negative");
+  assert(dimension < op.getNumLoops() && "dimension overflow");
+
+  StaticMultiSizeSpecification spec;
+  int64_t tripCount = op.getStaticLoopRanges()[dimension];
+  int64_t a = tripCount / divisor;
+  int64_t t = (targetSize + divisor - 1) / divisor;
+  int64_t totalTripCount = (a + t - 1) / t;
+  spec.lowTileSize = (a / totalTripCount) * divisor;
+  spec.highTileSize = spec.lowTileSize + divisor;
+  spec.highTripCount = a % totalTripCount;
+  spec.lowTripCount = totalTripCount - spec.highTripCount;
+  if (spec.lowTileSize * spec.lowTripCount +
+          spec.highTileSize * spec.highTripCount !=
+      tripCount) {
+    return failure();
+  }
+  return spec;
+}
+
 FailureOr<MultiSizeSpecification>
 mlir::linalg::computeMultiTileSizes(OpBuilder &builder, LinalgOp op,
                                     unsigned dimension, OpFoldResult targetSize,
                                     OpFoldResult divisor, bool emitAssertions) {
   // Bail out on dimension overflow.
   if (dimension >= op.getNumLoops())
     return failure();
 

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ReplicateOp::apply(transform::TransformResults &results,
                               transform::TransformState &state) {
   unsigned numRepetitions = state.getPayloadOps(getPattern()).size();
   for (const auto &en : llvm::enumerate(getHandles())) {
     Value handle = en.value();
+    if (handle.getType().isa<TransformHandleTypeInterface>()) {
       ArrayRef<Operation *> current = state.getPayloadOps(handle);
       SmallVector<Operation *> payload;
       payload.reserve(numRepetitions * current.size());
       for (unsigned i = 0; i < numRepetitions; ++i)
         llvm::append_range(payload, current);
       results.set(getReplicated()[en.index()].cast<OpResult>(), payload);
+    } else {
+      assert(handle.getType().isa<TransformParamTypeInterface>() &&
+             "expected param type");
+      ArrayRef<Attribute> current = state.getParams(handle);
+      SmallVector<Attribute> params;
+      params.reserve(numRepetitions * current.size());
+      for (unsigned i = 0; i < numRepetitions; ++i)
+        llvm::append_range(params, current);
+      results.setParams(getReplicated()[en.index()].cast<OpResult>(), params);
+    }
   }
   return DiagnosedSilenceableFailure::success();
 }
 
 void transform::ReplicateOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   onlyReadsHandle(getPattern(), effects);
   consumesHandle(getHandles(), effects);

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from ..ir import *
   from ._ods_common import get_op_result_or_value as _get_op_result_or_value
-  from ..dialects import pdl
+  from ..dialects import pdl, transform
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
-from typing import List, Optional, Sequence, Union
+from typing import List, Optional, Sequence, Union, overload
 
 IntOrAttrList = Sequence[Union[IntegerAttr, int]]
 OptionalIntList = Optional[Union[ArrayAttr, IntOrAttrList]]
 
 
 def _get_int64_attr(value: Union[int, Attribute]) -> IntegerAttr:
   if isinstance(value, int):
     return IntegerAttr.get(IntegerType.get_signless(64), value)
if values is None:
     return ArrayAttr.get([])
   if isinstance(values, ArrayAttr):
     return values
 
   return ArrayAttr.get([_get_int64_attr(v) for v in values])
 
 def _get_dense_int64_array_attr(
         values: Sequence[int]) -> DenseI64ArrayAttr:
-    """Creates a dense integer array from a sequence of integers.
     Expects the thread-local MLIR context to have been set by the context 
     manager.
     """
-    if values is None:
-        return DenseI64ArrayAttr.get([])
-    return DenseI64ArrayAttr.get(values)
 
 def _get_int_int_array_attr(
     values: Optional[Union[ArrayAttr, Sequence[Union[ArrayAttr,
                                                      IntOrAttrList]]]]
 ) -> ArrayAttr:
   """Creates an array attribute containing array attributes of integers.
 
   If the operand is already an array attribute, forwards it. Otherwise treats
return MatchOp(
         loc=loc,
         ip=ip)
 
 
 class MultiTileSizesOp:
   """Specialization for MultitileSizesOp class."""
 
   def __init__(self,
+               result_type: Type,
                target: Union[Operation, Value],
                *,
                dimension: Union[int, IntegerAttr],
                target_size: Union[int, IntegerAttr],
                divisor: Optional[Union[int, IntegerAttr]] = None,
                loc=None,
                ip=None):
     super().__init__(
-        pdl.OperationType.get(),
-        pdl.OperationType.get(),
-        pdl.OperationType.get(),
+        result_type,
+        result_type,
+        result_type,
         _get_op_result_or_value(target),
         dimension=_get_int64_attr(dimension),
         target_size=_get_int64_attr(target_size),
         divisor=_get_int64_attr(divisor if divisor else 1),
         loc=loc,
         ip=ip)
 
 
def __init__(self,
 
     if isinstance(split_point, Attribute):
       static_split_point = split_point
       dynamic_split_point = None
     else:
       static_split_point = _get_int64_attr(ShapedType.get_dynamic_size())
       dynamic_split_point = _get_op_result_or_value(split_point)
 
-    pdl_operation_type = pdl.OperationType.get()
+    target = _get_op_result_or_value(target)
+
     super().__init__(
-        pdl_operation_type,
-        pdl_operation_type,
-        _get_op_result_or_value(target),
+        target.type,
+        target.type,
+        target,
         dimension=dimension,
         static_split_point=static_split_point,
         dynamic_split_point=dynamic_split_point,
         loc=loc,
         ip=ip)
 
 
 class TileOp:
   """Specialization for TileOp class."""
 
+  @overload
   def __init__(self,
+               loop_types: Union[Type, List[Type]],
                target: Union[Operation, Value],
                *,
                sizes: Optional[Union[Sequence[Union[int, IntegerAttr, Operation,
                                                     Value]], ArrayAttr]] = None,
                interchange: OptionalIntList = None,
                loc=None,
                ip=None):
-    pdl_operation_type = pdl.OperationType.get()
-    i64_type = IntegerType.get_signless(64)
+    ...
+
+  @overload
+  def __init__(self,
+               target: Union[Operation, Value],
+               *,
+               sizes: Optional[Union[Sequence[Union[int, IntegerAttr, Operation,
+                                                    Value]], ArrayAttr]] = None,
+               interchange: OptionalIntList = None,
+               loc=None,
+               ip=None):
+    ...
 
+  def __init__(self,
+               loop_types_or_target: Union[Type, List[Type], Operation, Value],
+               target_or_none: Optional[Union[Operation, Value]] = None,
+               *,
+               sizes: Optional[Union[Sequence[Union[int, IntegerAttr, Operation,
+                                                    Value]], ArrayAttr]] = None,
+               interchange: OptionalIntList = None,
+               loc=None,
+               ip=None):
     if sizes is None:
       sizes = []
 
     static_sizes = []
     dynamic_sizes = []
     if isinstance(sizes, ArrayAttr):
       sizes_attr = sizes
     else:
       for size in sizes:
         if isinstance(size, int):
           static_sizes.append(size)
         else:
           static_sizes.append(ShapedType.get_dynamic_size())
           dynamic_sizes.append(_get_op_result_or_value(size))
       sizes_attr = DenseI64ArrayAttr.get(static_sizes)
 
     num_loops = sum(
         v if v == 0 else 1 for v in self.__extract_values(sizes_attr))
+
+    if isinstance(loop_types_or_target, (Operation, Value)):
+      loop_types = [transform.AnyOpType.get()] * num_loops
+      target = loop_types_or_target
+      assert target_or_none is None, "Cannot construct TileOp with two targets."
+    else:
+      loop_types = ([loop_types_or_target] * num_loops) if isinstance(
+          loop_types_or_target, Type) else loop_types_or_target
+      target = target_or_none
+
+    target = _get_op_result_or_value(target)
+
     super().__init__(
-        pdl_operation_type, [pdl_operation_type] * num_loops,
-        _get_op_result_or_value(target),
+        target.type,
+        loop_types,
+        target,
         dynamic_sizes=dynamic_sizes,
         static_sizes=sizes_attr,
-        interchange=_get_dense_int64_array_attr(interchange) if interchange else None,
+        interchange=_get_dense_int64_array_attr(interchange)
+        if interchange else None,
         loc=loc,
         ip=ip)
 
   def __extract_values(self, attr: Optional[DenseI64ArrayAttr]) -> List[int]:
     if not attr:
       return []
     return [element for element in attr]
 

mlir/test/Dialect/LLVM/transform-e2e.mlir

%0 = linalg.matmul  ins(%arg0, %arg1: tensor<2x4xf32>, tensor<4x6xf32>)
                      outs(%arg2: tensor<2x6xf32>)
     -> tensor<2x6xf32>
   return %0 : tensor<2x6xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op
-  %1, %loops:3 = transform.structured.tile %0 [2, 2, 2]
+  %1, %loops:3 = transform.structured.tile %0 [2, 2, 2] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
   transform.bufferization.one_shot_bufferize layout{IdentityLayoutMap} %module_op
     {bufferize_function_boundaries = true}
   %func = transform.structured.match ops{["func.func"]} in %module_op
   transform.vector.lower_vectors %func multireduction_lowering = "innerreduction"
 }

mlir/test/Dialect/Linalg/multisize-tiling-full.mlir

-// RUN: mlir-opt --test-transform-dialect-interpreter --scf-for-loop-canonicalization --canonicalize %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --scf-for-loop-canonicalization --canonicalize --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s --check-prefix=NOCANON
 
 // This implements a 2D multisize tiling with target sizes [3, 10].
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-  %1:3 = transform.structured.multitile_sizes %0 { dimension = 0, target_size = 3}
-  %t:3 = transform.structured.multitile_sizes %0 { dimension = 1, target_size = 10}
-  %2:2 = transform.structured.split %0 after %1#2 { dimension = 0 }
-  %3:2 = transform.structured.tile %2#0 [%1#0]
-  %4:2 = transform.structured.tile %2#1 [%1#1]
+  %1:3 = transform.structured.multitile_sizes %0 { dimension = 0, target_size = 3} : (!pdl.operation) -> !pdl.operation
+  %t:3 = transform.structured.multitile_sizes %0 { dimension = 1, target_size = 10} : (!pdl.operation) -> !pdl.operation
+  %2:2 = transform.structured.split %0 after %1#2 { dimension = 0 } : !pdl.operation, !pdl.operation
+  %3:2 = transform.structured.tile %2#0 [%1#0] : (!pdl.operation, !pdl.operation) -> (!pdl.operation, !pdl.operation)
+  %4:2 = transform.structured.tile %2#1 [%1#1] : (!pdl.operation, !pdl.operation) -> (!pdl.operation, !pdl.operation)
   %5 = merge_handles %3#0, %4#0 : !pdl.operation
   %tt:3 = replicate num(%5) %t#0, %t#1, %t#2 : !pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation
-  %6:2 = transform.structured.split %5 after %tt#2 { dimension = 1 }
-  transform.structured.tile %6#0 [0, %tt#0]
-  transform.structured.tile %6#1 [0, %tt#1]
+  %6:2 = transform.structured.split %5 after %tt#2 { dimension = 1 } : !pdl.operation, !pdl.operation
+  transform.structured.tile %6#0 [0, %tt#0] : (!pdl.operation, !pdl.operation) -> (!pdl.operation, !pdl.operation)
+  transform.structured.tile %6#1 [0, %tt#1] : (!pdl.operation, !pdl.operation) -> (!pdl.operation, !pdl.operation)
 }
 
 func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
 
+// Without canonicalization, tile sizes are computed dynamically as affine maps.
+// NOCANON-LABEL: @two_d
+// NOCANON-COUNT-8: affine.apply
+// NOCANON:         scf.for
+
 // CHECK-LABEL: @two_d
 // CHECK-SAME: %[[IN:.+]]: tensor<10x34xf32>, %[[OUT:.+]]: tensor<10x34xf32>
 func.func @two_d(%arg0: tensor<10x34xf32>,
                  %arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
   %0 = linalg.generic {
     indexing_maps = [affine_map<(i, j) -> (i, j)>,
                      affine_map<(i, j) -> (i, j)>],
     iterator_types = ["parallel", "parallel"]
func.func @two_d(%arg0: tensor<10x34xf32>,
   // CHECK:            scf.yield
   // CHECK-COUNT-2:  tensor.insert_slice
   // CHECK:          scf.yield
   // CHECK:        %[[RESULT:.+]] = tensor.insert_slice
   // CHECK:        return %[[RESULT]]
 
   return %0 : tensor<10x34xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
+  %1:3 = transform.structured.multitile_sizes %0 { dimension = 0, target_size = 3} : (!pdl.operation) -> !transform.param<i64>
+  %t:3 = transform.structured.multitile_sizes %0 { dimension = 1, target_size = 10} : (!pdl.operation) -> !transform.param<i64>
+  %2:2 = transform.structured.split %0 after %1#2 { dimension = 0 } : !pdl.operation, !transform.param<i64>
+  %3:2 = transform.structured.tile %2#0 [%1#0] : (!pdl.operation, !transform.param<i64>) -> (!pdl.operation, !pdl.operation)
+  %4:2 = transform.structured.tile %2#1 [%1#1] : (!pdl.operation, !transform.param<i64>) -> (!pdl.operation, !pdl.operation)
+  %5 = merge_handles %3#0, %4#0 : !pdl.operation
+  %tt:3 = replicate num(%5) %t#0, %t#1, %t#2 : !pdl.operation, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>
+  %6:2 = transform.structured.split %5 after %tt#2 { dimension = 1 } : !pdl.operation, !transform.param<i64>
+  transform.structured.tile %6#0 [0, %tt#0] : (!pdl.operation, !transform.param<i64>) -> (!pdl.operation, !pdl.operation)
+  transform.structured.tile %6#1 [0, %tt#1] : (!pdl.operation, !transform.param<i64>) -> (!pdl.operation, !pdl.operation)
+}
+
+func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
+
+// Even without canonicalization, tile sizes can be computed statically thanks
+// to parameters.
+// NOCANON-LABEL: @two_d
+// NOCANON-NOT:   affine.apply
+// NOCANON:       scf.for
+
+// CHECK-LABEL: @two_d_param
+// CHECK-SAME: %[[IN:.+]]: tensor<10x34xf32>, %[[OUT:.+]]: tensor<10x34xf32>
+func.func @two_d_param(%arg0: tensor<10x34xf32>,
+                       %arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
+  %0 = linalg.generic {
+    indexing_maps = [affine_map<(i, j) -> (i, j)>,
+                     affine_map<(i, j) -> (i, j)>],
+    iterator_types = ["parallel", "parallel"]
+  }
+  ins(%arg0: tensor<10x34xf32>)
+  outs(%arg1: tensor<10x34xf32>) {
+  ^bb0(%0: f32, %1: f32):
+    %i = linalg.index 0 : index
+    %j = linalg.index 1 : index
+    %call_res = func.call @elem(%0, %i, %j) : (f32, index, index) -> f32
+    linalg.yield %call_res : f32
+  } -> tensor<10x34xf32>
+
+  // CHECK:      %[[SLICE_1_IN:.+]] = tensor.extract_slice %[[IN]][0, 0] [4, 34] [1, 1]
+  // CHECK:      %[[SLICE_1:.+]] = tensor.extract_slice %[[OUT]][0, 0] [4, 34] [1, 1]
+  // CHECK:      scf.for %[[I1:.+]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ITERARG_1:.+]] = %[[SLICE_1]])
+  // CHECK:        %[[OUTSLICE_1_IN:.+]] = tensor.extract_slice %[[SLICE_1_IN]][%[[I1]], 0] [2, 34] [1, 1]
+  // CHECK:        %[[OUTSLICE_1:.+]] = tensor.extract_slice %[[ITERARG_1]][%[[I1]], 0] [2, 34] [1, 1]
+
+  // CHECK:        %[[SLICE_2_IN:.+]] = tensor.extract_slice %[[OUTSLICE_1_IN]][0, 0] [2, 16] [1, 1]
+  // CHECK:        %[[SLICE_2:.+]] = tensor.extract_slice %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1]
+  // CHECK:        %[[LOOPRES:.+]] = scf.for %[[I2:.+]] = %{{.*}} to %{{.*}} step %{{.*}} iter_args(%[[ITERARG_2:.+]] = %[[SLICE_2]])
+  // CHECK:          %[[INSLICE_2:.+]] = tensor.extract_slice %[[SLICE_2_IN]][0, %[[I2]]] [2, 8] [1, 1]
+  // CHECK:          %[[OUTSLICE_2:.+]] = tensor.extract_slice %[[ITERARG_2]][0, %[[I2]]] [2, 8] [1, 1]
+  // CHECK:          %[[RESSLICE_1:.+]] = linalg.generic {{.*}} ins(%[[INSLICE_2]] : tensor<2x8xf32>) outs(%[[OUTSLICE_2]] : tensor<2x8xf32>)
+  // CHECK:          %[[RESPARTIAL:.+]] = tensor.insert_slice %[[RESSLICE_1]] into %[[ITERARG_2]]
+  // CHECK:          scf.yield %[[RESPARTIAL]]
+
+  // CHECK:        %[[INSERTED:.+]] = tensor.insert_slice %[[LOOPRES]] into %[[OUTSLICE_1]][0, 0] [2, 16] [1, 1]
+  // CHECK:        %[[OUTSLICE_3:.+]] = tensor.extract_slice %[[INSERTED]][0, 16] [2, 18] [1, 1]
+  // CHECK:        scf.for %{{.*}} iter_args(%{{.*}} = %[[OUTSLICE_3]])
+  // CHECK-COUNT-2:  tensor.extract_slice
+  // CHECK:          linalg.generic {{.*}} ins(%{{.*}} : tensor<2x9xf32>)
+  // CHECK:          tensor.insert_slice
+  // CHECK:          scf.yield
+  // CHECK:        %[[INSERTED_2:.+]] = tensor.insert_slice %{{.*}} into %[[INSERTED]]
+  // CHECK:        %[[INSERTED_3:.+]] = tensor.insert_slice %[[INSERTED_2]] into %[[ITERARG_1]]
+  // CHECK:        scf.yield %[[INSERTED_3]]
+
+  // CHECK:        tensor.insert_slice
+  // CHECK:        tensor.extract_slice
+  // CHECK:        scf.for
+  // CHECK-COUNT-2:  tensor.extract_slice
+  // CHECK:          scf.for
+  // CHECK-COUNT-2:    tensor.extract_slice
+  // CHECK:            linalg.generic {{.*}} ins(%{{.*}} : tensor<3x8xf32>)
+  // CHECK:            tensor.insert_slice
+  // CHECK:            scf.yield
+  // CHECK:          tensor.insert_slice
+  // CHECK:          tensor.extract_slice
+  // CHECK:          scf.for
+  // CHECK-COUNT-2:    tensor.extract_slice
+  // CHECK:            linalg.generic {{.*}} ins(%{{.*}} : tensor<3x9xf32>)
+  // CHECK:            tensor.insert_slice
+  // CHECK:            scf.yield
+  // CHECK-COUNT-2:  tensor.insert_slice
+  // CHECK:          scf.yield
+  // CHECK:        %[[RESULT:.+]] = tensor.insert_slice
+  // CHECK:        return %[[RESULT]]
+
+  return %0 : tensor<10x34xf32>
+}

mlir/test/Dialect/Linalg/promotion_options.mlir

 //      CHECK:       linalg.matmul ins(%[[VA]], %[[svB]]{{.*}} outs(%[[VC]]
 //      CHECK:       memref.copy %[[svCC]], %[[svC]]
 //      CHECK:       memref.dealloc %[[tmpA]]
 //      CHECK:       memref.dealloc %[[tmpC]]
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-  %1, %loops:3 = transform.structured.tile %0 [16, 16, 16]
+  %1, %loops:3 = transform.structured.tile %0 [16, 16, 16] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
   %2 = transform.structured.promote %1 { operands_to_promote = [0, 2], force_full_tiles = [false, false], use_full_tiles_by_default }
 }

mlir/test/Dialect/Linalg/tile-conv.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize | FileCheck %s
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 2)>
 //  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0, 3)>
 //  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0] -> (d0 + s0 - 1)>
 
 func.func @conv(%arg0 : memref<?x?xf32>, %arg1 : memref<?x?xf32>, %arg2 : memref<?x?xf32>) {
   linalg.conv_2d ins(%arg0, %arg1 : memref<?x?xf32>, memref<?x?xf32>) outs(%arg2 : memref<?x?xf32>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.conv_2d"]} in %arg1
-    %1, %loop:2 = transform.structured.tile %0 [2, 3]
+    %1, %loop:2 = transform.structured.tile %0 [2, 3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 //       CHECK: func @conv
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?xf32>
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index

mlir/test/Dialect/Linalg/tile-indexed.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -canonicalize -split-input-file | FileCheck %s -check-prefix=TILE-10n25
 
 func.func @indexed_vector(%arg0: memref<50xindex>) {
   linalg.generic {indexing_maps = [affine_map<(i) -> (i)>],
                   iterator_types = ["parallel"]}
      outs(%arg0 : memref<50xindex>) {
     ^bb0(%a: index):
       %i = linalg.index 0 : index
       linalg.yield %i : index
   }
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-    %1, %loop = transform.structured.tile %0 [10]
+    %1, %loop = transform.structured.tile %0 [10] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
 }
 
 // TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
 // TILE-10n25-LABEL: func @indexed_vector
 // TILE-10n25: %[[C10:.*]] = arith.constant 10 : index
 // TILE-10n25: scf.for %[[J:.*]] = {{.*}} step %[[C10]]
 // TILE-10n25:   linalg.generic
 // TILE-10n25:     %[[I:.*]] = linalg.index 0 : index
^bb0(%a: index):
       linalg.yield %sum : index
   }
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-    %1, %loop:2 = transform.structured.tile %0 [10, 25]
+    %1, %loop:2 = transform.structured.tile %0 [10, 25] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // TILE-10n25-DAG: [[$MAP:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0 + d1)>
 // TILE-10n25-LABEL: func @indexed_matrix
 // TILE-10n25-DAG: %[[C25:.*]] = arith.constant 25 : index
 // TILE-10n25-DAG: %[[C10:.*]] = arith.constant 10 : index
 // TILE-10n25: scf.for %[[K:.*]] = {{.*}} step %[[C10]]
 // TILE-10n25:   scf.for %[[L:.*]] = {{.*}} step %[[C25]]
 // TILE-10n25:     linalg.generic
 // TILE-10n25:       %[[I:.*]] = linalg.index 0 : index
 // TILE-10n25:       %[[NEW_I:.*]] = affine.apply [[$MAP]](%[[I]], %[[K]])
 // TILE-10n25:       %[[J:.*]] = linalg.index 1 : index
 // TILE-10n25:       %[[NEW_J:.*]] = affine.apply [[$MAP]](%[[J]], %[[L]])
 // TILE-10n25:       %[[SUM:.*]] = arith.addi %[[NEW_I]], %[[NEW_J]] : index
 // TILE-10n25:       linalg.yield %[[SUM]] : index

mlir/test/Dialect/Linalg/tile-tensors.mlir

 
 //      CHECK: return %[[TD0]] : tensor<?x?xf32>
   return %0 : tensor<?x?xf32>
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4]
+    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // -----
 
 func.func @generic_op_tensors(
   %arg0 : tensor<?x?x?xf32>, %arg1 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
^bb0(%arg2 : f32, %arg3: f32, %arg4: f32):
       linalg.yield %5 : f32
     } -> tensor<?x?x?xf32>
   return %4 : tensor<?x?x?xf32>
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4]
+    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @generic_op_tensors
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>
 //       CHECK:   %[[INIT:.+]] = tensor.empty
 //       CHECK:   %[[TD0:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC0:.+]] = %[[INIT]]) -> (tensor<?x?x?xf32>) {
 //       CHECK:     %[[TD1:.+]] = scf.for %{{.+}} to %{{.+}} step %{{.+}} iter_args(%[[TC1:.+]] = %[[TC0]]) -> (tensor<?x?x?xf32>) {
^bb0(%arg3 : f32, %arg4: f32, %arg5: f32):
       linalg.yield %5 : f32
     } -> tensor<?x42xf32>
   return %2 : tensor<?x42xf32>
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4]
+    %1, %loops:3 = transform.structured.tile %0 [2, 3, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }

mlir/test/Dialect/Linalg/transform-op-fuse.mlir

//       CHECK: return %[[RES]]
   } -> tensor<12x25xf32>
   func.return %0 : tensor<12x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   %1, %loops:2 = transform.structured.fuse %0 {tile_sizes = [5, 0, 7], tile_interchange = [0, 2, 1]}
-  %2, %loops_2 = transform.structured.tile %1 [0, 4]
+  %2, %loops_2 = transform.structured.tile %1 [0, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
 }
 
 // -----
 
 // CHECK-LABEL: func.func @unpack_elemwise
 // CHECK:         %[[RES:.*]] = scf.for
 // CHECK:           scf.for
 // CHECK:             tensor.unpack

mlir/test/Dialect/Linalg/transform-op-multitile-sizes.mlir

-// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file | FileCheck %s
+// RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file --verify-diagnostics | FileCheck %s
 
 // CHECK-DAG: #[[$MAP13:.+]] = affine_map<() -> (13)>
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 }
+    transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 } : (!pdl.operation) -> !pdl.operation
 }
 
 // CHECK-LABEL: @multitile_sizes_static
 func.func @multitile_sizes_static(
   %arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>)
     -> tensor<13x42xf32> {
   %0 = linalg.matmul  ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>)
                      outs(%arg2: tensor<13x42xf32>)
     -> tensor<13x42xf32>
   // The first application computes the total size.
   // CHECK: %{{.*}} = affine.apply #[[$MAP13]]()
   // CHECK: %[[SIZE:.+]] = affine.apply #[[$MAP13]]()
   // CHECK: %[[COND:.+]] = arith.cmpi eq, %[[SIZE]], %{{.*}}
   // CHECK: cf.assert %[[COND]], "could not compute dynamic multi-size tile shapes"
 
   return %0 : tensor<13x42xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    transform.structured.multitile_sizes %0 { target_size = 3, divisor = 2, dimension = 0 }
+    %low_tile, %high_tile, %split_point =
+      transform.structured.multitile_sizes %0 { target_size = 3, dimension = 0 }
+      : (!pdl.operation) -> !transform.param<i64>
+    // expected-remark @below {{2 : i64}}
+    transform.test_print_param %low_tile : !transform.param<i64>
+    // expected-remark @below {{3 : i64}}
+    transform.test_print_param %high_tile : !transform.param<i64>
+    // expected-remark @below {{4 : i64}}
+    transform.test_print_param %split_point : !transform.param<i64>
+}
+
+// CHECK-LABEL: @multitile_sizes_static_gen
+func.func @multitile_sizes_static_gen(
+  %arg0: tensor<13x34xf32>, %arg1: tensor<34x42xf32>, %arg2: tensor<13x42xf32>)
+    -> tensor<13x42xf32> {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<13x34xf32>, tensor<34x42xf32>)
+                     outs(%arg2: tensor<13x42xf32>)
+    -> tensor<13x42xf32>
+
+  return %0 : tensor<13x42xf32>
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+    transform.structured.multitile_sizes %0 { target_size = 3, divisor = 2, dimension = 0 } : (!pdl.operation) -> !pdl.operation
 }
 
 // CHECK: #[[$MAP_A:.+]] = affine_map<()[s0] -> ([[A_IMPL:s0 floordiv 2]])>
 // CHECK: #[[$MAP_T:.+]] = affine_map<() -> (2)>
 // CHECK: #[[$MAP_D:.+]] = affine_map<()[s0] -> ([[D_IMPL:\(s0 floordiv 2 \+ 1\) floordiv 2]])>
 // CHECK: #[[$MAP_S:.+]] = affine_map<()[s0] -> ((([[A_IMPL]]) floordiv ([[D_IMPL]])) * 2)>
 // CHECK: #[[$MAP_V:.+]] = affine_map<()[s0] -> (([[A_IMPL]]) mod ([[D_IMPL]]))>
 // CHECK: #[[$MAP_U:.+]] = affine_map<()[s0] -> ([[D_IMPL]] - ([[A_IMPL]]) mod ([[D_IMPL]]))>
func.func @multitile_sizes_dynamic(
   %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
     -> tensor<?x?xf32> {
   %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                      outs(%arg2: tensor<?x?xf32>)
     -> tensor<?x?xf32>
 
   return %0 : tensor<?x?xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+    // expected-error @below {{cannot compute parametric tile sizes for dynamically shaped payload op}}
+    transform.structured.multitile_sizes %0 { target_size = 3, divisor = 2, dimension = 0 }
+      : (!pdl.operation) -> !transform.param<i64>
+}
+
+func.func @multitile_sizes_dynamic_gen(
+  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32> {
+  // expected-note @below {{payload op}}
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+                     outs(%arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32>
+
+  return %0 : tensor<?x?xf32>
+}

mlir/test/Dialect/Linalg/transform-op-scalarize.mlir

func.func @scalarize(%arg0: tensor<24x12xf32>,
 
   // CHECK: return %[[RES_LOOP_1]] : tensor<24x25xf32>
   func.return %0 : tensor<24x25xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-  %1, %loops = transform.structured.tile %0 [10, 0, 0]
+  %1, %loops = transform.structured.tile %0 [10, 0, 0] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
   %2 = transform.structured.scalarize %1
 }

mlir/test/Dialect/Linalg/transform-op-split.mlir

 // RUN: mlir-opt %s --test-transform-dialect-interpreter --split-input-file -verify-diagnostics | FileCheck %s
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-  %1:2 = transform.structured.split %0 after 42 { dimension = 0 }
+  %1:2 = transform.structured.split %0 after 42 { dimension = 0 } : !pdl.operation
 }
 
 func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
 
 // CHECK: #[[$ADD_42_MAP:.+]] = affine_map<(d0) -> (d0 + 42)>
 
 // CHECK-LABEL: @one_d_static
 // CHECK-SAME:  %[[IN:.+]]: tensor<100xf32>, %[[OUT:.+]]: tensor<100xf32>
func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   return %0 : tensor<100xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-  %1:2 = transform.structured.split %0 after 42 { dimension = 0 }
+  %1:2 = transform.structured.split %0 after 42 { dimension = 0 } : !pdl.operation
 }
 
 func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
 
 // CHECK-LABEL: @one_d_static_overflow
 // CHECK-SAME:  %[[IN:.+]]: tensor<10xf32>, %[[OUT:.+]]: tensor<10xf32>
 func.func @one_d_static_overflow(%arg0: tensor<10xf32>, %arg1: tensor<10xf32>) -> tensor<10xf32> {
   // Folding is sufficiently powerful to detect the static overflow and avoid

 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   %1 = transform.structured.match ops{["func.call"]} in %arg1
-  transform.structured.split %0 after %1 { dimension = 0 }
+  transform.structured.split %0 after %1 { dimension = 0 } : !pdl.operation, !pdl.operation
 }
 
 func.func private @get_size() -> index
 
 // CHECK: #[[$MAP_MIN_100:.+]] = affine_map<()[s0] -> (s0, 100)>
 // CHECK: #[[$MAP_S_MINUS_100:.+]] = affine_map<()[s0] -> (-s0 + 100)>
 
 // CHECK-LABEL: @dynamic
func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   return %1 : tensor<100xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
-  %1:2 = transform.structured.split %0 after 4 { dimension = 0}
-  %2:2 = transform.structured.split %1#1 after 16 { dimension = 1 }
+  %1:2 = transform.structured.split %0 after 4 { dimension = 0 } : !pdl.operation
+  %2:2 = transform.structured.split %1#1 after 16 { dimension = 1 } : !pdl.operation
 }
 
 func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
 
 // CHECK-LABEL: @two_d
 func.func @two_d(%arg0: tensor<10x34xf32>,
                  %arg1: tensor<10x34xf32>) -> tensor<10x34xf32> {
   // Check the overall structure: split along the dimension 0, and then split

 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   %1 = transform.structured.match ops{["func.call"]} in %arg1
   // expected-error @below {{expected dynamic split point handle to point to a single-result index-typed op}}
-  transform.structured.split %0 after %1 { dimension = 0 }
+  transform.structured.split %0 after %1 { dimension = 0 } : !pdl.operation, !pdl.operation
 }
 
 func.func private @get_size() -> i64
 
 func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   // expected-note @below {{dynamic split point}}
   %0 = func.call @get_size() : () -> i64
   %1 = linalg.generic {

 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   %1 = transform.structured.match ops{["func.call"]} in %arg1
   // expected-error @below {{expected the dynamic split point handle to point to as many operations (0) as the target handle (1)}}
-  transform.structured.split %0 after %1 { dimension = 0 }
+  transform.structured.split %0 after %1 { dimension = 0 } : !pdl.operation, !pdl.operation
 }
 
 func.func private @get_size() -> i64
 
 func.func @dynamic(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   %1 = linalg.generic {
     indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
     iterator_types = ["parallel"]
   }
   ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
   ^bb0(%3: f32, %4: f32):
     linalg.yield %3 : f32
   } -> tensor<100xf32>
   return %1 : tensor<100xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["func.return"]} in %arg1
   // expected-error @below {{only applies to structured ops}}
-  transform.structured.split %0 after 16 { dimension = 1 }
+  transform.structured.split %0 after 16 { dimension = 1 } : !pdl.operation
 }
 
 func.func @noop(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   // expected-note @below {{target op}}
   return %arg0 : tensor<100xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   // expected-error @below {{dimension 1 does not exist in target op}}
-  transform.structured.split %0 after 16 { dimension = 1 }
+  transform.structured.split %0 after 16 { dimension = 1 } : !pdl.operation
 }
 
 func.func @one_d_static(%arg0: tensor<100xf32>, %arg1: tensor<100xf32>) -> tensor<100xf32> {
   // expected-note @below {{target op}}
   %0 = linalg.generic {
     indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
     iterator_types = ["parallel"]
   }
   ins(%arg0: tensor<100xf32>) outs(%arg1: tensor<100xf32>) {
   ^bb0(%0: f32, %1: f32):
     linalg.yield %0 : f32
   } -> tensor<100xf32>
   return %0 : tensor<100xf32>
 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.generic"]} in %arg1
   // expected-error @below {{splitting does not produce the second part for a subset of targets}}
   // expected-note @below {{expected splitting to produce the second part of all or none of the targets}}
-  %1:2 = transform.structured.split %0 after 142 { dimension = 0 }
+  %1:2 = transform.structured.split %0 after 142 { dimension = 0 } : !pdl.operation
 }
 
 func.func private @elem(%arg0: f32, %arg1: index, %arg2: index) -> f32
 
 func.func @split_one_but_not_other(
     %arg0: tensor<100xf32>, %arg1: tensor<100xf32>,
     %arg2: tensor<200xf32>, %arg3: tensor<200xf32>)
     -> (tensor<100xf32>, tensor<200xf32>) {

mlir/test/Dialect/Linalg/transform-op-tile.mlir

-// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
+// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file --verify-diagnostics %s | FileCheck %s
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-  %1, %loops:3 = transform.structured.tile %0 [4, 4, 4]
+  %1, %loops:3 = transform.structured.tile %0 [4, 4, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @tile_linalg_matmul(
 // CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<128x128xf32>
 // CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<128x128xf32>
 // CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<128x128xf32>
 // CHECK-SAME:  -> tensor<128x128xf32> {
 func.func @tile_linalg_matmul(

 }
 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
   %1 = transform.structured.match ops{["func.call"]} in %arg1
-  %2, %loops:3 = transform.structured.tile %0 [%1, %1, 4]
+  %2, %loops:3 = transform.structured.tile %0 [%1, %1, 4] : (!pdl.operation, !pdl.operation, !pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 func.func private @get_dynamic_tile_size() -> index
 
 // CHECK-LABEL: func @tile_linalg_matmul_dynamic(
 // CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<128x128xf32>
 // CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<128x128xf32>
 // CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<128x128xf32>
//      CHECK:   scf.yield %[[TD1]] : tensor<128x128xf32>
   %sz = func.call @get_dynamic_tile_size() : () -> index
   %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
                      outs(%arg2: tensor<128x128xf32>)
     -> tensor<128x128xf32>
 
 //      CHECK: return %[[TD0]] : tensor<128x128xf32>
   return %0 : tensor<128x128xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+  // expected-note @below {{for this parameter}}
+  %1 = transform.test_produce_integer_param_with_type i64 : !transform.param<i64>
+  // expected-error @below {{expected as many parameter values (0) as target ops (2)}}
+  transform.structured.tile %0 [%1, %1, %1]
+    : (!pdl.operation, !transform.param<i64>, !transform.param<i64>, !transform.param<i64>)
+    -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+}
+
+func.func @tile_linalg_matmul(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> (tensor<128x128xf32>, tensor<128x128xf32>) {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  %1 = linalg.matmul  ins(%0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  return %0, %1 : tensor<128x128xf32>, tensor<128x128xf32>
+}
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg1: !pdl.operation):
+  %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
+  // expected-note @below {{for this handle}}
+  %1 = transform.structured.match ops{["arith.constant"]} in %arg1
+  // expected-error @below {{expected as many dynamic size-producing operations (0) as target ops (2)}}
+  transform.structured.tile %0 [%1, %1, 1]
+    : (!pdl.operation, !pdl.operation, !pdl.operation)
+    -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+}
+
+func.func @tile_linalg_matmul(
+  %arg0: tensor<128x128xf32>, %arg1: tensor<128x128xf32>, %arg2: tensor<128x128xf32>)
+    -> (tensor<128x128xf32>, tensor<128x128xf32>) {
+  %0 = linalg.matmul  ins(%arg0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  %1 = linalg.matmul  ins(%0, %arg1: tensor<128x128xf32>, tensor<128x128xf32>)
+                     outs(%arg2: tensor<128x128xf32>)
+    -> tensor<128x128xf32>
+  return %0, %1 : tensor<128x128xf32>, tensor<128x128xf32>
+}

mlir/test/Dialect/Linalg/transform-ops-invalid.mlir

 
 // -----
 
 transform.sequence failures(propagate) {
 ^bb0(%arg0: !pdl.operation):
   // expected-error@below {{'transform.structured.interchange' op attribute 'iterator_interchange' failed to satisfy constraint: i64 dense array attribute whose value is non-negative}}
   transform.structured.interchange %arg0 iterator_interchange = [-3, 1]
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+^bb0(%arg0: !pdl.operation):
+  // expected-error@below {{expects all results type to be the same}}
+  "transform.structured.multitile_sizes"(%arg0) { target_size = 3, divisor = 2, dimension = 0 }
+      : (!pdl.operation) -> (!transform.param<i64>, !transform.param<i64>, !transform.param<i32>)
+}

mlir/test/Dialect/Linalg/transform-ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 
 transform.sequence failures(propagate) {
 ^bb1(%arg0: !pdl.operation):
   // CHECK %{{.*}}, %{{.*}}:2 = transform.structured.tile
-  %0, %1:2 = transform.structured.tile %arg0 [2, 0, 3]
+  %0, %1:2 = transform.structured.tile %arg0 [2, 0, 3] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg0: !transform.any_op):
+  %0:2 = transform.structured.split %arg0 after 42 { dimension = 0 } : !transform.any_op
+  transform.structured.split %0#0 after %0#1 { dimension = 1 } : !transform.any_op, !transform.any_op
 }
 
 //===----------------------------------------------------------------------===//
 // Check that operations are registered correctly through the extension
 // mechanism. Their syntax is generated and requires no additional testing since
 // we test the generator.
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/Linalg/transform-patterns.mlir

 // RUN: mlir-opt %s -test-transform-dialect-interpreter -test-linalg-transform-patterns=test-patterns -split-input-file | FileCheck %s
 
 // -----
 
 func.func @dot(%x: memref<?xf32, strided<[1], offset: ?>>,
           %y: memref<?xf32, strided<[1], offset: ?>>,
           %v: memref<f32>) {
   linalg.dot ins(%x, %y: memref<?xf32, strided<[1], offset: ?>>,
                          memref<?xf32, strided<[1], offset: ?>>)
             outs(%v: memref<f32>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.dot"]} in %arg1
-    %1, %loop = transform.structured.tile %0 [8000]
+    %1, %loop = transform.structured.tile %0 [8000] : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @dot
 // CHECK-DAG:     %[[c0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[c8000:.*]] = arith.constant 8000 : index
 // CHECK:         scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c8000]] {
 // CHECK:           linalg.dot
 
 // -----
 
 func.func @matvec(%A: memref<?x?xf32, strided<[?, 1], offset: ?>>,
              %x: memref<?xf32, strided<[1], offset: ?>>,
              %y: memref<?xf32, strided<[1], offset: ?>>) {
   linalg.matvec
     ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
                 memref<?xf32, strided<[1], offset: ?>>)
     outs(%y: memref<?xf32, strided<[1], offset: ?>>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1
-    %1, %loops:2 = transform.structured.tile %0 [5, 6]
+    %1, %loops:2 = transform.structured.tile %0 [5, 6] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @matvec
 // CHECK-DAG:     %[[c0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[c5:.*]] = arith.constant 5 : index
 // CHECK-DAG:     %[[c6:.*]] = arith.constant 6 : index
 // CHECK:         scf.for {{.*}} step %[[c5]]
 // CHECK:           scf.for {{.*}} step %[[c6]]
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
                             memref<?x?xf32, strided<[?, 1], offset: ?>>)
                outs(%C: memref<?x?xf32, strided<[?, 1], offset: ?>>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000]
-    %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400]
-    %3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40]
-    %4, %loops_4:3 = transform.structured.tile %3 [2, 3, 4]
+    %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+    %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+    %3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+    %4, %loops_4:3 = transform.structured.tile %3 [2, 3, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @matmul
 // CHECK-DAG:     %[[c0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[c2:.*]] = arith.constant 2 : index
 // CHECK-DAG:     %[[c3:.*]] = arith.constant 3 : index
 // CHECK-DAG:     %[[c4:.*]] = arith.constant 4 : index
 // CHECK-DAG:     %[[c20:.*]] = arith.constant 20 : index
linalg.matvec ins(%A, %x: memref<?x?xf32, strided<[?, 1], offset: ?>>,
                             memref<?xf32, strided<[1], offset: ?>>)
                outs(%y: memref<?xf32, strided<[1], offset: ?>>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matvec"]} in %arg1
-    %1, %loops:2 = transform.structured.tile %0 [5, 6] {interchange = [1, 0]}
+    %1, %loops:2 = transform.structured.tile %0 [5, 6] {interchange = [1, 0]} : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @matvec_perm
 // CHECK-DAG:     %[[c0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[c5:.*]] = arith.constant 5 : index
 // CHECK-DAG:     %[[c6:.*]] = arith.constant 6 : index
 // CHECK:         scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c6]]
 // CHECK:           scf.for {{.*}} = %[[c0]] to {{.*}} step %[[c5]]
linalg.matmul ins(%A, %B: memref<?x?xf32, strided<[?, 1], offset: ?>>,
                             memref<?x?xf32, strided<[?, 1], offset: ?>>)
                outs(%C : memref<?x?xf32, strided<[?, 1], offset: ?>>)
   return
 }
 
 transform.sequence failures(propagate) {
   ^bb0(%arg1: !pdl.operation):
     %0 = transform.structured.match ops{["linalg.matmul"]} in %arg1
-    %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] {interchange = [1, 2, 0]}
-    %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] {interchange = [1, 0, 2]}
-    %3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40]
+    %1, %loops:3 = transform.structured.tile %0 [2000, 3000, 4000] {interchange = [1, 2, 0]} : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+    %2, %loops_2:3 = transform.structured.tile %1 [200, 300, 400] {interchange = [1, 0, 2]} : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
+    %3, %loops_3:3 = transform.structured.tile %2 [20, 30, 40] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
 }
 
 // CHECK-LABEL: func @matmul_perm
 // CHECK-DAG:     %[[c0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[c20:.*]] = arith.constant 20 : index
 // CHECK-DAG:     %[[c30:.*]] = arith.constant 30 : index
 // CHECK-DAG:     %[[c40:.*]] = arith.constant 40 : index
 // CHECK-DAG:     %[[c200:.*]] = arith.constant 200 : index

mlir/test/Dialect/Transform/selective-targeting.mlir

pdl.pattern @pdl_target_attrC : benefit(1) {
     %0 = operation "linalg.matmul"(%args : !pdl.range<value>) {"test.attrC" = %attr}-> (%results : !pdl.range<type>)
     // TODO: we don't want this, but it is the required terminator for pdl.pattern
     rewrite %0 with "transform.dialect"
   }
 
   transform.sequence %arg0 : !pdl.operation failures(propagate) {
   ^bb1(%arg1: !pdl.operation):
     %0 = pdl_match @pdl_target_attrA in %arg1 : (!pdl.operation) -> !pdl.operation
-    transform.structured.tile %0 [4, 4, 4]
+    transform.structured.tile %0 [4, 4, 4] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
     %1 = pdl_match @pdl_target_attrC in %arg1 : (!pdl.operation) -> !pdl.operation
     %2 = transform.get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
     transform.structured.vectorize %2
   }
 }
 
 // -----
 

mlir/test/Dialect/Vector/transform-vector.mlir

%0 = linalg.matmul  ins(%arg0, %arg1: tensor<8x16xf32>, tensor<16x32xf32>)
                      outs(%arg2: tensor<8x32xf32>)
     -> tensor<8x32xf32>
   return %0 : tensor<8x32xf32>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.matmul"]} in %module_op
-  %1, %loops:3 = transform.structured.tile %0 [8, 4, 2]
+  %1, %loops:3 = transform.structured.tile %0 [8, 4, 2] : (!pdl.operation) -> (!pdl.operation, !pdl.operation, !pdl.operation, !pdl.operation)
   %2 = get_closest_isolated_parent %1 : (!pdl.operation) -> !pdl.operation
   transform.structured.vectorize %2
   transform.bufferization.one_shot_bufferize %module_op
 
   %func = transform.structured.match ops{["func.func"]} in %module_op
   transform.vector.lower_vectors %func multireduction_lowering = "innerreduction"
 }

mlir/test/python/dialects/transform_structured_ext.py

def testInterchange():
   # CHECK: transform.structured.interchange
   # CHECK: iterator_interchange = [1, 0]
 
 
 @run
 def testMultitileSizes():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
-    structured.MultiTileSizesOp(
-        sequence.bodyTarget, dimension=1, target_size=42)
+    structured.MultiTileSizesOp(pdl.OperationType.get(),
+                                sequence.bodyTarget,
+                                dimension=1,
+                                target_size=42)
     transform.YieldOp()
   # CHECK-LABEL: TEST: testMultitileSizes
   # CHECK: transform.sequence
   # CHECK: transform.structured.multitile_sizes
   # CHECK-DAG: dimension = 1
   # CHECK-DAG: target_size = 42
 
 
def testSplit():
   # CHECK-LABEL: TEST: testSplit
   # CHECK: %[[F:.+]], %[[S:.+]] = transform.structured.split %{{.*}} after 42 {dimension = 1
   # CHECK: transform.structured.split %[[F]] after %[[S]] {dimension = 3
 
 @run
 def testTileCompact():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
-    structured.TileOp(sequence.bodyTarget, sizes=[4, 8], interchange=[0, 1])
+    structured.TileOp(sequence.bodyTarget,
+                      sizes=[4, 8],
+                      interchange=[0, 1])
     transform.YieldOp()
   # CHECK-LABEL: TEST: testTileCompact
   # CHECK: transform.sequence
   # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 8]
   # CHECK: interchange = [0, 1]
 
 @run
 def testTileAttributes():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   attr = DenseI64ArrayAttr.get([4, 8])
   ichange = DenseI64ArrayAttr.get([0, 1])
   with InsertionPoint(sequence.body):
-    structured.TileOp(sequence.bodyTarget, sizes=attr, interchange=ichange)
+    structured.TileOp(sequence.bodyTarget,
+                      sizes=attr,
+                      interchange=ichange)
     transform.YieldOp()
   # CHECK-LABEL: TEST: testTileAttributes
   # CHECK: transform.sequence
   # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 8]
   # CHECK: interchange = [0, 1]
 
 @run
 def testTileZero():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
-    structured.TileOp(
-        sequence.bodyTarget, sizes=[4, 0, 2, 0], interchange=[0, 1, 2, 3])
+    structured.TileOp(sequence.bodyTarget,
+                      sizes=[4, 0, 2, 0],
+                      interchange=[0, 1, 2, 3])
     transform.YieldOp()
   # CHECK-LABEL: TEST: testTileZero
   # CHECK: transform.sequence
   # CHECK: %{{.+}}, %{{.+}}:2 = transform.structured.tile %{{.*}}[4, 0, 2, 0]
   # CHECK: interchange = [0, 1, 2, 3]
 
 @run
 def testTileDynamic():
   with_pdl = transform.WithPDLPatternsOp(pdl.OperationType.get())
   with InsertionPoint(with_pdl.body):
     sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [],
                                     with_pdl.bodyTarget)
     with InsertionPoint(sequence.body):
       m1 = transform.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "first")
       m2 = transform.PDLMatchOp(pdl.OperationType.get(), sequence.bodyTarget, "second")
-      structured.TileOp(sequence.bodyTarget, sizes=[m1, 3, m2, 0])
+      structured.TileOp(sequence.bodyTarget,
+                        sizes=[m1, 3, m2, 0])
       transform.YieldOp()
   # CHECK-LABEL: TEST: testTileDynamic
   # CHECK: %[[FIRST:.+]] = pdl_match
   # CHECK: %[[SECOND:.+]] = pdl_match
   # CHECK: %{{.+}}, %{{.+}}:3 = transform.structured.tile %{{.*}}[%[[FIRST]], 3, %[[SECOND]], 0]
 
 
 @run
+def testTileExplicitLoopTypeSingle():
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.AnyOpType.get())
+  with InsertionPoint(sequence.body):
+    structured.TileOp(transform.OperationType.get("scf.for"),
+                      sequence.bodyTarget,
+                      sizes=[2, 3, 4])
+    transform.YieldOp()
+  # CHECK-LABEL: TEST: testTileExplicitLoopTypeSingle
+  # CHECK: = transform.structured.tile %{{.*}} : (!{{.*}}) ->
+  # CHECK-COUNT-3: !transform.op<"scf.for">
+
+
+
+@run
+def testTileExplicitLoopTypeAll():
+  sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
+                                  [], transform.AnyOpType.get())
+  types = [
+      transform.OperationType.get(x)
+      for x in ["scf.for", "scf.parallel", "scf.foreach_thread"]
+  ]
+  with InsertionPoint(sequence.body):
+    structured.TileOp(types, sequence.bodyTarget, sizes=[2, 3, 4])
+    transform.YieldOp()
+  # CHECK-LABEL: TEST: testTileExplicitLoopTypeAll
+  # CHECK: = transform.structured.tile
+  # CHECK-SAME : (!transform.any_op) -> (!transform.any_op, !transform.op<"scf.for">,
+  # CHECK-SAME: !transform.op<"scf.parallel">, !transform.op<"scf.foreach_thread">
+
+@run
 def testVectorize():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get())
   with InsertionPoint(sequence.body):
     structured.VectorizeOp(sequence.bodyTarget, vectorize_padding=True)
     transform.YieldOp()
   # CHECK-LABEL: TEST: testVectorize
   # CHECK: transform.sequence
   # CHECK: = transform.structured.vectorize
   # CHECK: {vectorize_padding}