fused_csc_sampling_graph.cc 52.6 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/fused_csc_sampling_graph.h>
8
#include <graphbolt/serialize.h>
9
10
#include <torch/torch.h>

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

19
#include "./random.h"
20
#include "./shared_memory_helper.h"
21

22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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

48
49
50
namespace graphbolt {
namespace sampling {

51
52
static const int kPickleVersion = 6199;

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

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

110
void FusedCSCSamplingGraph::Load(torch::serialize::InputArchive& archive) {
111
  const int64_t magic_num =
112
      read_from_archive<int64_t>(archive, "FusedCSCSamplingGraph/magic_num");
113
114
  TORCH_CHECK(
      magic_num == kCSCSamplingGraphSerializeMagic,
115
116
      "Magic numbers mismatch when loading FusedCSCSamplingGraph.");
  indptr_ =
117
118
119
120
121
122
123
      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");
124
  }
125
126
127
128
  if (read_from_archive<bool>(
          archive, "FusedCSCSamplingGraph/has_type_per_edge")) {
    type_per_edge_ = read_from_archive<torch::Tensor>(
        archive, "FusedCSCSamplingGraph/type_per_edge");
129
  }
130

131
132
133
134
  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");
135
136
  }

137
138
139
140
  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");
141
142
  }

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

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

204
void FusedCSCSamplingGraph::SetState(
205
206
207
208
209
210
211
212
213
    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})),
214
      "Version number mismatches when loading pickled FusedCSCSamplingGraph.")
215
216
217
218
219
220
221
222
223
  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");
  }
224
225
226
227
228
229
  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"));
  }
230
231
232
  if (state.find("node_attributes") != state.end()) {
    node_attributes_ = state.at("node_attributes");
  }
233
234
235
236
237
238
  if (state.find("edge_attributes") != state.end()) {
    edge_attributes_ = state.at("edge_attributes");
  }
}

torch::Dict<std::string, torch::Dict<std::string, torch::Tensor>>
239
FusedCSCSamplingGraph::GetState() const {
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
  // 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);
257
258
259
260
261
262
  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());
  }
263
264
265
  if (node_attributes_.has_value()) {
    state.insert("node_attributes", node_attributes_.value());
  }
266
267
268
269
270
271
  if (edge_attributes_.has_value()) {
    state.insert("edge_attributes", edge_attributes_.value());
  }
  return state;
}

272
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::InSubgraph(
273
274
275
    const torch::Tensor& nodes) const {
  using namespace torch::indexing;
  const int32_t kDefaultGrainSize = 100;
276
277
  const auto num_seeds = nodes.size(0);
  torch::Tensor indptr = torch::zeros({num_seeds + 1}, indptr_.dtype());
278
  std::vector<torch::Tensor> indices_arr(num_seeds);
279
280
  torch::Tensor original_column_node_ids =
      torch::zeros({num_seeds}, indptr_.dtype());
281
282
  std::vector<torch::Tensor> edge_ids_arr(num_seeds);
  std::vector<torch::Tensor> type_per_edge_arr(num_seeds);
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303

  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);
                }
              }
            });
      }));

304
  return c10::make_intrusive<FusedSampledSubgraph>(
305
      indptr.cumsum(0), torch::cat(indices_arr), original_column_node_ids,
306
307
308
309
310
311
      torch::arange(0, NumNodes()), torch::cat(edge_ids_arr),
      type_per_edge_
          ? torch::optional<torch::Tensor>{torch::cat(type_per_edge_arr)}
          : torch::nullopt);
}

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
/**
 * @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.
 *
328
329
330
331
332
 * @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.
333
334
335
336
337
338
339
340
 */
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](
341
             int64_t seed_offset, int64_t offset, int64_t num_neighbors) {
342
343
344
345
346
347
348
349
350
351
    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);
    }
  };
}

352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
/**
 * @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.
 *
369
370
371
372
373
 * @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.
374
 */
375
template <SamplerType S>
376
377
378
379
380
381
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](
382
383
             int64_t seed_offset, int64_t offset, int64_t num_neighbors,
             auto picked_data_ptr) {
384
385
386
    // 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.
387
388
389
    if (fanouts.size() > 1) {
      return PickByEtype(
          offset, num_neighbors, fanouts, replace, options,
390
          type_per_edge.value(), probs_or_mask, args, picked_data_ptr);
391
    } else {
392
      int64_t num_sampled = Pick(
393
          offset, num_neighbors, fanouts[0], replace, options, probs_or_mask,
394
          args, picked_data_ptr);
395
396
397
398
      if (type_per_edge) {
        std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
      }
      return num_sampled;
399
400
401
402
    }
  };
}

403
template <typename NumPickFn, typename PickFn>
404
405
c10::intrusive_ptr<FusedSampledSubgraph>
FusedCSCSamplingGraph::SampleNeighborsImpl(
406
407
    const torch::Tensor& nodes, bool return_eids, NumPickFn num_pick_fn,
    PickFn pick_fn) const {
408
  const int64_t num_nodes = nodes.size(0);
409
  const auto indptr_options = indptr_.options();
410
  torch::Tensor num_picked_neighbors_per_node =
411
      torch::empty({num_nodes + 1}, indptr_options);
412

413
414
415
  // Calculate GrainSize for parallel_for.
  // Set the default grain size to 64.
  const int64_t grain_size = 64;
416
417
418
419
420
  torch::Tensor picked_eids;
  torch::Tensor subgraph_indptr;
  torch::Tensor subgraph_indices;
  torch::optional<torch::Tensor> subgraph_type_per_edge = torch::nullopt;

421
  AT_DISPATCH_INTEGRAL_TYPES(
422
423
424
425
426
427
428
429
430
431
      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>();
432

433
434
435
436
437
438
439
440
441
442
443
444
              // 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;
445

446
447
448
                      num_picked_neighbors_data_ptr[i + 1] =
                          num_neighbors == 0
                              ? 0
449
                              : num_pick_fn(i, offset, num_neighbors);
450
451
                    }
                  });
452

453
454
455
456
              // 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());
457

458
459
460
461
462
463
464
465
466
467
              // 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());
              }
468

469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
              // 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(
484
                            i, offset, num_neighbors,
485
486
487
488
489
490
                            picked_eids_data_ptr + picked_offset);
                        TORCH_CHECK(
                            actual_picked_count == picked_number,
                            "Actual picked count doesn't match the calculated "
                            "pick "
                            "number.");
491

492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
                        // 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]];
                                }
                              }));
523
                        }
524
525
526
527
                      }
                    }
                  });
            }));
528
      }));
529

530
531
  torch::optional<torch::Tensor> subgraph_reverse_edge_ids = torch::nullopt;
  if (return_eids) subgraph_reverse_edge_ids = std::move(picked_eids);
532

533
  return c10::make_intrusive<FusedSampledSubgraph>(
534
      subgraph_indptr, subgraph_indices, nodes, torch::nullopt,
535
      subgraph_reverse_edge_ids, subgraph_type_per_edge);
536
537
}

538
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
539
540
541
    const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
    bool replace, bool layer, bool return_eids,
    torch::optional<std::string> probs_name) const {
542
543
  auto probs_or_mask = this->EdgeAttribute(probs_name);
  if (probs_name.has_value()) {
544
545
546
547
548
549
550
551
    // 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);
    }
  }
552

553
554
555
556
557
  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(
558
        nodes, return_eids,
559
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
560
561
562
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
563
564
565
  } else {
    SamplerArgs<SamplerType::NEIGHBOR> args;
    return SampleNeighborsImpl(
566
        nodes, return_eids,
567
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
568
569
570
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
571
572
573
  }
}

574
std::tuple<torch::Tensor, torch::Tensor>
575
FusedCSCSamplingGraph::SampleNegativeEdgesUniform(
576
577
578
579
580
581
582
583
584
585
    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);
}

586
587
static c10::intrusive_ptr<FusedCSCSamplingGraph>
BuildGraphFromSharedMemoryHelper(SharedMemoryHelper&& helper) {
588
589
590
591
592
  helper.InitializeRead();
  auto indptr = helper.ReadTorchTensor();
  auto indices = helper.ReadTorchTensor();
  auto node_type_offset = helper.ReadTorchTensor();
  auto type_per_edge = helper.ReadTorchTensor();
593
594
  auto node_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
  auto edge_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
595
  auto node_attributes = helper.ReadTorchTensorDict();
596
  auto edge_attributes = helper.ReadTorchTensorDict();
597
  auto graph = c10::make_intrusive<FusedCSCSamplingGraph>(
598
      indptr.value(), indices.value(), node_type_offset, type_per_edge,
599
      node_type_to_id, edge_type_to_id, node_attributes, edge_attributes);
600
601
602
  auto shared_memory = helper.ReleaseSharedMemory();
  graph->HoldSharedMemoryObject(
      std::move(shared_memory.first), std::move(shared_memory.second));
603
604
605
  return graph;
}

606
607
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::CopyToSharedMemory(
608
    const std::string& shared_memory_name) {
609
  SharedMemoryHelper helper(shared_memory_name);
610
611
612
613
  helper.WriteTorchTensor(indptr_);
  helper.WriteTorchTensor(indices_);
  helper.WriteTorchTensor(node_type_offset_);
  helper.WriteTorchTensor(type_per_edge_);
614
615
  helper.WriteTorchTensorDict(TensorizeDict(node_type_to_id_));
  helper.WriteTorchTensorDict(TensorizeDict(edge_type_to_id_));
616
  helper.WriteTorchTensorDict(node_attributes_);
617
618
619
  helper.WriteTorchTensorDict(edge_attributes_);
  helper.Flush();
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
620
621
}

622
623
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::LoadFromSharedMemory(
624
    const std::string& shared_memory_name) {
625
  SharedMemoryHelper helper(shared_memory_name);
626
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
627
628
}

629
void FusedCSCSamplingGraph::HoldSharedMemoryObject(
630
631
632
633
634
    SharedMemoryPtr tensor_metadata_shm, SharedMemoryPtr tensor_data_shm) {
  tensor_metadata_shm_ = std::move(tensor_metadata_shm);
  tensor_data_shm_ = std::move(tensor_data_shm);
}

635
636
637
638
int64_t NumPick(
    int64_t fanout, bool replace,
    const torch::optional<torch::Tensor>& probs_or_mask, int64_t offset,
    int64_t num_neighbors) {
639
640
641
642
643
644
645
646
647
648
649
  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);
        }));
  }
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
  if (num_valid_neighbors == 0 || fanout == -1) return num_valid_neighbors;
  return replace ? fanout : std::min(fanout, num_valid_neighbors);
}

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;
}

684
685
686
687
688
689
690
691
/**
 * @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.
692
693
694
 *  - 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).
695
696
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
697
 * @param replace Boolean indicating whether the sample is performed with or
698
699
700
 * 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.
701
702
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
703
 */
704
template <typename PickedType>
705
inline int64_t UniformPick(
706
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
707
    const torch::TensorOptions& options, PickedType* picked_data_ptr) {
708
  if ((fanout == -1) || (num_neighbors <= fanout && !replace)) {
709
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
710
    return num_neighbors;
711
  } else if (replace) {
712
713
714
715
716
    std::memcpy(
        picked_data_ptr,
        torch::randint(offset, offset + num_neighbors, {fanout}, options)
            .data_ptr<PickedType>(),
        fanout * sizeof(PickedType));
717
    return fanout;
718
  } else {
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
    // 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);
753
      return fanout;
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
    } 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;
      }
774
      return fanout;
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
    } 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);
798
      return picked_set.size();
799
    }
800
801
802
  }
}

803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
/**
 * @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.
821
822
823
824
 *  - 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).
825
826
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
827
 * @param replace Boolean indicating whether the sample is performed with or
828
829
830
831
832
833
834
 * 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.
835
836
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
837
 */
838
template <typename PickedType>
839
inline int64_t NonUniformPick(
840
841
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
842
843
    const torch::optional<torch::Tensor>& probs_or_mask,
    PickedType* picked_data_ptr) {
844
845
846
847
  auto local_probs =
      probs_or_mask.value().slice(0, offset, offset + num_neighbors);
  auto positive_probs_indices = local_probs.nonzero().squeeze(1);
  auto num_positive_probs = positive_probs_indices.size(0);
848
  if (num_positive_probs == 0) return 0;
849
  if ((fanout == -1) || (num_positive_probs <= fanout && !replace)) {
850
851
852
853
    std::memcpy(
        picked_data_ptr,
        (positive_probs_indices + offset).data_ptr<PickedType>(),
        num_positive_probs * sizeof(PickedType));
854
    return num_positive_probs;
855
856
  } else {
    if (!replace) fanout = std::min(fanout, num_positive_probs);
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
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
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
    if (fanout == 0) return 0;
    AT_DISPATCH_FLOATING_TYPES(
        local_probs.scalar_type(), "MultinomialSampling", ([&] {
          auto local_probs_data_ptr = local_probs.data_ptr<scalar_t>();
          auto positive_probs_indices_ptr =
              positive_probs_indices.data_ptr<PickedType>();

          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;
              PickedType 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 =
                    local_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];
                }
              }
              *picked_data_ptr = max_prob_index + offset;
            } else {
              // Return topk(p / q).
              std::vector<std::pair<scalar_t, PickedType>> q(
                  num_positive_probs);
              for (auto i = 0; i < num_positive_probs; ++i) {
                q[i].first =
                    local_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) {
                  picked_data_ptr[i] = q[i].second + offset;
                }
              } else {
                // Use nth_element.
                std::nth_element(
                    q.begin(), q.begin() + fanout - 1, q.end(), std::greater{});
                for (auto i = 0; i < fanout; ++i) {
                  picked_data_ptr[i] = q[i].second + offset;
                }
              }
            }
          } 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 += local_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();
              picked_data_ptr[i] =
                  positive_probs_indices_ptr[sampled_index] + offset;
            }
          }
        }));
942
    return fanout;
943
944
945
  }
}

946
template <typename PickedType>
947
int64_t Pick(
948
949
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
950
    const torch::optional<torch::Tensor>& probs_or_mask,
951
    SamplerArgs<SamplerType::NEIGHBOR> args, PickedType* picked_data_ptr) {
952
  if (probs_or_mask.has_value()) {
953
    return NonUniformPick(
954
955
        offset, num_neighbors, fanout, replace, options, probs_or_mask,
        picked_data_ptr);
956
  } else {
957
    return UniformPick(
958
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
959
960
961
  }
}

962
template <SamplerType S, typename PickedType>
963
int64_t PickByEtype(
964
965
    int64_t offset, int64_t num_neighbors, const std::vector<int64_t>& fanouts,
    bool replace, const torch::TensorOptions& options,
966
    const torch::Tensor& type_per_edge,
967
968
    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args,
    PickedType* picked_data_ptr) {
969
970
  int64_t etype_begin = offset;
  int64_t etype_end = offset;
971
  int64_t pick_offset = 0;
972
973
974
  AT_DISPATCH_INTEGRAL_TYPES(
      type_per_edge.scalar_type(), "PickByEtype", ([&] {
        const scalar_t* type_per_edge_data = type_per_edge.data_ptr<scalar_t>();
975
976
977
        const auto end = offset + num_neighbors;
        while (etype_begin < end) {
          scalar_t etype = type_per_edge_data[etype_begin];
978
          TORCH_CHECK(
979
              etype >= 0 && etype < (int64_t)fanouts.size(),
980
              "Etype values exceed the number of fanouts.");
981
          int64_t fanout = fanouts[etype];
982
983
984
985
          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;
986
987
          // Do sampling for one etype.
          if (fanout != 0) {
988
            int64_t picked_count = Pick(
989
                etype_begin, etype_end - etype_begin, fanout, replace, options,
990
991
                probs_or_mask, args, picked_data_ptr + pick_offset);
            pick_offset += picked_count;
992
993
994
995
          }
          etype_begin = etype_end;
        }
      }));
996
  return pick_offset;
997
998
}

999
template <typename PickedType>
1000
int64_t Pick(
1001
1002
1003
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1004
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1005
  if (fanout == 0) return 0;
1006
  if (probs_or_mask.has_value()) {
1007
    if (fanout < 0) {
1008
      return NonUniformPick(
1009
1010
1011
          offset, num_neighbors, fanout, replace, options, probs_or_mask,
          picked_data_ptr);
    } else {
1012
      int64_t picked_count;
1013
1014
1015
      AT_DISPATCH_FLOATING_TYPES(
          probs_or_mask.value().scalar_type(), "LaborPickFloatType", ([&] {
            if (replace) {
1016
              picked_count = LaborPick<true, true, scalar_t>(
1017
1018
1019
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            } else {
1020
              picked_count = LaborPick<true, false, scalar_t>(
1021
1022
1023
1024
                  offset, num_neighbors, fanout, options, probs_or_mask, args,
                  picked_data_ptr);
            }
          }));
1025
      return picked_count;
1026
1027
    }
  } else if (fanout < 0) {
1028
    return UniformPick(
1029
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
1030
  } else if (replace) {
1031
    return LaborPick<false, true, float>(
1032
        offset, num_neighbors, fanout, options,
1033
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1034
  } else {  // replace = false
1035
    return LaborPick<false, false, float>(
1036
        offset, num_neighbors, fanout, options,
1037
        /* probs_or_mask= */ torch::nullopt, args, picked_data_ptr);
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
  }
}

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.
1055
1056
1057
1058
 *  - 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).
1059
1060
1061
1062
1063
1064
1065
1066
 *  - 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.
1067
1068
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
1069
 */
1070
template <
1071
1072
    bool NonUniform, bool Replace, typename ProbsType, typename PickedType,
    int StackSize>
1073
inline int64_t LaborPick(
1074
1075
1076
    int64_t offset, int64_t num_neighbors, int64_t fanout,
    const torch::TensorOptions& options,
    const torch::optional<torch::Tensor>& probs_or_mask,
1077
    SamplerArgs<SamplerType::LABOR> args, PickedType* picked_data_ptr) {
1078
  fanout = Replace ? fanout : std::min(fanout, num_neighbors);
1079
  if (!NonUniform && !Replace && fanout >= num_neighbors) {
1080
    std::iota(picked_data_ptr, picked_data_ptr + num_neighbors, offset);
1081
    return num_neighbors;
1082
1083
  }
  // Assuming max_degree of a vertex is <= 4 billion.
1084
1085
1086
1087
1088
1089
1090
1091
1092
  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>());
  }
1093
1094
1095
  const ProbsType* local_probs_data =
      NonUniform ? probs_or_mask.value().data_ptr<ProbsType>() + offset
                 : nullptr;
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
  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))).
1122
1123
1124
1125
1126
1127
1128
1129
          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);
1130
1131
1132
1133
1134
          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(
1135
                    args.random_seed, t, args.num_nodes, j, remaining_data[i],
1136
1137
1138
1139
1140
1141
                    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);
1142
1143
1144
                  if (++heap_end >= heap_data + fanout) {
                    std::make_heap(heap_data, heap_data + fanout);
                  }
1145
1146
1147
1148
1149
1150
1151
                  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 {
1152
                  remaining_data[i] = -1;
1153
1154
1155
1156
                  return true;
                }
              };
          for (uint32_t i = 0; i < num_neighbors; ++i) {
1157
            const auto t = local_indices_data[i];
1158
1159
1160
1161
1162
            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) {
1163
            if (remaining_data[i] == -1) continue;
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
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
            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;
1215
1216
1217
1218
1219
1220
  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;
    }
  }
1221
  return num_sampled;
1222
1223
}

1224
1225
}  // namespace sampling
}  // namespace graphbolt