From 2b51b5196b9c6235fb07196cb1bd870bfb16e4ff Mon Sep 17 00:00:00 2001
From: anakinxc <103552181+anakinxc@users.noreply.github.com>
Date: Fri, 1 Dec 2023 07:04:15 +0000
Subject: [PATCH] port omp flags to yacl

---
 bazel/repositories.bzl | 13 +++++++++++++
 bazel/yacl.bzl         |  6 ++++++
 2 files changed, 19 insertions(+)

diff --git a/bazel/repositories.bzl b/bazel/repositories.bzl
index 8f67a936..d32665e3 100644
--- 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
index a516c2c3..14b3fd97 100644
--- 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,