-
Notifications
You must be signed in to change notification settings - Fork 1
/
WorkDistribution.cpp
95 lines (76 loc) · 2.75 KB
/
WorkDistribution.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
// SPDX-License-Identifier: GPL-3.0-or-later
#include <cassert>
// #define DEBUG_WORK_DISTRIBUTION
#ifdef DEBUG_WORK_DISTRIBUTION
#include <iostream>
#endif
#include "Matrix.h"
#include "WorkDistribution.h"
int WorkDistribution::findChunk(int row) const {
// Find corresponding chunk, linear search should be fine here...
for (int chunk = 0; chunk < numberOfChunks; chunk++) {
if (row < offsets[chunk] + lengths[chunk]) {
return chunk;
}
}
assert(0 && "Should have found a chunk!");
return -1;
}
WorkDistribution *WorkDistribution::calculateByRow(int N, int numberOfChunks) {
std::unique_ptr<int[]> offsets(new int[numberOfChunks]);
std::unique_ptr<int[]> lengths(new int[numberOfChunks]);
int chunkLength = N / numberOfChunks;
int remainder = N - chunkLength * numberOfChunks;
int currentOffset = 0;
for (int i = 0; i < numberOfChunks; i++) {
offsets[i] = currentOffset;
lengths[i] = chunkLength;
if (i < remainder) {
lengths[i]++;
}
#ifdef DEBUG_WORK_DISTRIBUTION
std::cout << "Chunk " << i << " of " << numberOfChunks << " with length "
<< lengths[i] << " from offset " << offsets[i] << std::endl;
#endif
currentOffset += lengths[i];
}
return new WorkDistribution(numberOfChunks, std::move(offsets),
std::move(lengths));
}
WorkDistribution *WorkDistribution::calculateByNz(const MatrixCOO &coo,
int numberOfChunks) {
std::unique_ptr<int[]> offsets(new int[numberOfChunks]);
std::unique_ptr<int[]> lengths(new int[numberOfChunks]);
int chunkLengthInNz = coo.nz / numberOfChunks;
int remainderInNz = coo.nz - chunkLengthInNz * numberOfChunks;
int currentOffset = 0;
int currentNz = 0;
#ifdef DEBUG_WORK_DISTRIBUTION
int lastNz = 0;
#endif
for (int i = 0; i < numberOfChunks; i++) {
offsets[i] = currentOffset;
int chunkEndInNz = (i + 1) * chunkLengthInNz;
if ((i + 1) < remainderInNz) {
chunkEndInNz += i + 1;
} else {
chunkEndInNz += remainderInNz;
}
while (currentNz < chunkEndInNz) {
// We might go behind chunkEndInNz here. This is ok because larger chunks
// in the beginning may hide the start of later chunks!
currentNz += coo.nzPerRow[currentOffset];
currentOffset++;
}
lengths[i] = currentOffset - offsets[i];
#ifdef DEBUG_WORK_DISTRIBUTION
std::cout << "Chunk " << i << " of " << numberOfChunks << " with length "
<< lengths[i] << " (" << (currentNz - lastNz)
<< " nonzeros) from offset " << offsets[i] << std::endl;
lastNz = currentNz;
#endif
}
assert(currentNz == coo.nz);
return new WorkDistribution(numberOfChunks, std::move(offsets),
std::move(lengths));
}