tkonolige commented on a change in pull request #9553:
URL: https://github.com/apache/tvm/pull/9553#discussion_r758809276
##########
File path: src/runtime/profiling.cc
##########
@@ -677,6 +677,56 @@
TVM_REGISTER_GLOBAL("runtime.profiling.FromJSON").set_body_typed(Report::FromJSO
TVM_REGISTER_GLOBAL("runtime.profiling.DeviceWrapper").set_body_typed([](Device
dev) {
return DeviceWrapper(dev);
});
+
+PackedFunc ProfileFunction(Module mod, std::string func_name, int device_type,
int device_id,
+ Array<MetricCollector> collectors) {
+ // Module::GetFunction is not const, so this lambda has to be mutable
+ return PackedFunc([=](TVMArgs args, TVMRetValue* ret) mutable {
+ PackedFunc f = mod.GetFunction(func_name);
+ Device dev{static_cast<DLDeviceType>(device_type), device_id};
+
+ // warmup
+ for (size_t i = 0; i < 10; i++) {
+ f.CallPacked(args, ret);
+ }
+
+ for (auto& collector : collectors) {
+ collector->Init({DeviceWrapper(dev)});
+ }
+ std::vector<ObjectRef> collector_data;
+ for (auto& collector : collectors) {
+ collector_data.push_back(collector->Start(dev));
+ }
+ // TODO(tkonolige): repeated calls if the runtime is small?
+ f.CallPacked(args, ret);
+ std::unordered_map<String, ObjectRef> results;
+ for (size_t i = 0; i < collectors.size(); i++) {
+ auto r = collectors[i]->Stop(collector_data[i]);
+ // We might want to do this in a separate loop to avoid unnecessary time
Review comment:
done
--
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]
