From bcc140d47a0d2a7b7bd117055274c68d9f632aa0 Mon Sep 17 00:00:00 2001
From: fpetrini15 <francescogpetrini@gmail.com>
Date: Fri, 30 Aug 2024 13:12:25 -0700
Subject: [PATCH] RHEL SBSA TF2 Backend Build

---
 CMakeLists.txt | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6adc487..1f4437e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -134,7 +134,11 @@ configure_file(src/libtriton_tensorflow.ldscript libtriton_tensorflow.ldscript C
 
 if (${TRITON_TENSORFLOW_DOCKER_BUILD})
   if (CMAKE_HOST_SYSTEM_PROCESSOR MATCHES "aarch64")
-    set(LIBS_ARCH "aarch64")
+    if(${RHEL_BUILD})
+      set(LIBS_ARCH "sbsa")
+    else()
+      set(LIBS_ARCH "aarch64")
+    endif()
   else()
     set(LIBS_ARCH "x86_64")
   endif()
@@ -155,7 +159,7 @@ if (${TRITON_TENSORFLOW_DOCKER_BUILD})
     COMMAND docker stop tensorflow_backend_deps || echo "error ignored..." || true
     COMMAND docker rm tensorflow_backend_deps || echo "error ignored..." || true
     COMMAND if [ "${TRITON_TENSORFLOW_INSTALL_EXTRA_DEPS}" = "ON" ] \; then mkdir tf_backend_deps && docker run -it -d --name tensorflow_backend_deps ${TRITON_TENSORFLOW_DOCKER_IMAGE} \; fi \;
-    COMMAND if [ "${TRITON_TENSORFLOW_INSTALL_EXTRA_DEPS}" = "ON" ] \; then docker exec tensorflow_backend_deps sh -c  "tar -cf - $<IF:$<BOOL:${RHEL_BUILD}>,/usr/local/cuda/targets/x86_64-linux/lib/,/usr/lib/${LIBS_ARCH}-linux-gnu/>libnccl.so*" | tar --strip-components=3 -xf - -C ./tf_backend_deps \; fi
+    COMMAND if [ "${TRITON_TENSORFLOW_INSTALL_EXTRA_DEPS}" = "ON" ] \; then docker exec tensorflow_backend_deps sh -c  "tar -cf - $<IF:$<BOOL:${RHEL_BUILD}>,/usr/local/cuda/targets/${LIBS_ARCH}-linux/lib/,/usr/lib/${LIBS_ARCH}-linux-gnu/>libnccl.so*" | tar --strip-components=3 -xf - -C ./tf_backend_deps \; fi
     COMMAND if [ "${TRITON_TENSORFLOW_INSTALL_EXTRA_DEPS}" = "ON" ] \; then docker stop tensorflow_backend_deps && docker rm tensorflow_backend_deps \; fi \;
 
     COMMENT "Extracting ${TRITON_TENSORFLOW_CC_LIBNAME}.${TRITON_TENSORFLOW_VERSION} and ${TRITON_TENSORFLOW_FW_LIBNAME}.${TRITON_TENSORFLOW_VERSION} from ${TRITON_TENSORFLOW_DOCKER_IMAGE}"