feat(profiler): custom event support device

GitOrigin-RevId: 4709c44b39

imperative/src/impl/interpreter/interpreter_impl.cpp

@@ -78,6 +78,20 @@ void imperative_log_profile(const char* message){
     imperative_log_profile_end(message);
 }
 
+SYMBOL_EXPORT
+void imperative_log_profile_begin(const char* message, const char* device) {
+    auto comp_node = CompNode::load(device);
+    MGB_RECORD_EVENT(CustomEvent, std::string{message}, {}, comp_node);
+    MGB_RECORD_EVENT(RecordDeviceEvent, EventPool::with_timer().alloc_shared(comp_node));
+}
+
+SYMBOL_EXPORT
+void imperative_log_profile_end(const char* message, const char* device) {
+    auto comp_node = CompNode::load(device);
+    MGB_RECORD_EVENT(RecordDeviceEvent, EventPool::with_timer().alloc_shared(comp_node));
+    MGB_RECORD_EVENT(CustomFinishEvent, std::string{message}, {}, comp_node);
+}
+
 }
 
 std::thread::id ChannelImpl::get_worker_tid() {

imperative/src/impl/profiler/chrome_timeline.cpp

@@ -350,8 +350,14 @@ struct ChromeTimelineEventVisitor: EventVisitor<ChromeTimelineEventVisitor> {
             new_host_event("StopProfile", 'E');
         } else if constexpr (std::is_same_v<TEvent, CustomEvent>) {
             new_host_event(event.title, 'B');
+            if (event.device.valid()) {
+                new_device_event(event.title, 'B', event.device);
+            }
         } else if constexpr (std::is_same_v<TEvent, CustomFinishEvent>) {
             new_host_event(event.title, 'E');
+            if (event.device.valid()) {
+                new_device_event(event.title, 'E', event.device);
+            }
         } else if constexpr (std::is_same_v<TEvent, AutoEvictEvent>) {
             new_host_event("AutoEvict", 'B');
         } else if constexpr (std::is_same_v<TEvent, AutoEvictFinishEvent>) {
@@ -378,12 +384,21 @@ struct ChromeTimelineEventVisitor: EventVisitor<ChromeTimelineEventVisitor> {
     }
 
     void name_threads(Profiler::thread_dict_t thread_dict) {
-        for (auto&& [tid, tname]: thread_dict) {
+        for (auto&& host: host_threads()) {
+            if (thread_dict.count(host)) {
+                trace_events.new_event()
+                        .name("thread_name")
+                        .pid('M')
+                        .tid(to_tid(host))
+                        .arg("name", thread_dict.at(host));
+            }
+        }
+        for (auto&& device: devices()) {
             trace_events.new_event()
                     .name("thread_name")
                     .pid('M')
-                    .tid(to_tid(tid))
-                    .arg("name", tname);
+                    .tid(to_tid(device))
+                    .arg("name", device.to_string_logical());
         }
     }
 };

imperative/src/impl/profiler/events.h

@@ -179,6 +179,7 @@ DEF_DUR_EVENT(AutoEvict, {});
 DEF_DUR_EVENT(Custom, {
     std::string title;
     std::string content;
+    CompNode device;
 });
 
 DEF_EVENT(RecordDevice, {

imperative/src/impl/profiler/states.h

@@ -218,6 +218,22 @@ protected:
         return m_device_tid_table.at(device);
     }
 
+    SmallVector<std::thread::id> host_threads() {
+        SmallVector<std::thread::id> host_threads;
+        for (auto&& [host, _]: m_host_tid_table) {
+            host_threads.push_back(host);
+        }
+        return host_threads;
+    }
+
+    SmallVector<CompNode> devices() {
+        SmallVector<CompNode> devices;
+        for (auto&& [device, _]: m_device_tid_table) {
+            devices.push_back(device);
+        }
+        return devices;
+    }
+
     void inc_counter(const char* key, int64_t delta) {
         if (!m_counter_table.count(key)) {
             m_counter_table[key] = 0;