[Mlir-commits] [mlir] a643bd3 - [mlir] add permutation utility

[Aart Bik]

Author: Aart Bik
Date: 2021-08-24T08:07:40-07:00
New Revision: a643bd3189aeee138187509a9bfec2f798798d76

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

LOG: [mlir] add permutation utility

I found myself typing this code several times at different places
by now, so time to make this a general utility instead. Given
a permutation, it returns the permuted position of the input,
for example (i,j,k) -> (k,i,j) yields position 1 for input 0.

Reviewed By: ftynse

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

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index f687253b36fe..906c53db4b32 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -162,6 +162,10 @@ class AffineMap {
   /// when the caller knows it is safe to do so.
   unsigned getDimPosition(unsigned idx) const;
 
+  /// Extracts the permuted position where given input index resides.
+  /// Fails when called on a non-permutation.
+  unsigned getPermutedPosition(unsigned input) const;
+
   /// Return true if any affine expression involves AffineDimExpr `position`.
   bool isFunctionOfDim(unsigned position) const {
     return llvm::any_of(getResults(), [&](AffineExpr e) {

diff  --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 2257617bd903..9c6f25d3c53e 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -336,6 +336,14 @@ unsigned AffineMap::getDimPosition(unsigned idx) const {
   return getResult(idx).cast<AffineDimExpr>().getPosition();
 }
 
+unsigned AffineMap::getPermutedPosition(unsigned input) const {
+  assert(isPermutation() && "invalid permutation request");
+  for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++)
+    if (getDimPosition(i) == input)
+      return i;
+  llvm_unreachable("incorrect permutation request");
+}
+
 /// Folds the results of the application of an affine map on the provided
 /// operands to a constant if possible. Returns false if the folding happens,
 /// true otherwise.