mbs-octoml commented on a change in pull request #9038:
URL: https://github.com/apache/tvm/pull/9038#discussion_r717071790
##########
File path: include/tvm/relay/attrs/annotation.h
##########
@@ -32,14 +32,55 @@ namespace tvm {
namespace relay {
/*!
- * \brief Options for the device annotation operators.
+ * \brief Attributes for the "on_device" operator.
+ *
+ * The relay call
+ * \code
+ * on_device(expr, device_type=2)
+ * \endcode
+ * denotes that the result of \p expr should be stored on the device with \p
DLDeviceType 2
+ * (i.e. \p kDLCuda). Semantically the operator is the identity function.
*/
struct OnDeviceAttrs : public tvm::AttrsNode<OnDeviceAttrs> {
+ // TODO(mbs): Replace device types with TargetDevice.
+ /*! \brief Device type on which argument expression should be evaluated. */
int device_type;
+ /*!
+ * \brief If true, the result device must also be \p device_type and device
planning should
Review comment:
This is because the association of expression to device is a
'constraint', and by default "on_device" on constrains the inner expr and not
the overall expr. However to make plan_devices `idempotent` with the current
approach I need some way to fully specify devices.
--
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]
