Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
275 changes: 266 additions & 9 deletions src/care/KeyValueSorter_decl.h
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,19 @@

#include "care/scan.h"

// Other library headers
#ifdef CARE_GPUCC
#if defined(__CUDACC__)
#include "cub/cub.cuh"
#undef CUB_NS_POSTFIX
#undef CUB_NS_PREFIX
#endif

#if defined(__HIPCC__)
#include "hipcub/hipcub.hpp"
#endif
#endif

#include <utility> // For std::move

namespace care {
Expand All @@ -42,6 +55,39 @@ class CARE_KEY_VALUE_SORTER_DLL_API KeyValueSorter;
template <typename KeyType, typename ValueType, typename Exec>
using LocalKeyValueSorter = KeyValueSorter<KeyType, ValueType, Exec> ;

template <typename KeyValueType>
inline bool cmpKeys(KeyValueType const & left, KeyValueType const & right);

template <typename KeyValueType>
inline bool cmpKeysThenValues(KeyValueType const & left, KeyValueType const & right);

namespace detail {
template <typename KeyT, typename ValueT>
CARE_INLINE void stableSortKeyValuePairs(host_device_ptr<KeyT> & keys,
host_device_ptr<ValueT> & values,
const size_t len,
const size_t start = 0) {

host_device_ptr<_kv<KeyT,ValueT>> keyValues(len);

CARE_SEQUENTIAL_LOOP(i, 0, (int) len) {
keyValues[i].key = keys[i+start];
keyValues[i].value = values[i+start];
} CARE_SEQUENTIAL_LOOP_END

CHAIDataGetter<_kv<KeyT, ValueT>, RAJA::seq_exec> getter {};
_kv<KeyT, ValueT> * rawData = getter.getRawArrayData(keyValues);
std::stable_sort(rawData, rawData + len, cmpKeys<_kv<KeyT,ValueT>>);

CARE_SEQUENTIAL_LOOP(i, 0, (int) len) {
keys[i+start] = keyValues[i].key;
values[i+start] = keyValues[i].value;
} CARE_SEQUENTIAL_LOOP_END

keyValues.free();
}
} // namespace detail



///////////////////////////////////////////////////////////////////////////
Expand All @@ -60,19 +106,230 @@ using LocalKeyValueSorter = KeyValueSorter<KeyType, ValueType, Exec> ;
///////////////////////////////////////////////////////////////////////////
template <typename Exec, typename KeyT, typename ValueT>
std::enable_if_t<std::is_arithmetic<typename CHAIDataGetter<KeyT, Exec>::raw_type>::value, void>
sortKeyValueArrays(host_device_ptr<KeyT> & keys,
host_device_ptr<ValueT> & values,
const size_t start, const size_t len,
const bool noCopy=false);
inline sortKeyValueArrays(host_device_ptr<KeyT> & keys,
host_device_ptr<ValueT> & values,
const size_t start, const size_t len,
const bool noCopy=false)
{
bool _noCopy ;
if (noCopy && start > 0) {
printf("[CARE] Warning: sortKeyValueArrays. noCopy should not be set if start > 0 (%d)\n", (int)start);
_noCopy = false;
}
else {
_noCopy = noCopy;
}

if constexpr (std::is_same_v<Exec, RAJA::seq_exec>) {
detail::stableSortKeyValuePairs(keys, values, len, start);
}
else {
// TODO openMP parallel implementation
#if defined(__HIPCC__) || (defined(__CUDACC__) && defined(CUB_MAJOR_VERSION) && defined(CUB_MINOR_VERSION) && (CUB_MAJOR_VERSION >= 2 || (CUB_MAJOR_VERSION == 1 && CUB_MINOR_VERSION >= 14)))

// Allocate space for the result
host_device_ptr<KeyT> keyResult{len};
host_device_ptr<ValueT> valueResult{len};

// Get the raw data to pass to cub
CHAIDataGetter<KeyT, Exec> keyGetter {};
CHAIDataGetter<ValueT, Exec> valueGetter {};

auto * rawKeyData = keyGetter.getRawArrayData(keys) + start;
auto * rawValueData = valueGetter.getRawArrayData(values) + start;

auto * rawKeyResult = keyGetter.getRawArrayData(keyResult);
auto * rawValueResult = valueGetter.getRawArrayData(valueResult);

// Get the temp storage length
char * d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;

// When called with a nullptr for temp storage, this returns how much
// temp storage should be allocated.
if (len > 0) {
#if defined(__CUDACC__)
cub::DeviceRadixSort::SortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData, rawKeyResult,
rawValueData, rawValueResult,
len);
#elif defined(__HIPCC__)
hipcub::DeviceRadixSort::SortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData, rawKeyResult,
rawValueData, rawValueResult,
len);
#endif
}

// Allocate the temp storage and get raw data to pass to cub
host_device_ptr<char> tmpManaged {temp_storage_bytes};

CHAIDataGetter<char, Exec> charGetter {};
d_temp_storage = charGetter.getRawArrayData(tmpManaged);

// Now sort
if (len > 0) {
#if defined(CHAI_THIN_GPU_ALLOCATE)
chai::ArrayManager::getInstance()->setExecutionSpace(chai::GPU);
#endif

#if defined(__CUDACC__)
cub::DeviceRadixSort::SortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData, rawKeyResult,
rawValueData, rawValueResult,
len);
#elif defined(__HIPCC__)
hipcub::DeviceRadixSort::SortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData, rawKeyResult,
rawValueData, rawValueResult,
len);
#endif

#if defined(CHAI_THIN_GPU_ALLOCATE)
chai::ArrayManager::getInstance()->setExecutionSpace(chai::NONE);
#endif

tmpManaged.free();
}

// Get the result
if (_noCopy) {
if (len > 0) {
keys.free();
values.free();
}

keys = keyResult;
values = valueResult;
}
else {
CARE_STREAM_LOOP(i, 0, len) {
keys[i+start] = keyResult[i];
values[i+start] = valueResult[i];
} CARE_STREAM_LOOP_END

if (len > 0) {
keyResult.free();
valueResult.free();
}
}

#else // defined(CARE_GPUCC)
detail::stableSortKeyValuePairs(keys, values, len, start);
#endif // defined(CARE_GPUCC)
}

}

template <typename Exec, typename KeyT, typename ValueT>
std::enable_if_t<!std::is_arithmetic<typename CHAIDataGetter<KeyT, Exec>::raw_type>::value, void>
sortKeyValueArrays(host_device_ptr<KeyT> & keys,
host_device_ptr<ValueT> & values,
const size_t start, const size_t len,
const bool noCopy=false);
inline sortKeyValueArrays(host_device_ptr<KeyT> & keys,
host_device_ptr<ValueT> & values,
const size_t start, const size_t len,
const bool noCopy=false)
{
bool _noCopy ;
if (noCopy && start > 0) {
printf("[CARE] Warning: sortKeyValueArrays. noCopy should not be set if start > 0 (%d)\n", (int)start);
_noCopy = false;
}
else {
_noCopy = noCopy;
}

if constexpr (std::is_same_v<Exec, RAJA::seq_exec>) {
detail::stableSortKeyValuePairs(keys, values, len, start);
}
else {
// TODO openMP parallel implementation
#if defined(__HIPCC__) || (defined(__CUDACC__) && defined(CUB_MAJOR_VERSION) && defined(CUB_MINOR_VERSION) && (CUB_MAJOR_VERSION >= 2 || (CUB_MAJOR_VERSION == 1 && CUB_MINOR_VERSION >= 14)))

// Allocate space for the result
host_device_ptr<KeyT> keyResult{len};
host_device_ptr<ValueT> valueResult{len};

// Get the raw data to pass to cub
CHAIDataGetter<KeyT, Exec> keyGetter {};
CHAIDataGetter<ValueT, Exec> valueGetter {};

auto * rawKeyData = keyGetter.getRawArrayData(keys) + start;
auto * rawValueData = valueGetter.getRawArrayData(values) + start;

auto * rawKeyResult = keyGetter.getRawArrayData(keyResult);
auto * rawValueResult = valueGetter.getRawArrayData(valueResult);

using RawKeyType = std::remove_reference_t<decltype(*rawKeyData)>;

auto custom_comparator = [] CARE_HOST_DEVICE (const RawKeyType& lhs,
const RawKeyType& rhs) {
return lhs < rhs;
};

// Get the temp storage length
char * d_temp_storage = nullptr;
size_t temp_storage_bytes = 0;

// When called with a nullptr for temp storage, this returns how much
// temp storage should be allocated.
if (len > 0) {
#if defined(__CUDACC__)
cub::DeviceMergeSort::StableSortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData,
rawValueData,
len, custom_comparator);
#elif defined(__HIPCC__)
hipcub::DeviceMergeSort::StableSortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData,
rawValueData,
len, custom_comparator);
#endif
}

// Allocate the temp storage and get raw data to pass to cub
host_device_ptr<char> tmpManaged {temp_storage_bytes};

CHAIDataGetter<char, Exec> charGetter {};
d_temp_storage = charGetter.getRawArrayData(tmpManaged);

// Now sort
if (len > 0) {
#if defined(CHAI_THIN_GPU_ALLOCATE)
chai::ArrayManager::getInstance()->setExecutionSpace(chai::GPU);
#endif

#if defined(__CUDACC__)
cub::DeviceMergeSort::StableSortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData,
rawValueData,
len,
custom_comparator);
#elif defined(__HIPCC__)
hipcub::DeviceMergeSort::StableSortPairs((void *)d_temp_storage, temp_storage_bytes,
rawKeyData,
rawValueData,
len,
custom_comparator);
#endif

#if defined(CHAI_THIN_GPU_ALLOCATE)
chai::ArrayManager::getInstance()->setExecutionSpace(chai::NONE);
#endif

tmpManaged.free();
}

// merge sort did an inplace sort, so the answer is already in keys and Values
if (len > 0) {
keyResult.free();
valueResult.free();
}

#else // defined(CARE_GPUCC)
detail::stableSortKeyValuePairs(keys, values, len, start);
#endif // defined(CARE_GPUCC)
}

}

#if defined(CARE_PARALLEL_DEVICE) || CARE_ENABLE_GPU_SIMULATION_MODE
///////////////////////////////////////////////////////////////////////////
Expand Down Expand Up @@ -617,7 +874,7 @@ class CARE_KEY_VALUE_SORTER_DLL_API KeyValueSorter<KeyType, ValueType, RAJADevic

// Use exclusive scan to compute output positions
host_device_ptr<int> positions(m_len+1);
care::exclusive_scan(RAJADeviceExec{}, isUnique, positions, m_len + 1, 0, false);
exclusive_scan(RAJADeviceExec{}, isUnique, positions, m_len + 1, 0, false);

// Get the total number of unique elements
int newSize = positions.pick(m_len);
Expand Down
Loading
Loading