Support ROCm 6.4 builds, pt 2 (#4151)

Summary:
X-link: facebookresearch/FBGEMM#1233

Pull Request resolved: #4151

- Support ROCm 6.4 builds, pt 2

Reviewed By: jwfromm, spcyppt

Differential Revision: D75008423

fbshipit-source-id: edf076f532ff2251ecf12ecd492bce918c022593

Author: q10

.github/scripts/fbgemm_gpu_build.bash

@@ -39,6 +39,7 @@ prepare_fbgemm_gpu_build () {
   echo "[BUILD] Running git submodules update ..."
   (exec_with_retries 3 git submodule sync) || return 1
   (exec_with_retries 3 git submodule update --init --recursive) || return 1
+  echo "[BUILD] Successfully ran git submodules update"
 
   # shellcheck disable=SC2155
   local env_prefix=$(env_name_or_prefix "${env_name}")
@@ -54,7 +55,16 @@ prepare_fbgemm_gpu_build () {
   # shellcheck disable=SC2086
   (test_python_import_package "${env_name}" skbuild) || return 1
 
-  echo "[BUILD] Successfully ran git submodules update"
+  # shellcheck disable=SC2155,SC2086
+  local conda_prefix=$(conda run ${env_prefix} printenv CONDA_PREFIX)
+
+  echo "[BUILD] Finding library paths ..."
+  # shellcheck disable=SC2061,SC2035
+  (find "${conda_prefix}" -name *omp.so* 2> /dev/null || true)
+  # shellcheck disable=SC2061,SC2035
+  (find "${conda_prefix}" -name *tbb.so* 2> /dev/null || true)
+  # shellcheck disable=SC2061,SC2035
+  (find "${conda_prefix}" -name *bolt.so* 2> /dev/null || true)
 }
 
 __configure_fbgemm_gpu_build_clang () {
@@ -599,7 +609,7 @@ __verify_library_symbols () {
       fbgemm_gpu::jagged_2d_to_dense
     )
 
-    if  [ "${fbgemm_build_variant}" == "cuda" ] &&
+    if  [ "${fbgemm_build_variant}" == "cuda" ] ||
         [ "${fbgemm_build_variant}" == "rocm" ]; then
       lib_symbols_to_check+=(
         fbgemm_gpu::asynchronous_inclusive_cumsum_gpu

.github/scripts/utils_build.bash

@@ -88,10 +88,6 @@ __conda_install_glibc () {
     (exec_with_retries 3 conda install ${env_prefix} -c conda-forge --override-channels -y \
       "libstdcxx-ng=${gcc_version}") || return 1
   fi
-
-  # Ensure libstdc++.so.6 is found
-  # shellcheck disable=SC2153
-  # (test_filepath "${env_name}" 'libstdc++.so.6') || return 1
 }
 
 __set_glibcxx_preload () {

.github/workflows/fbgemm_gpu_ci_rocm.yml

@@ -120,12 +120,6 @@ jobs:
       if: ${{ success() || failure() }}
       run: if . $PRELUDE && which conda; then collect_pytorch_env_info $BUILD_ENV; fi
 
-    - name: Find OpenMP Libraries
-      run: |
-        find / -name *omp.so* 2> /dev/null || true
-        find / -name *tbb.so* 2> /dev/null || true
-        find / -name *bolt.so* 2> /dev/null || true
-
     - name: Prepare FBGEMM_GPU Build
       run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
 
@@ -221,12 +215,6 @@ jobs:
     - name: Prepare FBGEMM_GPU Build
       run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
 
-    - name: Find OpenMP Libraries
-      run: |
-        find / -name *omp.so* 2> /dev/null || true
-        find / -name *tbb.so* 2> /dev/null || true
-        find / -name *bolt.so* 2> /dev/null || true
-
     - name: Install FBGEMM_GPU Wheel
       run: . $PRELUDE; install_fbgemm_gpu_wheel $BUILD_ENV *.whl
 

fbgemm_gpu/docs/src/fbgemm_gpu/development/BuildInstructions.rst

@@ -354,6 +354,7 @@ Install the other necessary build tools such as ``ninja``, ``cmake``, etc:
       ninja \
       numpy \
       scikit-build \
+      tbb \
       wheel
 
 .. _fbgemm-gpu.build.setup.pytorch.install:
@@ -675,6 +676,9 @@ presuming the toolchains have been properly installed.
 
   export ROCM_PATH=/path/to/rocm
 
+  # [OPTIONAL] If libtbb.so is missing, create the symlink (presuming libtbb.so.12 is present)
+  ln -s "${CONDA_PREFIX}/lib/libtbb.so.12" "${CONDA_PREFIX}/lib/libtbb.so"
+
   # [OPTIONAL] Enable verbose HIPCC logs
   export HIPCC_VERBOSE=1
 

fbgemm_gpu/setup.py

@@ -323,7 +323,7 @@ def _get_cxx11_abi():
                 + (
                     # Starting from ROCm 6.4, HIP clang complains about
                     # -fopenmp=libgomp being an invalid fopenmp-target
-                    ["-ltbb"]
+                    []
                     if self.variant() == "rocm"
                     else ["-fopenmp=libgomp"]
                 )
@@ -337,7 +337,23 @@ def _get_cxx11_abi():
             )
 
         if self.variant() == "rocm":
-            cxx_flags.extend([f"-DROCM_VERSION={RocmUtils.version_int()}"])
+            cxx_flags.extend(
+                [
+                    f"-DROCM_VERSION={RocmUtils.version_int()}",
+                ]
+            )
+
+            if self.nova_flag() is None:
+                cxx_flags.extend(
+                    [
+                        # For the ROCm case on non-Nova, an explicit link to
+                        # libtbb is required, or the following error is
+                        # encountered on library load:
+                        #
+                        # undefined symbol: _ZN3tbb6detail2r117deallocate_memoryEPv
+                        "-ltbb",
+                    ]
+                )
 
         cmake_args.extend(
             [