use openmp to accelerate nid mapping (dmlc#242)

* use openmp. * update cmake.
SunYa0 · Dec 5, 2018 · 40506ec · 40506ec
1 parent 378c264
commit 40506ec
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git 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
@@ -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);