Support tile-and-pad when padding doesn't span all dimension

[asaadaldien]

Summary

• Note: This only works if untiled dims are static. Otherwise we need to change the semantics of PadOp to accept [%size...] attributes.

[nicolasvasilache]

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

276–277

This revision breaks this invariant, why ?

mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir

46

This does not check = instead of -.
Also, CHECK-NOT: scf.for ?

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

87

tile-sizes is a bit too generic for this very specialized test.
Rename to tile-sizes-for-padding ?

[asaadaldien]

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

276–277

To alter the marker in any case otherwise this pattern will be stuck in in an infinite loop of tile -> pad -> fail -> tile -> pad -> fail -> ...
You can repro this by running only tile with 2 dims without this revision.

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

87

make sense! will change the name in the next revision

[asaadaldien]

Use tile-sizes-for-padding test argument

[asaadaldien]

Fix test

[hanchung]

Otherwise we need to change the semantics of PadOp to accept [%size...] attributes.

I think the pad op can accept [%size ...] ? See https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Linalg/roundtrip.mlir#L30-L37

func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
                  %pad_value: f32) -> tensor<6x?x?x?xf32> {
  %0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      linalg.yield %pad_value : f32
    } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32>
  return %0 : tensor<6x?x?x?xf32>
}

[hanchung]

I didn't intend to click approval...

[nicolasvasilache]

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

276–277

Ok, then this deserves a separate commit and a test case that exhibits the behavior, independently of your current CL (which is otherwise good to go from my perspective).

[asaadaldien]

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

276–277

done in https://reviews.llvm.org/D97720, retire this diff now..

[asaadaldien]

split into others started( https://reviews.llvm.org/D97720)

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

LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
   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.
-  bool succeeded = true;
   LinalgOp tiledOp = res->op;
   auto guard = llvm::make_scope_exit([&]() {
-    if (!succeeded)
-      return;
     // Return relevant information to derived pattern.
     result = *res;
     // Replace filter on both tiledOp and tiledAndPaddedOp, if necessary.
     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 (failed(rewriteAsPaddedOp(rewriter, *res, options))) {
     // Set so RAII guard does not propagate TiledLinalgOp to `result`.
-    succeeded = false;
     return failure();
   }
 
   // Do not perform replacement of `linalgOp`, let the derived patterns
   // do this as they see fit, from the resulting TiledLinalgOp.
   return success();
 }
 
 static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) {

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

IntegerAttr getSmallestBoundingIndex(Value size) {
   } else if (auto constIndexOp = size.getDefiningOp<ConstantOp>()) {
     if (constIndexOp.getType().isa<IndexType>())
       boundingConst = constIndexOp.value().cast<IntegerAttr>().getInt();
   } else if (auto affineApplyOp = size.getDefiningOp<AffineApplyOp>()) {
     if (auto cExpr = affineApplyOp.getAffineMap()
                          .getResult(0)
                          .dyn_cast<AffineConstantExpr>())
       boundingConst = cExpr.getValue();
+  } else if (auto dimOp = size.getDefiningOp<DimOp>()) {
+    // This dimension isn't tiled, so set extend bounding box to its static
+    // size.
+    auto shape = dimOp.memrefOrTensor().getType().dyn_cast<ShapedType>();
+    if (auto constOp = dimOp.index().getDefiningOp<ConstantOp>()) {
+      if (auto indexAttr = constOp.value().dyn_cast<IntegerAttr>()) {
+        auto dimIndex = indexAttr.getInt();
+        if (!shape.isDynamicDim(dimIndex))
+          boundingConst = shape.getShape()[dimIndex];
+      }
+    }
   }
   if (boundingConst && *boundingConst >= 0)
     return Builder(size.getContext()).getIndexAttr(*boundingConst);
   return nullptr;
 }
 
 /// Specialization to build an scf "for" nest.
 template <>

mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir

-// RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-and-pad-pattern -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-and-pad-pattern tile-sizes-for-padding=2,3,4" -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-and-pad-pattern tile-sizes-for-padding=2,3" -canonicalize | FileCheck %s -check-prefix=CHECK-1DIM-TILE
+
+func @matmul_tensors(
+  %arg0: tensor<?x?xi8>, %arg1: tensor<?x?xi8>, %arg2: tensor<?x?xi32>)
+    -> tensor<?x?xi32> {
+  %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile-and-pad"}
+      ins(%arg0, %arg1: tensor<?x?xi8>, tensor<?x?xi8>)
+     outs(%arg2: tensor<?x?xi32>)
+    -> tensor<?x?xi32>
+  return %0 : tensor<?x?xi32>
+}
 
 // CHECK-LABEL: func @matmul_tensors(
 // CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
-func @matmul_tensors(
-  %arg0: tensor<?x?xi8>, %arg1: tensor<?x?xi8>, %arg2: tensor<?x?xi32>)
-    -> tensor<?x?xi32> {
 //      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xi32>) {
 //      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xi32>) {
 //      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xi32>) {
 //      CHECK:       %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x?xi8> to tensor<?x?xi8>
 //      CHECK:       %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<?x?xi8> to tensor<?x?xi8>
 //      CHECK:       %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor<?x?xi32> to tensor<?x?xi32>
-
 // Dynamic op has been canonicalized away.
 //  CHECK-NOT:       linalg.matmul {{.*}} tensor<?x?xi8>
-
 // Padding injects static information.
 //      CHECK:       %[[pA:.*]] = linalg.pad_tensor %[[sTA]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
 //      CHECK:         : tensor<?x?xi8> to tensor<2x4xi8>
 //      CHECK:       %[[pB:.*]] = linalg.pad_tensor %[[sTB]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
 //      CHECK:         : tensor<?x?xi8> to tensor<4x3xi8>
 //      CHECK:       %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
 //      CHECK:         : tensor<?x?xi32> to tensor<2x3xi32>
 //      CHECK:       %[[pD:.*]] = linalg.matmul_i8_i8_i32 ins(%[[pA]], %[[pB]] : tensor<2x4xi8>, tensor<4x3xi8>)
 // CHECK-SAME:                                           outs(%[[pC]] : tensor<2x3xi32>)  -> tensor<2x3xi32>
 //      CHECK:       %[[sTD:.*]] = subtensor %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xi32> to tensor<?x?xi32>
 //      CHECK:       %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?xi32> into tensor<?x?xi32>
 //      CHECK:       scf.yield %[[TD]] : tensor<?x?xi32>
 //      CHECK:     scf.yield %[[TD2]] : tensor<?x?xi32>
 //      CHECK:   scf.yield %[[TD1]] : tensor<?x?xi32>
+//      CHECK: return %[[TD0]] : tensor<?x?xi32>
+
+// CHECK-1DIM-TILE: func @matmul_tensors(
+// CHECK-1DIM-TILE:    %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
+// CHECK-1DIM-TILE:    %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
+// CHECK-1DIM-TILE:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
+// CHECK=1DIM-TILE: linalg.matmul {{.*}} ins(%[[TA]], %[[TB]]: tensor<?x?xi8>, tensor<?x?xi8>)  outs(%[[TC]]: tensor<?x?xi32>) -> tensor<?x?xi32>
+
+func @matmul_partially_static_tensors(
+  %arg0: tensor<?x8xi8>, %arg1: tensor<8x?xi8>, %arg2: tensor<?x?xi32>)
+    -> tensor<?x?xi32> {
   %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile-and-pad"}
-      ins(%arg0, %arg1: tensor<?x?xi8>, tensor<?x?xi8>)
+      ins(%arg0, %arg1: tensor<?x8xi8>, tensor<8x?xi8>)
      outs(%arg2: tensor<?x?xi32>)
     -> tensor<?x?xi32>
-
-//      CHECK: return %[[TD0]] : tensor<?x?xi32>
   return %0 : tensor<?x?xi32>
 }
+// CHECK-LABEL: func @matmul_partially_static_tensors(
+// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x8xi8>
+// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<8x?xi8>
+// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
+//    CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xi32>) {
+//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xi32>) {
+//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xi32>) {
+//      CHECK:       %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x8xi8> to tensor<?x?xi8>
+//      CHECK:       %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<8x?xi8> to tensor<?x?xi8>
+//      CHECK:       %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor<?x?xi32> to tensor<?x?xi32>
+//      CHECK:       %[[pA:.*]] = linalg.pad_tensor %[[sTA]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK:         : tensor<?x?xi8> to tensor<2x4xi8>
+//      CHECK:       %[[pB:.*]] = linalg.pad_tensor %[[sTB]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK:         : tensor<?x?xi8> to tensor<4x3xi8>
+//      CHECK:       %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK:         : tensor<?x?xi32> to tensor<2x3xi32>
+//      CHECK:       %[[pD:.*]] = linalg.matmul_i8_i8_i32 ins(%[[pA]], %[[pB]] : tensor<2x4xi8>, tensor<4x3xi8>)
+// CHECK-SAME:                                           outs(%[[pC]] : tensor<2x3xi32>)  -> tensor<2x3xi32>
+//      CHECK: return %[[TD0]] : tensor<?x?xi32>
+
+
+// CHECK-1DIM-TILE: func @matmul_partially_static_tensors(
+// CHECK-1DIM-TILE:    %[[TA:[0-9a-z]+]]: tensor<?x8xi8>
+// CHECK-1DIM-TILE:    %[[TB:[0-9a-z]+]]: tensor<8x?xi8>
+// CHECK-1DIM-TILE:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
+//    CHECK-1DIM-TILE: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xi32>) {
+//      CHECK-1DIM-TILE:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xi32>) {
+//      CHECK-1DIM-TILE:   %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x8xi8> to tensor<?x8xi8>
+//      CHECK-1DIM-TILE:   %[[scTA:.*]] = tensor.cast %[[sTA]] : tensor<?x8xi8> to tensor<?x?xi8>
+//      CHECK-1DIM-TILE:   %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<8x?xi8> to tensor<8x?xi8>
+//      CHECK-1DIM-TILE:   %[[scTB:.*]] = tensor.cast %[[sTB]] : tensor<8x?xi8> to tensor<?x?xi8>
+//      CHECK-1DIM-TILE:   %[[sTC:.*]] = subtensor %[[TC1]][{{.*}}] : tensor<?x?xi32> to tensor<?x?xi32>
+//      CHECK-1DIM-TILE:   %[[pA:.*]] = linalg.pad_tensor %[[scTA]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK-1DIM-TILE:         : tensor<?x?xi8> to tensor<2x8xi8>
+//      CHECK-1DIM-TILE:   %[[pB:.*]] = linalg.pad_tensor %[[scTB]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK-1DIM-TILE:         : tensor<?x?xi8> to tensor<8x3xi8>
+//      CHECK-1DIM-TILE:   %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[%c0, %c0] high[%{{.*}}, %{{.*}}]
+//      CHECK-1DIM-TILE:         : tensor<?x?xi32> to tensor<2x3xi32>
+//      CHECK-1DIM-TILE:   %[[pD:.*]] = linalg.matmul_i8_i8_i32 ins(%[[pA]], %[[pB]] : tensor<2x8xi8>, tensor<8x3xi8>)
+//      CHECK-1DIM-TILE:                                           outs(%[[pC]] : tensor<2x3xi32>)  -> tensor<2x3xi32>
+//      CHECK-1DIM-TILE: return %[[TD0]] : tensor<?x?xi32>

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

Option<bool> testGenericToVectorPattern{
       llvm::cl::init(false)};
   Option<bool> testAffineMinSCFCanonicalizationPatterns{
       *this, "test-affine-min-scf-canonicalization-patterns",
       llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
       llvm::cl::init(false)};
   Option<bool> testTileAndPadPattern{
       *this, "test-tile-and-pad-pattern",
       llvm::cl::desc("Test tile and pad pattern"), llvm::cl::init(false)};
+  ListOption<int64_t> tileSizesForPadding{
+      *this, "tile-sizes-for-padding",
+      llvm::cl::desc("Linalg tile sizes when tile+pad"), llvm::cl::ZeroOrMore,
+      llvm::cl::MiscFlags::CommaSeparated};
   Option<bool> testHoistPadding2Levels{*this, "test-hoist-padding-2-level",
                                        llvm::cl::desc("Test hoist padding"),
                                        llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();

 
 // 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().getType());
   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
 }
 
-static void applyTileAndPadPattern(FuncOp funcOp) {
+static void applyTileAndPadPattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes) {
   MLIRContext *context = funcOp.getContext();
   OwningRewritePatternList tilingPattern;
   auto linalgTilingOptions =
       linalg::LinalgTilingOptions()
-          .setTileSizes({2, 3, 4})
+          .setTileSizes(tileSizes)
           .setPaddingValueComputationFunction(getNeutralOfLinalgOp);
   tilingPattern.insert<linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>>(
       context, linalgTilingOptions,
       linalg::LinalgTransformationFilter(
           Identifier::get("tile-and-pad", context)));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
 }
 
return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
   if (testAffineMinSCFCanonicalizationPatterns)
     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
   if (testTileAndPadPattern)
-    return applyTileAndPadPattern(getFunction());
+    return applyTileAndPadPattern(getFunction(), tileSizesForPadding);
   if (testHoistPadding2Levels) {
     getFunction().walk([](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, 2);
     });
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace test
 } // namespace mlir