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
|
#ifndef __COO_H__
#define __COO_H__
#include <algorithm>
/*
* Compute B = A for COO matrix A, CSR matrix B
*
*
* Input Arguments:
* I n_row - number of rows in A
* I n_col - number of columns in A
* I nnz - number of nonzeros in A
* I Ai[nnz(A)] - row indices
* I Aj[nnz(A)] - column indices
* T Ax[nnz(A)] - nonzeros
* Output Arguments:
* I Bp - row pointer
* I Bj - column indices
* T Bx - nonzeros
*
* Note:
* Output arrays Bp, Bj, and Bx must be preallocated
*
* Note:
* Input: row and column indices *are not* assumed to be ordered
*
* Note: duplicate entries are carried over to the CSR represention
*
* Complexity: Linear. Specifically O(nnz(A) + max(n_row,n_col))
*
*/
template <class I, class T>
void coo_tocsr(const I n_row,
const I n_col,
const I nnz,
const I Ai[],
const I Aj[],
const T Ax[],
I Bp[],
I Bj[],
T Bx[])
{
//compute number of non-zero entries per row of A
std::fill(Bp, Bp + n_row, 0);
for (I n = 0; n < nnz; n++){
Bp[Ai[n]]++;
}
//cumsum the nnz per row to get Bp[]
for(I i = 0, cumsum = 0; i < n_row; i++){
I temp = Bp[i];
Bp[i] = cumsum;
cumsum += temp;
}
Bp[n_row] = nnz;
//write Aj,Ax into Bj,Bx
for(I n = 0; n < nnz; n++){
I row = Ai[n];
I dest = Bp[row];
Bj[dest] = Aj[n];
Bx[dest] = Ax[n];
Bp[row]++;
}
for(I i = 0, last = 0; i <= n_row; i++){
I temp = Bp[i];
Bp[i] = last;
last = temp;
}
//now Bp,Bj,Bx form a CSR representation (with possible duplicates)
}
/*
* Compute B += A for COO matrix A, dense matrix B
*
* Input Arguments:
* I n_row - number of rows in A
* I n_col - number of columns in A
* npy_int64 nnz - number of nonzeros in A
* I Ai[nnz(A)] - row indices
* I Aj[nnz(A)] - column indices
* T Ax[nnz(A)] - nonzeros
* T Bx[n_row*n_col] - dense matrix
*
*/
template <class I, class T>
void coo_todense(const I n_row,
const I n_col,
const npy_int64 nnz,
const I Ai[],
const I Aj[],
const T Ax[],
T Bx[],
const int fortran)
{
if (!fortran) {
for(npy_int64 n = 0; n < nnz; n++){
Bx[ (npy_intp)n_col * Ai[n] + Aj[n] ] += Ax[n];
}
}
else {
for(npy_int64 n = 0; n < nnz; n++){
Bx[ (npy_intp)n_row * Aj[n] + Ai[n] ] += Ax[n];
}
}
}
/*
* Compute Y += A*X for COO matrix A and dense vectors X,Y
*
*
* Input Arguments:
* npy_int64 nnz - number of nonzeros in A
* I Ai[nnz] - row indices
* I Aj[nnz] - column indices
* T Ax[nnz] - nonzero values
* T Xx[n_col] - input vector
*
* Output Arguments:
* T Yx[n_row] - output vector
*
* Notes:
* Output array Yx must be preallocated
*
* Complexity: Linear. Specifically O(nnz(A))
*
*/
template <class I, class T>
void coo_matvec(const npy_int64 nnz,
const I Ai[],
const I Aj[],
const T Ax[],
const T Xx[],
T Yx[])
{
for(npy_int64 n = 0; n < nnz; n++){
Yx[Ai[n]] += Ax[n] * Xx[Aj[n]];
}
}
#endif
|
