"src/vscode:/vscode.git/clone" did not exist on "ec64f371b1ef52b1c2c0c146acea979c48b4a6db"
fused_csc_sampling_graph.cc 52.2 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<EdgeAttrMap>& edge_attributes)
60
    : indptr_(indptr),
61
      indices_(indices),
62
      node_type_offset_(node_type_offset),
63
      type_per_edge_(type_per_edge),
64
65
      node_type_to_id_(node_type_to_id),
      edge_type_to_id_(edge_type_to_id),
66
      edge_attributes_(edge_attributes) {
67
68
69
70
71
  TORCH_CHECK(indptr.dim() == 1);
  TORCH_CHECK(indices.dim() == 1);
  TORCH_CHECK(indptr.device() == indices.device());
}

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

102
void FusedCSCSamplingGraph::Load(torch::serialize::InputArchive& archive) {
103
  const int64_t magic_num =
104
      read_from_archive(archive, "FusedCSCSamplingGraph/magic_num").toInt();
105
106
  TORCH_CHECK(
      magic_num == kCSCSamplingGraphSerializeMagic,
107
108
109
110
111
112
      "Magic numbers mismatch when loading FusedCSCSamplingGraph.");
  indptr_ =
      read_from_archive(archive, "FusedCSCSamplingGraph/indptr").toTensor();
  indices_ =
      read_from_archive(archive, "FusedCSCSamplingGraph/indices").toTensor();
  if (read_from_archive(archive, "FusedCSCSamplingGraph/has_node_type_offset")
113
114
          .toBool()) {
    node_type_offset_ =
115
        read_from_archive(archive, "FusedCSCSamplingGraph/node_type_offset")
116
117
            .toTensor();
  }
118
  if (read_from_archive(archive, "FusedCSCSamplingGraph/has_type_per_edge")
119
120
          .toBool()) {
    type_per_edge_ =
121
122
        read_from_archive(archive, "FusedCSCSamplingGraph/type_per_edge")
            .toTensor();
123
  }
124

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
  if (read_from_archive(archive, "FusedCSCSamplingGraph/has_node_type_to_id")
          .toBool()) {
    torch::Dict<torch::IValue, torch::IValue> generic_dict =
        read_from_archive(archive, "FusedCSCSamplingGraph/node_type_to_id")
            .toGenericDict();
    NodeTypeToIDMap node_type_to_id;
    for (const auto& pair : generic_dict) {
      std::string key = pair.key().toStringRef();
      int64_t value = pair.value().toInt();
      node_type_to_id.insert(std::move(key), value);
    }
    node_type_to_id_ = std::move(node_type_to_id);
  }

  if (read_from_archive(archive, "FusedCSCSamplingGraph/has_edge_type_to_id")
          .toBool()) {
    torch::Dict<torch::IValue, torch::IValue> generic_dict =
        read_from_archive(archive, "FusedCSCSamplingGraph/edge_type_to_id")
            .toGenericDict();
    EdgeTypeToIDMap edge_type_to_id;
    for (const auto& pair : generic_dict) {
      std::string key = pair.key().toStringRef();
      int64_t value = pair.value().toInt();
      edge_type_to_id.insert(std::move(key), value);
    }
    edge_type_to_id_ = std::move(edge_type_to_id);
  }

153
154
155
  // Optional edge attributes.
  torch::IValue has_edge_attributes;
  if (archive.try_read(
156
          "FusedCSCSamplingGraph/has_edge_attributes", has_edge_attributes) &&
157
158
      has_edge_attributes.toBool()) {
    torch::Dict<torch::IValue, torch::IValue> generic_dict =
159
        read_from_archive(archive, "FusedCSCSamplingGraph/edge_attributes")
160
161
162
163
164
165
166
167
168
169
170
            .toGenericDict();
    EdgeAttrMap target_dict;
    for (const auto& pair : generic_dict) {
      std::string key = pair.key().toStringRef();
      torch::Tensor value = pair.value().toTensor();
      // Use move to avoid copy.
      target_dict.insert(std::move(key), std::move(value));
    }
    // Same as above.
    edge_attributes_ = std::move(target_dict);
  }
171
172
}

173
174
175
176
177
178
void FusedCSCSamplingGraph::Save(
    torch::serialize::OutputArchive& archive) const {
  archive.write(
      "FusedCSCSamplingGraph/magic_num", kCSCSamplingGraphSerializeMagic);
  archive.write("FusedCSCSamplingGraph/indptr", indptr_);
  archive.write("FusedCSCSamplingGraph/indices", indices_);
179
  archive.write(
180
181
      "FusedCSCSamplingGraph/has_node_type_offset",
      node_type_offset_.has_value());
182
183
  if (node_type_offset_) {
    archive.write(
184
        "FusedCSCSamplingGraph/node_type_offset", node_type_offset_.value());
185
186
  }
  archive.write(
187
      "FusedCSCSamplingGraph/has_type_per_edge", type_per_edge_.has_value());
188
  if (type_per_edge_) {
189
190
    archive.write(
        "FusedCSCSamplingGraph/type_per_edge", type_per_edge_.value());
191
  }
192
193
194
195
196
197
198
199
200
201
202
203
204
205
  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());
  }
206
  archive.write(
207
208
      "FusedCSCSamplingGraph/has_edge_attributes",
      edge_attributes_.has_value());
209
  if (edge_attributes_) {
210
211
    archive.write(
        "FusedCSCSamplingGraph/edge_attributes", edge_attributes_.value());
212
  }
213
214
}

215
void FusedCSCSamplingGraph::SetState(
216
217
218
219
220
221
222
223
224
    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})),
225
      "Version number mismatches when loading pickled FusedCSCSamplingGraph.")
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
  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");
  }
  if (state.find("edge_attributes") != state.end()) {
    edge_attributes_ = state.at("edge_attributes");
  }
}

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

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

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

297
  return c10::make_intrusive<FusedSampledSubgraph>(
298
      indptr.cumsum(0), torch::cat(indices_arr), original_column_node_ids,
299
300
301
302
303
304
      torch::arange(0, NumNodes()), torch::cat(edge_ids_arr),
      type_per_edge_
          ? torch::optional<torch::Tensor>{torch::cat(type_per_edge_arr)}
          : torch::nullopt);
}

305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
/**
 * @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.
 *
 * @return A lambda function (int64_t offset, int64_t num_neighbors) ->
 * torch::Tensor, which takes 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.
 */
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](
             int64_t offset, int64_t num_neighbors) {
    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);
    }
  };
}

344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
/**
 * @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.
 *
361
362
363
364
365
 * @return A lambda function: (int64_t offset, int64_t num_neighbors,
 * PickedType* picked_data_ptr) -> torch::Tensor, which takes 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.
366
 */
367
template <SamplerType S>
368
369
370
371
372
373
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](
374
375
376
377
             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.
378
379
380
    if (fanouts.size() > 1) {
      return PickByEtype(
          offset, num_neighbors, fanouts, replace, options,
381
          type_per_edge.value(), probs_or_mask, args, picked_data_ptr);
382
    } else {
383
      int64_t num_sampled = Pick(
384
          offset, num_neighbors, fanouts[0], replace, options, probs_or_mask,
385
          args, picked_data_ptr);
386
387
388
389
      if (type_per_edge) {
        std::sort(picked_data_ptr, picked_data_ptr + num_sampled);
      }
      return num_sampled;
390
391
392
393
    }
  };
}

394
template <typename NumPickFn, typename PickFn>
395
396
c10::intrusive_ptr<FusedSampledSubgraph>
FusedCSCSamplingGraph::SampleNeighborsImpl(
397
398
    const torch::Tensor& nodes, bool return_eids, NumPickFn num_pick_fn,
    PickFn pick_fn) const {
399
  const int64_t num_nodes = nodes.size(0);
400
  const auto indptr_options = indptr_.options();
401
  torch::Tensor num_picked_neighbors_per_node =
402
      torch::empty({num_nodes + 1}, indptr_options);
403

404
405
406
  // Calculate GrainSize for parallel_for.
  // Set the default grain size to 64.
  const int64_t grain_size = 64;
407
408
409
410
411
  torch::Tensor picked_eids;
  torch::Tensor subgraph_indptr;
  torch::Tensor subgraph_indices;
  torch::optional<torch::Tensor> subgraph_type_per_edge = torch::nullopt;

412
  AT_DISPATCH_INTEGRAL_TYPES(
413
414
415
416
417
418
419
420
421
422
      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>();
423

424
425
426
427
428
429
430
431
432
433
434
435
              // 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;
436

437
438
439
440
441
442
                      num_picked_neighbors_data_ptr[i + 1] =
                          num_neighbors == 0
                              ? 0
                              : num_pick_fn(offset, num_neighbors);
                    }
                  });
443

444
445
446
447
              // 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());
448

449
450
451
452
453
454
455
456
457
458
              // 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());
              }
459

460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
              // 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(
                            offset, num_neighbors,
                            picked_eids_data_ptr + picked_offset);
                        TORCH_CHECK(
                            actual_picked_count == picked_number,
                            "Actual picked count doesn't match the calculated "
                            "pick "
                            "number.");
482

483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
                        // 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]];
                                }
                              }));
514
                        }
515
516
517
518
                      }
                    }
                  });
            }));
519
      }));
520

521
522
  torch::optional<torch::Tensor> subgraph_reverse_edge_ids = torch::nullopt;
  if (return_eids) subgraph_reverse_edge_ids = std::move(picked_eids);
523

524
  return c10::make_intrusive<FusedSampledSubgraph>(
525
      subgraph_indptr, subgraph_indices, nodes, torch::nullopt,
526
      subgraph_reverse_edge_ids, subgraph_type_per_edge);
527
528
}

529
c10::intrusive_ptr<FusedSampledSubgraph> FusedCSCSamplingGraph::SampleNeighbors(
530
531
532
533
534
535
536
537
538
539
540
541
542
543
    const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
    bool replace, bool layer, bool return_eids,
    torch::optional<std::string> probs_name) const {
  torch::optional<torch::Tensor> probs_or_mask = torch::nullopt;
  if (probs_name.has_value() && !probs_name.value().empty()) {
    probs_or_mask = edge_attributes_.value().at(probs_name.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);
    }
  }
544

545
546
547
548
549
  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(
550
        nodes, return_eids,
551
        GetNumPickFn(fanouts, replace, type_per_edge_, probs_or_mask),
552
553
554
        GetPickFn(
            fanouts, replace, indptr_.options(), type_per_edge_, probs_or_mask,
            args));
555
556
557
  } else {
    SamplerArgs<SamplerType::NEIGHBOR> args;
    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
  }
}

566
std::tuple<torch::Tensor, torch::Tensor>
567
FusedCSCSamplingGraph::SampleNegativeEdgesUniform(
568
569
570
571
572
573
574
575
576
577
    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);
}

578
579
static c10::intrusive_ptr<FusedCSCSamplingGraph>
BuildGraphFromSharedMemoryHelper(SharedMemoryHelper&& helper) {
580
581
582
583
584
  helper.InitializeRead();
  auto indptr = helper.ReadTorchTensor();
  auto indices = helper.ReadTorchTensor();
  auto node_type_offset = helper.ReadTorchTensor();
  auto type_per_edge = helper.ReadTorchTensor();
585
586
  auto node_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
  auto edge_type_to_id = DetensorizeDict(helper.ReadTorchTensorDict());
587
  auto edge_attributes = helper.ReadTorchTensorDict();
588
  auto graph = c10::make_intrusive<FusedCSCSamplingGraph>(
589
      indptr.value(), indices.value(), node_type_offset, type_per_edge,
590
      node_type_to_id, edge_type_to_id, edge_attributes);
591
592
593
  auto shared_memory = helper.ReleaseSharedMemory();
  graph->HoldSharedMemoryObject(
      std::move(shared_memory.first), std::move(shared_memory.second));
594
595
596
  return graph;
}

597
598
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::CopyToSharedMemory(
599
    const std::string& shared_memory_name) {
600
601
602
603
604
  SharedMemoryHelper helper(shared_memory_name, SERIALIZED_METAINFO_SIZE_MAX);
  helper.WriteTorchTensor(indptr_);
  helper.WriteTorchTensor(indices_);
  helper.WriteTorchTensor(node_type_offset_);
  helper.WriteTorchTensor(type_per_edge_);
605
606
  helper.WriteTorchTensorDict(TensorizeDict(node_type_to_id_));
  helper.WriteTorchTensorDict(TensorizeDict(edge_type_to_id_));
607
608
609
  helper.WriteTorchTensorDict(edge_attributes_);
  helper.Flush();
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
610
611
}

612
613
c10::intrusive_ptr<FusedCSCSamplingGraph>
FusedCSCSamplingGraph::LoadFromSharedMemory(
614
    const std::string& shared_memory_name) {
615
616
  SharedMemoryHelper helper(shared_memory_name, SERIALIZED_METAINFO_SIZE_MAX);
  return BuildGraphFromSharedMemoryHelper(std::move(helper));
617
618
}

619
void FusedCSCSamplingGraph::HoldSharedMemoryObject(
620
621
622
623
624
    SharedMemoryPtr tensor_metadata_shm, SharedMemoryPtr tensor_data_shm) {
  tensor_metadata_shm_ = std::move(tensor_metadata_shm);
  tensor_data_shm_ = std::move(tensor_data_shm);
}

625
626
627
628
int64_t NumPick(
    int64_t fanout, bool replace,
    const torch::optional<torch::Tensor>& probs_or_mask, int64_t offset,
    int64_t num_neighbors) {
629
630
631
632
633
634
635
636
637
638
639
  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);
        }));
  }
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
  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;
}

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

793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
/**
 * @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.
811
812
813
814
 *  - 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).
815
816
 *  - When the value is a non-negative integer, it serves as a minimum
 * threshold for selecting neighbors.
817
 * @param replace Boolean indicating whether the sample is performed with or
818
819
820
821
822
823
824
 * 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.
825
826
 * @param picked_data_ptr The destination address where the picked neighbors
 * should be put. Enough memory space should be allocated in advance.
827
 */
828
template <typename PickedType>
829
inline int64_t NonUniformPick(
830
831
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
832
833
    const torch::optional<torch::Tensor>& probs_or_mask,
    PickedType* picked_data_ptr) {
834
835
836
837
  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);
838
  if (num_positive_probs == 0) return 0;
839
  if ((fanout == -1) || (num_positive_probs <= fanout && !replace)) {
840
841
842
843
    std::memcpy(
        picked_data_ptr,
        (positive_probs_indices + offset).data_ptr<PickedType>(),
        num_positive_probs * sizeof(PickedType));
844
    return num_positive_probs;
845
846
  } else {
    if (!replace) fanout = std::min(fanout, num_positive_probs);
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
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
    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;
            }
          }
        }));
932
    return fanout;
933
934
935
  }
}

936
template <typename PickedType>
937
int64_t Pick(
938
939
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
    const torch::TensorOptions& options,
940
    const torch::optional<torch::Tensor>& probs_or_mask,
941
    SamplerArgs<SamplerType::NEIGHBOR> args, PickedType* picked_data_ptr) {
942
  if (probs_or_mask.has_value()) {
943
    return NonUniformPick(
944
945
        offset, num_neighbors, fanout, replace, options, probs_or_mask,
        picked_data_ptr);
946
  } else {
947
    return UniformPick(
948
        offset, num_neighbors, fanout, replace, options, picked_data_ptr);
949
950
951
  }
}

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

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

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

1214
1215
}  // namespace sampling
}  // namespace graphbolt