[mlir][Linalg] emitLoopRanges and emitLoopRangesWithSymbols merged into one

[limo1996]

Summary

Right now there is a branching for 2 functions based on whether target map has
symbols or not. In this commit these functions are merged into one.
Furthermore, emitting does not require inverse and map applying as it computes
the correct Range in a single step and thus reduces unnecessary overhead.

[bondhugula]

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

115

Avoid auto here.

527

Avoid repeated evaluation of map.getNumSymbols().

[limo1996]

comments of bondhugula incorporated

[limo1996]

Comments of bondhugula marked as done

[bondhugula]

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

61–71

Nit: if it contains some -> if it contains any

[limo1996]

some -> any

[ftynse]

Can you please elaborate in the commit description and/or code comments _how_ this avoids inversion? The code seems to just pass in the direct map to the range calculation function, which takes result sizes in order. I don't see why this would work, but I may be missing something.

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

61–62

Could you please elaborate _which_ expressions are supported?

83

The change from d.getPosition to idx looks wrong. The former returns the index of the affine dimension in the expression and the latter is merely an index of the expression in the map. Has this not broken any tests?

100–110

Any chance of finding better names than m and n?

113

The RHS of floordiv is not necessarily 2 in the code, but it must be some constant AFAICS, please check for that.

525–526

Please terminate sentences with a full stop in comments.

[ftynse]

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

83

After a further investigation, the code on which this is based is wrong, but the updated code looks correct.

91

Please address the comments on the diff that introduces this code, then rebase this diff.

[limo1996]

ftynse's comments addressed

[limo1996]

evolve

[limo1996]

evolve

[limo1996]

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

91

Done

[ftynse]

Thanks for the cleanup!

[limo1996]

Changes evolved and merged from parent commit

[nicolasvasilache]

Looks great, thanks!

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

 static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> ivs,
                                         Optional<AffineMap> permutation) {
   return permutation ? applyMapToValues(ScopedContext::getBuilderRef(),
                                         ScopedContext::getLocation(),
                                         permutation.getValue(), ivs)
                      : SmallVector<Value, 4>(ivs.begin(), ivs.end());
 }
 
-// Creates a number of ranges equal to the number of results in `map`.
-// The returned ranges correspond to the loop ranges, in the proper order, for
-// which new loops will be created.
-static SmallVector<SubViewOp::Range, 4>
-emitLoopRanges(OpBuilder &b, Location loc, AffineMap map,
-               ArrayRef<Value> allViewSizes) {
-  // Apply `map` to get view sizes in loop order.
-  auto sizes = applyMapToValues(b, loc, map, allViewSizes);
-  // Create a new range with the applied tile sizes.
-  ScopedContext scope(b, loc);
-  SmallVector<SubViewOp::Range, 4> res;
-  for (unsigned idx = 0, e = map.getNumResults(); idx < e; ++idx) {
-    res.push_back(SubViewOp::Range{std_constant_index(0), sizes[idx],
-                                   std_constant_index(1)});
-  }
-  return res;
-}
-
 /// Creates a number of ranges equal to the number of dimensions in the `map`.
-/// The function supports for now only limited number of expressions inside
-/// map results. It expects a non-inverted, concatenated map and last values in
-/// viewSizes will be applied to the symbols in the map.
-static SmallVector<SubViewOp::Range, 4>
-emitLoopRangesWithSymbols(OpBuilder &b, Location loc, AffineMap map,
+/// The returned ranges correspond to the loop ranges, in the proper order, for
+/// which new loops will be created.
+/// The function supports only maps that are invertible and have results of type
+/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
+/// It expects a non-inverted, concatenated map and last values in
+/// allViewSizes will be applied to the symbols in the map if it contains any.
+static SmallVector<SubViewOp::Range, 4> emitLoopRanges(OpBuilder &b,
+                                                       Location loc,
+                                                       AffineMap map,
                                                        ValueRange viewSizes) {
   unsigned numDims = map.getNumDims(), numRes = map.getNumResults();
   unsigned numSym = map.getNumSymbols();
   assert(viewSizes.size() == numRes + numSym &&
          "viewSizes must contain sizes of all views and values for symbols");
   SmallVector<SubViewOp::Range, 4> res(numDims);
   for (unsigned idx = 0; idx < numRes; ++idx) {
     auto result = map.getResult(idx);
     if (auto d = result.dyn_cast<AffineDimExpr>()) {
       if (res[d.getPosition()].offset)
         continue;
       res[d.getPosition()] = SubViewOp::Range{
           std_constant_index(0), viewSizes[idx], std_constant_index(1)};
     }
 
     // If the access pattern is of form (m, n)[s] -> (m + n - s floordiv 2),
     // then the bounds are:
     //   (s floordiv 2) <= m <= (size(m) + s floordiv 2 - s + 1).
     // where size(n) is applied to the symbol s.
     // This is done statically now.
     if (auto binOp = result.dyn_cast<AffineBinaryOpExpr>()) {
       auto lhs = binOp.getLHS().dyn_cast<AffineBinaryOpExpr>();
       auto rhs = binOp.getRHS().dyn_cast<AffineBinaryOpExpr>();
       if (!lhs || !rhs || binOp.getKind() != AffineExprKind::Add ||
           lhs.getKind() != AffineExprKind::Add ||
           rhs.getKind() != mlir::AffineExprKind::Mul)
         continue;
 
       auto m = lhs.getLHS().dyn_cast<AffineDimExpr>();
       auto n = lhs.getRHS().dyn_cast<AffineDimExpr>();
       auto fDiv = rhs.getLHS().dyn_cast<AffineBinaryOpExpr>();
       auto minusOne = rhs.getRHS().dyn_cast<AffineConstantExpr>();
       if (!m || !n || !fDiv || !minusOne ||
           fDiv.getKind() != AffineExprKind::FloorDiv ||
           fDiv.getLHS().getKind() != AffineExprKind::SymbolId ||
           fDiv.getRHS().getKind() != AffineExprKind::Constant)
         continue;
 
       auto s = fDiv.getLHS().dyn_cast<AffineSymbolExpr>();
       if (minusOne.getValue() != -1)
         continue;
 
       int mPos = m.getPosition();
       AffineExpr one = getAffineConstantExpr(1, s.getContext());
       AffineExpr sizeOfM = getAffineSymbolExpr(numSym, s.getContext());
       // Construction of upper bound (size(m) + s floordiv 2 - s + 1).
       AffineExpr upperOffsetExpr = sizeOfM + fDiv + one - s;
       AffineMap fromMap = AffineMap::get(numDims, numSym + 1, fDiv);
       AffineMap toMap = AffineMap::get(numDims, numSym + 1, upperOffsetExpr);
       SmallVector<Value, 8> values(viewSizes.begin(),
                                    viewSizes.begin() + numDims);
       values.insert(values.end(), viewSizes.begin() + numRes, viewSizes.end());
       values.push_back(viewSizes[mPos]);
Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
   auto linalgOp = cast<ConcreteOpTy>(op);
   assert(linalgOp.hasBufferSemantics() &&
          "expected linalg op with buffer semantics");
   auto mapsRange =
       linalgOp.indexing_maps().template getAsRange<AffineMapAttr>();
   auto maps = llvm::to_vector<8>(
       llvm::map_range(mapsRange, [](AffineMapAttr a) { return a.getValue(); }));
   SmallVector<Value, 8> sizes = getViewSizes(builder, linalgOp);
   AffineMap map = concatAffineMaps(maps);
-  SmallVector<SubViewOp::Range, 4> loopRanges;
-
-  if (map.getNumSymbols()) {
-    loopRanges = emitLoopRangesWithSymbols(scope.getBuilderRef(),
-                                           scope.getLocation(), map, sizes);
-  } else {
-    AffineMap invertedMap = inversePermutation(map);
-    if (!invertedMap)
-      return {};
-    if (invertedMap.isEmpty()) {
-      emitScalarImplementation<IndexedValueTy>({}, linalgOp);
-      return LinalgLoops();
-    }
-
-    loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
-                                invertedMap, sizes);
-  }
+  auto loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
+                                   map, getViewSizes(builder, linalgOp));
   SmallVector<Value, 4> allIvs;
   GenerateLoopNest<LoopTy>::doit(
       loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) {
         allIvs.append(ivs.begin(), ivs.end());
         emitScalarImplementation<IndexedValueTy>(allIvs, linalgOp);
       });
   // Number of loop ops might be different from the number of ivs since some
   // loops like affine.parallel and scf.parallel have multiple ivs.