[mlir][RFC] Add scalable dimensions to VectorType

[jsetoain]

Summary

With VectorType supporting scalable dimensions, we don't need many of
the operations currently present in ArmSVE, like mask generation and
basic arithmetic instructions. Therefore, this patch also gets
rid of those.

Having built-in scalable vector support also simplifies the lowering of
scalable vector dialects down to LLVMIR.

Scalable dimensions are indicated with the scalable dimensions
between square brackets:

        vector<[4]xf32>

Is a scalable vector of 4 single precission floating point elements.

More generally, a VectorType can have a set of fixed-length dimensions
followed by a set of scalable dimensions:

        vector<2x[4x4]xf32>

Is a vector with 2 scalable 4x4 vectors of single precission floating
point elements.

The scale of the scalable dimensions can be obtained with the Vector
operation:

        %vs = vector.vscale

This change is being discussed in the discourse RFC:

https://llvm.discourse.group/t/rfc-add-built-in-support-for-scalable-vector-types/4484

Differential Revision: https://reviews.llvm.org/D111819

[mehdi_amini]

This direction makes a lot of sense to me, if we want to avoid code dup between the upcoming vector specific dialects (SVE and RISC-V at the moment).
Since this touches "core", however, I hope others chime in too.

My main concern is that the discussion on Discourse does not seem to have a conclusion right now.

[jsetoain]

In D111819#3081297, @mehdi_amini wrote:

This direction makes a lot of sense to me, if we want to avoid code dup between the upcoming vector specific dialects (SVE and RISC-V at the moment).
Since this touches "core", however, I hope others chime in too.

My main concern is that the discussion on Discourse does not seem to have a conclusion right now.

My apologies. We did have a conversation this week but I didn't update the RFC (I have now). Although some details are still in the air, the need for a built-in type is not in question. This is just the first, most basic implementation possible.

[jsetoain]

Address reviewers' comments

[jsetoain]

I've addressed all the comments.

[jsetoain]

Fixed formatting

[jsetoain]

Rebase on main

[jsetoain]

Rebase on main

[jsetoain]

Rebase on main

[rriddle]

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

487

The get here is a bit weird, why not something like isScalableVectorType?

mlir/include/mlir/Dialect/Vector/VectorOps.td

2404–2405

You could also drop the trailing type if you want, it can be inferred. (i.e. = "attr-dict")

mlir/include/mlir/IR/BuiltinTypes.td

941–942

Is this wrapped at 80 characters?

mlir/lib/Dialect/ArmSVE/IR/ArmSVEDialect.cpp

51

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

161

Why the extra spaces?

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

778–784

mlir/lib/IR/BuiltinTypes.cpp

296–299

319–320

Same here.

342–343

and here.

mlir/lib/Parser/TypeParser.cpp

456–458

[jsetoain]

Address review comments

[mehdi_amini]

What's the next step on this? Seems like the RFC discussion got to a resolution right?

@nicolasvasilache @ftynse : can you chime in on the change in mlir/include/mlir/IR/BuiltinTypes.td ?

[nicolasvasilache]

My apologies for the long delay, my biggest problem atm is I do not have a good mental model of how RVV and Arm SVE operate in detail.
I had started to read specs but it invariably gets pushed back on the stack as it is not high on my priority list ..

From a pure cleanup perspective, I generally like it.

From a composability perspective, I think I would prefer to have it spelled out as vector<<4>xf32> or vector<4*xf32> or vector<(4s)xf32>.
The rationale is that I think we still want to have n-D scalable vector types in MLIR to allow expressing a statically known number of 1-D scalable vectors that serves as an "unroll-and-jammed vector pack" vector<8x4*xf32>.
This would be more consistent with the design of the rest of the vector dialect.

One thing that is higher priority to me personally these days is that we are also exploring using the vector dialect as a programming model for GPUs.
In this context, vector<4x8x16*x32*xf32> would also make sense for us.

Bottom line, if we avoided anchoring on the current LLVM / HW implementation that only support 1-D scalable vectors and we made it future-proof in that direction, I am fine with proceeding.

[jsetoain]

Generalize the concept of scalable dimensions to support use cases unrelated to scalable vectors

[jsetoain]

mlir/include/mlir/Dialect/Vector/VectorOps.td

2399

I've seen this, I'll take care of it together with any other necessary fix.

2404–2405

It feels a bit "naked", but it might be because I'm used to see it with the return type attached. We can give it a go and see what people think, if people don't care, going "concise" is my preferred option. Is there a "good practices" manual for dialect syntax? I can't find one.

mlir/include/mlir/IR/BuiltinTypes.td

907

And this...

941–942

Not sure what happened there. Good catch, thanks!

mlir/lib/Parser/TypeParser.cpp

456–458

Arg! That was embarrassing... Sorry about that!

[nicolasvasilache]

Mostly LGTM, I added 3 areas of improvements.
Once these are addressed I'll happily accept.
Thanks for your hard work and patience!

mlir/include/mlir/Dialect/Vector/VectorOps.td

2390

Should we call this vector.vscale ?

2394

I would emphasize that this is for 1-D scalable vectors and that there is currently no way to extract the scale
for a >1-D scalable vectors.
This instruction may be extended in the future to take a position but I am unclear whether this is what we want atm.

I think the global vs local property of vscale should also be discussed here.

I'd maybe even go as far as spelling it vector.scale.global in the future?

Edit: as I read deeper through the PR, I am now unclear whether vector<[2x8]xf32> is the same as vector<[2]x[8]xf32> ?

I think vector<[2x8]xf32> would make sense for SVE in MLIR (and would then get flattened to 1-D going through LLVM).
In the future we may also want vector<[2]x[8]xf32> for GPUs but this is not the same representation?
Is this what you have in mind ?

In any case, please propose a few wording changes to integrate the relevant parts of my comments and disregard/add a TODO for the others :)

mlir/include/mlir/IR/BuiltinTypes.h

323

I fear this will prove annoying to use in practice ..

Could we go with unsigned numScalableDims?

Then you can just use APIS such as ArrayRef's shape.take_back(numScalableDims); and friends.

342

this would get nicer with numScalableDims.

mlir/include/mlir/IR/BuiltinTypes.td

928

Now that I read this, I am unclear whether vector<[2x8]xf32> is the same as vector<[2]x[8]xf32>, I would think not and the latter form could be a future extension (if so, add a TODO)?
This really depends on whether you think you can make use of vector<[2x8]xf32> in MLIR instead of having to represent as vector<[16]xf32>; I claim you would have a bunch of nice use cases for this (coupled with the shape_cast op once properly extended).

mlir/test/Dialect/ArmSVE/legalize-for-llvm.mlir

1

I would expect to see a test file (somwhere in the builtin stuff) where you have both:

• negative tests for various failure modes of misuses of scalable vectors (with appropriate error messages)
• positive tests with multi-dim multi-scale vector (atm everything I see is 0-dim 1-scale only).

In a followup PR, I'd love to see a 1-dim, 2-scale version of the neon 2d dot (or something equivalent) and see it lower to unrolled LLVM.

[nicolasvasilache]

Also signal boosting for @ThomasRaoux @dcaballe @springerm ; no need to review but be aware that this is coming.

[rriddle]

mlir/include/mlir/IR/BuiltinTypes.td

942

Can we use Optional here instead? -1 is a bit magic.

960–961

Why not just isScalable? The naming here is a bit weird.

mlir/lib/IR/AsmPrinter.cpp

1966

Please cache the end iterator to avoid recomputing it every iteration.

mlir/lib/Parser/TypeParser.cpp

522

Looks like this is missing test coverage.

[jsetoain]

mlir/include/mlir/IR/BuiltinTypes.td

928

It is, indeed, very much not the same. I find it useful to think about something like [2x8] as a series of 2x8 blocks, one after another. Therefore, even though they would have the same memory requirements, [2x8], [8x2], [4x4], and [16] can represent different data arrangements when you're loading your data from memory. From that point of view, [2]x[8] can't be the same as [2x8] even if the scale for both dimensions is the same. In fact, I don't think something like [2]x[8] makes sense in the context of scalable vectors. For GPU thread blocks, the situation is different. I'm not involved with that work so I can't come up with anything on the spot, but I intuit it could have potentially useful cases. As this work progresses, I suspect we will need to come back to it.

[jsetoain]

Addressed latest round of reviews

[jsetoain]

I've also moved a bunch of tests around. Scalable vector tests that are not SVE-specific have been moved either to Arithmetic or Vector, depending on their nature.

mlir/include/mlir/Dialect/Vector/VectorOps.td

2394

That's not exactly right. You can have a 2D scalable vector, and vscale represents its multiplicity, but you can't have a 2D scalable vector with two different scales (which we might want to have for GPUs). As it is, we can't represent those yet, so I don't think we need to clarify that in the description. I've added the multi-dimensional multi-scale vector and local/global scale to a TODO for future reference.

mlir/include/mlir/IR/BuiltinTypes.td

942

Not sure how to use Optional for Types, this is the only way I found to provide a default value in a type builder. In any case, I've changed it to "numScalableDims" as suggested by Nicolas. It makes code a bit less awkward and conveniently replaces a arguably ugly "first dimension = -1" to a more semantically sensible "number of dimensions = 0". If you still find this unacceptable, I can look into adding an "Optional" equivalent for types.

mlir/test/Dialect/ArmSVE/legalize-for-llvm.mlir

1

RE follow-up PR, that was in my low priority TODO list, I'll move it to the main TODO, it should be a quick and easy change.

[rriddle]

LG from my point-of-view, but also get an LGTM from Nicolas as well.

mlir/include/mlir/Dialect/Vector/VectorOps.td

2390

Can you move this into the documentation of the op? This seems useful to expose in the user facing docs.

2404–2405

I don't think we have a "good practices" manual, though that sounds useful.

2415

If a verifier isn't necessary, you can just ignore it.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

31

nit: Prefer pre-increment unless you need post increment behavior.

41

Same here.

mlir/lib/IR/AsmPrinter.cpp

1966

Unresolved.

[jsetoain]

Address reviewer comments

[jsetoain]

Addressed comments.

[nicolasvasilache]

Thanks for getting these through @jsetoain !

[jsetoain]

Rebase on main