[mlir][Linalg] Allow tiling of batch dimension for convolution ops with padding.

[mravishankar]

Summary

Existing tiling implementation of Linalg would still work for tiling
the batch dimensions of the convolution op.

Depends On D76415

[bondhugula]

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

340

Nit: Missing period at the end.

436

You don't need the getOperation().

[bondhugula]

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

436

Sorry, you do need it here.

[hanchung]

Add a test?

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

343

Use isZero(val)?

[bondhugula]

Shouldn't there be a test case here?

[mravishankar]

Addressing comments and adding test

[mravishankar]

Sorry for the delay. I thought I was blocked on pushing on this (but I really wasnt).

[hanchung]

LGTM, would be good to wait until Nicolas taking a look.

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

437–438

same here, missing a period at the end.

[mravishankar]

Addressing minor comment and rebasing

[asaadaldien]

LGTM!

[mravishankar]

@bondhugula : Gentle ping on this patch. I think it is ready to land now. PTAL.

[bondhugula]

In D76637#1951713, @mravishankar wrote:

@bondhugula : Gentle ping on this patch. I think it is ready to land now. PTAL.

Hi Mahesh, you don't have to wait for my approval; I just commented superficially on what I noticed. Please rely on the deeper reviews from others.

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

Optional<TiledLinalgOp> static tileLinalgOpImpl(OpBuilder &b, LinalgOp op,
   // dimension. This convention is significantly simpler to handle instead of
   // adjusting affine maps to account for missing dimensions.
   assert(op.getNumParallelLoops() + op.getNumReductionLoops() +
                  op.getNumWindowLoops() ==
              tileSizes.size() &&
          "expected matching number of tile sizes and loops");
 
   if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
-    // TODO(ntv): add a level of indirection to linalg.generic.
-    if (convOp.padding())
-      llvm_unreachable("Unexpected conv with padding");
+    // For conv op only support tiling along batch dimension (which is the first
+    // loop).
+    if (convOp.padding() &&
+        !llvm::all_of(tileSizes.drop_front(),
+                      [](Value val) { return isZero(val); }))
+      return llvm::None;
   }
 
   // If permutation is empty, use the identity. Build the permutation map
   // otherwise.
   auto invPermutationMap = AffineMap::getMultiDimIdentityMap(
       tileSizes.size(), ScopedContext::getContext());
   if (!permutation.empty())
     invPermutationMap = inversePermutation(

 template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ArrayRef<int64_t> tileSizes,
                  ArrayRef<unsigned> permutation, OperationFolder *folder) {
   assert(op.hasBufferSemantics() && "expected linalg op with buffer semantics");
   if (tileSizes.empty())
     return llvm::None;
 
-  if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
-    // TODO(ntv): add a level of indirection to linalg.generic.
-    if (convOp.padding())
-      llvm_unreachable("Unexpected conv with padding");
-  }
-
   // The following uses the convention that "tiling by zero" skips tiling a
   // particular dimension. This convention is significantly simpler to handle
   // instead of adjusting affine maps to account for missing dimensions.
   auto nLoops = op.getNumParallelLoops() + op.getNumReductionLoops() +
                 op.getNumWindowLoops();
   tileSizes = tileSizes.take_front(nLoops);
   // If only 0 tilings are left, then return.
   if (llvm::all_of(tileSizes, [](int64_t v) { return v == 0; }))
     return llvm::None;
 
+  if (auto convOp = dyn_cast<linalg::ConvOp>(op.getOperation())) {
+    // For conv op only support tiling along batch dimension (which is the first
+    // loop).
+    if (convOp.padding() && !llvm::all_of(tileSizes.drop_front(),
+                                          [](int64_t val) { return val == 0; }))
+      return llvm::None;
+  }
+
   // Create a builder for tile size constants.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPoint(op);
   ScopedContext scope(b, op.getLoc());
 
   // Materialize concrete tile size values to pass the generic tiling function.
   SmallVector<Value, 8> tileSizeValues;
   tileSizeValues.reserve(tileSizes.size());

mlir/test/Dialect/Linalg/tile_conv_padding.mlir

+// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,0,0,4" | FileCheck %s -check-prefix=TILE-23004
+// RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2" | FileCheck %s -check-prefix=TILE-20000
+
+// TILE-23004-DAG: #[[strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
+// TILE-20000-DAG: #[[strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
+// TILE-20000-DAG: #[[minmap:.*]] = affine_map<(d0, d1, d2) -> (d0, d1 - d2)>
+// TILE-20000-DAG: #[[subviewstride:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3, s4] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3 * s4)>
+
+func @conv_padding(%arg0: memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, %arg1: memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, %arg2: memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>) {
+  linalg.conv(%arg0, %arg1, %arg2) {dilations = [10, 20], padding = dense<[[1, 1], [0, 1]]> : tensor<2x2xi64>, strides = [30, 40]} : memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>, memref<?x?x?x?xf32, offset: ?, strides: [?, ?, ?, 1]>
+  return
+}
+// TILE-23004-LABEL: func @conv_padding(
+//  TILE-23004-SAME:   %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>
+//  TILE-23004-SAME:   %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>
+//  TILE-23004-SAME:   %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>)
+//       TILE-23004:         linalg.conv(%[[ARG0]], %[[ARG1]], %[[ARG2]])
+
+// TILE-20000-LABEL: func @conv_padding(
+//  TILE-20000-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>
+//  TILE-20000-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>
+//  TILE-20000-SAME: %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32, #[[strided4D]]>)
+//   TILE-20000-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE-20000-DAG:   %[[C1:.*]] = constant 1 : index
+//   TILE-20000-DAG:   %[[C2:.*]] = constant 2 : index
+//       TILE-20000:   %[[B:.*]] = dim %[[ARG1]], 0
+//       TILE-20000:   loop.for %[[ivI:.*]] = %[[C0]] to %[[B]] step %[[C2]] {
+//       TILE-20000:     %[[DIM10:.*]] = dim %[[ARG1]], 0
+//       TILE-20000:     %[[EXTENT:.*]] = affine.min #[[minmap]](%[[C2]], %[[DIM10]], %[[ivI]])
+//       TILE-20000:     %[[DIM11:.*]] = dim %[[ARG1]], 1
+//       TILE-20000:     %[[DIM12:.*]] = dim %[[ARG1]], 2
+//       TILE-20000:     %[[DIM13:.*]] = dim %[[ARG1]], 3
+//       TILE-20000:     %[[SUBVIEW1:.*]] = subview %[[ARG1]][%[[ivI]], %[[C0]], %[[C0]], %[[C0]]] [%[[EXTENT]], %[[DIM11]], %[[DIM12]], %[[DIM13]]]
+//       TILE-20000:     %[[DIM20:.*]] = dim %[[ARG2]], 0
+//       TILE-20000:     %[[EXTENT:.*]] = affine.min #[[minmap]](%[[C2]], %[[DIM20]], %[[ivI]])
+//       TILE-20000:     %[[DIM21:.*]] = dim %[[ARG2]], 1
+//       TILE-20000:     %[[DIM22:.*]] = dim %[[ARG2]], 2
+//       TILE-20000:     %[[DIM23:.*]] = dim %[[ARG2]], 3
+//       TILE-20000:     %[[SUBVIEW2:.*]] = subview %[[ARG2]][%[[ivI]], %[[C0]], %[[C0]], %[[C0]]] [%[[EXTENT]], %[[DIM21]], %[[DIM22]], %[[DIM23]]]
+//       TILE-20000:     linalg.conv(%[[ARG0]], %[[SUBVIEW1]], %[[SUBVIEW2]])