[XLA] Switch Python bindings to use pybind11 instead of SWIG.

* Simplifies the Python code by taking advantage of pybind11's greater flexibility in exposing a Pythonic API.
* Pybind11 is also better able to handle features of modern C++. SWIG is incapable of parsing many common C++ idioms in XLA's interface, which forced us to create intermediate wrapper classes with SWIG-friendly APIs that do little other than call other XLA APIs (e.g., ComputationBuilder). In many cases, we can now elide these. In combination with the preceding point, the code becomes significantly shorter and more readable.
* Pybind11's numpy integration makes it easy to construct xla Literals that alias numpy arrays. This reduces the number of copies when transferring a numpy array to, say, a GPU to 1 instead of 2; previously a NumPy<->Literal conversion required a copy.

PiperOrigin-RevId: 241067931

Author: hawkinsp

tensorflow/compiler/xla/python/BUILD

@@ -2,18 +2,16 @@ licenses(["notice"])  # Apache 2.0
 
 package(default_visibility = ["//tensorflow:internal"])
 
-load("//tensorflow:tensorflow.bzl", "tf_py_wrap_cc")
 load("//tensorflow/core:platform/default/build_config.bzl", "pyx_library")
 load("//tensorflow/compiler/xla:xla.bzl", "xla_python_default_plugins")
+load("//tensorflow:tensorflow.bzl", "tf_pybind_extension")
 
 py_library(
     name = "xla_client",
     srcs = ["xla_client.py"],
     srcs_version = "PY2AND3",
     visibility = ["//visibility:public"],
-    deps = [
-        ":pywrap_xla",
-    ],
+    deps = [":xla_extension"],
 )
 
 pyx_library(
@@ -37,31 +35,59 @@ py_test(
 )
 
 cc_library(
-    name = "numpy_bridge",
-    srcs = ["numpy_bridge.cc"],
-    hdrs = ["numpy_bridge.h"],
+    name = "types",
+    srcs = ["types.cc"],
+    hdrs = ["types.h"],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+        "-Wno-c++98-c++11-compat",
+    ],
+    features = ["-use_header_modules"],
     deps = [
         "//tensorflow/compiler/xla:literal",
-        "//tensorflow/compiler/xla:literal_util",
         "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:status",
+        "//tensorflow/compiler/xla:status_macros",
+        "//tensorflow/compiler/xla:statusor",
+        "//tensorflow/compiler/xla:types",
         "//tensorflow/compiler/xla:xla_data_proto",
+        "//tensorflow/compiler/xla/service:device_memory_allocator",
         "//tensorflow/core:lib",
-        "//tensorflow/python:numpy_lib",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/types:span",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/types:optional",
+        "@pybind11",
     ],
 )
 
-cc_library(
-    name = "local_computation_builder",
-    srcs = ["local_computation_builder.cc"],
-    hdrs = ["local_computation_builder.h"],
+tf_pybind_extension(
+    name = "xla_extension",
+    srcs = [
+        "local_client.cc",
+        "local_client.h",
+        "xla.cc",
+    ],
+    copts = [
+        "-fexceptions",
+        "-fno-strict-aliasing",
+        "-Wno-c++98-c++11-compat",
+    ],
+    features = ["-use_header_modules"],
+    module_name = "xla_extension",
     deps = [
+        ":types",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:optional",
+        "@com_google_absl//absl/types:span",
+        "@pybind11",
+        "//third_party/python_runtime:headers",  # buildcleaner: keep
         "//tensorflow/compiler/xla:executable_run_options",
         "//tensorflow/compiler/xla:literal",
         "//tensorflow/compiler/xla:literal_util",
         "//tensorflow/compiler/xla:shape_util",
+        "//tensorflow/compiler/xla:status",
+        "//tensorflow/compiler/xla:statusor",
         "//tensorflow/compiler/xla:util",
         "//tensorflow/compiler/xla:xla_data_proto",
         "//tensorflow/compiler/xla/client:client_library",
@@ -77,84 +103,14 @@ cc_library(
         "//tensorflow/compiler/xla/service:hlo_graph_dumper",
         "//tensorflow/compiler/xla/service:platform_util",
         "//tensorflow/compiler/xla/service:shaped_buffer",
+        "//tensorflow/compiler/xla/service:cpu_plugin",
         "//tensorflow/compiler/xla/service/cpu:custom_call_target_registry",
         "//tensorflow/core:lib",
-        "//third_party/python_runtime:headers",  # buildcleaner: keep
-        "@com_google_absl//absl/memory",
-        "@com_google_absl//absl/types:span",
-    ],
-)
-
-cc_library(
-    name = "xrt",
-    srcs = ["xrt.cc"],
-    hdrs = ["xrt.h"],
-    deps = [
-        "//tensorflow/cc:cc_ops",
-        "//tensorflow/cc:client_session",
-        "//tensorflow/cc:ops",
-        "//tensorflow/cc:scope",
-        "//tensorflow/compiler/xla:literal",
-        "//tensorflow/compiler/xla:literal_util",
-        "//tensorflow/compiler/xla:shape_util",
-        "//tensorflow/compiler/xla:util",
-        "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/service:hlo_proto",
-        "//tensorflow/compiler/xla/service:platform_util",
-        "//tensorflow/compiler/xrt:xrt_proto",
-        "//tensorflow/compiler/xrt/cc:xrt_ops",
-        "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
-        "@com_google_absl//absl/memory",
-        "@com_google_absl//absl/types:span",
-    ],
-)
-
-tf_py_wrap_cc(
-    name = "pywrap_xla",
-    srcs = [
-        "xla.i",
-    ],
-    swig_includes = [
-        "local_computation_builder.i",
-        "xla_data.i",
-        "//tensorflow/python:platform/base.i",
-    ],
-    version_script = select({
-        "//tensorflow:macos": "pywrap_xla_exported_symbols.lds",
-        "//tensorflow:windows": None,
-        "//conditions:default": "pywrap_xla_version_script.lds",
-    }),
-    deps = [
-        ":local_computation_builder",
-        ":numpy_bridge",
-        "//tensorflow/compiler/xla:literal",
-        "//tensorflow/compiler/xla:shape_util",
-        "//tensorflow/compiler/xla:xla_data_proto",
-        "//tensorflow/compiler/xla/service:cpu_plugin",
+        # Do NOT remove this dependency. The XLA Python extension must not
+        # depend on any part of TensorFlow at runtime, **including**
+        # libtensorflow_framework.so. The XLA module is deployed self-contained
+        # without any TF dependencies as "jaxlib" on Pypi, and "jaxlib" does
+        # not require Tensorflow.
+        "//tensorflow/core:lib_internal_impl",  # buildcleaner: keep
     ] + xla_python_default_plugins(),
 )
-
-tf_py_wrap_cc(
-    name = "pywrap_xrt",
-    srcs = [
-        "xrt.i",
-    ],
-    swig_includes = [
-        "xla_data.i",
-        "//tensorflow/python:platform/base.i",
-    ],
-    version_script = select({
-        "//tensorflow:macos": "pywrap_xla_exported_symbols.lds",
-        "//tensorflow:windows": None,
-        "//conditions:default": "pywrap_xla_version_script.lds",
-    }),
-    visibility = ["//visibility:public"],
-    deps = [
-        ":numpy_bridge",
-        ":xrt",
-        "//tensorflow/compiler/xla:literal",
-        "//tensorflow/compiler/xla:shape_util",
-        "//tensorflow/compiler/xla:xla_data_proto",
-    ],
-)