Draft of segmented reduce optimization

cub/agent/agent_reduce.cuh

@@ -80,6 +80,32 @@ struct AgentReducePolicy : ScalingType
   static constexpr CacheLoadModifier LOAD_MODIFIER = _LOAD_MODIFIER;
 };
 
+template <int BLOCK_THREADS,
+          int NOMINAL_WARP_THREADS_4B,
+          int NOMINAL_ITEMS_PER_THREAD_4B,
+          typename ComputeT,
+          int _VECTOR_LOAD_LENGTH,
+          CacheLoadModifier _LOAD_MODIFIER>
+struct AgentWarpReducePolicy 
+{
+  // TODO MemBoundScaling-like computation
+  static constexpr int ITEMS_PER_THREAD = NOMINAL_ITEMS_PER_THREAD_4B;
+
+  static constexpr int WARP_THREADS = NOMINAL_WARP_THREADS_4B;
+
+  /// Number of items per vectorized load
+  static constexpr int VECTOR_LOAD_LENGTH = _VECTOR_LOAD_LENGTH;
+
+  /// Cache load modifier for reading input elements
+  static constexpr CacheLoadModifier LOAD_MODIFIER = _LOAD_MODIFIER;
+
+  constexpr static int ITEMS_PER_TILE = ITEMS_PER_THREAD * WARP_THREADS; 
+
+  constexpr static int SEGMENTS_PER_BLOCK = BLOCK_THREADS / WARP_THREADS;  
+
+  static_assert((BLOCK_THREADS % WARP_THREADS) == 0, "Block should be multiple of warp");
+};
+
 /******************************************************************************
  * Thread block abstractions
  ******************************************************************************/
@@ -116,8 +142,10 @@ template <typename AgentReducePolicy,
           typename OutputIteratorT,
           typename OffsetT,
           typename ReductionOp,
-          typename AccumT>
-struct AgentReduce
+          typename AccumT,
+          typename CollectiveReduceT,
+          int THREADS>
+struct AgentReduceImpl
 {
   //---------------------------------------------------------------------
   // Types and constants
@@ -139,9 +167,8 @@ struct AgentReduce
     InputIteratorT>;
 
   /// Constants
-  static constexpr int BLOCK_THREADS      = AgentReducePolicy::BLOCK_THREADS;
   static constexpr int ITEMS_PER_THREAD   = AgentReducePolicy::ITEMS_PER_THREAD;
-  static constexpr int TILE_ITEMS         = BLOCK_THREADS * ITEMS_PER_THREAD;
+  static constexpr int TILE_ITEMS         = THREADS * ITEMS_PER_THREAD;
   static constexpr int VECTOR_LOAD_LENGTH =
     CUB_MIN(ITEMS_PER_THREAD, AgentReducePolicy::VECTOR_LOAD_LENGTH);
 
@@ -155,17 +182,10 @@ struct AgentReduce
   static constexpr CacheLoadModifier LOAD_MODIFIER =
     AgentReducePolicy::LOAD_MODIFIER;
 
-  static constexpr BlockReduceAlgorithm BLOCK_ALGORITHM =
-    AgentReducePolicy::BLOCK_ALGORITHM;
-
-  /// Parameterized BlockReduce primitive
-  using BlockReduceT =
-    BlockReduce<AccumT, BLOCK_THREADS, AgentReducePolicy::BLOCK_ALGORITHM>;
-
   /// Shared memory type required by this thread block
   struct _TempStorage
   {
-    typename BlockReduceT::TempStorage reduce;
+    typename CollectiveReduceT::TempStorage reduce;
   };
 
   /// Alias wrapper allowing storage to be unioned
@@ -177,9 +197,10 @@ struct AgentReduce
   //---------------------------------------------------------------------
 
   _TempStorage &temp_storage;         ///< Reference to temp_storage
-  InputIteratorT d_in;                ///< Input data to reduce
+  unsigned int lane_id;
   WrappedInputIteratorT d_wrapped_in; ///< Wrapped input data to reduce
   ReductionOp reduction_op;           ///< Binary reduction operator
+  InputIteratorT d_in;                ///< Input data to reduce
 
   //---------------------------------------------------------------------
   // Utility
@@ -213,13 +234,15 @@ struct AgentReduce
    * @param d_in Input data to reduce
    * @param reduction_op Binary reduction operator
    */
-  __device__ __forceinline__ AgentReduce(TempStorage &temp_storage,
-                                         InputIteratorT d_in,
-                                         ReductionOp reduction_op)
+  __device__ __forceinline__ AgentReduceImpl(TempStorage &temp_storage,
+                                             InputIteratorT d_in,
+                                             ReductionOp reduction_op,
+                                             unsigned int lane_id)
       : temp_storage(temp_storage.Alias())
       , d_in(d_in)
       , d_wrapped_in(d_in)
       , reduction_op(reduction_op)
+      , lane_id(lane_id)
   {}
 
   //---------------------------------------------------------------------
@@ -243,9 +266,9 @@ struct AgentReduce
     AccumT items[ITEMS_PER_THREAD];
 
     // Load items in striped fashion
-    LoadDirectStriped<BLOCK_THREADS>(threadIdx.x,
-                                     d_wrapped_in + block_offset,
-                                     items);
+    LoadDirectStriped<THREADS>(lane_id,
+                               d_wrapped_in + block_offset,
+                               items);
 
     // Reduce items within each thread stripe
     thread_aggregate =
@@ -276,7 +299,7 @@ struct AgentReduce
 
     // Fabricate a vectorized input iterator
     InputT *d_in_unqualified = const_cast<InputT *>(d_in) + block_offset +
-                               (threadIdx.x * VECTOR_LOAD_LENGTH);
+                               (lane_id * VECTOR_LOAD_LENGTH);
     CacheModifiedInputIterator<AgentReducePolicy::LOAD_MODIFIER, VectorT, OffsetT>
       d_vec_in(reinterpret_cast<VectorT *>(d_in_unqualified));
 
@@ -286,7 +309,7 @@ struct AgentReduce
 #pragma unroll
     for (int i = 0; i < WORDS; ++i)
     {
-      vec_items[i] = d_vec_in[BLOCK_THREADS * i];
+      vec_items[i] = d_vec_in[THREADS * i];
     }
 
     // Convert from input type to output type
@@ -320,13 +343,13 @@ struct AgentReduce
               Int2Type<CAN_VECTORIZE> /*can_vectorize*/)
   {
     // Partial tile
-    int thread_offset = threadIdx.x;
+    int thread_offset = lane_id;
 
     // Read first item
     if ((IS_FIRST_TILE) && (thread_offset < valid_items))
     {
       thread_aggregate = d_wrapped_in[block_offset + thread_offset];
-      thread_offset += BLOCK_THREADS;
+      thread_offset += THREADS;
     }
 
     // Continue reading items (block-striped)
@@ -335,7 +358,7 @@ struct AgentReduce
       InputT item(d_wrapped_in[block_offset + thread_offset]);
 
       thread_aggregate = reduction_op(thread_aggregate, item);
-      thread_offset += BLOCK_THREADS;
+      thread_offset += THREADS;
     }
   }
 
@@ -364,8 +387,12 @@ struct AgentReduce
                         valid_items,
                         Int2Type<false>(),
                         can_vectorize);
-      return BlockReduceT(temp_storage.reduce)
-        .Reduce(thread_aggregate, reduction_op, valid_items);
+
+      // TODO Extract clamping into the SFINAE to keep block version as is
+      int num_valid = (THREADS <= valid_items) ? THREADS : valid_items;
+
+      return CollectiveReduceT(temp_storage.reduce)
+        .Reduce(thread_aggregate, reduction_op, num_valid);
     }
 
     // At least one full block
@@ -399,7 +426,7 @@ struct AgentReduce
     }
 
     // Compute block-wide reduction (all threads have valid items)
-    return BlockReduceT(temp_storage.reduce)
+    return CollectiveReduceT(temp_storage.reduce)
       .Reduce(thread_aggregate, reduction_op);
   }
 
@@ -440,5 +467,76 @@ struct AgentReduce
   }
 };
 
+template <typename AgentReducePolicy,
+          typename InputIteratorT,
+          typename OutputIteratorT,
+          typename OffsetT,
+          typename ReductionOp,
+          typename AccumT>
+struct AgentReduce : AgentReduceImpl<AgentReducePolicy, 
+                                     InputIteratorT, 
+                                     OutputIteratorT, 
+                                     OffsetT, 
+                                     ReductionOp, 
+                                     AccumT, 
+                                     BlockReduce<AccumT, 
+                                                 AgentReducePolicy::BLOCK_THREADS, 
+                                                 AgentReducePolicy::BLOCK_ALGORITHM>,
+                                     AgentReducePolicy::BLOCK_THREADS>
+{
+  using base_t = AgentReduceImpl<AgentReducePolicy, 
+                                     InputIteratorT, 
+                                     OutputIteratorT, 
+                                     OffsetT, 
+                                     ReductionOp, 
+                                     AccumT, 
+                                     BlockReduce<AccumT, 
+                                                 AgentReducePolicy::BLOCK_THREADS, 
+                                                 AgentReducePolicy::BLOCK_ALGORITHM>,
+                                     AgentReducePolicy::BLOCK_THREADS>;
+
+  __device__ __forceinline__ AgentReduce(typename base_t::TempStorage &temp_storage,
+                                         InputIteratorT d_in,
+                                         ReductionOp reduction_op)
+    : base_t(temp_storage, d_in, reduction_op, threadIdx.x) 
+  {
+  }
+};
+
+template <typename AgentReducePolicy,
+          typename InputIteratorT,
+          typename OutputIteratorT,
+          typename OffsetT,
+          typename ReductionOp,
+          typename AccumT>
+struct AgentWarpReduce : AgentReduceImpl<AgentReducePolicy, 
+                                         InputIteratorT, 
+                                         OutputIteratorT, 
+                                         OffsetT, 
+                                         ReductionOp, 
+                                         AccumT, 
+                                         WarpReduce<AccumT, 
+                                                    AgentReducePolicy::WARP_THREADS>,
+                                         AgentReducePolicy::WARP_THREADS>
+{
+  using base_t = AgentReduceImpl<AgentReducePolicy, 
+                                     InputIteratorT, 
+                                     OutputIteratorT, 
+                                     OffsetT, 
+                                     ReductionOp, 
+                                     AccumT, 
+                                     WarpReduce<AccumT, 
+                                                AgentReducePolicy::WARP_THREADS>,
+                                     AgentReducePolicy::WARP_THREADS>;
+
+  __device__ __forceinline__ AgentWarpReduce(typename base_t::TempStorage &temp_storage,
+                                             InputIteratorT d_in,
+                                             ReductionOp reduction_op,
+                                             int lane_id)
+    : base_t(temp_storage, d_in, reduction_op, lane_id) 
+  {
+  }
+};
+
 CUB_NAMESPACE_END