port omp flags to yacl

secretflow · Dec 1, 2023 · 2b51b51 · 2b51b51
1 parent 3baea61
commit 2b51b51
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 0 deletions.
diff --git a/bazel/repositories.bzl b/bazel/repositories.bzl
@@ -62,6 +62,19 @@ def yacl_deps():
         remote = "{}/interconnection.git".format(SECRETFLOW_GIT),
     )
 
+    # Add homebrew openmp for macOS, somehow..homebrew installs to different location on Apple Silcon/Intel macs.. so we need two rules here
+    native.new_local_repository(
+        name = "macos_omp_x64",
+        build_file = "@yacl//bazel:local_openmp_macos.BUILD",
+        path = "/usr/local/opt/libomp",
+    )
+
+    native.new_local_repository(
+        name = "macos_omp_arm64",
+        build_file = "@yacl//bazel:local_openmp_macos.BUILD",
+        path = "/opt/homebrew/opt/libomp/",
+    )
+
 def _com_github_brpc_brpc():
     maybe(
         http_archive,

diff --git a/bazel/yacl.bzl b/bazel/yacl.bzl
@@ -39,6 +39,12 @@ AES_COPT_FLAGS = select({
     ],
 })
 
+OMP_LINK_FLAGS = select({
+    "@bazel_tools//src/conditions:darwin_x86_64": ["@macos_omp_x64//:openmp"],
+    "@bazel_tools//src/conditions:darwin_arm64": ["@macos_omp_arm64//:openmp"],
+    "//conditions:default": [],
+})
+
 def _yacl_copts():
     return select({
         "@yacl//bazel:yacl_build_as_release": RELEASE_FLAGS,