Page MenuHomePhabricator

[mlir][vector] Add scalable vectors support to OuterProductOp

Authored by WanderAway on Nov 25 2022, 7:33 AM.



This will probably be the first in a series of patches that tries to enable code generation for ARM SME (extension of SVE).

Since SME's core operation is the outer product instruction, I figured that it would probably be a good idea to enable the outer product operation to properly accept and generate scalable vectors.

Diff Detail

Event Timeline

WanderAway created this revision.Nov 25 2022, 7:33 AM
Herald added a project: Restricted Project. · View Herald Transcript
WanderAway requested review of this revision.Nov 25 2022, 7:33 AM
Matt added a subscriber: Matt.Nov 27 2022, 10:35 AM
kaitingwang accepted this revision.Dec 1 2022, 9:48 AM

This patch catches the corner case where RHS is non-scalable vector and LHS is scalable vector, throw an error. When RHS is scalable vector, but LHS is non-scalable vector, this case is supported and testcase shows it. This patch looks good to me.

This revision is now accepted and ready to land.Dec 1 2022, 9:48 AM

@jsetoain @nicolasvasilache @aartbik @dcaballe
Some additional feedback would be appreciated, thanks.

dcaballe requested changes to this revision.Dec 2 2022, 11:57 AM

Adding support for SME very well welcome! Thanks!
Just a couple of comments for now.


Vector types in MLIR may have multiple dimensions so there could be multiple scalable dims (see getNumScalableDims). It would be great if you could add a test with multiple scalar dimensions and make sure that that works.


I would expect this to be implemented (if it's not already?) with a trait similar to SameTypeOperands so that we can use it for all the vector operations that do not expect the same operand types but they do expect the same kind of vector. Perhaps something like SameVectorKindOperand?

This revision now requires changes to proceed.Dec 2 2022, 11:57 AM
WanderAway added inline comments.Dec 2 2022, 2:34 PM

Technically it is not necessary for both operands for the outer product op to be scalable or fixed:

<[2]xf64> x <2xf64> -> <[2]x2xf64>
<2xf64> x <[2]xf64> -> <2x[2]xf64>

However (if I understand correctly), the current implementation of scalable vectors only support the latter, with the scalable dimension on the inner dimension of the vector.

This can get kind of ugly with data layout issues etc. so I'm just emitting an error here in case the first case gets supported later on...

WanderAway added inline comments.Dec 5 2022, 6:46 AM

As specified on line 2358 in OuterProductOp::verify(), the input operands for a outer product op are expected to be 1d vectors (otherwise this would turn into sort of a gemm or gemv operation).

WanderAway updated this revision to Diff 480143.Dec 5 2022, 9:29 AM
WanderAway marked 2 inline comments as done.

Updated the test case to check for invalid multi-dimensional input vectors.

jsetoain added inline comments.Dec 6 2022, 2:35 AM

Indeed, only innermost dimensions are allowed to be scalable at the moment. It should not be too traumatic to extend it to allow per-dimension scalability but, also, I'm not sure that's going to be needed. Re.: type traits, we don't have any to indicate both have to be either scalable or fixed length, but requesting matching types should enforce it. vector<4xi32> and vector<[4]xi32> are two different types.

WanderAway planned changes to this revision.Dec 6 2022, 12:13 PM
WanderAway added inline comments.

The problem is that currently the OuterProductOp can take vectors of varying sizes as input
e.g. OuterProduct(<2xf32>, <3xf32>) => <2x3xf32>

Plus the operand can be a scalar as well
e.g. OuterProduct(f32, <3xf32>) => <3xf32>

It turns out that a constraint for these (in addition to the "scalability" check) is surprisingly hard to write in tablegen/ODS, and I think I will update the verify() code in C++ instead.

WanderAway updated this revision to Diff 480662.Dec 6 2022, 3:56 PM

Added verification of scalable operands.

WanderAway marked an inline comment as done.Dec 6 2022, 3:59 PM
kaitingwang resigned from this revision.Dec 8 2022, 8:42 AM
This comment was removed by kaitingwang.
This comment was removed by WanderAway.

@dcaballe do you mind taking another look? Thanks

dcaballe accepted this revision.Jan 12 2023, 9:45 PM

Sorry for the delay! I'm not sure why your ping didn't land on my inbox.
Ok, I understand the complexity. Let's go with this as is for now. Thanks!

This revision is now accepted and ready to land.Jan 12 2023, 9:45 PM
This revision was automatically updated to reflect the committed changes.