Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix DispatchScan

Browse Source
This commit is contained in:
amdsc21 2023-03-08 20:15:33 +01:00
parent bdcb036592
commit a45005863b
2 changed files with 37 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
src/common/device_helpers.hip.h
View File

@ -1,8 +1,10 @@
#include "hip/hip_runtime.h"
/**
* Copyright 2017-2023 XGBoost contributors
*/
#pragma once
#include "hip/hip_runtime.h"
#include <thrust/binary_search.h> // thrust::upper_bound
#include <thrust/device_malloc_allocator.h>
#include <thrust/device_ptr.h>
@ -22,8 +24,11 @@
#include <algorithm>
#include <chrono>
#include <cstddef> // for size_t
#include <hipcub/hipcub.hpp>
#include <hipcub/util_allocator.hpp>
#include <rocprim/rocprim.hpp>
#include <numeric>
#include <sstream>
#include <string>
@ -1153,6 +1158,7 @@ template <typename InputIteratorT, typename OutputIteratorT, typename ScanOpT,
void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
OffsetT num_items) {
size_t bytes = 0;
#if 0
#if THRUST_MAJOR_VERSION >= 2
safe_cuda((
hipcub::DispatchScan<InputIteratorT, OutputIteratorT, ScanOpT, hipcub::NullType,
@ -1165,7 +1171,14 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
hipcub::NullType(), num_items, nullptr,
false)));
#endif
#endif
safe_cuda((rocprim::inclusive_scan<InputIteratorT, OutputIteratorT, ScanOpT>(nullptr,
bytes, d_in, d_out, num_items, scan_op)));
TemporaryArray<char> storage(bytes);
#if 0
#if THRUST_MAJOR_VERSION >= 2
safe_cuda((
hipcub::DispatchScan<InputIteratorT, OutputIteratorT, ScanOpT, hipcub::NullType,
@ -1179,6 +1192,10 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
d_out, scan_op, hipcub::NullType(),
num_items, nullptr, false)));
#endif
#endif
safe_cuda((rocprim::inclusive_scan<InputIteratorT, OutputIteratorT, ScanOpT>(
storage.data().get(), bytes, d_in, d_out, num_items, scan_op)));
}
template <typename InIt, typename OutIt, typename Predicate>

19
src/data/ellpack_page.cu
View File

@ -13,6 +13,10 @@
#include "gradient_index.h"
#include "xgboost/data.h"
#if defined(__HIP_PLATFORM_AMD__)
#include <rocprim/rocprim.hpp>
#endif
namespace xgboost {
EllpackPage::EllpackPage() : impl_{new EllpackPageImpl()} {}
@ -235,6 +239,8 @@ void CopyDataToEllpack(const AdapterBatchT &batch,
// Go one level down into cub::DeviceScan API to set OffsetT as 64 bit
// So we don't crash on n > 2^31
size_t temp_storage_bytes = 0;
#if defined(__CUDACC__)
using DispatchScan =
cub::DispatchScan<decltype(key_value_index_iter), decltype(out),
TupleScanOp<Tuple>, cub::NullType, int64_t>;
@ -257,6 +263,19 @@ void CopyDataToEllpack(const AdapterBatchT &batch,
key_value_index_iter, out, TupleScanOp<Tuple>(),
cub::NullType(), batch.Size(), nullptr, false);
#endif
#elif defined (__HIP_PLATFORM_AMD__)
rocprim::inclusive_scan<decltype(key_value_index_iter), decltype(out), TupleScanOp<Tuple>>
(nullptr, temp_storage_bytes, key_value_index_iter, out, batch.Size(), TupleScanOp<Tuple>());
dh::TemporaryArray<char> temp_storage(temp_storage_bytes);
rocprim::inclusive_scan<decltype(key_value_index_iter), decltype(out), TupleScanOp<Tuple>>
(temp_storage.data().get(), temp_storage_bytes, key_value_index_iter, out, batch.Size(),
TupleScanOp<Tuple>());
#endif
}
void WriteNullValues(EllpackPageImpl* dst, int device_idx,