[mlir][Linalg] Wrong tile size for convolutions fixed

[limo1996]

Summary

Sizes of tiles (subviews) are bigger by 1 than they should. Let's consider
1D convolution without batches or channels. Furthermore let m iterate over
the output and n over the kernel then input is accessed with m + n. In tiling
subview sizes for convolutions are computed by applying requested tile size
together with kernel size to the above mentioned expression thus let's say
for tile size of 2 the subview size is 2 + size(n), which is bigger by one
than it should since we move kernel only once. The problem behind it is that
range is not turned into closed interval before the composition. This commit
fixes the problem by turning ranges first into closed intervals by substracting
1 and after the composition back to half open by adding 1.

PHAB_REVIEW=D86638

[mravishankar]

@limo1996 This is really nice find! I have been struggling with figuring out why convolutions are giving me errors in some cases, and I suspected something like, but wasnt able to nail this down.

There might be a better fix though. I'll take a look at this as well. For now I am going to just "Request changes" but I dont really have any changes to request. Just a placeholder for me to take a look again.

[limo1996]

In D86638#2239561, @mravishankar wrote:

@limo1996 This is really nice find! I have been struggling with figuring out why convolutions are giving me errors in some cases, and I suspected something like, but wasnt able to nail this down.

There might be a better fix though. I'll take a look at this as well. For now I am going to just "Request changes" but I dont really have any changes to request. Just a placeholder for me to take a look again.

Thank you @mravishankar! I know for sure there is a better fix but my goal here was to initiate a discussion.

You can try new Convolutions for different layouts and ranks that were recently added. I also saw that traditional linalg.conv is quite buggy but the plan is that it will be replaced so I did not invest time into fixing it.

[mravishankar]

Could we start a discourse post on this. Its more easy to document issues and discuss solutions there rather than on the patch.

Also, I am curious what other bugs you found in the convolution operation. Please add those to the discourse post if possible.

[limo1996]

In D86638#2242434, @mravishankar wrote:

Could we start a discourse post on this. Its more easy to document issues and discuss solutions there rather than on the patch.

Also, I am curious what other bugs you found in the convolution operation. Please add those to the discourse post if possible.

Here is the discourse post: https://llvm.discourse.group/t/wrong-tile-size-for-convolution-ops/1699

[limo1996]

Implementation changed as discussed here:
https://llvm.discourse.group/t/wrong-tile-size-for-convolution-ops/1699

[nicolasvasilache]

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

248

Please use mlir::edsc::op::operator+ and write it as size + std_constant_index(1) to get the proper affine_apply operation generated.
Without that things won't canonicalize nicely.
See e.g. mlir/test/EDSC/builder-api-test.cpp line 82 for some C++ and the generated IR.

289

same thing here re operator+

mlir/test/Dialect/Linalg/tile_simple_conv.mlir

34

these shouls not appear as subi / addi but instead use affine_apply with the properly canonicalized/simplified expressions.

[limo1996]

Comments of nicolas resolved

[nicolasvasilache]

Great, thank you @limo1996 !

[mravishankar]

Thanks!

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

static SmallVector<Value, 4> makeTiledViews(OpBuilder &b, Location loc,
 
   auto viewSizes = applyMapToValues(b, loc, map, allViewSizes);
   // Construct (potentially temporary) mins and maxes on which to apply maps
   // that define tile subviews.
   SmallVector<Value, 8> lbs, subViewSizes;
   for (unsigned idx = 0, idxIvs = 0, e = tileSizes.size(); idx < e; ++idx) {
     bool isTiled = !isZero(tileSizes[idx]);
     lbs.push_back(isTiled ? ivs[idxIvs++] : (Value)std_constant_index(0));
-    subViewSizes.push_back(isTiled ? tileSizes[idx] : viewSizes[idx]);
+    // Before composing, we need to make range a closed interval.
+    Value size = isTiled ? tileSizes[idx] : viewSizes[idx];
+    subViewSizes.push_back(std_subi(size, std_constant_index(1)));
   }
 
   auto *op = linalgOp.getOperation();
 
   SmallVector<Value, 4> res;
   res.reserve(op->getNumOperands());
   auto viewIteratorBegin = linalgOp.getInputsAndOutputBuffers().begin();
   for (unsigned viewIndex = 0; viewIndex < linalgOp.getNumInputsAndOutputs();
for (unsigned r = 0; r < rank; ++r) {
         continue;
       }
 
       // Tiling creates a new slice at the proper index, the slice step is 1
       // (i.e. the slice view does not subsample, stepping occurs in the loop).
       auto m = map.getSubMap({r});
       auto offset = applyMapToValues(b, loc, m, lbs).front();
       offsets.push_back(offset);
-      auto size = applyMapToValues(b, loc, m, subViewSizes).front();
+      auto closed_int_size = applyMapToValues(b, loc, m, subViewSizes).front();
+      // Resulting size needs to be made half open interval again.
+      auto size = (Value)std_addi(closed_int_size, std_constant_index(1));
 
       // The size of the subview should be trimmed to avoid out-of-bounds
       // accesses, unless we statically know the subview size divides the view
       // size evenly.
       int64_t viewSize = viewType.getDimSize(r);
       auto sizeCst = size.getDefiningOp<ConstantIndexOp>();
       if (ShapedType::isDynamic(viewSize) || !sizeCst ||
           (viewSize % sizeCst.getValue()) != 0) {

mlir/test/Dialect/Linalg/tile_conv.mlir

 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,0,0,4" | FileCheck %s -check-prefix=TILE-23004
 
 // TILE-23004-DAG: #[[$D0x30pS0x10:.*]] = affine_map<(d0) -> (d0 * 30)>
-// TILE-23004-DAG: #[[$S0x10p90D0x30pS1:.*]] = affine_map<(d0)[s0, s1] -> (s0 * 10 + 90, d0 * -30 + s1)>
+// TILE-23004-DAG: #[[$S0x10p90D0x30pS1:.*]] = affine_map<(d0) -> (d0 * 10 + 60)>
 // TILE-23004-DAG: #[[$strided4D:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
 // TILE-23004-DAG: #[[$bound_map_4:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
+// TILE-23004-DAG: #[[$step_map:.*]] = affine_map<(d0, d1)[s0] -> (d0, d1 * -30 + s0)>
 
 func @conv(%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], 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: func @conv(
 //  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-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE-23004-DAG:   %[[C1:.*]] = constant 1 : index
 //   TILE-23004-DAG:   %[[C2:.*]] = constant 2 : index
 //   TILE-23004-DAG:   %[[C3:.*]] = constant 3 : index
 //   TILE-23004-DAG:   %[[C4:.*]] = constant 4 : index
 //       TILE-23004:   %[[Z0:.*]] = dim %[[ARG0]], %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:   %[[Q:.*]] = dim %[[ARG0]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:   %[[B:.*]] = dim %[[ARG1]], %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:   %[[X0:.*]] = dim %[[ARG2]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:   scf.for %[[ivI:.*]] = %{{.*}} to %[[B]] step %{{.*}} {
 //       TILE-23004:     scf.for %[[ivJ:.*]] = %{{.*}} to %[[X0]] step %{{.*}} {
 //       TILE-23004:       scf.for %[[ivK:.*]] = %{{.*}} to %[[Q]] step %{{.*}} {
+//       TILE-23004:         %[[Z0_minus_1:.*]] = subi %[[Z0]], %[[C1]] : index
 //       TILE-23004:         %[[Z0_1:.*]] = dim %[[ARG0]], %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[Z1:.*]] = dim %[[ARG0]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[Z2:.*]] = dim %[[ARG0]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[szK:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[Z2]]]
 //       TILE-23004:         %[[K:.*]] = dim %[[ARG0]], %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0_1]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
 //
 //       TILE-23004:         %[[J1:.*]] = affine.apply #[[$D0x30pS0x10]](%[[ivJ]])
+//       TILE-23004:         %[[closed_int_size:.*]] = affine.apply #[[$S0x10p90D0x30pS1]](%[[Z0_minus_1]])
+//       TILE-23004:         %[[size:.*]] = addi %[[closed_int_size]], %[[C1]] : index
 //       TILE-23004:         %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[Z0]], %[[PaddedInput0b]]]
+//       TILE-23004:         %[[I1pStep:.*]] = affine.min #[[$step_map]](%[[size]], %[[ivJ]])[%[[PaddedInput0b]]]
 //       TILE-23004:         %[[SZ2:.*]] = dim %[[ARG1]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[dim3:.*]] = dim %[[ARG1]], %c3
 //       TILE-23004:         %[[sz3:.*]] = affine.min #[[$bound_map_4]](%[[ivK]])[%[[dim3]]]
 //       TILE-23004:         %[[InputView:.*]] = subview %{{.*}}[%[[ivI]], %[[J1]], 0, %[[ivK]]] [%{{.*}}, %{{.*}}, %[[SZ2]], %[[sz3]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
 //
 //       TILE-23004:         %[[X0:.*]] = dim %[[ARG2]], %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[X1:.*]] = dim %[[ARG2]], %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
 //       TILE-23004:         %[[OutputView:.*]] = subview %{{.*}}[%[[ivI]], %[[ivJ]], 0, 0] [%{{.*}}, %{{.*}}, %[[X0]], %[[X1]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
 //
 //       TILE-23004:         linalg.conv(%[[FilterView]], %[[InputView]], %[[OutputView]]) {dilations = [10, 20], strides = [30, 40]} : memref<?x?x?x?xf32, #[[$strided4D]]>, memref<?x?x?x?xf32, #[[$strided4D]]>, memref<?x?x?x?xf32, #[[$strided4D]]>

mlir/test/Dialect/Linalg/tile_simple_conv.mlir

 // RUN: mlir-opt %s -linalg-tile="linalg-tile-sizes=2,3,4" | FileCheck %s
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (2, -d0 + s0)>
-//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0)[s0, s1] -> (s0 + 3, -d0 + s1)>
-//  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0)[s0, s1] -> (s0 + 4, -d0 + s1)>
-//  CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)>
-//  CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
+//  CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0) -> (d0 + 2)>
+//  CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1)[s0] -> (d0, -d1 + s0)>
+//  CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0) -> (d0 + 3)>
+//  CHECK-DAG: #[[MAP4:.*]] = affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0 * s1 + s0 + d1 * s2 + d2 * s3 + d3)>
+//  CHECK-DAG: #[[MAP5:.*]] = affine_map<(d0)[s0] -> (3, -d0 + s0)>
+//  CHECK-DAG: #[[MAP6:.*]] = affine_map<(d0)[s0] -> (4, -d0 + s0)>
 
 func @conv(%arg0 : memref<?x?x?x?xf32>, %arg1 : memref<?x?x?x?xf32>, %arg2 : memref<?x?x?x?xf32>) {
   linalg.conv(%arg0, %arg1, %arg2) : memref<?x?x?x?xf32>, memref<?x?x?x?xf32>, memref<?x?x?x?xf32>
   return
 }
 
 //       CHECK: func @conv
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
 //  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]*]]: memref<?x?x?x?xf32>
 //   CHECK-DAG:   %[[C0:.*]] = constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = constant 1 : index
 //   CHECK-DAG:   %[[C2:.*]] = constant 2 : index
 //   CHECK-DAG:   %[[C3:.*]] = constant 3 : index
 //   CHECK-DAG:   %[[C4:.*]] = constant 4 : index
 //       CHECK:   %[[T0:.*]] = dim %[[ARG0]], %[[C0]]
 //       CHECK:   %[[T1:.*]] = dim %[[ARG0]], %[[C1]]
 //       CHECK:   %[[T2:.*]] = dim %[[ARG1]], %[[C0]]
 //       CHECK:   %[[T3:.*]] = dim %[[ARG2]], %[[C1]]
 //       CHECK:   %[[T4:.*]] = dim %[[ARG2]], %[[C2]]
 //       CHECK:   scf.for %[[ARG3:.*]] = %[[C0]] to %[[T2]] step %[[C2]]
 //       CHECK:     scf.for %[[ARG4:.*]] = %[[C0]] to %[[T3]] step %[[C3]]
 //       CHECK:       scf.for %[[ARG5:.*]] = %[[C0]] to %[[T4]] step %[[C4]]
-//       CHECK:         %[[T5:.*]] = dim %[[ARG1]], %[[C0]]
-//       CHECK:         %[[T6:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T5]]]
-//       CHECK:         %[[T7:.*]] = dim %[[ARG1]], %[[C1]]
-//       CHECK:         %[[T8:.*]] = affine.min #[[MAP1]](%[[ARG4]])[%[[T0]], %[[T7]]]
-//       CHECK:         %[[T9:.*]] = dim %[[ARG1]], %[[C2]]
-//       CHECK:         %[[T10:.*]] = affine.min #[[MAP2]](%[[ARG5]])[%[[T1]], %[[T9]]]
-//       CHECK:         %[[T11:.*]] = dim %[[ARG1]], %[[C3]]
+//       CHECK:         %[[T5:.*]] = subi %[[T0]], %[[C1]] : index
+//       CHECK:         %[[T6:.*]] = subi %[[T1]], %[[C1]] : index
+//       CHECK:         %[[T7:.*]] = dim %[[ARG1]], %[[C0]]
+//       CHECK:         %[[T8:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T7]]]
+//       CHECK:         %[[T9:.*]] = affine.apply #[[MAP1]](%[[T5]])
+//       CHECK:         %[[T10:.*]] = addi %[[T9]], %[[C1]] : index
+//       CHECK:         %[[T11:.*]] = dim %[[ARG1]], %[[C1]]
+//       CHECK:         %[[T12:.*]] = affine.min #[[MAP2]](%[[T10]], %[[ARG4]])[%[[T11]]]
+//       CHECK:         %[[T13:.*]] = affine.apply #[[MAP3]](%[[T6]])
+//       CHECK:         %[[T14:.*]] = addi %[[T13]], %[[C1]] : index
+//       CHECK:         %[[T15:.*]] = dim %[[ARG1]], %[[C2]]
+//       CHECK:         %[[T16:.*]] = affine.min #[[MAP2]](%[[T14]], %[[ARG5]])[%[[T15]]]
+//       CHECK:         %[[T17:.*]] = dim %[[ARG1]], %[[C3]]
 //       CHECK:         %[[SV1:.*]] = subview %[[ARG1]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0]
-//  CHECK-SAME:                                        [%[[T6]], %[[T8]], %[[T10]], %[[T11]]]
-//       CHECK:         %[[T13:.*]] = dim %[[ARG2]], %[[C0]]
-//       CHECK:         %[[T14:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T13]]]
-//       CHECK:         %[[T15:.*]] = dim %[[ARG2]], %[[C1]]
-//       CHECK:         %[[T16:.*]] = affine.min #[[MAP4]](%[[ARG4]])[%[[T15]]]
-//       CHECK:         %[[T17:.*]] = dim %[[ARG2]], %[[C2]]
-//       CHECK:         %[[T18:.*]] = affine.min #[[MAP5]](%[[ARG5]])[%[[T17]]]
-//       CHECK:         %[[T19:.*]] = dim %[[ARG2]], %[[C3]]
+//  CHECK-SAME:                                        [%[[T8]], %[[T12]], %[[T16]], %[[T17]]]
+//       CHECK:         %[[T18:.*]] = dim %[[ARG2]], %[[C0]]
+//       CHECK:         %[[T19:.*]] = affine.min #[[MAP0]](%[[ARG3]])[%[[T18]]]
+//       CHECK:         %[[T20:.*]] = dim %[[ARG2]], %[[C1]]
+//       CHECK:         %[[T21:.*]] = affine.min #[[MAP5]](%[[ARG4]])[%[[T20]]]
+//       CHECK:         %[[T22:.*]] = dim %[[ARG2]], %[[C2]]
+//       CHECK:         %[[T23:.*]] = affine.min #[[MAP6]](%[[ARG5]])[%[[T22]]]
+//       CHECK:         %[[T24:.*]] = dim %[[ARG2]], %[[C3]]
 //       CHECK:         %[[SV2:.*]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0]
-//  CHECK-SAME:                                        [%[[T14]], %[[T16]], %[[T18]], %[[T19]]]
-//       CHECK:         linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]])
+//  CHECK-SAME:                                        [%[[T19]], %[[T21]], %[[T23]], %[[T24]]]
+//       CHECK:         linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]])