[mlir][Linalg] Use reify for padded op shape derivation.

[nicolasvasilache]

Summary

Previously, we would insert a DimOp and rely on later canonicalizations.
Unfortunately, reifyShape kind of rewrites are not canonicalizations anymore.
This introduces undesirable pass dependencies.

Instead, immediately reify the result shape and avoid the DimOp altogether.
This is akin to a local folding, which avoids introducing more reliance on -resolve-shaped-type-result-dims (similar to compositions of affine.apply by construction to avoid chains of size > 1).

It does not completely get rid of the reliance on the pass as the process is merely local: calling the pass may still be necessary for global effects. Indeed, one of the tests still requires the pass.

[springerm]

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

192–193

can this be an assertion?

203–204

can be deleted?

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

for (OpOperand *opOperand : opToPad.getInputAndOutputOperands()) {
     // If padding was requested but the shape cannot be bounded statically then
     // the pattern fails to apply.
     if (failed(padOperandToSmallestStaticBoundingBox(
             rewriter, opToPad, opOperand, paddingFunc, paddedOperand)))
       return failure();
     newOperands.push_back(paddedOperand ? paddedOperand : opOperand->get());
   }
 
+  SmallVector<SmallVector<Value>> reifiedResultShapes;
+  if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
+                 .reifyResultShapes(rewriter, reifiedResultShapes)))
+    return failure();
+  assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
+         "expected same number of results");
+
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
   paddedOp = opToPad.clone(rewriter, loc, resultTensorTypes, newOperands);
 
   // Recover the slice out of the new static results. This keeps the original
   // linalg op around because it uses the dims of the original results.
-  // This later folds away.
   SmallVector<Value> paddedSubviewResults;
   paddedSubviewResults.reserve(opToPad->getNumResults());
-  SetVector<Operation *> newUsersOfOpToPad;
-  for (auto it : llvm::zip(opToPad->getResults(), paddedOp->getResults())) {
-    auto rank = std::get<0>(it).getType().cast<RankedTensorType>().getRank();
+  for (auto en : llvm::enumerate(paddedOp->getResults())) {
+    Value paddedResult = en.value();
+    int64_t resultNumber = en.index();
+    int64_t rank = paddedResult.getType().cast<RankedTensorType>().getRank();
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
-    auto sizes = llvm::to_vector<4>(llvm::map_range(
-        llvm::seq<unsigned>(0, rank), [&](unsigned d) -> OpFoldResult {
-          auto dimOp = rewriter.create<tensor::DimOp>(loc, std::get<0>(it), d);
-          newUsersOfOpToPad.insert(dimOp);
-          return dimOp.getResult();
-        }));
+    SmallVector<OpFoldResult> sizes;
+    for (Value v : reifiedResultShapes[resultNumber])
+      sizes.push_back(v);
     SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
     paddedSubviewResults.push_back(rewriter.create<tensor::ExtractSliceOp>(
-        loc, std::get<1>(it), offsets, sizes, strides));
+        loc, paddedResult, offsets, sizes, strides));
   }
-  // Replace the transient `opToPad` locally, except for uses that we just
-  // created for the purpose of extracting the dims.
-  rewriter.replaceOpWithIf(opToPad, paddedSubviewResults, [&](OpOperand &opOp) {
-    return !newUsersOfOpToPad.contains(opOp.getOwner());
-  });
+  rewriter.replaceOp(opToPad, paddedSubviewResults);
   return success();
 }
 
 /// Linalg base tiling pattern.
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(opName, benefit, context), filter(filter),
if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
 
   // Setup RAII guard to return properly.
+  LinalgOp paddedOp;
   LinalgOp tiledOp = res->op;
   auto guard = llvm::make_scope_exit([&]() {
     // Return relevant information to derived pattern.
     result = *res;
-    // Replace filter on both tiledOp and tiledAndPaddedOp, if necessary.
+    // Update filter.
+    if (paddedOp)
+      filter.replaceLinalgTransformationFilter(rewriter, paddedOp);
+    else
       filter.replaceLinalgTransformationFilter(rewriter, tiledOp);
-    if (tiledOp != res->op)
-      filter.replaceLinalgTransformationFilter(rewriter, res->op);
   });
 
   // Consider padding on the fly only if the op has tensor semantics.
   if (!options.paddingValueComputationFunction ||
       !linalgOp.hasTensorSemantics())
     return success();
 
   // Try to pad on the fly by rewriting res->op as a padded op. If successful,
   // `res.op` is rewritten in static form with padded operands.
-  LinalgOp paddedOp;
   if (succeeded(rewriteAsPaddedOp(rewriter, res->op,
                                   options.paddingValueComputationFunction,
                                   paddedOp))) {
     res->op = paddedOp;
     // Do not perform replacement of `linalgOp`, let the derived patterns
     // do this as they see fit, from the resulting TiledLinalgOp.
     return success();
   }

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

-// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3" -resolve-shaped-type-result-dims -cse -split-input-file | \
+// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3" -cse -split-input-file | \
 // RUN: FileCheck %s -check-prefix=TILE2
 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=0,3" -resolve-shaped-type-result-dims -cse -split-input-file | \
 // RUN: FileCheck %s -check-prefix=TILE1
 
 //  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>