Add folly support (attempt 2)

Summary: - Add folly support (attempt 2)

Differential Revision: D60616299

.github/scripts/fbgemm_gpu_build.bash

@@ -244,6 +244,19 @@ __configure_fbgemm_gpu_build_genai () {
   done
 }
 
+__configure_folly_flags () {
+  echo "[BUILD] Looking up Folly library filepath ..."
+  # shellcheck disable=SC2155,SC2086
+  local conda_prefix=$(conda run ${env_prefix} printenv CONDA_PREFIX)
+  # shellcheck disable=SC2155,SC2086
+  local folly_lib_path=$(conda run ${env_prefix} find ${conda_prefix} -name "libfolly.so")
+
+  echo "[BUILD] Setting CUDA build args ..."
+  build_args+=(
+    --folly_lib_path="${folly_lib_path}"
+  )
+}
+
 # shellcheck disable=SC2120
 __configure_fbgemm_gpu_build () {
   echo "################################################################################"
@@ -278,6 +291,9 @@ __configure_fbgemm_gpu_build () {
     __configure_fbgemm_gpu_build_clang
   fi
 
+  # Add build flags to support building against Folly
+  __configure_folly_flags
+
   # Set verbosity
   build_args+=(
     --verbose

fbgemm_gpu/setup.py

@@ -76,6 +76,12 @@ def from_args(cls, argv: List[str]):
             default=None,
             help="NCCL (libnccl.so.2) filepath. This is required for building certain targets.",
         )
+        parser.add_argument(
+            "--folly_lib_path",
+            type=str,
+            default=None,
+            help="Folly (libfolly.so) filepath. This is required for building certain targets.",
+        )
         parser.add_argument(
             "--cxxprefix",
             type=str,
@@ -276,6 +282,16 @@ def _get_cxx11_abi():
                 ]
             )
 
+        if self.args.folly_lib_path:
+            folly_root = os.path.dirname(os.path.dirname(self.args.folly_lib_path))
+            cxx_flags.extend([f"-L{folly_root}/lib"])
+            cmake_args.extend(
+                [
+                    f"-DFOLLY_INCLUDE_DIRS={folly_root}/include",
+                    f"-DFOLLY_LIBRARIES={self.args.folly_lib_path}",
+                ]
+            )
+
         if self.args.cxxprefix:
             print("[SETUP.PY] Setting CMake flags ...")
             path = self.args.cxxprefix