This is an archive of the discontinued LLVM Phabricator instance.

[mlir][VectorType] Allow arbitrary dimensions to be scalable
ClosedPublic

Authored by awarzynski on Jun 20 2023, 12:58 PM.

Details

Summary

At the moment, only the trailing dimensions in the vector type can be
scalable, i.e. this is supported:

vector<2x[4]xf32>

and this is not allowed:

vector<[2]x4xf32>

This patch extends the vector type so that arbitrary dimensions can be
scalable. To this end, an array of bool values is added to every vector
type to denote whether the corresponding dimensions are scalable or not.
For example, for this vector:

vector<[2]x[3]x4xf32>

the following array would be created:

{true, true, false}.

Additionally, the current syntax:

vector<[2x3]x4xf32>

is replaced with:

vector<[2]x[3]x4xf32>

This is primarily to simplify parsing (this way, the parser can easily
process one dimension at a time rather than e.g. tracking whether
"scalable block" has been entered/left).

NOTE: The isScalableDim parameter of VectorType (introduced in this patch) makes numScalableDims redundant. For the time being, numScalableDims is preserved to facilitate the transition between the two parameters. numScalableDims will be removed in one of the subsequent patches.

This change is a part of a larger effort to enable scalable
vectorisation in Linalg. See this RFC for more context:

Diff Detail

Event Timeline

awarzynski created this revision.Jun 20 2023, 12:58 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: ftynse. · View Herald Transcript
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: kuhar. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
awarzynski requested review of this revision.Jun 20 2023, 12:58 PM
jsetoain added inline comments.Jun 20 2023, 2:57 PM
mlir/lib/AsmParser/Parser.h
214–215

Is there a reason to keep this?

jsetoain added inline comments.Jun 20 2023, 2:58 PM
mlir/lib/AsmParser/Parser.h
214–215

Never mind. The blue colour in the "NOTE" tripped my mental parser 🙂

This is great! Thanks for taking care of this! A few comments

mlir/include/mlir/Bytecode/BytecodeImplementation.h
257

blob -> bool?

mlir/include/mlir/IR/BuiltinTypes.td
1084

Also replying to:

NOTE: The isScalableDim parameter of VectorType (introduced in this patch) makes numScalableDims redundant. For the time being, numScalableDims is preserved to facilitate the transition between the two parameters. numScalableDims will be removed in one of the subsequent patches.

Is this transition really that painful? The new representation is a superset of the previous one it should be a simple API change.
I'm worried about the complexity of having both representations in place and potential errors derived from that.

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
995–1003

I would make sure this is explicitly initialized to false

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
2799–2807

Not sure I follow the use case here but I think we shouldn't have any specific constraint at this level of abstraction regarding the number of scalable dimensions allowed. We should fail when we try to legalize this, if the scalable dimensions haven't been removed, unrolled (that's interesting, how do we unrolled a scalable vector dims?) or replaced in some way and there is no target representation for that (e.g., llvm intrinsics, SVE/SME ops, etc.)

3537

Ha, this is the part that required the bitmap representation (see comment in 3541).

I think we just have to apply the same permutation as to the shape so... we can do either...

invPermMap.compose(vecType.getScalableDims())  // not sure if we can pass a bool SmallVector.

or

applyPermutationMap(invPerMap, vecType.getScalableDims())
mlir/lib/IR/AsmPrinter.cpp
2463

commented code

Addressing Diego's comments, thanks for reviewing!

Thank you for reviewing!

mlir/include/mlir/IR/BuiltinTypes.td
1084

Is this transition really that painful?

Keeping both really helped with debugging and sanity-checking. Also, this is a rather intrusive change and I just like the idea of doing things like this in smaller, incremental steps.

Separately, I am worried that once this lands we'll see some "untested" code breaking (because, for example, I am making the verifiers a bit stricter). And smaller patches are easier to debug :)

The new representation is a superset of the previous one it should be a simple API change.

Yes, but the previous representation imposes some rather strong assumptions.

I'm worried about the complexity of having both representations in place and potential errors derived from that.

Fear not, I intend to remove it very soon :) I will try that now and see how complex that would be.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
2799–2807

Tl;Dr Please check the updated version :)

Not sure I follow the use case here

It's triggered by this test: vector-scalable-outerproduct.mlir. It leads to numScalableDims = 3.

I could just ignore that, but my changes to VectorType::verify mean that VectorType::get below (L2814) fails:

if (numScale != numScalableDims)
  return emitError() << "number of scalable dims must match, explicit: "
                     << numScalableDims << ", and bools:" << numScale;

This is why I'd like to keep numScalableDims for now - it helps capturing these cases. And I do have a "nice" fix for this specific example :)

I think we shouldn't have any specific constraint at this level of abstraction regarding the number of scalable dimensions allowed. We should fail when we try to legalize this, if the scalable dimensions haven't been removed, unrolled (that's interesting, how do we unrolled a scalable vector dims?) or replaced in some way and there is no target representation for that (e.g., llvm intrinsics, SVE/SME ops, etc.)

Agreed. In this particular case, we were setting numScalableDims to 3 for this result vector shape: vector<[4]x[4]xf32. That's incorrect and VectorType::verify should complain. At the moment, that would be ignored.

3537

Thanks! I'm inclined to preserve the assert from L3542 for now. WDYT?

awarzynski added inline comments.Jun 21 2023, 5:37 AM
mlir/include/mlir/IR/BuiltinTypes.td
1084

Uploaded here: https://reviews.llvm.org/D153412. Happy to merge with this patch if that's preferred. Both approaches (1 vs 2 patches) have their merits.

Update comments/documentation

kuhar added a comment.Jun 21 2023, 7:59 AM

Just some nits

mlir/include/mlir/IR/BuiltinTypes.td
1077

its

1080–1081

nit: the { should be go to the previous line

1082

nit: you can insert into the existing vector with isScalabeVec.resize(shape.size(), false) or isScalableVec.append(shape.size(), false)

mlir/lib/Bytecode/Reader/BytecodeReader.cpp
992–995

Would return reder.parseByte(result) work here?

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

nit: prefer C++ casts (static_cast)

mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp
60–62

nit: if you need an array ref, you can construct one from the single element (i.e., without any container)

Matt added a subscriber: Matt.Jun 21 2023, 12:41 PM

Addressing comments from Jakub

Thank you @kuhar, your pointers are much appreciated - I always learn a lot from them :)

dcaballe accepted this revision.Jun 22 2023, 5:11 PM

LGTM % pending comments. Thanks!

mlir/include/mlir/IR/BuiltinTypes.td
1084

Fair enough. We can fo with two patches then if both are landing soon

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
3537

SG but then we shouldn't need the permutation as the map would always be an identity?

This revision is now accepted and ready to land.Jun 22 2023, 5:11 PM
awarzynski updated this revision to Diff 533899.EditedJun 23 2023, 1:52 AM

Rename isScalableDim as scalableDims. This was suggested in https://reviews.llvm.org/D153412

Also updated "mlir/include/mlir/IR/BuiltinTypes.h" with some fairly straightforward changes that I forgot to include earlier.

awarzynski added inline comments.Jun 23 2023, 2:58 AM
mlir/lib/Bytecode/Reader/BytecodeReader.cpp
992–995

It would, parseByte also returns LogicalResult. Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
3537

Good point. I took another look and decided to remove the assert. It should be safe, though it's tricky to trigger this code path just now. IIUC, first I need to update the transform dialect to allow something like this vector_sizes [2, [4], 8].

kuhar accepted this revision.Jun 24 2023, 9:07 AM

LGTM

c-rhodes added inline comments.Jun 27 2023, 2:36 AM
mlir/include/mlir/IR/BuiltinTypes.h
314 ↗(On Diff #533899)

nit: this variable name is a bit confusing at first I thought new meant this is temporary until numScalableDims is removed, but then realised it's copied from setShape? The latter is doing an update so it is new, but that's not the case during construction. How about scalableDims but use this to refer to the class member?

317–322 ↗(On Diff #533899)

if/else like below?

328–332 ↗(On Diff #533899)

nit: reverse condition

mlir/include/mlir/IR/BuiltinTypes.td
1028

not sure this is right, it means dim must always be between square brackets?

1062
1069

nit: this can be inferred from the type

mlir/lib/AsmParser/Parser.h
215
mlir/lib/AsmParser/TypeParser.cpp
474

As above, not sure this is correct.

484–486
491–493

nit: drop braces

Thank you both for reviewing. I believe that this addresses all the outstanding comments. Do let me know if I missed anything. Otherwise, I will merge this shortly.

mlir/include/mlir/IR/BuiltinTypes.h
314 ↗(On Diff #533899)

Fair point, let me use this->scalableDims.

mlir/include/mlir/IR/BuiltinTypes.td
1028

Good catch! I will update in the next revision. Let me know whether it makes sense.

1069

Yeah, it was meant to be some clever comment :)

Address comments from Cullen

Thanks for reviewing @c-rhodes, lots of great suggestions!

c-rhodes accepted this revision.Jun 27 2023, 9:52 AM

I've left a few more minor nits but otherwise LGTM, cheers.

mlir/include/mlir/IR/BuiltinTypes.h
317–321 ↗(On Diff #535013)

nit: drop braces

327–331 ↗(On Diff #535013)

nit: drop braces

mlir/include/mlir/IR/BuiltinTypes.td
1079–1083

can this not be written as?

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
998–1000

nit: could be written as

Thanks Cullen, I will include your suggestions in the final version that I will merge shortly.

mlir/include/mlir/IR/BuiltinTypes.td
1079–1083

IIUC, you'd be creating an ArrayRef from a temporary, which is disabled.

This revision was landed with ongoing or failed builds.Jun 27 2023, 11:27 AM
This revision was automatically updated to reflect the committed changes.