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
158
159
160
161
162
163
164
165
166
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
|
#ifndef USCMATRIX_H
#define USCMATRIX_H
#include <Eigen/Dense>
#include "maybe_omp.h"
#include "util.h"
namespace nplm
{
// is this cheating?
using Eigen::Matrix;
using Eigen::MatrixBase;
using Eigen::Dynamic;
// USC = Uniform Sparse Columns. A USCMatrix is a sparse matrix in which
// each column has exactly k nonzero entries. This allows for a
// simpler and faster compressed representation.
// A USCMatrix can be converted into CSC format fairly easily, by
// adding a third array [0, k, 2k, ..., nk]. However, the indices will
// not be unique.
// We use:
// dense2 = dense1^T * sparse (output bProp, input fProp)
// dense1 = sparse * dense2^T (output computeGradient, input computeGradient)
// where:
// sparse is vocab_size x minibatch_size
// dense1 is vocab_size x embedding_dimension
// dense2 is embedding_dimension x minibatch_size
template <typename Scalar, typename Index=int> // should be EIGEN_DEFAULT_DENSE_INDEX_TYPE but int is smaller
class USCMatrix
{
public:
Matrix<Index,Dynamic,Dynamic> indexes;
Matrix<Scalar,Dynamic,Dynamic> values;
int m_rows;
USCMatrix() : m_rows(0) { }
template <typename Indexes, typename Values>
USCMatrix(Index rows, const MatrixBase<Indexes> &indexes, const MatrixBase<Values> &values)
:
indexes(indexes),
values(values),
m_rows(rows)
{ }
USCMatrix(Index rows, Index nnz, Index cols)
:
indexes(Matrix<Index,Dynamic,Dynamic>(nnz, cols)),
values(Matrix<Scalar,Dynamic,Dynamic>(nnz, cols)),
m_rows(rows)
{
this->indexes.fill(-1);
}
Index rows() const { return m_rows; }
Index cols() const { return indexes.cols(); }
void resize(Index rows, Index nnz, Index cols) {
indexes.resize(nnz, cols);
values.resize(nnz, cols);
m_rows = rows;
}
};
// Dense matrix - sparse matrix product
// a is presumably very wide
template <typename DerivedA, typename ScalarB, typename Index, typename DerivedC>
void uscgemm(user_data_t alpha, const MatrixBase<DerivedA> &a,
const USCMatrix<ScalarB,Index> &b,
const MatrixBase<DerivedC> &c_const)
{
UNCONST(DerivedC, c_const, c);
eigen_assert(a.rows() == c.rows());
eigen_assert(a.cols() == b.rows());
eigen_assert(b.cols() == c.cols());
#pragma omp parallel for
for (Index k=0; k<b.cols(); k++)
for (Index r=0; r<b.indexes.rows(); r++)
{
Index j = b.indexes(r,k);
eigen_assert(j >= 0);
eigen_assert(j < a.cols());
c.col(k) += alpha * a.col(j) * b.values(r,k);
}
}
// sparse matrix - dense matrix product
template <typename ScalarA, typename Index, typename DerivedB, typename DerivedC>
void uscgemm(user_data_t alpha,
const USCMatrix<ScalarA,Index> &a,
const MatrixBase<DerivedB> &b,
const MatrixBase<DerivedC> &c_const)
{
UNCONST(DerivedC, c_const, c);
eigen_assert(a.rows() == c.rows());
eigen_assert(a.cols() == b.rows());
eigen_assert(b.cols() == c.cols());
// This needs to be tuned for each system, unfortunately,
// and seems to vary a lot. A lot.
int i_blocks = omp_get_num_threads()*16;
// Assume only one block in k direction.
// We don't need to explicitly block in the j direction.
#pragma omp parallel for
for (Index ib=0; ib<i_blocks; ib++)
for (Index j=0; j<a.cols(); j++)
for (Index r=0; r<a.indexes.rows(); r++)
{
Index i = a.indexes(r,j);
eigen_assert(i >= 0);
eigen_assert(i < c.rows());
if (i % i_blocks == ib)
c.row(i) += alpha * a.values(r,j) * b.row(j);
}
/*
If c.cols() is really large, then theoretically it seems like we should do:
parallel for blocks in i direction
for blocks in j direction
pack block of a into smaller sparse matrix
for blocks in k direction
for k
for i (sparse)
for j
c(i,k) += a(i,j) * b(j,k)
However, the copying of blocks of a doesn't seem practical for any realistic
sizes of c.cols().
*/
}
// Dense matrix - dense matrix product, but masked by a sparse matrix,
// that is, compute a*b only for those positions in c.indexes, and put
// them in c.values.
// a is presumably a very tall matrix. Row-major order is preferred.
// For b, column-major is preferred.
template <typename DerivedA, typename DerivedB, typename ScalarC, typename Index>
void uscgemm_masked(user_data_t alpha,
const MatrixBase<DerivedA> &a,
const MatrixBase<DerivedB> &b,
USCMatrix<ScalarC,Index> &c)
{
eigen_assert(a.rows() == c.rows());
eigen_assert(a.cols() == b.rows());
eigen_assert(b.cols() == c.cols());
#pragma omp parallel for
for (Index k=0; k<b.cols(); k++)
for (Index r=0; r<c.indexes.rows(); r++)
{
Index i = c.indexes(r, k);
eigen_assert(i >= 0);
eigen_assert(i < a.rows());
c.values(r, k) += alpha * a.row(i) * b.col(k);
}
}
// sparse matrix - dense vector product
template <typename ScalarA, typename Index, typename DerivedB, typename DerivedC>
void uscgemv(user_data_t alpha,
const USCMatrix<ScalarA,Index> &a,
const MatrixBase<DerivedB> &b,
const MatrixBase<DerivedC> &c_const)
{
UNCONST(DerivedC, c_const, c);
eigen_assert(a.rows() == c.rows());
eigen_assert(a.cols() == b.rows());
eigen_assert(b.cols() == 1 && c.cols() == 1);
for (Index j=0; j<a.cols(); j++)
for (Index r=0; r<a.indexes.rows(); r++)
{
Index i = a.indexes(r,j);
eigen_assert(i >= 0);
eigen_assert(i < c.rows());
c(i) += alpha * a.values(r,j) * b(j);
}
}
}
#endif
|
