5
5
#include < dr/mp.hpp>
6
6
#include < fmt/core.h>
7
7
#include < random>
8
+ #include < ranges>
8
9
9
10
/* Sparse band matrix vector multiplication */
10
11
int main () {
@@ -21,6 +22,12 @@ int main() {
21
22
auto values = new V[nnz];
22
23
std::uniform_real_distribution<double > unif (0 , 1 );
23
24
std::default_random_engine re;
25
+ // x x 0 0 ... 0
26
+ // x x 0 0 ... 0
27
+ // x 0 x 0 ... 0
28
+ // x 0 0 x ... 0
29
+ // .............
30
+ // x ... 0 0 0 x
24
31
for (auto i = 0 ; i <= size; i++) {
25
32
rowInd[i] = i * 2 ; // two elements per row
26
33
}
@@ -40,17 +47,21 @@ int main() {
40
47
dr::mp::csr_eq_distribution<V, I, dr::mp::MpiBackend>>
41
48
matrix (local_data, root);
42
49
50
+ dr::mp::broadcasted_vector<double > broadcasted_b;
43
51
std::vector<double > b;
44
- b.reserve (matrix.shape ().second );
45
- std::vector<double > res (matrix.shape ().first );
46
- for (auto i = 0 ; i < matrix.shape ().second ; i++) {
47
- b.push_back (i);
48
- }
52
+ if (root == dr::mp::rank ()) {
53
+ b.resize (matrix.shape ().second );
54
+ std::iota (b.begin (), b.end (), 1 );
49
55
50
- dr::mp::broadcasted_vector<double > broadcasted_b;
51
56
broadcasted_b.broadcast_data (matrix.shape ().second , 0 , b,
52
57
dr::mp::default_comm ());
58
+ }
59
+ else {
60
+ broadcasted_b.broadcast_data (matrix.shape ().second , 0 , std::ranges::empty_view<V>(),
61
+ dr::mp::default_comm ());
62
+ }
53
63
64
+ std::vector<double > res (matrix.shape ().first );
54
65
gemv (root, res, matrix, broadcasted_b);
55
66
56
67
if (root == dr::mp::rank ()) {
0 commit comments