parent b77d3e0

author Lily Orth-Smith <lilyorthsmith@gmail.com> 1641339403 -0800
committer Lily Orth-Smith <lilyorthsmith@gmail.com> 1641949451 -0800

Make function result virtual_device_ stored in virtual_device_ field

include/tvm/ir/expr.h

@@ -182,6 +182,8 @@ class RelayExprNode : public BaseExprNode {
    *
    * Set to VirtualDevice::FullyUnconstrained by default.
    *
+   * Set to VirtualDevice::FullyUnconstrained by default.
+   *
    * \note Unfortunately, the type of virtual_device_ needs to be ObjectRef to avoid a circular
    * import.
    */

include/tvm/ir/function.h

@@ -200,16 +200,6 @@ constexpr const char* kGlobalSymbol = "global_symbol";
  */
 constexpr const char* kParamVirtualDevice = "param_virtual_devices";
 
-/*!
- * \brief The \p VirtualDevice which will hold the function result.
- *
- * Only supported on Relay \p Functions. Generally added by the \p PlanDevices pass, but
- * may be included as an annotation on user programs.
- *
- * Type: VirtualDevice
- */
-constexpr const char* kResultVirtualDevice = "result_virtual_device";
-
 }  // namespace attr
 }  // namespace tvm
 #endif  // TVM_IR_FUNCTION_H_

src/printer/relay_text_printer.cc

@@ -458,6 +458,7 @@ Doc RelayTextPrinter::PrintFunc(const Doc& prefix, const relay::Function& fn) {
   if (fn->ret_type.defined()) {
     doc << "-> " << Print(fn->ret_type) << " ";
   }
+  doc << "Virtual Device: " << Print(fn->virtual_device()) << " \n";
   doc << PrintBody(fn->body);
   return doc;
 }

src/relay/backend/graph_plan_memory.cc

@@ -279,6 +279,8 @@ class StorageAllocator : public StorageAllocaBaseVisitor {
       smap.Set(GetRef<Expr>(kv.first), storage_info);
     }
     // Either all or none of the nodes should be annotated.
+    VLOG(1) << "num annotated nodes / num_nodes: " << num_annotated_nodes << " / " << num_nodes
+            << std::endl;
     if (num_annotated_nodes != 0 && num_annotated_nodes != num_nodes) {
       LOG(FATAL) << num_annotated_nodes << " out of " << num_nodes
                  << "expressions are assigned with virtual device types. Either all "

src/relay/op/memory/on_device.cc

@@ -144,9 +144,11 @@ OnDeviceProps GetOnDeviceProps(const Expr& expr) {
 
 Function FunctionOnDevice(Function function, Array<VirtualDevice> param_virtual_devices,
                           VirtualDevice result_virtual_device) {
-  return WithAttrs(std::move(function),
-                   {{tvm::attr::kParamVirtualDevice, std::move(param_virtual_devices)},
-                    {tvm::attr::kResultVirtualDevice, std::move(result_virtual_device)}});
+  auto func = WithAttrs(
+      WithFields(std::move(function), {}, {}, {}, {}, {}, std::move(result_virtual_device)),
+      {{tvm::attr::kParamVirtualDevice, std::move(param_virtual_devices)}});
+  VLOG(1) << "Annotated func: " << PrettyPrint(func);
+  return func;
 }
 
 TVM_REGISTER_GLOBAL("relay.op.annotation._make.FunctionOnDevice").set_body_typed(FunctionOnDevice);
@@ -166,8 +168,7 @@ Function MaybeFunctionOnDevice(Function function, Array<VirtualDevice> param_vir
 }
 
 VirtualDevice GetFunctionResultVirtualDevice(const FunctionNode* function_node) {
-  auto opt_virtual_device = function_node->GetAttr<VirtualDevice>(tvm::attr::kResultVirtualDevice);
-  return opt_virtual_device.value_or(VirtualDevice::FullyUnconstrained());
+  return function_node->virtual_device();
 }
 
 VirtualDevice GetFunctionParamVirtualDevice(const FunctionNode* function_node, size_t i) {