From a6cfce24ef82a2c398b37d373f470ba249434d43 Mon Sep 17 00:00:00 2001 From: Chao Ma Date: Sun, 2 Dec 2018 22:49:11 +0800 Subject: [PATCH 01/24] Speedup and fix bug in dgl_csr_sampling op --- src/operator/contrib/dgl_graph.cc | 149 +++++++++++++++--------------- 1 file changed, 75 insertions(+), 74 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 6d586755c957..6453fef9dbcd 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -413,21 +413,6 @@ static bool CSRNeighborNonUniformSampleType(const nnvm::NodeAttrs& attrs, return success; } -/* - * Get src vertex and edge id for a destination vertex - */ -static void GetSrcList(const dgl_id_t* val_list, - const dgl_id_t* col_list, - const dgl_id_t* indptr, - const dgl_id_t dst_id, - std::vector* src_list, - std::vector* edge_list) { - for (dgl_id_t i = *(indptr+dst_id); i < *(indptr+dst_id+1); ++i) { - src_list->push_back(col_list[i]); - edge_list->push_back(val_list[i]); - } -} - static void RandomSample(size_t set_size, size_t num, std::vector* out, @@ -464,34 +449,38 @@ static void NegateSet(const std::vector &idxs, /* * Uniform sample */ -static void GetUniformSample(const std::vector& ver_list, - const std::vector& edge_list, +static void GetUniformSample(const dgl_id_t* val_list, + const dgl_id_t* col_list, + const dgl_id_t* indptr, + const dgl_id_t dst_id, const size_t max_num_neighbor, std::vector* out_ver, std::vector* out_edge, unsigned int* seed) { - CHECK_EQ(ver_list.size(), edge_list.size()); + size_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); // Copy ver_list to output - if (ver_list.size() <= max_num_neighbor) { - for (size_t i = 0; i < ver_list.size(); ++i) { - out_ver->push_back(ver_list[i]); - out_edge->push_back(edge_list[i]); + if (ver_len <= max_num_neighbor) { + for (dgl_id_t i = *(indptr+dst_id); i < *(indptr+dst_id+1); ++i) { + out_ver->push_back(col_list[i]); + out_edge->push_back(val_list[i]); } return; } // If we just sample a small number of elements from a large neighbor list. + const dgl_id_t* col_ptr = col_list + *(indptr + dst_id); + const dgl_id_t* val_ptr = val_list + *(indptr + dst_id); std::vector sorted_idxs; - if (ver_list.size() > max_num_neighbor * 2) { + if (ver_len > max_num_neighbor * 2) { sorted_idxs.reserve(max_num_neighbor); - RandomSample(ver_list.size(), max_num_neighbor, &sorted_idxs, seed); + RandomSample(ver_len, max_num_neighbor, &sorted_idxs, seed); std::sort(sorted_idxs.begin(), sorted_idxs.end()); } else { std::vector negate; - negate.reserve(ver_list.size() - max_num_neighbor); - RandomSample(ver_list.size(), ver_list.size() - max_num_neighbor, + negate.reserve(ver_len - max_num_neighbor); + RandomSample(ver_len, ver_len - max_num_neighbor, &negate, seed); std::sort(negate.begin(), negate.end()); - NegateSet(negate, ver_list.size(), &sorted_idxs); + NegateSet(negate, ver_len, &sorted_idxs); } // verify the result. CHECK_EQ(sorted_idxs.size(), max_num_neighbor); @@ -499,8 +488,8 @@ static void GetUniformSample(const std::vector& ver_list, CHECK_GT(sorted_idxs[i], sorted_idxs[i - 1]); } for (auto idx : sorted_idxs) { - out_ver->push_back(ver_list[idx]); - out_edge->push_back(edge_list[idx]); + out_ver->push_back(col_ptr[idx]); + out_edge->push_back(val_ptr[idx]); } } @@ -508,26 +497,30 @@ static void GetUniformSample(const std::vector& ver_list, * Non-uniform sample via ArrayHeap */ static void GetNonUniformSample(const float* probability, - const std::vector& ver_list, - const std::vector& edge_list, + const dgl_id_t* val_list, + const dgl_id_t* col_list, + const dgl_id_t* indptr, + const dgl_id_t dst_id, const size_t max_num_neighbor, std::vector* out_ver, std::vector* out_edge, unsigned int* seed) { - CHECK_EQ(ver_list.size(), edge_list.size()); + size_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); // Copy ver_list to output - if (ver_list.size() <= max_num_neighbor) { - for (size_t i = 0; i < ver_list.size(); ++i) { - out_ver->push_back(ver_list[i]); - out_edge->push_back(edge_list[i]); + if (ver_len <= max_num_neighbor) { + for (dgl_id_t i = *(indptr+dst_id); i < *(indptr+dst_id+1); ++i) { + out_ver->push_back(col_list[i]); + out_edge->push_back(val_list[i]); } return; } // Make sample + const dgl_id_t* col_ptr = col_list + *(indptr + dst_id); + const dgl_id_t* val_ptr = val_list + *(indptr + dst_id); std::vector sp_index(max_num_neighbor); - std::vector sp_prob(ver_list.size()); - for (size_t i = 0; i < ver_list.size(); ++i) { - sp_prob[i] = probability[ver_list[i]]; + std::vector sp_prob(ver_len); + for (size_t i = 0; i < ver_len; ++i) { + sp_prob[i] = probability[col_ptr[i]]; } ArrayHeap arrayHeap(sp_prob); arrayHeap.SampleWithoutReplacement(max_num_neighbor, &sp_index, seed); @@ -535,8 +528,8 @@ static void GetNonUniformSample(const float* probability, out_edge->resize(max_num_neighbor); for (size_t i = 0; i < max_num_neighbor; ++i) { size_t idx = sp_index[i]; - out_ver->at(i) = ver_list[idx]; - out_edge->at(i) = edge_list[idx]; + out_ver->at(i) = col_ptr[idx]; + out_edge->at(i) = val_ptr[idx]; } sort(out_ver->begin(), out_ver->end()); sort(out_edge->begin(), out_edge->end()); @@ -597,14 +590,15 @@ static void SampleSubgraph(const NDArray &csr, node.level = 0; node_queue.push(node); } - std::vector tmp_src_list; - std::vector tmp_edge_list; std::vector tmp_sampled_src_list; std::vector tmp_sampled_edge_list; - std::unordered_map neigh_mp; + // ver_id, position + std::unordered_map neigh_pos; + std::vector neighbor_list; size_t num_edges = 0; + while (!node_queue.empty() && - sub_vertices_count <= max_num_vertices ) { + sub_vertices_count < max_num_vertices) { ver_node& cur_node = node_queue.front(); dgl_id_t dst_id = cur_node.vertex_id; if (cur_node.level < num_hops) { @@ -613,35 +607,42 @@ static void SampleSubgraph(const NDArray &csr, node_queue.pop(); continue; } - tmp_src_list.clear(); - tmp_edge_list.clear(); tmp_sampled_src_list.clear(); tmp_sampled_edge_list.clear(); - GetSrcList(val_list, - col_list, - indptr, - dst_id, - &tmp_src_list, - &tmp_edge_list); if (probability == nullptr) { // uniform-sample - GetUniformSample(tmp_src_list, - tmp_edge_list, + GetUniformSample(val_list, + col_list, + indptr, + dst_id, num_neighbor, &tmp_sampled_src_list, &tmp_sampled_edge_list, &time_seed); } else { // non-uniform-sample GetNonUniformSample(probability, - tmp_src_list, - tmp_edge_list, + val_list, + col_list, + indptr, + dst_id, num_neighbor, &tmp_sampled_src_list, &tmp_sampled_edge_list, &time_seed); } - neigh_mp.insert(std::pair(dst_id, - neigh_list(tmp_sampled_src_list, - tmp_sampled_edge_list))); + CHECK_EQ(tmp_sampled_src_list.size(), + tmp_sampled_edge_list.size()); + size_t pos = neighbor_list.size(); + neigh_pos[dst_id] = pos; + // First we push the size of neighbor vector + neighbor_list.push_back(tmp_sampled_edge_list.size()); + // Then push the vertices + for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { + neighbor_list.push_back(tmp_sampled_src_list[i]); + } + // Finally we push the edge list + for (size_t i = 0; i < tmp_sampled_edge_list.size(); ++i) { + neighbor_list.push_back(tmp_sampled_edge_list[i]); + } num_edges += tmp_sampled_src_list.size(); sub_ver_mp[cur_node.vertex_id] = cur_node.level; for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { @@ -659,11 +660,9 @@ static void SampleSubgraph(const NDArray &csr, node_queue.pop(); continue; } - tmp_sampled_src_list.clear(); - tmp_sampled_edge_list.clear(); - neigh_mp.insert(std::pair(dst_id, - neigh_list(tmp_sampled_src_list, // empty vector - tmp_sampled_edge_list))); // empty vector + size_t pos = neighbor_list.size(); + neigh_pos[dst_id] = pos; + neighbor_list.push_back(0); sub_ver_mp[cur_node.vertex_id] = cur_node.level; } sub_vertices_count++; @@ -720,16 +719,18 @@ static void SampleSubgraph(const NDArray &csr, size_t collected_nedges = 0; for (size_t i = 0; i < num_vertices; i++) { dgl_id_t dst_id = *(out + i); - auto it = neigh_mp.find(dst_id); - const auto &edges = it->second.edges; - const auto &neighs = it->second.neighs; - CHECK_EQ(edges.size(), neighs.size()); - if (!edges.empty()) { - std::copy(edges.begin(), edges.end(), val_list_out + collected_nedges); - std::copy(neighs.begin(), neighs.end(), col_list_out + collected_nedges); - collected_nedges += edges.size(); + size_t pos = neigh_pos[dst_id]; + size_t edge_size = neighbor_list[pos]; + if (edge_size != 0) { + std::copy_n(neighbor_list.begin() + pos + 1, + edge_size, + col_list_out + collected_nedges); + std::copy_n(neighbor_list.begin() + pos + edge_size + 1, + edge_size, + val_list_out + collected_nedges); + collected_nedges += edge_size; } - indptr_out[i+1] = indptr_out[i] + edges.size(); + indptr_out[i+1] = indptr_out[i] + edge_size; } for (dgl_id_t i = num_vertices+1; i <= max_num_vertices; ++i) { indptr_out[i] = indptr_out[i-1]; From a33b6e019af2654a5cd3adc6c50aa0073a2c99ec Mon Sep 17 00:00:00 2001 From: Chao Ma Date: Sun, 2 Dec 2018 23:00:00 +0800 Subject: [PATCH 02/24] Update dgl_graph.cc --- src/operator/contrib/dgl_graph.cc | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 6453fef9dbcd..a875442c3790 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -629,7 +629,7 @@ static void SampleSubgraph(const NDArray &csr, &tmp_sampled_edge_list, &time_seed); } - CHECK_EQ(tmp_sampled_src_list.size(), + CHECK_EQ(tmp_sampled_src_list.size(), tmp_sampled_edge_list.size()); size_t pos = neighbor_list.size(); neigh_pos[dst_id] = pos; @@ -722,11 +722,11 @@ static void SampleSubgraph(const NDArray &csr, size_t pos = neigh_pos[dst_id]; size_t edge_size = neighbor_list[pos]; if (edge_size != 0) { - std::copy_n(neighbor_list.begin() + pos + 1, - edge_size, + std::copy_n(neighbor_list.begin() + pos + 1, + edge_size, col_list_out + collected_nedges); - std::copy_n(neighbor_list.begin() + pos + edge_size + 1, - edge_size, + std::copy_n(neighbor_list.begin() + pos + edge_size + 1, + edge_size, val_list_out + collected_nedges); collected_nedges += edge_size; } From 30415e956174a55ceb56a5032b259ea3d990cfd5 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 13:28:44 +0800 Subject: [PATCH 03/24] simplify functions. --- src/operator/contrib/dgl_graph.cc | 41 ++++++++++++------------------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index a875442c3790..91b21bbcbdcf 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -451,24 +451,20 @@ static void NegateSet(const std::vector &idxs, */ static void GetUniformSample(const dgl_id_t* val_list, const dgl_id_t* col_list, - const dgl_id_t* indptr, - const dgl_id_t dst_id, + const size_t ver_len, const size_t max_num_neighbor, std::vector* out_ver, std::vector* out_edge, unsigned int* seed) { - size_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); // Copy ver_list to output if (ver_len <= max_num_neighbor) { - for (dgl_id_t i = *(indptr+dst_id); i < *(indptr+dst_id+1); ++i) { + for (size_t i = 0; i < ver_len; ++i) { out_ver->push_back(col_list[i]); out_edge->push_back(val_list[i]); } return; } // If we just sample a small number of elements from a large neighbor list. - const dgl_id_t* col_ptr = col_list + *(indptr + dst_id); - const dgl_id_t* val_ptr = val_list + *(indptr + dst_id); std::vector sorted_idxs; if (ver_len > max_num_neighbor * 2) { sorted_idxs.reserve(max_num_neighbor); @@ -488,8 +484,8 @@ static void GetUniformSample(const dgl_id_t* val_list, CHECK_GT(sorted_idxs[i], sorted_idxs[i - 1]); } for (auto idx : sorted_idxs) { - out_ver->push_back(col_ptr[idx]); - out_edge->push_back(val_ptr[idx]); + out_ver->push_back(col_list[idx]); + out_edge->push_back(val_list[idx]); } } @@ -499,28 +495,24 @@ static void GetUniformSample(const dgl_id_t* val_list, static void GetNonUniformSample(const float* probability, const dgl_id_t* val_list, const dgl_id_t* col_list, - const dgl_id_t* indptr, - const dgl_id_t dst_id, + const size_t ver_len, const size_t max_num_neighbor, std::vector* out_ver, std::vector* out_edge, unsigned int* seed) { - size_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); // Copy ver_list to output if (ver_len <= max_num_neighbor) { - for (dgl_id_t i = *(indptr+dst_id); i < *(indptr+dst_id+1); ++i) { + for (size_t i = 0; i < ver_len; ++i) { out_ver->push_back(col_list[i]); out_edge->push_back(val_list[i]); } return; } // Make sample - const dgl_id_t* col_ptr = col_list + *(indptr + dst_id); - const dgl_id_t* val_ptr = val_list + *(indptr + dst_id); std::vector sp_index(max_num_neighbor); std::vector sp_prob(ver_len); for (size_t i = 0; i < ver_len; ++i) { - sp_prob[i] = probability[col_ptr[i]]; + sp_prob[i] = probability[col_list[i]]; } ArrayHeap arrayHeap(sp_prob); arrayHeap.SampleWithoutReplacement(max_num_neighbor, &sp_index, seed); @@ -528,8 +520,8 @@ static void GetNonUniformSample(const float* probability, out_edge->resize(max_num_neighbor); for (size_t i = 0; i < max_num_neighbor; ++i) { size_t idx = sp_index[i]; - out_ver->at(i) = col_ptr[idx]; - out_edge->at(i) = val_ptr[idx]; + out_ver->at(i) = col_list[idx]; + out_edge->at(i) = val_list[idx]; } sort(out_ver->begin(), out_ver->end()); sort(out_edge->begin(), out_edge->end()); @@ -609,21 +601,20 @@ static void SampleSubgraph(const NDArray &csr, } tmp_sampled_src_list.clear(); tmp_sampled_edge_list.clear(); + dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); if (probability == nullptr) { // uniform-sample - GetUniformSample(val_list, - col_list, - indptr, - dst_id, + GetUniformSample(val_list + *(indptr + dst_id), + col_list + *(indptr + dst_id), + ver_len, num_neighbor, &tmp_sampled_src_list, &tmp_sampled_edge_list, &time_seed); } else { // non-uniform-sample GetNonUniformSample(probability, - val_list, - col_list, - indptr, - dst_id, + val_list + *(indptr + dst_id), + col_list + *(indptr + dst_id), + ver_len, num_neighbor, &tmp_sampled_src_list, &tmp_sampled_edge_list, From 39995339f52cf3972cd61565c4eaee0a630516b3 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 13:41:13 +0800 Subject: [PATCH 04/24] avoid adding nodes in the last level in the queue. --- src/operator/contrib/dgl_graph.cc | 88 +++++++++++++++---------------- 1 file changed, 42 insertions(+), 46 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 91b21bbcbdcf..9b447072f05e 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -593,25 +593,24 @@ static void SampleSubgraph(const NDArray &csr, sub_vertices_count < max_num_vertices) { ver_node& cur_node = node_queue.front(); dgl_id_t dst_id = cur_node.vertex_id; - if (cur_node.level < num_hops) { - auto ret = sub_ver_mp.find(dst_id); - if (ret != sub_ver_mp.end()) { - node_queue.pop(); - continue; - } - tmp_sampled_src_list.clear(); - tmp_sampled_edge_list.clear(); - dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); - if (probability == nullptr) { // uniform-sample - GetUniformSample(val_list + *(indptr + dst_id), + auto ret = sub_ver_mp.find(dst_id); + if (ret != sub_ver_mp.end()) { + node_queue.pop(); + continue; + } + tmp_sampled_src_list.clear(); + tmp_sampled_edge_list.clear(); + dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); + if (probability == nullptr) { // uniform-sample + GetUniformSample(val_list + *(indptr + dst_id), col_list + *(indptr + dst_id), ver_len, num_neighbor, &tmp_sampled_src_list, &tmp_sampled_edge_list, &time_seed); - } else { // non-uniform-sample - GetNonUniformSample(probability, + } else { // non-uniform-sample + GetNonUniformSample(probability, val_list + *(indptr + dst_id), col_list + *(indptr + dst_id), ver_len, @@ -619,42 +618,37 @@ static void SampleSubgraph(const NDArray &csr, &tmp_sampled_src_list, &tmp_sampled_edge_list, &time_seed); - } - CHECK_EQ(tmp_sampled_src_list.size(), - tmp_sampled_edge_list.size()); - size_t pos = neighbor_list.size(); - neigh_pos[dst_id] = pos; - // First we push the size of neighbor vector - neighbor_list.push_back(tmp_sampled_edge_list.size()); - // Then push the vertices - for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { - neighbor_list.push_back(tmp_sampled_src_list[i]); - } - // Finally we push the edge list - for (size_t i = 0; i < tmp_sampled_edge_list.size(); ++i) { - neighbor_list.push_back(tmp_sampled_edge_list[i]); - } - num_edges += tmp_sampled_src_list.size(); - sub_ver_mp[cur_node.vertex_id] = cur_node.level; - for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { - auto ret = sub_ver_mp.find(tmp_sampled_src_list[i]); - if (ret == sub_ver_mp.end()) { - ver_node new_node; - new_node.vertex_id = tmp_sampled_src_list[i]; - new_node.level = cur_node.level + 1; + } + CHECK_EQ(tmp_sampled_src_list.size(), tmp_sampled_edge_list.size()); + size_t pos = neighbor_list.size(); + neigh_pos[dst_id] = pos; + // First we push the size of neighbor vector + neighbor_list.push_back(tmp_sampled_edge_list.size()); + // Then push the vertices + for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { + neighbor_list.push_back(tmp_sampled_src_list[i]); + } + // Finally we push the edge list + for (size_t i = 0; i < tmp_sampled_edge_list.size(); ++i) { + neighbor_list.push_back(tmp_sampled_edge_list[i]); + } + num_edges += tmp_sampled_src_list.size(); + sub_ver_mp[cur_node.vertex_id] = cur_node.level; + for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { + auto ret = sub_ver_mp.find(tmp_sampled_src_list[i]); + if (ret == sub_ver_mp.end()) { + ver_node new_node; + new_node.vertex_id = tmp_sampled_src_list[i]; + new_node.level = cur_node.level + 1; + if (new_node.level < num_hops) { node_queue.push(new_node); + } else { + size_t pos = neighbor_list.size(); + neigh_pos[new_node.vertex_id] = pos; + neighbor_list.push_back(0); + sub_ver_mp[new_node.vertex_id] = new_node.level; } } - } else { // vertex without any neighbor - auto ret = sub_ver_mp.find(dst_id); - if (ret != sub_ver_mp.end()) { - node_queue.pop(); - continue; - } - size_t pos = neighbor_list.size(); - neigh_pos[dst_id] = pos; - neighbor_list.push_back(0); - sub_ver_mp[cur_node.vertex_id] = cur_node.level; } sub_vertices_count++; node_queue.pop(); @@ -711,7 +705,9 @@ static void SampleSubgraph(const NDArray &csr, for (size_t i = 0; i < num_vertices; i++) { dgl_id_t dst_id = *(out + i); size_t pos = neigh_pos[dst_id]; + CHECK_LT(pos, neighbor_list.size()); size_t edge_size = neighbor_list[pos]; + CHECK_LE(pos + edge_size * 2 + 1, neighbor_list.size()); if (edge_size != 0) { std::copy_n(neighbor_list.begin() + pos + 1, edge_size, From 742a8b4581acff0549ae0488e14b2a7c58e168bd Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 14:25:17 +0800 Subject: [PATCH 05/24] remove a hashtable lookup in neigh_pos. --- src/operator/contrib/dgl_graph.cc | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 9b447072f05e..d45bd2e9b95f 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -643,9 +643,9 @@ static void SampleSubgraph(const NDArray &csr, if (new_node.level < num_hops) { node_queue.push(new_node); } else { - size_t pos = neighbor_list.size(); - neigh_pos[new_node.vertex_id] = pos; - neighbor_list.push_back(0); + // This vertex is in the last level. It doesn't have edges. + // If a vertex doesn't contain an edge, we don't need to add the vertex + // in neigh_pos or neighbor_list. sub_ver_mp[new_node.vertex_id] = new_node.level; } } @@ -704,11 +704,16 @@ static void SampleSubgraph(const NDArray &csr, size_t collected_nedges = 0; for (size_t i = 0; i < num_vertices; i++) { dgl_id_t dst_id = *(out + i); - size_t pos = neigh_pos[dst_id]; - CHECK_LT(pos, neighbor_list.size()); - size_t edge_size = neighbor_list[pos]; - CHECK_LE(pos + edge_size * 2 + 1, neighbor_list.size()); - if (edge_size != 0) { + // If a vertex is in sub_ver_mp but not in neigh_pos, this vertex must not + // have edges. + size_t edge_size = 0; + auto it = neigh_pos.find(dst_id); + if (it != neigh_pos.end()) { + size_t pos = it->second; + CHECK_LT(pos, neighbor_list.size()); + edge_size = neighbor_list[pos]; + CHECK_LE(pos + edge_size * 2 + 1, neighbor_list.size()); + std::copy_n(neighbor_list.begin() + pos + 1, edge_size, col_list_out + collected_nedges); From ffa4097e973a8914a27f4aa7bc2ab024a692af36 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 14:34:33 +0800 Subject: [PATCH 06/24] reduce a hashtable lookup in sub_ver_mp. --- src/operator/contrib/dgl_graph.cc | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index d45bd2e9b95f..f82e67cde71e 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -581,6 +581,7 @@ static void SampleSubgraph(const NDArray &csr, node.vertex_id = seed[i]; node.level = 0; node_queue.push(node); + sub_ver_mp[seed[i]] = 0; } std::vector tmp_sampled_src_list; std::vector tmp_sampled_edge_list; @@ -593,11 +594,6 @@ static void SampleSubgraph(const NDArray &csr, sub_vertices_count < max_num_vertices) { ver_node& cur_node = node_queue.front(); dgl_id_t dst_id = cur_node.vertex_id; - auto ret = sub_ver_mp.find(dst_id); - if (ret != sub_ver_mp.end()) { - node_queue.pop(); - continue; - } tmp_sampled_src_list.clear(); tmp_sampled_edge_list.clear(); dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); @@ -633,7 +629,6 @@ static void SampleSubgraph(const NDArray &csr, neighbor_list.push_back(tmp_sampled_edge_list[i]); } num_edges += tmp_sampled_src_list.size(); - sub_ver_mp[cur_node.vertex_id] = cur_node.level; for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { auto ret = sub_ver_mp.find(tmp_sampled_src_list[i]); if (ret == sub_ver_mp.end()) { @@ -642,12 +637,14 @@ static void SampleSubgraph(const NDArray &csr, new_node.level = cur_node.level + 1; if (new_node.level < num_hops) { node_queue.push(new_node); - } else { - // This vertex is in the last level. It doesn't have edges. - // If a vertex doesn't contain an edge, we don't need to add the vertex - // in neigh_pos or neighbor_list. - sub_ver_mp[new_node.vertex_id] = new_node.level; } + // We need to add the neighbor in the hashtable here. This ensures that + // the vertex in the queue is unique. If we see a vertex before, we don't + // need to add it to the queue again. + sub_ver_mp[new_node.vertex_id] = new_node.level; + // This vertex is in the last level. It doesn't have edges. + // If a vertex doesn't contain an edge, we don't need to add the vertex + // in neigh_pos or neighbor_list. } } sub_vertices_count++; From e0135361fcaac30954051cf503c4d8d6d541f8a5 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 14:51:14 +0800 Subject: [PATCH 07/24] merge copying vids and layers. --- src/operator/contrib/dgl_graph.cc | 46 +++++++++++++++---------------- 1 file changed, 22 insertions(+), 24 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index f82e67cde71e..eaa98bc52970 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -556,9 +556,9 @@ static void SampleSubgraph(const NDArray &csr, float* sub_prob, const NDArray &sub_layer, const float* probability, - dgl_id_t num_hops, - dgl_id_t num_neighbor, - dgl_id_t max_num_vertices) { + int num_hops, + size_t num_neighbor, + size_t max_num_vertices) { unsigned int time_seed = time(nullptr); size_t num_seeds = seed_arr.shape().Size(); CHECK_GE(max_num_vertices, num_seeds); @@ -571,7 +571,7 @@ static void SampleSubgraph(const NDArray &csr, dgl_id_t* out_layer = sub_layer.data().dptr(); // BFS traverse the graph and sample vertices - dgl_id_t sub_vertices_count = 0; + size_t sub_vertices_count = 0; // std::unordered_map sub_ver_mp; std::queue node_queue; @@ -652,23 +652,30 @@ static void SampleSubgraph(const NDArray &csr, } // Copy sub_ver_mp to output[0] - size_t idx = 0; - for (auto& data : sub_ver_mp) { - *(out+idx) = data.first; - idx++; - } + // Copy layer + std::vector > order_map; + order_map.reserve(sub_ver_mp.size()); + for (auto& data : sub_ver_mp) + order_map.push_back(std::pair(data.first, data.second)); size_t num_vertices = sub_ver_mp.size(); - std::sort(out, out + num_vertices); - // The rest data will be set to -1 - for (dgl_id_t i = idx; i < max_num_vertices; ++i) { - *(out+i) = -1; + std::sort(order_map.begin(), order_map.end(), [](const std::pair &a1, std::pair &a2) { + return a1.first < a2.second; + }); + for (size_t i = 0; i < order_map.size(); i++) { + out[i] = order_map[i].first; + out_layer[i] = order_map[i].second; + } + for (size_t i = order_map.size(); i < max_num_vertices; i++) { + out[i] = -1; + out_layer[i] = -1; } // The last element stores the actual // number of vertices in the subgraph. out[max_num_vertices] = sub_ver_mp.size(); + // Copy sub_probability if (sub_prob != nullptr) { - for (dgl_id_t i = 0; i < max_num_vertices; ++i) { + for (size_t i = 0; i < max_num_vertices; ++i) { dgl_id_t idx = out[i]; if (idx != -1) { sub_prob[i] = probability[idx]; @@ -677,15 +684,6 @@ static void SampleSubgraph(const NDArray &csr, } } } - // Copy layer - for (dgl_id_t i = 0; i < max_num_vertices; ++i) { - dgl_id_t idx = out[i]; - if (idx != -1) { - out_layer[i] = sub_ver_mp[idx]; - } else { - out_layer[i] = -1; - } - } // Construct sub_csr_graph TShape shape_1(1); TShape shape_2(1); @@ -721,7 +719,7 @@ static void SampleSubgraph(const NDArray &csr, } indptr_out[i+1] = indptr_out[i] + edge_size; } - for (dgl_id_t i = num_vertices+1; i <= max_num_vertices; ++i) { + for (size_t i = num_vertices+1; i <= max_num_vertices; ++i) { indptr_out[i] = indptr_out[i-1]; } } From 9cfb11c3212d4f1cf6c1e8dfc83dd0be5e04c754 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 16:21:00 +0800 Subject: [PATCH 08/24] fix a bug. --- src/operator/contrib/dgl_graph.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index eaa98bc52970..e5dda5b6180d 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -659,7 +659,7 @@ static void SampleSubgraph(const NDArray &csr, order_map.push_back(std::pair(data.first, data.second)); size_t num_vertices = sub_ver_mp.size(); std::sort(order_map.begin(), order_map.end(), [](const std::pair &a1, std::pair &a2) { - return a1.first < a2.second; + return a1.first < a2.first; }); for (size_t i = 0; i < order_map.size(); i++) { out[i] = order_map[i].first; From c2f99c55e8a4d9651bb973f08583c6229ed738ce Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 16:20:34 +0800 Subject: [PATCH 09/24] reduce hashtable lookup when writing to output csr. --- src/operator/contrib/dgl_graph.cc | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index e5dda5b6180d..e9b196b10ec9 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -697,14 +697,28 @@ static void SampleSubgraph(const NDArray &csr, dgl_id_t* indptr_out = sub_csr.aux_data(0).dptr(); indptr_out[0] = 0; size_t collected_nedges = 0; + + // This might look weird. However, it should be able to reduce hashtable + // lookup by a lot (depends on how many neighbors we sample). + // By putting vertices with neighbors in a sorted array, we can easily + // find which vertices have neighbors. + std::vector > order_map_with_neighs; + order_map_with_neighs.reserve(neigh_pos.size()); + for (auto& data : neigh_pos) + order_map_with_neighs.push_back(std::pair(data.first, data.second)); + std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), [](const std::pair &a1, std::pair &a2) { + return a1.first < a2.first; + }); + + size_t idx_with_neigh = 0; for (size_t i = 0; i < num_vertices; i++) { dgl_id_t dst_id = *(out + i); // If a vertex is in sub_ver_mp but not in neigh_pos, this vertex must not // have edges. size_t edge_size = 0; - auto it = neigh_pos.find(dst_id); - if (it != neigh_pos.end()) { - size_t pos = it->second; + if (idx_with_neigh < order_map_with_neighs.size() + && dst_id == order_map_with_neighs[idx_with_neigh].first) { + size_t pos = order_map_with_neighs[idx_with_neigh].second; CHECK_LT(pos, neighbor_list.size()); edge_size = neighbor_list[pos]; CHECK_LE(pos + edge_size * 2 + 1, neighbor_list.size()); @@ -716,6 +730,7 @@ static void SampleSubgraph(const NDArray &csr, edge_size, val_list_out + collected_nedges); collected_nedges += edge_size; + idx_with_neigh++; } indptr_out[i+1] = indptr_out[i] + edge_size; } From 05257e692f78f2f3bd270f92fce1fd6637c0249e Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 21:34:41 +0800 Subject: [PATCH 10/24] limit the number of sampled vertices. --- src/operator/contrib/dgl_graph.cc | 9 ++++++--- tests/python/unittest/test_dgl_graph.py | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index e9b196b10ec9..f9cd4ea7febb 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -571,7 +571,6 @@ static void SampleSubgraph(const NDArray &csr, dgl_id_t* out_layer = sub_layer.data().dptr(); // BFS traverse the graph and sample vertices - size_t sub_vertices_count = 0; // std::unordered_map sub_ver_mp; std::queue node_queue; @@ -591,7 +590,7 @@ static void SampleSubgraph(const NDArray &csr, size_t num_edges = 0; while (!node_queue.empty() && - sub_vertices_count < max_num_vertices) { + sub_ver_mp.size() < max_num_vertices) { ver_node& cur_node = node_queue.front(); dgl_id_t dst_id = cur_node.vertex_id; tmp_sampled_src_list.clear(); @@ -638,6 +637,9 @@ static void SampleSubgraph(const NDArray &csr, if (new_node.level < num_hops) { node_queue.push(new_node); } + // If we have sampled the max number of vertices, we have to stop. + if (sub_ver_mp.size() >= max_num_vertices) + break; // We need to add the neighbor in the hashtable here. This ensures that // the vertex in the queue is unique. If we see a vertex before, we don't // need to add it to the queue again. @@ -647,9 +649,10 @@ static void SampleSubgraph(const NDArray &csr, // in neigh_pos or neighbor_list. } } - sub_vertices_count++; node_queue.pop(); } + if (!node_queue.empty()) + LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices"; // Copy sub_ver_mp to output[0] // Copy layer diff --git a/tests/python/unittest/test_dgl_graph.py b/tests/python/unittest/test_dgl_graph.py index 069fef6e32f0..f7eac24de8b0 100644 --- a/tests/python/unittest/test_dgl_graph.py +++ b/tests/python/unittest/test_dgl_graph.py @@ -101,9 +101,9 @@ def test_uniform_sample(): check_compact(out[1], out[0], num_nodes) seed = mx.nd.array([0], dtype=np.int64) - out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=2, num_neighbor=1, max_num_vertices=4) + out = mx.nd.contrib.dgl_csr_neighbor_uniform_sample(a, seed, num_args=2, num_hops=2, num_neighbor=1, max_num_vertices=3) assert (len(out) == 3) - check_uniform(out, num_hops=2, max_num_vertices=4) + check_uniform(out, num_hops=2, max_num_vertices=3) num_nodes = out[0][-1].asnumpy() assert num_nodes > 0 assert num_nodes < len(out[0]) From f8e4bbba53ceb9188a9f55a113551d5999e110dc Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 22:39:01 +0800 Subject: [PATCH 11/24] fix lint. --- src/operator/contrib/dgl_graph.cc | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index f9cd4ea7febb..8a70915b16e3 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -651,8 +651,10 @@ static void SampleSubgraph(const NDArray &csr, } node_queue.pop(); } - if (!node_queue.empty()) - LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices"; + if (!node_queue.empty()) { + LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices\n" + << "Please use a smaller number of seeds or a small neighborhood"; + } // Copy sub_ver_mp to output[0] // Copy layer @@ -661,7 +663,8 @@ static void SampleSubgraph(const NDArray &csr, for (auto& data : sub_ver_mp) order_map.push_back(std::pair(data.first, data.second)); size_t num_vertices = sub_ver_mp.size(); - std::sort(order_map.begin(), order_map.end(), [](const std::pair &a1, std::pair &a2) { + std::sort(order_map.begin(), order_map.end(), + [](const std::pair &a1, std::pair &a2) { return a1.first < a2.first; }); for (size_t i = 0; i < order_map.size(); i++) { @@ -709,7 +712,8 @@ static void SampleSubgraph(const NDArray &csr, order_map_with_neighs.reserve(neigh_pos.size()); for (auto& data : neigh_pos) order_map_with_neighs.push_back(std::pair(data.first, data.second)); - std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), [](const std::pair &a1, std::pair &a2) { + std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), + [](const std::pair &a1, std::pair &a2) { return a1.first < a2.first; }); From 87b0039b394d9eb888381f51c136e5e260c91203 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 8 Dec 2018 22:50:46 +0800 Subject: [PATCH 12/24] fix a compile error. --- src/operator/contrib/dgl_graph.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 8a70915b16e3..ee562f5a9bed 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -664,7 +664,7 @@ static void SampleSubgraph(const NDArray &csr, order_map.push_back(std::pair(data.first, data.second)); size_t num_vertices = sub_ver_mp.size(); std::sort(order_map.begin(), order_map.end(), - [](const std::pair &a1, std::pair &a2) { + [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; }); for (size_t i = 0; i < order_map.size(); i++) { @@ -713,7 +713,7 @@ static void SampleSubgraph(const NDArray &csr, for (auto& data : neigh_pos) order_map_with_neighs.push_back(std::pair(data.first, data.second)); std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), - [](const std::pair &a1, std::pair &a2) { + [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; }); From 73042178a35bd99c95496a984a097c0a95c31c3e Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 09:43:29 +0800 Subject: [PATCH 13/24] fix compile error. --- src/operator/contrib/dgl_graph.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index ee562f5a9bed..d33587bfac0a 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -661,7 +661,7 @@ static void SampleSubgraph(const NDArray &csr, std::vector > order_map; order_map.reserve(sub_ver_mp.size()); for (auto& data : sub_ver_mp) - order_map.push_back(std::pair(data.first, data.second)); + order_map.emplace_back(data.first, data.second); size_t num_vertices = sub_ver_mp.size(); std::sort(order_map.begin(), order_map.end(), [](const std::pair &a1, const std::pair &a2) { @@ -711,7 +711,7 @@ static void SampleSubgraph(const NDArray &csr, std::vector > order_map_with_neighs; order_map_with_neighs.reserve(neigh_pos.size()); for (auto& data : neigh_pos) - order_map_with_neighs.push_back(std::pair(data.first, data.second)); + order_map_with_neighs.push_back(data.first, data.second); std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; From 3c9110a233b8bd7d1bb6618ce621944787b60ffa Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 09:47:20 +0800 Subject: [PATCH 14/24] fix compile. --- src/operator/contrib/dgl_graph.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index d33587bfac0a..fa3e276a96b9 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -711,7 +711,7 @@ static void SampleSubgraph(const NDArray &csr, std::vector > order_map_with_neighs; order_map_with_neighs.reserve(neigh_pos.size()); for (auto& data : neigh_pos) - order_map_with_neighs.push_back(data.first, data.second); + order_map_with_neighs.emplace_back(data.first, data.second); std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; From 6b4fa257b00c9043b44db4bc063a8d7c041bf1df Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 15:47:28 +0800 Subject: [PATCH 15/24] remove one hashtable lookup per vertex and hashtable iteration. --- src/operator/contrib/dgl_graph.cc | 51 ++++++++++++++++--------------- 1 file changed, 27 insertions(+), 24 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index fa3e276a96b9..cc22c8508923 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -572,15 +572,21 @@ static void SampleSubgraph(const NDArray &csr, // BFS traverse the graph and sample vertices // - std::unordered_map sub_ver_mp; + std::unordered_set sub_ver_mp; + std::vector > sub_vers; + sub_vers.reserve(num_seeds * 10); std::queue node_queue; // add seed vertices for (size_t i = 0; i < num_seeds; ++i) { - ver_node node; - node.vertex_id = seed[i]; - node.level = 0; - node_queue.push(node); - sub_ver_mp[seed[i]] = 0; + auto ret = sub_ver_mp.insert(seed[i]); + // If the vertex is inserted successfully. + if (ret.second) { + ver_node node; + node.vertex_id = seed[i]; + node.level = 0; + node_queue.push(node); + sub_vers.emplace_back(seed[i], 0); + } } std::vector tmp_sampled_src_list; std::vector tmp_sampled_edge_list; @@ -629,21 +635,22 @@ static void SampleSubgraph(const NDArray &csr, } num_edges += tmp_sampled_src_list.size(); for (size_t i = 0; i < tmp_sampled_src_list.size(); ++i) { - auto ret = sub_ver_mp.find(tmp_sampled_src_list[i]); - if (ret == sub_ver_mp.end()) { + // If we have sampled the max number of vertices, we have to stop. + if (sub_ver_mp.size() >= max_num_vertices) + break; + // We need to add the neighbor in the hashtable here. This ensures that + // the vertex in the queue is unique. If we see a vertex before, we don't + // need to add it to the queue again. + auto ret = sub_ver_mp.insert(tmp_sampled_src_list[i]); + // If the sampled neighbor is inserted to the map successfully. + if (ret.second) { ver_node new_node; new_node.vertex_id = tmp_sampled_src_list[i]; new_node.level = cur_node.level + 1; if (new_node.level < num_hops) { node_queue.push(new_node); } - // If we have sampled the max number of vertices, we have to stop. - if (sub_ver_mp.size() >= max_num_vertices) - break; - // We need to add the neighbor in the hashtable here. This ensures that - // the vertex in the queue is unique. If we see a vertex before, we don't - // need to add it to the queue again. - sub_ver_mp[new_node.vertex_id] = new_node.level; + sub_vers.emplace_back(new_node.vertex_id, new_node.level); // This vertex is in the last level. It doesn't have edges. // If a vertex doesn't contain an edge, we don't need to add the vertex // in neigh_pos or neighbor_list. @@ -658,20 +665,16 @@ static void SampleSubgraph(const NDArray &csr, // Copy sub_ver_mp to output[0] // Copy layer - std::vector > order_map; - order_map.reserve(sub_ver_mp.size()); - for (auto& data : sub_ver_mp) - order_map.emplace_back(data.first, data.second); size_t num_vertices = sub_ver_mp.size(); - std::sort(order_map.begin(), order_map.end(), + std::sort(sub_vers.begin(), sub_vers.end(), [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; }); - for (size_t i = 0; i < order_map.size(); i++) { - out[i] = order_map[i].first; - out_layer[i] = order_map[i].second; + for (size_t i = 0; i < sub_vers.size(); i++) { + out[i] = sub_vers[i].first; + out_layer[i] = sub_vers[i].second; } - for (size_t i = order_map.size(); i < max_num_vertices; i++) { + for (size_t i = sub_vers.size(); i < max_num_vertices; i++) { out[i] = -1; out_layer[i] = -1; } From b85a0151c82e2451a7defc42adbf96c9d278370b Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 16:13:35 +0800 Subject: [PATCH 16/24] remove queue. --- src/operator/contrib/dgl_graph.cc | 55 +++++++++++++------------------ 1 file changed, 23 insertions(+), 32 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index cc22c8508923..f130fb07aeff 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -527,14 +527,6 @@ static void GetNonUniformSample(const float* probability, sort(out_edge->begin(), out_edge->end()); } -/* - * This is used for BFS traversal - */ -struct ver_node { - dgl_id_t vertex_id; - int level; -}; - /* * Used for subgraph sampling */ @@ -575,16 +567,11 @@ static void SampleSubgraph(const NDArray &csr, std::unordered_set sub_ver_mp; std::vector > sub_vers; sub_vers.reserve(num_seeds * 10); - std::queue node_queue; // add seed vertices for (size_t i = 0; i < num_seeds; ++i) { auto ret = sub_ver_mp.insert(seed[i]); // If the vertex is inserted successfully. if (ret.second) { - ver_node node; - node.vertex_id = seed[i]; - node.level = 0; - node_queue.push(node); sub_vers.emplace_back(seed[i], 0); } } @@ -595,10 +582,21 @@ static void SampleSubgraph(const NDArray &csr, std::vector neighbor_list; size_t num_edges = 0; - while (!node_queue.empty() && + // sub_vers is used both as a node collection and a queue. + // In the while loop, we iterate over sub_vers and new nodes are added to the vector. + // A vertex in the vector only needs to be accessed once. If there is a vertex behind idx + // isn't in the last level, we will sample its neighbors. If not, the while loop terminates. + size_t idx = 0; + while (idx < sub_vers.size() && sub_ver_mp.size() < max_num_vertices) { - ver_node& cur_node = node_queue.front(); - dgl_id_t dst_id = cur_node.vertex_id; + dgl_id_t dst_id = sub_vers[idx].first; + int cur_node_level = sub_vers[idx].second; + idx++; + // If the node is in the last level, we don't need to sample neighbors + // from this node. + if (cur_node_level >= num_hops) + continue; + tmp_sampled_src_list.clear(); tmp_sampled_edge_list.clear(); dgl_id_t ver_len = *(indptr+dst_id+1) - *(indptr+dst_id); @@ -643,24 +641,17 @@ static void SampleSubgraph(const NDArray &csr, // need to add it to the queue again. auto ret = sub_ver_mp.insert(tmp_sampled_src_list[i]); // If the sampled neighbor is inserted to the map successfully. - if (ret.second) { - ver_node new_node; - new_node.vertex_id = tmp_sampled_src_list[i]; - new_node.level = cur_node.level + 1; - if (new_node.level < num_hops) { - node_queue.push(new_node); - } - sub_vers.emplace_back(new_node.vertex_id, new_node.level); - // This vertex is in the last level. It doesn't have edges. - // If a vertex doesn't contain an edge, we don't need to add the vertex - // in neigh_pos or neighbor_list. - } + if (ret.second) + sub_vers.emplace_back(tmp_sampled_src_list[i], cur_node_level + 1); } - node_queue.pop(); } - if (!node_queue.empty()) { - LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices\n" - << "Please use a smaller number of seeds or a small neighborhood"; + // Let's check if there is a vertex that we haven't sampled its neighbors. + for (; idx < sub_vers.size(); idx++) { + if (sub_vers[idx].second < num_hops) { + LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices\n" + << "Please use a smaller number of seeds or a small neighborhood"; + break; + } } // Copy sub_ver_mp to output[0] From 9f3adf47bd9a505b6539f6dcb9cc2b0336f086df Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 16:20:55 +0800 Subject: [PATCH 17/24] use vector for neigh_pos. --- src/operator/contrib/dgl_graph.cc | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index f130fb07aeff..fa98cd22943c 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -578,7 +578,8 @@ static void SampleSubgraph(const NDArray &csr, std::vector tmp_sampled_src_list; std::vector tmp_sampled_edge_list; // ver_id, position - std::unordered_map neigh_pos; + std::vector > neigh_pos; + neigh_pos.reserve(num_seeds); std::vector neighbor_list; size_t num_edges = 0; @@ -620,7 +621,7 @@ static void SampleSubgraph(const NDArray &csr, } CHECK_EQ(tmp_sampled_src_list.size(), tmp_sampled_edge_list.size()); size_t pos = neighbor_list.size(); - neigh_pos[dst_id] = pos; + neigh_pos.emplace_back(dst_id, pos); // First we push the size of neighbor vector neighbor_list.push_back(tmp_sampled_edge_list.size()); // Then push the vertices @@ -698,28 +699,20 @@ static void SampleSubgraph(const NDArray &csr, indptr_out[0] = 0; size_t collected_nedges = 0; - // This might look weird. However, it should be able to reduce hashtable - // lookup by a lot (depends on how many neighbors we sample). - // By putting vertices with neighbors in a sorted array, we can easily - // find which vertices have neighbors. - std::vector > order_map_with_neighs; - order_map_with_neighs.reserve(neigh_pos.size()); - for (auto& data : neigh_pos) - order_map_with_neighs.emplace_back(data.first, data.second); - std::sort(order_map_with_neighs.begin(), order_map_with_neighs.end(), - [](const std::pair &a1, const std::pair &a2) { + // Both the out array and neigh_pos are sorted. By scanning the two arrays, we can see + // which vertices have neighbors and which don't. + std::sort(neigh_pos.begin(), neigh_pos.end(), + [](const std::pair &a1, const std::pair &a2) { return a1.first < a2.first; }); - size_t idx_with_neigh = 0; for (size_t i = 0; i < num_vertices; i++) { dgl_id_t dst_id = *(out + i); // If a vertex is in sub_ver_mp but not in neigh_pos, this vertex must not // have edges. size_t edge_size = 0; - if (idx_with_neigh < order_map_with_neighs.size() - && dst_id == order_map_with_neighs[idx_with_neigh].first) { - size_t pos = order_map_with_neighs[idx_with_neigh].second; + if (idx_with_neigh < neigh_pos.size() && dst_id == neigh_pos[idx_with_neigh].first) { + size_t pos = neigh_pos[idx_with_neigh].second; CHECK_LT(pos, neighbor_list.size()); edge_size = neighbor_list[pos]; CHECK_LE(pos + edge_size * 2 + 1, neighbor_list.size()); From fc5f3e2cd591f56e1b9de48d28ed61c2d07e665c Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 17:11:28 +0800 Subject: [PATCH 18/24] fix lint --- src/operator/contrib/dgl_graph.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index fa98cd22943c..309ed920c6da 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -649,7 +649,8 @@ static void SampleSubgraph(const NDArray &csr, // Let's check if there is a vertex that we haven't sampled its neighbors. for (; idx < sub_vers.size(); idx++) { if (sub_vers[idx].second < num_hops) { - LOG(WARNING) << "The sampling is truncated because we have reached the max number of vertices\n" + LOG(WARNING) + << "The sampling is truncated because we have reached the max number of vertices\n" << "Please use a smaller number of seeds or a small neighborhood"; break; } From 043c1b4ac9caf5279789d4ec864d4e9e29a900f9 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sun, 9 Dec 2018 10:03:30 +0000 Subject: [PATCH 19/24] avoid init output arrays. --- src/operator/contrib/dgl_graph.cc | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 309ed920c6da..5a8851c4f578 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -667,23 +667,15 @@ static void SampleSubgraph(const NDArray &csr, out[i] = sub_vers[i].first; out_layer[i] = sub_vers[i].second; } - for (size_t i = sub_vers.size(); i < max_num_vertices; i++) { - out[i] = -1; - out_layer[i] = -1; - } // The last element stores the actual // number of vertices in the subgraph. out[max_num_vertices] = sub_ver_mp.size(); // Copy sub_probability if (sub_prob != nullptr) { - for (size_t i = 0; i < max_num_vertices; ++i) { + for (size_t i = 0; i < sub_ver_mp.size(); ++i) { dgl_id_t idx = out[i]; - if (idx != -1) { - sub_prob[i] = probability[idx]; - } else { - sub_prob[i] = -1; - } + sub_prob[i] = probability[idx]; } } // Construct sub_csr_graph From 834f56ea138924efc5456a8a62ac3edb6c5ae9e9 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Wed, 12 Dec 2018 17:23:17 +0800 Subject: [PATCH 20/24] fix tests. --- tests/python/unittest/test_dgl_graph.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/tests/python/unittest/test_dgl_graph.py b/tests/python/unittest/test_dgl_graph.py index f7eac24de8b0..75c877577a42 100644 --- a/tests/python/unittest/test_dgl_graph.py +++ b/tests/python/unittest/test_dgl_graph.py @@ -32,13 +32,10 @@ def check_uniform(out, num_hops, max_num_vertices): layer = out[2] # check sample_id assert (len(sample_id) == max_num_vertices+1) - count = 0 - for data in sample_id: - if data != -1: - count = count + 1 - assert (mx.nd.array([count-1], dtype=np.int64) == sample_id[-1]) + num_vertices = sample_id[-1].asnumpy()[0] # check sub_csr sub_csr.check_format(full_check=True) + assert np.all((sub_csr.indptr[num_vertices:] == sub_csr.indptr[num_vertices]).asnumpy()) # check layer for data in layer: assert(data <= num_hops) @@ -50,13 +47,10 @@ def check_non_uniform(out, num_hops, max_num_vertices): layer = out[3] # check sample_id assert (len(sample_id) == max_num_vertices+1) - count = 0 - for data in sample_id: - if data != -1: - count = count + 1 - assert (mx.nd.array([count-1], dtype=np.int64) == sample_id[-1]) + num_vertices = sample_id[-1].asnumpy()[0] # check sub_csr sub_csr.check_format(full_check=True) + assert np.all((sub_csr.indptr[num_vertices:] == sub_csr.indptr[num_vertices]).asnumpy()) # check prob assert (len(prob) == max_num_vertices) # check layer From be88caa4c2bcff8211d02ac339460d3def283277 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Sat, 15 Dec 2018 12:49:31 +0800 Subject: [PATCH 21/24] fix tests. --- tests/python/unittest/test_dgl_graph.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/python/unittest/test_dgl_graph.py b/tests/python/unittest/test_dgl_graph.py index 75c877577a42..e24cf4deb756 100644 --- a/tests/python/unittest/test_dgl_graph.py +++ b/tests/python/unittest/test_dgl_graph.py @@ -37,7 +37,7 @@ def check_uniform(out, num_hops, max_num_vertices): sub_csr.check_format(full_check=True) assert np.all((sub_csr.indptr[num_vertices:] == sub_csr.indptr[num_vertices]).asnumpy()) # check layer - for data in layer: + for data in layer[:num_vertices]: assert(data <= num_hops) def check_non_uniform(out, num_hops, max_num_vertices): @@ -54,7 +54,7 @@ def check_non_uniform(out, num_hops, max_num_vertices): # check prob assert (len(prob) == max_num_vertices) # check layer - for data in layer: + for data in layer[:num_vertices]: assert(data <= num_hops) def check_compact(csr, id_arr, num_nodes): From fa910ef733744d7fe19eab8601da12b42d49f14f Mon Sep 17 00:00:00 2001 From: Pingtian Date: Tue, 18 Dec 2018 16:34:58 +0800 Subject: [PATCH 22/24] update docs. --- src/operator/contrib/dgl_graph.cc | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/src/operator/contrib/dgl_graph.cc b/src/operator/contrib/dgl_graph.cc index 5a8851c4f578..a03cbef0b5ca 100644 --- a/src/operator/contrib/dgl_graph.cc +++ b/src/operator/contrib/dgl_graph.cc @@ -756,8 +756,16 @@ static void CSRNeighborUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs, } NNVM_REGISTER_OP(_contrib_dgl_csr_neighbor_uniform_sample) -.describe(R"code(This operator samples sub-graph from a csr graph via an -uniform probability. +.describe(R"code(This operator samples sub-graphs from a csr graph via an +uniform probability. The operator is designed for DGL. + +The operator outputs three sets of NDArrays to represent the sampled results +(the number of NDArrays in each set is the same as the number of seed NDArrays): +1) a set of 1D NDArrays containing the sampled vertices, 2) a set of CSRNDArrays representing +the sampled edges, 3) a set of 1D NDArrays indicating the layer where a vertex is sampled. +The first set of 1D NDArrays have a length of max_num_vertices+1. The last element in an NDArray +indicate the acutal number of vertices in a subgraph. The third set of NDArrays have a length +of max_num_vertices, and the valid number of vertices is the same as the ones in the first set. Example: @@ -843,7 +851,16 @@ static void CSRNeighborNonUniformSampleComputeExCPU(const nnvm::NodeAttrs& attrs NNVM_REGISTER_OP(_contrib_dgl_csr_neighbor_non_uniform_sample) .describe(R"code(This operator samples sub-graph from a csr graph via an -uniform probability. +non-uniform probability. The operator is designed for DGL. + +The operator outputs four sets of NDArrays to represent the sampled results +(the number of NDArrays in each set is the same as the number of seed NDArrays): +1) a set of 1D NDArrays containing the sampled vertices, 2) a set of CSRNDArrays representing +the sampled edges, 3) a set of 1D NDArrays with the probability that vertices are sampled, +4) a set of 1D NDArrays indicating the layer where a vertex is sampled. +The first set of 1D NDArrays have a length of max_num_vertices+1. The last element in an NDArray +indicate the acutal number of vertices in a subgraph. The third and fourth set of NDArrays have a length +of max_num_vertices, and the valid number of vertices is the same as the ones in the first set. Example: From 3bed4af8452fb5adf3f81ded93418feb0d61b5bd Mon Sep 17 00:00:00 2001 From: Pingtian Date: Tue, 18 Dec 2018 20:18:52 +0800 Subject: [PATCH 23/24] retrigger From 259c0f7fa997e71ae6531c449775ff6d74f77917 Mon Sep 17 00:00:00 2001 From: Pingtian Date: Wed, 19 Dec 2018 13:50:28 +0800 Subject: [PATCH 24/24] retrigger