Page MenuHomePhabricator

[mlir][sparse] add vectorization strategies to sparse compiler
ClosedPublic

Authored by aartbik on Jan 12 2021, 1:44 PM.

Details

Summary

Similar to the parallelization strategies, the vectorization strategies
provide control on what loops should be vectorize. Unlike the parallel
strategies, only innermost loops are considered, but including reductions,
with the control of vectorizing dense loops only or dense and sparse loops.

The vectorized loops are always controlled by a vector mask to avoid
overrunning the iterations, but subsequent vector operation folding removes
redundant masks and replaces the operations with more efficient counterparts.
Similarly, we will rely on subsequent loop optimizations to further optimize
masking, e.g. using an unconditional full vector loop and scalar cleanup loop.

The current strategy already demonstrates a nice interaction between the
sparse compiler and all prior optimizations that went into the vector dialect.

Ongoing discussion at:
https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/10

Diff Detail

Event Timeline

aartbik created this revision.Jan 12 2021, 1:44 PM
aartbik requested review of this revision.Jan 12 2021, 1:44 PM
aartbik updated this revision to Diff 316258.Jan 12 2021, 3:32 PM

fixed lint issue

aartbik edited the summary of this revision. (Show Details)Jan 12 2021, 4:28 PM
penpornk accepted this revision.Jan 12 2021, 10:06 PM

Yay for vectorization! :D

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
875

Nit: Upper-case T?

mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp
333

Nit: Should we make the name more meaningful, e.g., vLen, etc.?

586

Nit: int64_t hi here shares the name with the parameter Value hi. They are in different scopes so it works, but it bugs me a bit. Can we rename?

595

Typo.

600

Why the .. around rhs?

600

Nit: At first glance, I intepreted / as a division and that confused me. Maybe change it to "or"?

614

Missing a closing bracket. Also, same comment on /.

822

Nit: Forall doesn't really imply parallelization. How about isParallelFor?

834

A loop can be both parallelized and vectorized. Is the !isVector to simplify things for now? (I'm mostly thinking of a non-reduction loop -- I remember that we'll keep parallelizing reduction loops for later.) Should we mention this in the comment?

mlir/test/Dialect/Linalg/sparse_vector.mlir
24

Thank you for manually naming the captures! \dTvTb/

95–96

Just for my information, why do we have 2 and 3 after load %{{.*}}?

169

Do we need a %[[v0:.*]] = constant 0 : f32 so the %{{.*}} in iter_args can be %[[v0]]?

This revision is now accepted and ready to land.Jan 12 2021, 10:06 PM

On one hand, it's great to see minimal additional patterns translate to new capabilities, very nice!

OTOH, this has grown far beyond a few reasonable abstraction gap jumps, even for a prototype.
As time passes and the abstraction gap grows between the input and the output, the more difficult it will be to break down and refactor into composable pieces.

I am curious what the expected integration will look like and what other features you want to prototype before this integration can start?
For instance, it seems unlikely to me that the tests, as they are written, will be gradually ported to new abstractions.
One outcome that I would find unfortunate is whether people actually start to rely on all the functionality needing to be kept green at all times..

Thoughts?

ftynse added a subscriber: ftynse.Jan 13 2021, 3:59 AM
ftynse added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp
333

MLIR has a distinct preference for long verbose names, even vLen is frowned upon. Something currentVectorLength or curVecLength in the worst case. (We also tend to talk about vector size rather than length).

581–587

This can use matchPattern(value, m_Constant(attr))

On one hand, it's great to see minimal additional patterns translate to new capabilities, very nice!

OTOH, this has grown far beyond a few reasonable abstraction gap jumps, even for a prototype.
As time passes and the abstraction gap grows between the input and the output, the more difficult it will be to break down and refactor into composable pieces.

I am curious what the expected integration will look like and what other features you want to prototype before this integration can start?
For instance, it seems unlikely to me that the tests, as they are written, will be gradually ported to new abstractions.
One outcome that I would find unfortunate is whether people actually start to rely on all the functionality needing to be kept green at all times..

Thoughts?

I would like to understand how this connects to the rest of the code generation...

aartbik marked 12 inline comments as done.Jan 13 2021, 10:29 AM

Thanks for the comments. Nicolas, you will be happy to hear that vectorization was the last "computational" part I wanted to add to this first sparse compiler version, so that from here on this code will not grow much more (we could add support for sparse left-hand sides, but I don't even think we need that for the proof-of-concept phase). Alex and Nicolas, the next steps are auto-gen for code that reads the matrices and tensors from file into the proper bufferized format so that others can start evaluating the code without having to hack in I/O methods. This work will mainly be done in library support. After that we can start connecting the pieces with e.g. a Python front-end, I am looking forward to working with Mehdi on that! In parallel, we can start working on adding proper higher abstractions, introduce them and refactor the code.

Btw, although I hear you concerning avoiding large gaps in abstractions, I don't necessarily agree this code has jumped *too* far ahead yet, especially compared with some other parts of LingAlg. But, I share your love for breaking compiler passes into the smallest possible abstractions possible, so I look forward working with you on that next!

mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp
581–587

Ah, I was looking for something like that. Great!

834

Ah, for now *same* loop parallelization/vectorization is not supported (we could, I agree). So this logic makes sure that is not changed until proper support is added.

mlir/test/Dialect/Linalg/sparse_vector.mlir
95–96

The "l" was for load, and then numbering up. But I renamed without numbering, but referring back to the tensor names, hopefully more clear.

169

In this case we actually load the "tensor" scalar as initial value, and I did not want to commit to that pattern to rigidity yet.

aartbik updated this revision to Diff 316449.Jan 13 2021, 10:55 AM
aartbik marked 3 inline comments as done.

addressed detailed comments

aartbik added inline comments.Jan 13 2021, 10:57 AM
mlir/test/Dialect/Linalg/sparse_vector.mlir
95–96

Oh, I see now you did not refer to the label name. I will remove this too.

aartbik updated this revision to Diff 316452.Jan 13 2021, 10:58 AM

removed stray 2 and 3

Harbormaster completed remote builds in B85051: Diff 316452.