[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,7 +85,11 @@ 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) {
+  // 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;