Mousius commented on a change in pull request #9077:
URL: https://github.com/apache/tvm/pull/9077#discussion_r715441811
##########
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
+ * not insert any "device_copy" calls to respect this annotation.
+ *
+ * This is used by the device planning pass itself when annotating the
planned program.
+ */
+ bool is_fixed;
TVM_DECLARE_ATTRS(OnDeviceAttrs, "relay.attrs.OnDeviceAttrs") {
TVM_ATTR_FIELD(device_type)
- .describe("The virutal device/context type that an expression is
annotated with.")
+ .describe("The type of the virtual device which should hold the
expression result.")
+ .set_default(0);
+ TVM_ATTR_FIELD(is_fixed)
+ .describe("If true, do not insert a \"device_copy\" call to respect
this annotation.")
+ .set_default(false);
+ }
+};
+
+/*!
+ * \brief Attributes for Relay function definitions which capture the devices
for the
+ * function parameters and result.
+ */
+struct FunctionOnDeviceAttrs : public tvm::AttrsNode<FunctionOnDeviceAttrs> {
+ constexpr static const char* kFunctionAttrsKey = "on_device";
+
+ /*! \brief Device type on which each of the function's arguments already
resides. */
+ Array<Integer> param_device_types;
+ // TODO(mbs): Replace device types with TargetDevice.
Review comment:
I think this might be the start of quite a long discussion, I've seen
quite a lot of abandoned `TODO` comments so the best approach (in my opinion)
would be either GitHub issues to track all the work or just tracking
internally. In this case there may be contention as to whether we use
`TargetDevice` or something else, so if we leave this for a prolonged period
people coming to this code may be confused. One of my favourite examples of
this is bugs which get fixed in an entirely different area leaving
`TODO(Mousius) - Fix bug` which we will never get rid of.
Having said all of that, I think these will get cleaned up pretty quickly
and it's probably just my thinking - so let's not block anything on it
:smile_cat:
--
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]
