[TF] Build with GPU support

cms-sw · Feb 10, 2022 · 52fd78c · 52fd78c
1 parent c469bb3
commit 52fd78c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/tensorflow-requires.file b/tensorflow-requires.file
@@ -2,5 +2,5 @@ Requires: python3 py3-numpy py3-mock py3-typing py3-typing-extensions
 Requires: py3-keras-applications py3-keras-preprocessing py3-future py3-wrapt py3-gast py3-setuptools
 Requires: py3-cython py3-protobuf py3-astor py3-six py3-termcolor py3-absl-py
 Requires: py3-opt-einsum py3-flatbuffers
-Requires: eigen protobuf zlib libpng libjpeg-turbo curl giflib sqlite grpc flatbuffers py3-pybind11
+Requires: cuda cudnn eigen protobuf zlib libpng libjpeg-turbo curl giflib sqlite grpc flatbuffers py3-pybind11
 BuildRequires: py3-wheel
diff --git a/tensorflow-sources.file b/tensorflow-sources.file
@@ -3,8 +3,9 @@ BuildRequires: bazel java-env git
 #Keep all requires in separate file, so that can be included in py-tensorflow too
 ## INCLUDE tensorflow-requires
 ## INCLUDE compilation_flags
+## INCLUDE cuda-flags
 
-%define tag         719e00b6f9553de2662b6df2c353d6934e941103
+%define tag         0073b9fe65c29226311e4d5010176b837415cf2c
 %define branch      cms/v%{realversion}
 %define github_user cms-externals
 Source: git+https://github.com/%{github_user}/tensorflow.git?obj=%{branch}/%{tag}&export=tensorflow-%{realversion}&output=/tensorflow-%{realversion}.tgz
@@ -19,7 +20,6 @@ Source: git+https://github.com/%{github_user}/tensorflow.git?obj=%{branch}/%{tag
 %define majorversion %(echo %realversion | cut -d . -f 1)
 
 %prep
-
 %setup -q -n tensorflow-%{realversion}
 sed -i -e 's|lib/python[^/]*/site-packages/|lib/python%{cms_python3_major_minor_version}/site-packages/|'  third_party/systemlibs/pybind11.BUILD
 
@@ -41,14 +41,20 @@ BAZEL_OPTS="$BAZEL_OPTS --copt=-mcpu=native --copt=-mtune=native"
 %endif
 %endif
 BAZEL_OPTS="$BAZEL_OPTS --config=%{build_type} --cxxopt=-std=c++17 %{makeprocesses}"
-BAZEL_OPTS="$BAZEL_OPTS --config=noaws --config=nogcp --config=nohdfs --config=nonccl"
-
+BAZEL_OPTS="$BAZEL_OPTS --config=cuda --config=noaws --config=nogcp --config=nohdfs --config=nonccl"
+
+export GCC_HOST_COMPILER_PREFIX=${GCC_ROOT}/bin
+export GCC_HOST_COMPILER_PATH=$(which gcc)
+export TF_CUDA_COMPUTE_CAPABILITIES=$(echo "compute_%{cuda_arch}" | sed 's|\s\s*|,compute_|g')
+export TF_CUDA_VERSION="$(echo ${CUDA_VERSION} | cut -f1,2 -d.)"
+export TF_CUDA_PATHS="${CUDA_ROOT},${CUDNN_ROOT}"
+export TF_CUDA_CLANG=0
 export TF_NEED_JEMALLOC=0
 export TF_NEED_HDFS=0
 export TF_NEED_GCP=0
 export TF_ENABLE_XLA=0
 export TF_NEED_OPENCL=0
-export TF_NEED_CUDA=0
+export TF_NEED_CUDA=1
 export TF_NEED_VERBS=0
 export TF_NEED_MKL=0
 export TF_NEED_MPI=0
@@ -97,7 +103,7 @@ if [ -d ../build ] ; then
   chmod -R u+w  ../build
   rm -rf ../build
 fi
-./configure
+cuda=Y ./configure
 
 #Generate Python wrappers so that we can use PYTHON*PATH env
 rm -rf %{_builddir}/cms-pytool ; mkdir %{_builddir}/cms-pytool