Develop Stream: fix warp dpp reduction

CHANGELOG.md

@@ -38,6 +38,7 @@ Full documentation for rocPRIM is available at [https://rocm.docs.amd.com/projec
 * Fixed an issue where `rocprim::partial_sort_copy` would yield a compile error if the input iterator is const.
 * Fixed incorrect 128-bit signed and unsigned integers type traits.
 * Fixed compilation issue when `rocprim::radix_key_codec<...>` is specialized with a 128-bit integer.
+* Fixed the warp-level reduction `rocprim::warp_reduce.reduce` DPP implementation to avoid undefined intermediate values during the reduction.
 
 ### Upcoming changes
 * Using the initialisation constructor of `rocprim::reverse_iterator` will throw a deprecation warning. It will be marked as explicit in the next major release.

rocprim/include/rocprim/block/detail/block_reduce_warp_reduce.hpp

@@ -190,14 +190,15 @@ class block_reduce_warp_reduce
             }
             ::rocprim::syncthreads();
 
-            if(flat_tid < warps_no_)
+            if(warp_id == 0)
             {
                 // Use warp partial to calculate the final reduce results for every thread
-                auto warp_partial = storage_.warp_partials[lane_id];
+                auto warp_partial = storage_.warp_partials[lane_id % warps_no_];
 
-                warp_reduce<!warps_no_is_pow_of_two_, warp_reduce_output_type>(
-                    warp_partial, output, warps_no_, reduce_op
-                );
+                warp_reduce<!warps_no_is_pow_of_two_, warp_reduce_output_type>(warp_partial,
+                                                                               output,
+                                                                               warps_no_,
+                                                                               reduce_op);
             }
         }
     }

rocprim/include/rocprim/warp/detail/warp_reduce_dpp.hpp

@@ -1,4 +1,4 @@
-// Copyright (c) 2017-2023 Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2017-2024 Advanced Micro Devices, Inc. All rights reserved.
 //
 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to deal
@@ -49,7 +49,7 @@ class warp_reduce_dpp
 
     template<class BinaryFunction>
     ROCPRIM_DEVICE ROCPRIM_INLINE
-    void reduce(T input, T& output, BinaryFunction reduce_op)
+    void reduce_impl(T input, T& output, BinaryFunction reduce_op, std::false_type)
     {
         output = input;
 
@@ -65,13 +65,17 @@ class warp_reduce_dpp
         }
         if(WarpSize > 4)
         {
-            // row_shr:4
-            output = reduce_op(warp_move_dpp<T, 0x114>(output), output);
+            // row_ror:4
+            // Use rotation instead of shift to avoid leaving invalid values in the destination
+            // registers (asume warp size of at least hardware warp-size)
+            output = reduce_op(warp_move_dpp<T, 0x124>(output), output);
         }
         if(WarpSize > 8)
         {
-            // row_shr:8
-            output = reduce_op(warp_move_dpp<T, 0x118>(output), output);
+            // row_ror:8
+            // Use rotation instead of shift to avoid leaving invalid values in the destination
+            // registers (asume warp size of at least hardware warp-size)
+            output = reduce_op(warp_move_dpp<T, 0x128>(output), output);
         }
 #ifdef ROCPRIM_DETAIL_HAS_DPP_BROADCAST
         if(WarpSize > 16)
@@ -97,6 +101,23 @@ class warp_reduce_dpp
         output = warp_shuffle(output, WarpSize - 1, WarpSize);
     }
 
+    template<class BinaryFunction>
+    ROCPRIM_DEVICE ROCPRIM_INLINE
+    void reduce_impl(T input, T& output, BinaryFunction reduce_op, std::true_type)
+    {
+        warp_reduce_shuffle<T, WarpSize, UseAllReduce>().reduce(input, output, reduce_op);
+    }
+
+    template<class BinaryFunction>
+    ROCPRIM_DEVICE ROCPRIM_INLINE
+    void reduce(T input, T& output, BinaryFunction reduce_op)
+    {
+        reduce_impl(input,
+                    output,
+                    reduce_op,
+                    std::integral_constant<bool, (WarpSize < ::rocprim::device_warp_size())>{});
+    }
+
     template<class BinaryFunction>
     ROCPRIM_DEVICE ROCPRIM_INLINE
     void reduce(T input, T& output, storage_type& storage, BinaryFunction reduce_op)