A new index hoisting and CSE approach based on subexpression equivalence and dependency by zasdfgbnm · Pull Request #2235 · csarofeen/pytorch

zasdfgbnm · 2022-12-02T07:37:47Z

This is a complete rewrite of the index hoisting. I no longer use IterDomain mapping to do the analysis, instead, the analysis is just based on the indexing expressions, and their subexpression equivalence and dependency. Doing so, the code is shorter, cleaner, and it can find more hoisting opportunities. I don't have benchmark yet. Will run them this weekend.

This reverts commit 52c7ded.

zasdfgbnm · 2022-12-02T14:39:53Z

 }

-//! A utility class to check if an expression of a particular type exists
-class ExprFinder : kir::ConstIrVisitor {


moved to kernel_ir.cpp

zasdfgbnm · 2022-12-02T14:40:20Z

-  //! ParallelType::Group, but it isn't sufficient as the loop may be
-  //! for an initialization expression, for which the loop shold not
-  //! be grouped. Make sure a GroupedGridReduction is found.
-  bool isGroupedLoop(const kir::ForLoop* loop) {


Renamed as ForLoop::isGroup

zasdfgbnm · 2022-12-02T14:41:27Z

    auto root_ind_i =
        is_overriden ? override_it->second : index_map.at(root_dom[i]);

-    // index hoist must be done before the adjustments for halo


We no longer have this limitation anymore

zasdfgbnm · 2022-12-02T14:42:10Z

-// Version of hoisting without using reference tensor,
-//  should eventually deprecate the other one once reference
-//  tensor is completely deprecated.
-Val* hoistConsumerIndex(


Consumer, producer, and predicate hoisting now use a single interface.

zasdfgbnm · 2022-12-02T14:43:35Z

+      strided_inds[i] = GpuLower::current()->commonIndexMap().hoistIndex(
+          strided_inds[i], loops);


I think we should take the sum of strided_inds first, and hoist the resulting val. Done in #2234.

zasdfgbnm · 2022-12-02T14:45:13Z

+    start_index = GpuLower::current()->commonIndexMap().hoistIndex(
+        start_magic_zero_info.index, loops);
+    stop_index = GpuLower::current()->commonIndexMap().hoistIndex(
+        stop_magic_zero_info.index, loops);


Instead of hoisting start and stop index, we should be able to just hoist the entire predicate as a boolean expression. I will not do it in this PR, will investigate later.

zasdfgbnm · 2022-12-02T23:29:12Z

When can we lift out of two exactly mapped loops? Wouldn't that mean the index of the loop is involved so we can't lift anything on that index out of the loop that defines it right?

Generally, no we can't, but if those loops are not materialized, it's still possible to reuse.

https://github.com/csarofeen/pytorch/blob/devel/torch/csrc/jit/codegen/cuda/lower_index_hoist.cpp#L166-L169

For trivial loop, I am not considering its loop variable as a dependency, so it can be safely lifted out of it.

zasdfgbnm · 2022-12-02T23:37:19Z

-  if ((i165 < (T0.size[0] * (T0.size[1] * (T0.size[2] * T0.size[3]))))) {
+  int64_t i143;
+  i143 = (((nvfuser_index_t)blockIdx.x) * 128) + ((nvfuser_index_t)threadIdx.x);
+  int64_t i94;


i94 is used three times, so it is hoisted.

zasdfgbnm · 2022-12-02T23:52:10Z

+    auto create_fn = def->newObjectFunc();
+    create_fn(index->container(), inputs, {index}, def->attributes());


I am so happy that now I can create an arbitrary expression without knowing what type it is.

This is really cool!

Very cool. Love it.

naoyam · 2022-12-03T00:18:28Z

+    auto create_fn = def->newObjectFunc();
+    create_fn(index->container(), inputs, {index}, def->attributes());


This is really cool!

naoyam · 2022-12-03T00:19:10Z

+    if (current_level > 0) {
+      loop = loops[current_level - 1];
+    }
+    auto it = common_index_map_.find(loop);


A bit lost here. loop can be nullptr. Is it safe to do lookup with nullptr?

Why does the loop start with current_level=0? Doesn't level mean the loop position where all the loop dependency is resolved for this index? So, could there be an equivalent index in the loops from 0th to level-1-th positions? Isn't it just sufficient to check loops[level]?

I think nullptr lookup is fine. It is just an integer zero, it will not access the object.
I need something to refer to the top level exprs, which has no corresponding loop. I can either use 0 to refer to top level exprs and use 1 to refer to loops[0], or use -1 to refer to top level exprs and use 0 to refer to loops[0]

OK, I still don't follow why you need to loops that are outer than the level loop. Is it because some subexpression of index may be hoisted at those outer loops?

OK, I still don't follow why you need to loops that are outer than the level loop. Is it because some subexpression of index may be hoisted at those outer loops?

Oh, I think you are right. We don't need it. I originally wrote this thinking that some subexpression of index might be hoisted at outer loops. I thought that I would find subexpressions that are both in index and common_index_map_. But seems that I eventually didn't do that. All I am doing in this function is to find something sameAs index, I am not further breaking index down into pieces(the breakdown of index is done in hoistIndex). So I think checking loops[pos] is sufficient.

I updated this PR, did some cleanup with this new understanding. Now pos == -1 means top level exprs, pos == 0 means loops[0].

zasdfgbnm · 2022-12-03T02:43:43Z

Checked matmul perf:

I think it is within normal variation.

The diff in generate code https://www.diffchecker.com/ekzrijW0

Will also check batchnorm.

zasdfgbnm · 2022-12-03T06:58:33Z

This is the perf for batchnorm

csarofeen

This looks good to me, however I'm a bit fuzzy on how CommonIndexMap is produced so I think @naoyam should be the one to approve.

csarofeen · 2022-12-05T13:39:48Z

+    auto create_fn = def->newObjectFunc();
+    create_fn(index->container(), inputs, {index}, def->attributes());


Very cool. Love it.

naoyam · 2022-12-06T02:22:58Z

This looks good to me, however I'm a bit fuzzy on how CommonIndexMap is produced so I think @naoyam should be the one to approve.

Seems like there's a lot of changes since I looked at this PR last week. Revisiting now.

naoyam · 2022-12-06T04:08:16Z

+      int64_t position, // if `index` is given (i.e., is_give == true), then
+                        // this is the position of the outer-most loop nest that
+                        // contains all the dependencies of `index`. if `index`
+                        // is a subexpression of the given index, then this is


A bit confused here: index is a subexpression of the given index. Which index are you referring to by "the given index"?

"the given index" refers to the index given to the public method hoistIndex. I have renamed this function to hoistIndexImpl and updated the comment here to be more clear.

naoyam · 2022-12-06T04:09:39Z

+  //! try insert ((i1*1 + i2*2) + i3*3) to common_index_map_[FOR i3],
+  //! try insert ((i1*1 + i2*2) + i3*3) + i4*4 to common_index_map_[FOR i4],
+  //! Before insertion, this function recursively uses
+  //! eliminateCommonSubexpression to find existing expressions/subexpressions


The function was renamed?

Yes, it was renamed to reuseIndexIfAlreadyComputed. I have fixed the comment. Thanks for catching it!

naoyam

LGTM. Thanks for the cleanup and enhancement!

save

7d2fc7d

zasdfgbnm changed the title ~~[NOT READY FOR REVIEW] Expr based index hoisting~~ [NOT READY FOR REVIEW] Expr based index hoisting and CSE Dec 2, 2022

zasdfgbnm added 18 commits December 1, 2022 23:50

fix

13d37d3

fixes

ddff274

save

7aecd9f

fix

0957d60

fix

71b7bf5

real cse

c7622a3

fix

432e666

save

6d615c4

fix index select

c110a50

fix

5368856

FusionIssue1133_CUDA

484c688

fix mutator

cef2966

ForLoop::isGrouped()

52c7ded

Revert "ForLoop::isGrouped()"

7215fc7

This reverts commit 52c7ded.

fix

5035e69

fix grouped

c98a26a

isGroup

ff622a5

comments

66c73c4

zasdfgbnm changed the title ~~[NOT READY FOR REVIEW] Expr based index hoisting and CSE~~ Expr based index hoisting and CSE Dec 2, 2022

zasdfgbnm marked this pull request as ready for review December 2, 2022 14:37

zasdfgbnm requested review from csarofeen and naoyam December 2, 2022 14:38

zasdfgbnm commented Dec 2, 2022

View reviewed changes

zasdfgbnm mentioned this pull request Dec 2, 2022

Change TensorIndex to store a single Val* #2234

Merged

zasdfgbnm commented Dec 2, 2022

View reviewed changes

fix misaligned vectorize

4a719cf

zasdfgbnm commented Dec 2, 2022

View reviewed changes

naoyam reviewed Dec 3, 2022

View reviewed changes

Comment thread torch/csrc/jit/codegen/cuda/lower_index_hoist.cpp Outdated

resolve review comments

44fc0a2

zasdfgbnm added 2 commits December 2, 2022 20:58

more cleanup

7a26809

comment

346d3b9

fix hoisted_or_reused_ for hoisted index

6912d75

csarofeen reviewed Dec 5, 2022

View reviewed changes

zasdfgbnm added 6 commits December 5, 2022 13:28

reuseIndexIfAlreadyComputed

904331c

have

86c447f

return -1

5589bf4

better handling index select, improve readibility

0c2927b

comment

2174d71

rename

f6ae83d

naoyam reviewed Dec 6, 2022

View reviewed changes

zasdfgbnm added 2 commits December 5, 2022 22:47

clarification for hoistIndexImpl

13ac450

more comment cleanup

8b93be1

naoyam reviewed Dec 6, 2022

View reviewed changes

Comment thread torch/csrc/jit/codegen/cuda/lower_index_hoist.cpp Outdated

naoyam approved these changes Dec 6, 2022

View reviewed changes

[] -> at

398b157

zasdfgbnm merged commit 7f1bb3f into devel Dec 6, 2022

zasdfgbnm deleted the expr_based_index_hoist branch December 6, 2022 22:14

		strided_inds[i] = GpuLower::current()->commonIndexMap().hoistIndex(
		strided_inds[i], loops);

		auto create_fn = def->newObjectFunc();
		create_fn(index->container(), inputs, {index}, def->attributes());

Conversation

zasdfgbnm commented Dec 2, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zasdfgbnm commented Dec 2, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zasdfgbnm Dec 3, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

zasdfgbnm commented Dec 3, 2022

Uh oh!

zasdfgbnm commented Dec 3, 2022

Uh oh!

csarofeen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

naoyam commented Dec 6, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

naoyam left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

zasdfgbnm commented Dec 2, 2022 •

edited

Loading

zasdfgbnm Dec 3, 2022 •

edited

Loading