[mlir][Linalg] Add a TileToForeachThread transform.

[nicolasvasilache]

Summary

This revision adds a new transformation to tile a TilingInterface op to a tiled scf.foreach_thread, applying
tiling by num_threads.
If non-empty, the threadDimMapping is added as an attribute to the resulting scf.foreach_thread.
0-tile sizes (i.e. tile by the full size of the data) are used to encode
that a dimension is not tiled.

[ftynse]

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

568–569

574–575

Please use the appropriate header. Underscore === is a h1 and it wrecks the generated documentation.

592–593

I find it confusing to have "tile op" and "tiled op". Can one of these rather be "foreach_thread_op" ?

595

Non-blocking: could we rather adopt the view-like syntax $target[$num_threads] directly because we will have to switch to that to support dynamic number of threads as the TODO above indicates.

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

807

It looks like the code in tileToForeachThreadOp will have always reported an error, so it's better to just return DiagnosedSilenceableFailure::definiteFailure() here to avoid a second error message.

Also, this doesn't seem to match the description of the op that says it fails silently. reportUnknownTransformError produces a LogicalResult::failure(), which is converted into definite failure by DiagnosedSilenceableFailure and not into a silenceable one.

808

I don't see why we would ever need it. Patterns are only adding complexity here.

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

172–191

Please add top-level docs for these.

210

Nit: you no longer need to specify the size in llvm::to_vector. In this particular case, because you specified 4, it will force a copy from a 4-element vector on the RHS to the 6-element vector on the LHS.

250–253

I find the Range interface confusing. Isn't "size" intended as the distance between the lower bound (offset) and the upper bound (offset + size) of the iteration space dimension? If so, it should not be necessary to subtract the offset here. If not, it would be much better off called "upper bound" than "size". Note that it is impossible to catch now because all ops have zero iteration space offsets.

276

[nicolasvasilache]

Rebase and address review.

[nicolasvasilache]

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

807

Used return emitDefaultSilenceableFailure(target);
I find that nested error messages are usually not informative enough and like to see where in the transform dialect IR this happens.
Lmk if I am missing something ..

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

250–253

It is indeed the size, thanks!

[springerm]

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

579

properly as in "without failure"?

581–582

Not sure what this means. If no tiled op was produced, the transform must have failed. In that case we would not continue executing the sequence anyway.

585

Not necessarily a linalg op

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

527

maybe replace with e.g. as not every tiledOp is a linalg op

[nicolasvasilache]

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

581–582

The op may always be given handles to things that are all ignored and it may return and propagate the empty list.
This is the generic wording that we have been using to make the user aware of the fact that results may always be empty.

[nicolasvasilache]

Rebase and address last comments.

[nicolasvasilache]

Update doc comments.

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

 //===- LinalgTransformOps.h - Linalg transform ops --------------*- C++ -*-===//
 //
 // 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 MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
 #define MLIR_DIALECT_LINALG_TRANSFORMOPS_LINALGTRANSFORMOPS_H
 
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
+class TilingInterface;
 namespace linalg {
 class GenericOp;
 class LinalgOp;
 } // namespace linalg
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
 // Linalg Transform Operations

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

def TileOp : Op<Transform_Dialect, "structured.tile",
 
   let extraClassDeclaration = [{
     /// Returns the list of tile sizes, which may be static (Attribute) or
     /// dynamic (Value).
     SmallVector<OpFoldResult> getMixedSizes();
   }];
 }
 
+def TileToForeachThreadOp :
+    Op<Transform_Dialect, "structured.tile_to_foreach_thread_op",
+      [FunctionalStyleTransformOpTrait,
+       MemoryEffectsOpInterface,
+       TransformEachOpTrait,
+       TransformOpInterface]> {
+  let description = [{
+    Tile a TilingInterface `op` to a tiled `scf.foreach_thread`, applying
+    tiling by `num_threads`.
+    If non-empty, the `thread_dim_mapping` is added as an attribute to the 
+    resulting `scf.foreach_thread`.
+    Zero tile sizes indicate that the dimension is not tiled, and can be thought
+    of as tiling by the full size of data.
+    It is the user's responsibility to ensure that `num_threads` is a valid 
+    tiling specification (i.e. that only tiles parallel dimensions in the 
+    Linalg case).
+    
+    #### Return modes
+    
+    This operation ignores ops that do not implement the TilingInterface and
+    drops them in the return.
+
+    If all the operations referred to by the `target` PDLOperation tile
+    successfully, the transform succeeds.
+    Otherwise the transform silently fails.
+    
+    The 2 returned handles point to only the subset of successfully produced 
+    tiled operations, which can all be empty.
+    
+    These 2 returned handles point to:
+      - the new scf.foreach_thread op,
+      - the tiled op that implements TilingInterface.
+  }];
+
+  let arguments = (ins PDL_Operation:$target,
+                   // TODO: dynamic number of threads.
+                   DefaultValuedAttr<I64ArrayAttr, "{}">:$num_threads,
+                   OptionalAttr<I64ArrayAttr>:$thread_dim_mapping);
+  let results = (outs PDL_Operation:$foreach_thread_op,
+                      PDL_Operation:$tiled_op);
+
+  let assemblyFormat = [{
+    $target $num_threads (`(` `mapped` `to` `dims` $thread_dim_mapping^ `)`)? 
+      attr-dict
+  }];
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::TilingInterface target, 
+        ::llvm::SmallVectorImpl<::mlir::Operation *> &results, 
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 def VectorizeOp : Op<Transform_Dialect, "structured.vectorize",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformEachOpTrait, TransformOpInterface]> {
   let description = [{
     Indicates that the given `target` op all the ops it contains should be
     vectorized with the configuration specified by the attributes of this op.
     This vectorization only handles structured ops that operate on shaped types
     and does not vectorize loops or straight-line. Internally, it applies a

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

 ///
 /// 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);
 
+/// Rewrite a TilingInterface `op` to a tiled `scf.foreach_thread`, applying
+/// tiling by `numThreads`.
+/// If non-empty, the `threadDimMapping` is added as an attribute to the
+/// resulting `scf.foreach_thread`.
+/// 0-tile sizes (i.e. tile by the full size of the data) are used to encode
+/// that a dimension is not tiled.
+/// It is the user's responsibility to ensure that `numThreads` is a valid
+/// tiling specification (i.e. that only tiles parallel dimensions, e.g. in the
+/// Linalg case).
+struct ForeachThreadTilingResult {
+  Operation *tileOp;
+  Operation *tiledOp;
+};
+FailureOr<ForeachThreadTilingResult>
+tileToForeachThreadOp(OpBuilder &builder, TilingInterface op,
+                      ArrayRef<OpFoldResult> numThreads,
+                      ArrayRef<int64_t> threadDimMapping = {});
+
 /// All indices returned by IndexOp should be invariant with respect to tiling.
 /// Therefore, if an operation is tiled, we have to transform the indices
 /// accordingly, i.e. offset them by the values of the corresponding induction
 /// variables that are captured implicitly in the body of the op.
 ///
 /// Example. `linalg.generic` before tiling:
 ///
 /// #id_2d = (i, j) -> (i, j)

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

 /// Given a value, try to extract a constant Attribute. If this fails, return
 /// the original value.
 OpFoldResult getAsOpFoldResult(Value val);
 
 /// Given an array of values, try to extract a constant Attribute from each
 /// value. If this fails, return the original value.
 SmallVector<OpFoldResult> getAsOpFoldResult(ArrayRef<Value> values);
 
+/// Convert `arrayAttr` to a vector of OpFoldResult.
+SmallVector<OpFoldResult> getAsOpFoldResult(ArrayAttr arrayAttr);
+
 /// If ofr is a constant integer or an IntegerAttr, return the integer.
 Optional<int64_t> getConstantIntValue(OpFoldResult ofr);
 
 /// Return true if `ofr` is constant integer equal to `value`.
 bool isConstantIntValue(OpFoldResult ofr, int64_t value);
 
 /// Return true if ofr1 and ofr2 are the same integer constant attribute values
 /// or the same SSA value.
 /// Ignore integer bitwitdh and type mismatch that come from the fact there is
 /// no IndexAttr and that IndexType have no bitwidth.
 bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STATICVALUEUTILS_H

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

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.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/Transform/IR/TransformDialect.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Interfaces/TilingInterface.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::transform;
 
-/// Extracts a vector of int64_t from an array attribute. Asserts if the
-/// attribute contains values other than integers.
-static SmallVector<int64_t> extractI64Array(ArrayAttr attr) {
-  SmallVector<int64_t> result;
-  result.reserve(attr.size());
-  for (APInt value : attr.getAsValueRange<IntegerAttr>())
-    result.push_back(value.getSExtValue());
-  return result;
-}
-
 /// Extracts a vector of unsigned from an array attribute. Asserts if the
 /// attribute contains values other than intergers. May truncate.
 static SmallVector<unsigned> extractUIntArray(ArrayAttr attr) {
   SmallVector<unsigned> result;
   result.reserve(attr.size());
   for (APInt value : attr.getAsValueRange<IntegerAttr>())
     result.push_back(value.getZExtValue());
   return result;
if (!sizesAttr)
     return parser.emitError(opLoc)
            << "expected '" << sizesAttrName << "' attribute";
   auto sizesArrayAttr = sizesAttr.dyn_cast<ArrayAttr>();
   if (!sizesArrayAttr)
     return parser.emitError(opLoc)
            << "'" << sizesAttrName << "' attribute must be an array";
   Type pdlOpType = parser.getBuilder().getType<pdl::OperationType>();
   size_t numExpectedLoops =
-      sizesArrayAttr.size() - llvm::count(extractI64Array(sizesArrayAttr), 0);
+      sizesArrayAttr.size() -
+      llvm::count(extractFromI64ArrayAttr(sizesArrayAttr), 0);
   result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOpType));
   if (parser.resolveOperand(targetOperand, pdlOpType, result.operands))
     return failure();
   return success();
 }
 
 DiagnosedSilenceableFailure
 transform::FuseOp::apply(mlir::transform::TransformResults &transformResults,
                          mlir::transform::TransformState &state) {
   LinalgTilingAndFusionOptions fusionOptions;
-  fusionOptions.tileSizes = extractI64Array(getTileSizes());
-  fusionOptions.tileInterchange = extractI64Array(getTileInterchange());
+  fusionOptions.tileSizes = extractFromI64ArrayAttr(getTileSizes());
+  fusionOptions.tileInterchange = extractFromI64ArrayAttr(getTileInterchange());
 
   LogicalResult result = applyTilingToAll(
       getOperation(), state.getPayloadOps(getTarget()),
       fusionOptions.tileSizes.size() - llvm::count(fusionOptions.tileSizes, 0),
       transformResults, [&](LinalgOp linalgOp) -> FailureOr<TiledLinalgOp> {
         LinalgTileAndFuseTensorOpsPattern pattern(getContext(), fusionOptions);
         SimpleRewriter rewriter(getContext());
         rewriter.setInsertionPoint(linalgOp);

 
 void transform::FuseOp::print(OpAsmPrinter &p) {
   p << ' ';
   p << getTarget();
   p.printOptionalAttrDict((*this)->getAttrs());
 }
 
 LogicalResult transform::FuseOp::verify() {
-  SmallVector<int64_t> permutation = extractI64Array(getTileInterchange());
+  SmallVector<int64_t> permutation =
+      extractFromI64ArrayAttr(getTileInterchange());
   auto sequence = llvm::to_vector(llvm::seq<int64_t>(0, permutation.size()));
   if (!std::is_permutation(sequence.begin(), sequence.end(),
                            permutation.begin(), permutation.end())) {
     return emitOpError() << "expects interchange to be a permutation, found "
                          << getTileInterchange();
   }
   return success();
 }

 //===---------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::PadOp::applyToOne(linalg::LinalgOp target,
                              SmallVectorImpl<Operation *> &results,
                              transform::TransformState &state) {
   // Convert the integer packing flags to booleans.
   SmallVector<bool> packPaddings;
-  for (int64_t packPadding : extractI64Array(getPackPaddings()))
+  for (int64_t packPadding : extractFromI64ArrayAttr(getPackPaddings()))
     packPaddings.push_back(static_cast<bool>(packPadding));
 
   // Convert the padding values to attributes.
   SmallVector<Attribute> paddingValues;
   for (auto const &it :
        llvm::zip(getPaddingValues(), target->getOperandTypes())) {
     Attribute attr = std::get<0>(it);
     Type elementType = getElementTypeOrSelf(std::get<1>(it));
for (auto const &it :
     }
     paddingValues.push_back(attr);
   }
 
   // Extract the transpose vectors.
   SmallVector<SmallVector<int64_t>> transposePaddings;
   for (Attribute transposeVector : getTransposePaddings().cast<ArrayAttr>())
     transposePaddings.push_back(
-        extractI64Array(transposeVector.cast<ArrayAttr>()));
+        extractFromI64ArrayAttr(transposeVector.cast<ArrayAttr>()));
 
   LinalgPaddingOptions paddingOptions;
   paddingOptions.setPaddingValues(paddingValues);
-  paddingOptions.setPaddingDimensions(extractI64Array(getPaddingDimensions()));
+  paddingOptions.setPaddingDimensions(
+      extractFromI64ArrayAttr(getPaddingDimensions()));
   paddingOptions.setPackPaddings(packPaddings);
-  paddingOptions.setHoistPaddings(extractI64Array(getHoistPaddings()));
+  paddingOptions.setHoistPaddings(extractFromI64ArrayAttr(getHoistPaddings()));
   paddingOptions.setTransposePaddings(transposePaddings);
 
   FailureOr<LinalgOp> result =
       tryApply<LinalgPaddingPattern>(target, paddingOptions);
   if (succeeded(result)) {
     results.push_back(result->getOperation());
     return DiagnosedSilenceableFailure(success());
   }
 
   results.assign(1, nullptr);
   return emitDefaultSilenceableFailure(target);
 }
 
 LogicalResult transform::PadOp::verify() {
-  SmallVector<int64_t> packPaddings = extractI64Array(getPackPaddings());
+  SmallVector<int64_t> packPaddings =
+      extractFromI64ArrayAttr(getPackPaddings());
   if (any_of(packPaddings, [](int64_t packPadding) {
         return packPadding != 0 && packPadding != 1;
       })) {
     return emitOpError()
            << "expects pack_paddings to contain booleans (0/1), found "
            << getPackPaddings();
   }
 
   SmallVector<int64_t> paddingDimensions =
-      extractI64Array(getPaddingDimensions());
+      extractFromI64ArrayAttr(getPaddingDimensions());
   if (any_of(paddingDimensions,
              [](int64_t paddingDimension) { return paddingDimension < 0; })) {
     return emitOpError()
            << "expects padding_dimensions to contain positive integers, found "
            << getPaddingDimensions();
   }
 
-  SmallVector<int64_t> hoistPaddings = extractI64Array(getHoistPaddings());
+  SmallVector<int64_t> hoistPaddings =
+      extractFromI64ArrayAttr(getHoistPaddings());
   if (any_of(hoistPaddings,
              [](int64_t hoistPadding) { return hoistPadding < 0; })) {
     return emitOpError()
            << "expects hoist_paddings to contain positive integers, found "
            << getHoistPaddings();
   }
 
   ArrayAttr transposes = getTransposePaddings();

 //===----------------------------------------------------------------------===//
 // TileOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::TileOp::apply(TransformResults &transformResults,
                          TransformState &state) {
   LinalgTilingOptions tilingOptions;
-  SmallVector<int64_t> tileSizes = extractI64Array(getStaticSizes());
+  SmallVector<int64_t> tileSizes = extractFromI64ArrayAttr(getStaticSizes());
 
   ArrayRef<Operation *> targets = state.getPayloadOps(getTarget());
   SmallVector<ArrayRef<Operation *>> dynamicSizeProducers;
   dynamicSizeProducers.reserve(getDynamicSizes().size());
   for (Value dynamicSizeProducerHandle : getDynamicSizes()) {
     dynamicSizeProducers.push_back(
         state.getPayloadOps(dynamicSizeProducerHandle));
 
transform::TileOp::apply(TransformResults &transformResults,
   for (const auto &en : llvm::enumerate(loops))
     transformResults.set(getLoops()[en.index()].cast<OpResult>(), en.value());
 
   return DiagnosedSilenceableFailure::success();
 }
 
 SmallVector<OpFoldResult> transform::TileOp::getMixedSizes() {
   ValueRange dynamic = getDynamicSizes();
-  SmallVector<int64_t> tileSizes = extractI64Array(getStaticSizes());
+  SmallVector<int64_t> tileSizes = extractFromI64ArrayAttr(getStaticSizes());
   SmallVector<OpFoldResult> results;
   results.reserve(tileSizes.size());
   unsigned dynamicPos = 0;
   Builder builder(getContext());
   for (int64_t size : tileSizes) {
     if (size == ShapedType::kDynamicSize) {
       results.push_back(dynamic[dynamicPos++]);
     } else {
if (parser.parseOperand(target) ||
       parser.resolveOperand(target, pdlOperationType, result.operands) ||
       parseOperandsOrIntegersSizesList(parser, dynamicSizes, staticSizes) ||
       parser.resolveOperands(dynamicSizes, pdlOperationType, result.operands) ||
       parser.parseOptionalAttrDict(result.attributes))
     return ParseResult::failure();
 
   result.addAttribute(getStaticSizesAttrName(result.name), staticSizes);
   size_t numExpectedLoops =
-      staticSizes.size() - llvm::count(extractI64Array(staticSizes), 0);
+      staticSizes.size() - llvm::count(extractFromI64ArrayAttr(staticSizes), 0);
   result.addTypes(SmallVector<Type>(numExpectedLoops + 1, pdlOperationType));
   return success();
 }
 
 void TileOp::print(OpAsmPrinter &p) {
   p << ' ' << getTarget();
   printOperandsOrIntegersSizesList(p, getOperation(), getDynamicSizes(),
                                    getStaticSizes());
   p.printOptionalAttrDict((*this)->getAttrs(), {getStaticSizesAttrName()});
 }
 
 void transform::TileOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   consumesHandle(getTarget(), effects);
   onlyReadsHandle(getDynamicSizes(), effects);
   producesHandle(getTiledLinalgOp(), effects);
   producesHandle(getLoops(), effects);
   modifiesPayload(effects);
 }
 
 //===----------------------------------------------------------------------===//
+// TileToForeachThreadOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::TileToForeachThreadOp::applyToOne(
+    TilingInterface target, SmallVectorImpl<Operation *> &results,
+    transform::TransformState &state) {
+  IRRewriter rewriter(getContext());
+  rewriter.setInsertionPoint(target);
+  auto maybeThreadDimMappingAttr = getThreadDimMapping();
+  FailureOr<ForeachThreadTilingResult> tilingResult =
+      linalg::tileToForeachThreadOp(
+          rewriter, target, getAsOpFoldResult(getNumThreads()),
+          maybeThreadDimMappingAttr
+              ? extractFromI64ArrayAttr(*maybeThreadDimMappingAttr)
+              : ArrayRef<int64_t>{});
+  if (failed(tilingResult))
+    return emitDefaultSilenceableFailure(target);
+  rewriter.replaceOp(target, tilingResult->tileOp->getResults());
+  results.assign({tilingResult->tileOp, tilingResult->tiledOp});
+  return DiagnosedSilenceableFailure(success());
+}
+
+//===----------------------------------------------------------------------===//
 // VectorizeOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::VectorizeOp::applyToOne(Operation *target,
                                    SmallVectorImpl<Operation *> &results,
                                    transform::TransformState &state) {
   if (!target->hasTrait<OpTrait::IsIsolatedFromAbove>()) {

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

if (emitAssertions) {
     b.create<cf::AssertOp>(
         equals, builder.getStringAttr(
                     "could not compute dynamic multi-size tile shapes"));
   }
 
   return spec;
 }
 
+/// Given a `subsetExtractOp`, a `source` and a `dest`, create a new
+/// `ParallelInsertSlice` op of `source` into `dest` at the same subset location
+/// as `subsetExtractOp`.
+static void
+createMatchingParallelSubsetInsertOp(OpBuilder &b, Location loc,
+                                     tensor::ExtractSliceOp subsetExtractOp,
+                                     Value source, Value dest) {
+  b.create<tensor::ParallelInsertSliceOp>(
+      loc, source, dest, subsetExtractOp.getMixedOffsets(),
+      subsetExtractOp.getMixedSizes(), subsetExtractOp.getMixedStrides());
+}
+
+/// Build an `affine_max` of all the `vals`.
+static Value buildMax(OpBuilder &b, Location loc, ValueRange vals) {
+  return b.createOrFold<AffineMaxOp>(
+      loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
+      vals);
+}
+
+/// Build an `affine_min` of all the `vals`.
+static Value buildMin(OpBuilder &b, Location loc, ValueRange vals) {
+  return b.createOrFold<AffineMinOp>(
+      loc, AffineMap::getMultiDimIdentityMap(vals.size(), loc.getContext()),
+      vals);
+}
+
+FailureOr<ForeachThreadTilingResult>
+linalg::tileToForeachThreadOp(OpBuilder &b, TilingInterface op,
+                              ArrayRef<OpFoldResult> numThreads,
+                              ArrayRef<int64_t> threadDimMapping) {
+  Location loc = op->getLoc();
+  OpBuilder::InsertionGuard g(b);
+  SmallVector<Range> loopRanges = op.getIterationDomain(b);
+  if (loopRanges.empty())
+    return op->emitOpError("expected non-empty loop ranges");
+  auto hasStrideOne = [](Range r) { return !isConstantIntValue(r.stride, 1); };
+  if (llvm::any_of(loopRanges, hasStrideOne))
+    return op->emitOpError("only stride-1 supported atm");
+  // TODO: support `getTiledImplementation` with >1 produced tiled ops.
+  auto destOperands = op.getDestinationOperands(b);
+  if (destOperands.size() != 1)
+    return op->emitOpError("only single dest operand supported atm");
+
+  SmallVector<OpFoldResult> nonZeroNumThreads =
+      llvm::to_vector(llvm::make_filter_range(numThreads, [](OpFoldResult ofr) {
+        return !isConstantIntValue(ofr, 0);
+      }));
+  SmallVector<Value> materializedNonZeroNumThreads =
+      llvm::to_vector(llvm::map_range(nonZeroNumThreads, [&](OpFoldResult ofr) {
+        ImplicitLocOpBuilder ilocb(loc, b);
+        return materializeOpFoldResult(ilocb, ofr);
+      }));
+
+  Value zero = b.create<arith::ConstantIndexOp>(loc, 0);
+  Operation *tiledOp = nullptr;
+  scf::ForeachThreadOp foreachThreadOp = b.create<scf::ForeachThreadOp>(
+      loc, materializedNonZeroNumThreads, threadDimMapping,
+      [&](OpBuilder &b, Location loc, ValueRange threadIds) {
+        int64_t nLoops = loopRanges.size();
+        SmallVector<OpFoldResult> tiledOffsets, tiledSizes;
+        tiledOffsets.reserve(nLoops);
+        tiledSizes.reserve(nLoops);
+        for (unsigned loopIdx = 0, threadIdIdx = 0; loopIdx < nLoops;
+             ++loopIdx) {
+          bool overflow = loopIdx >= numThreads.size();
+          bool isZero = !overflow && isConstantIntValue(numThreads[loopIdx], 0);
+          // Degenerate case: take the whole domain.
+          if (overflow || isZero) {
+            tiledOffsets.push_back(loopRanges[loopIdx].offset);
+            tiledSizes.push_back(loopRanges[loopIdx].size);
+            continue;
+          }
+
+          // Tiled case: compute the offset and size.
+          AffineExpr i, j, M, N, O;
+          bindDims(b.getContext(), i, j);
+          bindSymbols(b.getContext(), M, N, O);
+          Value size = loopRanges[loopIdx].size;
+          Value offset = loopRanges[loopIdx].offset;
+          Value threadId = threadIds[threadIdIdx];
+          // TODO: more aggressive foldings.
+          // Symbolic fixed max size per thread.
+          // TODO: floor + 0/1 depending on case for better load-balancing.
+          Value maxSizePerThread = b.createOrFold<AffineApplyOp>(
+              loc, M.ceilDiv(N),
+              ValueRange{size, materializedNonZeroNumThreads[threadIdIdx]});
+          // Dynamic offset shifted by threadId * maxSizePerThread.
+          Value offsetPerThread = b.createOrFold<AffineApplyOp>(
+              loc, i + j * M, ValueRange{offset, threadId, maxSizePerThread});
+          // Dynamic upper-bound depending on the threadId.
+          Value sizeMinusOffsetPerThread = b.createOrFold<AffineApplyOp>(
+              loc, -i + M, ValueRange{offsetPerThread, size});
+          Value tileSizePerThread = buildMin(
+              b, loc, ValueRange{sizeMinusOffsetPerThread, maxSizePerThread});
+          tiledOffsets.push_back(offsetPerThread);
+          // TODO: if tileSizePerThread <= 0 early exit.
+          tiledSizes.push_back(
+              buildMax(b, loc, ValueRange{zero, tileSizePerThread}));
+          ++threadIdIdx;
+        }
+
+        SmallVector<Operation *> tiledOps =
+            op.getTiledImplementation(b, destOperands, tiledOffsets, tiledSizes,
+                                      /*tileDestOperands=*/true);
+        assert(tiledOps.size() == 1 && "expected a single produced tiled op");
+        tiledOp = tiledOps.front();
+
+        auto tilingInterfaceOp = dyn_cast<TilingInterface>(tiledOp);
+        assert(tilingInterfaceOp &&
+               "Tiled op does not implement TilingInterface");
+
+        auto tiledDestOperands = tilingInterfaceOp.getDestinationOperands(b);
+
+        // Create terminator with parallel subset insert operations.
+        auto performConcurrentlyOp = b.create<scf::PerformConcurrentlyOp>(loc);
+        OpBuilder::InsertionGuard g(b);
+        b.setInsertionPointToStart(performConcurrentlyOp.getBody());
+        for (auto it :
+             llvm::zip(tiledDestOperands, tilingInterfaceOp->getResults(),
+                       destOperands)) {
+          createMatchingParallelSubsetInsertOp(
+              b, loc,
+              cast<tensor::ExtractSliceOp>(std::get<0>(it).getDefiningOp()),
+              std::get<1>(it), std::get<2>(it));
+        }
+      });
+  return ForeachThreadTilingResult{foreachThreadOp, tiledOp};
+}
+
 // Insert a tile `source` into the destination tensor `dest`. The position at
 // which the tile is inserted (as well as size of tile) is taken from a given
 // ExtractSliceOp `sliceOp`.
 static Value insertSliceIntoTensor(RewriterBase &b, Location loc,
                                    tensor::ExtractSliceOp sliceOp, Value source,
                                    Value dest) {
   return b.create<tensor::InsertSliceOp>(
       loc, sliceOp.getSource().getType(), source, dest, sliceOp.getOffsets(),

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

 
 /// Given an array of values, try to extract a constant Attribute from each
 /// value. If this fails, return the original value.
 SmallVector<OpFoldResult> getAsOpFoldResult(ArrayRef<Value> values) {
   return llvm::to_vector<4>(
       llvm::map_range(values, [](Value v) { return getAsOpFoldResult(v); }));
 }
 
+/// Convert `arrayAttr` to a vector of OpFoldResult.
+SmallVector<OpFoldResult> getAsOpFoldResult(ArrayAttr arrayAttr) {
+  SmallVector<OpFoldResult> res;
+  res.reserve(arrayAttr.size());
+  for (Attribute a : arrayAttr)
+    res.push_back(a);
+  return res;
+}
+
 /// If ofr is a constant integer or an IntegerAttr, return the integer.
 Optional<int64_t> getConstantIntValue(OpFoldResult ofr) {
   // Case 1: Check for Constant integer.
   if (auto val = ofr.dyn_cast<Value>()) {
     APSInt intVal;
     if (matchPattern(val, m_ConstantInt(&intVal)))
       return intVal.getSExtValue();
     return llvm::None;

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

+// RUN: mlir-opt %s --test-transform-dialect-interpreter -canonicalize | FileCheck %s
+
+// Offset per thread:
+// CHECK-DAG: affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 10))>
+// Per thread tile size.
+// CHECK-DAG: affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 10)) + s0, s0 ceildiv 10)>
+// CHECK-DAG: affine_map<(d0)[s0] -> (d0 * (s0 ceildiv 20))>
+// CHECK-DAG: affine_map<(d0)[s0] -> (-(d0 * (s0 ceildiv 20)) + s0, s0 ceildiv 20)>
+
+module {
+// CHECK-LABEL: matmul(
+//  CHECK-SAME:   %[[A:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[B:[0-9a-z]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[C:[0-9a-z]+]]: tensor<?x?xf32>
+  func.func @matmul(%A: tensor<?x?xf32>, %B: tensor<?x?xf32>, %C: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  //  CHECK-DAG: %[[C10:.*]] = arith.constant 10 : index
+  //  CHECK-DAG: %[[C20:.*]] = arith.constant 20 : index
+  //      CHECK: scf.foreach_thread ({{.*}}) in (%[[C10]], %[[C20]]) -> (tensor<?x?xf32>) {
+  //      CHECK:   %[[tA:.*]] = tensor.extract_slice %[[A]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+  //      CHECK:   %[[tB:.*]] = tensor.extract_slice %[[B]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+  //      CHECK:   %[[tC:.*]] = tensor.extract_slice %[[C]]{{.*}} : tensor<?x?xf32> to tensor<?x?xf32>
+  //      CHECK:   %[[RES:.*]] = linalg.matmul
+  // CHECK-SAME:      ins(%[[tA]], %[[tB]] : tensor<?x?xf32>, tensor<?x?xf32>)
+  // CHECK-SAME:     outs(%[[tC]] : tensor<?x?xf32>) -> tensor<?x?xf32>
+  // CHECK-NEXT:   scf.foreach_thread.perform_concurrently {
+  // CHECK-NEXT:     tensor.parallel_insert_slice %[[RES]] into %[[C]]{{.*}} :
+  // CHECK-SAME:       tensor<?x?xf32> into tensor<?x?xf32>
+  // CHECK-NEXT:   }
+  // CHECK-NEXT: } {thread_dim_mapping = [1, 0]}
+    %0 = linalg.matmul ins(%A, %B : tensor<?x?xf32>, tensor<?x?xf32>)
+                      outs(%C : tensor<?x?xf32>) -> (tensor<?x?xf32>)
+    return %0 : tensor<?x?xf32>
+  }
+
+  transform.with_pdl_patterns {
+  ^bb0(%arg0: !pdl.operation):
+    pdl.pattern @match_linalg_matmul : benefit(1) {
+      %0 = operands
+      %1 = types
+      %2 = operation "linalg.matmul"(%0 : !pdl.range<value>)  -> (%1 : !pdl.range<type>)
+      rewrite %2 with "transform.dialect"
+    }
+    transform.sequence %arg0 {
+    ^bb1(%arg1: !pdl.operation):
+      %0 = pdl_match @match_linalg_matmul in %arg1
+      %1:2 = transform.structured.tile_to_foreach_thread_op %0 [10, 20] (mapped to dims [1, 0])
+    }
+  }
+}