Enable TF XLA, add tensorflow-xla-runtime.

cmssw-tool-conf.spec

@@ -97,6 +97,7 @@ Requires: hls
 Requires: opencv
 Requires: grpc
 Requires: onnxruntime
+Requires: tensorflow-xla-runtime
 Requires: TOoLLiP
 Requires: triton-inference-client
 Requires: hdf5

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

@@ -19,7 +19,8 @@ Requires: llvm
 %define PipPostInstall \
   for f in %{RelocateFiles} ; do \
     [ ! -f %{i}/$f ] || sed -i -e 's|/[^ ]*/%{cmsplatf}/external/|%{cmsroot}/%{cmsplatf}/external/|' %{i}/$f \
-  done
+  done ; \
+  sed -i 's/assert((arg_size(index)/assert(((size_t)arg_size(index)/' %{i}/lib/python%{cms_python3_major_minor_version}/site-packages/tensorflow/include/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.h
 
 %define PipPostPost \
   for f in %{RelocateFiles} ; do \

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/py3-tensorflow/py3-tensorflow.xml

@@ -0,0 +1,6 @@
+<tool name="py3-tensorflow" version="@TOOL_VERSION@">
+  <client>
+    <environment name="PY3_TENSORFLOW_BASE" default="@TOOL_ROOT@"/>
+  </client>
+  <runtime name="PATH" value="$PY3_TENSORFLOW_BASE/bin" type="path"/>
+</tool>

scram-tools.file/tools/py3-tensorflow/tensorflow-includes.xml

@@ -0,0 +1,6 @@
+<tool name="tensorflow-includes" version="@TOOL_VERSION@">
+  <client>
+    <environment name="TENSORFLOW_INCLUDES_BASE" default="@TOOL_ROOT@"/>
+    <environment name="INCLUDE" default="$TENSORFLOW_INCLUDES_BASE/@PYTHON3_LIB_SITE_PACKAGES@/tensorflow/include"/>
+  </client>
+</tool>

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

@@ -0,0 +1,10 @@
+<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"/>
+  <use name="tensorflow-includes"/>
+</tool>

scram-tools.file/tools/tensorflow-xla-runtime/vectorized.tmpl

@@ -0,0 +1,5 @@
+<tool name="tensorflow-xla-runtime_@TOOL_VECTORIZATION@" version="@TOOL_VERSION@">
+  <client>
+    <environment name="@TOOL_VECTORIZATION_KEY@_LIBDIR" default="@TOOL_ROOT@/lib"/>
+  </client>
+</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,23 @@
+### RPM external tensorflow-xla-runtime 2.12.0
+## INCLUDE cpp-standard
+
+Requires: eigen py3-tensorflow
+BuildRequires: cmake
+
+%prep
+
+cp -r ${PY3_TENSORFLOW_ROOT}/lib/python%{cms_python3_major_minor_version}/site-packages/tensorflow .
+
+%build
+
+export CPATH="${CPATH}:${EIGEN_ROOT}/include/eigen3"
+
+pushd tensorflow/xla_aot_runtime_src
+  cmake . -DCMAKE_CXX_FLAGS="-fPIC -msse3" -DCMAKE_CXX_STANDARD=%{cms_cxx_standard} -DBUILD_SHARED_LIBS=OFF
+  make %{makeprocesses}
+popd
+
+%install
+
+mkdir -p %{i}/lib/archive
+mv tensorflow/xla_aot_runtime_src/libtf_xla_runtime.a %{i}/lib/archive/libtf_xla_runtime-static.a

vectorization/cmsdist_packages.py

@@ -16,6 +16,7 @@
     "fastjet",
     "tensorflow-sources",
     "tensorflow",
+    "tensorflow-xla-runtime",
     "OpenBLAS",
     "rivet",
     "gbl",