From 93f4977c785fda05e9bd30b1254852e9716e116b Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 4 Dec 2020 15:45:41 +0800
Subject: [PATCH] feat(mge/imperative): add thread name

GitOrigin-RevId: 61eece55b8f4008882532fa4b50a94f1164efee4
---
 imperative/python/src/helper.h          | 4 ++++
 imperative/src/impl/interpreter_impl.h  | 7 ++++++-
 imperative/src/impl/physical_tensor.cpp | 3 +++
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/imperative/python/src/helper.h b/imperative/python/src/helper.h
index 05650a3d..05aaf922 100644
--- a/imperative/python/src/helper.h
+++ b/imperative/python/src/helper.h
@@ -69,6 +69,10 @@ class PyTaskDipatcher {
             pybind11::gil_scoped_acquire _;
             f();
         }
+
+        void on_async_queue_worker_thread_start() override {
+            mgb::sys::set_thread_name("py_task_worker");
+        }
     };
     Queue queue;
     bool finalized = false;
diff --git a/imperative/src/impl/interpreter_impl.h b/imperative/src/impl/interpreter_impl.h
index 6c894e95..5a8c974a 100644
--- a/imperative/src/impl/interpreter_impl.h
+++ b/imperative/src/impl/interpreter_impl.h
@@ -143,10 +143,15 @@ private:
     size_t m_enable_evict = 0;
 
     struct WorkQueue : AsyncQueueSC<Command, WorkQueue> {
-        WorkQueue(ChannelImpl* owner) : m_owner(owner) {}
+        WorkQueue(ChannelImpl* owner) : m_owner(owner) {
+            sys::set_thread_name("interpreter");
+        }
         void process_one_task(Command& cmd) {
             m_owner->process_one_task(cmd);
         }
+        void on_async_queue_worker_thread_start() override {
+               sys::set_thread_name("worker");
+        }
     private:
         ChannelImpl* m_owner;
     } m_worker;
diff --git a/imperative/src/impl/physical_tensor.cpp b/imperative/src/impl/physical_tensor.cpp
index 94a4219d..13faadab 100644
--- a/imperative/src/impl/physical_tensor.cpp
+++ b/imperative/src/impl/physical_tensor.cpp
@@ -44,6 +44,9 @@ class AsyncReleaser : public CompNodeDepedentObject {
             std::this_thread::sleep_for(1us);
             add_task(std::move(param));
         }
+        void on_async_queue_worker_thread_start() override {
+            sys::set_thread_name("releaser");
+        }
     };
     Waiter m_waiter{this};
 
-- 
GitLab