Enable multithread on OpenFHE with OpenMP support

Included suggested flags from discussion #1209. In addition, compile and link flags added to `tests/Examples/openfhe` to support OpenMP. Multithreaded execution is enabled by default with this setup. Clang based toolchain is supported currently.
google · Jan 26, 2025 · 1c42eb8 · 1c42eb8
1 parent dd91785
commit 1c42eb8
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 11 deletions.
diff --git a/bazel/openfhe/copts.bzl b/bazel/openfhe/copts.bzl
@@ -0,0 +1,19 @@
+"""Build settings for OpenFHE and OpenMP."""
+
+OPENFHE_COPTS = [
+    "-Wno-non-virtual-dtor",
+    "-Wno-shift-op-parentheses",
+    "-Wno-unused-private-field",
+    "-fexceptions",
+]
+
+OPENFHE_LINKOPTS = [
+    "-fopenmp",
+    "-lomp",
+]
+
+OPENMP_COPTS = [
+    "-fopenmp",
+    "-Xpreprocessor",
+    "-Wno-unused-command-line-argument",
+]
diff --git a/bazel/openfhe/openfhe.BUILD b/bazel/openfhe/openfhe.BUILD
@@ -1,4 +1,5 @@
 # BUILD file for a bazel-native OpenFHE build
+load("@heir//bazel/openfhe:copts.bzl", "OPENFHE_COPTS", "OPENFHE_LINKOPTS", "OPENMP_COPTS")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -18,17 +19,10 @@ OPENFHE_VERSION_PATCH = 3
 
 OPENFHE_VERSION = "{}.{}.{}".format(OPENFHE_VERSION_MAJOR, OPENFHE_VERSION_MINOR, OPENFHE_VERSION_PATCH)
 
-OPENFHE_COPTS = [
-    "-Wno-non-virtual-dtor",
-    "-Wno-shift-op-parentheses",
-    "-Wno-unused-private-field",
-    "-fexceptions",
-]
-
 OPENFHE_DEFINES = [
     "MATHBACKEND=2",
-    "OMP_NUM_THREADS=1",
     "OPENFHE_VERSION=" + OPENFHE_VERSION,
+    "PARALLEL",
 ]
 
 # This rule exists so that the python frontend can get access to the headers to
@@ -48,7 +42,7 @@ cc_library(
         "src/core/lib/**/*.c",
         "src/core/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS + [
+    copts = OPENFHE_COPTS + OPENMP_COPTS + [
         # /utils/blockAllocator/blockAllocator.cpp has misaligned-pointer-use
         "-fno-sanitize=alignment",
     ],
@@ -57,6 +51,7 @@ cc_library(
         "src/core/include",
         "src/core/lib",
     ],
+    linkopts = OPENFHE_LINKOPTS,
     textual_hdrs = glob([
         "src/core/include/**/*.h",
         "src/core/lib/**/*.cpp",
@@ -70,12 +65,13 @@ cc_library(
         "src/binfhe/lib/**/*.c",
         "src/binfhe/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS,
+    copts = OPENFHE_COPTS + OPENMP_COPTS,
     defines = OPENFHE_DEFINES,
     includes = [
         "src/binfhe/include",
         "src/binfhe/lib",
     ],
+    linkopts = OPENFHE_LINKOPTS,
     textual_hdrs = glob(["src/binfhe/include/**/*.h"]),
     deps = [
         "@openfhe//:core",
@@ -87,14 +83,15 @@ cc_library(
     srcs = glob([
         "src/pke/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS + [
+    copts = OPENFHE_COPTS + OPENMP_COPTS + [
         "-Wno-vla-extension",
     ],
     defines = OPENFHE_DEFINES,
     includes = [
         "src/pke/include",
         "src/pke/lib",
     ],
+    linkopts = OPENFHE_LINKOPTS,
     textual_hdrs = glob([
         "src/pke/include/**/*.h",
         "src/pke/lib/**/*.cpp",

diff --git a/tests/Examples/openfhe/test.bzl b/tests/Examples/openfhe/test.bzl
@@ -1,5 +1,6 @@
 """A macro providing an end-to-end test for OpenFHE codegen."""
 
+load("@heir//bazel/openfhe:copts.bzl", "OPENFHE_LINKOPTS", "OPENMP_COPTS")
 load("@heir//tools:heir-opt.bzl", "heir_opt")
 load("@heir//tools:heir-translate.bzl", "heir_translate")
 
@@ -55,6 +56,8 @@ def openfhe_end_to_end_test(name, mlir_src, test_src, generated_lib_header, heir
         hdrs = [":" + generated_lib_header],
         deps = deps + ["@openfhe//:pke"],
         tags = tags,
+        copts = OPENMP_COPTS,
+        linkopts = OPENFHE_LINKOPTS,
         **kwargs
     )
     native.cc_test(
@@ -68,5 +71,7 @@ def openfhe_end_to_end_test(name, mlir_src, test_src, generated_lib_header, heir
         ],
         tags = tags,
         data = data,
+        copts = OPENMP_COPTS,
+        linkopts = OPENFHE_LINKOPTS,
         **kwargs
     )