[mlir][tiling] Relax tiling to accept generating multiple operations.

[hanchung]

Summary

Some operations need to generate multiple operations when implementing
the tiling interface. Here is a sound example in IREE, see
https://github.com/iree-org/iree/pull/10905 for more details.

[hanchung]

delete debug code

[mravishankar]

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

65–68

Could you add a comment regarding what the ordering of operations here are expected to be. Long term I think we should just return SmallVector<Value>s here.

[hanchung]

add comments

mlir/include/mlir/Dialect/SCF/Transforms/TileUsingInterface.h

struct SCFTilingOptions {
   SCFTilingOptions &setInterchange(ArrayRef<int64_t> interchange) {
     interchangeVector = llvm::to_vector(interchange);
     return *this;
   }
 };
 
 /// Transformation information returned after tiling.
 struct SCFTilingResult {
-  /// The tiled operation generated.
-  Operation *tiledOp;
+  /// Tiled operations that are generated during tiling. The order does not
+  /// matter except the last op. The replacements are expected to be the results
+  /// of the last op.
+  SmallVector<Operation *> tiledOps;
   /// The `scf.for` operations that iterate over the tiles.
   SmallVector<scf::ForOp> loops;
   /// Values to use as replacements for the untiled op. Is the same size as the
   /// number of results of the untiled op.
   SmallVector<Value> replacements;
 };
 
 /// Method to tile an op that implements the `TilingInterface` using

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

transform::ScalarizeOp::applyToOne(linalg::LinalgOp target,
   SmallVector<int64_t> emptyTileSizes;
   SimpleRewriter rewriter(getContext());
   rewriter.setInsertionPoint(target);
   FailureOr<scf::SCFTilingResult> maybeTilingResult = tileUsingSCFForOp(
       rewriter, cast<TilingInterface>(target.getOperation()), tilingOptions);
   if (failed(maybeTilingResult))
     return DiagnosedSilenceableFailure(reportUnknownTransformError(target));
 
-  results.push_back(maybeTilingResult->tiledOp);
+  results.append(maybeTilingResult->tiledOps);
   return DiagnosedSilenceableFailure(success());
 }
 
 //===----------------------------------------------------------------------===//
 // SplitOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure SplitOp::apply(TransformResults &results,
if (failed(maybeTilingResult))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     if (linalgOp.hasBufferSemantics())
       rewriter.eraseOp(linalgOp);
     else
       rewriter.replaceOp(linalgOp,
                          maybeTilingResult->loops.front()->getResults());
 
-    tiled.push_back(maybeTilingResult->tiledOp);
+    tiled.append(maybeTilingResult->tiledOps);
     for (const auto &en2 : llvm::enumerate(maybeTilingResult->loops))
       loops[en2.index()].push_back(en2.value());
   }
 
   transformResults.set(getTiledLinalgOp().cast<OpResult>(), tiled);
   for (const auto &en : llvm::enumerate(loops))
     transformResults.set(getLoops()[en.index()].cast<OpResult>(), en.value());
 
for (auto &en : llvm::enumerate(targets)) {
     SimpleRewriter rewriter(tilingInterfaceOp.getContext());
     FailureOr<scf::SCFTilingResult> tilingResult =
         tileUsingSCFForOp(rewriter, tilingInterfaceOp, tilingOptions);
     if (failed(tilingResult))
       return DiagnosedSilenceableFailure::definiteFailure();
 
     rewriter.replaceOp(tilingInterfaceOp, tilingResult->replacements);
 
-    tiled.push_back(tilingResult->tiledOp);
+    tiled.append(tilingResult->tiledOps);
     for (const auto &en2 : llvm::enumerate(tilingResult->loops))
       loops[en2.index()].push_back(en2.value());
   }
 
   transformResults.set(getTiledLinalgOp().cast<OpResult>(), tiled);
   for (const auto &en : llvm::enumerate(loops))
     transformResults.set(getLoops()[en.index()].cast<OpResult>(), en.value());
 

mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp

mlir::scf::tileUsingSCFForOp(RewriterBase &rewriter, TilingInterface op,
   });
 
   // 4. Generate the tiled implementation within the inner most loop.
   if (!tilingResult.loops.empty())
     rewriter.setInsertionPoint(
         tilingResult.loops.back().getBody()->getTerminator());
   SmallVector<Operation *> tiledImplementation =
       op.getTiledImplementation(rewriter, offsets, sizes);
-  if (tiledImplementation.size() != 1) {
-    return rewriter.notifyMatchFailure(
-        op, "expected tiled implementation to return a single op");
-  }
-  tilingResult.tiledOp = tiledImplementation[0];
+  tilingResult.tiledOps.append(tiledImplementation);
   if (op->getNumResults() == 0) {
     // nothing more to do.
     return tilingResult;
   }
 
   // If loops are empty, the tiled op is used as the replacement for the untiled
   // op.
   if (tilingResult.loops.empty()) {
if (failed(op.getResultTilePosition(rewriter, result.index(), offsets,
                                         resultOffsetsList[result.index()],
                                         resultSizesList[result.index()]))) {
       return rewriter.notifyMatchFailure(
           op, "failed to get slice of result produced");
     }
   }
 
   FailureOr<SmallVector<Value>> replacementOr = yieldTiledValues(
-      rewriter, destinationTensors, tilingResult.tiledOp->getResults(),
+      rewriter, destinationTensors, tilingResult.tiledOps.back()->getResults(),
       resultOffsetsList, resultSizesList, tilingResult.loops);
   if (failed(replacementOr))
     return rewriter.notifyMatchFailure(op, "failed to yield replacement");
 
   if (auto dstOp =
-          dyn_cast<DestinationStyleOpInterface>(tilingResult.tiledOp)) {
+          dyn_cast<DestinationStyleOpInterface>(tilingResult.tiledOps.back())) {
     auto innerMostLoop = tilingResult.loops.back();
     SmallVector<Value> destinationTensors = dstOp.getDpsInitOperands();
     assert(destinationTensors.size() ==
                innerMostLoop.getRegionIterArgs().size() &&
            "unexpected number of outputs");
     updateDestinationOperandsForTiledOp(rewriter, destinationTensors,
                                         innerMostLoop.getRegionIterArgs());
   }
mlir::scf::tileConsumerAndFuseProducerGreedilyUsingSCFForOp(
   // 1. First tile the consumer.
   scf::SCFTileAndFuseResult tileAndFuseResult;
   llvm::SmallDenseMap<Value, int64_t> yieldedValueToResultNumber;
   {
     FailureOr<scf::SCFTilingResult> tilingResult =
         tileUsingSCFForOp(rewriter, consumer, options.tilingOptions);
     if (failed(tilingResult))
       return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
-    tileAndFuseResult.tiledAndFusedOps.insert(tilingResult->tiledOp);
+    for (auto tiledOp : tilingResult->tiledOps)
+      tileAndFuseResult.tiledAndFusedOps.insert(tiledOp);
     tileAndFuseResult.loops = std::move(tilingResult->loops);
     for (const auto &result : llvm::enumerate(
              llvm::zip(consumer->getResults(), tilingResult->replacements))) {
       tileAndFuseResult.replacements[std::get<0>(result.value())] =
           std::get<1>(result.value());
-      yieldedValueToResultNumber[tilingResult->tiledOp->getResult(
+      yieldedValueToResultNumber[tilingResult->tiledOps.back()->getResult(
           result.index())] = result.index();
     }
   }
 
   // If there are no loops generated, fusion is immaterial.
   if (tileAndFuseResult.loops.empty())
     return tileAndFuseResult;
 

mlir/test/lib/Interfaces/TilingInterface/TestTilingInterface.cpp

if (failed(tilingResult))
       return rewriter.notifyMatchFailure(op, "failed to tile operation");
 
     if (op->getNumResults()) {
       rewriter.replaceOp(op, tilingResult->replacements);
     } else {
       rewriter.eraseOp(op);
     }
 
-    filter.replaceLinalgTransformationFilter(rewriter, tilingResult->tiledOp);
+    for (auto tiledOp : tilingResult->tiledOps)
+      filter.replaceLinalgTransformationFilter(rewriter, tiledOp);
     return success();
   }
 
 private:
   scf::SCFTilingOptions options;
   LinalgTransformationFilter filter;
 };