Enable PJRT plugins by default (#7268)

pytorch · Jun 13, 2024 · 7bc82c5 · 7bc82c5
1 parent 5622efc
commit 7bc82c5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 2 deletions.
diff --git a/torch_xla/__init__.py b/torch_xla/__init__.py
@@ -212,7 +212,8 @@ def _init_xla_lazy_backend():
 from .experimental import plugins
 from ._internal import neuron, xpu  # Additional built-in plugins
 
-if os.getenv('XLA_REGISTER_INSTALLED_PLUGINS') == '1':
+if os.getenv('XLA_REGISTER_INSTALLED_PLUGINS',
+             '0' if _XLAC._has_cuda_support() else '1') == '1':
   plugins.use_dynamic_plugins()
   plugins.register_installed_plugins()
 

diff --git a/torch_xla/_internal/tpu.py b/torch_xla/_internal/tpu.py
@@ -342,7 +342,9 @@ def configure_multiprocess(self, local_rank, local_world_size):
     return configure_topology(local_rank, local_world_size)
 
   def physical_chip_count(self):
-    return num_available_chips()
+    # HACK: We may reduce the number of processes we spawn depending on TPU
+    # topology settings
+    return num_local_processes()
 
   def client_create_options(self):
     return {

diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
@@ -2423,6 +2423,13 @@ void InitXlaModuleBindings(py::module m) {
           return XlaCustomCall(inputs, payload, output_shapes, output_dtypes,
                                /*is_tpu=*/true);
         });
+  m.def("_has_cuda_support", []() {
+#ifdef GOOGLE_CUDA
+    return true;
+#else
+    return false;
+#endif
+  });
   m.def("_xla_gpu_custom_call",
         [](const std::vector<at::Tensor>& inputs, const std::string& payload,
            const std::vector<std::vector<int64_t>>& output_shapes,