I show here a way to use some new C++ features (coroutines and generators) to improve performance on composing sparse linear algebra algorithms by making it possible to share data between consecutive steps. I do this by replacing the more standard “callback-style” interface for sparse linear algebra algorithms with “reverse communication” which was once a common interface style for Fortran libraries (and still is e.g. for ARPACK). Reverse communication has always been possible in C++ but language features did not really make it a welcoming strategy. The introduction of coroutines (C++20) and later std::generator (C++23) considerably cleans up reverse communication and gives a uniform interface to it for all C++ programmers. I give first some context behind reverse communication, why it existed, and how callbacks later replaced it.
Reverse communication was devised for sparse linear algebra algorithms because of the following contention:
and especially in the early days of Fortran library implementation where it was very common to specify all interfaces in highly concrete terms strictly with only “plain old data” and arrays this posed a significant challenge to support all the possible sparse matrix formats a user might want to use. For example (Dongarra, Eijkhout, & Kahan, 1995) makes the case for reverse communication:
In this report we describe a reverse communication interface for the software implementing the iterative methods described in the Templates book [2]. Reverse communication is a technique by which we can hide the implementation details of various operations from the implementation of the iterative method. This allows us to (a) remove references to the user-prepared array or data structure containing the matrix within the iterative solver, and (b) uniformly take care of the various components that can be changed by a user. These include implementation details of matrix-vector operations, vector operations, stopping tests, and norm computations.”
or in (Lehoucq,Sorensen,Yang)
“ARPACK uses the reverse communication interface to make it easy for a user to work with any sparse matrix format or provide a customized sparse matrix vector multiplication subroutine. When a matrix operation is required it returns control to the calling program with a flag indicating what operation is required. The calling program must then perform the operation and call the ARPACK routine again to continue. The operations are typically matrix-vector products, and solving linear systems.”
the implementation of reverse communication usually worked as follows:
while loop repeatedly calling the main algorithm that uses reverse communicationthe diagram from (Dongarra, Eijkhout, & Kahan, 1995) illustrates this flow:
We will see shortly that this is very similar to “coroutines” which have very recently achieved official support in standard C++, but before discussing that I will show how the world kind of moved on from reverse communication to callbacks
The reason for the somewhat complex calling sequence was not that Fortran was universally unable to pass procedures—older Fortran already had EXTERNAL procedures—but that portable callback interfaces with user state, several operation types, and modern closure-like context were awkward. Reverse communication kept ownership of matrix storage, stopping tests, and auxiliary operations entirely in the caller. As language and library interfaces improved, callback-style designs became more prevalent because they were usually simpler for the caller and could have comparable overhead for an isolated algorithm.
That last remark on performance-parity kind of assumes that all these algorithms run in a vacuum. But actually they often run in parallel with each other and feed into one another as well - a complex sequence of linear algebra algorithms may be called in sequence used as building blocks to build up ever-more sophisticated algorithms. I have found that in this context separately scheduled callback algorithms can introduce an inefficiency because each algorithm does not know what the others are requesting. They may perform separate matrix-vector products, each requiring a scan of the sparse matrix. Returning control through reverse communication back to the caller affords the caller the opportunity to batch together multiple operations requiring the same data - usually their sparse matrix. This results in considerable reuse of matrix values across several right-hand sides, even though the vectors still have to be read and written.
Reverse communication has always been possible in C++, but coroutines make their implementation considerably easier and furthermore expose a universal interface that requires less specialized documentation for an experienced C++ to understand. I will show next how to implement the conjugate gradient algorithm using both callback style and reverse-communication style with coroutines.
Conjugate gradients is a nice algorithm for illustration because it is very easy to implement and the primary computational cost is found in the evaluation of matrix-vector products (and the preconditioner if you are using one, but just for simplicity I do not use them here).
The compact implementations below omit the usual convergence and breakdown checks in order to keep the interface example readable.
A classical implementation of conjugate gradients in C++ will use a callback for the matrix-vector product like below. Look for uses of the input lambda A.
template<typename T,typename Lambda>
std::vector<T> cg_plain(size_t nrows,Lambda A, std::span<const T> b,size_t maxiter){
bool verbose = false;
auto dot = [](std::span<const T> x, std::span<const T> y) -> T{
assert(x.size() == y.size());
T out = 0.0;
for(size_t i=0;i<x.size();i++){
out += x[i]*y[i];
}
return out;
};
std::vector<T> x(nrows,0.0);
std::vector<T> r(nrows,0.0);
std::vector<T> p(nrows,0.0);
std::vector<T> q(nrows,0.0);
/*Calculate initial residual.*/
A(x,r);
for(size_t i=0;i<r.size();i++){
r[i] = b[i] - r[i];
}
std::copy(r.begin(),r.end(),p.begin());
T beta = 0.0;
T rho_last = 0.0;
T rho_current = 0.0;
for(size_t it=0;it<maxiter;it++){
rho_current = dot(r,r);
if(verbose){
std::print("iteration: {}, residual: {}\n",it,std::sqrt(rho_current));
}
if(it>0){
beta = rho_current/rho_last;
for(size_t i=0;i<p.size();i++){
p[i] = r[i] + beta*p[i];
}
}
A(p,q);
T alpha = rho_current / dot(p,q);
/*Update.*/
for(size_t i=0;i<x.size();i++){
x[i] = x[i] + alpha*p[i];
r[i] = r[i] - alpha*q[i];
}
rho_last = rho_current;
}
return x;
}
In this way we have an independently testable algorithm such that the user can use any sparse matrix format they may want so long as they can expose it to the algorithm in the form of a callback. The calling sequence to use this is very simple I show a snippet for this below because we will later contrast it with reverse communication:
auto x = cg_plain(nrows,
[&A](std::span<const double> in, std::span<double> out){
/*A is a sparse matrix using CSR format, but algorithm only need `matvec`.*/
A.matvec(in,out);
},std::span<const double>(b),maxiter);
To implement reverse communication I use coroutines and generators. Instead of status flags as was used in Fortran reverse communication I use a std::variant. This requires a little extra setup but it’s surprisingly similar to the callback-style implementation thanks to elegant co_yield now supported by C++:
/* These are deferred operations that will be co_yielded from our reverse communication CG.*/
/* Deferred operation for a matrix-vector product.*/
template<typename T>
struct deferred_matvec{
std::span<const T> in;
std::span<T> out;
};
/*An indicator that the algorithm is complete, storing its result.*/
template<typename T>
struct done{
T result;
};
/* Note how the signature no longer includes _any_ reference to an input matrix, not even as a callback - it's now the users' responsibility! .*/
template<typename T>
std::generator< std::variant<deferred_matvec<T>,done<std::vector<T>>> > cg_revcomm(size_t nrows,std::span<const T> b,size_t maxiter){
bool verbose = false;
auto dot = [](std::span<const T> x, std::span<const T> y) -> T{
assert(x.size() == y.size());
T out = 0.0;
for(size_t i=0;i<x.size();i++){
out += x[i]*y[i];
}
return out;
};
std::vector<T> x(nrows,0.0);
std::vector<T> r(nrows,0.0);
std::vector<T> p(nrows,0.0);
std::vector<T> q(nrows,0.0);
/*Calculate initial residual.*/
co_yield deferred_matvec<T> { .in = x, .out = r };
for(size_t i=0;i<r.size();i++){
r[i] = b[i] - r[i];
}
std::copy(r.begin(),r.end(),p.begin());
T beta = 0.0;
T rho_last = 0.0;
T rho_current = 0.0;
for(size_t it=0;it<maxiter;it++){
rho_current = dot(r,r);
if(verbose){
std::print("iteration: {}, residual: {}\n",it,std::sqrt(rho_current));
}
if(it>0){
beta = rho_current/rho_last;
for(size_t i=0;i<p.size();i++){
p[i] = r[i] + beta*p[i];
}
}
co_yield deferred_matvec<T> { .in = p, .out = q };
T alpha = rho_current / dot(p,q);
/*Update.*/
for(size_t i=0;i<x.size();i++){
x[i] = x[i] + alpha*p[i];
r[i] = r[i] - alpha*q[i];
}
rho_last = rho_current;
}
co_yield done<std::vector<T>> { .result = x };
}
What has effectively happened here is that wherever a callback would have been called we just replaced it with a deferred evaluation and co_yield it back to the caller. Each co_yield returns a variant which either is deferred_matvec or done. The new calling sequence is more complicated just like in the Fortran case:
//Holds the output
std::vector<double> x;
//Elegant use of `std::generator` allows us to iterate it to completion as a range loop
for(auto deferred_op : cg_revcomm(nrows, std::span<const double>(b), maxiter)){
//Each `co_yield` from cg_revcomm returns a variant which is either a matvec or a done,
//thus we need `std::visit` to unpack the two possibilities
std::visit([&](auto&& arg){
using DT = std::decay_t<decltype(arg)>;
/*CG has requested ("reverse-communicated") a sparse matrix-vector product
* so now we evaluate that here.*/
if constexpr (std::is_same_v<DT, deferred_matvec<double>>){
auto [in, out] = arg;
A.matvec(in,out);
}
/*Algorithm has produced an answer which we capture.*/
else{
auto [tmp_x] = arg;
x = tmp_x;
}
}, deferred_op);
}
But this establishes that it is indeed possible to use coroutines to implement reverse communication, at the cost of a more complex calling sequence. I will show next however how we can use this to our advantage for performance.
I give a small contrived example (for simplicity) to illustrate that reverse communication enables potential data sharing between algorithms, saving on memory bandwidth and speeding up the algorithm:
Suppose we don’t just want to solve a sparse linear system but we also want an estimate of its dominant eigenvalue. Power iteration also only requires matrix-vector products. For the symmetric positive definite test matrix used here, the dominant eigenvalue is the spectral radius; for a general matrix, power iteration needs additional assumptions and may fail to converge to the spectral radius. We could run conjugate gradients in parallel with the power iteration but this effectively doubles up memory reads over the input matrix A because each evaluation of matvec requires a full scan of A’s datastructure. If we stick to the callback-style approach we’re kind of stuck with this inefficiency unless we custom-write a combined algorithm that can efficiently share A’s data. I think it should be obvious why we should prefer not to do this every time we want to compose linear algebra algorithms, so now I show how to avoid that with reverse communication.
I will omit power iteration implementation as it has almost exactly the same structure as conjugate-gradients above. The reverse-communication signature however looks as follows:
template<typename T>
std::generator< std::variant<deferred_matvec<T>, done<std::pair<T,std::vector<T>>> > >
power_revcomm(size_t nrows,std::span<const T> x,size_t maxiter);
Now we can co-iterate both of these algorithms in tandem and use a batched matrix-vector product whenever both algorithms call for it, allowing the two products to share one traversal of A:
/**
* The idea here is instead of simply sequentially iterating both coroutines
* to completion we can co-iterate them and batch up any matching
* matrix-vector products involving `A`, greatly reducing
* the redundant memory reads while still keeping
* the algorithms logically separate.
*/
auto cg_eval = cg_revcomm(nrows, std::span<const double>(b),maxiter);
auto power_eval = power_revcomm(nrows, std::span<const double>(b),maxiter);
auto cg_it = cg_eval.begin();
auto power_it = power_eval.begin();
std::vector<double> x;
double lambda;
std::vector<double> y;
while(cg_it != cg_eval.end() && power_it != power_eval.end()){
auto cg_op = *cg_it;
auto power_op = *power_it;
bool cg_is_matvec = std::visit([&](auto&& arg) -> bool{
using DT = std::decay_t<decltype(arg)>;
if constexpr (std::is_same_v<DT, deferred_matvec<double>>){
return true;
}
return false;
},cg_op);
bool power_is_matvec = std::visit([&](auto&& arg) -> bool{
using DT = std::decay_t<decltype(arg)>;
if constexpr (std::is_same_v<DT, deferred_matvec<double>>){
return true;
}
return false;
},power_op);
/*We have four posibilities:
* Either both CG and power iteration have requested a
* matrix-vector product,one of them is done,
* or both of them are done.
*
* But in all other cases we just handle them with the
* same visitor pattern as before.
* */
if(cg_is_matvec && power_is_matvec){
/*In this case since both algorithms want matrix-vector product
* with the same matrix `A` we batch them up.*/
auto [in0,out0] = std::get<deferred_matvec<double>>(cg_op);
auto [in1,out1] = std::get<deferred_matvec<double>>(power_op);
A.matvec_batch2(in0,in1,out0,out1);
}
else{
std::visit([&](auto&& arg){
using DT = std::decay_t<decltype(arg)>;
/*Power iteration has requested ("reverse-communicated") a sparse matrix-vector product
* so now we evaluate that here.*/
if constexpr (std::is_same_v<DT, deferred_matvec<double>>){
auto [in, out] = arg;
A.matvec(in,out);
}
/*Algorithm has produced an answer which we capture.*/
else{
auto [result] = arg;
auto [tmp_lambda,tmp_y] = result;
y = tmp_y;
lambda = tmp_lambda;
}
}, power_op);
std::visit([&](auto&& arg){
using DT = std::decay_t<decltype(arg)>;
/*CG has requested ("reverse-communicated") a sparse matrix-vector product
* so now we evaluate that here.*/
if constexpr (std::is_same_v<DT, deferred_matvec<double>>){
auto [in, out] = arg;
A.matvec(in,out);
}
/*Algorithm has produced an answer which we capture.*/
else{
auto [tmp_x] = arg;
x = tmp_x;
}
}, cg_op);
}
cg_it++;
power_it++;
}
//Must clean-up either CG or power may not have finished but I leave as
//exercise to the interested reader
That is a lot of code dealing with variants which I think can be cleaned up, but as a proof of concept it shows the scheduling idea. The listing also leaves coroutine cleanup and lifetime handling as an exercise. The key lines are the following where we call a batched matvec
/*We have four posibilities:
* Either both CG and power iteration have requested a
* matrix-vector product,one of them is done,
* or both of them are done.
*
* But in all other cases we just handle them with the
* same visitor pattern as before.
* */
if(cg_is_matvec && power_is_matvec){
/*In this case since both algorithms want matrix-vector product
* with the same matrix `A` we batch them up.*/
auto [in0,out0] = std::get<deferred_matvec<double>>(cg_op);
auto [in1,out1] = std::get<deferred_matvec<double>>(power_op);
A.matvec_batch2(in0,in1,out0,out1);
}
My performance results have been pretty interesting, which I discuss next.
I set up an example where I randomly generate a symmetric positive definite sparse matrix using the CSR format with 200000 rows, 30 nonzeros per row, and a graph-bandwidth of about 30. I then tested the following implementation strategies:
In general I expected that the batching versions should outperform non-batching versions, and that is indeed the case:
| Algorithm Type | Time (milliseconds) |
|---|---|
| Callback Sequential | 3766 |
| Reverse Communication Sequential | 3763 |
| Batched Reverse Communication | 3026 |
| Callback Multithreaded | 1837 |
| Batched Reverse Communication Multithreaded | 1618 |
With -O3, the sequential callback and reverse-communication measurements are nearly identical, so coroutine overhead is negligible here. The batched versions are faster.
This was an experiment to see if C++ coroutines might offer a new approach to interfaces for linear algebra algorithms. The idea here is to allow for the possibility that multiple algorithms could share data while still implementing those algorithms in logically self-contained functions. C++ coroutines have very cleanly enabled this pattern. The batching concept also shows a performance benefit, so there is real value in exposing requests to a scheduler. Reverse communication is not the only way to do this, but it makes cross-algorithm opportunities visible without fusing the implementations themselves.
The main drawback here was the added code complexity on the callers. So I think work is needed to simplify this new pattern. Perhaps linear algebra interfaces opting to use reverse communication should have a level-0 interface that exposes reverse communication and a level-1 interface that uses callbacks (but just uses the reverse communication implementation under-the-hood). Most of the time, especially in algorithm prototyping phases, the user will prefer the simpler callback interface, but in later optimizing stages they will want to seek opportunities for data sharing and reverse communication could offer a real possibility there.
You can find the (proof of concept) code for my results here