[Mlir-commits] [mlir] [mlir][vector] Support multi-dimensional vectors in VectorFromElementsLowering (PR #151175)

[Diego Caballero]

dcaballe wrote:

> I would much rather us do unrolling for lowering. While flattenting looks nice, it can sometimes generate shuffles (by generating extract_strided_slice/insert_strided_slice).

I don’t think we can make a call on unrolling vs linearization. Unrolling will bloat the code size when unrolling a large dimension whereas linearization will generate fewer ops (best case a single op). Vector shuffles will be generated anyways by LLVM regardless of what we do at MLIR level. The right call is probably project dependent. 

> my suggestion was to add a pattern there to split a rank-D from_elements op into a rank-1 from_elements op preceded by a shape_cast, and then populateVectorShapeCastLoweringPatterns will kick in and generate the inserts/extracts.

The shape cast implementation makes sense to me if we decouple it from the actual lowering to LLVM. It’s basically the vector linearization flavor so that probably should go into VectorLinearize.cpp. I think it’s important that we keep the linearization-like patterns focused on shape cast so that we implement an optimized lowering for linearization patterns in just one place.



> How we ideally want to structure conversion to backends is:
> 
> * Set of patterns to do unrolling from N-D vectors to 1-D vectors
> * Set of patterns to do flattening from N-D vectors to 1-D vectors (in case someone wants to do this, we dont have patterns for this today)
> * Set of patterns to cleanup boundary ops created by unrolling/flattening
> * Set of patterns to convert 1-D vector dialect operations to LLVM dialect
> * Set of patterns to convert 1-D vector (+ spirv restrictions) dialect operations to SPIRV dialect

This makes sense to me, with a twist. IMO, the main problem is that `ConvertVectorToLLVM` currently does far more than just converting to LLVM. It makes algorithmic decisions, such as how to lower a contract op or a transpose, which aren't necessarily driven by LLVM constraints. These choices were made out of early implementations available at the time, which can make the pass challenging to use in production nowadays.

Expanding on @Groverkss ' point, I believe we should decouple:
* **Algorithmic decision:** Move algorithmic decision out of the vector to LLVM conversion. These transformations could happen within a single or multiple configurable passes/transformations.
* **LLVM/SPIR-V legalization constraints:** Reframe ConvertVectorToLLVM/SPIRV as one (or two) dedicated legalization pass(es). These passes should strictly focus on applying vector-to-vector transformations based on LLVM/SPIRV constraints and should be configurable to support both unrolling and linearization.
* **Actual lowering to LLVM:** Leverage the LLVM conversion interface to perform the final, straightforward conversion to LLVM.

Deciding on the actual direction here is something we should prioritize as it would requite quite some work and coordination. This would be a great topic for the Tensor Compiler WG.

**For this PR specifically**, my suggestion is that we:
* Add the shape cast implementation to VectorLinearize.cpp (one PR).
* Add the unrolling implementation, using the direct implementation, similar to what other unrolling patterns do (one PR).

WDYT?


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