[TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position #12450

yincs-intellif · 2022-08-16T07:54:35Z

Current TIR "compute_at" primitive will compute at it's closest consumers. When a block has multiple producers, whoever
compute at later who is behind. But for some special hardware, we usually hope keep the a certain order whatever it's compute at early or late.
eg: block A and block B are producers of block C. block A compute at block C first and block B compute at block C later. We hope the result is block B->block A->block C under some loop var.

cc @Hzfengsy @junrushao1994

Hzfengsy

LGTM except for one nit

Hzfengsy · 2022-08-16T08:01:57Z

Do you have ideas about the naming of to_early_stage? cc @junrushao1994 @spectrometerHBH

tests/python/unittest/test_tir_schedule_compute_at.py

junrushao

This is definitely a valid concern, while we might want to deliberate on the following two questions:

Is it possible to generalize the boolean flag to_early_stage to an integer, and allow users to specify which produecr/consumer it should move to?
The naming to_early_stage itself is definitely a bit less informative. Let's consider a better name

Besides, given that we have reverse_compute_at symmetrically, shall we also support a similar flag for it?

Thanks in advance!

yincs-intellif · 2022-08-16T08:50:41Z

Thanks @junrushao1994 I quite agree with your suggestion. Currently I modify it according to scenario we have encountered. Do you have any suggestion about the better name?

yincs-intellif · 2022-08-16T09:35:20Z

Hi @junrushao1994 , what do you think of my follow suggestion?

we give a integer parameter instead of to_early_stage, which means block number, masked as N.
Default we will get a compute at position, which is after the last producer and before the first consumer, masked as [P, C].
when N in [P,C], we use N for the position of compute at, otherwise use C or report error.

The above rules apply to "reverse_compute_at" at same time

Thanks in advance!

yincs-intellif · 2022-08-22T03:19:17Z

@junrushao Hi, I have changed as you mentioned, could you please help review again？

Thanks in advance!

junrushao · 2022-08-22T04:58:27Z

Hey sorry for the late reply! Happy to do another round of review tomorrow!

yincs-intellif · 2022-08-24T01:06:45Z

Hey sorry for the late reply! Happy to do another round of review tomorrow!

@junrushao Hi, I don't receive any feedback information. Did I miss something?

include/tvm/tir/schedule/schedule.h

src/tir/schedule/primitive/compute_at.cc

Hzfengsy

LGTM. Thanks for the continuous improvement!

junrushao · 2022-08-26T04:44:33Z

python/tvm/tir/schedule/schedule.py

+            The block index of the loop body subtree blocks
+            -1 means inserted into the last possible insertion point
+            -2 means inserted into the first possible insertion point


Suggested change

The block index of the loop body subtree blocks

-1 means inserted into the last possible insertion point

-2 means inserted into the first possible insertion point

The block index of the loop body subtree blocks:

- `index = -1` means inserted into the last possible insertion point;

- `index = -2` means inserted into the first possible insertion point;

- Otherwise, `index` is a nonnegative number that indicates the insertion point

junrushao · 2022-08-26T04:44:52Z

src/tir/schedule/primitive.h

+ * \param index The block index of the loop body subtree blocks
+ * -1 means inserted into the last possible insertion point
+ * -2 means inserted into the first possible insertion point


Let's update the doc accordingly

junrushao · 2022-08-26T04:45:18Z

src/tir/schedule/primitive/compute_at.cc

+  // Step 4. Return the possible insertion point according to index
+  int insert_position;
+  if (index == -1) {
+    insert_position = split.first_consumer_position;
+  } else if (index == -2) {
+    insert_position = split.last_producer_position + 1;
+  } else if (index >= 0 && index >= split.last_producer_position + 1 &&
+             index <= split.first_consumer_position) {
+    insert_position = index;
+  } else {
+    LOG(FATAL) << "Valid index:(-1, -2, [" << split.last_producer_position + 1 << ", "
+               << split.first_consumer_position << "]), "
+               << "current index=" << index;
+    throw;
+  }
+  return insert_position;


Thanks for updating the logic here! It's much clearer now!

junrushao · 2022-08-26T04:45:50Z

src/tir/schedule/primitive/compute_at.cc

      /*subtrees=*/AsArray(loop->body),
      /*producer_srefs=*/producer_srefs,
-      /*consumer_srefs=*/consumer_srefs, /*block2realize=*/&block2realize);
+      /*consumer_srefs=*/consumer_srefs, /*block2realize=*/&block2realize,
+      /*index*/ index);


Suggested change

/*index*/ index);

/*index=*/index);

junrushao

Thanks @yincs-intellif for the update! The diff is in a pretty good shape now! Let's fix some final nitpicks and then get it merged!

…o choose possible position

yincs-intellif · 2022-08-26T11:11:39Z

Hi @junrushao @Hzfengsy The final modification has been updated.
Thank you for the high quality review!

Hzfengsy · 2022-08-26T11:44:33Z

Thanks @yincs-intellif for this PR and continuous update.

…o choose possible position (apache#12450) Current TIR "compute_at" primitive will compute at it's closest consumers. When a block has multiple producers, whoever compute at later who is behind. But for some special hardware, we usually hope keep the a certain order whatever it's compute at early or late. eg: block A and block B are producers of block C. block A compute at block C first and block B compute at block C later. We hope the result is block B->block A->block C under some loop var.

Hzfengsy approved these changes Aug 16, 2022

View reviewed changes

Hzfengsy reviewed Aug 16, 2022

View reviewed changes

tests/python/unittest/test_tir_schedule_compute_at.py Outdated Show resolved Hide resolved

junrushao requested changes Aug 16, 2022

View reviewed changes

yincs-intellif force-pushed the main branch from 26bf76c to 8e7382c Compare August 16, 2022 08:36

yincs-intellif force-pushed the main branch 3 times, most recently from 2e4c30e to ff77782 Compare August 17, 2022 10:13

yincs-intellif requested a review from junrushao August 18, 2022 01:46

junrushao reviewed Aug 24, 2022

View reviewed changes

include/tvm/tir/schedule/schedule.h Outdated Show resolved Hide resolved

junrushao reviewed Aug 24, 2022

View reviewed changes

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved

Hzfengsy requested changes Aug 24, 2022

View reviewed changes

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved

src/tir/schedule/primitive/compute_at.cc Outdated Show resolved Hide resolved

src/tir/schedule/primitive/compute_at.cc Show resolved Hide resolved

Hzfengsy self-assigned this Aug 24, 2022

yincs-intellif force-pushed the main branch from ff77782 to be5e34c Compare August 25, 2022 03:07

yincs-intellif changed the title ~~[TIR][Schedule] enhance compute_at primitive to choose proper position~~ [TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position Aug 25, 2022

yincs-intellif force-pushed the main branch from be5e34c to 3f0055f Compare August 25, 2022 03:13

yincs-intellif requested review from Hzfengsy and junrushao and removed request for Hzfengsy and junrushao August 25, 2022 03:16

Hzfengsy approved these changes Aug 25, 2022

View reviewed changes

yincs-intellif requested a review from junrushao August 25, 2022 06:51

junrushao reviewed Aug 26, 2022

View reviewed changes

junrushao approved these changes Aug 26, 2022

View reviewed changes

yincs-intellif force-pushed the main branch from 3f0055f to f0f74e4 Compare August 26, 2022 07:25

[TIR][Schedule] enhance compute_at and reverse_compute_at primitive t…

75082dc

…o choose possible position

yincs-intellif force-pushed the main branch from f0f74e4 to 75082dc Compare August 26, 2022 08:42

Hzfengsy merged commit e02f2f9 into apache:main Aug 26, 2022

AndrewZhaoLuo mentioned this pull request Oct 4, 2022

TVM v0.10.0.rc0 Release Candidate Notes #12979

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position #12450

[TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position #12450

yincs-intellif commented Aug 16, 2022 •

edited by github-actions bot

Loading

Hzfengsy left a comment

Hzfengsy commented Aug 16, 2022

junrushao left a comment

yincs-intellif commented Aug 16, 2022

yincs-intellif commented Aug 16, 2022 •

edited

Loading

yincs-intellif commented Aug 22, 2022

junrushao commented Aug 22, 2022

yincs-intellif commented Aug 24, 2022

Hzfengsy left a comment

junrushao Aug 26, 2022

junrushao Aug 26, 2022

junrushao Aug 26, 2022

junrushao Aug 26, 2022

junrushao left a comment

yincs-intellif commented Aug 26, 2022

Hzfengsy commented Aug 26, 2022

-            The block index of the loop body subtree blocks
-            -1 means inserted into the last possible insertion point
-            -2 means inserted into the first possible insertion point
+            The block index of the loop body subtree blocks:
+            - `index = -1` means inserted into the last possible insertion point;
+            - `index = -2` means inserted into the first possible insertion point;
+            - Otherwise, `index` is a nonnegative number that indicates the insertion point

[TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position #12450

[TIR][Schedule] enhance compute_at and reverse_compute_at primitive to choose possible position #12450

Conversation

yincs-intellif commented Aug 16, 2022 • edited by github-actions bot Loading

Hzfengsy left a comment

Choose a reason for hiding this comment

Hzfengsy commented Aug 16, 2022

junrushao left a comment

Choose a reason for hiding this comment

yincs-intellif commented Aug 16, 2022

yincs-intellif commented Aug 16, 2022 • edited Loading

yincs-intellif commented Aug 22, 2022

junrushao commented Aug 22, 2022

yincs-intellif commented Aug 24, 2022

Hzfengsy left a comment

Choose a reason for hiding this comment

junrushao Aug 26, 2022

Choose a reason for hiding this comment

junrushao Aug 26, 2022

Choose a reason for hiding this comment

junrushao Aug 26, 2022

Choose a reason for hiding this comment

junrushao Aug 26, 2022

Choose a reason for hiding this comment

junrushao left a comment

Choose a reason for hiding this comment

yincs-intellif commented Aug 26, 2022

Hzfengsy commented Aug 26, 2022

yincs-intellif commented Aug 16, 2022 •

edited by github-actions bot

Loading

yincs-intellif commented Aug 16, 2022 •

edited

Loading