-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathsparse_bind.cpp
158 lines (136 loc) · 4.9 KB
/
sparse_bind.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
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
117
118
119
120
121
122
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
153
154
155
156
157
// DO NOT EDIT: this file is generated
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/complex.h>
#include "sparse.h"
namespace py = pybind11;
template <class I, class T>
void _csr_matvec_omp(
const I n_row,
const I n_col,
py::array_t<I> & Ap,
py::array_t<I> & Aj,
py::array_t<T> & Ax,
py::array_t<T> & Xx,
py::array_t<T> & Yx
)
{
auto py_Ap = Ap.unchecked();
auto py_Aj = Aj.unchecked();
auto py_Ax = Ax.unchecked();
auto py_Xx = Xx.unchecked();
auto py_Yx = Yx.mutable_unchecked();
const I *_Ap = py_Ap.data();
const I *_Aj = py_Aj.data();
const T *_Ax = py_Ax.data();
const T *_Xx = py_Xx.data();
T *_Yx = py_Yx.mutable_data();
return csr_matvec_omp <I, T>(
n_row,
n_col,
_Ap, Ap.size(),
_Aj, Aj.size(),
_Ax, Ax.size(),
_Xx, Xx.size(),
_Yx, Yx.size()
);
}
template <class I, class T>
void _csr_matvec(
const I n_row,
const I n_col,
py::array_t<I> & Ap,
py::array_t<I> & Aj,
py::array_t<T> & Ax,
py::array_t<T> & Xx,
py::array_t<T> & Yx
)
{
auto py_Ap = Ap.unchecked();
auto py_Aj = Aj.unchecked();
auto py_Ax = Ax.unchecked();
auto py_Xx = Xx.unchecked();
auto py_Yx = Yx.mutable_unchecked();
const I *_Ap = py_Ap.data();
const I *_Aj = py_Aj.data();
const T *_Ax = py_Ax.data();
const T *_Xx = py_Xx.data();
T *_Yx = py_Yx.mutable_data();
return csr_matvec <I, T>(
n_row,
n_col,
_Ap, Ap.size(),
_Aj, Aj.size(),
_Ax, Ax.size(),
_Xx, Xx.size(),
_Yx, Yx.size()
);
}
PYBIND11_MODULE(sparse, m) {
m.doc() = R"pbdoc(
Pybind11 bindings for sparse.h
Methods
-------
csr_matvec_omp
csr_matvec
)pbdoc";
py::options options;
options.disable_function_signatures();
m.def("csr_matvec_omp", &_csr_matvec_omp<int, float>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec_omp", &_csr_matvec_omp<int, double>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec_omp", &_csr_matvec_omp<int, std::complex<float>>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec_omp", &_csr_matvec_omp<int, std::complex<double>>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert(),
R"pbdoc(
Threaded SpMV
y <- A * x
Parameters
----------
n_row, n_col : int
dimensions of the n_row x n_col matrix A
Ap, Aj, Ax : array
CSR pointer, index, and data vectors for matrix A
Xx : array
input vector
Yy : array
output vector (modified in-place)
See Also
--------
csr_matvec
Notes
-----
Requires GCC 4.9 for ivdep
Requires a compiler with OMP)pbdoc");
m.def("csr_matvec", &_csr_matvec<int, float>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec", &_csr_matvec<int, double>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec", &_csr_matvec<int, std::complex<float>>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert());
m.def("csr_matvec", &_csr_matvec<int, std::complex<double>>,
py::arg("n_row"), py::arg("n_col"), py::arg("Ap").noconvert(), py::arg("Aj").noconvert(), py::arg("Ax").noconvert(), py::arg("Xx").noconvert(), py::arg("Yx").noconvert(),
R"pbdoc(
Reference SpMV from scipy
y <- A * x
Parameters
----------
n_row, n_col : int
dimensions of the n_row x n_col matrix A
Ap, Aj, Ax : array
CSR pointer, index, and data vectors for matrix A
Xx : array
input vector
Yy : array
output vector (modified in-place)
See Also
--------
csr_matvec_omp
Notes
-----
Requires GCC 4.9 for ivdep
Requires a compiler with OMP
https://github.com/scipy/scipy/blob/master/scipy/sparse/sparsetools/csr.h#L1122)pbdoc");
}