[mlir][Linalg] Fixing bug in subview size computation in Linalg tiling.

[mravishankar]

Summary

The makeTiledViews did not use the sizes of the tiled views based on
the result of the loop bound inference computation. This manifested as
an error in computing tile sizes with convolution where not all the
result expression of concatenated affine maps are simple
AffineDimExpr.

[nicolasvasilache]

Thanks for the fix @maheshkhanwalkar

[mravishankar]

Fixing test run error

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

if (!map)
     return false;
   for (unsigned r = 0; r < map.getNumResults(); ++r)
     if (isTiled(map.getResult(r), tileSizes))
       return true;
   return false;
 }
 
 static SmallVector<Value, 4> makeTiledViews(OpBuilder &b, Location loc,
-                                            LinalgOp linalgOp,
+                                            LinalgOp linalgOp, AffineMap map,
                                             ArrayRef<Value> ivs,
                                             ArrayRef<Value> tileSizes,
-                                            ArrayRef<Value> viewSizes) {
+                                            ArrayRef<Value> allViewSizes) {
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   assert(ivs.size() == static_cast<size_t>(llvm::count_if(
                            llvm::make_range(tileSizes.begin(), tileSizes.end()),
                            [](Value v) { return !isZero(v); })) &&
          "expected as many ivs as non-zero sizes");
 
   using namespace edsc::op;
 
+  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]);
   }
auto invPermutationMap =
       AffineMap::getMultiDimIdentityMap(tileSizes.size(), b.getContext());
   if (!options.interchangeVector.empty())
     invPermutationMap = inversePermutation(AffineMap::getPermutationMap(
         options.interchangeVector, b.getContext()));
   if (!invPermutationMap)
     return llvm::None;
 
   // 2. Build the tiled loop ranges.
-  auto viewSizes = getViewSizes(b, op);
+  auto allViewSizes = getViewSizes(b, op);
   // The flattened loopToOperandRangesMaps is expected to be an invertible
   // permutation map (asserted in the inverse calculation).
   auto mapsRange = op.indexing_maps().getAsRange<AffineMapAttr>();
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   auto viewSizesToLoopsMap = inversePermutation(concatAffineMaps(maps));
   if (!viewSizesToLoopsMap)
     return llvm::None;
 
   SmallVector<SubViewOp::Range, 4> loopRanges;
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   std::tie(loopRanges, loopIndexToRangeIndex) = makeTiledLoopRanges(
-      b, scope.getLocation(), viewSizesToLoopsMap, viewSizes, tileSizes);
+      b, scope.getLocation(), viewSizesToLoopsMap, allViewSizes, tileSizes);
   if (!options.interchangeVector.empty())
     applyPermutationToVector(loopRanges, options.interchangeVector);
 
   // 3. Create the tiled loops.
   LinalgOp res = op;
   SmallVector<Value, 4> ivs;
   SmallVector<Attribute, 4> iteratorTypes =
       llvm::to_vector<4>(op.iterator_types().cast<ArrayAttr>().getValue());
GenerateLoopNest<LoopTy>::doit(
         // If we have to apply a permutation to the tiled loop nest, we have to
         // reorder the induction variables This permutation is the right one
         // assuming that loopRanges have previously been permuted by
         // (i,j,k)->(k,i,j) So this permutation should be the inversePermutation
         // of that one: (d0,d1,d2)->(d2,d0,d1)
         if (!options.interchangeVector.empty())
           ivValues = applyMapToValues(b, loc, invPermutationMap, ivValues);
 
-        auto views = makeTiledViews(b, loc, op, ivValues, tileSizes, viewSizes);
+        auto views = makeTiledViews(b, loc, op, viewSizesToLoopsMap, ivValues,
+                                    tileSizes, allViewSizes);
         auto operands = getAssumedNonViewOperands(op);
         views.append(operands.begin(), operands.end());
         res = op.clone(b, loc, views);
       });
 
   // 4. Transforms index arguments of `linalg.generic` w.r.t. to the tiling.
   transformIndexedGenericOpIndices(b, res, ivs, loopIndexToRangeIndex);
 

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: #[[$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)>
 
 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-LABEL: func @conv(
-//       TILE-23004:   %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>, %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>, %{{.*}}: memref<?x?x?x?xf32, #[[$strided4D]]>) {
-//       TILE-23004-DAG: %[[C0:.*]] = constant 0 : index
-//       TILE-23004-DAG: %[[C2:.*]] = constant 2 : index
-//       TILE-23004-DAG: %[[C3:.*]] = constant 3 : index
-//       TILE-23004-DAG: %[[C4:.*]] = constant 4 : index
-//       TILE-23004:   %[[Q:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:   %[[B:.*]] = dim %{{.*}}, %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:   %[[PaddedInput0:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:   %[[X0:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
+//       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:   %[[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:.*]] = dim %{{.*}}, %c0 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[Z1:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[Z2:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
+//       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 %{{.*}}, %c3 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[FilterView:.*]] = subview %{{.*}}[0, 0, %[[ivK]], 0] [%[[Z0]], %[[Z1]], %[[szK]], %[[K]]] [1, 1, 1, 1] : memref<?x?x?x?xf32, #[[$strided4D]]> to memref<?x?x?x?xf32, #[[$strided4D]]>
+//       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:         %[[PaddedInput0b:.*]] = dim %{{.*}}, %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[PaddedInput0]], %[[PaddedInput0b]]]
-//       TILE-23004:         %[[SZ2:.*]] = dim %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[dim3:.*]] = dim %{{.*}}, %c3
+//       TILE-23004:         %[[PaddedInput0b:.*]] = dim %[[ARG1]], %c1 : memref<?x?x?x?xf32, #[[$strided4D]]>
+//       TILE-23004:         %[[I1pStep:.*]] = affine.min #[[$S0x10p90D0x30pS1]](%[[ivJ]])[%[[Z0]], %[[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 %{{.*}}, %c2 : memref<?x?x?x?xf32, #[[$strided4D]]>
-//       TILE-23004:         %[[X1:.*]] = dim %{{.*}}, %c3 : 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)>
+
+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-LABEL: 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:         %[[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:         %[[SV2:.*]] = subview %[[ARG2]][%[[ARG3]], %[[ARG4]], %[[ARG5]], 0]
+//  CHECK-SAME:                                        [%[[T14]], %[[T16]], %[[T18]], %[[T19]]]
+//       CHECK:         linalg.conv(%[[ARG0]], %[[SV1]], %[[SV2]])
 No newline at end of file