[RELAY][AST] Add virtual device as a first class field to Relay expressions #45
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
electriclilies
changed the title
denise-k
reviewed
Dec 1, 2021
tqchen
changed the title
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 2, 2021
Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 2, 2021
Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45
|
Here's my proposed implementation: https://github.com/apache/tvm/pull/9641/files |
|
Thanks Lily, LGTM. Let me apologize for getting us into the current on_device + function attributes mess your proposal gets us out of. It seemed like a workable approach at the time which would not require both an AST change and careful treatment of AST node copies. But in retrospect the on_device + function attributes approach has caused at least four bugs due to the fragility of maintaining its annotation invariant. I'm hoping the experience we gain from adding virtual_device_ will inform improvements to checked_type_. In particular getting us out of the need to rerun InferType after every xform (and sometimes in strict succession). But one step at a time. |
comaniac
reviewed
Dec 3, 2021
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 4, 2021
- Allow device_copy into and out of PrimFuncs - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented.
comaniac
reviewed
Dec 6, 2021
There was a problem hiding this comment.
LGTM. I don't have other comments at this moment. Gentle ping @jroesch @zhiics @ZihengJiang @junrushao1994 for review, or I'll merge this RFC tomorrow.
|
@comaniac Thanks for reviewing! Looking forward to it going in :D |
|
Thanks @electriclilies and all reviewers. This RFC is merged. |
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 7, 2021
- Allow device_copy into and out of PrimFuncs - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented.
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 7, 2021
…pe constraints.
- Memory scope constraints can flow both out of and in to PrimFuncs
introduced by LowerTE. In TIR memory scopes are represented by
'storage scopes' on the PointerType type annotations on TIR Buffer data
variables.
- It is straightforward to extract memory scopes from PrimFuncs by
looking at the PrimFunc's buffer_map. We do this is 'phase 1' of
PlanDevices, which collects all the device constraints implied by
- However, pushing memory constraints in to PrimFuncs is more challenging
due to buffer aliasing. This aspect is still experimental.
- Allow device_copies to be inserted for both arguments and
results of PrimFunc calls, on the assumption PlanDevices has
already established a consistent device assignment prior to
lowering and any new mismatch is required to match up memory scopes.
- Realized both dead_code.cc and fold_constant.cc would
happily move values into a different lexical virtual
device context since device_planner.cc was being
'clever' and eliding on_devices for let-bound values
when there's no change. Fixed so that every let-bound
value has an on_device. Will be much better after
apache/tvm-rfcs#45 is implemented.
- Other misc fixups.
|
Thanks cody! |
mbs-octoml
added a commit
to mbs-octoml/mbs-tvm
that referenced
this pull request
Dec 13, 2021
…straints.
This PR:
1) Makes PlanDevices consider lowered calls when solving device domain constraints.
2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data
Var type annotation PointerTypes storage_scope fields) to the memory_scope
fields of the SEScopes which PlanDevices unifies over.
3) Allows new device_copies to be inserted on the arguments and results of lowered
calls so as to acount for any memory scope mismatches which are now apparent.
[device_planner.cc has main changes, rest is secondary.]
In the short term we'd like to use this machinery to flow memory scope choices made
during lowering back out into the overall Relay program. In the longer term we'd
also like to be able to use memory scopes to influence the lowering of
yet-to-be-lowered functions (or lowered functions which have yet to been scheduled,
a distinction now possible with TensorIR).
- Memory scope constraints can flow both out of and in to PrimFuncs
introduced by LowerTE. In TIR memory scopes are represented by
'storage scopes' on the PointerType type annotations on TIR Buffer data
variables.
- It is straightforward to extract memory scopes from PrimFuncs by
looking at the PrimFunc's buffer_map. We do this is 'phase 1' of
PlanDevices, which collects all the device constraints implied by
- However, pushing memory constraints in to PrimFuncs is more challenging
due to buffer aliasing. This aspect is still experimental.
- Allow device_copies to be inserted for both arguments and
results of PrimFunc calls, on the assumption PlanDevices has
already established a consistent device assignment prior to
lowering and any new mismatch is required to match up memory scopes.
We use the new 'free' on_device annotations to implement this.
Coming along for the ride:
- To make unit tests of mixed Relay/TIR functions possible needed
to be able to supply a checked_type to GlobalVar since that's currently
the only way to give a Relay type to PrimFuncs.
- Use GenSym to get unique var names in ANF & partial eval so easier
to diff debug output between passes and connect program fragments
back into the overall program. Relying on pretty-printing to
automagically unique-ify var names is certainly cute but until we
have better span support is very hard to work with.
- Realized both dead_code.cc and fold_constant.cc would
happily move values into a different lexical virtual
device context since device_planner.cc was being
'clever' and eliding on_devices for let-bound values
when there's no change. Fixed so that every let-bound
value has an on_device. Will be much better after
apache/tvm-rfcs#45 is implemented.
- Make build -Werror clean for clang-12 (mostly move fixups).
- Address post-submit comments from apache#9693.
mbrookhart
pushed a commit
to apache/tvm
that referenced
this pull request
Dec 14, 2021
…straints. (#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from #9693. * [checkpoint] thread safe GenSym
ylc
pushed a commit
to ylc/tvm
that referenced
this pull request
Jan 7, 2022
…straints. (apache#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from apache#9693. * [checkpoint] thread safe GenSym
yangulei
pushed a commit
to yangulei/tvm
that referenced
this pull request
Jan 11, 2022
…straints. (apache#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from apache#9693. * [checkpoint] thread safe GenSym
yangulei
pushed a commit
to yangulei/tvm
that referenced
this pull request
Jan 12, 2022
…straints. (apache#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from apache#9693. * [checkpoint] thread safe GenSym
ylc
pushed a commit
to ylc/tvm
that referenced
this pull request
Jan 13, 2022
…straints. (apache#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from apache#9693. * [checkpoint] thread safe GenSym
qsqqsqqsq-intellif
pushed a commit
to qsqqsqqsq-intellif/tvm
that referenced
this pull request
Apr 29, 2022
…straints. (apache#9613) * [Relay] Re-run PlanDevices after LowerTE to flow new memory scope constraints. This PR: 1) Makes PlanDevices consider lowered calls when solving device domain constraints. 2) Connects the storage scopes on PrimFunc parameters (encoded in their Buffer data Var type annotation PointerTypes storage_scope fields) to the memory_scope fields of the SEScopes which PlanDevices unifies over. 3) Allows new device_copies to be inserted on the arguments and results of lowered calls so as to acount for any memory scope mismatches which are now apparent. [device_planner.cc has main changes, rest is secondary.] In the short term we'd like to use this machinery to flow memory scope choices made during lowering back out into the overall Relay program. In the longer term we'd also like to be able to use memory scopes to influence the lowering of yet-to-be-lowered functions (or lowered functions which have yet to been scheduled, a distinction now possible with TensorIR). - Memory scope constraints can flow both out of and in to PrimFuncs introduced by LowerTE. In TIR memory scopes are represented by 'storage scopes' on the PointerType type annotations on TIR Buffer data variables. - It is straightforward to extract memory scopes from PrimFuncs by looking at the PrimFunc's buffer_map. We do this is 'phase 1' of PlanDevices, which collects all the device constraints implied by - However, pushing memory constraints in to PrimFuncs is more challenging due to buffer aliasing. This aspect is still experimental. - Allow device_copies to be inserted for both arguments and results of PrimFunc calls, on the assumption PlanDevices has already established a consistent device assignment prior to lowering and any new mismatch is required to match up memory scopes. We use the new 'free' on_device annotations to implement this. Coming along for the ride: - To make unit tests of mixed Relay/TIR functions possible needed to be able to supply a checked_type to GlobalVar since that's currently the only way to give a Relay type to PrimFuncs. - Use GenSym to get unique var names in ANF & partial eval so easier to diff debug output between passes and connect program fragments back into the overall program. Relying on pretty-printing to automagically unique-ify var names is certainly cute but until we have better span support is very hard to work with. - Realized both dead_code.cc and fold_constant.cc would happily move values into a different lexical virtual device context since device_planner.cc was being 'clever' and eliding on_devices for let-bound values when there's no change. Fixed so that every let-bound value has an on_device. Will be much better after apache/tvm-rfcs#45 is implemented. - Make build -Werror clean for clang-12 (mostly move fixups). - Address post-submit comments from apache#9693. * [checkpoint] thread safe GenSym
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is a proposal to add virtual devices as a first class field to Relay expressions.
cc @mbs-octoml @mikepapadim @jroesch