[Mlir-commits] [mlir] [MLIR][Vector] Added ToElementsOp::fold for broadcast->to_elements pattern rewrite. (PR #160318)

[Kunwar Grover]

================
@@ -2395,9 +2396,41 @@ foldToElementsFromElements(ToElementsOp toElementsOp,
   return success();
 }
 
+/// Folds vector.to_elements(vector.broadcast(%x)) for the scalar case only.
+///
+/// Cases handled:
+///  - %x is a scalar: replicate the scalar across all results.
+///
+/// The vector source case is handled by a canonicalization pattern.
+static LogicalResult
+foldToElementsOfBroadcast(ToElementsOp toElementsOp,
+                          SmallVectorImpl<OpFoldResult> &results) {
+  auto bcastOp = toElementsOp.getSource().getDefiningOp<BroadcastOp>();
+  if (!bcastOp)
+    return failure();
+
+  auto resultVecType = cast<VectorType>(toElementsOp.getSource().getType());
+  // Bail on scalable vectors, since the element count and per-dimension extents
+  // must be known at compile time.
+  if (resultVecType.getNumScalableDims() != 0)
+    return failure();
----------------
Groverkss wrote:

Does the toElements verifier actually allow scalable vectors? If not, we shouldn't check it here.

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