experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[TF] Add version_script option to tf_py_wrap_cc Bazel rule.

Browse Source 
[XLA] Hide all symbols except PyInit_* from the pywrap_xla module.

The XLA Python extension includes parts of TensorFlow core. If the TensorFlow symbols are exported by the XLA extension as well as TensorFlow itself, a segfault occurs if both tensorflow and jax are imported into the same Python interpreter.

JAX issue https://github.com/google/jax/issues/120 will be fixed when this change is integrated into the JAX repository.

PiperOrigin-RevId: 229962366
This commit is contained in:
Peter Hawkins 2019-01-18 11:24:02 -08:00 committed by TensorFlower Gardener
parent 691f6c2358
commit 7fce32e9be
6 changed files with 20 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/xla/python/BUILD
View File

@ -98,6 +98,11 @@ tf_py_wrap_cc(
"local_computation_builder.i",
"//tensorflow/python:platform/base.i",
],
version_script = select({
"//tensorflow:darwin": "pywrap_xla_exported_symbols.lds",
"//tensorflow:windows": None,
"//conditions:default": "pywrap_xla_version_script.lds",
}),
deps = [
":local_computation_builder",
":numpy_bridge",

1
tensorflow/compiler/xla/python/pywrap_xla_exported_symbols.lds Normal file
View File

@ -0,0 +1 @@
_PyInit__pywrap_xla

6
tensorflow/compiler/xla/python/pywrap_xla_version_script.lds Normal file
View File

@ -0,0 +1,6 @@
xla {
global:
PyInit_*;
local:
*;
};

12
tensorflow/tensorflow.bzl
View File

@ -1649,7 +1649,9 @@ def tf_py_wrap_cc(
swig_includes = [],
deps = [],
copts = [],
version_script = None,
**kwargs):
"""Builds a Python extension module."""
module_name = name.split("/")[-1]
# Convert a rule name such as foo/bar/baz to foo/bar/_baz.so
@ -1668,6 +1670,11 @@ def tf_py_wrap_cc(
toolchain_deps = ["@bazel_tools//tools/cpp:current_cc_toolchain"],
deps = deps + extra_deps,
)
if not version_script:
version_script = select({
"@local_config_cuda//cuda:darwin": clean_dep("//tensorflow:tf_exported_symbols.lds"),
"//conditions:default": clean_dep("//tensorflow:tf_version_script.lds"),
})
vscriptname = name + "_versionscript"
_append_init_to_versionscript(
name = vscriptname,
@ -1676,10 +1683,7 @@ def tf_py_wrap_cc(
"//conditions:default": True,
}),
module_name = module_name,
template_file = select({
"@local_config_cuda//cuda:darwin": clean_dep("//tensorflow:tf_exported_symbols.lds"),
"//conditions:default": clean_dep("//tensorflow:tf_version_script.lds"),
}),
template_file = version_script,
)
extra_linkopts = select({
"@local_config_cuda//cuda:darwin": [

1
tensorflow/tf_exported_symbols.lds
View File

@ -4,5 +4,4 @@
*TF_*
*TFE_*
*nsync_*
*pywrap_xla*
*stream_executor*

1
tensorflow/tf_version_script.lds
View File

@ -5,7 +5,6 @@ tensorflow {
*TF_*;
*TFE_*;
*nsync_*;
*pywrap_xla*;
*stream_executor*;
local:
*;