[mlir][interfaces] Return destination OpOperands instead of Values

[springerm]

Summary

There are now two methods for retrieving destinations: getDestinationOpOperands and getDestination. The former one returns the OpOperands of the op. The latter one is allowed to create new IR in case the op does not have destination operands.

This change is in preparation of explicitly capturing shared tensors in scf.foreach_thread. Inside of the foreach_thread body, all destinations must be expressed in terms of the block arguments of the foreach_thread body.

Note: getDestinationOpOperands is identical to DestinationStyleOpInterface::getOutputOperands, but we cannot use DestinationStyleOpInterface yet because it was not yet moved to mlir/Interfaces.

[mravishankar]

I am not sure this change composes well. Specifically, there is no way an transformation can know if it is supposed to use getDestinationOperands or getDestinationOpOperands. I think eventually even getDestinationOperands needs to be dropped from the TilingInterface.

[springerm]

In D133112#3764299, @mravishankar wrote:

Specifically, there is no way an transformation can know if it is supposed to use getDestinationOperands or getDestinationOpOperands.

A transformation should use getDestinationOperands (now called getDestination). getDestinationOpOperands should be called only if it is important to know which OpOperand is the destination.

I think eventually even getDestinationOperands needs to be dropped from the TilingInterface.

There is a new op interface that was added recently: DestinationStyleOpInterface. This interface would be sufficient for my needs and I think can proceed with that interface instead. (So I can probably abandon this change.)

I'm wondering if DestinationStyleOpInterface could also be a full replacement for getDestinationOperands/getDestinationOpOperands. The main problem is that getDestinationOperands takes an OpBuilder and creates a tensor if there is no destination operand (e.g., tensor.pad). Do we create the exact same IR for every op without a destination (i.e., reify result shape + linalg.init_tensor)? In that case, we can indeed drop getDestinationOperands. I'll see if that works (as a further cleanup of the TilingInterface), but let me know if you see any issues with that approach.

[mravishankar]

You are probably right that reify result shape + init Tensor is probably what you will be needed. But that would add dependence of scf to Linalg.

[mravishankar]

@springerm I actually did find the use for this. Do you want to land this while you explore the other path... Might be less churn if you dont change the name of the method. So keep getDestinationOperands instead of renaming to getDestination?

mlir/include/mlir/Interfaces/TilingInterface.td

let description = [{
     Interface for allowing operations to expose information needed to
     tile them (similar to LinalgOp, but without having access to
     indexing maps)
   }];
   let cppNamespace = "::mlir";
   let methods = [
       InterfaceMethod<
         /*desc=*/[{
-          Returns a list of operands into which the result of the
-          tiled implementation is written into. With `tensor`
-          operands, this will be used as the initial tensor into which
-          the tiled results are inserted into. With `memref` operands,
-          this will be the operand into which the result of the tiled
-          operation is written into.
+          Returns all destination (output) OpOperands of the operation.
+
+          TODO: Use DestinationStyleOpInterface in the future.
+        }],
+        /*retType=*/"SmallVector<OpOperand *>",
+        /*methodName=*/"getDestinationOpOperands",
+        /*args=*/(ins),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          llvm_unreachable("getDestinationOpOperands not implemented");
+        }]
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Returns a list of values into which the result of the tiled
+          implementation is written into.
+          
+          With `tensor` values, these will be used as the initial tensors into
+          which the tiled results are inserted into. With `memref` values, these
+          will be the buffers into which the results of the tiled operation is
+          written into.
+
+          By default, this method returns the destination operands of the
+          operation. Operations that do not have destination operands can
+          override this method to provide custom destination values.
         }],
         /*retType=*/"SmallVector<Value>",
-        /*methodName=*/"getDestinationOperands",
+        /*methodName=*/"getDestination",
         /*args=*/(ins "OpBuilder &":$b),
         /*methodBody=*/"",
-        /*defaultImplementation=*/"return ValueRange{};"
+        /*defaultImplementation=*/[{
+          auto tileableOp = cast<TilingInterface>($_op.getOperation());
+          SmallVector<OpOperand *> opOperands =
+              tileableOp.getDestinationOpOperands();
+          return llvm::to_vector(llvm::map_range(
+              opOperands, [](OpOperand *operand) { return operand->get(); }));
+        }]
       >,
       InterfaceMethod<
         /*desc=*/[{
           Returns a list of `StringRef`s that describe the number of
           loops and the iterator types of the operation. The list is
           expected to use
           `getParallelIteratorTypeName()`/`getReductionIteratorTypeName()`
           from MLIR Structured Op Utils.

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

linalg::splitOp(RewriterBase &rewriter, TilingInterface op, unsigned dimension,
   if (auto attr = remainingSize.dyn_cast<Attribute>()) {
     if (attr.cast<IntegerAttr>().getValue().isZero())
       return {op, TilingInterface()};
   }
 
   // Create the first part.
   SmallVector<Value> firstResults;
   TilingInterface firstPart = createSplitPart(
-      rewriter, op.getLoc(), op, offsets, sizes,
-      op.getDestinationOperands(rewriter), dimension, minSplitPoint,
-      iterationSpace[dimension].offset, firstResults);
+      rewriter, op.getLoc(), op, offsets, sizes, op.getDestination(rewriter),
+      dimension, minSplitPoint, iterationSpace[dimension].offset, firstResults);
 
   // Need to pretend that the original op now takes as operands firstResults,
   // otherwise tiling interface implementation will take the wrong value to
   // produce data tiles.
   rewriter.updateRootInPlace(op, [&]() {
     unsigned numTotalOperands = op->getNumOperands();
     unsigned numOutputOperands = firstResults.size();
     op->setOperands(numTotalOperands - numOutputOperands, numOutputOperands,

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

static FailureOr<ForeachThreadTilingResult> tileToForeachThreadOpImpl(
   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");
+  auto dest = op.getDestination(b);
+  if (dest.size() != 1)
+    return op->emitOpError("only single dest 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);
static FailureOr<ForeachThreadTilingResult> tileToForeachThreadOpImpl(
   SmallVector<Operation *> tiledOps =
       op.getTiledImplementation(b, tiledOffsets, tiledSizes);
   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");
   OpBuilder::InsertPoint insertPt = b.saveInsertionPoint();
-  for (auto it :
-       llvm::zip(llvm::seq(unsigned(0), unsigned(destOperands.size())),
-                 tilingInterfaceOp->getResults(), destOperands)) {
+  for (auto it : llvm::zip(llvm::seq(unsigned(0), unsigned(dest.size())),
+                           tilingInterfaceOp->getResults(), dest)) {
     b.setInsertionPoint(insertPt.getBlock(), insertPt.getPoint());
     SmallVector<OpFoldResult> resultOffsets, resultSizes;
     if (failed(op.getResultTilePosition(b, std::get<0>(it), tiledOffsets,
                                         tiledSizes, resultOffsets,
                                         resultSizes)))
       return op->emitOpError("output offsets couldn't be calculated");
     SmallVector<OpFoldResult> strides(resultSizes.size(), b.getIndexAttr(1));
     b.setInsertionPointToStart(foreachThreadOp.getTerminator().getBody());
for (int64_t i = 0; i < rank; ++i) {
     allDims.push_back(ranges[i].size);
     if (!isZero(tileSizes[i])) {
       lbs.push_back(materializeOpFoldResult(builder, loc, ranges[i].offset));
       dims.push_back(materializeOpFoldResult(builder, loc, ranges[i].size));
       steps.push_back(materializeOpFoldResult(builder, loc, tileSizes[i]));
     }
   }
   // Generate loop nest: One loop per dimension.
-  SmallVector<Value> destOperand =
-      tilingInterface.getDestinationOperands(builder);
+  SmallVector<Value> dest = tilingInterface.getDestination(builder);
   loopNest = mlir::scf::buildLoopNest(
-      builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
+      builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(dest),
       [&](OpBuilder &b, Location loc, ValueRange localIvs,
           ValueRange iterArgs) -> scf::ValueVector {
         // Compute offsets and sizes of ExtractSliceOp.
         SmallVector<Value> localIVVector = llvm::to_vector(localIvs);
         SmallVector<OpFoldResult> offsets = computeTileOffsets(
             b, loc, getAsOpFoldResult(localIVVector), tileSizes);
         SmallVector<OpFoldResult> sizes =
             computeTileSizes(b, loc, tileSizes, allDims);

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

 /// on-par with the current Linalg tiling + fusion patterns. Once it is
 /// maybe possible to move this into the op-definition (though there are
 /// advantages to leaving it as an external model)
 template <typename LinalgOpTy>
 struct LinalgOpTilingInterface
     : public TilingInterface::ExternalModel<LinalgOpTilingInterface<LinalgOpTy>,
                                             LinalgOpTy> {
   /// Return the destination operands.
-  SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
+  SmallVector<OpOperand *> getDestinationOpOperands(Operation *op) const {
     return cast<LinalgOp>(op).getOutputOperands();
   }
 
   /// Return the loop iterator type.
   SmallVector<StringRef> getLoopIteratorTypes(Operation *op) const {
     LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
     return llvm::to_vector(
         llvm::map_range(concreteOp.iterator_types(), [](Attribute strAttr) {

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

for (auto resultNum : llvm::seq<unsigned>(0, op->getNumResults())) {
           op->getLoc(), tilingResult.tiledOp->getResult(resultNum),
           newBBArgs[resultNum], resultTileOffsets, resultTileSizes,
           resultTileStrides);
       yieldedValues.push_back(yieldedValue);
     }
     return yieldedValues;
   };
   SmallVector<scf::ForOp> newLoops = replaceLoopNestWithNewYields(
-      rewriter, tilingResult.loops, op.getDestinationOperands(rewriter),
-      yieldValueFn);
+      rewriter, tilingResult.loops, op.getDestination(rewriter), yieldValueFn);
   for (const auto &loop : llvm::enumerate(tilingResult.loops)) {
     rewriter.eraseOp(loop.value());
     tilingResult.loops[loop.index()] = newLoops[loop.index()];
   }
   rewriter.replaceOp(op, tilingResult.loops.front().getResults());
   return tilingResult;
 }
 
while (!candidates.empty()) {
     //        ...
     //        %2 = linalg.fill ...
     //        %3 = linalg.matmul ... outs(%2: ...)...
     //     ```
     TilingInterface unfusedProducerOp =
         cast<TilingInterface>(fusableProducer->getOwner());
     scf::ForOp outerMostTiledLoop = tileAndFuseResult.loops.front();
     SmallVector<Value> unfusedProducerOpDestValues =
-        unfusedProducerOp.getDestinationOperands(rewriter);
+        unfusedProducerOp.getDestination(rewriter);
     for (OpOperand &uses : unfusedProducerOp->getUses()) {
       if (uses.getOwner() == outerMostTiledLoop.getOperation()) {
         unsigned resultNumber = uses.get().cast<OpResult>().getResultNumber();
         unsigned operandNumber = uses.getOperandNumber();
         outerMostTiledLoop->setOperand(
             operandNumber, unfusedProducerOpDestValues[resultNumber]);
       }
     }

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

 #include "mlir/Interfaces/TilingInterface.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 namespace {
 
 struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
+  SmallVector<OpOperand *> getDestinationOpOperands(Operation *op) const {
+    return {};
+  }
 
-  SmallVector<Value> getDestinationOperands(Operation *op, OpBuilder &b) const {
+  SmallVector<Value> getDestination(Operation *op, OpBuilder &b) const {
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
     ReifiedRankedShapedTypeDims reifiedShapes;
     ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
         dyn_cast<ReifyRankedShapedTypeOpInterface>(op);
     (void)reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes);
 
     auto padOp = cast<PadOp>(op);