[RFC] Unified device/target/memory scope planning #38

mbs-octoml · 2021-09-29T01:26:52Z

In rendered form:
https://github.com/mbs-octoml/mbs-tvm-rfcs/blob/mbs-target-and-device-planning/rfcs/0038-unified-device-target-and-memory-scope-planning.md

Some earlier discussion:
apache/tvm#8892

Tracking issue: apache/tvm#9327
(CORE-95 in the OctoML JIRA).

manupak · 2021-10-01T10:08:40Z

Hi @mbs-octoml ,

I may have put a related comment here : apache/tvm#8892 (comment)
However, partitioning for devices of same kind is a step forward from unifying the BYOC and Device annotations.

Is this the RFC intended to cover these all ?

mbs-octoml · 2021-10-01T14:21:36Z

Thanks @manupa-arm for the reminder there were some good comments on #8892. I see a work stream:

get the multi-target handling under control, and stop relying on the limiting device type -> target mapping
allow device planning to be re-run with additional memory scope constraints to further partition, those constraints may originate in already lowered PrimFuncs.
allow labeling of sets of target/devices a la BYOC target labeling (but probably continue to just pick the 'first')
bring BYOC target labeling / partitioning and device planning together
replace naive partitioning with something subject to optimization

Here I want to just focus on 1 -- everything beyond that really needs face-to-face discussion.

From 2 onward obviously overlaps your USMP. My vague thought was we can work from opposite ends and reconcile at 5. Ie 2-4 sets us up to work in a combined Relay+TIR world, then 5 is where everything we've learned from USMP could perhaps be replayed. Anyway, that's just a vague thought so I'd love to talk more about it.

mbs-octoml · 2021-10-01T18:36:28Z

Note to self: The With convention should probably also be removed by this work also, but I've not audited the code to see how pervasive it is.

Target already has a GetAttr<String>("device") convention, but I think the map from name to DLDeviceType is only on the py side, and the attribute is not consistently set.

areusch

@mbs-octoml thanks for the draft RFC! added some thoughts

rfcs/00xx-improved-multi-target-handling.md

areusch · 2021-10-04T19:28:46Z

rfcs/00xx-improved-multi-target-handling.md

+   ```
+   (We could also use a `Device` and accept the redundant `DLDeviceType` specification.) It is trivial
+   to go from an "on_device" label to a `TargetDevice` and back using the global `Target` registry.
+5. Remove all uses of `TargetMap`. For example, in `LowerTEPass` we simply use the `TargetDevice` associated with


do you propose any replacement in case we do need a map-like struct? Map<target_label, Target>?

I've removed this from the RFC. In #9313 you'll see I kept TargetMap but introduces a helper class to hid it. Over time I think we can replace TargetMap with just Array, but I feel it's not worth getting specific about that in an RFC and is more just a cleanup task. It may well come out of @Mousius ' work on tvmc target specification cleanup.

rfcs/00xx-improved-multi-target-handling.md

mbs-octoml · 2021-10-08T20:39:30Z

More notes to self:

You can get the device_type for a Target from target_kind->device_type. Somehow I missed that very obvious fact.
Eric has a very nice write-up explaining devices vs targets vs device api at docs/dev/device_target_interactions.rst

tqchen · 2021-10-08T22:34:55Z

cc @ zxybazh since you authored the Target system, cc @zhiics @comaniac as it is related to BYOC

…g in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

mbs-octoml · 2021-10-21T20:46:08Z

Thanks for the comments @areusch and @manupa-arm . Now that I've started working on this (with an emphasis on handling memory scopes) I've decided to shift focus a bit. In particular Manupa I'm now much more motivated to tackle the BYOC/device planning overlap aspect, which I think you're particularly interested in. PTAL.

…g in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

comaniac · 2021-10-29T22:39:12Z

For the proposed BYOC flow (i.e., MergeAndAnnotate/AnnotateSEScopes/PlanDevices/PartitionBySEScope), it doesn't clear to me that whether the developer programming model will change or not. Specifically, could we still use the current approaches (i.e., op-based annotation and pattern-based annotation)? Also how the approach that has a custom Relay pass to annotate compiler_begin/compiler_end would change? Thanks.

…g in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

…g in 'device' planning. (#9313) [Target] Adds SEScope (Storage/Execution Scope) for use as new unit of planning in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form. * Reworked to avoid global SEScopeCache. Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. So instead placed the cache in the CompilationConfig class. Since that class now has everything the device planner needs to do its job, promoted it to be an FFI-able Object, which is now in compilation_config.{h,cc}. I think we can do much better with CompilationConfig, but for now keeping it to the minimum I needed to prepare for device planning from all the executor compilation codepaths.

[checkpoint] rebase [checkpoint] fix merge [checkpoint] lint [checkpoint] rebase [checkpoint] Fixed stray use of kDLCPU in vm/profiler/vm.cc [checkpoint] lint trivia [checkpoint] fix unit tests [checkpoint] device planner unit tests passing again [checkpoint] Switch over to new CompilerOptions [checkpoint] include [checkpoint] Almost working again Need to move the SEScopeCache into CompilationConfig and pass that into DeviceDomains instead of just the Vector<Target>. Then the host_se_scope can be memoized so that direct uses of that scope downstream will match up with se_scopes already established by PlanDevices. Sigh. [checkpoint] Use cache in device domains. [checkpoint] more moves [checkpoint] lints [checkpoint] Fix merge with VM profiling changes. [checkpoint] trivial [checkpoint] rebase fix [checkpoint] More unit tests. Getting ready to fork out SEScope changes alone. [checkpoint] lints [checkpoint] All plan devices unit tests pass [checkpoint] First unit test passes [checkpoint] Another go at target management This at least centralizes all the hackery. Compiles. [commit] Start to rollback resolving to target in planner. Better is to do it as stand alone pass I think. Besides it doesn't work with the structural test for expected output. [checkpoint] Almost have first unit test going. About to merge Michalis' changes. target_host is still a mess. Starting to eliminate target_map. [checkpoint] Cleanup VM device matching [checkpoint] Compiles [checkpoint] First sweep replacing DLDeviceType with SEScope VM still not done. [checkpoint] Expose CompilationConfig ctor in py [checkpoint] CompilationConfig is nullable for default ctor [checkpoint] Don't use target:: namespace [checkpoint] Promote CompilationConfig to be FFI-friendly Object Also rework to never mix the host_target into the 'primitive' targets. [checkpoint] ResolveSEScope on CompilationConfig [checkpoint] hash_reduce using target's data ptr [checkpoint] Share FullyUnconstrained [checkpoint] Backtrack on using global memoization for SEScope Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. I'll instead tackle memoization of SEScopes directly in device_domains.cc. [checkpoint] Improve back compat for homogeneous case If no host target is given but we have a unique target of kDLCPU device type then also use that for the host. Reworked to avoid global SEScopeCache. Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. So instead placed the cache in the CompilationConfig class. Since that class now has everything the device planner needs to do its job, promoted it to be an FFI-able Object, which is now in compilation_config.{h,cc}. I think we can do much better with CompilationConfig, but for now keeping it to the minimum I needed to prepare for device planning from all the executor compilation codepaths. Adds SEScope (Storage/Execution Scope) for use as new unit of planning in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

[checkpoint] pretty printing fixes [checkpoint] Don't dup devices in executable, more unit tests [checkpoint] woops, left target str debug in Added Target::ToDebugString() so I can see the hosts since they were giving me a lot of trouble. [checkpoint] more pretty printing hackery, interpreter respects host devices Also try harder to integrate the existing target->host mechanism into CompilationConfig. [checkpoint] Almost working again - Unit test setup distinguishes CPU for prims from CPU for host. - Get pretty printing to use the SEScopeNode ReprPrinter. - Allow host and primitive to have same device types. test_dynamic_input failing [checkpoint] rebase [checkpoint] fix merge [checkpoint] lint [checkpoint] rebase [checkpoint] Fixed stray use of kDLCPU in vm/profiler/vm.cc [checkpoint] lint trivia [checkpoint] fix unit tests [checkpoint] device planner unit tests passing again [checkpoint] Switch over to new CompilerOptions [checkpoint] include [checkpoint] Almost working again Need to move the SEScopeCache into CompilationConfig and pass that into DeviceDomains instead of just the Vector<Target>. Then the host_se_scope can be memoized so that direct uses of that scope downstream will match up with se_scopes already established by PlanDevices. Sigh. [checkpoint] Use cache in device domains. [checkpoint] more moves [checkpoint] lints [checkpoint] Fix merge with VM profiling changes. [checkpoint] trivial [checkpoint] rebase fix [checkpoint] More unit tests. Getting ready to fork out SEScope changes alone. [checkpoint] lints [checkpoint] All plan devices unit tests pass [checkpoint] First unit test passes [checkpoint] Another go at target management This at least centralizes all the hackery. Compiles. [commit] Start to rollback resolving to target in planner. Better is to do it as stand alone pass I think. Besides it doesn't work with the structural test for expected output. [checkpoint] Almost have first unit test going. About to merge Michalis' changes. target_host is still a mess. Starting to eliminate target_map. [checkpoint] Cleanup VM device matching [checkpoint] Compiles [checkpoint] First sweep replacing DLDeviceType with SEScope VM still not done. [checkpoint] Expose CompilationConfig ctor in py [checkpoint] CompilationConfig is nullable for default ctor [checkpoint] Don't use target:: namespace [checkpoint] Promote CompilationConfig to be FFI-friendly Object Also rework to never mix the host_target into the 'primitive' targets. [checkpoint] ResolveSEScope on CompilationConfig [checkpoint] hash_reduce using target's data ptr [checkpoint] Share FullyUnconstrained [checkpoint] Backtrack on using global memoization for SEScope Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. I'll instead tackle memoization of SEScopes directly in device_domains.cc. [checkpoint] Improve back compat for homogeneous case If no host target is given but we have a unique target of kDLCPU device type then also use that for the host. Reworked to avoid global SEScopeCache. Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. So instead placed the cache in the CompilationConfig class. Since that class now has everything the device planner needs to do its job, promoted it to be an FFI-able Object, which is now in compilation_config.{h,cc}. I think we can do much better with CompilationConfig, but for now keeping it to the minimum I needed to prepare for device planning from all the executor compilation codepaths. Adds SEScope (Storage/Execution Scope) for use as new unit of planning in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

[checkpoint] bad rebase [checkpoint] pretty printing fixes [checkpoint] Don't dup devices in executable, more unit tests [checkpoint] woops, left target str debug in Added Target::ToDebugString() so I can see the hosts since they were giving me a lot of trouble. [checkpoint] more pretty printing hackery, interpreter respects host devices Also try harder to integrate the existing target->host mechanism into CompilationConfig. [checkpoint] Almost working again - Unit test setup distinguishes CPU for prims from CPU for host. - Get pretty printing to use the SEScopeNode ReprPrinter. - Allow host and primitive to have same device types. test_dynamic_input failing [checkpoint] rebase [checkpoint] fix merge [checkpoint] lint [checkpoint] rebase [checkpoint] Fixed stray use of kDLCPU in vm/profiler/vm.cc [checkpoint] lint trivia [checkpoint] fix unit tests [checkpoint] device planner unit tests passing again [checkpoint] Switch over to new CompilerOptions [checkpoint] include [checkpoint] Almost working again Need to move the SEScopeCache into CompilationConfig and pass that into DeviceDomains instead of just the Vector<Target>. Then the host_se_scope can be memoized so that direct uses of that scope downstream will match up with se_scopes already established by PlanDevices. Sigh. [checkpoint] Use cache in device domains. [checkpoint] more moves [checkpoint] lints [checkpoint] Fix merge with VM profiling changes. [checkpoint] trivial [checkpoint] rebase fix [checkpoint] More unit tests. Getting ready to fork out SEScope changes alone. [checkpoint] lints [checkpoint] All plan devices unit tests pass [checkpoint] First unit test passes [checkpoint] Another go at target management This at least centralizes all the hackery. Compiles. [commit] Start to rollback resolving to target in planner. Better is to do it as stand alone pass I think. Besides it doesn't work with the structural test for expected output. [checkpoint] Almost have first unit test going. About to merge Michalis' changes. target_host is still a mess. Starting to eliminate target_map. [checkpoint] Cleanup VM device matching [checkpoint] Compiles [checkpoint] First sweep replacing DLDeviceType with SEScope VM still not done. [checkpoint] Expose CompilationConfig ctor in py [checkpoint] CompilationConfig is nullable for default ctor [checkpoint] Don't use target:: namespace [checkpoint] Promote CompilationConfig to be FFI-friendly Object Also rework to never mix the host_target into the 'primitive' targets. [checkpoint] ResolveSEScope on CompilationConfig [checkpoint] hash_reduce using target's data ptr [checkpoint] Share FullyUnconstrained [checkpoint] Backtrack on using global memoization for SEScope Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. I'll instead tackle memoization of SEScopes directly in device_domains.cc. [checkpoint] Improve back compat for homogeneous case If no host target is given but we have a unique target of kDLCPU device type then also use that for the host. Reworked to avoid global SEScopeCache. Realized while working through unit tests in the sequel that it's reasonable for folks to call build multiple times with distinct Target objects, in which case the global cache would grow without bound. So instead placed the cache in the CompilationConfig class. Since that class now has everything the device planner needs to do its job, promoted it to be an FFI-able Object, which is now in compilation_config.{h,cc}. I think we can do much better with CompilationConfig, but for now keeping it to the minimum I needed to prepare for device planning from all the executor compilation codepaths. Adds SEScope (Storage/Execution Scope) for use as new unit of planning in 'device' planning This is the first step in apache/tvm-rfcs#38 to bring devices and targets together when doing device planning. I've gone ahead and also included a memory scope in this object since we will also need to propagate memory scopes across Relay expressions once this basic preparation is in place. In the meantime that field will be left as "". Once device planning works in units of SEScopes it will be possible to directly read off the device and target for any Relay sub-expression without the need for TargetMaps ort the construction of default Targets. SEScopes also support 'Join' and 'Default' operations needed when constraint solving in the device planner. You can see those in use in my scratchpad branch: https://github.com/mbs-octoml/mbs-tvm/tree/mbs-scopes This PR also brings some duplicated and the ad-hoc 'default target' handling logic together into a CompilationConfig class. (Again, see the scratchpad branch for how that will end up being used). I've placed that next to SEScope since it's main purpose is to a) establish the default SEScope for primitive ops b) establish the SEScope for the 'host' c) feed a definitive vector of Targets into device planning so it can resolve all "on_device" and "device_copy" device references to their full SEScope form.

areusch

@mbs-octoml some comments/clarifications, can you also remove/address the boilerplate?

areusch · 2021-11-08T16:44:52Z