[Enhancement]Set default graph dataloader thread number (#5479)

Co-authored-by: Ubuntu <ubuntu@ip-172-31-16-19.ap-northeast-1.compute.internal>

include/dgl/env_variable.h

+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file dgl/env_variable.h
+ * @brief Class about envrionment variables.
+ */
+#ifndef DGL_ENV_VARIABLE_H_
+#define DGL_ENV_VARIABLE_H_
+
+#include <cstdlib>
+
+namespace dgl {
+
+static const char* kDGLParallelForGrainSize =
+    std::getenv("DGL_PARALLEL_FOR_GRAIN_SIZE");
+
+}  // namespace dgl
+
+#endif  // DGL_ENV_VARIABLE_H_

include/dgl/runtime/parallel_for.h

@@ -6,6 +6,7 @@
 #ifndef DGL_RUNTIME_PARALLEL_FOR_H_
 #define DGL_RUNTIME_PARALLEL_FOR_H_
 
+#include <dgl/env_variable.h>
 #include <dmlc/omp.h>
 
 #include <algorithm>
@@ -26,13 +27,15 @@ namespace {
 struct DefaultGrainSizeT {
   size_t grain_size;
 
-  DefaultGrainSizeT() {
-    auto var = std::getenv("DGL_PARALLEL_FOR_GRAIN_SIZE");
+  DefaultGrainSizeT() : DefaultGrainSizeT(1) {}
 
-    if (!var) {
-      grain_size = 1;
-    } else {
+  explicit DefaultGrainSizeT(size_t default_grain_size) {
+    auto var = dgl::kDGLParallelForGrainSize;
+
+    if (var) {
       grain_size = std::stoul(var);
+    } else {
+      grain_size = default_grain_size;
     }
   }
 

src/array/cpu/disjoint_union.cc

@@ -85,32 +85,36 @@ COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) {
   auto res_dst_data = result_dst.Ptr<IdType>();
   auto res_dat_data = result_dat.Ptr<IdType>();
 
-  dgl::runtime::parallel_for(0, coos.size(), [&](IdType b, IdType e) {
-    for (IdType i = b; i < e; ++i) {
-      const aten::COOMatrix& coo = coos[i];
-      if (!coo.row_sorted) row_sorted = false;
-      if (!coo.col_sorted) col_sorted = false;
-
-      auto edges_src = coo.row.Ptr<IdType>();
-      auto edges_dst = coo.col.Ptr<IdType>();
-      auto edges_dat = coo.data.Ptr<IdType>();
-
-      for (IdType j = 0; j < coo.row->shape[0]; j++) {
-        res_src_data[prefix_elm[i] + j] = edges_src[j] + prefix_src[i];
-      }
-
-      for (IdType j = 0; j < coo.row->shape[0]; j++) {
-        res_dst_data[prefix_elm[i] + j] = edges_dst[j] + prefix_dst[i];
-      }
-
-      if (has_data) {
-        for (IdType j = 0; j < coo.row->shape[0]; j++) {
-          const auto d = (!COOHasData(coo)) ? j : edges_dat[j];
-          res_dat_data[prefix_elm[i] + j] = d + prefix_elm[i];
+  // 32 is a number obtained from experience. If a user set the grain size
+  // explicitly via env, use that value instead.
+  size_t grain_size = dgl::runtime::DefaultGrainSizeT(32)();
+  dgl::runtime::parallel_for(
+      0, coos.size(), grain_size, [&](IdType b, IdType e) {
+        for (IdType i = b; i < e; ++i) {
+          const aten::COOMatrix& coo = coos[i];
+          if (!coo.row_sorted) row_sorted = false;
+          if (!coo.col_sorted) col_sorted = false;
+
+          auto edges_src = coo.row.Ptr<IdType>();
+          auto edges_dst = coo.col.Ptr<IdType>();
+          auto edges_dat = coo.data.Ptr<IdType>();
+
+          for (IdType j = 0; j < coo.row->shape[0]; j++) {
+            res_src_data[prefix_elm[i] + j] = edges_src[j] + prefix_src[i];
+          }
+
+          for (IdType j = 0; j < coo.row->shape[0]; j++) {
+            res_dst_data[prefix_elm[i] + j] = edges_dst[j] + prefix_dst[i];
+          }
+
+          if (has_data) {
+            for (IdType j = 0; j < coo.row->shape[0]; j++) {
+              const auto d = (!COOHasData(coo)) ? j : edges_dat[j];
+              res_dat_data[prefix_elm[i] + j] = d + prefix_elm[i];
+            }
+          }
         }
-      }
-    }
-  });
+      });
   return COOMatrix(
       prefix_src[coos.size()], prefix_dst[coos.size()], result_src, result_dst,
       result_dat, row_sorted, col_sorted);