Enable TF XLA, add tensorflow-xla-runtime.

cmssw-vectorization.file

@@ -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
 %{expand:%(for t in %{vectorized_packages} ; do echo Requires: $t; for v in %{package_vectorization}; do echo Requires: ${t}_${v}; done; done)}

pip/cmsml.file

@@ -1 +1 @@
-Requires: py3-six
+Requires: py3-six py3-tabulate py3-pytest py3-pytest-xdist

pip/execnet.file

@@ -0,0 +1 @@
+BuildRequires: py3-hatchling py3-hatch-vcs

pip/pytest-xdist.file

@@ -0,0 +1 @@
+Requires: py3-pytest py3-execnet

pip/requirements.txt

@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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

pip/tabulate.file

@@ -0,0 +1 @@
+BuildRequires: py3-hatchling py3-hatch-vcs

pip/tensorflow.file

@@ -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'
 %define PipDepsPost | grep -v libclang
 ## INCLUDE tensorflow-requires
 Requires: py3-funcsigs py3-pbr py3-packaging py3-appdirs py3-pyparsing py3-mock py3-Werkzeug

python_tools.spec

@@ -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

scram-tools.file/tools/tensorflow-xla-runtime/tensorflow-xla-runtime.xml

@@ -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"/>
+  </client>
+  <lib name="tf_xla_runtime-static"/>
+  <use name="eigen"/>
+</tool>

tensorflow-sources.file

@@ -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
@@ -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
@@ -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

tensorflow-xla-runtime.spec

@@ -0,0 +1,32 @@
+### RPM external tensorflow-xla-runtime 2.12.0
+
+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"
+  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}
+
+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