in graphlearn_torch/csrc/cpu/random_sampler.cc [114:135]
void CPURandomSampler::CSRRowWiseSample(
const int64_t* nodes,
const int64_t* nbrs_offset,
const int32_t bs,
const int32_t req_num,
const int64_t row_count,
const int64_t* row_ptr,
const int64_t* col_idx,
const int64_t* edge_ids,
int64_t* out_nbrs,
int64_t* out_eid) {
at::parallel_for(0, bs, 1, [&](int32_t start, int32_t end){
for(int32_t i = start; i < end; ++i) {
auto v = nodes[i];
if (v < row_count) {
UniformSample(col_idx + row_ptr[v], col_idx + row_ptr[v+1],
edge_ids + row_ptr[v], edge_ids + row_ptr[v+1],
req_num, out_nbrs + nbrs_offset[i], out_eid + nbrs_offset[i]);
}
}
});
}