[relay][pass] Annotation for heterogeneous compilation #2361

zhiics · 2019-01-02T19:33:25Z

This PR adds passes in Relay to annotate expressions to indicate which device/context each operator should be executed at. #2296 is the proposed RFC. The following changes are made in this PR.

Extend target to accept a dictionary of device/context to target and add the fallback_device argument in build of build_module.py
build(func, target=None, target_host=None, params=None, fallback_device=None)
Slightly Modify compile engine to stop lowering and generating schedules for device copy operators since the real data transferring will be only performed at runtime and no schedule is needed for this type of operators.
Modify memory plan to return both storage ids and device ids.
Add on_device(expr, dev_id) and device_copy(expr, src_dev_id, dst_dev_id) operators as synthetic op as @jroesch suggested. The former takes an expr and a device_id as inputs which indicate where an expression should be annotated. The latter will be used to perform data copy between different devices.
Write several passes to validate the annotated program, rewrite the program (e.g. insert device copy operators), and propagate the device information from device copy operators to the other expression, etc.
Add unit tests to test the functionality of different annotation schemes.

@tqchen @jroesch @yidawang @yzhliu @tmoreau89

jroesch · 2019-01-03T21:57:41Z

I just got back from vacation, will review this PR tonight, looks like great work 👍

zhiics · 2019-01-06T02:31:10Z

@jroesch Thank you. Please take a look when you have time.

jroesch · 2019-01-06T23:35:38Z

@zhiics sorry I had done a partial review and forgot to hit submit, will finish rest and post comments right now.

jroesch

Overall looks like a good first pass, most of my comments are small nits about the comments. It would be good to solicit review from someone who is familiar with heterogeneous-execution.

python/tvm/relay/op/tensor.py

src/relay/backend/compile_engine.cc

python/tvm/relay/op/tensor.py

python/tvm/relay/ir_pass.py

python/tvm/relay/build_module.py

python/tvm/relay/backend/_backend.py

tmoreau89

Thank you @zhiics for this excellent contribution. Quick question: how hard would it be to construct a heterogeneous test case (as test_pass_annotation.py) to run on the VTA simulator? It would be interesting to be able to provide explicit control over what components of the graphs get offloaded to CPU vs. VTA with this approach.

zhiics · 2019-01-07T17:44:13Z

@tmoreau89 It shouldn't be hard if you only want to offload most of ops to VTA and only keep a few to CPU. Otherwise, I think it might be a little tedious to traverse the program and add on_device ops. Another way might make annotation easier is probably to allow users to just provide relay op names instead of adding on_device primitives directly, and we can add these primitives in the backend. How do you think?

icemelon · 2019-01-07T18:59:49Z

@zhiics Could you use git rebase so that it will include your commits?

zhiics · 2019-01-07T19:01:27Z

oops. Sorry. I will rebase now.

yzhliu

haven't finished yet, will continue tonight. A high level question: for now users do need to write relay.on_device explicitly for annotation?

tests/python/relay/test_pass_annotation.py

python/tvm/relay/backend/graph_runtime_codegen.py

python/tvm/relay/build_module.py

zhiics · 2019-01-08T00:55:34Z

@yzhliu Yes, we need to use relay.on_device to do annotation. I think we can also allow users to provide the op names for annotation as well, but that probably is worthy a separate PR.

tqchen · 2019-01-08T17:32:23Z

related RFC #2391

jroesch · 2019-01-08T23:40:04Z

@zhiics @yzhliu my suggestion to use the explicit op is to give us flexibility. We can build approaches which are automated/use user input/etc by just writing a pass which annotates the program with the appropriate on_device calls.

zhiics · 2019-01-08T23:47:22Z

@jroesch Yes, I actually agree with you. We can add other passes for different annotation schemes, but users should have the flexibility to annotate expressions from the language directly.

yzhliu · 2019-01-09T22:20:04Z

src/relay/backend/graph_plan_memory.cc

    std::vector<StorageToken*> tokens;
+    int device_id = node_device_map_.count(GetRef<Expr>(op))
+                        ? node_device_map_[GetRef<Expr>(op)]->value
+                        : 0;


shall we add 0 as a placeholder in DLDeviceType so that others will not use it for other special purpose by mistake.
also device_id -> device_type

@yzhliu Yes, I also thought that it's probably necessary to have a field in DLDeviceType, like 'kDLUNDEFINED = 0'. Let's keep this for now. I will send a RFC later to hear from more people because it needs a slight change in dlpack.

src/relay/pass/device_annotation.cc

add fallback cpptest fix lint accept both nn.op_name and op_name accept both nn.op_name and op_name use expr annotation instead of op names fix test_back_graph_runtime unit test

@jroesch

add fallback cpptest fix lint accept both nn.op_name and op_name accept both nn.op_name and op_name use expr annotation instead of op names fix test_back_graph_runtime unit test address @jroesch's comments

zhiics · 2019-01-11T06:15:46Z

@jroesch Please take another look. We can probably bring it in if everything looks good to you. Thanks.

jroesch

Reviewed the delta, looks good to me 👍

yzhliu added the status: need review label Jan 3, 2019

jroesch requested changes Jan 6, 2019

View reviewed changes

tmoreau89 reviewed Jan 7, 2019

View reviewed changes

zhiics force-pushed the relay_annotation branch from 07b5ce3 to 789d510 Compare January 7, 2019 18:57

zhiics force-pushed the relay_annotation branch from 789d510 to 074f95d Compare January 7, 2019 19:04

yzhliu reviewed Jan 8, 2019

View reviewed changes

yzhliu reviewed Jan 9, 2019

View reviewed changes

Zhi Chen and others added 4 commits January 9, 2019 17:55

[relay] heterogeneous graph annotation

4d36348

add fallback cpptest fix lint accept both nn.op_name and op_name accept both nn.op_name and op_name use expr annotation instead of op names fix test_back_graph_runtime unit test

[relay] heterogeneous graph annotation

7ac6cb2

add fallback cpptest fix lint accept both nn.op_name and op_name accept both nn.op_name and op_name use expr annotation instead of op names fix test_back_graph_runtime unit test address @jroesch's comments

fix target

8648495

fix @yzhliu's comments

7e6be33

zhiics force-pushed the relay_annotation branch from 72121db to 7e6be33 Compare January 10, 2019 02:09

zhiics added 2 commits January 10, 2019 10:22

refactor device_id to device_type to keep consistent to DLContext

b5d2e45

move on_device op to annotation namespace

a7b86a5

yzhliu approved these changes Jan 11, 2019

View reviewed changes

jroesch approved these changes Jan 11, 2019

View reviewed changes

icemelon merged commit 09236bf into apache:master Jan 11, 2019

icemelon removed the status: need review label Jan 12, 2019

yzhliu added the status: accepted label Jan 12, 2019

zhiics deleted the relay_annotation branch January 12, 2019 04:56

This was referenced Jan 17, 2019

[RFC] Streamline the relay build and optimize interface #2449

Closed

[Relay][RFC] Compilation for heterogeneous execution #2296

Closed

ZihengJiang mentioned this pull request Feb 1, 2019

TVM 0.5 Release Note #2448

Closed

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

[relay][pass] Annotation for heterogeneous compilation (apache#2361)

7c180e3

wweic pushed a commit to neo-ai/tvm that referenced this pull request Feb 20, 2019

[relay][pass] Annotation for heterogeneous compilation (apache#2361)

5555674

wweic mentioned this pull request Oct 22, 2019

[RFC][VM] Heterogeneous execution in Relay VM #4178

Closed

4 tasks

[relay][pass] Annotation for heterogeneous compilation #2361

[relay][pass] Annotation for heterogeneous compilation #2361

Uh oh!

Conversation

zhiics commented Jan 2, 2019 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jroesch commented Jan 3, 2019

Uh oh!

zhiics commented Jan 6, 2019

Uh oh!

jroesch commented Jan 6, 2019

Uh oh!

jroesch 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!

Uh oh!

Uh oh!

tmoreau89 left a comment

Choose a reason for hiding this comment

Uh oh!

zhiics commented Jan 7, 2019

Uh oh!

icemelon commented Jan 7, 2019

Uh oh!

zhiics commented Jan 7, 2019

Uh oh!

yzhliu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhiics commented Jan 8, 2019

Uh oh!

tqchen commented Jan 8, 2019

Uh oh!

jroesch commented Jan 8, 2019

Uh oh!

zhiics commented Jan 8, 2019

Uh oh!

yzhliu Jan 9, 2019

Choose a reason for hiding this comment

Uh oh!

zhiics Jan 9, 2019 • 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!

Uh oh!

Uh oh!

zhiics commented Jan 11, 2019

Uh oh!

jroesch left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

zhiics commented Jan 2, 2019 •

edited

Loading

zhiics Jan 9, 2019 •

edited

Loading