Skip to content

[Relay] Make DeviceAnalyzer a mixed mode visitor - #10248

Merged
tmoreau89 merged 4 commits into
apache:mainfrom
masahi:device-planner-mixed-mode
Feb 15, 2022
Merged

tmoreau89 merged 4 commits into
apache:mainfrom
masahi:device-planner-mixed-mode

Conversation

@masahi

@masahi masahi commented Feb 15, 2022

Copy link
Copy Markdown
Member

Closes #10135

There are about 5 visitors / mutators in device_planner.cc, but only one of them is mixed mode for now following @mei-ye's original PR.

@mei-ye Please confirm if this fixes the issue on windows.

cc @mbrookhart @mbs-octoml @tmoreau89

One issue with making DeviceAnalyzer inherit from MixedModeVisitor is the intermediate tuple in call_lowered op, as in

%0 = on_device(%x, constrain_body=False) /* from_string */ /* ty=Tensor[(128, 128), float32] */;
%1 = on_device(%y, constrain_body=False) /* from_string */ /* ty=Tensor[(128, 128), float32] */;
%2 = on_device(%z, constrain_body=False) /* from_string */ /* ty=Tensor[(128, 128), float32] */;
%3 = (%0, %1, %2);
call_lowered(@gem, %3) /* from_string */ /* ty=Tensor[(128, 128), float32] */ 

If we do pure post-DFS traversal, the tuple will be visited before call_lowered, and that makes all original CallNode args unify to the domain of the tuple. The solution we agreed upon is to update ExpandDataflow function to ignore that tuple when it expands call_lowered op.

*
* TODO(mbs): Rename VirtualDevice |-> VirtualDevice, and use 'virtual device' (or just 'device')
* throughout.
*

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mbs-octoml I assume this is ok to remove.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh yeah, thanks.

@masahi
masahi force-pushed the device-planner-mixed-mode branch from bf9c637 to 0708bf7 Compare February 15, 2022 08:48

@mbrookhart mbrookhart left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@tmoreau89
tmoreau89 merged commit 2b00835 into apache:main Feb 15, 2022
@tmoreau89

Copy link
Copy Markdown
Contributor

Thank you @mbrookhart and @masahi - PR has been merged.

@mei-ye

mei-ye commented Feb 15, 2022

Copy link
Copy Markdown
Contributor

I confirmed that this fixes the issue I saw on windows. Thanks!

ylc pushed a commit to ylc/tvm that referenced this pull request Feb 16, 2022
* hack to ExpandDataflow

* add test from mei

* Update DeviceAnalyzer to inherit from MixedModeVisitor

* indent
pfk-beta pushed a commit to pfk-beta/tvm that referenced this pull request Apr 11, 2022
* hack to ExpandDataflow

* add test from mei

* Update DeviceAnalyzer to inherit from MixedModeVisitor

* indent
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Bug] Update DevicePlanner to use MixedModeVisitor

5 participants