[mlir][interfaces] Remove getDestinationOperands from TilingInterface

[springerm]

Summary

getDestinationOperands was almost a duplicate of DestinationStyleOpInterface::getOutputOperands. Now that the interface has been moved to mlir/Interfaces, it is no longer needed.

[springerm]

update

[pifon2a]

Nice

[mravishankar]

Thanks for doing this Matthias. Have a comment on the API/requirements to use the API though.

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

135

I see that this methods asserts if the defining op does not implement one of two interfaces.... Maybe we can make this return an Optional<Value> of FailureOr<Value> to avoid the assert.
At the very least please add the assert requirement in the comments here.

139

This one is a bit harder to avoid the assert. One options is

LogicalResult getOrCreateDestinations(OpBuilder &b, Location loc, Operation *op, SmallVectorImpl<Value> &destinationValues);

Makes the two APIs incosistent, but I am thinking loud here. Having an assert though is the least desirable solution.

[mravishankar]

Also please drop the NFC tag. It is a breaking change...

[springerm]

address comments

[mravishankar]

Have a few more comments. With that this should be good to go.

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

141

Nit: Usually return by value is at the end of the arguments list.

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

382

Do you mind moving this to the start. Basically make this an early failure before any transformations are done. Without this it can get the pattern rewrite driver into an infinite loop (triggering the fails safe of returning failure after 10 iterations or so).

393

This doesnt need to be a failure I think. You can just replace lines 390 onwards to

if (auto destinationPassingOp = dyn_cast<DestinationStyleOpInterface>(...)) {
  SmallVector<Value> destination;
  if (succeeded(getOrCreateDestination(...)) {
    updateDestinationOperandsForTiledOp(...)
  }
}

577

Replace lines 578 onwards with the same schema as line 390 onwards above.

581

Replace lines 588 onwards with the same schema as line 390 onwards above.

[springerm]

address comments

[mravishankar]

Thanks!

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

 
 /// Create a rank-reducing InsertSliceOp @[0 .. 0] with strides [1 .. 1] and
 /// appropriate sizes (i.e. `dest.getSizes()`). The result is a new tensor with
 /// rank increased to that of `dest`, obtained by inserting `tensor` into `dest`
 /// at the canonical [0 .. 0] position.
 Value createCanonicalRankReducingInsertSliceOp(OpBuilder &b, Location loc,
                                                Value tensor, Value dest);
 
+/// This is a helper function DestinationStyleOpInterface. If there is a
+/// destination operand for the given OpResult, return that operand. Otherwise,
+/// return an empty tensor (`tensor.empty`) with the shape of the OpResult.
+Value getOrCreateDestination(OpBuilder &b, Location loc, OpResult opResult);
+
+/// This is a helper function DestinationStyleOpInterface. Get or create
+/// destinations for every tensor OpResult of the given op.
+SmallVector<Value> getOrCreateDestinations(OpBuilder &b, Location loc,
+                                           Operation *op);
+
 /// Function to control the folding of constant and extract slice
 using ControlConstantExtractSliceFusionFn = std::function<bool(ExtractSliceOp)>;
 
 /// Patterns to fold the extract slice op with its constant operand
 void populateFoldConstantExtractSlicePatterns(
     RewritePatternSet &patterns,
     const ControlConstantExtractSliceFusionFn &controlFn =
         [](ExtractSliceOp op) {

mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td

let methods = [
     >,
     InterfaceMethod<
       /*desc=*/"Return true if the `opOperand` is a scalar value.",
       /*retTy=*/"bool",
       /*methodName=*/"isScalar",
       /*args=*/(ins "OpOperand *":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        assert(opOperand->getOwner() == this->getOperation());
+        assert(opOperand->getOwner() == $_op.getOperation());
         return !opOperand->get().getType().template isa<ShapedType>();
       }]
     >,
     InterfaceMethod<
-      /*desc=*/"Return the result tied to `opOperand`.",
+      /*desc=*/"Return the OpResult that is tied to the given OpOperand.",
       /*retTy=*/"OpResult",
       /*methodName=*/"getTiedOpResult",
       /*args=*/(ins "OpOperand *":$opOperand),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        assert(opOperand->getOwner() == this->getOperation());
+        assert(opOperand->getOwner() == $_op.getOperation());
 
         auto [start, end] = $_op.getOutputsPositionRange();
         int64_t resultIndex = opOperand->getOperandNumber() - start;
         assert(resultIndex >= 0 &&
                resultIndex < $_op->getNumResults() );
         return $_op->getResult(resultIndex);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/"Return the OpOperand that is tied to the given OpResult.",
+      /*retTy=*/"OpOperand *",
+      /*methodName=*/"getTiedOpOperand",
+      /*args=*/(ins "OpResult":$opResult),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        assert(opResult.getDefiningOp() == $_op.getOperation());
+        return $_op.getOutputOperand(opResult.getResultNumber());
+      }]
+    >,
     //===------------------------------------------------------------------===//
     // Other interface methods.
     //===------------------------------------------------------------------===//
     InterfaceMethod<
       /*desc=*/"Return whether the op has only MemRef input and outputs.",
       /*retTy=*/"bool",
       /*methodName=*/"hasBufferSemantics",
       /*args=*/(ins),

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.
-        }],
-        /*retType=*/"SmallVector<Value>",
-        /*methodName=*/"getDestinationOperands",
-        /*args=*/(ins "OpBuilder &":$b),
-        /*methodBody=*/"",
-        /*defaultImplementation=*/"return ValueRange{};"
-      >,
-      InterfaceMethod<
-        /*desc=*/[{
           Returns a list of iterator types that describe the number of loops.
         }],
         /*retType=*/"SmallVector<utils::IteratorType>",
         /*methodName=*/"getLoopIteratorTypes",
         /*args=*/(ins),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return {};"
       >,

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

static Operation *tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(
   OpBuilder::InsertionGuard guard(rewriter);
   rewriter.setInsertionPoint(sliceOpToTile);
 
   // Replace the use in the tileableProducer before tiling: clone, replace and
   // then tile.
   int64_t resultNumber = pUse->get().cast<OpResult>().getResultNumber();
   LLVM_DEBUG(llvm::dbgs() << "resultNumber: " << resultNumber << "\n");
 
-  auto destinationOperands = tileableProducer.getDestinationOperands(rewriter);
+  // Gather destination tensors.
+  SmallVector<Value> destinationTenors = tensor::getOrCreateDestinations(
+      rewriter, tileableProducer->getLoc(), tileableProducer);
 
   BlockAndValueMapping bvm;
-  bvm.map(destinationOperands[resultNumber], bbArg);
+  bvm.map(destinationTenors[resultNumber], bbArg);
   auto tileableProducerClone =
       cast<TilingInterface>(rewriter.clone(*tileableProducer, bvm));
   auto scopeGuard =
       llvm::make_scope_exit([&]() { rewriter.eraseOp(tileableProducerClone); });
 
   // Tile the producer.
   FailureOr<Value> tiledProducer =
       tileableProducerClone.generateResultTileValue(
           rewriter, resultNumber, sliceOpToTile.getMixedOffsets(),
           sliceOpToTile.getMixedSizes());
   if (failed(tiledProducer)) {
     diag.attachNote(tileableProducer->getLoc())
         << "failed to tile producer op: " << *tileableProducer;
     return nullptr;
   }
   LLVM_DEBUG(llvm::dbgs() << "tiledProducer: " << *tiledProducer << "\n");
 
   // Replace the extract op.
   Operation *fusedOp = tiledProducer->getDefiningOp();
   rewriter.replaceOp(sliceOpToTile, fusedOp->getResult(resultNumber));
 
   // Replace the use in containingOp.
   rewriter.updateRootInPlace(containingOp, [&]() {
     containingOp->setOperand(pUse->getOperandNumber(),
-                             destinationOperands.front());
+                             destinationTenors.front());
   });
 
   return fusedOp;
 }
 
 static Operation *cloneAndFuseFirstUse(RewriterBase &rewriter, Diagnostic &diag,
                                        Operation *producerOp,
                                        Operation *containingOp) {

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

OpFoldResult remainingSize = makeComposedFoldedAffineApply(
       rewriter, op.getLoc(), d0 + d1 - d2,
       {iterationSpace[dimension].offset, iterationSpace[dimension].size,
        minSplitPoint});
   if (auto attr = remainingSize.dyn_cast<Attribute>()) {
     if (attr.cast<IntegerAttr>().getValue().isZero())
       return {op, TilingInterface()};
   }
 
+  // Compute destination tensors.
+  SmallVector<Value> destinationTensors =
+      tensor::getOrCreateDestinations(rewriter, op.getLoc(), op);
+
   // 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, destinationTensors, 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(
   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");
-  auto dest = op.getDestinationOperands(b);
+
+  // Gather destination tensors.
+  SmallVector<Value> dest = tensor::getOrCreateDestinations(b, loc, op);
 
   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) {
         return getValueOrCreateConstantIndexOp(b, loc, ofr);
if (!isZero(tileSizes[i])) {
       lbs.push_back(
           getValueOrCreateConstantIndexOp(builder, loc, ranges[i].offset));
       dims.push_back(
           getValueOrCreateConstantIndexOp(builder, loc, ranges[i].size));
       steps.push_back(
           getValueOrCreateConstantIndexOp(builder, loc, tileSizes[i]));
     }
   }
-  // Generate loop nest: One loop per dimension.
-  SmallVector<Value> destOperand =
-      tilingInterface.getDestinationOperands(builder);
+  SmallVector<Value> destinationTensors =
+      tensor::getOrCreateDestinations(builder, loc, tilingInterface);
+
   loopNest = mlir::scf::buildLoopNest(
-      builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destOperand),
+      builder, loc, lbs, /*ubs=*/dims, steps, ValueRange(destinationTensors),
       [&](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

 /// model implementation is used for now till the use of `TilingInterface` is
 /// 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 {
-    return cast<LinalgOp>(op).getOutputOperands();
-  }
-
   /// Return the loop iterator type.
   SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     LinalgOpTy concreteOp = cast<LinalgOpTy>(op);
     return llvm::to_vector(llvm::map_range(
         concreteOp.getIteratorTypesArray(), [](StringRef iteratorType) {
           return utils::symbolizeIteratorType(iteratorType).value();
         }));
   }

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

if (failed(op.getResultTilePosition(rewriter, result.index(), offsets,
                                         sizes,
                                         resultOffsetsList[result.index()],
                                         resultSizesList[result.index()]))) {
       return rewriter.notifyMatchFailure(
           op, "failed to get slice of result produced");
     }
   }
 
+  SmallVector<Value> dest =
+      tensor::getOrCreateDestinations(rewriter, op.getLoc(), op);
   FailureOr<SmallVector<Value>> replacementOr =
-      yieldTiledValues(rewriter, op.getDestinationOperands(rewriter),
-                       tilingResult.tiledOp->getResults(), resultOffsetsList,
-                       resultSizesList, tilingResult.loops);
+      yieldTiledValues(rewriter, dest, tilingResult.tiledOp->getResults(),
+                       resultOffsetsList, resultSizesList, tilingResult.loops);
   if (failed(replacementOr))
     return rewriter.notifyMatchFailure(op, "failed to yield replacement");
+
   if (auto tiledInterfaceOp = dyn_cast<TilingInterface>(tilingResult.tiledOp)) {
     auto innerMostLoop = tilingResult.loops.back();
-    updateDestinationOperandsForTiledOp(
-        rewriter, tiledInterfaceOp.getDestinationOperands(rewriter),
+    // Gather destination tensors.
+    SmallVector<Value> dest = tensor::getOrCreateDestinations(
+        rewriter, tiledInterfaceOp.getLoc(), tiledInterfaceOp);
+    updateDestinationOperandsForTiledOp(rewriter, dest,
                                         innerMostLoop.getRegionIterArgs());
   }
 
   tilingResult.replacements = replacementOr.value();
 
   LLVM_DEBUG({
     if (!tilingResult.loops.empty()) {
       llvm::dbgs() << "After tiled implementation :\n";
       tilingResult.loops.front().dump();
while (!candidates.empty()) {
     if (destinationIterArg) {
       iterArgNumber = outerMostLoop.getIterArgNumberForOpOperand(
           *destinationIterArg.value());
     }
     if (iterArgNumber) {
       int64_t resultNumber = fusableProducer.getResultNumber();
       if (auto producerOp =
               dyn_cast<TilingInterface>(fusableProducer.getOwner())) {
-        SmallVector<Value> destination =
-            producerOp.getDestinationOperands(rewriter);
+        SmallVector<Value> destination = tensor::getOrCreateDestinations(
+            rewriter, producerOp.getLoc(), producerOp);
         outerMostLoop.setIterArg(iterArgNumber.value(),
                                  destination[resultNumber]);
       }
       if (auto tiledAndFusedInterfaceOp =
               fusedProducerValue.value().getDefiningOp<TilingInterface>()) {
         scf::ForOp innerMostLoop = tileAndFuseResult.loops.back();
-        SmallVector<Value> destination =
-            tiledAndFusedInterfaceOp.getDestinationOperands(rewriter);
+        SmallVector<Value> destination = tensor::getOrCreateDestinations(
+            rewriter, tiledAndFusedInterfaceOp.getLoc(),
+            tiledAndFusedInterfaceOp);
         updateDestinationOperandsForTiledOp(
             rewriter, destination[resultNumber],
             innerMostLoop.getRegionIterArgs()[iterArgNumber.value()]);
       }
     }
   }
   return tileAndFuseResult;
 }

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRTensorDialect
   Core
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRArithDialect
   MLIRArithUtils
   MLIRCastInterfaces
   MLIRComplexDialect
+  MLIRDestinationStyleOpInterface
   MLIRDialectUtils
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRParallelCombiningOpInterface
   MLIRShapedOpInterfaces
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRViewLikeInterface

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/StringRef.h"
 #include <algorithm>
 
 using namespace mlir;
 using namespace mlir::tensor;
if (tensorType.isDynamicDim(i)) {
       result.push_back(size);
     } else {
       result.push_back(builder.getIndexAttr(tensorType.getDimSize(i)));
     }
   }
   return result;
 }
 
+Value tensor::getOrCreateDestination(OpBuilder &b, Location loc,
+                                     OpResult opResult) {
+  auto tensorType = opResult.getType().dyn_cast<TensorType>();
+  assert(tensorType && "expected tensor type");
+
+  // If the op has a destination, it implements DestinationStyleOpInterface and
+  // we can query the destination operand from that interface.
+  auto destOp = opResult.getDefiningOp<DestinationStyleOpInterface>();
+  if (destOp)
+    return destOp.getTiedOpOperand(opResult)->get();
+
+  // Otherwise, create a new destination tensor with the same shape.
+  OpBuilder::InsertionGuard g(b);
+  b.setInsertionPoint(opResult.getDefiningOp());
+  ReifiedRankedShapedTypeDims reifiedShapes;
+  ReifyRankedShapedTypeOpInterface reifyShapedTypeInterface =
+      dyn_cast<ReifyRankedShapedTypeOpInterface>(opResult.getDefiningOp());
+  assert(reifyShapedTypeInterface &&
+         "creating a destination requires reifying the shape");
+  if (failed(reifyShapedTypeInterface.reifyResultShapes(b, reifiedShapes)))
+    llvm_unreachable("failed to reify result shapes");
+
+  SmallVector<OpFoldResult> mixedSizes =
+      getAsOpFoldResult(reifiedShapes[opResult.getResultNumber()]);
+  Value initTensor =
+      b.create<tensor::EmptyOp>(loc, mixedSizes, tensorType.getElementType());
+  return initTensor;
+}
+
+SmallVector<Value> tensor::getOrCreateDestinations(OpBuilder &b, Location loc,
+                                                   Operation *op) {
+  SmallVector<Value> result;
+  for (OpResult opResult : op->getResults())
+    if (opResult.getType().isa<TensorType>())
+      result.push_back(getOrCreateDestination(b, loc, opResult));
+  return result;
+}
+
 //===----------------------------------------------------------------------===//
 // CastOp
 //===----------------------------------------------------------------------===//
 
 void CastOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "cast");
 }
 

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

 
 using namespace mlir;
 using namespace mlir::tensor;
 
 namespace {
 
 struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
 
-  SmallVector<Value> getDestinationOperands(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);
-    SmallVector<OpFoldResult> mixedSizes = getAsOpFoldResult(reifiedShapes[0]);
-    Value emptyTensor = b.create<EmptyOp>(
-        op->getLoc(), mixedSizes, padOp.getResultType().getElementType());
-    return {emptyTensor};
-  }
-
   SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
     auto padOp = cast<PadOp>(op);
     SmallVector<utils::IteratorType> iteratorTypes(
         padOp.getResultType().getRank(), utils::IteratorType::parallel);
     return iteratorTypes;
   }
 
   SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":ArithDialect",
         ":ArithUtils",
         ":CastOpInterfaces",
         ":ComplexDialect",
         ":ControlFlowInterfaces",
+        ":DestinationStyleOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":ParallelCombiningOpInterface",
         ":ShapedOpInterfaces",
         ":SideEffectInterfaces",
         ":TensorOpsIncGen",
         ":TilingInterface",