Skip to content

Commit 2a8a50f

Browse files
sleeepyjacksleeepyjack
committed
Added reduction functors that can be used by static_reduction_map.
1 parent 2998d9f commit 2a8a50f

4 files changed

Lines changed: 502 additions & 0 deletions

File tree

‎include/cuco/detail/reduction_functor_impl.cuh‎

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
* Copyright (c) 2021-2022, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#pragma once
18+
19+
#include <atomic>
20+
#include <cuda/atomic>
21+
#include <type_traits>
22+
23+
namespace cuco {
24+
namespace detail {
25+
26+
/**
27+
* @brief Base class of all reduction functors.
28+
*
29+
* @warning This class should not be used directly.
30+
*
31+
*/
32+
class reduction_functor_base {};
33+
34+
template <typename T, typename Enable = void>
35+
struct reduce_add_impl {
36+
template <cuda::thread_scope Scope>
37+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
38+
return lhs.fetch_add(rhs) + rhs;
39+
}
40+
};
41+
42+
template <typename T, typename Enable = void>
43+
struct reduce_min_impl {
44+
template <cuda::thread_scope Scope>
45+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
46+
return min(lhs.fetch_min(rhs), rhs);
47+
}
48+
};
49+
50+
template <typename T, typename Enable = void>
51+
struct reduce_max_impl {
52+
template <cuda::thread_scope Scope>
53+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
54+
return max(lhs.fetch_max(rhs), rhs);
55+
}
56+
};
57+
58+
template <typename T, typename Enable = void>
59+
struct reduce_count_impl {
60+
template <cuda::thread_scope Scope>
61+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& /* rhs */) const noexcept {
62+
return ++lhs;
63+
}
64+
};
65+
66+
// remove the following WAR once libcu++ extends FP atomics support and fixes signed integer atomics
67+
// https://github.com/NVIDIA/libcudacxx/pull/286
68+
template <typename T>
69+
struct reduce_add_impl<T, typename cuda::std::enable_if<cuda::std::is_floating_point<T>::value>::type> {
70+
template <cuda::thread_scope Scope>
71+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T rhs) const noexcept {
72+
if constexpr (Scope == cuda::thread_scope_system)
73+
return atomicAdd_system(reinterpret_cast<T*>(&lhs), rhs) + rhs;
74+
else if constexpr (Scope == cuda::thread_scope_device)
75+
return atomicAdd(reinterpret_cast<T*>(&lhs), rhs) + rhs;
76+
else
77+
return atomicAdd_block(reinterpret_cast<T*>(&lhs), rhs) + rhs;
78+
}
79+
};
80+
81+
template <typename T>
82+
struct reduce_min_impl<T, typename cuda::std::enable_if<cuda::std::is_integral<T>::value && cuda::std::is_signed<T>::value>::type> {
83+
template <cuda::thread_scope Scope>
84+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
85+
using InternalT = typename cuda::std::conditional<sizeof(T) == 8, long long int, int>::type;
86+
InternalT * ptr = reinterpret_cast<InternalT*>(&lhs);
87+
InternalT value = rhs;
88+
if constexpr (Scope == cuda::thread_scope_system)
89+
return min(atomicMin_system(ptr, value), value);
90+
else if constexpr (Scope == cuda::thread_scope_device)
91+
return min(atomicMin(ptr, value), value);
92+
else
93+
return min(atomicMin_block(ptr, value), value);
94+
}
95+
};
96+
97+
template <typename T>
98+
struct reduce_max_impl<T, typename cuda::std::enable_if<cuda::std::is_integral<T>::value && cuda::std::is_signed<T>::value>::type> {
99+
template <cuda::thread_scope Scope>
100+
__device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
101+
using InternalT = typename cuda::std::conditional<sizeof(T) == 8, long long int, int>::type;
102+
InternalT * ptr = reinterpret_cast<InternalT*>(&lhs);
103+
InternalT value = rhs;
104+
if constexpr (Scope == cuda::thread_scope_system)
105+
return max(atomicMax_system(ptr, value), value);
106+
else if constexpr (Scope == cuda::thread_scope_device)
107+
return max(atomicMax(ptr, value), value);
108+
else
109+
return max(atomicMax_block(ptr, value), value);
110+
}
111+
};
112+
113+
template <typename T>
114+
struct reduce_min_impl<T, typename cuda::std::enable_if<cuda::std::is_floating_point<T>::value>::type> {
115+
__device__ T operator()(T lhs, T rhs) const noexcept {
116+
return min(lhs, rhs);
117+
}
118+
};
119+
120+
template <typename T>
121+
struct reduce_max_impl<T, typename cuda::std::enable_if<cuda::std::is_floating_point<T>::value>::type> {
122+
__device__ T operator()(T lhs, T rhs) const noexcept {
123+
return max(lhs, rhs);
124+
}
125+
};
126+
127+
} // namespace detail
128+
} // namespace cuco

‎include/cuco/reduction_functors.cuh‎

Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* Copyright (c) 2022, NVIDIA CORPORATION.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
#pragma once
18+
19+
#include <atomic>
20+
#include <cuco/detail/reduction_functor_impl.cuh>
21+
22+
#include <cuda/atomic>
23+
#include <limits>
24+
#include <type_traits>
25+
26+
namespace cuco {
27+
28+
/**
29+
* @brief Wrapper for reduction identity value.
30+
*
31+
* @tparam T The underlying value type used for reduction
32+
*/
33+
template <typename T>
34+
class identity_value {
35+
public:
36+
using type = T;
37+
constexpr identity_value(T const& identity) noexcept : identity_(identity) {}
38+
constexpr T value() const noexcept { return identity_; }
39+
private:
40+
T identity_;
41+
};
42+
43+
/**
44+
* @brief Wrapper for a user-defined custom reduction operator.
45+
*
46+
* External synchronization, if required,
47+
* is established via an atomic compare-and-swap loop.
48+
*
49+
* Example:
50+
* \code{.cpp}
51+
* template <typename T>
52+
* struct custom_plus {
53+
* __device__ T operator()(T const& lhs, T const& rhs) const noexcept {
54+
* return lhs + rhs;
55+
* }
56+
* };
57+
*
58+
* template <typename T>
59+
* struct custom_plus_sync {
60+
* template <cuda::thread_scope Scope>
61+
* __device__ T operator()(cuda::atomic<T, Scope>& lhs, T const& rhs) const noexcept {
62+
* return lhs.fetch_add(rhs) + rhs;
63+
* }
64+
* };
65+
*
66+
* int main() {
67+
* cuco::identity_value<int> identity{0}; // define the identity value for the given reduction operation, i.e., op(identity, x) == x
68+
*
69+
* auto f1 = cuco::reduction_functor<custom_plus<int>, int>(identity); // synchronized via CAS-loop
70+
* auto f2 = cuco::reduction_functor<custom_plus_sync<int>, int>(identity); // implicitly synchronized
71+
*
72+
* auto custom_plus_lambda = [] __device__ (int lhs, int rhs) noexcept { return lhs + rhs; };
73+
* auto f3 = cuco::reduction_functor<decltype(custom_plus_lambda), int>(identity, custom_plus_lambda);
74+
* }
75+
* \endcode
76+
*
77+
* @tparam Func The user-defined reduction functor
78+
* @tparam Value The value type used for reduction
79+
*/
80+
template <typename Func, typename Value>
81+
class reduction_functor : detail::reduction_functor_base {
82+
public:
83+
using value_type = Value;
84+
85+
reduction_functor(cuco::identity_value<Value> identity, Func functor = Func{}) noexcept : identity_(identity), functor_(functor) {}
86+
87+
template <cuda::thread_scope Scope>
88+
__device__ value_type operator()(cuda::atomic<value_type, Scope>& lhs, value_type const& rhs) const noexcept
89+
{
90+
if constexpr (uses_external_sync()) {
91+
value_type old = lhs.load(cuda::memory_order_relaxed);
92+
value_type desired;
93+
94+
do {
95+
desired = functor_(old, rhs);
96+
} while (!lhs.compare_exchange_weak(old, desired, cuda::memory_order_release, cuda::memory_order_relaxed));
97+
98+
return desired;
99+
} else {
100+
return functor_(lhs, rhs);
101+
}
102+
}
103+
104+
__host__ __device__ value_type identity() const noexcept {
105+
return identity_.value();
106+
}
107+
108+
__host__ __device__ static constexpr bool uses_external_sync() noexcept {
109+
return !atomic_invocable_ || naive_invocable_;
110+
}
111+
112+
private:
113+
cuco::identity_value<value_type> identity_;
114+
Func functor_;
115+
static constexpr bool naive_invocable_ = std::is_invocable_r<value_type, Func, value_type, value_type>::value;
116+
static constexpr bool atomic_invocable_ =
117+
std::is_invocable_r<value_type, Func, cuda::atomic<value_type, cuda::thread_scope_system>&, value_type>::value ||
118+
std::is_invocable_r<value_type, Func, cuda::atomic<value_type, cuda::thread_scope_device>&, value_type>::value ||
119+
std::is_invocable_r<value_type, Func, cuda::atomic<value_type, cuda::thread_scope_block>&, value_type>::value ||
120+
std::is_invocable_r<value_type, Func, cuda::atomic<value_type, cuda::thread_scope_thread>&, value_type>::value;
121+
122+
static_assert(atomic_invocable_ || naive_invocable_, "Invalid operator signature.");
123+
};
124+
125+
/**
126+
* @brief Synchronized `+` reduction functor.
127+
*
128+
* @tparam T The value type used for reduction
129+
*/
130+
template <typename T>
131+
auto reduce_add() { return reduction_functor(identity_value<T>{0}, detail::reduce_add_impl<T>{}); };
132+
133+
/**
134+
* @brief Synchronized `min` reduction functor.
135+
*
136+
* @tparam T The value type used for reduction
137+
*/
138+
template <typename T>
139+
auto reduce_min() { return reduction_functor(identity_value{cuda::std::numeric_limits<T>::max()}, detail::reduce_min_impl<T>{}); };
140+
141+
/**
142+
* @brief Synchronized `max` reduction functor.
143+
*
144+
* @tparam T The value type used for reduction
145+
*/
146+
template <typename T>
147+
auto reduce_max() { return reduction_functor(identity_value{cuda::std::numeric_limits<T>::lowest()}, detail::reduce_max_impl<T>{}); };
148+
149+
/**
150+
* @brief Synchronized `count` reduction functor.
151+
*
152+
* @tparam T The value type used for reduction
153+
*/
154+
template <typename T>
155+
auto reduce_count() { return reduction_functor(identity_value<T>{0}, detail::reduce_count_impl<T>{}); };
156+
157+
} // namespace cuco

‎tests/CMakeLists.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,3 +86,8 @@ ConfigureTest(STATIC_MULTIMAP_TEST
8686
static_multimap/multiplicity_test.cu
8787
static_multimap/non_match_test.cu
8888
static_multimap/pair_function_test.cu)
89+
90+
###################################################################################################
91+
# - static_reduction_map tests --------------------------------------------------------------------
92+
ConfigureTest(STATIC_REDUCTION_MAP_TEST
93+
static_reduction_map/reduction_functors_test.cu)

0 commit comments

Comments
 (0)