Some more utilities

Author: zou3519

functorch/csrc/DynamicLayer.cpp

@@ -445,6 +445,12 @@ static DispatchKeySet keysForEnteringDynamicLayer(DispatchKey key) {
   }
 }
 
+static void dump_local_tls() {
+  auto tls = c10::impl::tls_local_dispatch_key_set();
+  std::cout << "[Local Include] " << tls.included_ << std::endl;
+  std::cout << "[Local Exclude] " << tls.excluded_ << std::endl;
+}
+
 static DispatchKeySet keysToExcludeWhenEnteringDynamicLayer(DispatchKey key) {
   DispatchKeySet exclude = all_dynlayer_keyset;
   exclude = exclude.remove(kDynamicLayerBackModeKey);
@@ -457,6 +463,7 @@ void dynamicLayerFrontFallback(const c10::OperatorHandle& op, torch::jit::Stack*
 #ifdef HAS_TORCH_SHOW_DISPATCH_TRACE
   if (c10::show_dispatch_trace_enabled()) {
     std::cout << dynamicLayerStack << std::endl;
+    dump_local_tls();
   }
 #endif
   if (dynamicLayerStack.size() == 0) {
@@ -498,6 +505,12 @@ void dynamicLayerFrontFallback(const c10::OperatorHandle& op, torch::jit::Stack*
   local_keyset.included_ = local_keyset.included_ | hacky_include;
   ForceLocalDispatchKeySet guard(local_keyset);
 
+#ifdef HAS_TORCH_SHOW_DISPATCH_TRACE
+  if (c10::show_dispatch_trace_enabled()) {
+    dump_local_tls();
+  }
+#endif
+
   // Re-dispatch
   op.callBoxed(stack);
 }
@@ -594,6 +607,12 @@ void dynamicLayerBackFallback(const c10::OperatorHandle& op, torch::jit::Stack*
   c10::impl::_force_tls_local_dispatch_key_set(keyset);
   setDynamicLayerFrontBackKeysIncluded(true);
 
+#ifdef HAS_TORCH_SHOW_DISPATCH_TRACE
+  if (c10::show_dispatch_trace_enabled()) {
+    dump_local_tls();
+  }
+#endif
+
   // Re-dispatch
   if (cur_key == DispatchKey::Autograd && *prev_grad_mode == false) {
     // See NOTE [grad and vjp interaction with no_grad]

functorch/csrc/init.cpp

@@ -233,6 +233,12 @@ static void dump_dls() {
   std::cout << getDynamicLayerStack() << std::endl;
 }
 
+static void dump_local_tls() {
+  auto tls = c10::impl::tls_local_dispatch_key_set();
+  std::cout << "[Local Include] " << tls.included_ << std::endl;
+  std::cout << "[Local Exclude] " << tls.excluded_ << std::endl;
+}
+
 } // namespace functorch
 }
 
@@ -268,6 +274,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("tls_set_vmap_excluded", &at::functorch::tls_set_vmap_excluded);
   m.def("tls_set_is_included", &at::functorch::tls_set_is_included);
   m.def("dump_dls", &at::functorch::dump_dls);
+  m.def("dump_local_tls", &at::functorch::dump_local_tls);
   at::functorch::initPointwiseOperatorCompileCacheBindings(m.ptr());
   at::functorch::initCompileCacheBindings(m.ptr());
   initDispatchBindings(m.ptr());