derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
216 lines
13 KiB
C++
216 lines
13 KiB
C++
/*
|
|
* Copyright 2018 NVIDIA Corporation
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include <thrust/detail/config.h>
|
|
|
|
#include <thrust/detail/integer_math.h>
|
|
#include <thrust/detail/integer_traits.h>
|
|
#include <thrust/detail/preprocessor.h>
|
|
|
|
#ifndef ST_DEVICE_ERROR
|
|
# if defined(__HIP_DEVICE_COMPILE__)
|
|
# include <hip/hip_runtime.h>
|
|
# define THRUST_DEVICE_ERROR(message) (printf("%s\n", (messa
|
|
# elif defined(__CUDA_ARCH__)
|
|
# include <cuda/std/detail/libcxx/include/stdexcept>
|
|
# define THRUST_DEVICE_ERROR(message) ::cuda::std::__throw_runtime_error(message)
|
|
# else
|
|
# include <stdexcept>
|
|
# define THRUST_DEVICE_ERROR(message) throw std::runtime_error(message)
|
|
# endif
|
|
#endif
|
|
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <string>
|
|
#include <type_traits>
|
|
|
|
#if defined(THRUST_FORCE_32_BIT_OFFSET_TYPE) && defined(THRUST_FORCE_64_BIT_OFFSET_TYPE)
|
|
# error "Only THRUST_FORCE_32_BIT_OFFSET_TYPE or THRUST_FORCE_64_BIT_OFFSET_TYPE may be defined!"
|
|
#endif // THRUST_FORCE_32_BIT_OFFSET_TYPE && THRUST_FORCE_64_BIT_OFFSET_TYPE
|
|
|
|
#define _THRUST_INDEX_TYPE_DISPATCH(index_type, status, call, count, arguments) \
|
|
{ \
|
|
auto THRUST_PP_CAT2(count, _fixed) = static_cast<index_type>(count); \
|
|
status = call arguments; \
|
|
}
|
|
|
|
#define _THRUST_INDEX_TYPE_DISPATCH2(index_type, status, call, count1, count2, arguments) \
|
|
{ \
|
|
auto THRUST_PP_CAT2(count1, _fixed) = static_cast<index_type>(count1); \
|
|
auto THRUST_PP_CAT2(count2, _fixed) = static_cast<index_type>(count2); \
|
|
status = call arguments; \
|
|
}
|
|
|
|
#define _THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
if (thrust::detail::is_negative(count)) \
|
|
{ \
|
|
THRUST_DEVICE_ERROR("Invalid input range, passed negative size"); \
|
|
}
|
|
|
|
#define _THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count1) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count2)
|
|
|
|
#if defined(THRUST_FORCE_64_BIT_OFFSET_TYPE)
|
|
//! @brief Always dispatches to 64 bit offset version of an algorithm
|
|
# define THRUST_INDEX_TYPE_DISPATCH(status, call, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int64_t, status, call, count, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH but with two counts
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, call, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int64_t, status, call, count1, count2, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH but with two different call implementations
|
|
# define THRUST_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int64_t, status, call_64, count, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH2 but uses two counts.
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int64_t, status, call_64, count1, count2, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH2 but always dispatching to uint64_t. `count` must not be negative.
|
|
# define THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::uint64_t, status, call_64, count, arguments)
|
|
|
|
#elif defined(THRUST_FORCE_32_BIT_OFFSET_TYPE)
|
|
|
|
//! @brief Ensures that the size of the input does not overflow the offset type
|
|
# define _THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW(index_type, count) \
|
|
if (static_cast<std::uint64_t>(count) \
|
|
> static_cast<std::uint64_t>(thrust::detail::integer_traits<index_type>::const_max)) \
|
|
{ \
|
|
THRUST_DEVICE_ERROR( \
|
|
"Input size exceeds the maximum allowable value for " #index_type \
|
|
". It was used because the macro THRUST_FORCE_32_BIT_OFFSET_TYPE was defined. " \
|
|
"To handle larger input sizes, either remove this macro to dynamically dispatch " \
|
|
"between 32-bit and 64-bit index types, or define THRUST_FORCE_64_BIT_OFFSET_TYPE."); \
|
|
}
|
|
|
|
//! @brief Ensures that the sizes of the inputs do not overflow the offset type, but two counts
|
|
# define _THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW2(index_type, count1, count2) \
|
|
if (static_cast<std::uint64_t>(count1) + static_cast<std::uint64_t>(count2) \
|
|
> static_cast<std::uint64_t>(thrust::detail::integer_traits<index_type>::const_max)) \
|
|
{ \
|
|
THRUST_DEVICE_ERROR( \
|
|
"Input size exceeds the maximum allowable value for " #index_type \
|
|
". It was used because the macro THRUST_FORCE_32_BIT_OFFSET_TYPE was defined. " \
|
|
"To handle larger input sizes, either remove this macro to dynamically dispatch " \
|
|
"between 32-bit and 64-bit index types, or define THRUST_FORCE_64_BIT_OFFSET_TYPE."); \
|
|
}
|
|
|
|
//! @brief Always dispatches to 32 bit offset version of an algorithm but throws if count would overflow
|
|
# define THRUST_INDEX_TYPE_DISPATCH(status, call, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW(std::int32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int32_t, status, call, count, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH but with two counts
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, call, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW2(std::int32_t, count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int32_t, status, call, count1, count2, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH but with two different call implementations
|
|
# define THRUST_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW(std::int32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int32_t, status, call_32, count, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH2 but uses two counts.
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW2(std::int32_t, count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int32_t, status, call_32, count1, count2, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH but always dispatching to uint64_t. `count` must not be negative.
|
|
# define THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_OVERFLOW(std::uint32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::uint32_t, status, call_32, count, arguments)
|
|
|
|
#else // ^^^ THRUST_FORCE_32_BIT_OFFSET_TYPE ^^^ / vvv !THRUST_FORCE_32_BIT_OFFSET_TYPE vvv
|
|
|
|
# define _THRUST_INDEX_TYPE_DISPATCH_SELECT(index_type, count) \
|
|
(static_cast<std::uint64_t>(count) \
|
|
<= static_cast<std::uint64_t>(thrust::detail::integer_traits<index_type>::const_max))
|
|
|
|
# define _THRUST_INDEX_TYPE_DISPATCH_SELECT2(index_type, count1, count2) \
|
|
(static_cast<std::uint64_t>(count1) + static_cast<std::uint64_t>(count2) \
|
|
<= static_cast<std::uint64_t>(thrust::detail::integer_traits<index_type>::const_max))
|
|
|
|
//! Dispatch between 32-bit and 64-bit index_type based versions of the same algorithm implementation. This version
|
|
//! assumes that callables for both branches consist of the same tokens, and is intended to be used with Thrust-style
|
|
//! dispatch interfaces, that always deduce the size type from the arguments.
|
|
# define THRUST_INDEX_TYPE_DISPATCH(status, call, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
if _THRUST_INDEX_TYPE_DISPATCH_SELECT (std::int32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int32_t, status, call, count, arguments) \
|
|
else \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int64_t, status, call, count, arguments)
|
|
|
|
//! Dispatch between 32-bit and 64-bit index_type based versions of the same algorithm implementation. This version
|
|
//! assumes that callables for both branches consist of the same tokens, and is intended to be used with Thrust-style
|
|
//! dispatch interfaces, that always deduce the size type from the arguments.
|
|
//!
|
|
//! This version of the macro supports providing two count variables, which is necessary for set algorithms.
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH(status, call, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
if _THRUST_INDEX_TYPE_DISPATCH_SELECT2 (std::int32_t, count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int32_t, status, call, count1, count2, arguments) \
|
|
else \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int64_t, status, call, count1, count2, arguments)
|
|
|
|
//! Dispatch between 32-bit and 64-bit index_type based versions of the same algorithm implementation. This version
|
|
//! allows using different token sequences for callables in both branches, and is intended to be used with CUB-style
|
|
//! dispatch interfaces, where the "simple" interface always forces the size to be `int` (making it harder for us to
|
|
//! use), but the complex interface that we end up using doesn't actually provide a way to fully deduce the type from
|
|
//! just the call, making the size type appear in the token sequence of the callable.
|
|
//!
|
|
//! See reduce_n_impl to see an example of how this is meant to be used.
|
|
# define THRUST_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
if _THRUST_INDEX_TYPE_DISPATCH_SELECT (std::int32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int32_t, status, call_32, count, arguments) \
|
|
else \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::int64_t, status, call_64, count, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH2 but uses two counts.
|
|
# define THRUST_DOUBLE_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count1, count2, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW2(count1, count2) \
|
|
if _THRUST_INDEX_TYPE_DISPATCH_SELECT2 (std::int32_t, count1, count2) \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int32_t, status, call_32, count1, count2, arguments) \
|
|
else \
|
|
_THRUST_INDEX_TYPE_DISPATCH2(std::int64_t, status, call_64, count1, count2, arguments)
|
|
|
|
//! Like \ref THRUST_INDEX_TYPE_DISPATCH2 but dispatching to uint32_t and uint64_t, respectively, depending on the
|
|
//! `count` argument. `count` must not be negative.
|
|
# define THRUST_UNSIGNED_INDEX_TYPE_DISPATCH2(status, call_32, call_64, count, arguments) \
|
|
_THRUST_INDEX_TYPE_DISPATCH_GUARD_UNDERFLOW(count) \
|
|
if _THRUST_INDEX_TYPE_DISPATCH_SELECT (std::uint32_t, count) \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::uint32_t, status, call_32, count, arguments) \
|
|
else \
|
|
_THRUST_INDEX_TYPE_DISPATCH(std::uint64_t, status, call_64, count, arguments)
|
|
|
|
#endif // !THRUST_FORCE_32_BIT_OFFSET_TYPE
|