[mlir][NFC] Extend 'getPermutedPosition' to support projected permutations

[dcaballe]

Summary

Small change to support projected permutations in the
'getPermutedPosition' utility.

[nicolasvasilache]

I am sorry but this API does not make sense to me, I think it was mistakenly added, IMO

Optional<int64_t> AffineMap::getResultPosition(AffineExpr)

makes much more sense and does not require Permutation, ProjectedPermutation, PermutationWithMaybeSomeBroadcast or other characterizations; if works with any AffineMap.

The client is responsible for assertions (it already is in this case).

Use unsigned if you really must (for maintaining consistency with neighbouring code) but we know this is not a good use case of unsigned by now and should be globally cleaned up in the future.

[dcaballe]

I'm sorry, Nicolas, I'm not sure I follow what the AffineExpr input is.
I changed the name to getResultPos and removed the assert so that it works for any case.
Please, let me know if that is not what you wanted!

[nicolasvasilache]

In D138946#3958762, @dcaballe wrote:

I'm sorry, Nicolas, I'm not sure I follow what the AffineExpr input is.
I changed the name to getResultPos and removed the assert so that it works for any case.
Please, let me know if that is not what you wanted!

The AffineExpr is what the client passes, it can be an AffineDimExpr(pos) or not.

Client-side:

assert(order.isPermutation());
auto maybePos = order.getFirstResultPosition(AffineDimExpr::get(d, ctx));
assert(maybePos.has_value());
return *maybePos;

AffineMap-side:

Optional<unsigned> AffineMap::getFirstResultPosition(AffineExpr desiredExpr) const {
  for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {
    if (getResult(i) == desiredExpr)
      return i;
  }
  return llvm::None;
}

Actually, now that I write it, it looks a lot like a C++ find, so findResult ?

[dcaballe]

I think that is computing something different. The input in the original utility is an input dimension position and we want to figure out what the position of that dimension is in the result (assuming a projected permutation). We don't have the result expression.
For example, for (d0, d1, d2, d3) -> (d1, d0), the input to the utility would be 0, 1, 2 or 3, and the utility would return 1, 0, None, None, respectively, for those inputs. It's kind of applying the map to an input dimension position. Does it make sense?

[nicolasvasilache]

In D138946#3959227, @dcaballe wrote:

I think that is computing something different. The input in the original utility is an input dimension position and we want to figure out what the position of that dimension is in the result (assuming a projected permutation). We don't have the result expression.
For example, for (d0, d1, d2, d3) -> (d1, d0), the input to the utility would be 0, 1, 2 or 3, and the utility would return 1, 0, None, None, respectively, for those inputs. It's kind of applying the map to an input dimension position. Does it make sense?

Yes absolutely, and taking your example, what I am suggesting is that the input to the utility would be AffineDimExpr::get(0, ctx), AffineDimExpr::get(1, ctx),AffineDimExpr::get(2, ctx),AffineDimExpr::get(3, ctx) and the utility would return 1, 0, None, None, respectively, for those inputs.
This makes the API take a generic AffineExpr input and avoid passing magic integers that get a special meaning.
This is more general and less surprising to use.

[dcaballe]

Addressed feedback. Thanks!

[nicolasvasilache]

mlir/include/mlir/IR/AffineMap.h

169–172

Please update the comment.

Extracts the first result position where `input` resides.
Return llvm::None if `input` cannot be found.

mlir/include/mlir/IR/AffineMap.h

public:
 
   ArrayRef<AffineExpr> getResults() const;
   AffineExpr getResult(unsigned idx) const;
 
   /// Extracts the position of the dimensional expression at the given result,
   /// 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;
+  /// Extracts the first result position where `input` dimension resides.
+  /// Returns `llvm::None` if `input` is not a dimension expression or cannot be
+  /// found in results.
+  Optional<unsigned> getResultPosition(AffineExpr input) const;
 
   /// Return true if any affine expression involves AffineDimExpr `position`.
   bool isFunctionOfDim(unsigned position) const {
     return llvm::any_of(getResults(), [&](AffineExpr e) {
       return e.isFunctionOfDim(position);
     });
   }
 

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 }
 
 uint64_t mlir::sparse_tensor::toStoredDim(const SparseTensorEncodingAttr &enc,
                                           uint64_t d) {
   if (enc) {
     auto order = enc.getDimOrdering();
     if (order) {
       assert(order.isPermutation());
-      return order.getPermutedPosition(d);
+      auto maybePos =
+          order.getResultPosition(getAffineDimExpr(d, enc.getContext()));
+      assert(maybePos.has_value());
+      return *maybePos;
     }
   }
   return d;
 }
 
 uint64_t mlir::sparse_tensor::toOrigDim(RankedTensorType type, uint64_t d) {
   assert(d < static_cast<uint64_t>(type.getRank()));
   return toOrigDim(getSparseTensorEncoding(type), d);

mlir/lib/IR/AffineMap.cpp

 AffineExpr AffineMap::getResult(unsigned idx) const {
   return getResults()[idx];
 }
 
 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)
+Optional<unsigned> AffineMap::getResultPosition(AffineExpr input) const {
+  if (!input.isa<AffineDimExpr>())
+    return llvm::None;
+
+  for (unsigned i = 0, numResults = getNumResults(); i < numResults; i++) {
+    if (getResult(i) == input)
       return i;
-  llvm_unreachable("incorrect permutation request");
+  }
+
+  return llvm::None;
 }
 
 /// 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.
 LogicalResult
 AffineMap::constantFold(ArrayRef<Attribute> operandConstants,
                         SmallVectorImpl<Attribute> &results) const {