[GitHub] [tvm] mbs-octoml commented on a change in pull request #9077: [Relay] Prepare for merging context_analysis.cc and device_annotation.cc

[GitBox <gi...@apache.org>]

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: commits-unsubscribe@tvm.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org