Fix segmented sort compilation in case of windows.h

NVIDIA · Jan 19, 2022 · c6e338d · c6e338d
1 parent 5d31d2d
commit c6e338d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/cmake/header_test.in b/cmake/header_test.in
@@ -42,5 +42,6 @@
 // projects build with NOMINMAX this doesn't seem to be high priority to fix.
 //#define min(...) CUB_MACRO_CHECK('min', windows.h)
 //#define max(...) CUB_MACRO_CHECK('max', windows.h)
+#define small CUB_MACRO_CHECK('small', windows.h)
 
 #include <cub/${header}>
diff --git a/cub/device/dispatch/dispatch_segmented_sort.cuh b/cub/device/dispatch/dispatch_segmented_sort.cuh
@@ -334,8 +334,11 @@ __global__ void DeviceSegmentedSortKernelSmall(
 
   __shared__ union
   {
-    typename MediumAgentWarpMergeSortT::TempStorage medium[segments_per_medium_block];
-    typename SmallAgentWarpMergeSortT::TempStorage small[segments_per_small_block];
+    typename MediumAgentWarpMergeSortT::TempStorage
+      medium_storage[segments_per_medium_block];
+
+    typename SmallAgentWarpMergeSortT::TempStorage
+      small_storage[segments_per_small_block];
   } temp_storage;
 
   if (bid < medium_blocks)
@@ -353,7 +356,7 @@ __global__ void DeviceSegmentedSortKernelSmall(
       const OffsetT segment_end   = d_end_offsets[global_segment_id];
       const OffsetT num_items     = segment_end - segment_begin;
 
-      MediumAgentWarpMergeSortT(temp_storage.medium[sid_within_block])
+      MediumAgentWarpMergeSortT(temp_storage.medium_storage[sid_within_block])
         .ProcessSegment(num_items,
                         d_keys_in + segment_begin,
                         d_keys_out + segment_begin,
@@ -376,7 +379,7 @@ __global__ void DeviceSegmentedSortKernelSmall(
       const OffsetT segment_end   = d_end_offsets[global_segment_id];
       const OffsetT num_items     = segment_end - segment_begin;
 
-      SmallAgentWarpMergeSortT(temp_storage.small[sid_within_block])
+      SmallAgentWarpMergeSortT(temp_storage.small_storage[sid_within_block])
         .ProcessSegment(num_items,
                         d_keys_in + segment_begin,
                         d_keys_out + segment_begin,