[PROFILING] Profiling over RPC

include/tvm/runtime/profiling.h

@@ -251,6 +251,12 @@ class Report : public ObjectRef {
  */
  explicit Report(Array<Map<String, ObjectRef>> calls,
  Map<String, Map<String, ObjectRef>> device_metrics);
+
+ /*! Deserialize a Report from a JSON object. Needed for sending the report over RPC.
+ * \param json Serialized json report from `ReportNode::AsJSON`.
+ * \returns A Report.
+ */
+ static Report FromJSON(String json);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Report, ObjectRef, ReportNode);
 };
 

python/tvm/contrib/debugger/debug_executor.py

@@ -25,6 +25,7 @@
 from tvm._ffi.base import string_types
 from tvm.contrib import graph_executor
 from . import debug_result
+from ...runtime.profiling import Report
 
 _DUMP_ROOT_PREFIX = "tvmdbg_"
 _DUMP_PATH_PREFIX = "_tvmdbg_"
@@ -102,6 +103,7 @@ def __init__(self, module, device, graph_json_str, dump_root):
  self._execute_node = module["execute_node"]
  self._get_node_output = module["get_node_output"]
  self._profile = module["profile"]
+ self._profile_rpc = module["profile_rpc"]
  graph_executor.GraphModule.__init__(self, module)
  self._create_debug_env(graph_json_str, device)
 
@@ -274,7 +276,7 @@ def profile(self, collectors=None, **input_dict):
  Parameters
  ----------
  collectors : Optional[Sequence[MetricCollector]]
- Extra metrics to collect.
+ Extra metrics to collect. If profiling over RPC, collectors must be `None`.
 
  input_dict : dict of str to NDArray
  List of input values to be feed to
@@ -284,10 +286,13 @@ def profile(self, collectors=None, **input_dict):
  timing_results : str
  Per-operator and whole graph timing results in a table format.
  """
- collectors = [] if collectors is None else collectors
  if input_dict:
  self.set_input(**input_dict)
 
+ if self.module.type_key == "rpc":
+ # We cannot serialize MetricCollectors over RPC
+ assert collectors is None, "Profiling with collectors is not supported over RPC"
+ return Report.from_json(self._profile_rpc())
  return self._profile(collectors)
 
  def exit(self):

python/tvm/runtime/profiler_vm.py

@@ -22,7 +22,9 @@
 """
 import warnings
 from tvm.runtime import _ffi_api
+from tvm.rpc import base as rpc_base
 from . import vm
+from .profiling import Report
 
 
 def enabled():
@@ -35,10 +37,18 @@ class VirtualMachineProfiler(vm.VirtualMachine):
 
  def __init__(self, exe, device, memory_cfg=None):
  super(VirtualMachineProfiler, self).__init__(exe, device, memory_cfg)
- self.module = _ffi_api._VirtualMachineDebug(exe.module)
+
+ # Make sure the constructor of the VM module is on the proper device
+ # Remote devices have device_type of their actual device_type + RPC_SESS_MASK
+ if device.device_type >= rpc_base.RPC_SESS_MASK:
+ self.module = device._rpc_sess.get_function("runtime._VirtualMachineDebug")(exe)
+ else:
+ self.module = _ffi_api._VirtualMachineDebug(exe.module)
+
  self._init = self.module["init"]
  self._invoke = self.module["invoke"]
  self._profile = self.module["profile"]
+ self._profile_rpc = self.module["profile_rpc"]
  self._set_input = self.module["set_input"]
  self._setup_device(device, memory_cfg)
 
@@ -59,7 +69,7 @@ def profile(self, *args, func_name="main", collectors=None, **kwargs):
  The name of the function.
 
  collectors : Optional[Sequence[MetricCollector]]
- Extra metrics to collect.
+ Extra metrics to collect. If profiling over RPC, collectors must be `None`.
 
  args : list[tvm.runtime.NDArray] or list[np.ndarray]
  The arguments to the function.
@@ -72,7 +82,10 @@ def profile(self, *args, func_name="main", collectors=None, **kwargs):
  timing_results : str
  Overall and per-op timing results formatted in a table.
  """
- collectors = [] if collectors is None else collectors
  if args or kwargs:
  self.set_input(func_name, *args, **kwargs)
+ if self.module.type_key == "rpc":
+ # We cannot serialize MetricCollectors over RPC
+ assert collectors is None, "Profiling with collectors is not supported over RPC"
+ return Report.from_json(self._profile_rpc(func_name))
  return self._profile(func_name, collectors)

python/tvm/runtime/profiling/__init__.py

@@ -104,6 +104,22 @@ def json(self):
  """
  return _ffi_api.AsJSON(self)
 
+ @classmethod
+ def from_json(cls, s):
+ """Deserialize a report from JSON.
+
+ Parameters
+ ----------
+ s : str
+ Report serialize via :py:meth:`json`.
+
+ Returns
+ -------
+ report : Report
+ The deserialized report.
+ """
+ return _ffi_api.FromJSON(s)
+
 
 @_ffi.register_object("runtime.profiling.MetricCollector")
 class MetricCollector(Object):

python/tvm/runtime/vm.py

@@ -48,6 +48,8 @@ def _convert(arg, cargs):
  dtype = "int32" if isinstance(arg, (_base.integer_types, bool)) else "float32"
  value = tvm.nd.array(np.array(arg, dtype=dtype), device=tvm.cpu(0))
  cargs.append(value)
+ elif isinstance(arg, str):
+ cargs.append(arg)
  else:
  raise TypeError("Unsupported type: %s" % (type(arg)))
 

src/runtime/graph_executor/debug/graph_executor_debug.cc

@@ -365,8 +365,23 @@ PackedFunc GraphExecutorDebug::GetFunction(const std::string& name,
  } else if (name == "profile") {
  return TypedPackedFunc<profiling::Report(Array<profiling::MetricCollector>)>(
  [sptr_to_self, this](Array<profiling::MetricCollector> collectors) {
- return this->Profile(collectors);
+ // We cannot send Arrays over rpc, so in order to support profiling
+ // on remotes, we accept a nullptr for collectors.
+ if (collectors.defined()) {
+ return this->Profile(collectors);
+ } else {
+ return this->Profile({});
+ }
  });
+ } else if (name == "profile_rpc") {
+ // We cannot return a Report over RPC because TMV RPC mechanism only
+ // supports a subset of Object classes. Instead we serialize it on the
+ // remote (here) and deserialize it on the other end.
+ return TypedPackedFunc<std::string()>([sptr_to_self, this]() {
+ PackedFunc profile = GetFunction("profile", sptr_to_self);
+ profiling::Report report = profile(Array<profiling::MetricCollector>());
+ return report->AsJSON();
+ });
  } else {
  return GraphExecutor::GetFunction(name, sptr_to_self);
  }

src/runtime/profiling.cc

@@ -22,6 +22,7 @@
  * \brief Runtime profiling including timers.
  */
 
+#include <dmlc/json.h>
 #include <tvm/ir/expr.h>
 #include <tvm/runtime/c_backend_api.h>
 #include <tvm/runtime/packed_func.h>
@@ -231,7 +232,9 @@ String ReportNode::AsCSV() const {
 namespace {
 void print_metric(std::ostream& os, ObjectRef o) {
  if (o.as<StringObj>()) {
- os << "\"" << Downcast<String>(o) << "\"";
+ os << "{\"string\":"
+ << "\"" << Downcast<String>(o) << "\""
+ << "}";
  } else if (const CountNode* n = o.as<CountNode>()) {
  os << "{\"count\":" << std::to_string(n->value) << "}";
  } else if (const DurationNode* n = o.as<DurationNode>()) {
@@ -540,6 +543,72 @@ Report::Report(Array<Map<String, ObjectRef>> calls,
  data_ = std::move(node);
 }
 
+Map<String, ObjectRef> parse_metrics(dmlc::JSONReader* reader) {
+ reader->BeginObject();
+ std::string metric_name, metric_value_name;
+ Map<String, ObjectRef> metrics;
+ while (reader->NextObjectItem(&metric_name)) {
+ ObjectRef o;
+ reader->BeginObject();
+ reader->NextObjectItem(&metric_value_name);
+ if (metric_value_name == "microseconds") {
+ double microseconds;
+ reader->Read(&microseconds);
+ o = ObjectRef(make_object<DurationNode>(microseconds));
+ } else if (metric_value_name == "percent") {
+ double percent;
+ reader->Read(&percent);
+ o = ObjectRef(make_object<PercentNode>(percent));
+ } else if (metric_value_name == "count") {
+ int64_t count;
+ reader->Read(&count);
+ o = ObjectRef(make_object<CountNode>(count));
+ } else if (metric_value_name == "string") {
+ std::string s;
+ reader->Read(&s);
+ o = String(s);
+ } else {
+ LOG(FATAL) << "Cannot parse metric of type " << metric_value_name
+ << " valid types are microseconds, percent, count.";
+ }
+ metrics.Set(metric_name, o);
+ // Necessary to make sure that the parser hits the end of the object.
+ ICHECK(!reader->NextObjectItem(&metric_value_name));
+ // EndObject does not exist, leaving this here for clarity
+ // reader.EndObject();
+ }
+ // reader.EndObject();
+ return metrics;
+}
+
+Report Report::FromJSON(String json) {
+ std::stringstream input(json.operator std::string());
+ dmlc::JSONReader reader(&input);
+ std::string key;
+ Array<Map<String, ObjectRef>> calls;
+ Map<String, Map<String, ObjectRef>> device_metrics;
+
+ reader.BeginObject();
+ while (reader.NextObjectItem(&key)) {
+ if (key == "calls") {
+ reader.BeginArray();
+ while (reader.NextArrayItem()) {
+ calls.push_back(parse_metrics(&reader));
+ }
+ // reader.EndArray();
+ } else if (key == "device_metrics") {
+ reader.BeginObject();
+ std::string device_name;
+ while (reader.NextObjectItem(&device_name)) {
+ device_metrics.Set(device_name, parse_metrics(&reader));
+ }
+ // reader.EndObject();
+ }
+ }
+
+ return Report(calls, device_metrics);
+}
+
 TVM_REGISTER_OBJECT_TYPE(DurationNode);
 TVM_REGISTER_OBJECT_TYPE(PercentNode);
 TVM_REGISTER_OBJECT_TYPE(CountNode);
@@ -551,6 +620,7 @@ TVM_REGISTER_GLOBAL("runtime.profiling.AsCSV").set_body_typed([](Report n) { ret
 TVM_REGISTER_GLOBAL("runtime.profiling.AsJSON").set_body_typed([](Report n) {
  return n->AsJSON();
 });
+TVM_REGISTER_GLOBAL("runtime.profiling.FromJSON").set_body_typed(Report::FromJSON);
 TVM_REGISTER_GLOBAL("runtime.profiling.DeviceWrapper").set_body_typed([](Device dev) {
  return DeviceWrapper(dev);
 });

src/runtime/vm/profiler/vm.cc

@@ -52,8 +52,14 @@ PackedFunc VirtualMachineDebug::GetFunction(const std::string& name,
  }
  }
 
- std::vector<profiling::MetricCollector> cs(collectors.begin(), collectors.end());
- prof_ = profiling::Profiler(devices, cs);
+ // We cannot send Arrays over rpc, so in order to support profiling
+ // on remotes, we accept a nullptr for collectors.
+ if (collectors.defined()) {
+ std::vector<profiling::MetricCollector> cs(collectors.begin(), collectors.end());
+ prof_ = profiling::Profiler(devices, cs);
+ } else {
+ prof_ = profiling::Profiler(devices, {});
+ }
 
  auto invoke = VirtualMachine::GetFunction("invoke", sptr_to_self);
  // warmup
@@ -68,6 +74,15 @@ PackedFunc VirtualMachineDebug::GetFunction(const std::string& name,
  prof_ = dmlc::optional<profiling::Profiler>(); // releases hardware counters
  return report;
  });
+ } else if (name == "profile_rpc") {
+ // We cannot return a Report over RPC because TMV RPC mechanism only
+ // supports a subset of Object classes. Instead we serialize it on the
+ // remote (here) and deserialize it on the other end.
+ return TypedPackedFunc<std::string(std::string)>([sptr_to_self, this](std::string arg_name) {
+ PackedFunc profile = GetFunction("profile", sptr_to_self);
+ profiling::Report report = profile(arg_name, Array<profiling::MetricCollector>());
+ return report->AsJSON();
+ });
  } else {
  return VirtualMachine::GetFunction(name, sptr_to_self);
  }

tests/python/unittest/test_runtime_profiling.py

@@ -26,6 +26,9 @@
 from tvm import relay
 from tvm.relay.testing import mlp
 from tvm.contrib.debugger import debug_executor
+from tvm import rpc
+from tvm.contrib import utils
+from tvm.runtime.profiling import Report
 
 
 def read_csv(report):
@@ -102,7 +105,6 @@ def test_papi(target, dev):
  func_name="main",
  collectors=[tvm.runtime.profiling.PAPIMetricCollector({dev: [metric]})],
  )
- print(report)
  assert metric in str(report)
 
  csv = read_csv(report)
@@ -126,10 +128,60 @@ def test_json():
  assert "microseconds" in parsed["calls"][0]["Duration (us)"]
  assert len(parsed["calls"]) > 0
  for call in parsed["calls"]:
- assert isinstance(call["Name"], str)
+ assert isinstance(call["Name"]["string"], str)
  assert isinstance(call["Count"]["count"], int)
  assert isinstance(call["Duration (us)"]["microseconds"], float)
 
 
+@tvm.testing.requires_llvm
+def test_rpc_vm():
+ server = rpc.Server(key="profiling")
+ remote = rpc.connect("127.0.0.1", server.port, key="profiling")
+
+ mod, params = mlp.get_workload(1)
+ exe = relay.vm.compile(mod, "llvm", params=params)
+ temp = utils.tempdir()
+ path = temp.relpath("lib.tar")
+ exe.mod.export_library(path)
+ remote.upload(path)
+ rexec = remote.load_module("lib.tar")
+ vm = profiler_vm.VirtualMachineProfiler(rexec, remote.cpu())
+ report = vm.profile(tvm.nd.array(np.ones((1, 1, 28, 28), dtype="float32"), device=remote.cpu()))
+ assert len(report.calls) > 0
+
+
+def test_rpc_graph():
+ server = rpc.Server(key="profiling")
+ remote = rpc.connect("127.0.0.1", server.port, key="profiling")
+
+ mod, params = mlp.get_workload(1)
+ exe = relay.build(mod, "llvm", params=params)
+ temp = utils.tempdir()
+ path = temp.relpath("lib.tar")
+ exe.export_library(path)
+ remote.upload(path)
+ rexec = remote.load_module("lib.tar")
+
+ gr = debug_executor.create(exe.get_graph_json(), rexec, remote.cpu())
+
+ data = np.random.rand(1, 1, 28, 28).astype("float32")
+ report = gr.profile(data=data)
+ assert len(report.calls) > 0
+
+
+def test_report_serialization():
+ mod, params = mlp.get_workload(1)
+
+ exe = relay.vm.compile(mod, "llvm", params=params)
+ vm = profiler_vm.VirtualMachineProfiler(exe, tvm.cpu())
+
+ data = np.random.rand(1, 1, 28, 28).astype("float32")
+ report = vm.profile(data, func_name="main")
+
+ report2 = Report.from_json(report.json())
+ # equality on reports compares pointers, so we compare the printed results instead.
+ assert str(report) == str(report2)
+
+
 if __name__ == "__main__":
  test_papi("llvm", tvm.cpu())