[mlir][Linalg] Relax PadTensor tiling constraints and expose it to strategies.

[nicolasvasilache]

Rebase.

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

 /// reshapes.
 bool skipUnitDimReshape(const OpResult &producer, OpOperand &consumer);
 
 //===----------------------------------------------------------------------===//
 // Transformations exposed as function calls.
 //===----------------------------------------------------------------------===//
 using LinalgLoops = SmallVector<Operation *, 4>;
 
+void populatePadTensorTilingPatterns(RewritePatternSet &patterns,
+                                     const LinalgTilingOptions &options);
+
 /// Populate patterns for vectorizing low-D convolution ops. This is a step in
 /// progressive lowering for convolution ops, it assume high-D convolution ops
 /// were decomposed previously.
 void populateConvolutionVectorizationPatterns(RewritePatternSet &patterns,
                                               PatternBenefit benefit = 1);
 
 /// Populate patterns that convert `ElementwiseMappable` ops to linalg
 /// parallel loops.

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

void runOnFunction() override {
 
     MLIRContext *ctx = funcOp.getContext();
     RewritePatternSet tilingPattern(ctx);
     if (!anchorOpName.empty())
       tilingPattern.add<LinalgTilingPattern>(anchorOpName, ctx, options,
                                              filter);
     else
       tilingPattern.add<LinalgTilingPattern>(ctx, options, filter);
+    if (anchorOpName == linalg::PadTensorOp::getOperationName())
+      populatePadTensorTilingPatterns(tilingPattern, options);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
   }
 
   LinalgTilingOptions options;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to apply hoisting and padding.

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

static LogicalResult tilePadTensorOp(RewriterBase &builder, PadTensorOp op,
   builder.setInsertionPoint(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);
+  // Normalize untiled padding dimensions to 0.
+  Value zero = builder.create<arith::ConstantIndexOp>(loc, 0);
+  tileSizes.append(rank - tileSizes.size(), zero);
   // Compute lower and upper bounds of the loop nest.
   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);
LinalgTransformationFilter f(ArrayRef<StringAttr>{},
                                StringAttr::get(ctx, "tiled"));
   TilingPatterns<GenericOp,
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
                  >::insert(patterns, options, f);
   patterns.add<PadTensorOpTilingPattern>(ctx, options);
 }
 
+void mlir::linalg::populatePadTensorTilingPatterns(
+    RewritePatternSet &patterns, const LinalgTilingOptions &options) {
+  auto *ctx = patterns.getContext();
+  patterns.add<PadTensorOpTilingPattern>(ctx, options);
+}
+
 static void applyExtractSliceOfPadTensorSwapPattern(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
   patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
 }

mlir/test/Dialect/Linalg/tile-pad-tensor-op.mlir

 // RUN: mlir-opt %s -linalg-tile="tile-sizes=2,3" -cse -split-input-file | \
 // RUN: FileCheck %s -check-prefix=TILE2
 // RUN: mlir-opt %s -linalg-tile="tile-sizes=0,3" -resolve-shaped-type-result-dims -cse -split-input-file | \
 // RUN: FileCheck %s -check-prefix=TILE1
+// This test only checks that tiling does not crash.
+// RUN: mlir-opt %s -linalg-tile="tile-sizes=2" -resolve-shaped-type-result-dims -cse -split-input-file
 
 //  TILE2-DAG:  #[[MAP0:.*]] = affine_map<()[s0] -> (s0 + 8)>
 //  TILE2-DAG:  #[[MAP1:.*]] = affine_map<()[s0] -> (s0 + 7)>
 //       TILE2: func @dynamic_pad_tensor(
 //  TILE2-SAME:     %[[IN:.*]]: tensor<?x?xf32>
 //   TILE2-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   TILE2-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   TILE2-DAG:   %[[C2:.*]] = arith.constant 2 : index