mbs-octoml commented on a change in pull request #9077:
URL: https://github.com/apache/tvm/pull/9077#discussion_r714937453
##########
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'd like to be able to freely write "TODO(issues/1234)" so everything is
cross linked. That would help anyone picking up the issue get a head start on
where to look in the code, as well as just remind us all where things are
likely to shift. That got a bit tangled up with questions about how we should
reconcile GitHub and our internal OctoML Jira issues, and the granularity of
GitHub issues, but I think we've decided to just switch to GitHub. I'm guessing
you have the same dilemma at Arm? Anyway these "TODO(mbs)" is just me leaving
enough cookie crumbs so I can go back and do it properly.
--
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]
