Merge branch 'master' of github.com:dumerrill/PrivateCub

Former-commit-id: 0e44d55572ec87852c16c9bd79442c9b88773a3c
GaryShen2008 · Mar 18, 2015 · 8842512 · 8842512
2 parents 130fe97 + b597468
commit 8842512
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 12 deletions.
diff --git a/cub/agent/agent_radix_sort_downsweep.cuh b/cub/agent/agent_radix_sort_downsweep.cuh
@@ -415,6 +415,22 @@ struct AgentRadixSortDownsweep
     }
 
 
+    /**
+     * Load a tile of items (specialized for full tile)
+     */
+    template <typename BlockLoadT, typename T, typename InputIteratorT>
+    __device__ __forceinline__ void LoadItems(
+        BlockLoadT      &block_loader,
+        T               (&items)[ITEMS_PER_THREAD],
+        InputIteratorT  d_in,
+        OffsetT         valid_items,
+        T               oob_item,
+        Int2Type<true>  is_full_tile)
+    {
+        block_loader.Load(d_in, items);
+    }
+
+
     /**
      * Load a tile of items (specialized for partial tile)
      */
@@ -429,6 +445,21 @@ struct AgentRadixSortDownsweep
         block_loader.Load(d_in, items, valid_items);
     }
 
+    /**
+     * Load a tile of items (specialized for partial tile)
+     */
+    template <typename BlockLoadT, typename T, typename InputIteratorT>
+    __device__ __forceinline__ void LoadItems(
+        BlockLoadT      &block_loader,
+        T               (&items)[ITEMS_PER_THREAD],
+        InputIteratorT  d_in,
+        OffsetT         valid_items,
+        T               oob_item,
+        Int2Type<false> is_full_tile)
+    {
+        block_loader.Load(d_in, items, valid_items, oob_item);
+    }
+
 
     /**
      * Truck along associated values
@@ -487,12 +518,8 @@ struct AgentRadixSortDownsweep
         int             ranks[ITEMS_PER_THREAD];                    // For each key, the local rank within the CTA
         OffsetT         relative_bin_offsets[ITEMS_PER_THREAD];     // For each key, the global scatter base offset of the corresponding digit
 
-        // Assign max-key to all keys
-        #pragma unroll
-        for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
-        {
-            keys[ITEM] = (DESCENDING) ? MIN_KEY : MAX_KEY;
-        }
+        // Assign default (min/max) value to all keys
+        UnsignedBits default_key = (DESCENDING) ? MIN_KEY : MAX_KEY;
 
         // Load tile of keys
         BlockLoadKeys loader(temp_storage.load_keys);
@@ -501,6 +528,7 @@ struct AgentRadixSortDownsweep
             keys,
             d_keys_in + block_offset,
             valid_items, 
+            default_key,
             Int2Type<FULL_TILE>());
 
         __syncthreads();

diff --git a/cub/device/dispatch/dispatch_radix_sort.cuh b/cub/device/dispatch/dispatch_radix_sort.cuh
@@ -243,14 +243,9 @@ __global__ void DeviceRadixSortSingleKernel(
 
     // Assign default (min/max) value to all keys
     UnsignedBits default_key = (DESCENDING) ? Traits<KeyT>::MIN_KEY : Traits<KeyT>::MAX_KEY;
-    #pragma unroll
-    for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM)
-    {
-        keys[ITEM] = reinterpret_cast<KeyT&>(default_key);
-    }
 
     // Load keys
-    BlockLoadKeys(temp_storage.load_keys).Load(d_keys_in, keys, num_items);
+    BlockLoadKeys(temp_storage.load_keys).Load(d_keys_in, keys, num_items, default_key);
 
     __syncthreads();