Revert D34604068: [PyTorch] [Model Tracer] Use c10::Synchronized<T> for kernel dtype tracer

Test Plan: revert-hammer

Differential Revision:
D34604068 (pytorch@6fd6fe0)

Original commit changeset: 1ec50ada8112

Original Phabricator Diff: D34604068 (pytorch@6fd6fe0)

fbshipit-source-id: 8b80bfd947c96108306e4472505c1af62c4fe8cb
(cherry picked from commit 62d6d29)

Author: suo

torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.cpp

@@ -15,9 +15,8 @@ KernelDTypeTracer::KernelDTypeTracer() {
     std::string kernel_tag = name.substr(0, dollar_pos);
     std::string dtype = name.substr(dollar_pos + 1);
 
-    getCalledKernelTags().withLock([&](kernel_tags_type& kernel_tags) {
-      kernel_tags[kernel_tag].insert(dtype);
-    });
+    std::lock_guard<std::mutex> guard(getMutex());
+    getCalledKernelTags()[kernel_tag].insert(dtype);
     return nullptr;
   };
 
@@ -26,12 +25,16 @@ KernelDTypeTracer::KernelDTypeTracer() {
           .scopes({at::RecordScope::KERNEL_FUNCTION_DTYPE}));
 }
 
-c10::Synchronized<KernelDTypeTracer::kernel_tags_type>& KernelDTypeTracer::
-    getCalledKernelTags() {
-  static c10::Synchronized<kernel_tags_type> called_kernel_tags;
+KernelDTypeTracer::kernel_tags_type& KernelDTypeTracer::getCalledKernelTags() {
+  static kernel_tags_type called_kernel_tags;
   return called_kernel_tags;
 }
 
+std::mutex& KernelDTypeTracer::getMutex() {
+  static std::mutex m;
+  return m;
+}
+
 } // namespace mobile
 } // namespace jit
 } // namespace torch

torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h

@@ -1,8 +1,8 @@
 #pragma once
 
 #include <ATen/record_function.h>
-#include <c10/util/Synchronized.h>
 #include <map>
+#include <mutex>
 #include <set>
 #include <string>
 
@@ -30,7 +30,9 @@ struct KernelDTypeTracer final {
   typedef std::map<std::string, std::set<std::string>> kernel_tags_type;
 
   KernelDTypeTracer();
-  static c10::Synchronized<kernel_tags_type>& getCalledKernelTags();
+  static kernel_tags_type& getCalledKernelTags();
+  /* Protect concurrent writes into the map. */
+  static std::mutex& getMutex();
 
   ~KernelDTypeTracer() {
     at::removeCallback(handle_);

torch/csrc/jit/mobile/model_tracer/TracerRunner.cpp

@@ -295,10 +295,12 @@ TracerResult trace_run(const std::string& input_module_path) {
 
   recordCustomClassesFromOpSchemas(root_ops, traced_operators, loaded_classes);
 
-  kdtype_tracer.getCalledKernelTags().withLock(
-      [&](KernelDTypeTracer::kernel_tags_type& kernel_tags) {
-        called_kernel_tags.insert(kernel_tags.begin(), kernel_tags.end());
-      });
+  {
+    std::lock_guard<std::mutex> guard(KernelDTypeTracer::getMutex());
+    called_kernel_tags.insert(
+        kdtype_tracer.getCalledKernelTags().begin(),
+        kdtype_tracer.getCalledKernelTags().end());
+  }
 
   traced_operators.insert(
       always_included_traced_ops.begin(), always_included_traced_ops.end());