Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Enable TF XLA, add tensorflow-xla-runtime. #9005

Merged
merged 8 commits into from
Mar 14, 2024
Merged

Enable TF XLA, add tensorflow-xla-runtime. #9005

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0f8b473
Port XLA/AOT from 14_0_X.
riga Dec 19, 2023
cf12ea3
Add tensorflow-xla-runtime to vectorization/cmsdist_packages.py.
riga Feb 12, 2024
e40b72c
Add "-msse3" flag in tensorflow-xla-runtime.spec.
riga Feb 12, 2024
cc33bb3
Add tensorflow-xla-runtime to cmssw-tool-conf.spec.
riga Feb 12, 2024
f77258c
Address review comments by @smuzaffar.
riga Feb 14, 2024
12f9496
Remove include.
riga Feb 14, 2024
a1019af
Merge branch 'IB/CMSSW_14_1_X/master' into IB/CMSSW_14_1_X/tf_xla_aot
smuzaffar Mar 13, 2024
c2096d0
Remove tensorflow-xla-runtime from vectorization file.
riga Mar 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cmssw-vectorization.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1,2 @@
%define vectorized_packages zlib fastjet tensorflow-sources tensorflow OpenBLAS rivet gbl
%define vectorized_packages zlib fastjet tensorflow-sources tensorflow tensorflow-xla-runtime OpenBLAS rivet gbl
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@riga, please also update https://github.com/cms-sw/cmsdist/blob/IB/CMSSW_14_1_X/master/vectorization/cmsdist_packages.py#L14-L22 and include tensorflow-xla-runtime in the list

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added 👍

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ah sorry, I forgot to mention, please also add scram-tools.file/tools/tensorflow-xla-runtime/vectorized.tmpl file same as https://github.com/cms-sw/cmsdist/blob/IB/CMSSW_14_1_X/master/scram-tools.file/tools/tensorflow/vectorized.tmpl except that the tool name should be tensorflow-xla-runtime_@TOOL_VECTORIZATION@

%{expand:%(for t in %{vectorized_packages} ; do echo Requires: $t; for v in %{package_vectorization}; do echo Requires: ${t}_${v}; done; done)}
2 changes: 1 addition & 1 deletion pip/cmsml.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1 +1 @@
Requires: py3-six
Requires: py3-six py3-tabulate py3-pytest py3-pytest-xdist
1 change: 1 addition & 0 deletions pip/execnet.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
BuildRequires: py3-hatchling py3-hatch-vcs
1 change: 1 addition & 0 deletions pip/pytest-xdist.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Requires: py3-pytest py3-execnet
14 changes: 9 additions & 5 deletions pip/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ charset-normalizer==3.1.0
cleo==2.0.1
click==8.1.3
clikit==0.6.2
cmsml==0.1.2
cmsml==0.2.2
contourpy==1.0.7
correctionlib==2.2.2
crashtest==0.4.1
Expand Down Expand Up @@ -94,6 +94,7 @@ docutils==0.19
dulwich==0.21.3
editables==0.3
entrypoints==0.4
execnet==2.0.2
executing==1.2.0
exceptiongroup==1.1.1
fasteners==0.18
Expand Down Expand Up @@ -170,13 +171,13 @@ keras-preprocessing==1.1.2
keras2onnx==1.7.0
keyring==23.13.1
kiwisolver==1.4.4
law==0.1.12
law==0.1.18
lazy-object-proxy==1.9.0
lizard==1.17.10
#llvmlite==v0.41.0dev0
locket==1.0.0
lockfile==0.12.2
luigi==3.2.1
luigi==3.5.0
lxml==4.9.2
lz4==4.3.2
make==0.1.6.post2
Expand Down Expand Up @@ -284,6 +285,7 @@ pysqlite3==0.5.0
pytest==7.4.0
pytest-cov==4.0.0
pytest-runner==6.0.0
pytest-xdist==3.5.0
python-daemon==3.0.1
python-dateutil==2.8.2
python-json-logger==2.0.7
Expand Down Expand Up @@ -313,7 +315,7 @@ schema==0.7.5
scikit-build==0.16.7
scikit-build-core==0.3.0
scikit-learn==1.2.2
scinum==1.4.6
scinum==2.0.2
scipy==1.10.1
secretstorage==3.3.3
semantic-version==2.10.0
Expand All @@ -337,6 +339,7 @@ stevedore==5.0.0
subprocess32==3.5.4
sympy==1.11.1
tables==3.8.0
tabulate==0.9.0
tenacity==8.2.2
#NO_AUTO_UPDATE:1: Force to use tensorflow 2.6.4; this should match the version in tensorflow-sources.spec
tensorflow==2.12.0
Expand Down Expand Up @@ -379,7 +382,8 @@ Werkzeug==3.0.1
#NO_AUTO_UPDATE:1: you need wheel to build wheel
wheel==0.40.0
widgetsnbextension==4.0.7
wrapt==1.15.0
#NO_AUTO_UPDATE: wrapt >=1.15.0 breaks tensorflow saved model, https://github.com/GrahamDumpleton/wrapt/issues/231
wrapt==1.14.1
xgboost==1.7.5
yarl==1.8.2
zipp==3.15.0
1 change: 1 addition & 0 deletions pip/tabulate.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
BuildRequires: py3-hatchling py3-hatch-vcs
2 changes: 1 addition & 1 deletion pip/tensorflow.file
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
Requires: py3-tensorboard py3-opt-einsum py3-tensorflow-estimator py3-wrapt py3-google-pasta py3-scipy
Requires: tensorflow-sources
%define PipPreBuildPy PIPFILE=${TENSORFLOW_SOURCES_ROOT}/tensorflow-%{realversion}-cp%{cms_python3_major_minor}-cp%{cms_python3_major_minor}-linux_%{_arch}.whl
%define PipPostBuild rm -f %{i}/bin/tensorboard* ; ls %{i}/bin/* | xargs -i mv '{}' '{}3'
%define PipPostBuild rm -f %{i}/bin/tensorboard* ; rm -r %{i}/lib/python%{cms_python3_major_minor_version}/site-packages/tensorflow/xla_aot_runtime_src ; ls %{i}/bin/* | xargs -i mv '{}' '{}3'
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@riga any specific reason to delete it? Looks like this is what you get out of tensorflow-sources wheel file. May be if we do not delete it here then you can just add dependency on py3-tensorflow and copy $PY3_TENSORFLOW_ROOT/$PYTHON3_LIB_SITE_PACKAGES/tensorflow/xla_aot_runtime_src to build tensorflow-xla-runtime ... right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Indeed. We changed this in f77258c.

%define PipDepsPost | grep -v libclang
## INCLUDE tensorflow-requires
Requires: py3-funcsigs py3-pbr py3-packaging py3-appdirs py3-pyparsing py3-mock py3-Werkzeug
Expand Down
1 change: 1 addition & 0 deletions python_tools.spec
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ Requires: py3-law
Requires: py3-protobuf

Requires: py3-tables
Requires: py3-tabulate
Requires: py3-numexpr
Requires: py3-histogrammar
Requires: py3-pandas
Expand Down
9 changes: 9 additions & 0 deletions scram-tools.file/tools/tensorflow-xla-runtime/tensorflow-xla-runtime.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
<tool name="tensorflow-xla-runtime" version="@TOOL_VERSION@">
<client>
<environment name="TENSORFLOW_XLA_RUNTIME_BASE" default="@TOOL_ROOT@"/>
<environment name="LIBDIR" default="$TENSORFLOW_XLA_RUNTIME_BASE/lib/archive"/>
<environment name="INCLUDE" default="$TENSORFLOW_XLA_RUNTIME_BASE/include"/>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@riga, I guess this is not needed any more ... right?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right. It's removed now.

</client>
<lib name="tf_xla_runtime-static"/>
<use name="eigen"/>
</tool>
6 changes: 4 additions & 2 deletions tensorflow-sources.file
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ export TF_NEED_CUDA=%{enable_gpu}
export TF_NEED_JEMALLOC=0
export TF_NEED_HDFS=0
export TF_NEED_GCP=0
export TF_ENABLE_XLA=0
export TF_ENABLE_XLA=1
export TF_NEED_OPENCL=0
export TF_NEED_VERBS=0
export TF_NEED_MKL=0
Expand Down Expand Up @@ -150,6 +150,7 @@ bazel $BAZEL_OPTS //tensorflow/compiler/tf2xla:xla_compiled_cpu_function
#bazel $BAZEL_OPTS //tensorflow/compiler/aot:tfcompile
bazel $BAZEL_OPTS //tensorflow/core/profiler
bazel $BAZEL_OPTS //tensorflow:install_headers
bazel $BAZEL_OPTS //tensorflow/compiler/tf2xla:tf2xla_supported_ops

# rebuild *.pb.{h|cc} files using the external protobuf compiler
chmod -R a+rwX $PWD/bazel-bin/tensorflow/include
Expand Down Expand Up @@ -192,8 +193,9 @@ for l in tensorflow_cc tensorflow_framework tensorflow ; do
done

#cp -p $srcdir/compiler/aot/tfcompile $bindir
cp -p $srcdir/compiler/tf2xla/tf2xla_supported_ops $bindir
for name in tensorflow absl re2 third_party ; do
cp -r -p $srcdir/include/$name $incdir
cp -r -p $srcdir/include/$name $incdir
done

# copy headers from downloaded dependencies
Expand Down
32 changes: 32 additions & 0 deletions tensorflow-xla-runtime.spec
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
### RPM external tensorflow-xla-runtime 2.12.0
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@riga, we need tensorflow-xla-runtime to be part of cmssw dependency chain, so please add it in https://github.com/cms-sw/cmsdist/blob/IB/CMSSW_14_1_X/master/cmssw-tool-conf.spec

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added 👍


Requires: eigen tensorflow-sources
BuildRequires: cmake

%prep
unzip -q -o ${TENSORFLOW_SOURCES_ROOT}/tensorflow-%{realversion}-cp%{cms_python3_major_minor}-cp%{cms_python3_major_minor}-linux_%{_arch}.whl

%build

# fix a single comparison between size_t and int
sed -i -r 's/assert\(\(arg_size\(index\)/assert\(\(\(size_t\)arg_size\(index\)/' tensorflow/include/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h

export CPATH="${CPATH}:${EIGEN_ROOT}/include/eigen3"
export CPATH="${CPATH}:%{i}/tensorflow/include"
export CPATH="${CPATH}:%{i}/tensorflow/include/third_party/eigen3"

pushd tensorflow/xla_aot_runtime_src
cmake . -DCMAKE_CXX_FLAGS="-fPIC"
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In order to properly build for multi-vectorization please add -msse3 here e.g. -DCMAKE_CXX_FLAGS="-fPIC -msse3"

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added 👍

make %{makeprocesses}
# this builds a shared library, but when used some symbols are missig (e.g `tsl::mutex::unlock()')
# so it does not seem intended to be used as a shared lib
gcc -shared -o libtf_xla_runtime.so -Wl,--whole-archive libtf_xla_runtime.a -Wl,--no-whole-archive
popd

%install

mv tensorflow/include %{i}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@riga, the contents of tensorflow/include are same as the one in $PY3_TENSORFLOW_ROOT/$PYTHON3_LIB_SITE_PACKAGES/tensorflow/include . I would suggest that instead of duplicating this why not we get the include from py3-tensorflow ( we can add an extra toolfile so that these headers are accessible in cmssw)

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In f77258c we changed the dependency of tensorflow-xla-runtime from tensorflow-sources to py3-tensorflow and added a separate tool file py3-tensorflow/tensorflow-includes.xml that we used in the toolfile of the xla runtime.


mkdir -p %{i}/lib/archive
mv tensorflow/xla_aot_runtime_src/libtf_xla_runtime.a %{i}/lib/archive/libtf_xla_runtime-static.a
mv tensorflow/xla_aot_runtime_src/libtf_xla_runtime.so %{i}/lib/libtf_xla_runtime.so
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need the shared lib?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not for now. We removed it in f77258c.

1 change: 1 addition & 0 deletions vectorization/cmsdist_packages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
"fastjet",
"tensorflow-sources",
"tensorflow",
"tensorflow-xla-runtime",
"OpenBLAS",
"rivet",
"gbl",
Expand Down