Use constexpr in dh::CopyIf. (#6828)

This commit is contained in:
Jiaming Yuan 2021-04-08 07:37:47 +08:00 committed by GitHub
parent b65e3c4444
commit f294c4e023
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -1292,14 +1292,14 @@ void InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op,
template <typename InIt, typename OutIt, typename Predicate>
void CopyIf(InIt in_first, InIt in_second, OutIt out_first, Predicate pred) {
// We loop over batches because thrust::copy_if cant deal with sizes > 2^31
// See thrust issue #1302, #6822
size_t max_copy_size = std::numeric_limits<int>::max() / 2;
// We loop over batches because thrust::copy_if can't deal with sizes > 2^31
// See thrust issue #1302, XGBoost #6822
size_t constexpr kMaxCopySize = std::numeric_limits<int>::max() / 2;
size_t length = std::distance(in_first, in_second);
XGBCachingDeviceAllocator<char> alloc;
for (size_t offset = 0; offset < length; offset += max_copy_size) {
for (size_t offset = 0; offset < length; offset += kMaxCopySize) {
auto begin_input = in_first + offset;
auto end_input = in_first + std::min(offset + max_copy_size, length);
auto end_input = in_first + std::min(offset + kMaxCopySize, length);
out_first = thrust::copy_if(thrust::cuda::par(alloc), begin_input,
end_input, out_first, pred);
}