diff --git a/onnxruntime/contrib_ops/cpu/activations.h b/onnxruntime/contrib_ops/cpu/activations.h index 4d2baee3bd..f127ea2fba 100644 --- a/onnxruntime/contrib_ops/cpu/activations.h +++ b/onnxruntime/contrib_ops/cpu/activations.h @@ -40,20 +40,35 @@ class Gelu : public OpKernel { const auto* X = context->Input(0); Tensor* Y = context->Output(0, X->Shape()); concurrency::ThreadPool* tp = context->GetOperatorThreadPool(); - const int64_t input_size = X->Shape().Size(); - std::ptrdiff_t batch_size = static_cast(input_size); - //The cost comes from microbenchmark(manual tunning). - const double cost = 10.0; - const T* data = X->template Data(); - T* output = Y->template MutableData(); - concurrency::ThreadPool::TryParallelFor(tp, batch_size, cost, [data, output](ptrdiff_t first, ptrdiff_t last) { - ptrdiff_t len = last - first; - onnxruntime::ConstEigenVectorArrayMap xm(data + first, len); - onnxruntime::EigenVectorArrayMap ym(output + first, len); - ym = xm * static_cast(M_SQRT1_2); - MlasComputeErf(output, output, len); - ym = xm * 0.5f * (ym + 1.0f); - }); + if (nullptr != tp) { + const T* input = X->template Data(); + T* output = Y->template MutableData(); + int task_count = tp->NumThreads(); + int64_t elem_count = X->Shape().Size(); + if (elem_count > task_count) { + tp->SimpleParallelFor(task_count, [input, + output, + elem_count, + task_count](std::ptrdiff_t i) { + int64_t elem_inx_start = i * elem_count / task_count; + int64_t elem_inx_end = (i + 1) * elem_count / task_count; + for (int64_t elem_inx = elem_inx_start; elem_inx < elem_inx_end; elem_inx++) { + output[elem_inx] = input[elem_inx] * static_cast(M_SQRT1_2); + } + MlasComputeErf(output + elem_inx_start, output + elem_inx_start, elem_inx_end - elem_inx_start); + for (int64_t elem_inx = elem_inx_start; elem_inx < elem_inx_end; elem_inx++) { + output[elem_inx] = 0.5f * input[elem_inx] * (output[elem_inx] + 1.0f); + } + }); + return Status::OK(); + } + } + + EIGEN_X_VAR(xm); + EIGEN_Y_VAR(ym); + ym = xm * static_cast(M_SQRT1_2); + MlasComputeErf(Y->template MutableData(), Y->template MutableData(), X->Shape().Size()); + ym = xm * 0.5f * (ym + 1.0f); return Status::OK(); } };