[mlir] Fix `AffineMap.dropResults`.

[olegshyshkov]

Summary

AffineMap.dropResult erases one result from the array and it changes indexing. Calling dropResult is a loop with increasing indexes does not produce a desired result.

[akuegel]

Thanks for fixing this :)

mlir/include/mlir/IR/AffineMap.h

252

Personally I would pass std::greater<int>() instead of using rbegin()/rend()
I think this is easier readable.

[olegshyshkov]

Replace rbegin()/rend() with std::greater<int>().

[olegshyshkov]

Thank you!

mlir/include/mlir/IR/AffineMap.h

252

Fair point! And with llvm::sort it's even shorter.

mlir/include/mlir/IR/AffineMap.h

return AffineMap::get(getNumDims(), getNumSymbols() + shift,
                                 return e.shiftSymbols(getNumSymbols(), shift,
                                                       offset);
                               })),
                           getContext());
   }
 
   /// Returns a new AffineMap with the same number of dims and symbols and one
   /// less result at `pos`, dropped.
-  AffineMap dropResult(int64_t pos) {
-    auto exprs = llvm::to_vector<4>(getResults());
-    exprs.erase(exprs.begin() + pos);
-    return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
-  }
+  AffineMap dropResult(int64_t pos) { return dropResults({pos}); }
 
   // Returns a new AffineMap with the same number of dims and symbols, but all
   // positions in `positions` dropped from results.
   AffineMap dropResults(ArrayRef<int64_t> positions) {
-    AffineMap resultMap = *this;
-    for (int64_t pos : positions)
-      resultMap = resultMap.dropResult(pos);
-    return resultMap;
+    SmallVector<int64_t> reverse_sorted_positions = llvm::to_vector(positions);
+    llvm::sort(reverse_sorted_positions, std::greater<int64_t>());
+
+    auto exprs = llvm::to_vector<4>(getResults());
+    for (int64_t pos : reverse_sorted_positions)
+      exprs.erase(exprs.begin() + pos);
+    return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
   }
 
   /// Returns a new AffineMap with the same number of dims and symbols and an
   /// extra result inserted at `pos`.
   AffineMap insertResult(AffineExpr expr, unsigned pos) {
     auto exprs = llvm::to_vector<4>(getResults());
     exprs.insert(exprs.begin() + pos, expr);
     return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());