[mlir] refactor common idiom into AffineMap method

[aartbik]

Summary

motivated by a refactoring in the new sparse code (yet to be merged), this avoids some lengthy code dup

[mehdi_amini]

mlir/include/mlir/IR/AffineMap.h

128

Can you add an API doc?

mlir/lib/IR/AffineMap.cpp

232

You could save the assert with:

getResult(idx).cast<AffineDimExpr>().getPosition();

[aartbik]

Thanks Medhi.

mlir/lib/IR/AffineMap.cpp

232

We could do that for getResult itself to and call the index less one (didn't do that in this CL though)

[aartbik]

addressed comments

mlir/include/mlir/IR/AffineMap.h

public:
   unsigned getNumDims() const;
   unsigned getNumSymbols() const;
   unsigned getNumResults() const;
   unsigned getNumInputs() const;
 
   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;
+
   /// Walk all of the AffineExpr's in this mapping. Each node in an expression
   /// tree is visited in postorder.
   void walkExprs(std::function<void(AffineExpr)> callback) const;
 
   /// This method substitutes any uses of dimensions and symbols (e.g.
   /// dim#0 with dimReplacements[0]) in subexpressions and returns the modified
   /// expression mapping.  Because this can be used to eliminate dims and
   /// symbols, the client needs to specify the number of dims and symbols in

mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp

fuseWithReshapeByExpansion(LinalgOp linalgOp, TensorReshapeOp reshapeOp,
   SmallVector<unsigned, 4> numFoldedDims(foldedType.getRank(), 0);
   SmallVector<SmallVector<int64_t, 4>, 4> expandedDimsShape(
       foldedType.getRank());
   auto reassociationMaps = reshapeOp.getReassociationMaps();
   for (auto resultExpr : llvm::enumerate(fusedIndexMap.getResults())) {
     unsigned pos = resultExpr.value().cast<AffineDimExpr>().getPosition();
     AffineMap foldedDims = reassociationMaps[resultExpr.index()];
     numFoldedDims[pos] = foldedDims.getNumResults();
-    ArrayRef<int64_t> shape = expandedShape.slice(
-        foldedDims.getResult(0).cast<AffineDimExpr>().getPosition(),
-        numFoldedDims[pos]);
+    ArrayRef<int64_t> shape =
+        expandedShape.slice(foldedDims.getDimPosition(0), numFoldedDims[pos]);
     expandedDimsShape[pos].assign(shape.begin(), shape.end());
   }
 
   if (isa<IndexedGenericOp>(linalgOp.getOperation())) {
     // For indexed generic op, the region contains arguments that represent the
     // induction variable value of the loops. In the fused op these values are
     // obtained by linearizing the expanded dimensions. For now just check that
     // the extents used in the linearization (all the expanded dims except the

mlir/lib/Dialect/Vector/VectorOps.cpp

if (expectedResultDims.size() == 0) {
     AffineMap lhsMap = op.getIndexingMaps()[0];
     AffineMap rhsMap = op.getIndexingMaps()[1];
     SmallVector<AffineExpr, 4> extents(lhsMap.getNumInputs());
     for (auto pair :
          {std::make_pair(lhsType, lhsMap), std::make_pair(rhsType, rhsMap)}) {
       VectorType v = pair.first;
       auto map = pair.second;
       for (unsigned idx = 0, e = v.getRank(); idx < e; ++idx) {
-        unsigned pos = map.getResult(idx).cast<AffineDimExpr>().getPosition();
+        unsigned pos = map.getDimPosition(idx);
         if (!extents[pos])
           extents[pos] = getAffineConstantExpr(v.getShape()[idx], ctx);
       }
     }
     assert(llvm::all_of(extents, [](AffineExpr e) { return e; }) &&
            "expected extent along all dimensions.");
 
     AffineMap resMap = op.getIndexingMaps()[2];
while (insertOp || transposeOp) {
     // However, the rank of `insertedPos` may be different from that of the
     // `permutationMap`. To support such case, we need to:
     //   1. apply on the `insertedPos.size()` major dimensions
     //   2. check the other dimensions of the permutation form a minor identity.
     assert(permutationMap.isPermutation() && "expected a permutation");
     if (insertedPos.size() == extractedPos.size()) {
       bool fold = true;
       for (unsigned idx = 0, sz = extractedPos.size(); idx < sz; ++idx) {
-        auto pos =
-            permutationMap.getResult(idx).cast<AffineDimExpr>().getPosition();
+        auto pos = permutationMap.getDimPosition(idx);
         if (pos >= sz || insertedPos[pos] != extractedPos[idx]) {
           fold = false;
           break;
         }
       }
       if (fold) {
         assert(permutationMap.getNumResults() >= insertedPos.size() &&
                "expected map of rank larger than insert indexing");

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 #define DEBUG_TYPE "vector-to-vector"
 
 using namespace mlir;
 using llvm::dbgs;
 
 // Helper to find an index in an affine map.
 static Optional<int64_t> getResultIndex(AffineMap map, int64_t index) {
   for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
-    int64_t idx = map.getResult(i).cast<AffineDimExpr>().getPosition();
+    int64_t idx = map.getDimPosition(i);
     if (idx == index)
       return i;
   }
   return None;
 }
 
 // Helper to construct iterator types with one index removed.
 static SmallVector<Attribute, 4> adjustIter(ArrayAttr iteratorTypes,

 }
 
 // Helper to construct an affine map with one index removed.
 static AffineMap adjustMap(AffineMap map, int64_t index,
                            PatternRewriter &rewriter) {
   auto *ctx = rewriter.getContext();
   SmallVector<AffineExpr, 4> results;
   for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
-    int64_t idx = map.getResult(i).cast<AffineDimExpr>().getPosition();
+    int64_t idx = map.getDimPosition(i);
     if (idx == index)
       continue;
     // Re-insert remaining indices, but renamed when occurring
     // after the removed index.
     auto targetExpr = getAffineDimExpr(idx < index ? idx : idx - 1, ctx);
     results.push_back(targetExpr);
   }
   return AffineMap::get(map.getNumDims() - 1, 0, results, ctx);
Value ContractionOpLowering::lowerParallel(vector::ContractionOp op,
   VectorType lhsType = op.getLhsType();
   VectorType rhsType = op.getRhsType();
   VectorType resType = op.getResultType().cast<VectorType>();
   // Find the iterator type index and result index.
   SmallVector<AffineMap, 4> iMap = op.getIndexingMaps();
   int64_t iterIndex = -1;
   int64_t dimSize = -1;
   if (lhsIndex >= 0) {
-    iterIndex = iMap[0].getResult(lhsIndex).cast<AffineDimExpr>().getPosition();
-    assert(
-        (rhsIndex < 0 ||
-         iterIndex ==
-             iMap[1].getResult(rhsIndex).cast<AffineDimExpr>().getPosition()) &&
+    iterIndex = iMap[0].getDimPosition(lhsIndex);
+    assert((rhsIndex < 0 || iterIndex == iMap[1].getDimPosition(rhsIndex)) &&
            "parallel index should be free in LHS or batch in LHS/RHS");
     dimSize = lhsType.getDimSize(lhsIndex);
   } else {
     assert(rhsIndex >= 0 && "missing parallel index");
-    iterIndex = iMap[1].getResult(rhsIndex).cast<AffineDimExpr>().getPosition();
+    iterIndex = iMap[1].getDimPosition(rhsIndex);
     dimSize = rhsType.getDimSize(rhsIndex);
   }
   assert(iterIndex >= 0 && "parallel index not listed in operand mapping");
   Optional<int64_t> lookup = getResultIndex(iMap[2], iterIndex);
   assert(lookup.hasValue() && "parallel index not listed in reduction");
   int64_t resIndex = lookup.getValue();
   // Construct new iterator types and affine map array attribute.
   std::array<AffineMap, 3> lowIndexingMaps = {

mlir/lib/IR/AffineMap.cpp

ArrayRef<AffineExpr> AffineMap::getResults() const {
   assert(map && "uninitialized map storage");
   return map->results;
 }
 AffineExpr AffineMap::getResult(unsigned idx) const {
   assert(map && "uninitialized map storage");
   return map->results[idx];
 }
 
+unsigned AffineMap::getDimPosition(unsigned idx) const {
+  return getResult(idx).cast<AffineDimExpr>().getPosition();
+}
+
 /// 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 {
   // Attempt partial folding.
   SmallVector<int64_t, 2> integers;