fused_csc_sampling_graph.cc 65 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
  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr_) &&
278
279
280
281
282
283
284
      utils::is_accessible_from_gpu(indices_) &&
      (!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_);
    });
  }
285
286
  using namespace torch::indexing;
  const int32_t kDefaultGrainSize = 100;
287
288
  const auto num_seeds = nodes.size(0);
  torch::Tensor indptr = torch::zeros({num_seeds + 1}, indptr_.dtype());
289
  std::vector<torch::Tensor> indices_arr(num_seeds);
290
291
  torch::Tensor original_column_node_ids =
      torch::zeros({num_seeds}, indptr_.dtype());
292
293
  std::vector<torch::Tensor> edge_ids_arr(num_seeds);
  std::vector<torch::Tensor> type_per_edge_arr(num_seeds);
294

295
  AT_DISPATCH_INDEX_TYPES(
296
297
298
299
      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) {
300
301
302
                const auto node_id = nodes[i].item<index_t>();
                const auto start_idx = indptr_[node_id].item<index_t>();
                const auto end_idx = indptr_[node_id + 1].item<index_t>();
303
304
305
306
307
308
309
310
311
312
313
314
                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);
                }
              }
            });
      }));

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

323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
/**
 * @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.
 *
339
340
341
342
343
 * @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.
344
345
346
347
348
349
350
351
 */
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](
352
             int64_t seed_offset, int64_t offset, int64_t num_neighbors) {
353
354
355
356
357
358
359
360
361
362
    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);
    }
  };
}

363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
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);
    }
  };
}

388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
/**
 * @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.
 *
405
406
407
408
409
 * @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.
410
 */
411
template <SamplerType S>
412
413
414
415
416
417
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](
418
419
             int64_t seed_offset, int64_t offset, int64_t num_neighbors,
             auto picked_data_ptr) {
420
421
422
    // 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.
423
424
425
    if (fanouts.size() > 1) {
      return PickByEtype(
          offset, num_neighbors, fanouts, replace, options,
426
          type_per_edge.value(), probs_or_mask, args, picked_data_ptr);
427
    } else {
428
      int64_t num_sampled = Pick(
429
          offset, num_neighbors, fanouts[0], replace, options, probs_or_mask,
430
          args, picked_data_ptr);
431
432
433
434
      if (type_per_edge) {
        std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
      }
      return num_sampled;
435
436
437
438
    }
  };
}

439
template <SamplerType S>
440
441
442
443
444
445
446
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,
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
    const torch::optional<torch::Tensor>& edge_timestamp, SamplerArgs<S> args) {
  return
      [&seed_timestamp, &csc_indices, &fanouts, replace, &options,
       &type_per_edge, &probs_or_mask, &node_timestamp, &edge_timestamp, args](
          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, args, 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, args, picked_data_ptr);
          if (type_per_edge.has_value()) {
            std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
          }
          return num_sampled;
        }
      };
472
473
}

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

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

492
  AT_DISPATCH_INDEX_TYPES(
493
      indptr_.scalar_type(), "SampleNeighborsImplWrappedWithIndptr", ([&] {
494
495
        using indptr_t = index_t;
        AT_DISPATCH_INDEX_TYPES(
496
            nodes.scalar_type(), "SampleNeighborsImplWrappedWithNodes", ([&] {
497
              using nodes_t = index_t;
498
499
500
501
502
              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>();
503

504
505
506
507
508
509
510
511
512
513
514
515
              // 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;
516

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

524
525
526
527
              // 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());
528

529
530
531
532
533
534
535
536
537
538
              // 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());
              }
539

540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
              // 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(
555
                            i, offset, num_neighbors,
556
557
558
559
560
561
                            picked_eids_data_ptr + picked_offset);
                        TORCH_CHECK(
                            actual_picked_count == picked_number,
                            "Actual picked count doesn't match the calculated "
                            "pick "
                            "number.");
562

563
564
                        // Step 5. Calculate other attributes and return the
                        // subgraph.
565
                        AT_DISPATCH_INDEX_TYPES(
566
567
568
                            subgraph_indices.scalar_type(),
                            "IndexSelectSubgraphIndices", ([&] {
                              auto subgraph_indices_data_ptr =
569
                                  subgraph_indices.data_ptr<index_t>();
570
                              auto indices_data_ptr =
571
                                  indices_.data_ptr<index_t>();
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
                              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]];
                                }
                              }));
594
                        }
595
596
597
598
                      }
                    }
                  });
            }));
599
      }));
600

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

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

609
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
610
611
612
    const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
    bool replace, bool layer, bool return_eids,
    torch::optional<std::string> probs_name) const {
613
614
615
616
617
  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);
  }

618
619
  if (!replace && utils::is_on_gpu(nodes) &&
      utils::is_accessible_from_gpu(indptr_) &&
620
621
      utils::is_accessible_from_gpu(indices_) &&
      (!probs_or_mask.has_value() ||
622
623
624
       utils::is_accessible_from_gpu(probs_or_mask.value())) &&
      (!type_per_edge_.has_value() ||
       utils::is_accessible_from_gpu(type_per_edge_.value()))) {
625
626
627
628
629
630
631
632
633
    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()) {
634
635
636
637
638
639
640
641
    // 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);
    }
  }
642

643
644
645
646
647
  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(
648
        nodes, return_eids,
649
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
650
651
652
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
653
654
655
  } else {
    SamplerArgs<SamplerType::NEIGHBOR> args;
    return SampleNeighborsImpl(
656
        nodes, return_eids,
657
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
658
659
660
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
661
662
663
  }
}

664
665
666
667
c10::intrusive_ptr<FusedSampledSubgraph>
FusedCSCSamplingGraph::TemporalSampleNeighbors(
    const torch::Tensor& input_nodes,
    const torch::Tensor& input_nodes_timestamp,
668
669
    const std::vector<int64_t>& fanouts, bool replace, bool layer,
    bool return_eids, torch::optional<std::string> probs_name,
670
671
672
    torch::optional<std::string> node_timestamp_attr_name,
    torch::optional<std::string> edge_timestamp_attr_name) const {
  // 1. Get probs_or_mask.
673
674
675
676
677
678
679
680
681
682
  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);
    }
  }
683
  // 2. Get the timestamp attribute for nodes of the graph
684
  auto node_timestamp = this->NodeAttribute(node_timestamp_attr_name);
685
  // 3. Get the timestamp attribute for edges of the graph
686
687
  auto edge_timestamp = this->EdgeAttribute(edge_timestamp_attr_name);
  // 4. Call SampleNeighborsImpl
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
  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(
        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, args));
  } else {
    SamplerArgs<SamplerType::NEIGHBOR> args;
    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, args));
  }
713
714
}

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

735
736
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::CopyToSharedMemory(
737
    const std::string& shared_memory_name) {
738
  SharedMemoryHelper helper(shared_memory_name);
739
740
741
742
  helper.WriteTorchTensor(indptr_);
  helper.WriteTorchTensor(indices_);
  helper.WriteTorchTensor(node_type_offset_);
  helper.WriteTorchTensor(type_per_edge_);
743
744
  helper.WriteTorchTensorDict(TensorizeDict(node_type_to_id_));
  helper.WriteTorchTensorDict(TensorizeDict(edge_type_to_id_));
745
  helper.WriteTorchTensorDict(node_attributes_);
746
747
748
  helper.WriteTorchTensorDict(edge_attributes_);
  helper.Flush();
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
749
750
}

751
752
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::LoadFromSharedMemory(
753
    const std::string& shared_memory_name) {
754
  SharedMemoryHelper helper(shared_memory_name);
755
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
756
757
}

758
void FusedCSCSamplingGraph::HoldSharedMemoryObject(
759
760
761
762
763
    SharedMemoryPtr tensor_metadata_shm, SharedMemoryPtr tensor_data_shm) {
  tensor_metadata_shm_ = std::move(tensor_metadata_shm);
  tensor_data_shm_ = std::move(tensor_data_shm);
}

764
765
766
767
int64_t NumPick(
    int64_t fanout, bool replace,
    const torch::optional<torch::Tensor>& probs_or_mask, int64_t offset,
    int64_t num_neighbors) {
768
769
770
771
772
773
774
775
776
777
778
  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);
        }));
  }
779
780
781
782
  if (num_valid_neighbors == 0 || fanout == -1) return num_valid_neighbors;
  return replace ? fanout : std::min(fanout, num_valid_neighbors);
}

783
784
785
786
787
788
789
790
791
792
793
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));
794
    mask &= neighbor_timestamp < seed_timestamp;
795
796
  }
  if (edge_timestamp.has_value()) {
797
    mask &= edge_timestamp.value().slice(0, l, r) < seed_timestamp;
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
  }
  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);
}

820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
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;
}

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
877
878
879
880
881
882
883
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;
}

884
885
886
887
888
889
890
891
/**
 * @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.
892
893
894
 *  - 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).
895
896
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
897
 * @param replace Boolean indicating whether the sample is performed with or
898
899
900
 * 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.
901
902
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
903
 */
904
template <typename PickedType>
905
inline int64_t UniformPick(
906
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
907
    const torch::TensorOptions& options, PickedType* picked_data_ptr) {
908
  if ((fanout == -1) || (num_neighbors <= fanout && !replace)) {
909
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
910
    return num_neighbors;
911
  } else if (replace) {
912
913
914
915
916
    std::memcpy(
        picked_data_ptr,
        torch::randint(offset, offset + num_neighbors, {fanout}, options)
            .data_ptr<PickedType>(),
        fanout * sizeof(PickedType));
917
    return fanout;
918
  } else {
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
946
947
948
949
950
951
952
    // 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);
953
      return fanout;
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
    } 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;
      }
974
      return fanout;
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
    } 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);
998
      return picked_set.size();
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
1093
1094
1095
1096
1097
1098
1099
/** @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;
}

1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
/**
 * @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.
1118
1119
1120
1121
 *  - 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).
1122
1123
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
1124
 * @param replace Boolean indicating whether the sample is performed with or
1125
1126
1127
1128
1129
1130
1131
 * 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.
1132
1133
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
1134
 */
1135
template <typename PickedType>
1136
inline int64_t NonUniformPick(
1137
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
1138
    const torch::TensorOptions& options, const torch::Tensor& probs_or_mask,
1139
    PickedType* picked_data_ptr) {
1140
  auto local_probs =
1141
1142
1143
      probs_or_mask.size(0) > num_neighbors
          ? probs_or_mask.slice(0, offset, offset + num_neighbors)
          : probs_or_mask;
1144
1145
1146
1147
1148
  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;
1149
  }
1150
  return picked_indices.numel();
1151
1152
}

1153
template <typename PickedType>
1154
int64_t Pick(
1155
1156
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
1157
    const torch::optional<torch::Tensor>& probs_or_mask,
1158
    SamplerArgs<SamplerType::NEIGHBOR> args, PickedType* picked_data_ptr) {
1159
  if (probs_or_mask.has_value()) {
1160
    return NonUniformPick(
1161
        offset, num_neighbors, fanout, replace, options, probs_or_mask.value(),
1162
        picked_data_ptr);
1163
  } else {
1164
    return UniformPick(
1165
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
1166
1167
1168
  }
}

1169
template <SamplerType S, typename PickedType>
1170
1171
1172
1173
1174
1175
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,
1176
    const torch::optional<torch::Tensor>& edge_timestamp, SamplerArgs<S> args,
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
    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);
  }
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
  if constexpr (S == SamplerType::NEIGHBOR) {
    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();
  }
  if constexpr (S == SamplerType::LABOR) {
    return Pick(
        offset, num_neighbors, fanout, replace, options, masked_prob, args,
        picked_data_ptr);
1202
1203
1204
  }
}

1205
template <SamplerType S, typename PickedType>
1206
int64_t PickByEtype(
1207
1208
    int64_t offset, int64_t num_neighbors, const std::vector<int64_t>& fanouts,
    bool replace, const torch::TensorOptions& options,
1209
    const torch::Tensor& type_per_edge,
1210
1211
    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args,
    PickedType* picked_data_ptr) {
1212
1213
  int64_t etype_begin = offset;
  int64_t etype_end = offset;
1214
  int64_t pick_offset = 0;
1215
1216
1217
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "PickByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
1218
1219
1220
        const auto end = offset + num_neighbors;
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
1221
          TORCH_CHECK(
1222
              etype >= 0 && etype < (int64_t)fanouts.size(),
1223
              "Etype values exceed the number of fanouts.");
1224
          int64_t fanout = fanouts[etype];
1225
1226
1227
1228
          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;
1229
1230
          // Do sampling for one etype.
          if (fanout != 0) {
1231
            int64_t picked_count = Pick(
1232
                etype_begin, etype_end - etype_begin, fanout, replace, options,
1233
1234
                probs_or_mask, args, picked_data_ptr + pick_offset);
            pick_offset += picked_count;
1235
1236
1237
1238
          }
          etype_begin = etype_end;
        }
      }));
1239
  return pick_offset;
1240
1241
}

1242
template <SamplerType S, typename PickedType>
1243
1244
1245
1246
1247
1248
1249
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,
1250
    const torch::optional<torch::Tensor>& edge_timestamp, SamplerArgs<S> args,
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
    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,
1274
                probs_or_mask, node_timestamp, edge_timestamp, args,
1275
1276
1277
1278
1279
1280
1281
1282
1283
                picked_data_ptr + pick_offset);
            pick_offset += picked_count;
          }
          etype_begin = etype_end;
        }
      }));
  return pick_offset;
}

1284
template <typename PickedType>
1285
int64_t Pick(
1286
1287
1288
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1289
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1290
  if (fanout == 0) return 0;
1291
  if (probs_or_mask.has_value()) {
1292
    if (fanout < 0) {
1293
      return NonUniformPick(
1294
1295
          offset, num_neighbors, fanout, replace, options,
          probs_or_mask.value(), picked_data_ptr);
1296
    } else {
1297
      int64_t picked_count;
1298
1299
1300
      AT_DISPATCH_FLOATING_TYPES(
          probs_or_mask.value().scalar_type(), "LaborPickFloatType", ([&] {
            if (replace) {
1301
              picked_count = LaborPick<true, true, scalar_t>(
1302
1303
1304
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            } else {
1305
              picked_count = LaborPick<true, false, scalar_t>(
1306
1307
1308
1309
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            }
          }));
1310
      return picked_count;
1311
1312
    }
  } else if (fanout < 0) {
1313
    return UniformPick(
1314
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
1315
  } else if (replace) {
1316
    return LaborPick<false, true, float>(
1317
        offset, num_neighbors, fanout, options,
1318
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1319
  } else {  // replace = false
1320
    return LaborPick<false, false, float>(
1321
        offset, num_neighbors, fanout, options,
1322
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
  }
}

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.
1340
1341
1342
1343
 *  - 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).
1344
1345
1346
1347
1348
1349
1350
1351
 *  - 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.
1352
1353
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
1354
 */
1355
template <
1356
1357
    bool NonUniform, bool Replace, typename ProbsType, typename PickedType,
    int StackSize>
1358
inline int64_t LaborPick(
1359
1360
1361
    int64_t offset, int64_t num_neighbors, int64_t fanout,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1362
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1363
  fanout = Replace ? fanout : std::min(fanout, num_neighbors);
1364
  if (!NonUniform && !Replace && fanout >= num_neighbors) {
1365
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
1366
    return num_neighbors;
1367
1368
  }
  // Assuming max_degree of a vertex is <= 4 billion.
1369
1370
1371
1372
1373
1374
1375
1376
1377
  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>());
  }
1378
1379
1380
  const ProbsType* local_probs_data =
      NonUniform ? probs_or_mask.value().data_ptr<ProbsType>() + offset
                 : nullptr;
1381
1382
1383
  if (NonUniform && probs_or_mask.value().size(0) <= num_neighbors) {
    local_probs_data -= offset;
  }
1384
  AT_DISPATCH_INDEX_TYPES(
1385
      args.indices.scalar_type(), "LaborPickMain", ([&] {
1386
1387
        const index_t* local_indices_data =
            args.indices.data_ptr<index_t>() + offset;
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
        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))).
1410
1411
1412
1413
1414
1415
1416
1417
          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);
1418
1419
1420
          auto heap_end = heap_data;
          const auto init_count = (num_neighbors + fanout - 1) / num_neighbors;
          auto sample_neighbor_i_with_index_t_jth_time =
1421
              [&](index_t t, int64_t j, uint32_t i) {
1422
                auto rnd = labor::jth_sorted_uniform_random(
1423
                    args.random_seed, t, args.num_nodes, j, remaining_data[i],
1424
1425
1426
1427
1428
1429
                    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);
1430
1431
1432
                  if (++heap_end >= heap_data + fanout) {
                    std::make_heap(heap_data, heap_data + fanout);
                  }
1433
1434
1435
1436
1437
1438
1439
                  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 {
1440
                  remaining_data[i] = -1;
1441
1442
1443
1444
                  return true;
                }
              };
          for (uint32_t i = 0; i < num_neighbors; ++i) {
1445
            const auto t = local_indices_data[i];
1446
1447
1448
1449
1450
            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) {
1451
            if (remaining_data[i] == -1) continue;
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
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
            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;
1503
1504
1505
1506
1507
1508
  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;
    }
  }
1509
  return num_sampled;
1510
1511
}

1512
1513
}  // namespace sampling
}  // namespace graphbolt