[Mlir-commits] [mlir] [mlir][vector] Fix FlattenGather for scalable vectors (PR #96074)

[Benjamin Maxwell]

================
@@ -64,6 +66,12 @@ struct FlattenGather : OpRewritePattern<vector::GatherOp> {
     if (resultTy.getRank() < 2)
       return rewriter.notifyMatchFailure(op, "already flat");
 
+    // Unrolling doesn't take vscale into account. Pattern is disabled for
+    // vectors with leading scalable dim(s).
+    if (resultTy.isScalable() && !isTrailingDimScalable(resultTy))
+      return rewriter.notifyMatchFailure(
+          op, "vector type must be fixed-width or scalable in trailing dim");
----------------
MacDue wrote:

I think this may be overly restrictive. This rewrite (applied once) only unrolls the leading dimension, so we just need to check `if (resultTy.getScalableDims().front())`. 

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