Remove OpenXLA CUDA fallback and _XLAC_cuda_functions.so extension. (#9581)

This PR removes the OpenXLA CUDA fallback implementation, and also the
`_XLA_cuda_functions.so` extension, completely. Starting on this PR, the
fallback shall be run only on CPU.

**Key Changes:**

- Remove _aten_cuda_functions.cpp_ and _aten_cuda_functions.h_
- Remove the OpenXLA CUDA fallback functions from _aten_fallback.cpp_
- Remove the `_XLAC_cuda_functions.so` library from _BUILD_
- Remove the Python `_XLAC_cuda_functions.so`  extension from _setup.py_
- Remove the conditional loading of `_XLAC_cuda_functions.so` Python
extension from _torch_xla/__init__.py_

Author: ysiraichi

BUILD

@@ -46,22 +46,6 @@ cc_binary(
     ]),
 )
 
-cc_binary(
-    name = "_XLAC_cuda_functions.so",
-    copts = [
-        "-fopenmp",
-        "-fPIC",
-    ],
-    linkopts = [
-        "-Wl,-soname,_XLAC_cuda_functions.so",
-    ],
-    linkshared = 1,
-    visibility = ["//visibility:public"],
-    deps = [
-        "//torch_xla/csrc:aten_cuda_functions",
-    ],
-)
-
 test_suite(
     name = "cpp_tests",
     # testonly = True,

configuration.yaml

@@ -397,8 +397,3 @@ variables:
           your code.
       type: bool
       default_value: false
-    XLA_FALLBACK_CPU:
-      description:
-        - Forces CPU OpenXLA fallback. By default, PyTorch/XLA will run any operation
-          that doesn't have a lowering using PyTorch CUDA as fallback. Setting this
-          flag will force PyTorch/XLA to use PyTorch CPU as fallback.

setup.py

@@ -455,7 +455,6 @@ def _get_jax_install_requirements():
     package_dir=package_dir_mapping,
     ext_modules=[
         BazelExtension('//:_XLAC.so'),
-        BazelExtension('//:_XLAC_cuda_functions.so'),
     ],
     install_requires=[
         'absl-py>=1.0.0',

test/cpp/BUILD

@@ -54,7 +54,6 @@ ptxla_cc_test(
         ":cpp_test_util",
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -65,7 +64,6 @@ ptxla_cc_test(
     deps = [
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
         "@xla//xla:shape_util",
     ],
@@ -81,7 +79,6 @@ ptxla_cc_test(
         "//torch_xla/csrc/runtime:debug_macros",
         "//torch_xla/csrc:status",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "//torch_xla/csrc:thread_pool",
         "@com_google_absl//absl/synchronization",
         "@com_google_googletest//:gtest_main",
@@ -100,7 +97,6 @@ ptxla_cc_test(
         ":cpp_test_util",
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -127,7 +123,6 @@ ptxla_cc_test(
         "//torch_xla/csrc/runtime:sys_util",
         "//torch_xla/csrc:status",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
         "@xla//xla:xla_data_proto_cc",
         "@xla//xla/tsl/profiler/utils:session_manager",
@@ -146,7 +141,6 @@ ptxla_cc_test(
             ":torch_xla_test",
             "//torch_xla/csrc/runtime:metrics",
             "//torch_xla/csrc:tensor",
-            "//torch_xla/csrc:aten_cuda_functions",
             "@com_google_googletest//:gtest_main",
             "@xla//xla:permutation_util",
         ],
@@ -212,7 +206,6 @@ ptxla_cc_test(
         ":cpp_test_util",
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
-        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
     ],
-)
+)

torch_xla/__init__.py

@@ -7,16 +7,6 @@
 
 import torch
 
-if not torch.cuda.is_available():
-  # Load _XLAC_cuda_functions to RTLD_GLOBAL, so that it can be used by _XLAC.
-  flags = sys.getdlopenflags()
-  sys.setdlopenflags(flags | os.RTLD_NOW | os.RTLD_GLOBAL)
-
-  import _XLAC_cuda_functions
-
-  # Then, restore the original flags.
-  sys.setdlopenflags(flags)
-
 import _XLAC
 from ._internal import tpu
 from .version import __version__

torch_xla/csrc/BUILD

@@ -77,7 +77,6 @@ ptxla_cc_library(
     hdrs = [
         "aten_autograd_ops.h",
         "aten_fallback.h",
-        "aten_cuda_functions.h",
         "aten_xla_bridge.h",
         "batch_norm.h",
         "convert_ops.h",
@@ -364,16 +363,6 @@ ptxla_cc_library(
     ],
 )
 
-ptxla_cc_library(
-    name = "aten_cuda_functions",
-    srcs = ["aten_cuda_functions.cpp"],
-    hdrs = ["aten_cuda_functions.h"],
-    deps = [
-        "@local_config_python//:python_headers",
-        "@pybind11//:pybind11_embed",
-    ],
-)
-
 cc_library(
     name = "status",
     srcs = ["status.cpp"],