Support for multi-dim vectorization #4873

vksnk · 2020-04-24T21:27:02Z

This PR adds support for multi-dimensional vectorization as described in #4616.

There are two main pieces here:

nested ramps/broadcasts:
- Ramp and Broadcast can be nested now, i.e. it is allowed for them to have a vector base, stride and value;
- updates simplifier rules to account for that + extended fuzz_simplify to test it.
- changes to Deinterleave, so it supports extract_lane and others on nested vector expressions;
- smaller changes in some other place.
actual multi-dimensional vectorization:
- this is the place where nested ramps/broadcasts can be introduced;
- I have added a number of tests, which exercise most of the code branches inside of the VectorizeLoops to ensure it works correctly;
- extra pass in the end of the lowering which flattens nested ramps/broadcasts, so existing codegens can work without changes

Notes/questions:

I updated every simplifier rule which involves ramp/broadcast, such that they are only applied when types of the arguments are matching. For some rules, which explicitly use lanes (for example, trying to prove ranges of the ramps), I've added an extra check if argument is scalar to make sure it's only applied in the case when lanes == type.lanes(). I think in order to make it general, we'd need to turn lanes into wildcard.
In some cases lanes argument is passed to the result of the rewriting rule, even though it's not present in the input. I've removed it and it works everywhere except Simplify_LT, I don't quite understand why, so would greatly appreciate any suggestions.
fuzz_simplify is failing for the type UInt(1) when vector width is more than 2 even before the changes, so I've disabled it for now. The error message is "Integer constant 2 will be implicitly coerced to type uint1, which changes its value to (uint1)0.. I am going to look into it next, but figured I ask in case this is known or anyone has any suggestions.
I am not sure how to trigger VectorizeLoops::scalarize(Expr), so put an assert if it's called from nested vectorization.

Otherwise, I think, it should be ready for review. Thanks!

…such that they are only applied to the Expr with same number of lanes.

* fuzz_simpify generates nested ramps/broadcasts * changes is_same_lane_num to is_same_type * some rules are disabled for now

* added is_scalar * fixed ramp lane count * enabled back all rules by guarding some which use 'lanes' with is_scalar

test/correctness/fuzz_simplify.cpp

steven-johnson · 2020-07-08T16:55:18Z

Where do we stand on this -- are we still pending other PRs?

vksnk · 2020-07-08T17:06:57Z

I already merged in within-vector reductions PR, so it should be good as soon as tests failing on build bots are fixed. I was going to take a look at it today.

steven-johnson · 2020-07-08T17:12:22Z

Cool -- no rush, just checking :-)

src/Deinterleave.cpp

abadams · 2020-07-08T17:42:57Z

src/VectorizeLoops.cpp

@@ -191,7 +191,7 @@ struct InterleavedRamp {
    int lanes, repetitions;
 };

-bool is_interleaved_ramp(const Expr &e, const Scope<Expr> &scope, InterleavedRamp *result) {
+bool is_interleaved_ramp(const Expr &e, const std::map<string, Expr> &widened_vars, InterleavedRamp *result) {


Why was this changed from a Scope to a map?

For the case when there are multiple vectorized loops, it's possible that the same let variable may be vectorized differently depending on the loop level it's used and the exact form may not be known when we first see it (for example, if it's used inside of the inner vectorized loop), so map is used to keep track of all possible variant which were used.

I'm confused because a Scope is a map where you can have multiple values per key, and the most recently added value for the key hides the older ones. So a Scope should be strictly more general than a map.

Huh, I guess I misunderstood the way it works; based on name and some usages I saw, I assumed that it's more like named stack and push/pop should go in last-in-first-out order, but in this case push and pop may happen in different functions (and maybe in different orders with other lets), so I thought it won't work correctly. However, if it's map of stacks instead, I guess this shouldn't be an issue. I can change map to scope and track widened names for Let/LetStmt separately, would this be preferable?

The reason to use a Scope instead is if the same name might show up multiple times in a way that a Scope would just handle correctly. Currently vectorization just happens on code with unique var names, but maybe it would be a useful to someone to be able to run it on stmts that don't obey that constraint. But if you're pushing/popping in non-stack order, then maybe that wouldn't work transparently. I need to figure out why you're pushing/popping these in non-stack order though to make sure I understand this code. I added a question about it elsewhere.

Thanks for the explanation, I agree this is a good reason to use Scope! The reason I did it this way is because widened value of Let may change in the case when there is a vectorized loop inside of the Let body and this Let is used in there, so instead of mutating it in advance, it's mutated (sort of) lazily only when needed inside of the visit(const Variable *op) (see get_mutated_var_from_scope() function) to make sure that we've seen all of the vectorized loops. Maybe, as an alternative, it would be better to update all variables in scope every time we encounter a new vectorized loop, replace mutated variables in the scope, mutate loop body, and on return from the mutate(body), create widened variables if needed. That might be quite a bit of a change, but sounds like it's worth doing. What do you think?

I've updated it to use Scope again using approach described above, and I think the code is better now (there are fewer changes for sure), thanks for the the suggestions!

src/VectorizeLoops.cpp

steven-johnson · 2020-07-15T16:15:57Z

Periodic ping to see where this PR stands

abadams · 2020-07-15T17:05:24Z

I believe it's stalled on me continuing my review. Sorry for the delay, I'm currently overloaded with things I have told people I will do.

steven-johnson · 2020-07-15T17:07:34Z

No worries, just making sure there isn't stuff that needed something trivial to unblock :-)

vksnk · 2020-09-17T16:15:25Z

All tests are green again.

src/Lower.cpp

src/VectorizeLoops.cpp

abadams · 2020-09-19T17:12:05Z

lgtm apart from my two minor comments.

vksnk · 2020-09-21T16:51:48Z

Thank you for reviewing!

vksnk added 26 commits March 31, 2020 15:25

Add a test for nested vectorization (doesn't work now)

dcd9fac

First try on nested vectorization

93fcaa9

Ouch, wrong loop

2c52d35

Update some of the simplifier rules which involve ramp and broadcast …

eb95f22

…such that they are only applied to the Expr with same number of lanes.

Move FlattenNestedRamps into separate file

3e727e4

Remove stray var which was actually used in few places

8f13c6e

Support GuardWithIf for the nested vectorization

d4d16f5

Remove the warning test for double vectorization

2125abd

Move replacements construction into VectorSubs

1282e75

Now that bounds_of_lanes work, inner For loop should be allowed

6eef291

Handle Allocate nodes in the multi-dim vectorization

239ac23

Clean-up the test

54727bf

One more test

7f29d38

Refactor to handle vectorized compute_at

c186654

Oops, fixed test

30c9959

Merge branch 'master' into vksnk/vector-ramp

fc708dc

A better way to handle multi-dim alloc

7bc6135

Support arbitrary nested ramps in Deinterleave

12d436a

Use element_of instead of full type for ramp

c2bcec4

Handle lets

ecef86c

More simplifier updates

bf99f33

* fuzz_simpify generates nested ramps/broadcasts * changes is_same_lane_num to is_same_type * some rules are disabled for now

Another vectorize_nested test

c6e4e91

Forgot to add changes to IRMatch

98db7f0

More updates to simplifier

d9cc53d

* added is_scalar * fixed ramp lane count * enabled back all rules by guarding some which use 'lanes' with is_scalar

Clean up

1f8c737

Misc clean-up

73aedc4

dsharletg reviewed Apr 24, 2020

View reviewed changes

test/correctness/fuzz_simplify.cpp Show resolved Hide resolved

vksnk and others added 3 commits April 24, 2020 14:49

Fixed cmake

501fff5

Restore shortcuts for exprs in fuzz_simplify

99a923e

clang-format

3e79e89

abadams reviewed Jul 8, 2020

View reviewed changes

src/Deinterleave.cpp Outdated Show resolved Hide resolved

abadams reviewed Jul 8, 2020

View reviewed changes

src/VectorizeLoops.cpp Outdated Show resolved Hide resolved

vksnk added 2 commits July 8, 2020 12:38

oof, fix access to moved object

1ed694a

Address review comments

60bffe8

abadams reviewed Jul 9, 2020

View reviewed changes

src/VectorizeLoops.cpp Outdated Show resolved Hide resolved

vksnk added 9 commits July 23, 2020 09:26

Merge branch 'master' into vksnk/vector-ramp

6efdabb

Merge branch 'master' into vksnk/vector-ramp

c9377bd

Update Allocate handler to directly use replacement list

b38c34f

Scope is back

0c338cf

Comment

738f8b1

Merge branch 'master' into vksnk/vector-ramp

1587eed

make format

55a722d

Update comment / trigger build bots

600d8b4

Merge branch 'master' into vksnk/vector-ramp

afd6e93

abadams reviewed Sep 19, 2020

View reviewed changes

src/Lower.cpp Outdated Show resolved Hide resolved

abadams reviewed Sep 19, 2020

View reviewed changes

src/VectorizeLoops.cpp Outdated Show resolved Hide resolved

Address review comments

f256e8f

abadams approved these changes Sep 19, 2020

View reviewed changes

vksnk merged commit f5a764f into master Sep 21, 2020

vksnk deleted the vksnk/vector-ramp branch October 8, 2020 02:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support for multi-dim vectorization #4873

Support for multi-dim vectorization #4873

vksnk commented Apr 24, 2020 •

edited

Loading

steven-johnson commented Jul 8, 2020

vksnk commented Jul 8, 2020

steven-johnson commented Jul 8, 2020

abadams Jul 8, 2020

vksnk Jul 8, 2020

abadams Jul 8, 2020 •

edited

Loading

vksnk Jul 9, 2020

abadams Jul 9, 2020

vksnk Jul 10, 2020

vksnk Sep 12, 2020

steven-johnson commented Jul 15, 2020

abadams commented Jul 15, 2020

steven-johnson commented Jul 15, 2020

vksnk commented Sep 17, 2020

abadams commented Sep 19, 2020

vksnk commented Sep 21, 2020

Support for multi-dim vectorization #4873

Support for multi-dim vectorization #4873

Conversation

vksnk commented Apr 24, 2020 • edited Loading

steven-johnson commented Jul 8, 2020

vksnk commented Jul 8, 2020

steven-johnson commented Jul 8, 2020

abadams Jul 8, 2020

Choose a reason for hiding this comment

vksnk Jul 8, 2020

Choose a reason for hiding this comment

abadams Jul 8, 2020 • edited Loading

Choose a reason for hiding this comment

vksnk Jul 9, 2020

Choose a reason for hiding this comment

abadams Jul 9, 2020

Choose a reason for hiding this comment

vksnk Jul 10, 2020

Choose a reason for hiding this comment

vksnk Sep 12, 2020

Choose a reason for hiding this comment

steven-johnson commented Jul 15, 2020

abadams commented Jul 15, 2020

steven-johnson commented Jul 15, 2020

vksnk commented Sep 17, 2020

abadams commented Sep 19, 2020

vksnk commented Sep 21, 2020

vksnk commented Apr 24, 2020 •

edited

Loading

abadams Jul 8, 2020 •

edited

Loading