79 lines
2.4 KiB
Plaintext
79 lines
2.4 KiB
Plaintext
#include "../../../src/common/compressed_iterator.h"
|
|
#include "../../../src/common/device_helpers.cuh"
|
|
#include "gtest/gtest.h"
|
|
#include <algorithm>
|
|
#include <thrust/device_vector.h>
|
|
|
|
namespace xgboost {
|
|
namespace common {
|
|
|
|
struct WriteSymbolFunction {
|
|
CompressedBufferWriter cbw;
|
|
unsigned char* buffer_data_d;
|
|
int* input_data_d;
|
|
WriteSymbolFunction(CompressedBufferWriter cbw, unsigned char* buffer_data_d,
|
|
int* input_data_d)
|
|
: cbw(cbw), buffer_data_d(buffer_data_d), input_data_d(input_data_d) {}
|
|
|
|
__device__ void operator()(size_t i) {
|
|
cbw.AtomicWriteSymbol(buffer_data_d, input_data_d[i], i);
|
|
}
|
|
};
|
|
|
|
struct ReadSymbolFunction {
|
|
CompressedIterator<int> ci;
|
|
int* output_data_d;
|
|
ReadSymbolFunction(CompressedIterator<int> ci, int* output_data_d)
|
|
: ci(ci), output_data_d(output_data_d) {}
|
|
|
|
__device__ void operator()(size_t i) {
|
|
output_data_d[i] = ci[i];
|
|
}
|
|
};
|
|
|
|
TEST(CompressedIterator, TestGPU) {
|
|
#if defined(XGBOOST_USE_CUDA)
|
|
dh::safe_cuda(cudaSetDevice(0));
|
|
#elif defined(XGBOOST_USE_HIP)
|
|
dh::safe_cuda(hipSetDevice(0));
|
|
#endif
|
|
std::vector<int> test_cases = {1, 3, 426, 21, 64, 256, 100000, INT32_MAX};
|
|
int num_elements = 1000;
|
|
int repetitions = 1000;
|
|
srand(9);
|
|
|
|
for (auto alphabet_size : test_cases) {
|
|
for (int i = 0; i < repetitions; i++) {
|
|
std::vector<int> input(num_elements);
|
|
std::generate(input.begin(), input.end(),
|
|
[=]() { return rand() % alphabet_size; });
|
|
CompressedBufferWriter cbw(alphabet_size);
|
|
thrust::device_vector<int> input_d(input);
|
|
|
|
thrust::device_vector<unsigned char> buffer_d(
|
|
CompressedBufferWriter::CalculateBufferSize(input.size(),
|
|
alphabet_size));
|
|
|
|
// write the data on device
|
|
auto input_data_d = input_d.data().get();
|
|
auto buffer_data_d = buffer_d.data().get();
|
|
dh::LaunchN(input_d.size(),
|
|
WriteSymbolFunction(cbw, buffer_data_d, input_data_d));
|
|
|
|
// read the data on device
|
|
CompressedIterator<int> ci(buffer_d.data().get(), alphabet_size);
|
|
thrust::device_vector<int> output_d(input.size());
|
|
auto output_data_d = output_d.data().get();
|
|
dh::LaunchN(output_d.size(), ReadSymbolFunction(ci, output_data_d));
|
|
|
|
std::vector<int> output(output_d.size());
|
|
thrust::copy(output_d.begin(), output_d.end(), output.begin());
|
|
|
|
ASSERT_TRUE(input == output);
|
|
}
|
|
}
|
|
}
|
|
|
|
} // namespace common
|
|
} // namespace xgboost
|