[PyTorch] Make tls_local_dispatch_key_set inlineable (reapply)

FLAGS_disable_variable_dispatch had to go, but it looks like the only user was some benchmarks anyway.

Differential Revision: [D25547962](https://our.internmc.facebook.com/intern/diff/D25547962/)

**NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D25547962/)!

ghstack-source-id: 118624599
Pull Request resolved: #49412

c10/core/impl/LocalDispatchKeySet.cpp

@@ -5,10 +5,6 @@
 namespace c10 {
 namespace impl {
 
-C10_DEFINE_bool(disable_variable_dispatch, false, "This flag forcibly disables the Variable code paths from executing, which currently breaks profiling in the process.");
-
-namespace {
-
 /// In the CAFFE2_FB_LIMITED_MOBILE_CAPABILITY build setting,
 /// thread_local is not supported.
 #ifndef CAFFE2_FB_LIMITED_MOBILE_CAPABILITY
@@ -18,26 +14,10 @@ thread_local PODLocalDispatchKeySet raw_local_dispatch_key_set;
 
 #else // defined(CAFFE2_FB_LIMITED_MOBILE_CAPABILITY)
 
-static PODLocalDispatchKeySet raw_local_dispatch_key_set;
+PODLocalDispatchKeySet raw_local_dispatch_key_set;
 
 #endif
 
-} // anonymous namespace
-
-LocalDispatchKeySet tls_local_dispatch_key_set() {
-  // Hack until variable performance is fixed
-  //
-  // ezyang: I'm pretty unhappy about this implementation, it looks wrong
-  // to me, as it seems to be performing a mutation on
-  // raw_local_dispatch_key_set.  I can't conveniently test the correct
-  // version though...
-  if (FLAGS_disable_variable_dispatch) {
-    raw_local_dispatch_key_set.set_excluded(
-      raw_local_dispatch_key_set.excluded() | autograd_dispatch_keyset);
-  }
-  return raw_local_dispatch_key_set;
-}
-
 void _force_tls_local_dispatch_key_set(LocalDispatchKeySet key_set) {
   raw_local_dispatch_key_set = PODLocalDispatchKeySet {
     key_set.included_.raw_repr(),

c10/core/impl/LocalDispatchKeySet.h

@@ -23,8 +23,6 @@
 namespace c10 {
 namespace impl {
 
-C10_DECLARE_bool(disable_variable_dispatch);
-
 // POD version of LocalDispatchKeySet.  Declared here just so that
 // we can put it in the guards.
 struct C10_API PODLocalDispatchKeySet {
@@ -54,7 +52,19 @@ struct C10_API LocalDispatchKeySet {
   DispatchKeySet excluded_;
 };
 
-C10_API LocalDispatchKeySet tls_local_dispatch_key_set();
+/// In the CAFFE2_FB_LIMITED_MOBILE_CAPABILITY build setting,
+/// thread_local is not supported.
+#ifndef CAFFE2_FB_LIMITED_MOBILE_CAPABILITY
+  extern C10_API thread_local PODLocalDispatchKeySet raw_local_dispatch_key_set;
+#else // defined(CAFFE2_FB_LIMITED_MOBILE_CAPABILITY)
+  extern C10_API PODLocalDispatchKeySet raw_local_dispatch_key_set;
+#endif
+
+inline C10_API LocalDispatchKeySet tls_local_dispatch_key_set() {
+  // Don't let people fiddle with the thread_local directly just
+  // because they include this header.
+  return raw_local_dispatch_key_set;
+}
 
 // Internal, use ThreadLocalStateGuard
 C10_API void _force_tls_local_dispatch_key_set(LocalDispatchKeySet key_set);