Add a 'rolling_buffer' scheduling primitive #7925

mbaret · 2021-04-26T13:54:54Z

Support declaring a tensor as a rolling buffer so that it will be realized as a circular buffer without the need to recompute elements. For further detail you can take a look at this RFC: https://discuss.tvm.apache.org/t/rfc-introducing-a-rolling-buffer-scheduling-primitive/9836

mbaret · 2021-04-28T15:21:27Z

@junrushao1994 @merrymercy could you take a look at this?

junrushao · 2021-05-01T00:07:31Z

CC @Hzfengsy @spectrometerHBH @jinhongyii @MasterJH5574

mbaret · 2021-05-07T14:33:04Z

Can anyone take a look at this?

mbaret · 2021-05-10T15:03:13Z

@manupa-arm @giuseros

Hzfengsy · 2021-05-12T02:08:10Z

It would be great if you can fix CI errors before requesting reviewing :)

giuseros

Only two minor comments, otherwise LGTM

python/tvm/tir/transform/inject_rolling_buffer.py

src/te/schedule/schedule_ops.cc

manupak

Hi @mbaret ,

Thanks for work! Did an initial pass with a lookout for documentation and coding style.

I ll do a technical pass later.

python/tvm/tir/transform/inject_rolling_buffer.py

mbaret · 2021-05-19T20:16:23Z

@Hzfengsy took 3 tries, but CI is now passing :)

manupak

Broadly looks good.

Just few nits and clarification questions.

friendly ping! @junrushao1994 , It would be appreciated to get some feedback and to know what more needs to be done going forward to get this in.

python/tvm/tir/transform/inject_rolling_buffer.py

manupak · 2021-06-01T09:22:26Z

python/tvm/tir/transform/inject_rolling_buffer.py

+                roll_axis = -1
+                for loop in iter_vars:
+                    iter_var = loop.loop_var
+                    if iter_var in bound_iter_vars:


Clarification question : why cant we just look at bound_iter_vars directly ? Is there non-outermost iter_vars identified?

It's because we don't necessarily iterate over a tensor in the same order as its bounds (e.g. we don't have to go axis 0, 1, 2...)

manupak · 2021-06-04T12:32:44Z

adding cc : @jcf94
A friendly ping! We would really appreciate a review on this :)

jcf94

Overall looks good to me. Thanks! @mbaret This is really an interesting work. And also thanks to your remind @manupa-arm !

By the way, this PR didn't extent the support of relay integration, does that means currently we're not able to use this feature in an end to end model? Since in my understanding the fusion part of relay does not support to fuse multiple pool to one subgraph.

jcf94 · 2021-06-07T06:48:55Z

python/tvm/tir/transform/inject_rolling_buffer.py

+    return tvm.tir.transform.prim_func_pass(
+        _ftransform, opt_level=0, name="tir.InjectRollingBuffer"
+    )


Seems this is a pretty complex pass, would you consider to rewrite it as a C++ implementation? (not necessary in this PR, we can add a TODO here if the C++ migration is planned)

Yes I'd definitely consider doing that rewrite at some point. I was actually waiting to see how the new 'scheduling passes' would look for TensorIR so that I could potentially follow any such pattern there. Unfortunately I don't have time to currently so I'd appreciate taking this in with the TODO. Perhaps we can revisit once TensorIR is complete?

jcf94 · 2021-06-07T06:50:23Z

src/te/operation/compute_op.cc

+    bool skip_ivar_domain = !stage->rolling_buffer;
+    ret.init_predicates =
+        MakeBoundCheck(stage, dom_map, ret.init_vmap, skip_ivar_domain, skip_iter);


Suggested change

bool skip_ivar_domain = !stage->rolling_buffer;

ret.init_predicates =

MakeBoundCheck(stage, dom_map, ret.init_vmap, skip_ivar_domain, skip_iter);

ret.init_predicates =

MakeBoundCheck(stage, dom_map, ret.init_vmap, !stage->rolling_buffer, skip_iter);

Nit: Skipping the domain check of IterVar here looks more likely a hack to me, though I don't hava any better suggestion. Seems it's hard to process some more check to guard this operation.

I agree this is a slight hack. It's because in a rolling buffer we need to expand the size and scope of the intermediate buffers and if we drop the bound checks they end up getting corrupted. To do this more accurately we'd need to replicate something closer to Halide's store_at I think to formally change the realization point of the intermediate tensor. However this is my current workaround.

python/tvm/tir/transform/inject_rolling_buffer.py

mbaret · 2021-06-07T15:46:16Z

By the way, this PR didn't extent the support of relay integration, does that means currently we're not able to use this feature in an end to end model? Since in my understanding the fusion part of relay does not support to fuse multiple pool to one subgraph.

This doesn't currently improve the Relay build flow but rather adds a new scheduling primitive primarily for use with tvm::build at the moment. We do however intend to introduce some inter-operating scheduling as part of this design https://discuss.tvm.apache.org/t/rfc-cascade-scheduling/8119 and hope to make use of this primitive as part of that.

jcf94

@mbaret Thanks for your explanation, this looks good to me.

Support declaring a tensor as a rolling buffer so that it will be realized as a circular buffer without the need to recompute elements. Change-Id: I32e0878bb1402ff0276adf3da3f9a4aaac46dd30

Change-Id: I160a2f95fb31beedb9e6ac8c8b45d51d6ec7ebce

Change-Id: I2973475413331cb9ef044407f7771a06491c390d

Change-Id: Ig6133dd822f33a8d32f3ddd8b8ce22b92490694e

mbaret · 2021-06-14T15:23:32Z

ping @manupa-arm, could you please take another look at this patch?

manupak

LGTM

mbaret · 2021-06-15T11:19:06Z

ping @junrushao1994, could you take a look?

junrushao

LGTM :-)

jcf94 · 2021-06-21T02:06:48Z

@mbaret Seems this PR is ready for merge once the conflict has been solved?

Do we have a github issue to track the progress of your Cascade scheduling & RFC? If not, I suggest to have one since this schedule primitive can only work in a special condition. 😄

mbaret · 2021-06-23T14:05:52Z

It would appear the conflict is quite non-trivial (the removal of the Python driver). I don't especially want to work around that by registering my pass into the global registry and then calling it from the C++ API so as not to pollute that with a Python dependency. Given this problem, I shall close this PR for now and when I find the time rewrite the pass in C++.

tqchen added the status: need review label Apr 28, 2021

jroesch requested a review from junrushao May 1, 2021 00:04

jroesch assigned junrushao May 1, 2021

giuseros approved these changes May 12, 2021

View reviewed changes

python/tvm/tir/transform/inject_rolling_buffer.py Outdated Show resolved Hide resolved

src/te/schedule/schedule_ops.cc Outdated Show resolved Hide resolved

mbaret force-pushed the rolling-buffer branch from 0898c7e to eb2894c Compare May 12, 2021 16:42

manupak requested changes May 13, 2021

View reviewed changes

mbaret force-pushed the rolling-buffer branch 2 times, most recently from db157f9 to beaa427 Compare May 19, 2021 12:30

manupak reviewed Jun 1, 2021

View reviewed changes

jcf94 self-assigned this Jun 7, 2021

jcf94 reviewed Jun 7, 2021

View reviewed changes

jcf94 approved these changes Jun 8, 2021

View reviewed changes

mbaret added 3 commits June 9, 2021 10:20

Add a 'rolling_buffer' scheduling primitive

bb294e3

Support declaring a tensor as a rolling buffer so that it will be realized as a circular buffer without the need to recompute elements. Change-Id: I32e0878bb1402ff0276adf3da3f9a4aaac46dd30

Fixed linting

c16c1a9

Change-Id: I160a2f95fb31beedb9e6ac8c8b45d51d6ec7ebce

Fixed review comments

53b99b1

Change-Id: I2973475413331cb9ef044407f7771a06491c390d

mbaret force-pushed the rolling-buffer branch from beaa427 to 2b61414 Compare June 9, 2021 11:05

Fixed more review comments

3ba82ed

Change-Id: Ig6133dd822f33a8d32f3ddd8b8ce22b92490694e

mbaret force-pushed the rolling-buffer branch from 2b61414 to 3ba82ed Compare June 9, 2021 19:43

manupak approved these changes Jun 14, 2021

View reviewed changes

junrushao approved these changes Jun 15, 2021

View reviewed changes

mbaret closed this Jun 23, 2021

NicolaLancellotti mentioned this pull request Nov 4, 2021

Add a 'rolling_buffer' scheduling primitive #9444

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add a 'rolling_buffer' scheduling primitive #7925

Add a 'rolling_buffer' scheduling primitive #7925

mbaret commented Apr 26, 2021 •

edited

Loading

mbaret commented Apr 28, 2021

junrushao commented May 1, 2021

mbaret commented May 7, 2021

mbaret commented May 10, 2021

Hzfengsy commented May 12, 2021

giuseros left a comment

manupak left a comment

mbaret commented May 19, 2021

manupak left a comment

manupak Jun 1, 2021

mbaret Jun 9, 2021

manupak commented Jun 4, 2021

jcf94 left a comment •

edited

Loading

jcf94 Jun 7, 2021 •

edited

Loading

mbaret Jun 7, 2021

jcf94 Jun 7, 2021

mbaret Jun 9, 2021

mbaret commented Jun 7, 2021

jcf94 left a comment

mbaret commented Jun 14, 2021

manupak left a comment

mbaret commented Jun 15, 2021

junrushao left a comment

jcf94 commented Jun 21, 2021 •

edited

Loading

mbaret commented Jun 23, 2021

Add a 'rolling_buffer' scheduling primitive #7925

Add a 'rolling_buffer' scheduling primitive #7925

Conversation

mbaret commented Apr 26, 2021 • edited Loading

mbaret commented Apr 28, 2021

junrushao commented May 1, 2021

mbaret commented May 7, 2021

mbaret commented May 10, 2021

Hzfengsy commented May 12, 2021

giuseros left a comment

Choose a reason for hiding this comment

manupak left a comment

Choose a reason for hiding this comment

mbaret commented May 19, 2021

manupak left a comment

Choose a reason for hiding this comment

manupak Jun 1, 2021

Choose a reason for hiding this comment

mbaret Jun 9, 2021

Choose a reason for hiding this comment

manupak commented Jun 4, 2021

jcf94 left a comment • edited Loading

Choose a reason for hiding this comment

jcf94 Jun 7, 2021 • edited Loading

Choose a reason for hiding this comment

mbaret Jun 7, 2021

Choose a reason for hiding this comment

jcf94 Jun 7, 2021

Choose a reason for hiding this comment

mbaret Jun 9, 2021

Choose a reason for hiding this comment

mbaret commented Jun 7, 2021

jcf94 left a comment

Choose a reason for hiding this comment

mbaret commented Jun 14, 2021

manupak left a comment

Choose a reason for hiding this comment

mbaret commented Jun 15, 2021

junrushao left a comment

Choose a reason for hiding this comment

jcf94 commented Jun 21, 2021 • edited Loading

mbaret commented Jun 23, 2021

mbaret commented Apr 26, 2021 •

edited

Loading

jcf94 left a comment •

edited

Loading

jcf94 Jun 7, 2021 •

edited

Loading

jcf94 commented Jun 21, 2021 •

edited

Loading