wip

pytorch · Feb 17, 2024 · dad6a13 · dad6a13
1 parent 7713cc7
commit dad6a13
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/.github/scripts/fbgemm_gpu_build.bash b/.github/scripts/fbgemm_gpu_build.bash
@@ -66,7 +66,9 @@ __configure_compiler_flags () {
     local conda_prefix=$(conda run ${env_prefix} printenv CONDA_PREFIX)
 
     # shellcheck disable=SC2155,SC2086
-    local cpp_path=$(conda run ${env_prefix} which c++)
+    local cxx_path=$(conda run ${env_prefix} which c++)
+    # shellcheck disable=SC2155,SC2086
+    local cxx_dir=$(dirname "$cxx_path")
 
     # echo "[BUILD] Setting Clang (should already be symlinked as c++) as the host compiler: ${cpp_path}"
     # # NOTE: There appears to be no ROCm equivalent for NVCC_PREPEND_FLAGS:
@@ -77,7 +79,7 @@ __configure_compiler_flags () {
     # echo "[BUILD] Configuring for Clang build ..."
     # shellcheck disable=SC2206
     build_args+=(
-      --compiler ${cpp_path}
+      --cxxdir ${cxx_dir}
     )
   fi
 }

diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py
@@ -53,7 +53,7 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
         " this in a custom location (through cudatoolkit-dev), provide the path here.",
     )
     parser.add_argument(
-        "--compiler",
+        "--cxxdir",
         type=str,
         default=None,
         # nargs=2,
@@ -178,13 +178,13 @@ def _get_cxx11_abi():
     if setup_py_args.nvml_lib_path:
         cmake_args.append(f"-DNVML_LIB_PATH={setup_py_args.nvml_lib_path}")
 
-    if setup_py_args.compiler:
+    if setup_py_args.cxxdir:
         print("[SETUP.PY] Setting compiler flags ...")
         # name, path = setup_py_args.openmp
         cmake_args.extend(
             [
-                f"-DCMAKE_C_COMPILER={setup_py_args.compiler}",
-                f"-DCMAKE_CXX_COMPILER={setup_py_args.compiler}",
+                f"-DCMAKE_C_COMPILER={setup_py_args.cxxdir}/cc",
+                f"-DCMAKE_CXX_COMPILER={setup_py_args.cxxdir}/c++",
                 f"-DCMAKE_C_FLAGS='-fopenmp -stdlib=libstdc++'",
                 f"-DCMAKE_CXX_FLAGS='-fopenmp -stdlib=libstdc++'",
                 # f"-DOpenMP_C_LIB_NAMES={name}",