Page MenuHomePhabricator

[mlir][Vector][Affine] Improve affine vectorizer algorithm
ClosedPublic

Authored by dcaballe on Feb 24 2021, 10:47 PM.

Details

Summary

This patch replaces the root-terminal vectorization approach implemented in the
Affine vectorizer with a topological order approach that vectorizes all the
operations within the target loop nest. This effort is aimed at simplifying
the existing vectorization algorithm so that we can introduce support for
'iter_args' (reductions) [1] and more advanced vectorization scenarios in the
future. More detailed information about the vectorization approach is included
in the patch. These are the most important changes introduced by the new
algorithm:

  • Removed tracking of root and terminal ops. Existing vectorization functionality is preserved and extended so that loop nests without root-terminal chains can be vectorized.
  • Vectorizing a loop nest now only requires a single topological traversal.
  • A new vector loop nest is incrementally built along the vectorization process. The original scalar loop is kept intact. No cloning guard is needed to recover the scalar loop if vectorization fails. This approach also simplifies the challenging task of replacing a loop operation amid the vectorization process without invalidating the analysis information that depends on the original loop.
  • Vectorization of specific operations has been implemented as independent, preparing them to be moved to a potential vectorization interface.

Initial support for 'iter_args' will be introduced in follow-up patches. I already
the first scenario [1] working locally.

It's very likely that you also find more improvement opportunities while looking
at the code since this pass is a bit "old". Feel free to bring them up. We can
address those that are strictly related to the new code in this review and follow
up with further changes. I would prefer to keep this patch as focused as possible
since it's already quite large.

Thanks!
Diego

[1] https://llvm.discourse.group/t/rfc-adding-reduction-support-iter-args-to-the-affine-vectorizer/2807.

Diff Detail

Event Timeline

dcaballe created this revision.Feb 24 2021, 10:47 PM
dcaballe requested review of this revision.Feb 24 2021, 10:47 PM
Herald added a project: Restricted Project. · View Herald TranscriptFeb 24 2021, 10:47 PM

Kind ping :)

sgrechanik added inline comments.Mar 2 2021, 3:16 PM
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
702

Word repetition

828

I think there should be a more concise builder for ConstantOp, maybe something like state.builder.create<ConstantOp>(constOp.getLoc(), vecAttr) will work?

dcaballe updated this revision to Diff 327881.Mar 3 2021, 12:13 PM
dcaballe marked an inline comment as done.
  • Addressed Sergei's feedback.
  • Minor fix to 'iter_args' tests.
  • Minor changes to registration APIs.
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
828

Thanks! Yes, this was probably some old code.

bondhugula resigned from this revision.Mar 7 2021, 10:47 PM

Thanks for refactoring and modernizing this!
This LGTM so far, left a few comments for now, I'll finish reviewing tomorrow.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
646

A lot of this implementation predates BlockAndValueMapping.
Can we use this more idiomatic data structure in all places?

679

It should be possible to replace this by a simple BlockAndValueMapping call to bvm.map(replaced->getResults(), replacement->getResults()).

881

bvm.lookupOrNull / bvm.lookupOrDefault are the more modern ways of doing some of this.

aartbik added inline comments.Mar 8 2021, 6:02 PM
mlir/include/mlir/Dialect/Vector/VectorUtils.h
157–160

is this layout picked by lint (or can we have return type on same line and break args earlier)?

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
255

I would call out "for now" with explicit TODO. otherwise such subtle comments are easily overlooked and forgotten later

258

can you be more specific what this TODO covers? it is unclear to me what future cases you have in mind

603

Here and several times below, would you mind breaking the line so that the example jumps out.
So, I mean:

.... text ...

Example:

.. example code...
887

Spell out "That could mean..." -> "Reaching this point could mean ..."

908

nit, this test now only matters for the debug message, since all cases return null henceforth

1069

typo: counterparts

dcaballe updated this revision to Diff 329209.Mar 8 2021, 9:37 PM
dcaballe marked 5 inline comments as done.

Addressed feedback. Thanks!

mlir/include/mlir/Dialect/Vector/VectorUtils.h
157–160

Yes, this is the clang-format layout.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
258

Good point! I'm actually adding support for some loops in the follow-up patch up for review and forgot to update this doc.

646

Sure! I wasn't aware of this class. I'll use it for the Value map. Is there anything similar for Operation?

679

I replaced the mapping in the registerValueVectorReplacementImpl utility. It's important to use the different utilities since each one has different sanity checks that ensure that we don't mess up the mapping.

881

Neat! Thanks!

908

Good catch!

nicolasvasilache accepted this revision.Mar 9 2021, 2:49 AM
nicolasvasilache added inline comments.
mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp
881
if (Value vecRepl = ... ) {
  ...
}
926

Type elementType = loadOp.getMemRefType().getElementType();

alternatively: MemRefType memRefType = loadOp.getMemRefType(); and reuse it everywhere below.

This revision is now accepted and ready to land.Mar 9 2021, 2:49 AM
dcaballe updated this revision to Diff 329468.Mar 9 2021, 2:36 PM
dcaballe marked 2 inline comments as done.

Addressed feedback, thanks!
I'll commit this tomorrow if no more comments.