From b2c7ccced1c3965bd52692917a8fdd31eddd15d0 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Fri, 10 Mar 2023 15:17:32 -0500
Subject: [PATCH] enable openmp for gelu

Fix #2373.
---
 source/lib/src/gelu.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/source/lib/src/gelu.cc b/source/lib/src/gelu.cc
index c9ff7c9f2b..f352fcd1b5 100644
--- a/source/lib/src/gelu.cc
+++ b/source/lib/src/gelu.cc
@@ -6,6 +6,7 @@
 
 template <typename FPTYPE>
 void deepmd::gelu_cpu(FPTYPE* out, const FPTYPE* xx, const int_64 size) {
+#pragma omp parallel for
   for (int ii = 0; ii < size; ii++) {
     out[ii] = xx[ii] * (FPTYPE)0.5 *
               ((FPTYPE)1.0 +
@@ -19,6 +20,7 @@ void deepmd::gelu_grad_cpu(FPTYPE* out,
                            const FPTYPE* xx,
                            const FPTYPE* dy,
                            const int_64 size) {
+#pragma omp parallel for
   for (int ii = 0; ii < size; ii++) {
     const FPTYPE var =
         tanh((FPTYPE)SQRT_2_PI *
@@ -36,6 +38,7 @@ void deepmd::gelu_grad_grad_cpu(FPTYPE* out,
                                 const FPTYPE* dy,
                                 const FPTYPE* dy_2,
                                 const int_64 size) {
+#pragma omp parallel for
   for (int ii = 0; ii < size; ii++) {
     const FPTYPE var1 =
         tanh((FPTYPE)SQRT_2_PI *