Back out "Revert D23494065: Refactor CallbackManager as a friend class of RecordFunction." (#44699)

louisfeng · facebook-github-bot · commit 09f2c6a94c7b · 2020-09-18T11:08:27.000-07:00
Summary: Pull Request resolved: #44699 Original commit changeset: 3b1ec928e3db Previous revert (D23698861) was on the wrong diff stack. Backing out the revert. Test Plan: Passed unit tests and previously landed. Reviewed By: mruberry Differential Revision: D23702258 fbshipit-source-id: 5c3e197bca412f454db5a7e86251ec85faf621c1
diff --git a/aten/src/ATen/record_function.cpp b/aten/src/ATen/record_function.cpp
@@ -24,6 +24,34 @@ thread_local RecordFunctionCallbacks sorted_tls_callbacks_;
 
 std::atomic<int64_t> defaultNodeId(-1);
 
+// Enumerates thread ids logically;
+// note: std::this_thread::get_id may return potentially
+// reused thread id
+std::atomic<uint64_t> next_thread_id_ {0};
+thread_local uint64_t current_thread_id_ = 0;
+
+thread_local bool tls_record_function_enabled_ = true;
+
+// Low probability constant
+const double kLowProb = 0.001;
+thread_local int tries_left_ = 0;
+
+int sample_geometric() {
+  static thread_local auto gen =
+      std::make_unique<std::mt19937>(std::random_device()());
+  std::geometric_distribution<int> dist(kLowProb);
+  return dist(*gen);
+}
+
+double sample_zero_one() {
+  static thread_local auto gen =
+      std::make_unique<std::mt19937>(std::random_device()());
+  std::uniform_real_distribution<double> dist(0.0, 1.0);
+  return dist(*gen);
+}
+
+} // namespace
+
 class CallbackManager {
  public:
   CallbackHandle addThreadLocalCallback(RecordFunctionCallback cb) {
@@ -212,37 +240,12 @@ class CallbackManager {
   RecordFunctionCallbacks sorted_global_callbacks_;
 };
 
-// Enumerates thread ids logically;
-// note: std::this_thread::get_id may return potentially
-// reused thread id
-std::atomic<uint64_t> next_thread_id_ {0};
-thread_local uint64_t current_thread_id_ = 0;
-
-inline CallbackManager& manager() {
-  static CallbackManager _manager;
-  return _manager;
-}
-
-thread_local bool tls_record_function_enabled_ = true;
-
-// Low probability constant
-const double kLowProb = 0.001;
-thread_local int tries_left_ = 0;
-
-int sample_geometric() {
-  static thread_local auto gen =
-      std::make_unique<std::mt19937>(std::random_device()());
-  std::geometric_distribution<int> dist(kLowProb);
-  return dist(*gen);
-}
-
-double sample_zero_one() {
-  static thread_local auto gen =
-      std::make_unique<std::mt19937>(std::random_device()());
-  std::uniform_real_distribution<double> dist(0.0, 1.0);
-  return dist(*gen);
-}
-
+namespace {
+  // Keeping this static manager local.
+  CallbackManager& manager() {
+    static CallbackManager _manager;
+    return _manager;
+  }
 } // namespace
 
 bool RecordFunctionCallback::shouldRun(RecordScope scope) const {
diff --git a/aten/src/ATen/record_function.h b/aten/src/ATen/record_function.h
@@ -168,9 +168,17 @@ struct TORCH_API RecordFunction {
     handle_ = handle;
   }
 
+  // Whether this RecordFunction runs any callbacks
+  bool active = false;
+  // Whether any of the picked callbacks require inputs
+  bool needs_inputs = false;
+
+ private:
+  // Allows the modification of some internal states for callbacks.
+  friend class CallbackManager;
+
   // Used internally to keep track of thread local and global callbacks
   // that were picked to run; must be sorted;
-  // public because of anonymous "friend" class
   CallbackHandles sorted_active_tls_handles_;
   CallbackHandles sorted_active_global_handles_;
 
@@ -182,17 +190,11 @@ struct TORCH_API RecordFunction {
   // callbacks.
   ObserverContextList global_ctx_;
 
-  // Whether this RecordFunction runs any callbacks
-  bool active = false;
-  /// Whether any of the picked callbacks require inputs
-  bool needs_inputs = false;
-
   // In cases when RecordFunction might be active but we chose not to
   // use the observers (e.g. operator is not observed), this boolean
   // flag is used to check whether the start callbacks were called
   bool called_start_callbacks_ = false;
 
- private:
   StringView name_;
   int64_t sequence_nr_ = -1;
   std::vector<c10::IValue> inputs_;
