Progress on parsing & printing

include/tvm/target/virtual_device.h

@@ -162,6 +162,7 @@ using MemoryScope = String;
  *
  * These operations are needed during device planning.
  */
+
 class VirtualDeviceNode : public AttrsNode<VirtualDeviceNode> {
  private:
   /*!
@@ -361,6 +362,13 @@ class VirtualDeviceCache {
   std::unordered_set<VirtualDevice, StructuralHash, StructuralEqual> cache_;
 };
 
+/*! brief The attribute key for the virtual device. This key will be promoted to first class on
+ * functions. For use in the parser and printer only.
+ *
+ * Type: VirtualDevice
+ */
+constexpr const char* kVirtualDevice = "result_virtual_device";
+
 }  // namespace tvm
 
 #endif  //  TVM_TARGET_VIRTUAL_DEVICE_H_

src/parser/parser.cc

@@ -31,6 +31,7 @@
 #include <tvm/runtime/logging.h>
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/registry.h>
+#include <tvm/target/virtual_device.h>
 
 #include <fstream>
 
@@ -230,13 +231,6 @@ GlobalTypeVar AddOrGet(InternTable<GlobalTypeVar>* table, const std::string& nam
   }
 }
 
-/*! brief The attribute key for the virtual device. This key will be promoted to first class on
- * functions.
- *
- * Type: VirtualDevice
- */
-constexpr const char* kVirtualDevice = "result_virtual_device";
-
 /*! \brief The parser class is the main interface to the parser.
  * the parser is not currently exposed beyond this .cc file.
  *
@@ -1145,8 +1139,7 @@ class Parser {
       // TODO(@jroesch): attributes should never be null, they should always be empty.
       if (raw_attrs.size()) {
         // Promote kVirtualDevice to first-class
-        String vid_key = kVirtualDevice;
-        if (raw_attrs.count(vid_key)) {
+        if (raw_attrs.count(kVirtualDevice)) {
           ObjectRef vid = raw_attrs.at(kVirtualDevice);
           ICHECK(vid.as<VirtualDeviceNode>())
               << "Expected the " << kVirtualDevice << " to have type VirtualDeviceNode, but got "

src/printer/relay_text_printer.cc

@@ -445,11 +445,16 @@ Doc RelayTextPrinter::PrintFunc(const Doc& prefix, const relay::Function& fn) {
   for (const Doc& d : PrintDictAttrs(fn->attrs)) {
     params.push_back(d);
   }
+  if (fn->virtual_device() != VirtualDevice::FullyUnconstrained()) {
+    Doc vid_doc;
+    vid_doc << kVirtualDevice << "=" << Print(fn->virtual_device());
+    params.push_back(vid_doc);
+  }
   doc << Doc::Concat(params) << ") ";
   if (fn->ret_type.defined()) {
     doc << "-> " << Print(fn->ret_type) << " ";
   }
-  doc << "Virtual Device: " << Print(fn->virtual_device()) << " \n";
+
   doc << PrintBody(fn->body);
   return doc;
 }