From 7ab21cdc208491c5f935b9fbc18c3436a6e14c1f Mon Sep 17 00:00:00 2001 From: Oleg Shyshkov Date: Thu, 27 Oct 2022 13:32:52 +0200 Subject: [PATCH] [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 --- mlir/include/mlir/IR/AffineMap.h | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h index 0c1d6bd..ccd167d 100644 --- a/mlir/include/mlir/IR/AffineMap.h +++ b/mlir/include/mlir/IR/AffineMap.h @@ -243,19 +243,18 @@ public: /// 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 positions) { - AffineMap resultMap = *this; - for (int64_t pos : positions) - resultMap = resultMap.dropResult(pos); - return resultMap; + SmallVector reverse_sorted_positions = llvm::to_vector(positions); + llvm::sort(reverse_sorted_positions, std::greater()); + + 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 -- 2.7.4