[Profiler] Limit calls to recordThreadInfo (#74888)

Taylor Robie · pytorchmergebot · commit d92fd2dd0396 · 2022-03-29T17:38:29.000Z
Summary: Pull Request resolved: #74888 So far as I can tell, `recordThreadInfo` only needs to be called once per thread. Once we have thread local subqueues we can easily manage this by simply calling it in the subqueue constructor. Test Plan: The effect on single threaded overhead is pretty minimal, but it improves stress test overhead from ~6.1 us to ~1.4us since we're no contending over the lock in Kineto. Reviewed By: chaekit Differential Revision: D34811694 fbshipit-source-id: da1047f7ae43af048773610a0f250fa514c67989 (cherry picked from commit 9a5b926)
diff --git a/torch/csrc/autograd/profiler_kineto.cpp b/torch/csrc/autograd/profiler_kineto.cpp
@@ -631,7 +631,6 @@ void pushProfilingCallbacks(const std::unordered_set<at::RecordScope>& scopes) {
             }
 
             torch::profiler::impl::kineto::popCorrelationId();
-            torch::profiler::impl::kineto::recordThreadInfo();
           })
           .needsInputs(registration_state_ptr->config().report_input_shapes)
           .scopes(scopes));
@@ -667,8 +666,6 @@ void reportBackendEventToActiveKinetoProfiler(
     ctx_ptr->dtypes = inputTypes(fn);
   }
   */
-
-  torch::profiler::impl::kineto::recordThreadInfo();
 }
 
 void prepareProfiler(
diff --git a/torch/csrc/profiler/collection.cpp b/torch/csrc/profiler/collection.cpp
@@ -141,7 +141,9 @@ uint64_t Result::correlation_id() const {
 ThreadLocalSubqueue::ThreadLocalSubqueue(
     const uint64_t tid,
     const ProfilerConfig& config)
-    : tid_{tid}, config_{config}, kineto_info_{kineto::kineto_ids()} {}
+    : tid_{tid}, config_{config}, kineto_info_{kineto::kineto_ids()} {
+  torch::profiler::impl::kineto::recordThreadInfo();
+}
 
 std::unique_ptr<KinetoObserverContext> ThreadLocalSubqueue::begin_op(
     const at::RecordFunction& fn,

Original file line number	Diff line number	Diff line change
`@@ -631,7 +631,6 @@ void pushProfilingCallbacks(const std::unordered_set<at::RecordScope>& scopes) {`
`631`	`631`	`}`
`632`	`632`
`633`	`633`	`torch::profiler::impl::kineto::popCorrelationId();`
`634`		`- torch::profiler::impl::kineto::recordThreadInfo();`
`635`	`634`	`})`
`636`	`635`	`.needsInputs(registration_state_ptr->config().report_input_shapes)`
`637`	`636`	`.scopes(scopes));`
`@@ -667,8 +666,6 @@ void reportBackendEventToActiveKinetoProfiler(`
`667`	`666`	`ctx_ptr->dtypes = inputTypes(fn);`
`668`	`667`	`}`
`669`	`668`	`*/`
`670`		`-`
`671`		`- torch::profiler::impl::kineto::recordThreadInfo();`
`672`	`669`	`}`
`673`	`670`
`674`	`671`	`void prepareProfiler(`