[Mlir-commits] [mlir] [MLIR][Vector] Relax shape_cast unrolling to per-reassociation-group contiguity (PR #205684)

[Jianhui Li]

================
@@ -1335,20 +1390,44 @@ struct UnrollShapeCastPattern : public OpRewritePattern<vector::ShapeCastOp> {
     ArrayRef<int64_t> sourceShape = sourceType.getShape();
     ArrayRef<int64_t> resultShape = resultType.getShape();
 
-    if (!isContiguous(*targetShape, resultShape))
+    // The cast factors into reassociation groups; the target tile only needs to
+    // be contiguous within each group, not in the whole result vector.
+    std::optional<SmallVector<ShapeCastReassociationGroup>> groups =
+        computeShapeCastGroups(sourceShape, resultShape);
+    if (!groups)
       return rewriter.notifyMatchFailure(
-          shapeCastOp, "Only supports cases where target shape is "
-                       "contiguous in result vector shape");
-
-    int64_t targetElements = ShapedType::getNumElements(*targetShape);
-
-    // Calculate the shape to extract from source.
-    std::optional<SmallVector<int64_t>> extractShape =
-        calculateSourceExtractShape(sourceShape, targetElements);
-    if (!extractShape)
-      return rewriter.notifyMatchFailure(
-          shapeCastOp,
-          "cannot extract target number of elements contiguously from source");
+          shapeCastOp, "cannot align source and result reassociation groups");
+
+    // The tile is right-aligned against the result; left-pad with 1s so it can
+    // be indexed per group.
+    SmallVector<int64_t> paddedTarget(resultShape.size(), 1);
+    llvm::copy(*targetShape,
+               paddedTarget.end() - static_cast<int64_t>(targetShape->size()));
+
+    // Validate per-group contiguity and build the source extract shape.
+    SmallVector<int64_t> extractShapeStorage;
+    for (const ShapeCastReassociationGroup &g : *groups) {
+      ArrayRef<int64_t> resSub =
+          resultShape.slice(g.resBegin, g.resEnd - g.resBegin);
+      ArrayRef<int64_t> tgtSub = ArrayRef<int64_t>(paddedTarget)
+                                     .slice(g.resBegin, g.resEnd - g.resBegin);
+      if (!isContiguous(tgtSub, resSub))
+        return rewriter.notifyMatchFailure(
+            shapeCastOp, "target shape is not contiguous within a "
+                         "reassociation group of the result vector shape");
+
+      ArrayRef<int64_t> srcSub =
+          sourceShape.slice(g.srcBegin, g.srcEnd - g.srcBegin);
+      int64_t groupTargetElements = ShapedType::getNumElements(tgtSub);
+      std::optional<SmallVector<int64_t>> groupExtract =
+          calculateSourceExtractShape(srcSub, groupTargetElements);
+      if (!groupExtract)
+        return rewriter.notifyMatchFailure(
+            shapeCastOp, "cannot extract the target number of elements "
+                         "contiguously from a source reassociation group");
----------------
Jianhui-Li wrote:

added

https://github.com/llvm/llvm-project/pull/205684