electriclilies commented on pull request #9848: URL: https://github.com/apache/tvm/pull/9848#issuecomment-1018125295
@tqchen Thanks for the comment. You're right, it is a subtle difference that @mbs-octoml and I are very aware of. Currently, device planning assumes that 'stored on' is exactly the same as 'executes on'-- so `fn->body->virtual_device()` is always equivalent to `fn->virtual_device()`. This subtlety is documented here: https://github.com/apache/tvm/blob/main/src/relay/transforms/device_planner.cc#L180-L184 One downside (from an implementation standpoint) of storing the result virtual device on the function's body is that the body can be any type of Relay expression. Unfortunately, people reconstruct Relay exprs everywhere, willy-nilly. Virtual devices are not propagated through these constructors. This is why I implemented WithFields. I haven't replaced all the constructors with WithFields, but I did just go through and replace all the Function constructors with WithFields so that this would work. If we did decide to switch to `function->body->virtual_device()`, _we'd have to replace every Relay constructor with WithFields_ so that the virtual devices are properly propagated. Since `function->virtual_device()` and `function->body->virtual_device()` are always going to be the same, and there's a significant engineering overhead for storing it on the body of the function, I'd prefer to keep it where it is. If you'd like to take on replacing all the Relay constructors with WithFields, then I'd be happy to make the change, just let me know. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
