[Mlir-commits] [mlir] [mlir][vector] Folder: shape_cast(extract) -> extract (PR #146368)

[Andrzej Warzyński]

================
@@ -1696,59 +1696,68 @@ static bool hasZeroDimVectors(Operation *op) {
          llvm::any_of(op->getResultTypes(), hasZeroDimVectorType);
 }
 
+/// All BroadcastOps and SplatOps, and ShapeCastOps that only prepends 1s, are
+/// considered 'broadcastlike'.
+static bool isBroadcastLike(Operation *op) {
+  if (isa<BroadcastOp, SplatOp>(op))
+    return true;
+
+  auto shapeCast = dyn_cast<ShapeCastOp>(op);
+  if (!shapeCast)
+    return false;
+
+  VectorType srcType = shapeCast.getSourceVectorType();
+  ArrayRef<int64_t> srcShape = srcType.getShape();
+  uint64_t srcRank = srcType.getRank();
+  ArrayRef<int64_t> dstShape = shapeCast.getType().getShape();
+  return dstShape.size() <= srcRank && dstShape.take_back(srcRank) == srcShape;
----------------
banach-space wrote:

Since this is looking for `shape_cast` that prepends 1s, shouldn't the destination rank be larger than the source rank?

Also, since the comment above mentions "prepends 1s", I'd find this more intuitive:
```cpp
llvm::all_of(dstShape.take_front(dstRank - srcRank), [](int64_t dim){return dim == 1;})
```

THis is a nit :)

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