diff --git a/.github/scripts/fbgemm_gpu_build.bash b/.github/scripts/fbgemm_gpu_build.bash index 0b979c9a14..00b546d4f5 100644 --- a/.github/scripts/fbgemm_gpu_build.bash +++ b/.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 diff --git a/.github/scripts/utils_build.bash b/.github/scripts/utils_build.bash index 6e6f9c2e4c..d1d79bbd30 100644 --- a/.github/scripts/utils_build.bash +++ b/.github/scripts/utils_build.bash @@ -249,6 +249,7 @@ install_build_tools () { build \ click \ cmake \ + folly \ hypothesis \ jinja2 \ make \ diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py index cdb131f558..0656a20bb1 100644 --- a/fbgemm_gpu/setup.py +++ b/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 diff --git a/fbgemm_gpu/src/metric_ops/metric_ops_host.cpp b/fbgemm_gpu/src/metric_ops/metric_ops_host.cpp index c34bfcea9b..0003004a46 100644 --- a/fbgemm_gpu/src/metric_ops/metric_ops_host.cpp +++ b/fbgemm_gpu/src/metric_ops/metric_ops_host.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include "fbgemm_gpu/sparse_ops_utils.h" @@ -14,6 +15,10 @@ namespace fbgemm_gpu { +uint32_t example_folly_code(const std::vector& buffer) { + return folly::crc32c(buffer.data(), buffer.size()); +} + TORCH_LIBRARY_FRAGMENT(fbgemm, m) { m.def( "batch_auc(int num_tasks, Tensor indices, Tensor laebls, Tensor weights) -> Tensor");