From 40506ecc2005d5d8fd7ab1cea39f5e2135160473 Mon Sep 17 00:00:00 2001
From: Da Zheng <zhengda1936@gmail.com>
Date: Wed, 5 Dec 2018 09:55:59 -0800
Subject: [PATCH] use openmp to accelerate nid mapping (#242)

* use openmp.

* update cmake.
---
 CMakeLists.txt        | 4 ++++
 src/graph/graph_op.cc | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index edeb39873ad5..7f5ff2f3a945 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -77,6 +77,10 @@ else(MSVC)
   check_cxx_compiler_flag("-std=c++11"    SUPPORT_CXX11)
   set(CMAKE_C_FLAGS "-O2 -Wall -fPIC ${CMAKE_C_FLAGS}")
   set(CMAKE_CXX_FLAGS "-O2 -Wall -fPIC -std=c++11 ${CMAKE_CXX_FLAGS}")
+  if (${CMAKE_SYSTEM_NAME} MATCHES "Linux")
+    set(CMAKE_C_FLAGS "-fopenmp ${CMAKE_C_FLAGS}")
+    set(CMAKE_CXX_FLAGS "-fopenmp ${CMAKE_CXX_FLAGS}")
+  endif()
 endif(MSVC)
 
 # add source group
diff --git a/src/graph/graph_op.cc b/src/graph/graph_op.cc
index fdb9c704be04..bf47d1b9ec4a 100644
--- a/src/graph/graph_op.cc
+++ b/src/graph/graph_op.cc
@@ -116,6 +116,7 @@ IdArray GraphOp::MapParentIdToSubgraphId(IdArray parent_vids, IdArray query) {
 
   const bool is_sorted = std::is_sorted(parent_data, parent_data + parent_len);
   if (is_sorted) {
+#pragma omp parallel for
     for (int64_t i = 0; i < query_len; i++) {
       const dgl_id_t id = query_data[i];
       const auto it = std::find(parent_data, parent_data + parent_len, id);
@@ -132,6 +133,7 @@ IdArray GraphOp::MapParentIdToSubgraphId(IdArray parent_vids, IdArray query) {
       const dgl_id_t id = parent_data[i];
       parent_map[id] = i;
     }
+#pragma omp parallel for
     for (int64_t i = 0; i < query_len; i++) {
       const dgl_id_t id = query_data[i];
       auto it = parent_map.find(id);