Skip to content

Commit 8f64f0b

Browse files
keichikeichi
authored
Merge pull request #81 from keichi/cpu-sort
Add CPU sort functions and optimize GPU radix sort with CUB
2 parents cc27c3c + 1072102 commit 8f64f0b

4 files changed

Lines changed: 376 additions & 44 deletions

File tree

‎src/ccm.cpp‎

Lines changed: 166 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,16 @@
1+
#include <algorithm>
2+
#include <numeric>
13
#include <random>
4+
#include <stdexcept>
25
#include <unordered_set>
36

47
#include <Kokkos_Bitset.hpp>
58
#include <Kokkos_Core.hpp>
69
#include <Kokkos_NestedSort.hpp>
10+
11+
#ifdef KOKKOS_ENABLE_CUDA
12+
#include <cub/cub.cuh>
13+
#endif
714
#include <boost/math/distributions/binomial.hpp>
815
#include <pcg_random.hpp>
916

@@ -16,8 +23,8 @@
1623
namespace edm
1724
{
1825

19-
void full_sort(TmpDistances distances, TmpIndices indices, int n_lib,
20-
int n_pred, int n_partial, int Tp)
26+
void full_sort_kokkos(TmpDistances distances, TmpIndices indices, int n_lib,
27+
int n_pred, int n_partial, int Tp)
2128
{
2229
Kokkos::parallel_for(
2330
"EDM::ccm::sort", Kokkos::TeamPolicy<>(n_pred, Kokkos::AUTO),
@@ -37,6 +44,26 @@ void full_sort(TmpDistances distances, TmpIndices indices, int n_lib,
3744
});
3845
}
3946

47+
void full_sort(TmpDistances distances, TmpIndices indices, int n_lib,
48+
int n_pred, int n_partial, int Tp)
49+
{
50+
bool use_scratch =
51+
ScratchDistances1D::shmem_size(distances.extent(1)) +
52+
ScratchIndices1D::shmem_size(indices.extent(1)) <
53+
Kokkos::TeamPolicy<>(n_pred, Kokkos::AUTO).scratch_size_max(0);
54+
55+
if (use_scratch) {
56+
full_sort_with_scratch(distances, indices, n_lib, n_pred, n_partial,
57+
Tp);
58+
} else {
59+
#ifdef KOKKOS_ENABLE_CUDA
60+
full_sort_radix(distances, indices, n_lib, n_pred, n_partial, Tp);
61+
#else
62+
full_sort_kokkos(distances, indices, n_lib, n_pred, n_partial, Tp);
63+
#endif
64+
}
65+
}
66+
4067
void full_sort_with_scratch(TmpDistances distances, TmpIndices indices,
4168
int n_lib, int n_pred, int n_partial, int Tp)
4269
{
@@ -74,6 +101,102 @@ void full_sort_with_scratch(TmpDistances distances, TmpIndices indices,
74101
});
75102
}
76103

104+
void full_sort_cpu(TmpDistances distances, TmpIndices indices, int n_lib,
105+
int n_pred, int n_partial, int Tp)
106+
{
107+
auto distances_h =
108+
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), distances);
109+
auto indices_h = Kokkos::create_mirror_view(Kokkos::HostSpace(), indices);
110+
111+
Kokkos::parallel_for(
112+
"EDM::ccm::full_sort_cpu",
113+
Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n_pred),
114+
[=](int i) {
115+
float *dist_row = &distances_h(i, 0);
116+
int *ind_row = &indices_h(i, 0);
117+
118+
std::iota(ind_row, ind_row + n_lib, 0);
119+
120+
std::sort(ind_row, ind_row + n_lib, [dist_row](int a, int b) {
121+
return dist_row[a] < dist_row[b];
122+
});
123+
124+
std::vector<float> sorted_dist(n_lib);
125+
for (int j = 0; j < n_lib; j++) {
126+
sorted_dist[j] = std::sqrt(dist_row[ind_row[j]]);
127+
}
128+
129+
for (int j = 0; j < n_lib; j++) {
130+
distances_h(i, j) = sorted_dist[j];
131+
indices_h(i, j) = ind_row[j] + n_partial + Tp;
132+
}
133+
});
134+
135+
Kokkos::deep_copy(distances, distances_h);
136+
Kokkos::deep_copy(indices, indices_h);
137+
}
138+
139+
void full_sort_radix(TmpDistances distances, TmpIndices indices, int n_lib,
140+
int n_pred, int n_partial, int Tp)
141+
{
142+
#ifdef KOKKOS_ENABLE_CUDA
143+
// Initialize: apply sqrt and set indices
144+
Kokkos::parallel_for(
145+
"EDM::ccm::radix_init", Kokkos::TeamPolicy<>(n_pred, Kokkos::AUTO),
146+
KOKKOS_LAMBDA(const Kokkos::TeamPolicy<>::member_type &member) {
147+
int row = member.league_rank();
148+
Kokkos::parallel_for(
149+
Kokkos::TeamThreadRange(member, n_lib), [=](int j) {
150+
distances(row, j) = sqrt(distances(row, j));
151+
indices(row, j) = j + n_partial + Tp;
152+
});
153+
});
154+
155+
// Create segment offsets array: [0, n_lib, 2*n_lib, ..., n_pred*n_lib]
156+
Kokkos::View<int *, DevSpace> offsets("offsets", n_pred + 1);
157+
Kokkos::parallel_for(
158+
"EDM::ccm::init_offsets", n_pred + 1,
159+
KOKKOS_LAMBDA(int i) { offsets(i) = i * n_lib; });
160+
161+
// Allocate temporary buffers for double-buffering
162+
TmpDistances dist_temp("dist_temp", n_pred, n_lib);
163+
TmpIndices idx_temp("idx_temp", n_pred, n_lib);
164+
165+
// Set up double buffers for CUB
166+
cub::DoubleBuffer<float> d_keys(distances.data(), dist_temp.data());
167+
cub::DoubleBuffer<int> d_values(indices.data(), idx_temp.data());
168+
169+
// Determine temporary storage requirements
170+
size_t temp_storage_bytes = 0;
171+
cub::DeviceSegmentedRadixSort::SortPairs(
172+
nullptr, temp_storage_bytes, d_keys, d_values, n_pred * n_lib, n_pred,
173+
offsets.data(), offsets.data() + 1);
174+
175+
// Allocate temporary storage
176+
Kokkos::View<char *, DevSpace> temp_storage("temp_storage",
177+
temp_storage_bytes);
178+
179+
// Run sorting operation
180+
cub::DeviceSegmentedRadixSort::SortPairs(
181+
temp_storage.data(), temp_storage_bytes, d_keys, d_values,
182+
n_pred * n_lib, n_pred, offsets.data(), offsets.data() + 1);
183+
184+
// Copy results back if needed (CUB may have swapped buffers)
185+
if (d_keys.Current() != distances.data()) {
186+
Kokkos::deep_copy(distances, dist_temp);
187+
Kokkos::deep_copy(indices, idx_temp);
188+
}
189+
#else
190+
(void)distances;
191+
(void)indices;
192+
(void)n_lib;
193+
(void)n_pred;
194+
(void)n_partial;
195+
(void)Tp;
196+
throw std::runtime_error("full_sort_radix requires CUDA");
197+
#endif
198+
}
199+
77200
const unsigned int RADIX_BITS = 8;
78201
const unsigned int RADIX_SIZE = 1 << RADIX_BITS;
79202
const unsigned int RADIX_MASK = RADIX_SIZE - 1;
@@ -194,6 +317,46 @@ void partial_sort(TmpDistances distances, TmpIndices indices, int k, int n_lib,
194317
});
195318
}
196319

320+
void partial_sort_cpu(TmpDistances distances, TmpIndices indices, int k,
321+
int n_lib, int n_pred, int n_partial, int Tp)
322+
{
323+
auto distances_h =
324+
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), distances);
325+
auto indices_h = Kokkos::create_mirror_view(Kokkos::HostSpace(), indices);
326+
327+
Kokkos::parallel_for(
328+
"EDM::ccm::partial_sort_cpu",
329+
Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace>(0, n_pred),
330+
[=](int i) {
331+
float *dist_row = &distances_h(i, 0);
332+
int *ind_row = &indices_h(i, 0);
333+
334+
std::iota(ind_row, ind_row + n_lib, 0);
335+
336+
std::partial_sort(
337+
ind_row, ind_row + k, ind_row + n_lib,
338+
[dist_row](int a, int b) { return dist_row[a] < dist_row[b]; });
339+
340+
std::vector<float> topk_dist(k);
341+
for (int j = 0; j < k; j++) {
342+
topk_dist[j] = std::sqrt(dist_row[ind_row[j]]);
343+
}
344+
345+
for (int j = 0; j < k; j++) {
346+
distances_h(i, j) = topk_dist[j];
347+
indices_h(i, j) = ind_row[j] + n_partial + Tp;
348+
}
349+
350+
for (int j = k; j < n_lib; j++) {
351+
distances_h(i, j) = FLT_MAX;
352+
indices_h(i, j) = -1;
353+
}
354+
});
355+
356+
Kokkos::deep_copy(distances, distances_h);
357+
Kokkos::deep_copy(indices, indices_h);
358+
}
359+
197360
std::vector<float> ccm(TimeSeries lib, TimeSeries target,
198361
const std::vector<int> &lib_sizes, int sample, int E,
199362
int tau, int Tp, int seed, float accuracy)
@@ -217,15 +380,6 @@ std::vector<float> ccm(TimeSeries lib, TimeSeries target,
217380
// Compute pairwise distance matrix
218381
calc_distances(lib, lib, tmp_dist, n_lib, n_pred, E, tau);
219382

220-
bool use_scratch =
221-
#ifdef KOKKOS_ENABLE_CUDA
222-
ScratchDistances1D::shmem_size(tmp_dist.extent(1)) +
223-
ScratchIndices1D::shmem_size(tmp_ind.extent(1)) <
224-
Kokkos::TeamPolicy<>(n_pred, Kokkos::AUTO).scratch_size_max(0);
225-
#else
226-
false;
227-
#endif
228-
229383
// (Partially) Sort each row of the distance matrix
230384
if (accuracy < 1.0f) {
231385
// Calculate the probability of a row to be sampled
@@ -245,12 +399,7 @@ std::vector<float> ccm(TimeSeries lib, TimeSeries target,
245399

246400
partial_sort(tmp_dist, tmp_ind, k, n_lib, n_pred, n_partial, Tp);
247401
} else {
248-
if (use_scratch) {
249-
full_sort_with_scratch(tmp_dist, tmp_ind, n_lib, n_pred, n_partial,
250-
Tp);
251-
} else {
252-
full_sort(tmp_dist, tmp_ind, n_lib, n_pred, n_partial, Tp);
253-
}
402+
full_sort(tmp_dist, tmp_ind, n_lib, n_pred, n_partial, Tp);
254403
}
255404

256405
pcg32 rng;

‎src/ccm.hpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,9 +11,24 @@ namespace edm
1111
void full_sort(TmpDistances distances, TmpIndices indices, int n_lib,
1212
int n_pred, int n_partial, int Tp);
1313

14+
void full_sort_kokkos(TmpDistances distances, TmpIndices indices, int n_lib,
15+
int n_pred, int n_partial, int Tp);
16+
17+
void full_sort_with_scratch(TmpDistances distances, TmpIndices indices,
18+
int n_lib, int n_pred, int n_partial, int Tp);
19+
20+
void full_sort_cpu(TmpDistances distances, TmpIndices indices, int n_lib,
21+
int n_pred, int n_partial, int Tp);
22+
23+
void full_sort_radix(TmpDistances distances, TmpIndices indices, int n_lib,
24+
int n_pred, int n_partial, int Tp);
25+
1426
void partial_sort(TmpDistances distances, TmpIndices indices, int k, int n_lib,
1527
int n_pred, int n_partial, int Tp);
1628

29+
void partial_sort_cpu(TmpDistances distances, TmpIndices indices, int k,
30+
int n_lib, int n_pred, int n_partial, int Tp);
31+
1732
std::vector<float> ccm(TimeSeries library, TimeSeries target,
1833
const std::vector<int> &lib_sizes, int samples, int E,
1934
int tau, int Tp, int seed, float accuracy = 1.0f);

‎src/partial_sort_bench.cpp‎

Lines changed: 69 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,11 @@ void usage(const std::string &app_name)
3636
" -k, --topk arg Number of top-k neighbors (default: 21)\n"
3737
" -i, --iteration arg Number of iterations (default: 10)\n"
3838
" -f, --full-sort Use full sort instead of partial sort\n"
39+
" -c, --cpu-sort Use CPU sort (std::sort/std::partial_sort)\n"
40+
" -r, --radix-sort Use CUB radix sort (GPU only, with -f)\n"
41+
" -s, --scratch-sort Use scratch memory sort (GPU only, with "
42+
"-f)\n"
43+
" -K, --kokkos-sort Use Kokkos sort (with -f)\n"
3944
" -v, --verbose Enable verbose logging (default: false)\n"
4045
" -h, --help Show this help";
4146

@@ -54,6 +59,10 @@ int main(int argc, char *argv[])
5459
int iterations;
5560
cmdl({"i", "iteration"}, 10) >> iterations;
5661
bool full = cmdl[{"-f", "--full-sort"}];
62+
bool cpu = cmdl[{"-c", "--cpu-sort"}];
63+
bool radix = cmdl[{"-r", "--radix-sort"}];
64+
bool scratch = cmdl[{"-s", "--scratch-sort"}];
65+
bool kokkos_sort = cmdl[{"-K", "--kokkos-sort"}];
5766
bool verbose = cmdl[{"-v", "--verbose"}];
5867

5968
if (cmdl[{"-h", "--help"}]) {
@@ -68,6 +77,12 @@ int main(int argc, char *argv[])
6877
std::cout << "k: " << k << std::endl;
6978
std::cout << "iterations: " << iterations << std::endl;
7079
std::cout << "full_sort: " << (full ? "true" : "false") << std::endl;
80+
std::cout << "cpu_sort: " << (cpu ? "true" : "false") << std::endl;
81+
std::cout << "radix_sort: " << (radix ? "true" : "false") << std::endl;
82+
std::cout << "scratch_sort: " << (scratch ? "true" : "false")
83+
<< std::endl;
84+
std::cout << "kokkos_sort: " << (kokkos_sort ? "true" : "false")
85+
<< std::endl;
7186
}
7287

7388
edm::TmpDistances distances("distances", N, N);
@@ -87,7 +102,12 @@ int main(int argc, char *argv[])
87102
LIKWID_MARKER_THREADINIT;
88103

89104
LIKWID_MARKER_REGISTER("partial_sort");
105+
LIKWID_MARKER_REGISTER("partial_sort_cpu");
90106
LIKWID_MARKER_REGISTER("full_sort");
107+
LIKWID_MARKER_REGISTER("full_sort_cpu");
108+
LIKWID_MARKER_REGISTER("full_sort_radix");
109+
LIKWID_MARKER_REGISTER("full_sort_scratch");
110+
LIKWID_MARKER_REGISTER("full_sort_kokkos");
91111
}
92112

93113
for (auto i = 0; i < iterations; i++) {
@@ -101,15 +121,35 @@ int main(int argc, char *argv[])
101121
#pragma omp parallel
102122
#endif
103123
{
104-
if (full) {
124+
if (full && cpu) {
125+
LIKWID_MARKER_START("full_sort_cpu");
126+
} else if (full && radix) {
127+
LIKWID_MARKER_START("full_sort_radix");
128+
} else if (full && scratch) {
129+
LIKWID_MARKER_START("full_sort_scratch");
130+
} else if (full && kokkos_sort) {
131+
LIKWID_MARKER_START("full_sort_kokkos");
132+
} else if (full) {
105133
LIKWID_MARKER_START("full_sort");
134+
} else if (cpu) {
135+
LIKWID_MARKER_START("partial_sort_cpu");
106136
} else {
107137
LIKWID_MARKER_START("partial_sort");
108138
}
109139
}
110140

111-
if (full) {
141+
if (full && cpu) {
142+
edm::full_sort_cpu(distances, indices, N, N, 0, 0);
143+
} else if (full && radix) {
144+
edm::full_sort_radix(distances, indices, N, N, 0, 0);
145+
} else if (full && scratch) {
146+
edm::full_sort_with_scratch(distances, indices, N, N, 0, 0);
147+
} else if (full && kokkos_sort) {
148+
edm::full_sort_kokkos(distances, indices, N, N, 0, 0);
149+
} else if (full) {
112150
edm::full_sort(distances, indices, N, N, 0, 0);
151+
} else if (cpu) {
152+
edm::partial_sort_cpu(distances, indices, k, N, N, 0, 0);
113153
} else {
114154
edm::partial_sort(distances, indices, k, N, N, 0, 0);
115155
}
@@ -120,8 +160,18 @@ int main(int argc, char *argv[])
120160
#pragma omp parallel
121161
#endif
122162
{
123-
if (full) {
163+
if (full && cpu) {
164+
LIKWID_MARKER_STOP("full_sort_cpu");
165+
} else if (full && radix) {
166+
LIKWID_MARKER_STOP("full_sort_radix");
167+
} else if (full && scratch) {
168+
LIKWID_MARKER_STOP("full_sort_scratch");
169+
} else if (full && kokkos_sort) {
170+
LIKWID_MARKER_STOP("full_sort_kokkos");
171+
} else if (full) {
124172
LIKWID_MARKER_STOP("full_sort");
173+
} else if (cpu) {
174+
LIKWID_MARKER_STOP("partial_sort_cpu");
125175
} else {
126176
LIKWID_MARKER_STOP("partial_sort");
127177
}
@@ -134,9 +184,24 @@ int main(int argc, char *argv[])
134184

135185
std::cout << "elapsed: " << timer.seconds() << " [s]" << std::endl;
136186

137-
if (full) {
187+
if (full && cpu) {
188+
std::cout << "full_sort_cpu " << timer_sort.elapsed() / iterations
189+
<< std::endl;
190+
} else if (full && radix) {
191+
std::cout << "full_sort_radix " << timer_sort.elapsed() / iterations
192+
<< std::endl;
193+
} else if (full && scratch) {
194+
std::cout << "full_sort_scratch " << timer_sort.elapsed() / iterations
195+
<< std::endl;
196+
} else if (full && kokkos_sort) {
197+
std::cout << "full_sort_kokkos " << timer_sort.elapsed() / iterations
198+
<< std::endl;
199+
} else if (full) {
138200
std::cout << "full_sort " << timer_sort.elapsed() / iterations
139201
<< std::endl;
202+
} else if (cpu) {
203+
std::cout << "partial_sort_cpu " << timer_sort.elapsed() / iterations
204+
<< std::endl;
140205
} else {
141206
std::cout << "partial_sort " << timer_sort.elapsed() / iterations
142207
<< std::endl;

0 commit comments

Comments
 (0)