[rfc] Reduce number of coin flips in RecordFunction (pytorch#40758)

Summary:
Pull Request resolved: pytorch#40758

Currently we flip a coin for each sampled callback each time
we run RecordFunction, this PR is an attempt to skip most of the coin
flips (for the low-probability observers) and keep the distribution
close to the original one

Test Plan:
CI and record_function_benchmark
```
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 30108 us.
Time per iteration (1x1): 1496.78 us.
Time per iteration (16x16): 2142.46 us.
Pure RecordFunction runtime of 10000000 iterations 687929 us, number of callback invocations: 978
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 19051 us.
Time per iteration (1x1): 1581.89 us.
Time per iteration (16x16): 2195.67 us.
Pure RecordFunction runtime of 10000000 iterations 682402 us, number of callback invocations: 1023
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 18715 us.
Time per iteration (1x1): 1566.11 us.
Time per iteration (16x16): 2131.17 us.
Pure RecordFunction runtime of 10000000 iterations 693571 us, number of callback invocations: 963
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$

(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 18814 us.
Time per iteration (1x1): 1536.2 us.
Time per iteration (16x16): 1985.82 us.
Pure RecordFunction runtime of 10000000 iterations 944959 us, number of callback invocations: 1015
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 18278 us.
Time per iteration (1x1): 1526.32 us.
Time per iteration (16x16): 2093.77 us.
Pure RecordFunction runtime of 10000000 iterations 985307 us, number of callback invocations: 1013
(python_venv) iliacher@devgpu151:~/local/pytorch  (reduce_coin_flops)$ ./build/bin/record_function_benchmark
Warmup time: 18545 us.
Time per iteration (1x1): 1524.65 us.
Time per iteration (16x16): 2080 us.
Pure RecordFunction runtime of 10000000 iterations 952835 us, number of callback invocations: 1048
```

Reviewed By: dzhulgakov

Differential Revision: D22320879

Pulled By: ilia-cher

fbshipit-source-id: 2193f07d2f7625814fe7bc3cc85ba4092fe036bc

Author: Ilia Cherniavskii

aten/src/ATen/record_function.cpp

@@ -212,16 +212,58 @@ inline CallbackManager& manager() {
   return _manager;
 }
 
-} // namespace
+// Low probability constant
+const double kLowProb = 0.001;
+thread_local int tries_left_ = 0;
 
-/* static */
-double RecordFunctionCallback::sample_zero_one() {
+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
+
+bool RecordFunctionCallback::shouldRun(RecordScope scope) const {
+  // first check whether this callback is interested in
+  // the given scope type
+  if (!checkScope(scope)) {
+    return false;
+  }
+  // if we have registered should_run_ function, use it
+  if (should_run_) {
+    return should_run_(*this);
+  }
+  // otherwise potentially do the uniform sampling
+  if (sampling_prob_ != 1.0) {
+    // model the low probability events as events happening
+    // with prob. kLowProb followed by another sampling with
+    // prob. (sampling_prob_ / kLowProb), then replace the coin
+    // flip for kLowProb with a thread local number of tries tries_left_
+    // sampled from the geometric distribution
+    if (sampling_prob_ < kLowProb) {
+      if (tries_left_ == 0) {
+        tries_left_ = sample_geometric();
+        return (sample_zero_one() < sampling_prob_ / kLowProb);
+      } else {
+        --tries_left_;
+        return false;
+      }
+    } else {
+      return (sample_zero_one() < sampling_prob_);
+    }
+  }
+  return true;
+}
+
 RecordFunctionCallbacks _getTLSCallbacks() {
   return sorted_tls_callbacks_;
 }

aten/src/ATen/record_function.h

@@ -303,23 +303,8 @@ class TORCH_API RecordFunctionCallback {
     return end_;
   }
 
-  // whether this callbacks should run in the given scope
-  inline bool shouldRun(RecordScope scope) const {
-    // first check whether this callback is interested in
-    // the given scope type
-    if (!checkScope(scope)) {
-      return false;
-    }
-    // if we have registered should_run_ function, use it
-    if (should_run_) {
-      return should_run_(*this);
-    }
-    // otherwise potentially do the uniform sampling
-    if (sampling_prob_ != 1.0) {
-      return (sample_zero_one() < sampling_prob_);
-    }
-    return true;
-  }
+  // whether the callbacks should run in the given scope
+  bool shouldRun(RecordScope scope) const;
 
  private:
   std::function<void(const RecordFunction&)> start_;
@@ -329,8 +314,6 @@ class TORCH_API RecordFunctionCallback {
   bool needs_ids_ = false;
   double sampling_prob_ = 1.0;
   std::array<bool, static_cast<size_t>(RecordScope::NUM_SCOPES)> scopes_ = {};
-
-  static double sample_zero_one();
 };
 
 // Using macro to minimize inputs copies,

binaries/record_function_benchmark.cc

@@ -8,32 +8,31 @@
 #include <ctime>
 
 C10_DEFINE_int(iter, 100, "Number of iterations");
-C10_DEFINE_int(warmup_iter, 10, "Number of warmup iterations")
+C10_DEFINE_int(warmup_iter, 10, "Number of warmup iterations");
+C10_DEFINE_int(rec_fn_iter, 10e6,
+    "Number of iterations for the pure RecordFunction benchmark");
 
 namespace {
 const int kInnerIter = 100;
 const int kNumSampledCb = 2;
 const int kTensorSize = 16;
 const int kSmallTensorSize = 1;
 const float kSampingProb = 0.1;
-}
 
+const float kLowSamplingProb = 0.0001;
+}
 
-void setupCallbacks() {
+void setupBenchmarkCallbacks() {
   // non-sampled callback
   at::addGlobalCallback(at::RecordFunctionCallback(
-      [&](const at::RecordFunction& fn) {
-        return true;
-      },
+      [&](const at::RecordFunction& fn) {},
       [](const at::RecordFunction&) {})
     .needsInputs(true));
 
   // sampled
   for (auto idx = 0; idx < kNumSampledCb; ++idx) {
     at::addGlobalCallback(at::RecordFunctionCallback(
-        [](const at::RecordFunction& fn) {
-          return true;
-        },
+        [](const at::RecordFunction& fn) {},
         [](const at::RecordFunction&) {})
       .needsInputs(true)
       .samplingProb(kSampingProb)
@@ -61,7 +60,8 @@ int main(int argc, char** argv) {
     return -1;
   }
 
-  setupCallbacks();
+  at::enableRecordFunction();
+  setupBenchmarkCallbacks();
 
   auto duration = runBench(kSmallTensorSize, FLAGS_warmup_iter);
   std::cout << "Warmup time: " << duration << " us." << std::endl;
@@ -76,5 +76,32 @@ int main(int argc, char** argv) {
               << " us." << std::endl;
   }
 
+  at::clearCallbacks();
+
+  int cb_count = 0;
+  at::addGlobalCallback(at::RecordFunctionCallback(
+      [&](const at::RecordFunction& fn) {
+        ++cb_count;
+      },
+      [](const at::RecordFunction&) {})
+    .needsInputs(true)
+    .samplingProb(kLowSamplingProb)
+  );
+
+  typedef std::chrono::high_resolution_clock clock;
+  typedef std::chrono::microseconds us;
+  std::chrono::time_point<clock> start_time = clock::now();
+  for (auto n = 0; n < FLAGS_rec_fn_iter; ++n) {
+    RECORD_USER_SCOPE("test");
+  }
+  duration = static_cast<float>(
+      std::chrono::duration_cast<us>(clock::now() - start_time).count());
+  std::cout << "Pure RecordFunction runtime of " << FLAGS_rec_fn_iter
+            << " iterations " << duration
+            << " us, number of callback invocations: " << cb_count
+            << ", expected number: ~" << (int)(FLAGS_rec_fn_iter * kLowSamplingProb)
+            << " invocations" << std::endl;
+
+  at::clearCallbacks();
   return 0;
 }