[mlir][linalg] New tiling option: Scalarize dynamic dims

[springerm]

Summary

This tiling option scalarizes all dynamic dimensions, i.e., it tiles all dynamic dimensions by 1.

This option is useful for linalg ops with partly dynamic tensor dimensions. E.g., such ops can appear in the partial iteration after loop peeling. After scalarizing dynamic dims, those ops can be vectorized.

Depends On D109043

[springerm]

rebase

[nicolasvasilache]

I'll let this one simmer until you had a chance to decide whether you update scf.for peeling to be consistent with linalg.tiled_loop peeling.
Ping me if you want me to review.

[springerm]

rebase

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

LinalgTilingOptions &setTileSizes(SmallVector<Value, 4> ts) {
     tileSizeComputationFunction = [=](OpBuilder &, Operation *) { return ts; };
     return *this;
   }
   /// Convenience function to set the `tileSizeComputationFunction` to a
   /// function that computes tile sizes at the point they are needed. Allows
   /// proper interaction with folding.
   LinalgTilingOptions &setTileSizes(ArrayRef<int64_t> ts);
 
+  /// Tile all dynamic dimensions by 1. I.e., scalarize those dimensions.
+  /// Note: `scalarizeDynamicDims` and `setTileSizes` cannot be used together.
+  LinalgTilingOptions &scalarizeDynamicDims();
+
   /// The interchange vector to reorder the tiled loops.
   SmallVector<unsigned, 4> interchangeVector = {};
 
   LinalgTilingOptions &setInterchange(ArrayRef<unsigned> interchange) {
     interchangeVector.assign(interchange.begin(), interchange.end());
     return *this;
   }
 

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

op->setAttr(LinalgTransforms::kLinalgTransformMarker,
                 rewriter.getStringAttr(replacement.getValue().strref()));
   else
     op->removeAttr(Identifier::get(LinalgTransforms::kLinalgTransformMarker,
                                    rewriter.getContext()));
 }
 
 LinalgTilingOptions &
 mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef<int64_t> ts) {
+  assert(!tileSizeComputationFunction && "tile sizes already set");
   SmallVector<int64_t, 4> tileSizes(ts.begin(), ts.end());
   tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
     OpBuilder::InsertionGuard guard(b);
     b.setInsertionPointToStart(
         &op->getParentOfType<FuncOp>().getBody().front());
     return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
       Value v = b.create<ConstantIndexOp>(op->getLoc(), s);
       return v;
     }));
   };
   return *this;
 }
 
+LinalgTilingOptions &mlir::linalg::LinalgTilingOptions::scalarizeDynamicDims() {
+  assert(!tileSizeComputationFunction && "tile sizes already set");
+  tileSizeComputationFunction = [](OpBuilder &b, Operation *op) {
+    SmallVector<Value, 4> tileSizes;
+    auto linalgOp = dyn_cast<LinalgOp>(op);
+    if (!linalgOp)
+      return tileSizes;
+    Location loc = linalgOp.getLoc();
+    auto allShapeSizes = linalgOp.createFlatListOfOperandDims(b, loc);
+    AffineMap map = linalgOp.getShapesToLoopsMap();
+    if (!map)
+      return tileSizes;
+    auto shapeSizes = applyMapToValues(b, loc, map, allShapeSizes);
+    // If the shape size is dynamic, tile by 1. Otherwise, do not tile (tile
+    // size 0).
+    for (Value shapeSize : shapeSizes)
+      tileSizes.push_back(getConstantIntValue(shapeSize).hasValue()
+                              ? b.create<ConstantIndexOp>(loc, 0)
+                              : b.create<ConstantIndexOp>(loc, 1));
+    return tileSizes;
+  };
+  return *this;
+}
+
 /// Try to compute a static bounding box for `operand`
 /// Return success if either:
 ///   1. The operand is already statically shaped, `result` is left unchanged.
 ///   2. The operand is (partially) dynamic, `result` is the result of a freshly
 ///      created PadTensorOp.
 /// Return failure if the operand cannot be padded to a static shape.
 static LogicalResult padOperandToSmallestStaticBoundingBox(
     PatternRewriter &rewriter, linalg::LinalgOp opToPad, OpOperand *opOperand,

mlir/test/Dialect/Linalg/tile-scalarize-dynamic-dims.mlir

+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-scalarize-dynamic-dims" -for-loop-canonicalization -canonicalize -split-input-file | \
+// RUN:     FileCheck %s
+
+// CHECK-LABEL: func @matmul_partly_dynamic_tensor(
+//  CHECK-SAME:     %[[ARG0:.*]]: tensor<?x?xf32>, %[[ARG1:.*]]: tensor<?x2000xf32>
+//   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.*]] = constant 1 : index
+//       CHECK:   tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>
+//       CHECK:   %[[UB1:.*]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?xf32>
+//       CHECK:   %[[UB2:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>
+//       CHECK:   scf.for %[[IV0:.*]] = %[[C0]] to %[[UB1]] step %[[C1]]
+//       CHECK:     scf.for %[[IV1:.*]] = %[[C0]] to %[[UB2]] step %[[C1]]
+//       CHECK:       %[[S1:.*]] = tensor.extract_slice %[[ARG0]][%[[IV0]], %[[IV1]]] [1, 1] [1, 1] : tensor<?x?xf32> to tensor<1x1xf32>
+//       CHECK:       %[[S2:.*]] = tensor.extract_slice %[[ARG1]][%[[IV1]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32>
+//       CHECK:       %[[S3:.*]] = tensor.extract_slice %{{.*}}[%[[IV0]], 0] [1, 2000] [1, 1] : tensor<?x2000xf32> to tensor<1x2000xf32>
+//       CHECK:       linalg.matmul ins(%[[S1]], %[[S2]] : tensor<1x1xf32>, tensor<1x2000xf32>) outs(%[[S3]] : tensor<1x2000xf32>) -> tensor<1x2000xf32>
+func @matmul_partly_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x2000xf32>)
+    -> tensor<?x2000xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %out = linalg.init_tensor [%d0, 2000] : tensor<?x2000xf32>
+  %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
+      ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x2000xf32>)
+      outs(%out: tensor<?x2000xf32>) -> tensor<?x2000xf32>
+  return %r : tensor<?x2000xf32>
+}

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

struct TestLinalgTransforms
   Option<bool> testGenericToVectorPattern{
       *this, "test-linalg-to-vector-patterns",
       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
                      "in vector.contract form"),
       llvm::cl::init(false)};
   Option<bool> testTilePattern{*this, "test-tile-pattern",
                                llvm::cl::desc("Test tile pattern"),
                                llvm::cl::init(false)};
+  Option<bool> testTileScalarizeDynamicDims{
+      *this, "test-tile-scalarize-dynamic-dims",
+      llvm::cl::desc("Test tiling of dynamic dims by 1"),
+      llvm::cl::init(false)};
   Option<int> testHoistPadding{*this, "test-hoist-padding",
                                llvm::cl::desc("Test hoist padding"),
                                llvm::cl::init(0)};
   Option<bool> testTransformPadTensor{
       *this, "test-transform-pad-tensor",
       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
       llvm::cl::init(false)};
   Option<bool> testGeneralizePadTensor{

 // For now, just assume it is the zero of type.
 // In the future, it should be the zero of type + op.
 static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
   auto t = getElementTypeOrSelf(op.get());
   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
 }
 
 static void applyTilePattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes,
-                             bool padTiles, ArrayRef<int64_t> peeledLoops) {
+                             bool padTiles, ArrayRef<int64_t> peeledLoops,
+                             bool scalarizeDynamicDims) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet tilingPattern(context);
   auto linalgTilingOptions =
-      linalg::LinalgTilingOptions().setTileSizes(tileSizes).setPeeledLoops(
-          peeledLoops);
+      linalg::LinalgTilingOptions().setPeeledLoops(peeledLoops);
+  if (scalarizeDynamicDims) {
+    linalgTilingOptions.scalarizeDynamicDims();
+    assert(tileSizes.empty() &&
+           "tileSizes and scalarizeDynamicDims is mutually exclusive");
+  } else {
+    linalgTilingOptions.setTileSizes(tileSizes);
+  }
   if (padTiles)
     linalgTilingOptions.setPaddingValueComputationFunction(
         getNeutralOfLinalgOp);
 
   tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulOp>,
                     linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>,
                     linalg::LinalgTilingPattern<linalg::GenericOp>>(
       context, linalgTilingOptions,
if (testTransformPadTensor)
     return applyPadTensorToGenericPatterns(getFunction());
   if (testGeneralizePadTensor)
     return applyGeneralizePadTensorPatterns(getFunction());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getFunction());
   if (testTiledLoopPeeling.hasValue())
     return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling);
   if (testTilePattern)
-    return applyTilePattern(getFunction(), tileSizes, padTiles, peeledLoops);
+    return applyTilePattern(getFunction(), tileSizes, padTiles, peeledLoops,
+                            /*scalarizeDynamicDims=*/false);
+  if (testTileScalarizeDynamicDims)
+    return applyTilePattern(getFunction(), tileSizes, padTiles,
+                            /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
   if (testHoistPadding) {
     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
     });
   }
   if (testInterchangePattern.hasValue())
     return applyInterchangePattern(getFunction(), testInterchangePattern);
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir