"tests/vscode:/vscode.git/clone" did not exist on "531e719163d2d7cf0d725bb685c1e8fe3393b9da"
index_select.cc 1.56 KB
Newer Older
1
2
3
4
5
/**
 *  Copyright (c) 2023 by Contributors
 * @file index_select.cc
 * @brief Index select operators.
 */
6
7
#include "./index_select.h"

8
9
#include <graphbolt/cuda_ops.h>
#include <graphbolt/fused_csc_sampling_graph.h>
10
11

#include "./macro.h"
12
#include "./utils.h"
13
14
15
16
17

namespace graphbolt {
namespace ops {

torch::Tensor IndexSelect(torch::Tensor input, torch::Tensor index) {
18
  if (utils::is_on_gpu(index) && input.is_pinned()) {
19
20
21
22
23
24
25
    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
        c10::DeviceType::CUDA, "UVAIndexSelect",
        { return UVAIndexSelectImpl(input, index); });
  }
  return input.index({index.to(torch::kLong)});
}

26
std::tuple<torch::Tensor, torch::Tensor> IndexSelectCSC(
27
28
    torch::Tensor indptr, torch::Tensor indices, torch::Tensor nodes,
    torch::optional<int64_t> output_size) {
29
30
  TORCH_CHECK(
      indices.sizes().size() == 1, "IndexSelectCSC only supports 1d tensors");
31
32
  if (utils::is_on_gpu(nodes) && utils::is_accessible_from_gpu(indptr) &&
      utils::is_accessible_from_gpu(indices)) {
33
    GRAPHBOLT_DISPATCH_CUDA_ONLY_DEVICE(
34
        c10::DeviceType::CUDA, "IndexSelectCSCImpl",
35
        { return IndexSelectCSCImpl(indptr, indices, nodes, output_size); });
36
37
38
39
40
  }
  // @todo: The CPU supports only integer dtypes for indices tensor.
  TORCH_CHECK(
      c10::isIntegralType(indices.scalar_type(), false),
      "IndexSelectCSC is not implemented to slice noninteger types yet.");
41
  sampling::FusedCSCSamplingGraph g(indptr, indices);
42
43
44
45
  const auto res = g.InSubgraph(nodes);
  return std::make_tuple(res->indptr, res->indices);
}

46
47
}  // namespace ops
}  // namespace graphbolt