fused_csc_sampling_graph.cc 64.1 KB
Newer Older
1
2
/**
 *  Copyright (c) 2023 by Contributors
3
 * @file fused_csc_sampling_graph.cc
4
5
6
 * @brief Source file of sampling graph.
 */

7
#include <graphbolt/cuda_sampling_ops.h>
8
#include <graphbolt/fused_csc_sampling_graph.h>
9
#include <graphbolt/serialize.h>
10
11
#include <torch/torch.h>

12
13
#include <algorithm>
#include <array>
14
15
#include <cmath>
#include <limits>
16
#include <numeric>
17
18
#include <tuple>
#include <vector>
19

20
#include "./macro.h"
21
#include "./random.h"
22
#include "./shared_memory_helper.h"
23
#include "./utils.h"
24

25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
namespace {
torch::optional<torch::Dict<std::string, torch::Tensor>> TensorizeDict(
    const torch::optional<torch::Dict<std::string, int64_t>>& dict) {
  if (!dict.has_value()) {
    return torch::nullopt;
  }
  torch::Dict<std::string, torch::Tensor> result;
  for (const auto& pair : dict.value()) {
    result.insert(pair.key(), torch::tensor(pair.value(), torch::kInt64));
  }
  return result;
}

torch::optional<torch::Dict<std::string, int64_t>> DetensorizeDict(
    const torch::optional<torch::Dict<std::string, torch::Tensor>>& dict) {
  if (!dict.has_value()) {
    return torch::nullopt;
  }
  torch::Dict<std::string, int64_t> result;
  for (const auto& pair : dict.value()) {
    result.insert(pair.key(), pair.value().item<int64_t>());
  }
  return result;
}
}  // namespace

51
52
53
namespace graphbolt {
namespace sampling {

54
55
static const int kPickleVersion = 6199;

56
FusedCSCSamplingGraph::FusedCSCSamplingGraph(
57
    const torch::Tensor& indptr, const torch::Tensor& indices,
58
    const torch::optional<torch::Tensor>& node_type_offset,
59
    const torch::optional<torch::Tensor>& type_per_edge,
60
61
    const torch::optional<NodeTypeToIDMap>& node_type_to_id,
    const torch::optional<EdgeTypeToIDMap>& edge_type_to_id,
62
    const torch::optional<NodeAttrMap>& node_attributes,
63
    const torch::optional<EdgeAttrMap>& edge_attributes)
64
    : indptr_(indptr),
65
      indices_(indices),
66
      node_type_offset_(node_type_offset),
67
      type_per_edge_(type_per_edge),
68
69
      node_type_to_id_(node_type_to_id),
      edge_type_to_id_(edge_type_to_id),
70
      node_attributes_(node_attributes),
71
      edge_attributes_(edge_attributes) {
72
73
74
75
76
  TORCH_CHECK(indptr.dim() == 1);
  TORCH_CHECK(indices.dim() == 1);
  TORCH_CHECK(indptr.device() == indices.device());
}

77
c10::intrusive_ptr<FusedCSCSamplingGraph> FusedCSCSamplingGraph::Create(
78
    const torch::Tensor& indptr, const torch::Tensor& indices,
79
    const torch::optional<torch::Tensor>& node_type_offset,
80
    const torch::optional<torch::Tensor>& type_per_edge,
81
82
    const torch::optional<NodeTypeToIDMap>& node_type_to_id,
    const torch::optional<EdgeTypeToIDMap>& edge_type_to_id,
83
    const torch::optional<NodeAttrMap>& node_attributes,
84
    const torch::optional<EdgeAttrMap>& edge_attributes) {
85
86
87
  if (node_type_offset.has_value()) {
    auto& offset = node_type_offset.value();
    TORCH_CHECK(offset.dim() == 1);
88
89
90
91
    TORCH_CHECK(node_type_to_id.has_value());
    TORCH_CHECK(
        offset.size(0) ==
        static_cast<int64_t>(node_type_to_id.value().size() + 1));
92
93
94
95
  }
  if (type_per_edge.has_value()) {
    TORCH_CHECK(type_per_edge.value().dim() == 1);
    TORCH_CHECK(type_per_edge.value().size(0) == indices.size(0));
96
    TORCH_CHECK(edge_type_to_id.has_value());
97
  }
98
99
100
101
102
  if (node_attributes.has_value()) {
    for (const auto& pair : node_attributes.value()) {
      TORCH_CHECK(pair.value().size(0) == indptr.size(0) - 1);
    }
  }
103
104
105
106
107
  if (edge_attributes.has_value()) {
    for (const auto& pair : edge_attributes.value()) {
      TORCH_CHECK(pair.value().size(0) == indices.size(0));
    }
  }
108
  return c10::make_intrusive<FusedCSCSamplingGraph>(
109
      indptr, indices, node_type_offset, type_per_edge, node_type_to_id,
110
      edge_type_to_id, node_attributes, edge_attributes);
111
112
}

113
void FusedCSCSamplingGraph::Load(torch::serialize::InputArchive& archive) {
114
  const int64_t magic_num =
115
      read_from_archive<int64_t>(archive, "FusedCSCSamplingGraph/magic_num");
116
117
  TORCH_CHECK(
      magic_num == kCSCSamplingGraphSerializeMagic,
118
119
      "Magic numbers mismatch when loading FusedCSCSamplingGraph.");
  indptr_ =
120
121
122
123
124
125
126
      read_from_archive<torch::Tensor>(archive, "FusedCSCSamplingGraph/indptr");
  indices_ = read_from_archive<torch::Tensor>(
      archive, "FusedCSCSamplingGraph/indices");
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_node_type_offset")) {
    node_type_offset_ = read_from_archive<torch::Tensor>(
        archive, "FusedCSCSamplingGraph/node_type_offset");
127
  }
128
129
130
131
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_type_per_edge")) {
    type_per_edge_ = read_from_archive<torch::Tensor>(
        archive, "FusedCSCSamplingGraph/type_per_edge");
132
  }
133

134
135
136
137
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_node_type_to_id")) {
    node_type_to_id_ = read_from_archive<NodeTypeToIDMap>(
        archive, "FusedCSCSamplingGraph/node_type_to_id");
138
139
  }

140
141
142
143
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_edge_type_to_id")) {
    edge_type_to_id_ = read_from_archive<EdgeTypeToIDMap>(
        archive, "FusedCSCSamplingGraph/edge_type_to_id");
144
145
  }

146
147
148
149
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_node_attributes")) {
    node_attributes_ = read_from_archive<NodeAttrMap>(
        archive, "FusedCSCSamplingGraph/node_attributes");
150
  }
151
152
153
154
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_edge_attributes")) {
    edge_attributes_ = read_from_archive<EdgeAttrMap>(
        archive, "FusedCSCSamplingGraph/edge_attributes");
155
  }
156
157
}

158
159
160
161
162
163
void FusedCSCSamplingGraph::Save(
    torch::serialize::OutputArchive& archive) const {
  archive.write(
      "FusedCSCSamplingGraph/magic_num", kCSCSamplingGraphSerializeMagic);
  archive.write("FusedCSCSamplingGraph/indptr", indptr_);
  archive.write("FusedCSCSamplingGraph/indices", indices_);
164
  archive.write(
165
166
      "FusedCSCSamplingGraph/has_node_type_offset",
      node_type_offset_.has_value());
167
168
  if (node_type_offset_) {
    archive.write(
169
        "FusedCSCSamplingGraph/node_type_offset", node_type_offset_.value());
170
171
  }
  archive.write(
172
      "FusedCSCSamplingGraph/has_type_per_edge", type_per_edge_.has_value());
173
  if (type_per_edge_) {
174
175
    archive.write(
        "FusedCSCSamplingGraph/type_per_edge", type_per_edge_.value());
176
  }
177
178
179
180
181
182
183
184
185
186
187
188
189
190
  archive.write(
      "FusedCSCSamplingGraph/has_node_type_to_id",
      node_type_to_id_.has_value());
  if (node_type_to_id_) {
    archive.write(
        "FusedCSCSamplingGraph/node_type_to_id", node_type_to_id_.value());
  }
  archive.write(
      "FusedCSCSamplingGraph/has_edge_type_to_id",
      edge_type_to_id_.has_value());
  if (edge_type_to_id_) {
    archive.write(
        "FusedCSCSamplingGraph/edge_type_to_id", edge_type_to_id_.value());
  }
191
192
193
194
195
196
197
  archive.write(
      "FusedCSCSamplingGraph/has_node_attributes",
      node_attributes_.has_value());
  if (node_attributes_) {
    archive.write(
        "FusedCSCSamplingGraph/node_attributes", node_attributes_.value());
  }
198
  archive.write(
199
200
      "FusedCSCSamplingGraph/has_edge_attributes",
      edge_attributes_.has_value());
201
  if (edge_attributes_) {
202
203
    archive.write(
        "FusedCSCSamplingGraph/edge_attributes", edge_attributes_.value());
204
  }
205
206
}

207
void FusedCSCSamplingGraph::SetState(
208
209
210
211
212
213
214
215
216
    const torch::Dict<std::string, torch::Dict<std::string, torch::Tensor>>&
        state) {
  // State is a dict of dicts. The tensor-type attributes are stored in the dict
  // with key "independent_tensors". The dict-type attributes (edge_attributes)
  // are stored directly with the their name as the key.
  const auto& independent_tensors = state.at("independent_tensors");
  TORCH_CHECK(
      independent_tensors.at("version_number")
          .equal(torch::tensor({kPickleVersion})),
217
      "Version number mismatches when loading pickled FusedCSCSamplingGraph.")
218
219
220
221
222
223
224
225
226
  indptr_ = independent_tensors.at("indptr");
  indices_ = independent_tensors.at("indices");
  if (independent_tensors.find("node_type_offset") !=
      independent_tensors.end()) {
    node_type_offset_ = independent_tensors.at("node_type_offset");
  }
  if (independent_tensors.find("type_per_edge") != independent_tensors.end()) {
    type_per_edge_ = independent_tensors.at("type_per_edge");
  }
227
228
229
230
231
232
  if (state.find("node_type_to_id") != state.end()) {
    node_type_to_id_ = DetensorizeDict(state.at("node_type_to_id"));
  }
  if (state.find("edge_type_to_id") != state.end()) {
    edge_type_to_id_ = DetensorizeDict(state.at("edge_type_to_id"));
  }
233
234
235
  if (state.find("node_attributes") != state.end()) {
    node_attributes_ = state.at("node_attributes");
  }
236
237
238
239
240
241
  if (state.find("edge_attributes") != state.end()) {
    edge_attributes_ = state.at("edge_attributes");
  }
}

torch::Dict<std::string, torch::Dict<std::string, torch::Tensor>>
242
FusedCSCSamplingGraph::GetState() const {
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
  // State is a dict of dicts. The tensor-type attributes are stored in the dict
  // with key "independent_tensors". The dict-type attributes (edge_attributes)
  // are stored directly with the their name as the key.
  torch::Dict<std::string, torch::Dict<std::string, torch::Tensor>> state;
  torch::Dict<std::string, torch::Tensor> independent_tensors;
  // Serialization version number. It indicates the serialization method of the
  // whole state.
  independent_tensors.insert("version_number", torch::tensor({kPickleVersion}));
  independent_tensors.insert("indptr", indptr_);
  independent_tensors.insert("indices", indices_);
  if (node_type_offset_.has_value()) {
    independent_tensors.insert("node_type_offset", node_type_offset_.value());
  }
  if (type_per_edge_.has_value()) {
    independent_tensors.insert("type_per_edge", type_per_edge_.value());
  }
  state.insert("independent_tensors", independent_tensors);
260
261
262
263
264
265
  if (node_type_to_id_.has_value()) {
    state.insert("node_type_to_id", TensorizeDict(node_type_to_id_).value());
  }
  if (edge_type_to_id_.has_value()) {
    state.insert("edge_type_to_id", TensorizeDict(edge_type_to_id_).value());
  }
266
267
268
  if (node_attributes_.has_value()) {
    state.insert("node_attributes", node_attributes_.value());
  }
269
270
271
272
273
274
  if (edge_attributes_.has_value()) {
    state.insert("edge_attributes", edge_attributes_.value());
  }
  return state;
}

275
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::InSubgraph(
276
    const torch::Tensor& nodes) const {
277
278
279
280
281
282
283
284
285
  if (utils::is_accessible_from_gpu(indptr_) &&
      utils::is_accessible_from_gpu(indices_) &&
      utils::is_accessible_from_gpu(nodes) &&
      (!type_per_edge_.has_value() ||
       utils::is_accessible_from_gpu(type_per_edge_.value()))) {
    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(c10::DeviceType::CUDA, "InSubgraph", {
      return ops::InSubgraph(indptr_, indices_, nodes, type_per_edge_);
    });
  }
286
287
  using namespace torch::indexing;
  const int32_t kDefaultGrainSize = 100;
288
289
  const auto num_seeds = nodes.size(0);
  torch::Tensor indptr = torch::zeros({num_seeds + 1}, indptr_.dtype());
290
  std::vector<torch::Tensor> indices_arr(num_seeds);
291
292
  torch::Tensor original_column_node_ids =
      torch::zeros({num_seeds}, indptr_.dtype());
293
294
  std::vector<torch::Tensor> edge_ids_arr(num_seeds);
  std::vector<torch::Tensor> type_per_edge_arr(num_seeds);
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315

  AT_DISPATCH_INTEGRAL_TYPES(
      indptr_.scalar_type(), "InSubgraph", ([&] {
        torch::parallel_for(
            0, num_seeds, kDefaultGrainSize, [&](size_t start, size_t end) {
              for (size_t i = start; i < end; ++i) {
                const auto node_id = nodes[i].item<scalar_t>();
                const auto start_idx = indptr_[node_id].item<scalar_t>();
                const auto end_idx = indptr_[node_id + 1].item<scalar_t>();
                indptr[i + 1] = end_idx - start_idx;
                original_column_node_ids[i] = node_id;
                indices_arr[i] = indices_.slice(0, start_idx, end_idx);
                edge_ids_arr[i] = torch::arange(start_idx, end_idx);
                if (type_per_edge_) {
                  type_per_edge_arr[i] =
                      type_per_edge_.value().slice(0, start_idx, end_idx);
                }
              }
            });
      }));

316
  return c10::make_intrusive<FusedSampledSubgraph>(
317
      indptr.cumsum(0), torch::cat(indices_arr), original_column_node_ids,
318
319
320
321
322
323
      torch::arange(0, NumNodes()), torch::cat(edge_ids_arr),
      type_per_edge_
          ? torch::optional<torch::Tensor>{torch::cat(type_per_edge_arr)}
          : torch::nullopt);
}

324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
/**
 * @brief Get a lambda function which counts the number of the neighbors to be
 * sampled.
 *
 * @param fanouts The number of edges to be sampled for each node with or
 * without considering edge types.
 * @param replace Boolean indicating whether the sample is performed with or
 * without replacement. If True, a value can be selected multiple times.
 * Otherwise, each value can be selected only once.
 * @param type_per_edge A tensor representing the type of each edge, if
 * present.
 * @param probs_or_mask Optional tensor containing the (unnormalized)
 * probabilities associated with each neighboring edge of a node in the original
 * graph. It must be a 1D floating-point tensor with the number of elements
 * equal to the number of edges in the graph.
 *
340
341
342
343
344
 * @return A lambda function (int64_t seed_offset, int64_t offset, int64_t
 * num_neighbors) -> torch::Tensor, which takes seed offset (the offset of the
 * seed to sample), offset (the starting edge ID of the given node) and
 * num_neighbors (number of neighbors) as params and returns the pick number of
 * the given node.
345
346
347
348
349
350
351
352
 */
auto GetNumPickFn(
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::optional<torch::Tensor>& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask) {
  // If fanouts.size() > 1, returns the total number of all edge types of the
  // given node.
  return [&fanouts, replace, &probs_or_mask, &type_per_edge](
353
             int64_t seed_offset, int64_t offset, int64_t num_neighbors) {
354
355
356
357
358
359
360
361
362
363
    if (fanouts.size() > 1) {
      return NumPickByEtype(
          fanouts, replace, type_per_edge.value(), probs_or_mask, offset,
          num_neighbors);
    } else {
      return NumPick(fanouts[0], replace, probs_or_mask, offset, num_neighbors);
    }
  };
}

364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
auto GetTemporalNumPickFn(
    torch::Tensor seed_timestamp, torch::Tensor csc_indices,
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::optional<torch::Tensor>& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp) {
  // If fanouts.size() > 1, returns the total number of all edge types of the
  // given node.
  return [&seed_timestamp, &csc_indices, &fanouts, replace, &probs_or_mask,
          &type_per_edge, &node_timestamp, &edge_timestamp](
             int64_t seed_offset, int64_t offset, int64_t num_neighbors) {
    if (fanouts.size() > 1) {
      return TemporalNumPickByEtype(
          seed_timestamp, csc_indices, fanouts, replace, type_per_edge.value(),
          probs_or_mask, node_timestamp, edge_timestamp, seed_offset, offset,
          num_neighbors);
    } else {
      return TemporalNumPick(
          seed_timestamp, csc_indices, fanouts[0], replace, probs_or_mask,
          node_timestamp, edge_timestamp, seed_offset, offset, num_neighbors);
    }
  };
}

389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
/**
 * @brief Get a lambda function which contains the sampling process.
 *
 * @param fanouts The number of edges to be sampled for each node with or
 * without considering edge types.
 * @param replace Boolean indicating whether the sample is performed with or
 * without replacement. If True, a value can be selected multiple times.
 * Otherwise, each value can be selected only once.
 * @param options Tensor options specifying the desired data type of the result.
 * @param type_per_edge A tensor representing the type of each edge, if
 * present.
 * @param probs_or_mask Optional tensor containing the (unnormalized)
 * probabilities associated with each neighboring edge of a node in the original
 * graph. It must be a 1D floating-point tensor with the number of elements
 * equal to the number of edges in the graph.
 * @param args Contains sampling algorithm specific arguments.
 *
406
407
408
409
410
 * @return A lambda function: (int64_t seed_offset, int64_t offset, int64_t
 * num_neighbors, PickedType* picked_data_ptr) -> torch::Tensor, which takes
 * seed_offset (the offset of the seed to sample), offset (the starting edge ID
 * of the given node) and num_neighbors (number of neighbors) as params and puts
 * the picked neighbors at the address specified by picked_data_ptr.
411
 */
412
template <SamplerType S>
413
414
415
416
417
418
auto GetPickFn(
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args) {
  return [&fanouts, replace, &options, &type_per_edge, &probs_or_mask, args](
419
420
             int64_t seed_offset, int64_t offset, int64_t num_neighbors,
             auto picked_data_ptr) {
421
422
423
    // If fanouts.size() > 1, perform sampling for each edge type of each
    // node; otherwise just sample once for each node with no regard of edge
    // types.
424
425
426
    if (fanouts.size() > 1) {
      return PickByEtype(
          offset, num_neighbors, fanouts, replace, options,
427
          type_per_edge.value(), probs_or_mask, args, picked_data_ptr);
428
    } else {
429
      int64_t num_sampled = Pick(
430
          offset, num_neighbors, fanouts[0], replace, options, probs_or_mask,
431
          args, picked_data_ptr);
432
433
434
435
      if (type_per_edge) {
        std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
      }
      return num_sampled;
436
437
438
439
    }
  };
}

440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
auto GetTemporalPickFn(
    torch::Tensor seed_timestamp, torch::Tensor csc_indices,
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp) {
  return [&seed_timestamp, &csc_indices, &fanouts, replace, &options,
          &type_per_edge, &probs_or_mask, &node_timestamp, &edge_timestamp](
             int64_t seed_offset, int64_t offset, int64_t num_neighbors,
             auto picked_data_ptr) {
    // If fanouts.size() > 1, perform sampling for each edge type of each
    // node; otherwise just sample once for each node with no regard of edge
    // types.
    if (fanouts.size() > 1) {
      return TemporalPickByEtype(
          seed_timestamp, csc_indices, seed_offset, offset, num_neighbors,
          fanouts, replace, options, type_per_edge.value(), probs_or_mask,
          node_timestamp, edge_timestamp, picked_data_ptr);
    } else {
      int64_t num_sampled = TemporalPick(
          seed_timestamp, csc_indices, seed_offset, offset, num_neighbors,
          fanouts[0], replace, options, probs_or_mask, node_timestamp,
          edge_timestamp, picked_data_ptr);
      if (type_per_edge) {
        std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
      }
      return num_sampled;
    }
  };
}

473
template <typename NumPickFn, typename PickFn>
474
475
c10::intrusive_ptr<FusedSampledSubgraph>
FusedCSCSamplingGraph::SampleNeighborsImpl(
476
477
    const torch::Tensor& nodes, bool return_eids, NumPickFn num_pick_fn,
    PickFn pick_fn) const {
478
  const int64_t num_nodes = nodes.size(0);
479
  const auto indptr_options = indptr_.options();
480
  torch::Tensor num_picked_neighbors_per_node =
481
      torch::empty({num_nodes + 1}, indptr_options);
482

483
484
485
  // Calculate GrainSize for parallel_for.
  // Set the default grain size to 64.
  const int64_t grain_size = 64;
486
487
488
489
490
  torch::Tensor picked_eids;
  torch::Tensor subgraph_indptr;
  torch::Tensor subgraph_indices;
  torch::optional<torch::Tensor> subgraph_type_per_edge = torch::nullopt;

491
  AT_DISPATCH_INTEGRAL_TYPES(
492
493
494
495
496
497
498
499
500
501
      indptr_.scalar_type(), "SampleNeighborsImplWrappedWithIndptr", ([&] {
        using indptr_t = scalar_t;
        AT_DISPATCH_INTEGRAL_TYPES(
            nodes.scalar_type(), "SampleNeighborsImplWrappedWithNodes", ([&] {
              using nodes_t = scalar_t;
              const auto indptr_data = indptr_.data_ptr<indptr_t>();
              auto num_picked_neighbors_data_ptr =
                  num_picked_neighbors_per_node.data_ptr<indptr_t>();
              num_picked_neighbors_data_ptr[0] = 0;
              const auto nodes_data_ptr = nodes.data_ptr<nodes_t>();
502

503
504
505
506
507
508
509
510
511
512
513
514
              // Step 1. Calculate pick number of each node.
              torch::parallel_for(
                  0, num_nodes, grain_size, [&](int64_t begin, int64_t end) {
                    for (int64_t i = begin; i < end; ++i) {
                      const auto nid = nodes_data_ptr[i];
                      TORCH_CHECK(
                          nid >= 0 && nid < NumNodes(),
                          "The seed nodes' IDs should fall within the range of "
                          "the "
                          "graph's node IDs.");
                      const auto offset = indptr_data[nid];
                      const auto num_neighbors = indptr_data[nid + 1] - offset;
515

516
517
518
                      num_picked_neighbors_data_ptr[i + 1] =
                          num_neighbors == 0
                              ? 0
519
                              : num_pick_fn(i, offset, num_neighbors);
520
521
                    }
                  });
522

523
524
525
526
              // Step 2. Calculate prefix sum to get total length and offsets of
              // each node. It's also the indptr of the generated subgraph.
              subgraph_indptr = num_picked_neighbors_per_node.cumsum(
                  0, indptr_.scalar_type());
527

528
529
530
531
532
533
534
535
536
537
              // Step 3. Allocate the tensor for picked neighbors.
              const auto total_length =
                  subgraph_indptr.data_ptr<indptr_t>()[num_nodes];
              picked_eids = torch::empty({total_length}, indptr_options);
              subgraph_indices =
                  torch::empty({total_length}, indices_.options());
              if (type_per_edge_.has_value()) {
                subgraph_type_per_edge = torch::empty(
                    {total_length}, type_per_edge_.value().options());
              }
538

539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
              // Step 4. Pick neighbors for each node.
              auto picked_eids_data_ptr = picked_eids.data_ptr<indptr_t>();
              auto subgraph_indptr_data_ptr =
                  subgraph_indptr.data_ptr<indptr_t>();
              torch::parallel_for(
                  0, num_nodes, grain_size, [&](int64_t begin, int64_t end) {
                    for (int64_t i = begin; i < end; ++i) {
                      const auto nid = nodes_data_ptr[i];
                      const auto offset = indptr_data[nid];
                      const auto num_neighbors = indptr_data[nid + 1] - offset;
                      const auto picked_number =
                          num_picked_neighbors_data_ptr[i + 1];
                      const auto picked_offset = subgraph_indptr_data_ptr[i];
                      if (picked_number > 0) {
                        auto actual_picked_count = pick_fn(
554
                            i, offset, num_neighbors,
555
556
557
558
559
560
                            picked_eids_data_ptr + picked_offset);
                        TORCH_CHECK(
                            actual_picked_count == picked_number,
                            "Actual picked count doesn't match the calculated "
                            "pick "
                            "number.");
561

562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
                        // Step 5. Calculate other attributes and return the
                        // subgraph.
                        AT_DISPATCH_INTEGRAL_TYPES(
                            subgraph_indices.scalar_type(),
                            "IndexSelectSubgraphIndices", ([&] {
                              auto subgraph_indices_data_ptr =
                                  subgraph_indices.data_ptr<scalar_t>();
                              auto indices_data_ptr =
                                  indices_.data_ptr<scalar_t>();
                              for (auto i = picked_offset;
                                   i < picked_offset + picked_number; ++i) {
                                subgraph_indices_data_ptr[i] =
                                    indices_data_ptr[picked_eids_data_ptr[i]];
                              }
                            }));
                        if (type_per_edge_.has_value()) {
                          AT_DISPATCH_INTEGRAL_TYPES(
                              subgraph_type_per_edge.value().scalar_type(),
                              "IndexSelectTypePerEdge", ([&] {
                                auto subgraph_type_per_edge_data_ptr =
                                    subgraph_type_per_edge.value()
                                        .data_ptr<scalar_t>();
                                auto type_per_edge_data_ptr =
                                    type_per_edge_.value().data_ptr<scalar_t>();
                                for (auto i = picked_offset;
                                     i < picked_offset + picked_number; ++i) {
                                  subgraph_type_per_edge_data_ptr[i] =
                                      type_per_edge_data_ptr
                                          [picked_eids_data_ptr[i]];
                                }
                              }));
593
                        }
594
595
596
597
                      }
                    }
                  });
            }));
598
      }));
599

600
601
  torch::optional<torch::Tensor> subgraph_reverse_edge_ids = torch::nullopt;
  if (return_eids) subgraph_reverse_edge_ids = std::move(picked_eids);
602

603
  return c10::make_intrusive<FusedSampledSubgraph>(
604
      subgraph_indptr, subgraph_indices, nodes, torch::nullopt,
605
      subgraph_reverse_edge_ids, subgraph_type_per_edge);
606
607
}

608
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
609
610
611
    const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
    bool replace, bool layer, bool return_eids,
    torch::optional<std::string> probs_name) const {
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
  torch::optional<torch::Tensor> probs_or_mask = torch::nullopt;
  if (probs_name.has_value() && !probs_name.value().empty()) {
    probs_or_mask = this->EdgeAttribute(probs_name);
  }

  if (!replace && utils::is_accessible_from_gpu(indptr_) &&
      utils::is_accessible_from_gpu(indices_) &&
      utils::is_accessible_from_gpu(nodes) &&
      (!probs_or_mask.has_value() ||
       utils::is_accessible_from_gpu(probs_or_mask.value()))) {
    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
        c10::DeviceType::CUDA, "SampleNeighbors", {
          return ops::SampleNeighbors(
              indptr_, indices_, nodes, fanouts, replace, layer, return_eids,
              type_per_edge_, probs_or_mask);
        });
  }

  if (probs_or_mask.has_value()) {
631
632
633
634
635
636
637
638
    // Note probs will be passed as input for 'torch.multinomial' in deeper
    // stack, which doesn't support 'torch.half' and 'torch.bool' data types. To
    // avoid crashes, convert 'probs_or_mask' to 'float32' data type.
    if (probs_or_mask.value().dtype() == torch::kBool ||
        probs_or_mask.value().dtype() == torch::kFloat16) {
      probs_or_mask = probs_or_mask.value().to(torch::kFloat32);
    }
  }
639

640
641
642
643
644
  if (layer) {
    const int64_t random_seed = RandomEngine::ThreadLocal()->RandInt(
        static_cast<int64_t>(0), std::numeric_limits<int64_t>::max());
    SamplerArgs<SamplerType::LABOR> args{indices_, random_seed, NumNodes()};
    return SampleNeighborsImpl(
645
        nodes, return_eids,
646
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
647
648
649
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
650
651
652
  } else {
    SamplerArgs<SamplerType::NEIGHBOR> args;
    return SampleNeighborsImpl(
653
        nodes, return_eids,
654
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
655
656
657
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
658
659
660
  }
}

661
662
663
664
665
666
667
668
669
c10::intrusive_ptr<FusedSampledSubgraph>
FusedCSCSamplingGraph::TemporalSampleNeighbors(
    const torch::Tensor& input_nodes,
    const torch::Tensor& input_nodes_timestamp,
    const std::vector<int64_t>& fanouts, bool replace, bool return_eids,
    torch::optional<std::string> probs_name,
    torch::optional<std::string> node_timestamp_attr_name,
    torch::optional<std::string> edge_timestamp_attr_name) const {
  // 1. Get probs_or_mask.
670
671
672
673
674
675
676
677
678
679
  auto probs_or_mask = this->EdgeAttribute(probs_name);
  if (probs_name.has_value()) {
    // Note probs will be passed as input for 'torch.multinomial' in deeper
    // stack, which doesn't support 'torch.half' and 'torch.bool' data types. To
    // avoid crashes, convert 'probs_or_mask' to 'float32' data type.
    if (probs_or_mask.value().dtype() == torch::kBool ||
        probs_or_mask.value().dtype() == torch::kFloat16) {
      probs_or_mask = probs_or_mask.value().to(torch::kFloat32);
    }
  }
680
  // 2. Get the timestamp attribute for nodes of the graph
681
  auto node_timestamp = this->NodeAttribute(node_timestamp_attr_name);
682
  // 3. Get the timestamp attribute for edges of the graph
683
684
685
686
687
688
689
690
691
692
693
  auto edge_timestamp = this->EdgeAttribute(edge_timestamp_attr_name);
  // 4. Call SampleNeighborsImpl
  return SampleNeighborsImpl(
      input_nodes, return_eids,
      GetTemporalNumPickFn(
          input_nodes_timestamp, this->indices_, fanouts, replace,
          type_per_edge_, probs_or_mask, node_timestamp, edge_timestamp),
      GetTemporalPickFn(
          input_nodes_timestamp, this->indices_, fanouts, replace,
          indptr_.options(), type_per_edge_, probs_or_mask, node_timestamp,
          edge_timestamp));
694
695
}

696
std::tuple<torch::Tensor, torch::Tensor>
697
FusedCSCSamplingGraph::SampleNegativeEdgesUniform(
698
699
700
701
702
703
704
705
706
707
    const std::tuple<torch::Tensor, torch::Tensor>& node_pairs,
    int64_t negative_ratio, int64_t max_node_id) const {
  torch::Tensor pos_src;
  std::tie(pos_src, std::ignore) = node_pairs;
  auto neg_len = pos_src.size(0) * negative_ratio;
  auto neg_src = pos_src.repeat(negative_ratio);
  auto neg_dst = torch::randint(0, max_node_id, {neg_len}, pos_src.options());
  return std::make_tuple(neg_src, neg_dst);
}

708
709
static c10::intrusive_ptr<FusedCSCSamplingGraph>
BuildGraphFromSharedMemoryHelper(SharedMemoryHelper&& helper) {
710
711
712
713
714
  helper.InitializeRead();
  auto indptr = helper.ReadTorchTensor();
  auto indices = helper.ReadTorchTensor();
  auto node_type_offset = helper.ReadTorchTensor();
  auto type_per_edge = helper.ReadTorchTensor();
715
716
  auto node_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
  auto edge_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
717
  auto node_attributes = helper.ReadTorchTensorDict();
718
  auto edge_attributes = helper.ReadTorchTensorDict();
719
  auto graph = c10::make_intrusive<FusedCSCSamplingGraph>(
720
      indptr.value(), indices.value(), node_type_offset, type_per_edge,
721
      node_type_to_id, edge_type_to_id, node_attributes, edge_attributes);
722
723
724
  auto shared_memory = helper.ReleaseSharedMemory();
  graph->HoldSharedMemoryObject(
      std::move(shared_memory.first), std::move(shared_memory.second));
725
726
727
  return graph;
}

728
729
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::CopyToSharedMemory(
730
    const std::string& shared_memory_name) {
731
  SharedMemoryHelper helper(shared_memory_name);
732
733
734
735
  helper.WriteTorchTensor(indptr_);
  helper.WriteTorchTensor(indices_);
  helper.WriteTorchTensor(node_type_offset_);
  helper.WriteTorchTensor(type_per_edge_);
736
737
  helper.WriteTorchTensorDict(TensorizeDict(node_type_to_id_));
  helper.WriteTorchTensorDict(TensorizeDict(edge_type_to_id_));
738
  helper.WriteTorchTensorDict(node_attributes_);
739
740
741
  helper.WriteTorchTensorDict(edge_attributes_);
  helper.Flush();
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
742
743
}

744
745
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::LoadFromSharedMemory(
746
    const std::string& shared_memory_name) {
747
  SharedMemoryHelper helper(shared_memory_name);
748
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
749
750
}

751
void FusedCSCSamplingGraph::HoldSharedMemoryObject(
752
753
754
755
756
    SharedMemoryPtr tensor_metadata_shm, SharedMemoryPtr tensor_data_shm) {
  tensor_metadata_shm_ = std::move(tensor_metadata_shm);
  tensor_data_shm_ = std::move(tensor_data_shm);
}

757
758
759
760
int64_t NumPick(
    int64_t fanout, bool replace,
    const torch::optional<torch::Tensor>& probs_or_mask, int64_t offset,
    int64_t num_neighbors) {
761
762
763
764
765
766
767
768
769
770
771
  int64_t num_valid_neighbors = num_neighbors;
  if (probs_or_mask.has_value()) {
    // Subtract the count of zeros in probs_or_mask.
    AT_DISPATCH_ALL_TYPES(
        probs_or_mask.value().scalar_type(), "CountZero", ([&] {
          scalar_t* probs_data_ptr = probs_or_mask.value().data_ptr<scalar_t>();
          num_valid_neighbors -= std::count(
              probs_data_ptr + offset, probs_data_ptr + offset + num_neighbors,
              0);
        }));
  }
772
773
774
775
  if (num_valid_neighbors == 0 || fanout == -1) return num_valid_neighbors;
  return replace ? fanout : std::min(fanout, num_valid_neighbors);
}

776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
torch::Tensor TemporalMask(
    int64_t seed_timestamp, torch::Tensor csc_indices,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp,
    std::pair<int64_t, int64_t> edge_range) {
  auto [l, r] = edge_range;
  torch::Tensor mask = torch::ones({r - l}, torch::kBool);
  if (node_timestamp.has_value()) {
    auto neighbor_timestamp =
        node_timestamp.value().index_select(0, csc_indices.slice(0, l, r));
    mask &= neighbor_timestamp <= seed_timestamp;
  }
  if (edge_timestamp.has_value()) {
    mask &= edge_timestamp.value().slice(0, l, r) <= seed_timestamp;
  }
  if (probs_or_mask.has_value()) {
    mask &= probs_or_mask.value().slice(0, l, r) != 0;
  }
  return mask;
}

int64_t TemporalNumPick(
    torch::Tensor seed_timestamp, torch::Tensor csc_indics, int64_t fanout,
    bool replace, const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp, int64_t seed_offset,
    int64_t offset, int64_t num_neighbors) {
  auto mask = TemporalMask(
      utils::GetValueByIndex<int64_t>(seed_timestamp, seed_offset), csc_indics,
      probs_or_mask, node_timestamp, edge_timestamp,
      {offset, offset + num_neighbors});
  int64_t num_valid_neighbors = utils::GetValueByIndex<int64_t>(mask.sum(), 0);
  if (num_valid_neighbors == 0 || fanout == -1) return num_valid_neighbors;
  return replace ? fanout : std::min(fanout, num_valid_neighbors);
}

813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
int64_t NumPickByEtype(
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::Tensor& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask, int64_t offset,
    int64_t num_neighbors) {
  int64_t etype_begin = offset;
  const int64_t end = offset + num_neighbors;
  int64_t total_count = 0;
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "NumPickFnByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
          TORCH_CHECK(
              etype >= 0 && etype < (int64_t)fanouts.size(),
              "Etype values exceed the number of fanouts.");
          auto etype_end_it = std::upper_bound(
              type_per_edge_data + etype_begin, type_per_edge_data + end,
              etype);
          int64_t etype_end = etype_end_it - type_per_edge_data;
          // Do sampling for one etype.
          total_count += NumPick(
              fanouts[etype], replace, probs_or_mask, etype_begin,
              etype_end - etype_begin);
          etype_begin = etype_end;
        }
      }));
  return total_count;
}

843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
int64_t TemporalNumPickByEtype(
    torch::Tensor seed_timestamp, torch::Tensor csc_indices,
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::Tensor& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp, int64_t seed_offset,
    int64_t offset, int64_t num_neighbors) {
  int64_t etype_begin = offset;
  const int64_t end = offset + num_neighbors;
  int64_t total_count = 0;
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "TemporalNumPickFnByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
          TORCH_CHECK(
              etype >= 0 && etype < (int64_t)fanouts.size(),
              "Etype values exceed the number of fanouts.");
          auto etype_end_it = std::upper_bound(
              type_per_edge_data + etype_begin, type_per_edge_data + end,
              etype);
          int64_t etype_end = etype_end_it - type_per_edge_data;
          // Do sampling for one etype.
          total_count += TemporalNumPick(
              seed_timestamp, csc_indices, fanouts[etype], replace,
              probs_or_mask, node_timestamp, edge_timestamp, seed_offset,
              etype_begin, etype_end - etype_begin);
          etype_begin = etype_end;
        }
      }));
  return total_count;
}

877
878
879
880
881
882
883
884
/**
 * @brief Perform uniform sampling of elements and return the sampled indices.
 *
 * @param offset The starting edge ID for the connected neighbors of the sampled
 * node.
 * @param num_neighbors The number of neighbors to pick.
 * @param fanout The number of edges to be sampled for each node. It should be
 * >= 0 or -1.
885
886
887
 *  - When the value is -1, all neighbors will be sampled once regardless of
 * replacement. It is equivalent to selecting all neighbors when the fanout is
 * >= the number of neighbors (and replacement is set to false).
888
889
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
890
 * @param replace Boolean indicating whether the sample is performed with or
891
892
893
 * without replacement. If True, a value can be selected multiple times.
 * Otherwise, each value can be selected only once.
 * @param options Tensor options specifying the desired data type of the result.
894
895
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
896
 */
897
template <typename PickedType>
898
inline int64_t UniformPick(
899
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
900
    const torch::TensorOptions& options, PickedType* picked_data_ptr) {
901
  if ((fanout == -1) || (num_neighbors <= fanout && !replace)) {
902
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
903
    return num_neighbors;
904
  } else if (replace) {
905
906
907
908
909
    std::memcpy(
        picked_data_ptr,
        torch::randint(offset, offset + num_neighbors, {fanout}, options)
            .data_ptr<PickedType>(),
        fanout * sizeof(PickedType));
910
    return fanout;
911
  } else {
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
    // We use different sampling strategies for different sampling case.
    if (fanout >= num_neighbors / 10) {
      // [Algorithm]
      // This algorithm is conceptually related to the Fisher-Yates
      // shuffle.
      //
      // [Complexity Analysis]
      // This algorithm's memory complexity is O(num_neighbors), but
      // it generates fewer random numbers (O(fanout)).
      //
      // (Compare) Reservoir algorithm is one of the most classical
      // sampling algorithms. Both the reservoir algorithm and our
      // algorithm offer distinct advantages, we need to compare to
      // illustrate our trade-offs.
      // The reservoir algorithm is memory-efficient (O(fanout)) but
      // creates many random numbers (O(num_neighbors)), which is
      // costly.
      //
      // [Practical Consideration]
      // Use this algorithm when `fanout >= num_neighbors / 10` to
      // reduce computation.
      // In this scenarios above, memory complexity is not a concern due
      // to the small size of both `fanout` and `num_neighbors`. And it
      // is efficient to allocate a small amount of memory. So the
      // algorithm performence is great in this case.
      std::vector<PickedType> seq(num_neighbors);
      // Assign the seq with [offset, offset + num_neighbors].
      std::iota(seq.begin(), seq.end(), offset);
      for (int64_t i = 0; i < fanout; ++i) {
        auto j = RandomEngine::ThreadLocal()->RandInt(i, num_neighbors);
        std::swap(seq[i], seq[j]);
      }
      // Save the randomly sampled fanout elements to the output tensor.
      std::copy(seq.begin(), seq.begin() + fanout, picked_data_ptr);
946
      return fanout;
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
    } else if (fanout < 64) {
      // [Algorithm]
      // Use linear search to verify uniqueness.
      //
      // [Complexity Analysis]
      // Since the set of numbers is small (up to 64), so it is more
      // cost-effective for the CPU to use this algorithm.
      auto begin = picked_data_ptr;
      auto end = picked_data_ptr + fanout;

      while (begin != end) {
        // Put the new random number in the last position.
        *begin = RandomEngine::ThreadLocal()->RandInt(
            offset, offset + num_neighbors);
        // Check if a new value doesn't exist in current
        // range(picked_data_ptr, begin). Otherwise get a new
        // value until we haven't unique range of elements.
        auto it = std::find(picked_data_ptr, begin, *begin);
        if (it == begin) ++begin;
      }
967
      return fanout;
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
    } else {
      // [Algorithm]
      // Use hash-set to verify uniqueness. In the best scenario, the
      // time complexity is O(fanout), assuming no conflicts occur.
      //
      // [Complexity Analysis]
      // Let K = (fanout / num_neighbors), the expected number of extra
      // sampling steps is roughly K^2 / (1-K) * num_neighbors, which
      // means in the worst case scenario, the time complexity is
      // O(num_neighbors^2).
      //
      // [Practical Consideration]
      // In practice, we set the threshold K to 1/10. This trade-off is
      // due to the slower performance of std::unordered_set, which
      // would otherwise increase the sampling cost. By doing so, we
      // achieve a balance between theoretical efficiency and practical
      // performance.
      std::unordered_set<PickedType> picked_set;
      while (static_cast<int64_t>(picked_set.size()) < fanout) {
        picked_set.insert(RandomEngine::ThreadLocal()->RandInt(
            offset, offset + num_neighbors));
      }
      std::copy(picked_set.begin(), picked_set.end(), picked_data_ptr);
991
      return picked_set.size();
992
    }
993
994
995
  }
}

996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
/** @brief An operator to perform non-uniform sampling. */
static torch::Tensor NonUniformPickOp(
    torch::Tensor probs, int64_t fanout, bool replace) {
  auto positive_probs_indices = probs.nonzero().squeeze(1);
  auto num_positive_probs = positive_probs_indices.size(0);
  if (num_positive_probs == 0) return torch::empty({0}, torch::kLong);
  if ((fanout == -1) || (num_positive_probs <= fanout && !replace)) {
    return positive_probs_indices;
  }
  if (!replace) fanout = std::min(fanout, num_positive_probs);
  if (fanout == 0) return torch::empty({0}, torch::kLong);
  auto ret_tensor = torch::empty({fanout}, torch::kLong);
  auto ret_ptr = ret_tensor.data_ptr<int64_t>();
  AT_DISPATCH_FLOATING_TYPES(
      probs.scalar_type(), "MultinomialSampling", ([&] {
        auto probs_data_ptr = probs.data_ptr<scalar_t>();
        auto positive_probs_indices_ptr =
            positive_probs_indices.data_ptr<int64_t>();

        if (!replace) {
          // The algorithm is from gumbel softmax.
          // s = argmax( logp - log(-log(eps)) ) where eps ~ U(0, 1).
          // Here we can apply exp to the formula which will not affect result
          // of argmax or topk. Then we have
          // s = argmax( p / (-log(eps)) ) where eps ~ U(0, 1).
          // We can also simplify the formula above by
          // s = argmax( p / q ) where q ~ Exp(1).
          if (fanout == 1) {
            // Return argmax(p / q).
            scalar_t max_prob = 0;
            int64_t max_prob_index = -1;
            // We only care about the neighbors with non-zero probability.
            for (auto i = 0; i < num_positive_probs; ++i) {
              // Calculate (p / q) for the current neighbor.
              scalar_t current_prob =
                  probs_data_ptr[positive_probs_indices_ptr[i]] /
                  RandomEngine::ThreadLocal()->Exponential(1.);
              if (current_prob > max_prob) {
                max_prob = current_prob;
                max_prob_index = positive_probs_indices_ptr[i];
              }
            }
            ret_ptr[0] = max_prob_index;
          } else {
            // Return topk(p / q).
            std::vector<std::pair<scalar_t, int64_t>> q(num_positive_probs);
            for (auto i = 0; i < num_positive_probs; ++i) {
              q[i].first = probs_data_ptr[positive_probs_indices_ptr[i]] /
                           RandomEngine::ThreadLocal()->Exponential(1.);
              q[i].second = positive_probs_indices_ptr[i];
            }
            if (fanout < num_positive_probs / 64) {
              // Use partial_sort.
              std::partial_sort(
                  q.begin(), q.begin() + fanout, q.end(), std::greater{});
              for (auto i = 0; i < fanout; ++i) {
                ret_ptr[i] = q[i].second;
              }
            } else {
              // Use nth_element.
              std::nth_element(
                  q.begin(), q.begin() + fanout - 1, q.end(), std::greater{});
              for (auto i = 0; i < fanout; ++i) {
                ret_ptr[i] = q[i].second;
              }
            }
          }
        } else {
          // Calculate cumulative sum of probabilities.
          std::vector<scalar_t> prefix_sum_probs(num_positive_probs);
          scalar_t sum_probs = 0;
          for (auto i = 0; i < num_positive_probs; ++i) {
            sum_probs += probs_data_ptr[positive_probs_indices_ptr[i]];
            prefix_sum_probs[i] = sum_probs;
          }
          // Normalize.
          if ((sum_probs > 1.00001) || (sum_probs < 0.99999)) {
            for (auto i = 0; i < num_positive_probs; ++i) {
              prefix_sum_probs[i] /= sum_probs;
            }
          }
          for (auto i = 0; i < fanout; ++i) {
            // Sample a probability mass from a uniform distribution.
            double uniform_sample =
                RandomEngine::ThreadLocal()->Uniform(0., 1.);
            // Use a binary search to find the index.
            int sampled_index = std::lower_bound(
                                    prefix_sum_probs.begin(),
                                    prefix_sum_probs.end(), uniform_sample) -
                                prefix_sum_probs.begin();
            ret_ptr[i] = positive_probs_indices_ptr[sampled_index];
          }
        }
      }));
  return ret_tensor;
}

1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
/**
 * @brief Perform non-uniform sampling of elements based on probabilities and
 * return the sampled indices.
 *
 * If 'probs_or_mask' is provided, it indicates that the sampling is
 * non-uniform. In such cases:
 * - When the number of neighbors with non-zero probability is less than or
 * equal to fanout, all neighbors with non-zero probability will be selected.
 * - When the number of neighbors with non-zero probability exceeds fanout, the
 * sampling process will select 'fanout' elements based on their respective
 * probabilities. Higher probabilities will increase the chances of being chosen
 * during the sampling process.
 *
 * @param offset The starting edge ID for the connected neighbors of the sampled
 * node.
 * @param num_neighbors The number of neighbors to pick.
 * @param fanout The number of edges to be sampled for each node. It should be
 * >= 0 or -1.
1111
1112
1113
1114
 *  - When the value is -1, all neighbors with non-zero probability will be
 * sampled once regardless of replacement. It is equivalent to selecting all
 * neighbors with non-zero probability when the fanout is >= the number of
 * neighbors (and replacement is set to false).
1115
1116
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
1117
 * @param replace Boolean indicating whether the sample is performed with or
1118
1119
1120
1121
1122
1123
1124
 * without replacement. If True, a value can be selected multiple times.
 * Otherwise, each value can be selected only once.
 * @param options Tensor options specifying the desired data type of the result.
 * @param probs_or_mask Optional tensor containing the (unnormalized)
 * probabilities associated with each neighboring edge of a node in the original
 * graph. It must be a 1D floating-point tensor with the number of elements
 * equal to the number of edges in the graph.
1125
1126
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
1127
 */
1128
template <typename PickedType>
1129
inline int64_t NonUniformPick(
1130
1131
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
1132
1133
    const torch::optional<torch::Tensor>& probs_or_mask,
    PickedType* picked_data_ptr) {
1134
1135
  auto local_probs =
      probs_or_mask.value().slice(0, offset, offset + num_neighbors);
1136
1137
1138
1139
1140
  auto picked_indices = NonUniformPickOp(local_probs, fanout, replace);
  auto picked_indices_ptr = picked_indices.data_ptr<int64_t>();
  for (int i = 0; i < picked_indices.numel(); ++i) {
    picked_data_ptr[i] =
        static_cast<PickedType>(picked_indices_ptr[i]) + offset;
1141
  }
1142
  return picked_indices.numel();
1143
1144
}

1145
template <typename PickedType>
1146
int64_t Pick(
1147
1148
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
1149
    const torch::optional<torch::Tensor>& probs_or_mask,
1150
    SamplerArgs<SamplerType::NEIGHBOR> args, PickedType* picked_data_ptr) {
1151
  if (probs_or_mask.has_value()) {
1152
    return NonUniformPick(
1153
1154
        offset, num_neighbors, fanout, replace, options, probs_or_mask,
        picked_data_ptr);
1155
  } else {
1156
    return UniformPick(
1157
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
1158
1159
1160
  }
}

1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
template <typename PickedType>
int64_t TemporalPick(
    torch::Tensor seed_timestamp, torch::Tensor csc_indices,
    int64_t seed_offset, int64_t offset, int64_t num_neighbors, int64_t fanout,
    bool replace, const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp,
    PickedType* picked_data_ptr) {
  auto mask = TemporalMask(
      utils::GetValueByIndex<int64_t>(seed_timestamp, seed_offset), csc_indices,
      probs_or_mask, node_timestamp, edge_timestamp,
      {offset, offset + num_neighbors});
  torch::Tensor masked_prob;
  if (probs_or_mask.has_value()) {
    masked_prob =
        probs_or_mask.value().slice(0, offset, offset + num_neighbors) * mask;
  } else {
    masked_prob = mask.to(torch::kFloat32);
  }
  auto picked_indices = NonUniformPickOp(masked_prob, fanout, replace);
  auto picked_indices_ptr = picked_indices.data_ptr<int64_t>();
  for (int i = 0; i < picked_indices.numel(); ++i) {
    picked_data_ptr[i] =
        static_cast<PickedType>(picked_indices_ptr[i]) + offset;
  }
  return picked_indices.numel();
}

1190
template <SamplerType S, typename PickedType>
1191
int64_t PickByEtype(
1192
1193
    int64_t offset, int64_t num_neighbors, const std::vector<int64_t>& fanouts,
    bool replace, const torch::TensorOptions& options,
1194
    const torch::Tensor& type_per_edge,
1195
1196
    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args,
    PickedType* picked_data_ptr) {
1197
1198
  int64_t etype_begin = offset;
  int64_t etype_end = offset;
1199
  int64_t pick_offset = 0;
1200
1201
1202
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "PickByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
1203
1204
1205
        const auto end = offset + num_neighbors;
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
1206
          TORCH_CHECK(
1207
              etype >= 0 && etype < (int64_t)fanouts.size(),
1208
              "Etype values exceed the number of fanouts.");
1209
          int64_t fanout = fanouts[etype];
1210
1211
1212
1213
          auto etype_end_it = std::upper_bound(
              type_per_edge_data + etype_begin, type_per_edge_data + end,
              etype);
          etype_end = etype_end_it - type_per_edge_data;
1214
1215
          // Do sampling for one etype.
          if (fanout != 0) {
1216
            int64_t picked_count = Pick(
1217
                etype_begin, etype_end - etype_begin, fanout, replace, options,
1218
1219
                probs_or_mask, args, picked_data_ptr + pick_offset);
            pick_offset += picked_count;
1220
1221
1222
1223
          }
          etype_begin = etype_end;
        }
      }));
1224
  return pick_offset;
1225
1226
}

1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
template <typename PickedType>
int64_t TemporalPickByEtype(
    torch::Tensor seed_timestamp, torch::Tensor csc_indices,
    int64_t seed_offset, int64_t offset, int64_t num_neighbors,
    const std::vector<int64_t>& fanouts, bool replace,
    const torch::TensorOptions& options, const torch::Tensor& type_per_edge,
    const torch::optional<torch::Tensor>& probs_or_mask,
    const torch::optional<torch::Tensor>& node_timestamp,
    const torch::optional<torch::Tensor>& edge_timestamp,
    PickedType* picked_data_ptr) {
  int64_t etype_begin = offset;
  int64_t etype_end = offset;
  int64_t pick_offset = 0;
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "TemporalPickByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
        const auto end = offset + num_neighbors;
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
          TORCH_CHECK(
              etype >= 0 && etype < (int64_t)fanouts.size(),
              "Etype values exceed the number of fanouts.");
          int64_t fanout = fanouts[etype];
          auto etype_end_it = std::upper_bound(
              type_per_edge_data + etype_begin, type_per_edge_data + end,
              etype);
          etype_end = etype_end_it - type_per_edge_data;
          // Do sampling for one etype.
          if (fanout != 0) {
            int64_t picked_count = TemporalPick(
                seed_timestamp, csc_indices, seed_offset, etype_begin,
                etype_end - etype_begin, fanout, replace, options,
                probs_or_mask, node_timestamp, edge_timestamp,
                picked_data_ptr + pick_offset);
            pick_offset += picked_count;
          }
          etype_begin = etype_end;
        }
      }));
  return pick_offset;
}

1269
template <typename PickedType>
1270
int64_t Pick(
1271
1272
1273
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1274
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1275
  if (fanout == 0) return 0;
1276
  if (probs_or_mask.has_value()) {
1277
    if (fanout < 0) {
1278
      return NonUniformPick(
1279
1280
1281
          offset, num_neighbors, fanout, replace, options, probs_or_mask,
          picked_data_ptr);
    } else {
1282
      int64_t picked_count;
1283
1284
1285
      AT_DISPATCH_FLOATING_TYPES(
          probs_or_mask.value().scalar_type(), "LaborPickFloatType", ([&] {
            if (replace) {
1286
              picked_count = LaborPick<true, true, scalar_t>(
1287
1288
1289
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            } else {
1290
              picked_count = LaborPick<true, false, scalar_t>(
1291
1292
1293
1294
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            }
          }));
1295
      return picked_count;
1296
1297
    }
  } else if (fanout < 0) {
1298
    return UniformPick(
1299
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
1300
  } else if (replace) {
1301
    return LaborPick<false, true, float>(
1302
        offset, num_neighbors, fanout, options,
1303
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1304
  } else {  // replace = false
1305
    return LaborPick<false, false, float>(
1306
        offset, num_neighbors, fanout, options,
1307
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
  }
}

template <typename T, typename U>
inline void safe_divide(T& a, U b) {
  a = b > 0 ? (T)(a / b) : std::numeric_limits<T>::infinity();
}

/**
 * @brief Perform uniform-nonuniform sampling of elements depending on the
 * template parameter NonUniform and return the sampled indices.
 *
 * @param offset The starting edge ID for the connected neighbors of the sampled
 * node.
 * @param num_neighbors The number of neighbors to pick.
 * @param fanout The number of edges to be sampled for each node. It should be
 * >= 0 or -1.
1325
1326
1327
1328
 *  - When the value is -1, all neighbors (with non-zero probability, if
 * weighted) will be sampled once regardless of replacement. It is equivalent to
 * selecting all neighbors with non-zero probability when the fanout is >= the
 * number of neighbors (and replacement is set to false).
1329
1330
1331
1332
1333
1334
1335
1336
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
 * @param options Tensor options specifying the desired data type of the result.
 * @param probs_or_mask Optional tensor containing the (unnormalized)
 * probabilities associated with each neighboring edge of a node in the original
 * graph. It must be a 1D floating-point tensor with the number of elements
 * equal to the number of edges in the graph.
 * @param args Contains labor specific arguments.
1337
1338
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
1339
 */
1340
template <
1341
1342
    bool NonUniform, bool Replace, typename ProbsType, typename PickedType,
    int StackSize>
1343
inline int64_t LaborPick(
1344
1345
1346
    int64_t offset, int64_t num_neighbors, int64_t fanout,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1347
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1348
  fanout = Replace ? fanout : std::min(fanout, num_neighbors);
1349
  if (!NonUniform && !Replace && fanout >= num_neighbors) {
1350
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
1351
    return num_neighbors;
1352
1353
  }
  // Assuming max_degree of a vertex is <= 4 billion.
1354
1355
1356
1357
1358
1359
1360
1361
1362
  std::array<std::pair<float, uint32_t>, StackSize> heap;
  auto heap_data = heap.data();
  torch::Tensor heap_tensor;
  if (fanout > StackSize) {
    constexpr int factor = sizeof(heap_data[0]) / sizeof(int32_t);
    heap_tensor = torch::empty({fanout * factor}, torch::kInt32);
    heap_data = reinterpret_cast<std::pair<float, uint32_t>*>(
        heap_tensor.data_ptr<int32_t>());
  }
1363
1364
1365
  const ProbsType* local_probs_data =
      NonUniform ? probs_or_mask.value().data_ptr<ProbsType>() + offset
                 : nullptr;
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
  AT_DISPATCH_INTEGRAL_TYPES(
      args.indices.scalar_type(), "LaborPickMain", ([&] {
        const scalar_t* local_indices_data =
            args.indices.data_ptr<scalar_t>() + offset;
        if constexpr (Replace) {
          // [Algorithm] @mfbalin
          // Use a max-heap to get rid of the big random numbers and filter the
          // smallest fanout of them. Implements arXiv:2210.13339 Section A.3.
          // Unlike sampling without replacement below, the same item can be
          // included fanout times in our sample. Thus, we sort and pick the
          // smallest fanout random numbers out of num_neighbors * fanout of
          // them. Each item has fanout many random numbers in the race and the
          // smallest fanout of them get picked. Instead of generating
          // fanout * num_neighbors random numbers and increase the complexity,
          // I devised an algorithm to generate the fanout numbers for an item
          // in a sorted manner on demand, meaning we continue generating random
          // numbers for an item only if it has been sampled that many times
          // already.
          // https://gist.github.com/mfbalin/096dcad5e3b1f6a59ff7ff2f9f541618
          //
          // [Complexity Analysis]
          // Will modify the heap at most linear in O(num_neighbors + fanout)
          // and each modification takes O(log(fanout)). So the total complexity
          // is O((fanout + num_neighbors) log(fanout)). It is possible to
          // decrease the logarithmic factor down to
          // O(log(min(fanout, num_neighbors))).
1392
1393
1394
1395
1396
1397
1398
1399
          std::array<float, StackSize> remaining;
          auto remaining_data = remaining.data();
          torch::Tensor remaining_tensor;
          if (num_neighbors > StackSize) {
            remaining_tensor = torch::empty({num_neighbors}, torch::kFloat32);
            remaining_data = remaining_tensor.data_ptr<float>();
          }
          std::fill_n(remaining_data, num_neighbors, 1.f);
1400
1401
1402
1403
1404
          auto heap_end = heap_data;
          const auto init_count = (num_neighbors + fanout - 1) / num_neighbors;
          auto sample_neighbor_i_with_index_t_jth_time =
              [&](scalar_t t, int64_t j, uint32_t i) {
                auto rnd = labor::jth_sorted_uniform_random(
1405
                    args.random_seed, t, args.num_nodes, j, remaining_data[i],
1406
1407
1408
1409
1410
1411
                    fanout - j);  // r_t
                if constexpr (NonUniform) {
                  safe_divide(rnd, local_probs_data[i]);
                }  // r_t / \pi_t
                if (heap_end < heap_data + fanout) {
                  heap_end[0] = std::make_pair(rnd, i);
1412
1413
1414
                  if (++heap_end >= heap_data + fanout) {
                    std::make_heap(heap_data, heap_data + fanout);
                  }
1415
1416
1417
1418
1419
1420
1421
                  return false;
                } else if (rnd < heap_data[0].first) {
                  std::pop_heap(heap_data, heap_data + fanout);
                  heap_data[fanout - 1] = std::make_pair(rnd, i);
                  std::push_heap(heap_data, heap_data + fanout);
                  return false;
                } else {
1422
                  remaining_data[i] = -1;
1423
1424
1425
1426
                  return true;
                }
              };
          for (uint32_t i = 0; i < num_neighbors; ++i) {
1427
            const auto t = local_indices_data[i];
1428
1429
1430
1431
1432
            for (int64_t j = 0; j < init_count; j++) {
              sample_neighbor_i_with_index_t_jth_time(t, j, i);
            }
          }
          for (uint32_t i = 0; i < num_neighbors; ++i) {
1433
            if (remaining_data[i] == -1) continue;
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
            const auto t = local_indices_data[i];
            for (int64_t j = init_count; j < fanout; ++j) {
              if (sample_neighbor_i_with_index_t_jth_time(t, j, i)) break;
            }
          }
        } else {
          // [Algorithm]
          // Use a max-heap to get rid of the big random numbers and filter the
          // smallest fanout of them. Implements arXiv:2210.13339 Section A.3.
          //
          // [Complexity Analysis]
          // the first for loop and std::make_heap runs in time O(fanouts).
          // The next for loop compares each random number to the current
          // minimum fanout numbers. For any given i, the probability that the
          // current random number will replace any number in the heap is fanout
          // / i. Summing from i=fanout to num_neighbors, we get f * (H_n -
          // H_f), where n is num_neighbors and f is fanout, H_f is \sum_j=1^f
          // 1/j. In the end H_n - H_f = O(log n/f), there are n - f iterations,
          // each heap operation takes time log f, so the total complexity is
          // O(f + (n - f)
          // + f log(n/f) log f) = O(n + f log(f) log(n/f)). If f << n (f is a
          // constant in almost all cases), then the average complexity is
          // O(num_neighbors).
          for (uint32_t i = 0; i < fanout; ++i) {
            const auto t = local_indices_data[i];
            auto rnd =
                labor::uniform_random<float>(args.random_seed, t);  // r_t
            if constexpr (NonUniform) {
              safe_divide(rnd, local_probs_data[i]);
            }  // r_t / \pi_t
            heap_data[i] = std::make_pair(rnd, i);
          }
          if (!NonUniform || fanout < num_neighbors) {
            std::make_heap(heap_data, heap_data + fanout);
          }
          for (uint32_t i = fanout; i < num_neighbors; ++i) {
            const auto t = local_indices_data[i];
            auto rnd =
                labor::uniform_random<float>(args.random_seed, t);  // r_t
            if constexpr (NonUniform) {
              safe_divide(rnd, local_probs_data[i]);
            }  // r_t / \pi_t
            if (rnd < heap_data[0].first) {
              std::pop_heap(heap_data, heap_data + fanout);
              heap_data[fanout - 1] = std::make_pair(rnd, i);
              std::push_heap(heap_data, heap_data + fanout);
            }
          }
        }
      }));
  int64_t num_sampled = 0;
1485
1486
1487
1488
1489
1490
  for (int64_t i = 0; i < fanout; ++i) {
    const auto [rnd, j] = heap_data[i];
    if (!NonUniform || rnd < std::numeric_limits<float>::infinity()) {
      picked_data_ptr[num_sampled++] = offset + j;
    }
  }
1491
  return num_sampled;
1492
1493
}

1494
1495
}  // namespace sampling
}  // namespace graphbolt