[mlir][linalg] Use affine apply in im2col gather index calculations

[ThomasRaoux]

Thanks!

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

 //===- ConvertConv2DToImg2Col.cpp - im2col implementation -----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"

 
 static Value createMul(Location loc, Value x, Value y, OpBuilder &builder) {
   bool isInt = x.getType().isa<IntegerType>();
   if (isInt)
     return builder.create<arith::MulIOp>(loc, x, y);
   return builder.create<arith::MulFOp>(loc, x, y);
 }
 
-// Unrolls the given composite `index` into a set of subindices with maximum
-// iteration ranges specified by `factors` according to the following
-// assumptions:
-//   1. The iteration range for `index` is [0, f1 * f2 * ... * fn] i.e. the
-//   product of the given list of factors
-//   2. The iterators corresponding to the entries in `factors` are ordered from
-//   slowest to fastest varying
-// Each subindex is then computed as:
-//    subindex[i] = floor( (index % (fi * ... * fn)) / (fi-1 * ... * fn) )
-static SmallVector<Value, 3> unrollIndex(OpBuilder &b, Location loc,
-                                         Value index,
-                                         ArrayRef<int64_t> factors) {
+// Delinearizes the given composite `index` by the basis specified in `factors`.
+static SmallVector<Value> unrollIndex(OpBuilder &b, Location loc, Value index,
   assert(factors.size() >= 1 && "empty factor list");
-  SmallVector<Value, 3> indices(factors.size());
-  int64_t runningProd = 1;
-  for (int i = factors.size() - 1, end = 0; i >= end; i--) {
-    Value unrolledIndex = index;
-    if (i > 0) {
-      Value modBase = b.create<arith::ConstantOp>(
-          loc, b.getIndexAttr(runningProd * factors[i]));
-      unrolledIndex = b.create<arith::RemUIOp>(loc, unrolledIndex, modBase);
-    }
-    if (runningProd > 1) {
-      Value divDenom =
-          b.create<arith::ConstantOp>(loc, b.getIndexAttr(runningProd));
-      unrolledIndex = b.create<arith::DivUIOp>(loc, unrolledIndex, divDenom);
-    }
-    runningProd *= factors[i];
-    indices[i] = unrolledIndex;
-  }
-  return indices;
+  SmallVector<Value> basis;
+  for (int64_t f : factors)
+    basis.push_back(b.create<arith::ConstantOp>(loc, b.getIndexAttr(f)));
+  FailureOr<SmallVector<Value>> multiIndex =
+      delinearizeIndex(b, loc, index, basis);
+  assert(!failed(multiIndex) && "Failed to linearize img2col index");
+  return *multiIndex;
 }
 
 // Given indices corresponding to iterators in the output (oIndex) and filter
 // (fIndex) for a convolution, compute the convolved index for the
 // input as `oIndex * stride + fIndex`.
 static Value getConvolvedIndex(OpBuilder &b, Location loc, Value oIndex,
                                Value fIndex, int64_t stride) {
-  Value strideVal = b.create<arith::ConstantOp>(loc, b.getIndexAttr(stride));
-  Value convIndex = b.create<arith::MulIOp>(loc, oIndex, strideVal);
-  return b.create<arith::AddIOp>(loc, convIndex, fIndex);
+  AffineExpr oExpr, fExpr;
+  bindSymbols(b.getContext(), oExpr, fExpr);
+  AffineMap convMap = AffineMap::get(0, 2, stride * oExpr + fExpr);
+  return makeComposedAffineApply(b, loc, convMap, ValueRange{oIndex, fIndex});
 }
 
 FailureOr<std::pair<Operation *, Operation *>>
 rewriteInIm2Col(RewriterBase &rewriter, linalg::Conv2DNhwcHwcfOp convOp) {
   auto inputType = convOp.getInputs()[0].getType().cast<ShapedType>();
   auto filterType = convOp.getInputs()[1].getType().cast<ShapedType>();
   auto outputType = convOp.getOutputs()[0].getType().cast<ShapedType>();
 
auto img2ColTensor = rewriter.create<linalg::GenericOp>(
       img2colIterators,
       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
         // Get the iterators named based on the matmul (batch, m, k).
         Value bIndex = nestedBuilder.create<linalg::IndexOp>(loc, 0);
         Value mIndex = nestedBuilder.create<linalg::IndexOp>(loc, 1);
         Value kIndex = nestedBuilder.create<linalg::IndexOp>(loc, 2);
 
         // Recover the original iteration indices from the problem/input sizes.
-        SmallVector<Value, 3> mIndices = unrollIndex(
+        SmallVector<Value> mIndices = unrollIndex(
             nestedBuilder, nestedLoc, mIndex, ArrayRef<int64_t>{oh, ow});
         auto ohIndex = mIndices[0];
         auto owIndex = mIndices[1];
 
-        SmallVector<Value, 3> kIndices = unrollIndex(
+        SmallVector<Value> kIndices = unrollIndex(
             nestedBuilder, nestedLoc, kIndex, ArrayRef<int64_t>{fh, fw, ic});
         auto fhIndex = kIndices[0];
         auto fwIndex = kIndices[1];
         auto icIndex = kIndices[2];
 
         // Extract the input element corresponding to the expanded indices.
         Value hIndex =
             getConvolvedIndex(nestedBuilder, nestedLoc, ohIndex, fhIndex,
auto img2ColTensor = rewriter.create<linalg::GenericOp>(
       img2colIterators,
       [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
         // Get the iterators named based on the matmul (batch, m, k).
         Value bIndex = nestedBuilder.create<linalg::IndexOp>(loc, 0);
         Value kIndex = nestedBuilder.create<linalg::IndexOp>(loc, 1);
         Value nIndex = nestedBuilder.create<linalg::IndexOp>(loc, 2);
 
         // Recover the original iteration indices from the problem/input sizes.
-        SmallVector<Value, 3> kIndices = unrollIndex(
+        SmallVector<Value> kIndices = unrollIndex(
             nestedBuilder, nestedLoc, kIndex, ArrayRef<int64_t>{ic, fh, fw});
         auto icIndex = kIndices[0];
         auto fhIndex = kIndices[1];
         auto fwIndex = kIndices[2];
 
-        SmallVector<Value, 3> nIndices = unrollIndex(
+        SmallVector<Value> nIndices = unrollIndex(
             nestedBuilder, nestedLoc, nIndex, ArrayRef<int64_t>{oh, ow});
         auto ohIndex = nIndices[0];
         auto owIndex = nIndices[1];
 
         // Extract the input element corresponding to the expanded indices.
         Value hIndex =
             getConvolvedIndex(nestedBuilder, nestedLoc, ohIndex, fhIndex,
                               convOp.getStrides().getValues<int64_t>()[0]);

mlir/test/Dialect/Linalg/convert-conv2d-to-img2col.mlir

 // CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>]
 // CHECK: ^bb0(%[[OUT_DATA:.+]]: f32)
 
 // Collapsed indices.
 // CHECK: %[[BINDEX:.+]] = linalg.index 0 : index
 // CHECK: %[[MINDEX:.+]] = linalg.index 1 : index
 // CHECK: %[[KINDEX:.+]] = linalg.index 2 : index
 
-// Unrolled output shape indices.
-// CHECK: %[[C14:.+]] = arith.constant 14 : index
-// CHECK: %[[OWINDEX:.+]] = arith.remui %[[MINDEX]], %[[C14]] : index
-// CHECK: %[[C14_1:.+]] = arith.constant 14 : index
-// CHECK: %[[OHINDEX:.+]] = arith.divui %[[MINDEX]], %[[C14_1]] : index
-
-// Unrolled filter shape indices.
-// CHECK: %[[C4:.+]] = arith.constant 4 : index
-// CHECK: %[[ICINDEX:.+]] = arith.remui %[[KINDEX]], %[[C4]] : index
-// CHECK: %[[C12:.+]] = arith.constant 12 : index
-// CHECK: %[[FWREM:.+]] = arith.remui %[[KINDEX]], %[[C12]] : index
-// CHECK: %[[C4_2:.+]] = arith.constant 4 : index
-// CHECK: %[[FWINDEX:.+]] = arith.divui %[[FWREM]], %[[C4_2]] : index
-// CHECK: %[[C12_3:.+]] = arith.constant 12 : index
-// CHECK: %[[FHINDEX:.+]] = arith.divui %[[KINDEX]], %[[C12_3]] : index
-
-// Compute input indices.
-// CHECK: %[[SH:.+]] = arith.constant 1 : index
-// CHECK: %[[STRIDEDOH:.+]] = arith.muli %[[OHINDEX]], %[[SH]] : index
-// CHECK: %[[CONVH:.+]] = arith.addi %[[STRIDEDOH]], %[[FHINDEX]] : index
-// CHECK: %[[SW:.+]] = arith.constant 1 : index
-// CHECK: %[[STRIDEDOW:.+]] = arith.muli %[[OWINDEX]], %[[SW]] : index
-// CHECK: %[[CONVW:.+]] = arith.addi %[[STRIDEDOW]], %[[FWINDEX]] : index
+// Compute input channel/convolved indices.
+// CHECK: %[[ICINDEX:.+]] = affine.apply affine_map<(d0) -> (d0 mod 4)>(%[[KINDEX]])
+// CHECK: %[[CONVH:.+]] = affine.apply affine_map<(d0, d1) -> (d0 floordiv 14 + d1 floordiv 12)>(%[[MINDEX]], %[[KINDEX]])
+// CHECK: %[[CONVW:.+]] = affine.apply affine_map<(d0, d1) -> (d0 mod 14 + (d1 mod 12) floordiv 4)>(%[[MINDEX]], %[[KINDEX]])
+
+// Extract from the input tensor.
 // CHECK: %[[EXTRACTED_INPUT:.+]] = tensor.extract
 // CHECK-SAME: %{{.+}}{{\[}}%[[BINDEX]], %[[CONVH]], %[[CONVW]], %[[ICINDEX]]] : tensor<1x16x16x4xf32>
 // CHECK: linalg.yield %[[EXTRACTED_INPUT]] : f32
 
 // CHECK: IR printer: transformed
 // CHECK: tensor.expand_shape %{{[^ ]*}} {{\[}}[0], [1, 2], [3]] : tensor<1x196x16xf32> into tensor<1x14x14x16xf32>
 
 // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>

 ^bb1(%arg1: !pdl.operation):
   %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!pdl.operation) -> !pdl.operation
   %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!pdl.operation) -> (!pdl.operation, !pdl.operation)
 }
 
 // -----
 
 //  CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+
+//  Im2col maps
+//  CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0 floordiv 9)>
+//  CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1) -> (d0 floordiv 14 + (d1 mod 9) floordiv 3)>
+//  CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1) -> (d0 + d1 - (d0 floordiv 14) * 14 - (d1 floordiv 3) * 3)>
+
+
 //  CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
 //  CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
 //  CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
 
 //      CHECK: func.func @batch_nchw_conv
 // CHECK-SAME: (%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.+]]: tensor<16x4x3x3xf32>, %[[INIT:.+]]: tensor<8x16x14x14xf32>)
 //  CHECK-DAG:   %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xf32> into tensor<16x36xf32>
 //  CHECK-DAG:   %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]] : tensor<8x16x14x14xf32> into tensor<8x16x196xf32>
 //      CHECK:   %[[IT:.+]] = tensor.empty() : tensor<8x36x196xf32>
 //      CHECK:   %[[IMG2COL:.+]] = linalg.generic
 // CHECK-SAME:      indexing_maps = [#[[MAP]]]
 // CHECK-SAME:      iterator_types = ["parallel", "parallel", "parallel"]
 // CHECK-SAME:   outs(%[[IT]] : tensor<8x36x196xf32>)
 //      Collapsed indices.
 //      CHECK:       %[[BINDEX:.+]] = linalg.index 0 : index
 //      CHECK:       %[[KINDEX:.+]] = linalg.index 1 : index
 //      CHECK:       %[[NINDEX:.+]] = linalg.index 2 : index
 
-//      Unrolled filter shape indices.
-//      CHECK:       %[[C3:.+]] = arith.constant 3 : index
-//      CHECK:       %[[FWINDEX:.+]] = arith.remui %[[KINDEX]], %[[C3]] : index
-//      CHECK:       %[[C9:.+]] = arith.constant 9 : index
-//      CHECK:       %[[FHREM:.+]] = arith.remui %[[KINDEX]], %[[C9]] : index
-//      CHECK:       %[[C3_1:.+]] = arith.constant 3 : index
-//      CHECK:       %[[FHINDEX:.+]] = arith.divui %[[FHREM]], %[[C3_1]] : index
-//      CHECK:       %[[C9_2:.+]] = arith.constant 9 : index
-//      CHECK:       %[[ICINDEX:.+]] = arith.divui %[[KINDEX]], %[[C9_2]] : index
-
-//      Unrolled output shape indices.
-//      CHECK:       %[[C14:.+]] = arith.constant 14 : index
-//      CHECK:       %[[OWINDEX:.+]] = arith.remui %[[NINDEX]], %[[C14]] : index
-//      CHECK:       %[[C14_3:.+]] = arith.constant 14 : index
-//      CHECK:       %[[OHINDEX:.+]] = arith.divui %[[NINDEX]], %[[C14_3]] : index
-
-//      Compute input indices.
-//      CHECK:       %[[SH:.+]] = arith.constant 1 : index
-//      CHECK:       %[[STRIDEDOH:.+]] = arith.muli %[[OHINDEX]], %[[SH]] : index
-//      CHECK:       %[[CONVH:.+]] = arith.addi %[[STRIDEDOH]], %[[FHINDEX]] : index
-//      CHECK:       %[[SW:.+]] = arith.constant 1 : index
-//      CHECK:       %[[STRIDEDOW:.+]] = arith.muli %[[OWINDEX]], %[[SW]] : index
-//      CHECK:       %[[CONVW:.+]] = arith.addi %[[STRIDEDOW]], %[[FWINDEX]] : index
+//      Compute input channel/convolved indices.
+//      CHECK:       %[[ICINDEX:.+]] = affine.apply #[[MAP1]](%[[KINDEX]])
+//      CHECK:       %[[CONVH:.+]] = affine.apply #[[MAP7]](%[[NINDEX]], %[[KINDEX]])
+//      CHECK:       %[[CONVW:.+]] = affine.apply #[[MAP8]](%[[NINDEX]], %[[KINDEX]])
+
+//      Extract from the input tensor.
 //      CHECK:       %[[EXTRACTED_INPUT:.+]] = tensor.extract
 //      CHECK-SAME:  %[[INPUT]]{{\[}}%[[BINDEX]], %[[ICINDEX]], %[[CONVH]], %[[CONVW]]] : tensor<8x4x16x16xf32>
 //      CHECK: linalg.yield %[[EXTRACTED_INPUT]] : f32
 //      CHECK:   %[[MATMUL:.+]] = linalg.generic
 // CHECK-SAME:      indexing_maps = [#[[LHSMAP]], #[[RHSMAP]], #[[RESMAP]]],
 // CHECK-SAME:      iterator_types = ["parallel", "parallel", "parallel", "reduction"]
 // CHECK-SAME:   ins(%[[CS_FILTER]], %[[IMG2COL]] : tensor<16x36xf32>, tensor<8x36x196xf32>)
 // CHECK-SAME:   outs(%[[CS_RESULT]] : tensor<8x16x196xf32>)