sparse.h 3.6 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
#pragma once

3
#include "extensions.h"
Matthias Fey's avatar
Matthias Fey committed
4
#include "macros.h"
rusty1s's avatar
rusty1s committed
5

Matthias Fey's avatar
Matthias Fey committed
6
7
8
9
10
11
12
namespace sparse {
SPARSE_API int64_t cuda_version() noexcept;

namespace detail {
SPARSE_INLINE_VARIABLE int64_t _cuda_version = cuda_version();
} // namespace detail
} // namespace sparse
Daniel Falbel's avatar
Daniel Falbel committed
13
14
15
16

SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M);
SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E);

Matthias Fey's avatar
Matthias Fey committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
SPARSE_API torch::Tensor
partition(torch::Tensor rowptr, torch::Tensor col,
          torch::optional<torch::Tensor> optional_value, int64_t num_parts,
          bool recursive);

SPARSE_API torch::Tensor
partition2(torch::Tensor rowptr, torch::Tensor col,
           torch::optional<torch::Tensor> optional_value,
           torch::optional<torch::Tensor> optional_node_weight,
           int64_t num_parts, bool recursive);

SPARSE_API torch::Tensor
mt_partition(torch::Tensor rowptr, torch::Tensor col,
             torch::optional<torch::Tensor> optional_value,
             torch::optional<torch::Tensor> optional_node_weight,
             int64_t num_parts, bool recursive, int64_t num_workers);

Daniel Falbel's avatar
Daniel Falbel committed
34
SPARSE_API std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
Matthias Fey's avatar
Matthias Fey committed
35
                                                            torch::Tensor idx);
rusty1s's avatar
rusty1s committed
36

Matthias Fey's avatar
Matthias Fey committed
37
38
SPARSE_API std::tuple<torch::Tensor, torch::Tensor,
                      torch::optional<torch::Tensor>, torch::Tensor>
rusty1s's avatar
rusty1s committed
39
40
relabel_one_hop(torch::Tensor rowptr, torch::Tensor col,
                torch::optional<torch::Tensor> optional_value,
Daniel Falbel's avatar
Daniel Falbel committed
41
                torch::Tensor idx, bool bipartite);
rusty1s's avatar
rusty1s committed
42

Daniel Falbel's avatar
Daniel Falbel committed
43
SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
Matthias Fey's avatar
Matthias Fey committed
44
                                     torch::Tensor start, int64_t walk_length);
rusty1s's avatar
rusty1s committed
45

Daniel Falbel's avatar
Daniel Falbel committed
46
SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
rusty1s's avatar
rusty1s committed
47
48
49
subgraph(torch::Tensor idx, torch::Tensor rowptr, torch::Tensor row,
         torch::Tensor col);

Matthias Fey's avatar
Matthias Fey committed
50
51
SPARSE_API
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
Daniel Falbel's avatar
Daniel Falbel committed
52
53
sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
           int64_t num_neighbors, bool replace);
rusty1s's avatar
rusty1s committed
54

Daniel Falbel's avatar
Daniel Falbel committed
55
SPARSE_API torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
Matthias Fey's avatar
Matthias Fey committed
56
57
58
59
60
                                  torch::Tensor rowptr, torch::Tensor col,
                                  torch::optional<torch::Tensor> opt_value,
                                  torch::optional<torch::Tensor> opt_colptr,
                                  torch::optional<torch::Tensor> opt_csr2csc,
                                  torch::Tensor mat);
rusty1s's avatar
rusty1s committed
61

Daniel Falbel's avatar
Daniel Falbel committed
62
SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
Matthias Fey's avatar
Matthias Fey committed
63
64
65
66
67
68
                                   torch::Tensor rowptr, torch::Tensor col,
                                   torch::optional<torch::Tensor> opt_value,
                                   torch::optional<torch::Tensor> opt_rowcount,
                                   torch::optional<torch::Tensor> opt_colptr,
                                   torch::optional<torch::Tensor> opt_csr2csc,
                                   torch::Tensor mat);
rusty1s's avatar
rusty1s committed
69

Daniel Falbel's avatar
Daniel Falbel committed
70
SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
rusty1s's avatar
rusty1s committed
71
72
73
spmm_min(torch::Tensor rowptr, torch::Tensor col,
         torch::optional<torch::Tensor> opt_value, torch::Tensor mat);

Daniel Falbel's avatar
Daniel Falbel committed
74
SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
rusty1s's avatar
rusty1s committed
75
76
77
spmm_max(torch::Tensor rowptr, torch::Tensor col,
         torch::optional<torch::Tensor> opt_value, torch::Tensor mat);

Matthias Fey's avatar
Matthias Fey committed
78
79
SPARSE_API
std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
rusty1s's avatar
rusty1s committed
80
81
82
spspmm_sum(torch::Tensor rowptrA, torch::Tensor colA,
           torch::optional<torch::Tensor> optional_valueA,
           torch::Tensor rowptrB, torch::Tensor colB,
Matthias Fey's avatar
Matthias Fey committed
83
           torch::optional<torch::Tensor> optional_valueB, int64_t K);