fix(dtensor): support torch dispatch mode exclude

JsBlueCat · JsBlueCat · commit 3bdfde29cfd4 · 2024-04-04T18:07:08.000+08:00
support torch dispatch python mode to tls excluded which will cause dispatch hang
diff --git a/src/cc/torchdistx/deferred_init.cc b/src/cc/torchdistx/deferred_init.cc
@@ -944,7 +944,7 @@ void runDeferredInitHandler(const OperatorHandle& op, DispatchKeySet key_set, St
 }
 
 void enableDeferredInitHandler(bool value) noexcept {
-  c10::impl::tls_set_dispatch_key_included(DispatchKey::DeferredInit, value);
+  c10::impl::tls_set_dispatch_key_excluded(DispatchKey::DeferredInit, value);
 }
 
 bool isDeferredInitEnabled() noexcept {

Original file line number	Diff line number	Diff line change
`@@ -944,7 +944,7 @@ void runDeferredInitHandler(const OperatorHandle& op, DispatchKeySet key_set, St`
`944`	`944`	`}`
`945`	`945`
`946`	`946`	`void enableDeferredInitHandler(bool value) noexcept {`
`947`		`- c10::impl::tls_set_dispatch_key_included(DispatchKey::DeferredInit, value);`
	`947`	`+ c10::impl::tls_set_dispatch_key_excluded(DispatchKey::DeferredInit, value);`
`948`	`948`	`}`
`949`	`949`
`950`	`950`	`bool isDeferredInitEnabled() noexcept {`