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

7
8
#include <graphbolt/csc_sampling_graph.h>
#include <graphbolt/serialize.h>
9
10
11
12
#include <torch/torch.h>

#include <tuple>
#include <vector>
13

14
15
#include "./shared_memory_utils.h"

16
17
18
19
namespace graphbolt {
namespace sampling {

CSCSamplingGraph::CSCSamplingGraph(
20
    const torch::Tensor& indptr, const torch::Tensor& indices,
21
22
23
    const torch::optional<torch::Tensor>& node_type_offset,
    const torch::optional<torch::Tensor>& type_per_edge)
    : indptr_(indptr),
24
      indices_(indices),
25
26
      node_type_offset_(node_type_offset),
      type_per_edge_(type_per_edge) {
27
28
29
30
31
32
  TORCH_CHECK(indptr.dim() == 1);
  TORCH_CHECK(indices.dim() == 1);
  TORCH_CHECK(indptr.device() == indices.device());
}

c10::intrusive_ptr<CSCSamplingGraph> CSCSamplingGraph::FromCSC(
33
    const torch::Tensor& indptr, const torch::Tensor& indices,
34
35
36
37
38
39
40
41
42
43
    const torch::optional<torch::Tensor>& node_type_offset,
    const torch::optional<torch::Tensor>& type_per_edge) {
  if (node_type_offset.has_value()) {
    auto& offset = node_type_offset.value();
    TORCH_CHECK(offset.dim() == 1);
  }
  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));
  }
44
45

  return c10::make_intrusive<CSCSamplingGraph>(
46
      indptr, indices, node_type_offset, type_per_edge);
47
48
}

49
void CSCSamplingGraph::Load(torch::serialize::InputArchive& archive) {
50
51
  const int64_t magic_num =
      read_from_archive(archive, "CSCSamplingGraph/magic_num").toInt();
52
53
54
  TORCH_CHECK(
      magic_num == kCSCSamplingGraphSerializeMagic,
      "Magic numbers mismatch when loading CSCSamplingGraph.");
55
56
  indptr_ = read_from_archive(archive, "CSCSamplingGraph/indptr").toTensor();
  indices_ = read_from_archive(archive, "CSCSamplingGraph/indices").toTensor();
57
58
59
60
61
62
63
64
65
66
67
  if (read_from_archive(archive, "CSCSamplingGraph/has_node_type_offset")
          .toBool()) {
    node_type_offset_ =
        read_from_archive(archive, "CSCSamplingGraph/node_type_offset")
            .toTensor();
  }
  if (read_from_archive(archive, "CSCSamplingGraph/has_type_per_edge")
          .toBool()) {
    type_per_edge_ =
        read_from_archive(archive, "CSCSamplingGraph/type_per_edge").toTensor();
  }
68
69
70
}

void CSCSamplingGraph::Save(torch::serialize::OutputArchive& archive) const {
71
  archive.write("CSCSamplingGraph/magic_num", kCSCSamplingGraphSerializeMagic);
72
73
  archive.write("CSCSamplingGraph/indptr", indptr_);
  archive.write("CSCSamplingGraph/indices", indices_);
74
75
76
77
78
79
80
81
82
83
84
  archive.write(
      "CSCSamplingGraph/has_node_type_offset", node_type_offset_.has_value());
  if (node_type_offset_) {
    archive.write(
        "CSCSamplingGraph/node_type_offset", node_type_offset_.value());
  }
  archive.write(
      "CSCSamplingGraph/has_type_per_edge", type_per_edge_.has_value());
  if (type_per_edge_) {
    archive.write("CSCSamplingGraph/type_per_edge", type_per_edge_.value());
  }
85
86
}

87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
c10::intrusive_ptr<SampledSubgraph> CSCSamplingGraph::InSubgraph(
    const torch::Tensor& nodes) const {
  using namespace torch::indexing;
  const int32_t kDefaultGrainSize = 100;
  torch::Tensor indptr = torch::zeros_like(indptr_);
  const size_t num_seeds = nodes.size(0);
  std::vector<torch::Tensor> indices_arr(num_seeds);
  std::vector<torch::Tensor> edge_ids_arr(num_seeds);
  std::vector<torch::Tensor> type_per_edge_arr(num_seeds);
  torch::parallel_for(
      0, num_seeds, kDefaultGrainSize, [&](size_t start, size_t end) {
        for (size_t i = start; i < end; ++i) {
          const int64_t node_id = nodes[i].item<int64_t>();
          const int64_t start_idx = indptr_[node_id].item<int64_t>();
          const int64_t end_idx = indptr_[node_id + 1].item<int64_t>();
          indptr[node_id + 1] = end_idx - start_idx;
          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);
          }
        }
      });

  const auto& nonzero_idx = torch::nonzero(indptr).reshape(-1);
  torch::Tensor compact_indptr =
      torch::zeros({nonzero_idx.size(0) + 1}, indptr_.dtype());
  compact_indptr.index_put_({Slice(1, None)}, indptr.index({nonzero_idx}));
  return c10::make_intrusive<SampledSubgraph>(
117
      compact_indptr.cumsum(0), torch::cat(indices_arr), nonzero_idx - 1,
118
119
120
121
122
123
      torch::arange(0, NumNodes()), torch::cat(edge_ids_arr),
      type_per_edge_
          ? torch::optional<torch::Tensor>{torch::cat(type_per_edge_arr)}
          : torch::nullopt);
}

124
c10::intrusive_ptr<SampledSubgraph> CSCSamplingGraph::SampleNeighbors(
125
    const torch::Tensor& nodes, int64_t fanout, bool replace) const {
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
  const int64_t num_nodes = nodes.size(0);

  std::vector<torch::Tensor> picked_neighbors_per_node(num_nodes);
  torch::Tensor num_picked_neighbors_per_node =
      torch::zeros({num_nodes + 1}, indptr_.options());

  torch::parallel_for(0, num_nodes, 32, [&](size_t b, size_t e) {
    for (size_t i = b; i < e; ++i) {
      const auto nid = nodes[i].item<int64_t>();
      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_[nid].item<int64_t>();
      const auto num_neighbors = indptr_[nid + 1].item<int64_t>() - offset;

      if (num_neighbors == 0) {
        // Initialization is performed here because all tensors will be
        // concatenated in the master thread, and having an undefined tensor
        // during concatenation can result in a crash.
        picked_neighbors_per_node[i] = torch::tensor({}, indptr_.options());
        continue;
      }

      picked_neighbors_per_node[i] =
151
          Pick(offset, num_neighbors, fanout, replace, indptr_.options());
152
153
      num_picked_neighbors_per_node[i + 1] =
          picked_neighbors_per_node[i].size(0);
154
155
156
157
158
159
160
161
162
163
    }
  });  // End of the thread.

  torch::Tensor subgraph_indptr =
      torch::cumsum(num_picked_neighbors_per_node, 0);

  torch::Tensor picked_eids = torch::cat(picked_neighbors_per_node);
  torch::Tensor subgraph_indices =
      torch::index_select(indices_, 0, picked_eids);

164
  return c10::make_intrusive<SampledSubgraph>(
165
166
      subgraph_indptr, subgraph_indices, nodes, torch::nullopt, torch::nullopt,
      torch::nullopt);
167
168
}

169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
c10::intrusive_ptr<CSCSamplingGraph>
CSCSamplingGraph::BuildGraphFromSharedMemoryTensors(
    std::tuple<
        SharedMemoryPtr, SharedMemoryPtr,
        std::vector<torch::optional<torch::Tensor>>>&& shared_memory_tensors) {
  auto& optional_tensors = std::get<2>(shared_memory_tensors);
  auto graph = c10::make_intrusive<CSCSamplingGraph>(
      optional_tensors[0].value(), optional_tensors[1].value(),
      optional_tensors[2], optional_tensors[3]);
  graph->tensor_meta_shm_ = std::move(std::get<0>(shared_memory_tensors));
  graph->tensor_data_shm_ = std::move(std::get<1>(shared_memory_tensors));
  return graph;
}

c10::intrusive_ptr<CSCSamplingGraph> CSCSamplingGraph::CopyToSharedMemory(
    const std::string& shared_memory_name) {
  auto optional_tensors = std::vector<torch::optional<torch::Tensor>>{
      indptr_, indices_, node_type_offset_, type_per_edge_};
  auto shared_memory_tensors = CopyTensorsToSharedMemory(
      shared_memory_name, optional_tensors, SERIALIZED_METAINFO_SIZE_MAX);
  return BuildGraphFromSharedMemoryTensors(std::move(shared_memory_tensors));
}

c10::intrusive_ptr<CSCSamplingGraph> CSCSamplingGraph::LoadFromSharedMemory(
    const std::string& shared_memory_name) {
  auto shared_memory_tensors = LoadTensorsFromSharedMemory(
      shared_memory_name, SERIALIZED_METAINFO_SIZE_MAX);
  return BuildGraphFromSharedMemoryTensors(std::move(shared_memory_tensors));
}

199
torch::Tensor Pick(
200
    int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
201
202
    const torch::TensorOptions& options) {
  torch::Tensor picked_neighbors;
203
  if ((fanout == -1) || (num_neighbors <= fanout && !replace)) {
204
205
    picked_neighbors = torch::arange(offset, offset + num_neighbors, options);
  } else {
206
207
208
209
210
211
212
    if (replace) {
      picked_neighbors =
          torch::randint(offset, offset + num_neighbors, {fanout}, options);
    } else {
      picked_neighbors = torch::randperm(num_neighbors, options) + offset;
      picked_neighbors = picked_neighbors.slice(0, 0, fanout);
    }
213
214
215
216
  }
  return picked_neighbors;
}

217
218
}  // namespace sampling
}  // namespace graphbolt