enable OpenMP for OpenFHE & python backend

Author: AlexanderViand

BUILD.bazel

@@ -28,8 +28,7 @@ package_group(
 # OpenMP
 string_flag(
     name = "enable_openmp",
-    # TODO(#1361): re-enable when it's compatible with the Python frontend
-    build_setting_default = "0",
+    build_setting_default = "1",
 )
 
 config_setting(

docs/content/en/docs/getting_started.md

@@ -42,7 +42,8 @@ heir-translate --help
 
 - [Git](https://git-scm.com/)
 - A C++ compiler and linker ([clang](https://clang.llvm.org/) and
-  [lld](https://lld.llvm.org/) or a recent version of `gcc`).
+  [lld](https://lld.llvm.org/) or a recent version of `gcc`). If you want to run
+  OpenFHE with parallelism (enabled by default), you'll also need OpenMP.
 - Bazel via [bazelisk](https://github.com/bazelbuild/bazelisk). The precise
   Bazel version used is in `.bazelversion` in the repository root.
 
@@ -53,7 +54,7 @@ heir-translate --help
   For example, on Ubuntu, these can be installed with
 
 ```bash
-sudo apt-get update && sudo apt-get install clang lld
+sudo apt-get update && sudo apt-get install clang lld libomp-dev
 ```
 
 You can download the latest Bazelisk release, e.g., for linux-amd64 (see the

lib/Target/OpenFhePke/OpenFhePkeTemplates.h

@@ -109,7 +109,13 @@ PYBIND11_MODULE({0}, m) {{
 )cpp";
 // clang-format on
 
-constexpr std::string_view kPybindFunctionTemplate = "m.def(\"{0}\", &{0});";
+// Emit a pybind11 binding that releases the GIL for the duration of the C++
+// function call.  This enables multi-threaded C++ code (e.g. OpenMP parallel
+// regions inside OpenFHE) to run concurrently with the Python interpreter.
+// The `py::call_guard<py::gil_scoped_release>()` helper ensures the GIL is
+// relinquished on entry and re-acquired on exit.
+constexpr std::string_view kPybindFunctionTemplate =
+    "m.def(\"{0}\", &{0}, py::call_guard<py::gil_scoped_release>());";
 
 }  // namespace openfhe
 }  // namespace heir

tests/Emitter/Openfhe/emit_pybind.mlir

@@ -28,9 +28,9 @@
 
 // CHECK: PYBIND11_MODULE(_heir_foo, m) {
 // CHECK:   bind_common(m);
-// CHECK:   m.def("simple_sum", &simple_sum);
-// CHECK:   m.def("simple_sum__encrypt", &simple_sum__encrypt);
-// CHECK:   m.def("simple_sum__decrypt", &simple_sum__decrypt);
+// CHECK:   m.def("simple_sum", &simple_sum, py::call_guard<py::gil_scoped_release>());
+// CHECK:   m.def("simple_sum__encrypt", &simple_sum__encrypt, py::call_guard<py::gil_scoped_release>());
+// CHECK:   m.def("simple_sum__decrypt", &simple_sum__decrypt, py::call_guard<py::gil_scoped_release>());
 // CHECK: }
 
 !Z1095233372161_i64_ = !mod_arith.int<1095233372161 : i64>

tests/Regression/issue_1929.mlir

@@ -2,39 +2,6 @@
 // TODO (#1929): Improve test with better CHECKs
 // TODO (#1641): enable i32-based tests once CKKS pipeline allows integer types
 
-// // CHECK: func.func @secret_loop_index_step
-// func.func @secret_loop_index_step(%arg0: i32 {secret.secret}, %arg1: i32 {secret.secret}) -> i32 {
-//     %c0 = arith.constant 0 : index
-//     %c1 = arith.constant 1 : index
-//     %c1_i32 = arith.constant 1 : i32
-//     %0 = arith.index_cast %arg0 : i32 to index
-//     // CHECK-NOT: scf.for
-//     // CHECK-NOT: arith.cmpi
-//     // CHECK-NOT: math_ext.sign
-//     // CHECK: polynomial.eval
-//     %1 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %c1_i32) -> (i32) {
-//         %2 = arith.muli %arg3, %arg0 : i32
-//         scf.yield %2 : i32
-//     } {upper = 6}
-//     return %1 : i32
-// }
-
-// // When scf.for uses a signless integer step, --convert-secret-for-to-static-for would forget to issue arith.index_cast
-// // CHECK: func.func @secret_loop_signless_integer_step
-// func.func @secret_loop_signless_integer_step(%arg0: i32 {secret.secret}, %arg1: i32 {secret.secret}) -> i32 {
-//     %c0_i32 = arith.constant 0 : i32
-//     %c1_i32 = arith.constant 1 : i32
-//     // CHECK-NOT: scf.for
-//     // CHECK-NOT: arith.cmpi
-//     // CHECK-NOT: math_ext.sign
-//     // CHECK: polynomial.eval
-//     %1 = scf.for %arg2 = %c0_i32 to %arg0 step %c1_i32 iter_args(%arg3 = %c1_i32) -> (i32) : i32 {
-//         %2 = arith.muli %arg3, %arg0 : i32
-//         scf.yield %2 : i32
-//     } {upper = 6}
-//     return %1 : i32
-// }
-
 // CHECK: func.func @float_secret_loop_index_step
 func.func @float_secret_loop_index_step(%arg0: f32 {secret.secret}, %arg1: f32 {secret.secret}) -> f32 {
     %c0 = arith.constant 0 : index
@@ -45,7 +12,7 @@ func.func @float_secret_loop_index_step(%arg0: f32 {secret.secret}, %arg1: f32 {
     // CHECK-NOT: scf.for
     // CHECK-NOT: arith.cmpi
     // CHECK-NOT: math_ext.sign
-    // CHECK: polynomial.eval
+    // CHECK-NOT: polynomial.eval
     %2 = scf.for %arg2 = %c0 to %1 step %c1 iter_args(%arg3 = %c1_f32) -> (f32) {
         %2 = arith.mulf %arg3, %arg0 : f32
         scf.yield %2 : f32
@@ -63,10 +30,33 @@ func.func @float_secret_loop_signless_integer_step(%arg0: f32 {secret.secret}, %
     // CHECK-NOT: scf.for
     // CHECK-NOT: arith.cmpi
     // CHECK-NOT: math_ext.sign
-    // CHECK: polynomial.eval
+    // CHECK-NOT: polynomial.eval
     %1 = scf.for %arg2 = %c0_i32 to %0 step %c1_i32 iter_args(%arg3 = %c1_f32) -> (f32) : i32 {
         %2 = arith.mulf %arg3, %arg0 : f32
         scf.yield %2 : f32
     } {upper = 6}
     return %1 : f32
 }
+
+// func.func @secret_loop_index_step(%arg0: i32 {secret.secret}, %arg1: i32 {secret.secret}) -> i32 {
+//     %c0 = arith.constant 0 : index
+//     %c1 = arith.constant 1 : index
+//     %c1_i32 = arith.constant 1 : i32
+//     %0 = arith.index_cast %arg0 : i32 to index
+//     %1 = scf.for %arg2 = %c0 to %0 step %c1 iter_args(%arg3 = %c1_i32) -> (i32) {
+//         %2 = arith.muli %arg3, %arg0 : i32
+//         scf.yield %2 : i32
+//     } {upper = 6}
+//     return %1 : i32
+// }
+
+// // When scf.for uses a signless integer step, --convert-secret-for-to-static-for would forget to issue arith.index_cast
+// func.func @secret_loop_signless_integer_step(%arg0: i32 {secret.secret}, %arg1: i32 {secret.secret}) -> i32 {
+//     %c0_i32 = arith.constant 0 : i32
+//     %c1_i32 = arith.constant 1 : i32
+//     %1 = scf.for %arg2 = %c0_i32 to %arg0 step %c1_i32 iter_args(%arg3 = %c1_i32) -> (i32) : i32 {
+//         %2 = arith.muli %arg3, %arg0 : i32
+//         scf.yield %2 : i32
+//     } {upper = 6}
+//     return %1 : i32
+// }