[Mlir-commits] [mlir] 7ab21cd - [mlir] Fix `AffineMap.dropResults`.

[Oleg Shyshkov]

Author: Oleg Shyshkov
Date: 2022-10-27T13:34:55+02:00
New Revision: 7ab21cdc208491c5f935b9fbc18c3436a6e14c1f

URL: https://github.com/llvm/llvm-project/commit/7ab21cdc208491c5f935b9fbc18c3436a6e14c1f
DIFF: https://github.com/llvm/llvm-project/commit/7ab21cdc208491c5f935b9fbc18c3436a6e14c1f.diff

LOG: [mlir] Fix `AffineMap.dropResults`.

`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.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/AffineMap.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index 0c1d6bdde9af..ccd167dfa5ab 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -243,19 +243,18 @@ class AffineMap {
 
   /// 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