[mlir][Linalg] NFC - Extend the TilingInterface to allow better composition with out-of-tree dialects.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

def Linalg_InitTensorOp : Linalg_Op<"init_tensor",
 
   let hasCanonicalizer = 1;
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
     [AttrSizedOperandSegments, NoSideEffect,
      DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      DeclareOpInterfaceMethods<TilingInterface,
-         ["getDestinationOperands", "getLoopIteratorTypes", "getLoopBounds",
+         ["getDestinationOperands", "getLoopIteratorTypes", "getIterationDomain",
           "getTiledImplementation"]>]> {
   let summary = "tensor pad operation";
   let description = [{
     `linalg.pad_tensor` is an operation that pads the `source` tensor
     with given `low` and `high` padding config.
 
     The PadTensor operation supports the following arguments:
 

mlir/include/mlir/Interfaces/TilingInterface.td

let methods = [
         /*defaultImplementation=*/"return {};"
       >,
       InterfaceMethod<
         /*desc=*/[{
           Returns a list of ranges that describe the loop bounds and
           step for the loops of the operation.
         }],
         /*retTy=*/"SmallVector<Range>",
-        /*methodName=*/"getLoopBounds",
+        /*methodName=*/"getIterationDomain",
         /*args=*/(ins "OpBuilder &":$b),
         /*methodBody=*/"",
         /*defaultImplementation=*/"return {};"
       >,
       InterfaceMethod<
         /*desc=*/[{
           Method to generate the tiled implementation of an operation.
 
           The iteration space of the operation is returned by
-          `getLoopBounds`. The caller provides the information of the
+          `getIterationDomain`. The caller provides the information of the
           tile within this iteration space whose implementation the
           caller needs.
           - `dest` are the Value into which the result of the tiled
             operation is to be inserted into. The type of the `dest`
             Values is same as the types returned by
             `getDestinationOperands` method.
           - `offsets` provides the offset of the tile within the
             iteration space
           - `sizes` provides the size of the tile.
+          - `tileDestOperands` specifies whether to also tile `dest` operands
+            or not. Avoiding tiling `dest` operands can be useful for 
+            composition with various looping container ops.
 
           The method returns the operation that is the tiled
           implementation.
         }],
-        /*retType=*/"Operation *",
+        /*retType=*/"SmallVector<Operation *>",
         /*methodName=*/"getTiledImplementation",
         /*args=*/(ins
             "OpBuilder &":$b,
             "ValueRange ":$dest,
             "ArrayRef<OpFoldResult> ":$offsets,
-            "ArrayRef<OpFoldResult> ":$sizes),
+            "ArrayRef<OpFoldResult> ":$sizes,
+            "bool ":$tileDestOperands),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
-          return nullptr;
+          return {};
         }]
       >
   ];
 }
 #endif // MLIR_TILINGINTERFACE

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 }
 
 SmallVector<StringRef> PadTensorOp::getLoopIteratorTypes() {
   SmallVector<StringRef> iteratorTypes(getResultType().getRank(),
                                        getParallelIteratorTypeName());
   return iteratorTypes;
 }
 
-SmallVector<Range> PadTensorOp::getLoopBounds(OpBuilder &b) {
+SmallVector<Range> PadTensorOp::getIterationDomain(OpBuilder &b) {
   ReifiedRankedShapedTypeDims reifiedShapes;
   (void)reifyResultShapes(b, reifiedShapes);
   Value zero = b.create<arith::ConstantIndexOp>(getLoc(), 0);
   Value one = b.create<arith::ConstantIndexOp>(getLoc(), 1);
   // Initialize all the ranges to {zero, one, one}. All the `ub`s are
   // overwritten.
   SmallVector<Range> loopRanges(reifiedShapes[0].size(), {zero, one, one});
   for (auto ub : enumerate(reifiedShapes[0]))
     loopRanges[ub.index()].size = ub.value();
   return loopRanges;
 }
 
-Operation *PadTensorOp::getTiledImplementation(OpBuilder &b, ValueRange dest,
-                                               ArrayRef<OpFoldResult> offsets,
-                                               ArrayRef<OpFoldResult> sizes) {
+SmallVector<Operation *> PadTensorOp::getTiledImplementation(
+    OpBuilder &b, ValueRange dest, ArrayRef<OpFoldResult> offsets,
+    ArrayRef<OpFoldResult> sizes, bool /*tileDestOperands*/) {
   // Only constant padding value supported.
   Value padValue = getConstantPaddingValue();
   if (!padValue)
-    return nullptr;
+    return {};
 
   // Helper variables and functions for various arithmetic operations. These are
   // used extensively for computing new offset/length and padding values.
   Location loc = getLoc();
   AffineExpr dim0, dim1;
   bindDims(b.getContext(), dim0, dim1);
   // Add two integers.
   auto addMap = AffineMap::get(2, 0, {dim0 + dim1});
auto createPadTensorOfSubTensor = [&]() {
 
     // Cast result and return.
     return castResult(newPadTensorOp);
   };
 
   // Rewrite subtensor(pad_tensor(x)) into a GenerateOp it is statically known
   // that the original data source x is not used.
   if (hasZeroLen) {
-    return createGenerateOp();
+    return {createGenerateOp()};
   }
 
   // If there are dynamic dimensions: Generate an scf.if check to avoid creating
   // SliceOps with result dimensions of size 0 at runtime.
   if (dynHasZeroLenCond) {
     auto result = b.create<scf::IfOp>(
         loc, resultType, dynHasZeroLenCond,
         /*thenBuilder=*/
         [&](OpBuilder &b, Location loc) {
           b.create<scf::YieldOp>(loc, createGenerateOp()->getResult(0));
         },
         /*elseBuilder=*/
         [&](OpBuilder &b, Location loc) {
           b.create<scf::YieldOp>(loc,
                                  createPadTensorOfSubTensor()->getResult(0));
         });
-    return result;
+    return {result};
   }
-  return createPadTensorOfSubTensor();
+  return {createPadTensorOfSubTensor()};
 }
 
 namespace {
 // Folds linalg.pad_tensor when padding is static zeros and the attribute
 // doesn't request otherwise.
 struct FoldStaticZeroPadding : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 

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

static LogicalResult tilePadTensorOp(OpBuilder &builder, PadTensorOp op,
   // Clone PadTensorOp so that the existing op can be replaced more easily.
   newPadOp = cast<PadTensorOp>(builder.clone(*op.getOperation()));
   // Get rank and tile sizes.
   int64_t rank = op.getResultType().getRank();
   SmallVector<Value> tileSizes =
       options.tileSizeComputationFunction(builder, op);
   assert(static_cast<int64_t>(tileSizes.size()) == rank);
   // Compute lower and upper bounds of the loop nest.
-  SmallVector<Range> ranges = op.getLoopBounds(builder);
+  SmallVector<Range> ranges = op.getIterationDomain(builder);
   SmallVector<Value> lbs, dims, allDims, steps;
   for (int64_t i = 0; i < rank; ++i) {
     allDims.push_back(ranges[i].size);
     if (!isZero(tileSizes[i])) {
       lbs.push_back(ranges[i].offset);
       dims.push_back(ranges[i].size);
       steps.push_back(tileSizes[i]);
     }

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

LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(
     tensor::ExtractSliceOp sliceOp, PatternRewriter &rewriter) const {
   auto padOp = sliceOp.source().getDefiningOp<PadTensorOp>();
   if (!padOp)
     return failure();
   // Only unit stride supported.
   if (!sliceOp.hasUnitStride())
     return failure();
 
-  Operation *tiledPadOp = padOp.getTiledImplementation(
+  Operation *tiledPadOp =
+      padOp
+          .getTiledImplementation(
               rewriter, /*dest=*/ValueRange{}, sliceOp.getMixedOffsets(),
-      sliceOp.getMixedSizes());
+              sliceOp.getMixedSizes(), /*tileDestOperands=*/false)
+          .front();
   // All shapes are static and the data source is actually used. Rewrite into
   // pad_tensor(subtensor(x)).
   rewriter.replaceOp(sliceOp, tiledPadOp->getResults());
   return success();
 }
 
 namespace {
 // The following are patterns for downscaling convolution ops with size-1