Merge pull request #1314 from eymay:openmp_for_openfhe

PiperOrigin-RevId: 722648966

.github/workflows/build_and_test.yml

@@ -68,11 +68,11 @@ jobs:
       # requires the `PATH` to find the `cargo` binary.
       - name: "Run `bazel build`"
         run: |
-          bazel build --noincompatible_strict_action_env  -c fastbuild //...
+          bazel build --noincompatible_strict_action_env --//:enable_openmp=0 -c fastbuild //...
 
       - name: "Run `bazel test`"
         run: |
-          bazel test --noincompatible_strict_action_env -c fastbuild //...
+          bazel test --noincompatible_strict_action_env --//:enable_openmp=0 -c fastbuild //...
 
       # Tests specifically for the tfhe-rs codegen
       - name: rustup toolchain install

BUILD

@@ -21,8 +21,33 @@ package_group(
     packages = [],
 )
 
-# Disables Yosys deps for CLI tools and tests.
-# use by passing `--//:enable_yosys=0` to `bazel build` or `bazel test`
+# Disables deps for CI tools and tests.
+# use by passing `--//:enable_openmp=0` or `--//:enable_yosys=0`
+# to `bazel build` or `bazel test`
+
+# OpenMP
+string_flag(
+    name = "enable_openmp",
+    build_setting_default = "1",
+)
+
+config_setting(
+    name = "config_enable_openmp",
+    flag_values = {
+        "//:enable_openmp": "1",
+    },
+    visibility = ["//visibility:public"],
+)
+
+config_setting(
+    name = "config_disable_openmp",
+    flag_values = {
+        "//:enable_openmp": "0",
+    },
+    visibility = ["//visibility:public"],
+)
+
+# Yosys
 string_flag(
     name = "enable_yosys",
     build_setting_default = "1",

bazel/openfhe/copts.bzl

@@ -0,0 +1,45 @@
+"""Build settings for OpenFHE and OpenMP."""
+
+OPENFHE_VERSION_MAJOR = 1
+
+OPENFHE_VERSION_MINOR = 11
+
+OPENFHE_VERSION_PATCH = 3
+
+OPENFHE_VERSION = "{}.{}.{}".format(OPENFHE_VERSION_MAJOR, OPENFHE_VERSION_MINOR, OPENFHE_VERSION_PATCH)
+
+OPENFHE_DEFINES = [
+    "MATHBACKEND=2",
+    "OPENFHE_VERSION=" + OPENFHE_VERSION,
+] + select({
+    "@heir//:config_enable_openmp": ["PARALLEL"],
+    "@heir//:config_disable_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",
+]
+
+MAYBE_OPENFHE_LINKOPTS = select({
+    "@heir//:config_enable_openmp": _OPENFHE_LINKOPTS,
+    "@heir//:config_disable_openmp": [],
+})
+
+MAYBE_OPENMP_COPTS = select({
+    "@heir//:config_enable_openmp": _OPENMP_COPTS,
+    "@heir//:config_disable_openmp": [],
+})

bazel/openfhe/openfhe.BUILD

@@ -1,4 +1,5 @@
 # BUILD file for a bazel-native OpenFHE build
+load("@heir//bazel/openfhe:copts.bzl", "MAYBE_OPENFHE_LINKOPTS", "MAYBE_OPENMP_COPTS", "OPENFHE_COPTS", "OPENFHE_DEFINES")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -10,27 +11,6 @@ package(
 
 licenses(["notice"])
 
-OPENFHE_VERSION_MAJOR = 1
-
-OPENFHE_VERSION_MINOR = 11
-
-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,
-]
-
 # This rule exists so that the python frontend can get access to the headers to
 # pass dynamically to clang when building compiled code.
 filegroup(
@@ -48,7 +28,7 @@ cc_library(
         "src/core/lib/**/*.c",
         "src/core/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS + [
+    copts = OPENFHE_COPTS + MAYBE_OPENMP_COPTS + [
         # /utils/blockAllocator/blockAllocator.cpp has misaligned-pointer-use
         "-fno-sanitize=alignment",
     ],
@@ -57,6 +37,7 @@ cc_library(
         "src/core/include",
         "src/core/lib",
     ],
+    linkopts = MAYBE_OPENFHE_LINKOPTS,
     textual_hdrs = glob([
         "src/core/include/**/*.h",
         "src/core/lib/**/*.cpp",
@@ -70,12 +51,13 @@ cc_library(
         "src/binfhe/lib/**/*.c",
         "src/binfhe/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS,
+    copts = OPENFHE_COPTS + MAYBE_OPENMP_COPTS,
     defines = OPENFHE_DEFINES,
     includes = [
         "src/binfhe/include",
         "src/binfhe/lib",
     ],
+    linkopts = MAYBE_OPENFHE_LINKOPTS,
     textual_hdrs = glob(["src/binfhe/include/**/*.h"]),
     deps = [
         "@openfhe//:core",
@@ -87,14 +69,15 @@ cc_library(
     srcs = glob([
         "src/pke/lib/**/*.cpp",
     ]),
-    copts = OPENFHE_COPTS + [
+    copts = OPENFHE_COPTS + MAYBE_OPENMP_COPTS + [
         "-Wno-vla-extension",
     ],
     defines = OPENFHE_DEFINES,
     includes = [
         "src/pke/include",
         "src/pke/lib",
     ],
+    linkopts = MAYBE_OPENFHE_LINKOPTS,
     textual_hdrs = glob([
         "src/pke/include/**/*.h",
         "src/pke/lib/**/*.cpp",

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", "MAYBE_OPENFHE_LINKOPTS", "MAYBE_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 = MAYBE_OPENMP_COPTS,
+        linkopts = MAYBE_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 = MAYBE_OPENMP_COPTS,
+        linkopts = MAYBE_OPENFHE_LINKOPTS,
         **kwargs
     )