Introduce CUDA OpenXLA fallback. (#7318)

pytorch · Jul 3, 2024 · c782e0d · c782e0d
1 parent 172097b
commit c782e0d
Show file tree

Hide file tree

Showing 14 changed files with 1,037 additions and 513 deletions.
diff --git a/BUILD b/BUILD
@@ -46,6 +46,22 @@ 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,

diff --git a/setup.py b/setup.py
@@ -286,6 +286,7 @@ def run(self):
     packages=find_packages(include=['torch_xla*']),
     ext_modules=[
         BazelExtension('//:_XLAC.so'),
+        BazelExtension('//:_XLAC_cuda_functions.so'),
     ],
     install_requires=[
         'absl-py>=1.0.0',

diff --git a/test/cpp/BUILD b/test/cpp/BUILD
@@ -53,6 +53,7 @@ ptxla_cc_test(
         ":cpp_test_util",
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
+        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -63,6 +64,7 @@ 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",
     ],
@@ -77,6 +79,7 @@ ptxla_cc_test(
         "//torch_xla/csrc/runtime:runtime",
         "//torch_xla/csrc/runtime:debug_macros",
         "//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",
@@ -95,6 +98,7 @@ ptxla_cc_test(
         ":cpp_test_util",
         ":torch_xla_test",
         "//torch_xla/csrc:tensor",
+        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -119,6 +123,7 @@ ptxla_cc_test(
         "//torch_xla/csrc/runtime:env_vars",
         "//torch_xla/csrc/runtime:sys_util",
         "//torch_xla/csrc:tensor",
+        "//torch_xla/csrc:aten_cuda_functions",
         "@com_google_googletest//:gtest_main",
         "@xla//xla:xla_data_proto_cc",
         "@tsl//tsl/profiler/utils:session_manager",
@@ -137,6 +142,7 @@ 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",
         ],