cms-sw · cmsbuild · Mar 31, 2024 · Mar 20, 2024 · Mar 20, 2024 · Mar 21, 2024
diff --git a/cmssw-tool-conf.spec b/cmssw-tool-conf.spec
@@ -183,6 +183,7 @@ Requires: xtl
 Requires: xgboost
 Requires: pytorch
 
+## INCLUDE tfaot-models
 ## INCLUDE cmssw-vectorization
 ## INCLUDE cmssw-drop-tools
 ## INCLUDE scram-tool-conf
diff --git a/pip/cms-tfaot.file b/pip/cms-tfaot.file
@@ -0,0 +1,12 @@
+BuildRequires: py3-pip py3-setuptools py3-wheel
+Requires: py3-PyYAML py3-cmsml
+
+%define github_user cms-externals
+%define tag a03ccef81c610e051d3d7d172d675fe6c1ff17e2
+%define branch master
+%define source0 git+https://github.com/%{github_user}/cms-tfaot.git?obj=%{branch}/%{tag}&export=%{n}-%{realversion}&output=/%{n}-%{realversion}-%{tag}.tgz
+
+# copy test models
+%define PipPostInstall \
+  mkdir -p %{i}/share; \
+  cp -r cmsdist-tmp/pip-req-build-*/test_models %{i}/share/
diff --git a/pip/requirements.txt b/pip/requirements.txt
@@ -56,7 +56,8 @@ charset-normalizer==3.1.0
 cleo==2.0.1
 click==8.1.3
 clikit==0.6.2
-cmsml==0.2.2
+cmsml==0.2.5
+cms-tfaot==1.0.0
 contourpy==1.0.7
 correctionlib==2.2.2
 crashtest==0.4.1

diff --git a/python_tools.spec b/python_tools.spec
@@ -12,6 +12,7 @@ Requires: py3-keras
 Requires: py3-scikit-learn
 #save for the end
 Requires: py3-tensorflow
+Requires: py3-cms-tfaot
 Requires: py3-cmsml
 Requires: py3-law
 Requires: py3-protobuf

diff --git a/scram-tools.file/tools/tensorflow-xla-runtime/tensorflow-xla-runtime.xml b/scram-tools.file/tools/tensorflow-xla-runtime/tensorflow-xla-runtime.xml
@@ -1,9 +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="LIBDIR" default="$TENSORFLOW_XLA_RUNTIME_BASE/lib"/>
   </client>
-  <lib name="tf_xla_runtime-static"/>
+  <lib name="tf_xla_runtime"/>
+
   <use name="eigen"/>
   <use name="tensorflow-includes"/>
 </tool>
diff --git a/tensorflow-xla-runtime-absl.patch b/tensorflow-xla-runtime-absl.patch
@@ -0,0 +1,12 @@
+--- tensorflow/xla_aot_runtime_src/CMakeLists.txt    2024-03-24 08:28:34.000000000 +0100
++++ tensorflow/xla_aot_runtime_src/CMakeLists.txt      2024-03-25 11:17:58.108587945 +0100
+@@ -14,6 +14,8 @@
+   -Wno-sign-compare
+ )
+
+-add_library(tf_xla_runtime STATIC
++find_package(absl REQUIRED)
++add_library(tf_xla_runtime SHARED
+   $<TARGET_OBJECTS:tf_xla_runtime_objects>
+ )
++target_link_libraries(tf_xla_runtime absl::strings absl::str_format_internal)
diff --git a/tensorflow-xla-runtime.spec b/tensorflow-xla-runtime.spec
@@ -4,28 +4,39 @@
 
 Source99: scram-tools.file/tools/eigen/env
 
-Requires: eigen py3-tensorflow
+Patch0: tensorflow-xla-runtime-absl
+
+Requires: eigen py3-tensorflow abseil-cpp
 BuildRequires: cmake
 
 %prep
 
 cp -r ${PY3_TENSORFLOW_ROOT}/lib/python%{cms_python3_major_minor_version}/site-packages/tensorflow .
+%patch -p0
 
 %build
+
 source %{_sourcedir}/env
 export CPATH="${CPATH}:${EIGEN_ROOT}/include/eigen3"
 
-CXXFLAGS="-fPIC %{arch_build_flags} ${CMS_EIGEN_CXX_FLAGS}"
+CXXFLAGS="-fPIC -Wl,-z,defs %{arch_build_flags} ${CMS_EIGEN_CXX_FLAGS}"
 %ifarch x86_64
-    CXXFLAGS="${CXXFLAGS} -msse3"
+  CXXFLAGS="${CXXFLAGS} -msse3"
 %endif
 
 pushd tensorflow/xla_aot_runtime_src
-  cmake . -DCMAKE_CXX_FLAGS="${CXXFLAGS}" -DCMAKE_CXX_STANDARD=%{cms_cxx_standard} -DBUILD_SHARED_LIBS=OFF
+  # remove unnecessary implementations that use symbols that are not even existing
+  rm tensorflow/compiler/xla/service/cpu/runtime_fork_join.cc
+
+  cmake . \
+    -DCMAKE_CXX_FLAGS="${CXXFLAGS}" \
+    -DCMAKE_CXX_STANDARD=%{cms_cxx_standard} \
+    -DCMAKE_PREFIX_PATH=${ABSEIL_CPP_ROOT} \
+    -DBUILD_SHARED_LIBS=ON
   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
+mkdir -p %{i}/lib
+mv tensorflow/xla_aot_runtime_src/libtf_xla_runtime.so %{i}/lib/
diff --git a/tfaot-compile.file b/tfaot-compile.file
@@ -0,0 +1,41 @@
+## tfaot common compilation and requirement file
+## specs including this file should provide:
+## 1. a variable %{aot_config}, pointing to the aot config file of the model to compile (required)
+## 2. a variable %{aot_source}, referring to a fetched source to unpack during %prep (optional)
+##    (in this case a "Source" should be defined and %{aot_source} is likely %{n}-%{realversion})
+
+BuildRequires: py3-cms-tfaot
+Requires: tensorflow-xla-runtime
+
+%ifarch ppc64le
+%define build_arch powerpc64le-unknown-linux-gnu
+%else
+%define build_arch %{_arch}-unknown-linux-gnu
+%endif
+
+%prep
+%if "%{?aot_source}"
+%setup -n %{aot_source}
+%endif
+
+%build
+cms_tfaot_compile \
+  --aot-config "%{aot_config}" \
+  --tool-name "%{n}" \
+  --tool-base "%{i}" \
+  --output-directory compiled_model \
+  --additional-flags="--target_triple %{build_arch}"
+
+%install
+mkdir -p %{i}/lib
+mv compiled_model/*.o %{i}/lib/
+
+mkdir -p %{i}/include/%{n}
+mv compiled_model/*.h %{i}/include/%{n}
+
+mkdir -p %{i}/etc/scram.d
+mv compiled_model/%{n}.xml %{i}/etc/scram.d/
+
+%post
+%{relocateConfig}etc/scram.d/%{n}.xml
+%relocateConfigAll include/%{n} *.h
diff --git a/tfaot-model-test-multi.spec b/tfaot-model-test-multi.spec
@@ -0,0 +1,5 @@
+### RPM external tfaot-model-test-multi 1.0.0
+
+%define aot_config $PY3_CMS_TFAOT_ROOT/share/test_models/multi/aot_config.yaml
+
+## INCLUDE tfaot-compile
diff --git a/tfaot-model-test-simple.spec b/tfaot-model-test-simple.spec
@@ -0,0 +1,5 @@
+### RPM external tfaot-model-test-simple 1.0.0
+
+%define aot_config $PY3_CMS_TFAOT_ROOT/share/test_models/simple/aot_config.yaml
+
+## INCLUDE tfaot-compile
diff --git a/tfaot-models.file b/tfaot-models.file
@@ -0,0 +1,3 @@
+# test models needed by unit tests in PhysicsTools/TensorFlowAOT
+Requires: tfaot-model-test-simple
+Requires: tfaot-model-test-multi