[mlir][vector] Add actualRank output parameter to createUnrollIterator()

This provides an easy way of finding the actual rank the vector type
will/can be unrolled to (which may be > the `targetRank`).

Author: MacDue

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

@@ -86,20 +86,24 @@ bool isContiguousSlice(MemRefType memrefType, VectorType vectorType);
 ///
 /// If no leading dimensions can be unrolled an empty optional will be returned.
 ///
+/// The actual rank the vector type can be unrolled to can be discovered by
+/// passing a pointer (to an int64_t) to the optional `actualRank` parameter.
+///
 /// Examples:
 ///
 ///   For vType = vector<2x3x4> and targetRank = 1
 ///
 ///   The resulting iterator will yield:
-///     [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2]
+///     [0, 0], [0, 1], [0, 2], [1, 0], [1, 1], [1, 2] (actualRank = 1)
 ///
 ///   For vType = vector<3x[4]x5> and targetRank = 0
 ///
 ///   The scalable dimension blocks unrolling so the iterator yields only:
-///     [0], [1], [2]
+///     [0], [1], [2] (actualRank = 2)
 ///
 std::optional<StaticTileOffsetRange>
-createUnrollIterator(VectorType vType, int64_t targetRank = 1);
+createUnrollIterator(VectorType vType, int64_t targetRank = 1,
+                     int64_t *actualRank = nullptr);
 
 /// A wrapper for getMixedSizes for vector.transfer_read and
 /// vector.transfer_write Ops (for source and destination, respectively).

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -285,24 +285,35 @@ bool vector::isContiguousSlice(MemRefType memrefType, VectorType vectorType) {
 }
 
 std::optional<StaticTileOffsetRange>
-vector::createUnrollIterator(VectorType vType, int64_t targetRank) {
-  if (vType.getRank() <= targetRank)
+vector::createUnrollIterator(VectorType vType, int64_t targetRank,
+                             int64_t *actualRank) {
+  auto reportActualRank = [&](int64_t rank) {
+    if (actualRank)
+      *actualRank = rank;
+  };
+  auto vectorRank = vType.getRank();
+  if (vectorRank <= targetRank) {
+    reportActualRank(vectorRank);
     return {};
+  }
   // Attempt to unroll until targetRank or the first scalable dimension (which
   // cannot be unrolled).
   auto shapeToUnroll = vType.getShape().drop_back(targetRank);
   auto scalableDimsToUnroll = vType.getScalableDims().drop_back(targetRank);
   auto it =
       std::find(scalableDimsToUnroll.begin(), scalableDimsToUnroll.end(), true);
   auto firstScalableDim = it - scalableDimsToUnroll.begin();
-  if (firstScalableDim == 0)
+  if (firstScalableDim == 0) {
+    reportActualRank(vectorRank);
     return {};
+  }
   // All scalable dimensions should be removed now.
   scalableDimsToUnroll = scalableDimsToUnroll.slice(0, firstScalableDim);
   assert(!llvm::is_contained(scalableDimsToUnroll, true) &&
          "unexpected leading scalable dimension");
   // Create an unroll iterator for leading dimensions.
   shapeToUnroll = shapeToUnroll.slice(0, firstScalableDim);
+  reportActualRank(vectorRank - shapeToUnroll.size());
   return StaticTileOffsetRange(shapeToUnroll, /*unrollStep=*/1);
 }