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

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.

Differential Revision: https://reviews.llvm.org/D83756

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

@@ -58,31 +58,17 @@ static SmallVector<Value, 4> permuteIvs(ArrayRef<Value> 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
+/// The returned ranges correspond to the loop ranges, in the proper order, for
-/// map results. It expects a non-inverted, concatenated map and last values in
+/// which new loops will be created.
-/// viewSizes will be applied to the symbols in the map.
+/// The function supports only maps that are invertible and have results of type
-static SmallVector<SubViewOp::Range, 4>
+/// DimExpr or (DimExpr + DimExpr - SymbolExpr floordiv ConstExpr).
-emitLoopRangesWithSymbols(OpBuilder &b, Location loc, AffineMap map,
+/// It expects a non-inverted, concatenated map and last values in
-                          ValueRange viewSizes) {
+/// 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 &&
@@ -537,23 +523,8 @@ Optional<LinalgLoops> linalgOpToLoopsImpl(Operation *op, OpBuilder &builder) {
       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;
+  auto loopRanges = emitLoopRanges(scope.getBuilderRef(), scope.getLocation(),
-
+                                   map, getViewSizes(builder, linalgOp));
-  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);
-  }
   SmallVector<Value, 4> allIvs;
   GenerateLoopNest<LoopTy>::doit(
       loopRanges, linalgOp.iterator_types().getValue(), [&](ValueRange ivs) {