OMP pragma was added to support hetero_graph loop. (dmlc#2148)

* OMP pragma was added to support hetero_graph loop, as follows: pragma omp parallel for * OMP pragma will be enabled iff the gpu is used for Co-authored-by: Zihao Ye <[email protected]>
feixian15 · Oct 9, 2020 · 7d19b33 · 7d19b33
1 parent 8a4520d
commit 7d19b33
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/graph/heterograph_capi.cc b/src/graph/heterograph_capi.cc
@@ -605,7 +605,10 @@ DGL_REGISTER_GLOBAL("heterograph_index._CAPI_DGLHeteroCreateFormat")
 .set_body([] (DGLArgs args, DGLRetValue* rv) {
     HeteroGraphRef hg = args[0];
     dgl_format_code_t code = hg->GetRelationGraph(0)->GetAllowedFormats();
-    for (dgl_type_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
+#if !defined(DGL_USE_CUDA)
+#pragma omp parallel for
+#endif
+    for (int64_t etype = 0; etype < hg->NumEdgeTypes(); ++etype) {
       auto bg = std::dynamic_pointer_cast<UnitGraph>(hg->GetRelationGraph(etype));
       for (auto format : CodeToSparseFormats(code))
         bg->GetFormat(format);