From 7aecadedd8df7fa04923457fc19639964bd6e00c Mon Sep 17 00:00:00 2001
From: esp0r <james2002chen@126.com>
Date: Mon, 24 Jun 2024 00:18:04 +0800
Subject: [PATCH] add omp support

---
 setup.py         |  8 ++++++++
 src/omp.cpp      | 18 ++++++++++++++++++
 tests/omptest.py |  2 ++
 tests/test.py    |  1 +
 4 files changed, 29 insertions(+)
 create mode 100644 src/omp.cpp
 create mode 100644 tests/omptest.py

diff --git a/setup.py b/setup.py
index 205144f6..f313415d 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,14 @@
         # Example: passing in the version to the compiled code
         define_macros=[("VERSION_INFO", __version__)],
     ),
+    Pybind11Extension(
+        "ompybind11",
+        ["src/omp.cpp"],
+        # Example: passing in the version to the compiled code
+        extra_compile_args=['-fopenmp'],
+        extra_link_args=['-fopenmp'],
+        define_macros=[("VERSION_INFO", __version__)],
+    ),
 ]
 
 setup(
diff --git a/src/omp.cpp b/src/omp.cpp
new file mode 100644
index 00000000..8f3f5bcc
--- /dev/null
+++ b/src/omp.cpp
@@ -0,0 +1,18 @@
+#include <pybind11/pybind11.h>
+#include <omp.h>
+
+namespace py = pybind11;
+
+int sum_thread_ids() {
+    int sum = 0;
+    #pragma omp parallel shared(sum)
+    {
+        #pragma omp critical
+        sum += omp_get_thread_num();
+    }
+    return sum;
+}
+
+PYBIND11_MODULE(ompybind11, m) {
+    m.def("sum_thread_ids", &sum_thread_ids, "Sum the thread IDs");
+}
\ No newline at end of file
diff --git a/tests/omptest.py b/tests/omptest.py
new file mode 100644
index 00000000..76835c19
--- /dev/null
+++ b/tests/omptest.py
@@ -0,0 +1,2 @@
+import ompybind11
+print(ompybind11.sum_thread_ids())
\ No newline at end of file
diff --git a/tests/test.py b/tests/test.py
index 5a72e5f2..307c5e52 100644
--- a/tests/test.py
+++ b/tests/test.py
@@ -3,3 +3,4 @@
 assert m.__version__ == "0.0.1"
 assert m.add(1, 2) == 3
 assert m.subtract(1, 2) == -1
+