Extend fallback coverage for copy_if (#512)

We recently added a "fallback" implementation for thrust::copy_if that
is invoked when copying a custom type that's too large to fit in
shared memory.

This change extends the fallback slightly so that it can be used with
an overload of copy_if that accepts a stencil buffer (to copy by
key). It also adds a unit test to cover this case.

It also fixes a small bug in the fallback implementation that could
cause the scan accumulator type to overflow when the results are
compacted.

Author: umfranzw

test/test_copy.cpp

@@ -444,6 +444,66 @@ TYPED_TEST(CopyIntegerTests, TestCopyIf)
     }
 }
 
+TEST(CopyLargeTypesTests, TestCopyIfStencilLargeType)
+{
+    using T = large_data;
+
+    SCOPED_TRACE(testing::Message() << "with device_id= " << test::set_device_from_ctest());
+
+    for(auto size : get_sizes())
+    {
+        SCOPED_TRACE(testing::Message() << "with size= " << size);
+
+        thrust::host_vector<T> h_data(size);
+        thrust::sequence(h_data.begin(), h_data.end());
+        thrust::device_vector<T> d_data(size);
+        thrust::sequence(d_data.begin(), d_data.end());
+
+        for(auto seed : get_seeds())
+        {
+            SCOPED_TRACE(testing::Message() << "with seed= " << seed);
+
+            thrust::host_vector<T> h_stencil = get_random_data<int>(size, std::numeric_limits<int>::min(), std::numeric_limits<int>::max(), seed);;
+            thrust::device_vector<T> d_stencil = h_stencil;
+
+            typename thrust::host_vector<T>::iterator   h_new_end;
+            typename thrust::device_vector<T>::iterator d_new_end;
+
+            // test with Predicate that returns a bool
+            {
+                thrust::host_vector<T>   h_result(size);
+                thrust::device_vector<T> d_result(size);
+
+                h_new_end
+                   = thrust::copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_result.begin(), is_even<T>());
+                d_new_end
+                    = thrust::copy_if(d_data.begin(), d_data.end(), d_stencil.begin(), d_result.begin(), is_even<T>());
+
+                h_result.resize(h_new_end - h_result.begin());
+                d_result.resize(d_new_end - d_result.begin());
+
+                ASSERT_EQ(h_result, d_result);
+            }
+
+            // test with Predicate that returns a non-bool
+            {
+                thrust::host_vector<T>   h_result(size);
+                thrust::device_vector<T> d_result(size);
+
+                h_new_end
+                    = thrust::copy_if(h_data.begin(), h_data.end(), h_stencil.begin(), h_result.begin(), mod_3<T>());
+                d_new_end
+                    = thrust::copy_if(d_data.begin(), d_data.end(), d_stencil.begin(), d_result.begin(), mod_3<T>());
+
+                h_result.resize(h_new_end - h_result.begin());
+                d_result.resize(d_new_end - d_result.begin());
+
+                ASSERT_EQ(h_result, d_result);
+            }
+        }
+    }
+}
+
 TYPED_TEST(CopyIntegerTests, TestCopyIfStencil)
 {
     using T = typename TestFixture::input_type;

thrust/system/hip/detail/copy_if.h

@@ -156,49 +156,99 @@ namespace __copy_if
         return output + num_selected;
     }
 
-    template <typename Derived, typename InputIt, typename OutputIt, typename Predicate>
+    template <typename Derived, typename InputIt, typename OutputIt, typename Predicate, typename PredicateInputIt>
     THRUST_HIP_RUNTIME_FUNCTION auto
-    copy_if(execution_policy<Derived>& policy, InputIt first, InputIt last, OutputIt output, Predicate predicate)
+    copy_if_common(execution_policy<Derived>& policy, InputIt first, InputIt last, OutputIt output, Predicate predicate, PredicateInputIt predicate_input)
     -> std::enable_if_t<!(sizeof(typename std::iterator_traits<InputIt>::value_type) < 512), OutputIt>
     {
         using namespace thrust::system::hip_rocprim::temp_storage;
         using size_type = typename iterator_traits<InputIt>::difference_type;
+        using pos_type = thrust::detail::uint32_t;
+        using flag_type = thrust::detail::uint8_t;
 
         size_type   num_items  = thrust::distance(first, last);
         hipStream_t stream     = hip_rocprim::stream(policy);
         bool        debug_sync = THRUST_HIP_DEBUG_SYNC_FLAG;
 
-        thrust::detail::temporary_array<thrust::detail::uint8_t, Derived> flags(policy, num_items);
+        if(num_items == 0)
+            return output;
+
+        // Note: although flags can be stored in a uint8_t, in the inclusive scan performed on flags below,
+        // the scan accumulator type to something larger (flag_type) to prevent overflow.
+        // For this reason, we call rocprim::inclusive_scan directly here and pass in the accumulator type as template argument.
+        thrust::detail::temporary_array<flag_type, Derived> flags(policy, num_items);
 
-        hip_rocprim::throw_on_error(rocprim::transform(first,
+        hip_rocprim::throw_on_error(rocprim::transform(predicate_input,
                                                        flags.begin(),
                                                        num_items,
                                                        [predicate] __host__ __device__ (auto const & val){ return predicate(val) ? 1 : 0; },
                                                        stream,
                                                        debug_sync),
                                     "copy_if failed on transform");
 
-        thrust::detail::temporary_array<thrust::detail::uint32_t, Derived> pos(policy, num_items);
+        thrust::detail::temporary_array<pos_type, Derived> pos(policy, num_items);
+
+        // Get the required temporary storage size.
+        size_t storage_size = 0;
+        hip_rocprim::throw_on_error(rocprim::inclusive_scan<rocprim::default_config,
+                                    typename thrust::detail::temporary_array<flag_type, Derived>::iterator,
+                                    typename thrust::detail::temporary_array<pos_type, Derived>::iterator,
+                                    rocprim::plus<pos_type>,
+                                    pos_type>(nullptr, storage_size, flags.begin(), pos.begin(), num_items, rocprim::plus<pos_type>{}, stream, debug_sync),
+            "copy_if failed while determining inclusive scan storage size");
 
-        thrust::inclusive_scan(policy, flags.begin(), flags.end(), pos.begin());
+        // Allocate temporary storage.
+        thrust::detail::temporary_array<thrust::detail::uint8_t, Derived> tmp(policy, storage_size);
+        void *ptr = static_cast<void*>(tmp.data().get());
 
+        // Perform a scan on the positions.
+        hip_rocprim::throw_on_error(rocprim::inclusive_scan<rocprim::default_config,
+                                    typename thrust::detail::temporary_array<flag_type, Derived>::iterator,
+                                    typename thrust::detail::temporary_array<pos_type, Derived>::iterator,
+                                    rocprim::plus<pos_type>,
+                                    pos_type>(ptr, storage_size, flags.begin(), pos.begin(), num_items, rocprim::plus<pos_type>{}, stream, debug_sync),
+            "copy_if failed on inclusive scan");
+
+        // Pull out the values for which the predicate evaluated to true and compact them into the output array.
         constexpr static size_t items_per_thread = 16;
         constexpr static size_t threads_per_block = 256;
-        const size_t block_size = std::ceil(static_cast<float>(num_items) / 16 / threads_per_block);
+        const size_t block_size = std::ceil(static_cast<float>(num_items) / items_per_thread / threads_per_block);
 
         copy_if_kernel<items_per_thread><<<block_size, threads_per_block>>>(first, flags.begin(), pos.begin(), num_items, output);
 
-        return output + pos[num_items-1];
+        return output + pos[num_items - 1];
+    }
+
+    template <typename Derived, typename InputIt, typename OutputIt, typename Predicate>
+    THRUST_HIP_RUNTIME_FUNCTION auto
+    copy_if(execution_policy<Derived>& policy, InputIt first, InputIt last, OutputIt output, Predicate predicate)
+    -> std::enable_if_t<!(sizeof(typename std::iterator_traits<InputIt>::value_type) < 512), OutputIt>
+    {
+        return copy_if_common(policy, first, last, output, predicate, first);
+    }
+
+    template <typename Derived, typename InputIt, typename StencilIt, typename OutputIt, typename Predicate>
+    THRUST_HIP_RUNTIME_FUNCTION auto
+    copy_if(execution_policy<Derived>& policy,
+            InputIt                    first,
+            InputIt                    last,
+            StencilIt                  stencil,
+            OutputIt                   output,
+            Predicate                  predicate)
+    -> std::enable_if_t<!(sizeof(typename std::iterator_traits<InputIt>::value_type) < 512), OutputIt>
+    {
+        return copy_if_common(policy, first, last, output, predicate, stencil);
     }
 
     template <typename Derived, typename InputIt, typename StencilIt, typename OutputIt, typename Predicate>
-    THRUST_HIP_RUNTIME_FUNCTION OutputIt
+    THRUST_HIP_RUNTIME_FUNCTION auto
     copy_if(execution_policy<Derived>& policy,
             InputIt                    first,
             InputIt                    last,
             StencilIt                  stencil,
             OutputIt                   output,
             Predicate                  predicate)
+        -> std::enable_if_t<(sizeof(typename std::iterator_traits<InputIt>::value_type) < 512), OutputIt>
     {
         using namespace thrust::system::hip_rocprim::temp_storage;
         typedef typename iterator_traits<InputIt>::difference_type size_type;
@@ -264,7 +314,6 @@ namespace __copy_if
 
         return output + num_selected;
     }
-
 } // namespace __copy_if
 
 //-------------------------