Unverified Commit c51cc82e authored by peizhou001's avatar peizhou001 Committed by GitHub
Browse files

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


Co-authored-by: default avatarUbuntu <ubuntu@ip-172-31-16-19.ap-northeast-1.compute.internal>
parent a1f74982
/**
* 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_
...@@ -6,6 +6,7 @@ ...@@ -6,6 +6,7 @@
#ifndef DGL_RUNTIME_PARALLEL_FOR_H_ #ifndef DGL_RUNTIME_PARALLEL_FOR_H_
#define DGL_RUNTIME_PARALLEL_FOR_H_ #define DGL_RUNTIME_PARALLEL_FOR_H_
#include <dgl/env_variable.h>
#include <dmlc/omp.h> #include <dmlc/omp.h>
#include <algorithm> #include <algorithm>
...@@ -26,13 +27,15 @@ namespace { ...@@ -26,13 +27,15 @@ namespace {
struct DefaultGrainSizeT { struct DefaultGrainSizeT {
size_t grain_size; size_t grain_size;
DefaultGrainSizeT() { DefaultGrainSizeT() : DefaultGrainSizeT(1) {}
auto var = std::getenv("DGL_PARALLEL_FOR_GRAIN_SIZE");
if (!var) { explicit DefaultGrainSizeT(size_t default_grain_size) {
grain_size = 1; auto var = dgl::kDGLParallelForGrainSize;
} else {
if (var) {
grain_size = std::stoul(var); grain_size = std::stoul(var);
} else {
grain_size = default_grain_size;
} }
} }
......
...@@ -85,32 +85,36 @@ COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) { ...@@ -85,32 +85,36 @@ COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos) {
auto res_dst_data = result_dst.Ptr<IdType>(); auto res_dst_data = result_dst.Ptr<IdType>();
auto res_dat_data = result_dat.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
for (IdType i = b; i < e; ++i) { // explicitly via env, use that value instead.
const aten::COOMatrix& coo = coos[i]; size_t grain_size = dgl::runtime::DefaultGrainSizeT(32)();
if (!coo.row_sorted) row_sorted = false; dgl::runtime::parallel_for(
if (!coo.col_sorted) col_sorted = false; 0, coos.size(), grain_size, [&](IdType b, IdType e) {
for (IdType i = b; i < e; ++i) {
auto edges_src = coo.row.Ptr<IdType>(); const aten::COOMatrix& coo = coos[i];
auto edges_dst = coo.col.Ptr<IdType>(); if (!coo.row_sorted) row_sorted = false;
auto edges_dat = coo.data.Ptr<IdType>(); if (!coo.col_sorted) col_sorted = false;
for (IdType j = 0; j < coo.row->shape[0]; j++) { auto edges_src = coo.row.Ptr<IdType>();
res_src_data[prefix_elm[i] + j] = edges_src[j] + prefix_src[i]; auto edges_dst = coo.col.Ptr<IdType>();
} auto edges_dat = coo.data.Ptr<IdType>();
for (IdType j = 0; j < coo.row->shape[0]; j++) { for (IdType j = 0; j < coo.row->shape[0]; j++) {
res_dst_data[prefix_elm[i] + j] = edges_dst[j] + prefix_dst[i]; res_src_data[prefix_elm[i] + j] = edges_src[j] + prefix_src[i];
} }
if (has_data) { for (IdType j = 0; j < coo.row->shape[0]; j++) {
for (IdType j = 0; j < coo.row->shape[0]; j++) { res_dst_data[prefix_elm[i] + j] = edges_dst[j] + prefix_dst[i];
const auto d = (!COOHasData(coo)) ? j : edges_dat[j]; }
res_dat_data[prefix_elm[i] + j] = d + prefix_elm[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( return COOMatrix(
prefix_src[coos.size()], prefix_dst[coos.size()], result_src, result_dst, prefix_src[coos.size()], prefix_dst[coos.size()], result_src, result_dst,
result_dat, row_sorted, col_sorted); result_dat, row_sorted, col_sorted);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment