[PyTorch] remove convenience RecordFunctionCallback interface (#48620)

swolchok · facebook-github-bot · commit 900aa4ee9757 · 2020-12-14T20:03:15.000-08:00
Summary: Pull Request resolved: #48620 In preparation for storing bare function pointer (8 bytes) instead of std::function (32 bytes). ghstack-source-id: 118568242 Test Plan: CI Reviewed By: ezyang Differential Revision: D25132183 fbshipit-source-id: 3790cfb5d98479a46cf665b14eb0041a872c13da
diff --git a/android/pytorch_android/src/main/cpp/pytorch_jni_jit.cpp b/android/pytorch_android/src/main/cpp/pytorch_jni_jit.cpp
@@ -90,13 +90,13 @@ class PytorchJni : public facebook::jni::HybridClass<PytorchJni> {
 #endif
 
 #ifdef TRACE_ENABLED
-  static bool onFunctionEnter(
+  static std::unique_ptr<at::ObserverContext> onFunctionEnter(
       const at::RecordFunction& fn) {
     Trace::beginSection(fn.name().str());
-    return true;
+    return nullptr;
   }
 
-  static void onFunctionExit(const at::RecordFunction&) {
+  static void onFunctionExit(const at::RecordFunction&, at::ObserverContext*) {
     Trace::endSection();
   }
 #endif
diff --git a/aten/src/ATen/record_function.h b/aten/src/ATen/record_function.h
@@ -316,17 +316,6 @@ class TORCH_API RecordFunctionCallback {
     scopes_.fill(true);
   }
 
-  // This interface is for observers that do not pass an ObserverContext object
-  // between start and end callbacks.
-  explicit RecordFunctionCallback(
-      std::function<void(const RecordFunction&)> start,
-      std::function<void(const RecordFunction&)> end =
-        [](const RecordFunction&) {}):
-      start_{[start](const RecordFunction& rf) { start(rf); return nullptr; }},
-      end_{[end](const RecordFunction& rf, ObserverContext*) { end(rf); }} {
-    scopes_.fill(true);
-  }
-
   RecordFunctionCallback& needsInputs(bool needs_inputs) {
     needs_inputs_ = needs_inputs;
     return *this;
diff --git a/binaries/record_function_benchmark.cc b/binaries/record_function_benchmark.cc
@@ -19,11 +19,11 @@ const float kLowSamplingProb = 0.0001;
 
 void addTestCallback(
     double sampling_prob = 1.0,
-    std::function<void(const at::RecordFunction&)> fn =
-        [](const at::RecordFunction&) {}) {
+    std::function<std::unique_ptr<at::ObserverContext>(const at::RecordFunction&)> fn =
+        [](const at::RecordFunction&) { return nullptr; }) {
   auto cb = at::RecordFunctionCallback(
       std::move(fn),
-      [](const at::RecordFunction&) {})
+      [](const at::RecordFunction&, at::ObserverContext*) {})
     .needsInputs(false);
   if (sampling_prob < 1.0) {
     cb.samplingProb(sampling_prob);
@@ -111,6 +111,7 @@ int main(int argc, char** argv) {
       kLowSamplingProb,
       [&](const at::RecordFunction& fn) {
         ++cb_count;
+        return nullptr;
       }
   );
 
diff --git a/test/cpp/jit/test_misc.cpp b/test/cpp/jit/test_misc.cpp
@@ -739,8 +739,8 @@ void checkScopeCallbacks() {
             std::string(fn.name().str()) == "test_user_scope") {
           found_user_scope = true;
         }
-      },
-      [](const at::RecordFunction&) {}));
+        return nullptr;
+      }));
 
   bool bad_scope = false;
   auto pushScopedCallback = [&](at::RecordScope scope, size_t& cnt) {
@@ -752,9 +752,8 @@ void checkScopeCallbacks() {
               } else {
                 bad_scope = true;
               }
-              return true;
-            },
-            [](const at::RecordFunction&) {})
+              return nullptr;
+            })
             .scopes({scope}));
   };
 
@@ -813,8 +812,8 @@ TEST(RecordFunctionTest, Basic) {
             } else if (fn.scope() == RecordScope::TORCHSCRIPT_FUNCTION) {
               ts_names.insert(fn.name().str());
             }
-          },
-          [](const RecordFunction&) {})
+            return nullptr;
+          })
           .needsInputs(true));
 
   TracedTestInputs eager_inputs, jit_inputs;
@@ -851,9 +850,8 @@ TEST(RecordFunctionTest, Basic) {
                                    if (std::string(fn.name().str()) == "test") {
                                      ++sampled_cb_ctr;
                                    }
-                                   return true;
-                                 },
-                                 [](const RecordFunction&) {})
+                                   return nullptr;
+                                 })
                                  .samplingProb(sampling_prob));
   };
 
@@ -863,9 +861,8 @@ TEST(RecordFunctionTest, Basic) {
         if (std::string(fn.name().str()) == "test") {
           ++non_sampled_cb_ctr;
         }
-        return true;
-      },
-      [](const RecordFunction&) {}));
+        return nullptr;
+      }));
 
   auto handle = setup_sampled_callback(0.5);
 
@@ -908,9 +905,8 @@ TEST(RecordFunctionTest, Basic) {
       [&fn_names, &mtx](const RecordFunction& fn) {
         std::lock_guard<std::mutex> lock(mtx);
         fn_names.push_back(fn.name().str());
-        return true;
-      },
-      [](const RecordFunction&) {}));
+        return nullptr;
+      }));
   {
     RecordFunctionGuard g1(false);
     {
@@ -934,8 +930,10 @@ TEST(RecordFunctionTest, Basic) {
   std::vector<size_t> ids;
   auto add_remove_test_add_cb = [&ids](size_t id) {
     return addGlobalCallback(RecordFunctionCallback(
-        [&ids, id](const RecordFunction& fn) { ids.push_back(id); },
-        [](const RecordFunction&) {}));
+        [&ids, id](const RecordFunction& fn) {
+          ids.push_back(id);
+          return nullptr ;
+        }));
   };
 
   auto h1 = add_remove_test_add_cb(1);
@@ -972,8 +970,7 @@ TEST(RecordFunctionTest, Basic) {
 
   ids.clear();
   addGlobalCallback(RecordFunctionCallback(
-      [&ids](const RecordFunction& fn) { ids.push_back(1); },
-      [](const RecordFunction&) {}));
+      [&ids](const RecordFunction& fn) { ids.push_back(1); return nullptr; }));
 
   { RECORD_USER_SCOPE("test"); }
 
@@ -983,8 +980,7 @@ TEST(RecordFunctionTest, Basic) {
 
   auto th = std::thread([&ids]() {
     addThreadLocalCallback(RecordFunctionCallback(
-        [&ids](const RecordFunction& fn) { ids.push_back(2); },
-        [](const RecordFunction&) {}));
+        [&ids](const RecordFunction& fn) { ids.push_back(2); return nullptr; }));
 
     { RECORD_USER_SCOPE("test_thread"); }
   });
@@ -1070,8 +1066,7 @@ TEST(RecordFunctionTest, Basic) {
   bool ran = false;
   should_run = false;
   addGlobalCallback(RecordFunctionCallback(
-                        [&ran](const RecordFunction& fn) { ran = true; },
-                        [](const RecordFunction&) {})
+                        [&ran](const RecordFunction& fn) { ran = true; return nullptr; })
                         .setShouldRun(shouldRunCallback));
 
   { RECORD_USER_SCOPE("test"); }
@@ -1093,8 +1088,8 @@ TEST(RecordFunctionTest, Basic) {
     auto handle = addThreadLocalCallback(RecordFunctionCallback(
         [&recorded_op](const RecordFunction& fn) {
           recorded_op = fn.name().str();
-        },
-        [](const RecordFunction&) {}));
+          return nullptr;
+        }));
     ThreadLocalState state;
     std::thread t_child([state]() {
       ThreadLocalStateGuard g_tls(state);
@@ -1111,16 +1106,20 @@ TEST(RecordFunctionTest, Basic) {
   bool has_ids = false;
   addGlobalCallback(
       RecordFunctionCallback(
-          [&has_ids](const RecordFunction& fn) { has_ids = fn.handle() > 0; },
-          [](const RecordFunction&) {})
+          [&has_ids](const RecordFunction& fn) {
+            has_ids = fn.handle() > 0;
+            return nullptr;
+          })
           .needsIds(true));
   { RECORD_USER_SCOPE("test"); }
   TORCH_CHECK(has_ids);
   clearCallbacks();
   has_ids = false;
   addGlobalCallback(RecordFunctionCallback(
-      [&has_ids](const RecordFunction& fn) { has_ids = fn.handle() > 0; },
-      [](const RecordFunction&) {}));
+      [&has_ids](const RecordFunction& fn) {
+        has_ids = fn.handle() > 0;
+        return nullptr;
+      }));
   { RECORD_USER_SCOPE("test"); }
   TORCH_CHECK(!has_ids);
   clearCallbacks();
@@ -1138,6 +1137,7 @@ TEST(RecordFunctionTest, OperatorNameOverload) {
                               } else {
                                 operator_names.insert("No Operator Name");
                               }
+                              return nullptr;
                             })
                             .scopes({at::RecordScope::FUNCTION}));
   auto t = torch::randn({1, 2, 3}, at::kCPU);
@@ -1209,9 +1209,8 @@ TEST(ThreadLocalDebugInfoTest, Basic) {
       [&done](const RecordFunction&) {
         checkDebugInfo(c10::DebugInfoKind::TEST_INFO, 42);
         done = true;
-        return true;
-      },
-      [](const RecordFunction&) {}));
+        return nullptr;
+      }));
   {
     c10::DebugInfoGuard guard(c10::DebugInfoKind::TEST_INFO, debug_info);
     auto t = torch::randn({1, 2, 3}, at::kCPU);
diff --git a/torch/csrc/autograd/init.cpp b/torch/csrc/autograd/init.cpp
@@ -173,8 +173,8 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject *unused) {
   });
   m.def("_set_empty_test_observer", [](bool is_global, double sampling_prob) {
     auto cb = at::RecordFunctionCallback(
-        [](const at::RecordFunction&) {},
-        [](const at::RecordFunction&) {})
+        [](const at::RecordFunction&) { return nullptr; },
+        [](const at::RecordFunction&, at::ObserverContext*) {})
       .needsInputs(true)
       .samplingProb(sampling_prob);
     if (is_global) {
diff --git a/torch/csrc/autograd/profiler_legacy.cpp b/torch/csrc/autograd/profiler_legacy.cpp
@@ -417,7 +417,7 @@ void pushProfilingCallbacksLegacy() {
       [](const at::RecordFunction& fn) {
         auto state_ptr = getProfilerTLSState();
         if (!state_ptr || state_ptr->config().state == ProfilerState::Disabled) {
-          return;
+          return nullptr;
         }
         bool record_cuda =
             state_ptr->config().state == ProfilerState::CUDA;
@@ -432,8 +432,10 @@ void pushProfilingCallbacksLegacy() {
         } else {
           state_ptr->pushRange(fn, record_cuda, msg);
         }
+
+        return nullptr;
       },
-      [](const at::RecordFunction& fn) {
+      [](const at::RecordFunction& fn, at::ObserverContext*) {
         auto state_ptr = getProfilerTLSState();
         if (!state_ptr || state_ptr->config().state == ProfilerState::Disabled) {
           return;
