[Mlir-commits] [mlir] [mli][vector] canonicalize vector.from_elements from ascending extracts (PR #139819)

[Andrzej Warzyński]

================
@@ -2385,9 +2386,105 @@ static LogicalResult rewriteFromElementsAsSplat(FromElementsOp fromElementsOp,
   return success();
 }
 
+/// Rewrite vector.from_elements as vector.shape_cast, if possible.
+///
+/// Example:
+///   %0 = vector.extract %source[0, 0] : i8 from vector<1x2xi8>
+///   %1 = vector.extract %source[0, 1] : i8 from vector<1x2xi8>
+///   %2 = vector.from_elements %0, %1 : vector<2xi8>
+///
+/// becomes
+///   %2 = vector.shape_cast %source : vector<1x2xi8> to vector<2xi8>
+///
+/// The requirements for this to be valid are
+/// i) source and from_elements result have the same number of elements,
+/// ii) all elements are extracted from the same vector (%source),
+/// iii) the elements are extracted in ascending order.
+///
+/// It might be possible to rewrite vector.from_elements as a single
+/// vector.extract if (i) is not satisifed, or in some cases as a
+/// a single vector_extract_strided_slice if (i) and (iii) are not satisfied,
+/// this is left for future consideration.
----------------
banach-space wrote:

I know we already have quite a few TODOs/FIXMEs that basically mean "let’s look at this later." But the phrasing "It might be possible…” feels particularly vague here - I’d suggest omitting it unless we can be more specific.

If we do want to leave a note, maybe something like:
> “Consider extending to use a single vector.extract when (i) does not hold.”

Also, just a general thought: extending this pattern could quickly become quite complex. If we're seeing bad code that would benefit from such a complicated rewrite, it might be worth checking whether the producer of that code could be improved instead.

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