warp-lang 1.6.1__py3-none-win_amd64.whl → 1.7.0__py3-none-win_amd64.whl

warp/native/sort.cpp

@@ -1,9 +1,18 @@
-/** Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
- * NVIDIA CORPORATION and its licensors retain all intellectual property
- * and proprietary rights in and to this software, related documentation
- * and any modifications thereto.  Any use, reproduction, disclosure or
- * distribution of this software and related documentation without an express
- * license agreement from NVIDIA CORPORATION is strictly prohibited.
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 #include "warp.h"
@@ -12,69 +21,75 @@
 
 #include <cstdint>
 
-void radix_sort_pairs_host(int* keys, int* values, int n)
+//Only integer keys (bit count 32 or 64) are supported. Floats need to get converted into int first. see radix_float_to_int.
+template <typename KeyType>
+void radix_sort_pairs_host(KeyType* keys, int* values, int n, int offset_to_scratch_memory)
 {
-	static int tables[2][1 << 16];
+	const int numPasses = sizeof(KeyType) / 2;
+	static int tables[numPasses][1 << 16];
 	memset(tables, 0, sizeof(tables));
-		
-	int* auxKeys = keys + n;
-	int* auxValues = values + n;
-
+	
 	// build histograms
-	for (int i=0; i < n; ++i)
-	{
-		const unsigned short low = keys[i] & 0xffff;
-		const unsigned short high = keys[i] >> 16;
-		
-		++tables[0][low];
-		++tables[1][high];
+	for (int p = 0; p < numPasses; ++p)
+    {
+		for (int i=0; i < n; ++i)
+		{
+			const int shift = p * 16;
+			const int b = (keys[i] >> shift) & 0xffff;
+
+			++tables[p][b];
+		}
 	}
 	
-	// convert histograms to offset tables in-place
-	int offlow = 0;
-	int offhigh = 0;
-	
-	for (int i=0; i < 65536; ++i)
+	// convert histograms to offset tables in-place	
+	for (int p = 0; p < numPasses; ++p)
 	{
-		const int newofflow = offlow + tables[0][i];
-		const int newoffhigh = offhigh + tables[1][i];
-		
-		tables[0][i] = offlow;
-		tables[1][i] = offhigh;
-		
-		offlow = newofflow;
-		offhigh = newoffhigh;
+		int off = 0;
+		for (int i = 0; i < 65536; ++i)
+		{
+			const int newoff = off + tables[p][i];
+			
+			tables[p][i] = off;
+			
+			off = newoff;
+		}
 	}
-		
-	// pass 1 - sort by low 16 bits
-	for (int i=0; i < n; ++i)
-	{
-		// lookup offset of input
-		const int k = keys[i];
-		const int v = values[i];
-		const int b = k & 0xffff;
-		
-		// find offset and increment
-		const int offset = tables[0][b]++;
-		
-		auxKeys[offset] = k;
-		auxValues[offset] = v;
-	}	
-		
-	// pass 2 - sort by high 16 bits
-	for (int i=0; i < n; ++i)
-	{
-		// lookup offset of input
-		const int k = auxKeys[i];
-		const int v = auxValues[i];
+	
+    for (int p = 0; p < numPasses; ++p)
+    {
+		int flipFlop = p % 2;
+		KeyType* readKeys = keys + offset_to_scratch_memory * flipFlop;
+		int* readValues = values + offset_to_scratch_memory * flipFlop;
+		KeyType* writeKeys = keys + offset_to_scratch_memory * (1 - flipFlop);
+		int* writeValues = values + offset_to_scratch_memory * (1 - flipFlop);
+
+		// pass 1 - sort by low 16 bits
+		for (int i=0; i < n; ++i)
+		{
+			// lookup offset of input
+			const KeyType k = readKeys[i];
+			const int v = readValues[i];
+			
+			const int shift = p * 16;
+			const int b = (k >> shift) & 0xffff;
+			
+			// find offset and increment
+			const int offset = tables[p][b]++;
+			
+			writeKeys[offset] = k;
+			writeValues[offset] = v;
+		}
+	}
+}
 
-		const int b = k >> 16;
-		
-		const int offset = tables[1][b]++;
-		
-		keys[offset] = k;
-		values[offset] = v;
-	}	
+void radix_sort_pairs_host(int* keys, int* values, int n)
+{
+	radix_sort_pairs_host<int>(keys, values, n, n);
+}
+
+void radix_sort_pairs_host(int64_t* keys, int* values, int n)
+{
+	radix_sort_pairs_host<int64_t>(keys, values, n, n);
 }
 
  //http://stereopsis.com/radix.html
@@ -85,13 +100,13 @@ inline unsigned int radix_float_to_int(float f)
 	return i ^ mask;
 }
 
-void radix_sort_pairs_host(float* keys, int* values, int n)
+void radix_sort_pairs_host(float* keys, int* values, int n, int offset_to_scratch_memory)
 {
 	static unsigned int tables[2][1 << 16];
 	memset(tables, 0, sizeof(tables));
 		
-	float* auxKeys = keys + n;
-	int* auxValues = values + n;
+	float* auxKeys = keys + offset_to_scratch_memory;
+	int* auxValues = values + offset_to_scratch_memory;
 
 	// build histograms
 	for (int i=0; i < n; ++i)
@@ -153,14 +168,46 @@ void radix_sort_pairs_host(float* keys, int* values, int n)
 	}	
 }
 
+void radix_sort_pairs_host(float* keys, int* values, int n)
+{
+	radix_sort_pairs_host(keys, values, n, n);
+}
+
+void segmented_sort_pairs_host(float* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments)
+{
+	for (int i = 0; i < num_segments; ++i)
+	{
+		const int start = segment_start_indices[i];
+		const int end = segment_end_indices[i];
+		radix_sort_pairs_host(keys + start, values + start, end - start, n);
+	}
+}
+
+void segmented_sort_pairs_host(int* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments)
+{
+	for (int i = 0; i < num_segments; ++i)
+	{
+		const int start = segment_start_indices[i];
+		const int end = segment_end_indices[i];
+		radix_sort_pairs_host(keys + start, values + start, end - start, n);
+	}
+}
+
+
 #if !WP_ENABLE_CUDA
 
 void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out) {}
 
 void radix_sort_pairs_int_device(uint64_t keys, uint64_t values, int n) {}
 
+void radix_sort_pairs_int64_device(uint64_t keys, uint64_t values, int n) {}
+
 void radix_sort_pairs_float_device(uint64_t keys, uint64_t values, int n) {}
 
+void segmented_sort_pairs_float_device(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments) {}
+
+void segmented_sort_pairs_int_device(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments) {}
+
 #endif // !WP_ENABLE_CUDA
 
 
@@ -171,9 +218,34 @@ void radix_sort_pairs_int_host(uint64_t keys, uint64_t values, int n)
         reinterpret_cast<int *>(values), n);
 }
 
+void radix_sort_pairs_int64_host(uint64_t keys, uint64_t values, int n)
+{
+    radix_sort_pairs_host(
+        reinterpret_cast<int64_t *>(keys),
+        reinterpret_cast<int *>(values), n);
+}
+
 void radix_sort_pairs_float_host(uint64_t keys, uint64_t values, int n)
 {
     radix_sort_pairs_host(
         reinterpret_cast<float *>(keys),
         reinterpret_cast<int *>(values), n);
-}
+}
+
+void segmented_sort_pairs_float_host(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments)
+{
+    segmented_sort_pairs_host(
+        reinterpret_cast<float *>(keys),
+        reinterpret_cast<int *>(values), n,
+        reinterpret_cast<int *>(segment_start_indices),
+        reinterpret_cast<int *>(segment_end_indices), num_segments);
+}
+
+void segmented_sort_pairs_int_host(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments)
+{
+    segmented_sort_pairs_host(
+        reinterpret_cast<int *>(keys),
+        reinterpret_cast<int *>(values), n,
+        reinterpret_cast<int *>(segment_start_indices),
+        reinterpret_cast<int *>(segment_end_indices), num_segments);
+}

warp/native/sort.cu

@@ -1,9 +1,18 @@
-/** Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
- * NVIDIA CORPORATION and its licensors retain all intellectual property
- * and proprietary rights in and to this software, related documentation
- * and any modifications thereto.  Any use, reproduction, disclosure or
- * distribution of this software and related documentation without an express
- * license agreement from NVIDIA CORPORATION is strictly prohibited.
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 #include "warp.h"
@@ -27,11 +36,12 @@ struct RadixSortTemp
 static std::map<void*, RadixSortTemp> g_radix_sort_temp_map;
 
 
-void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out)
+template <typename KeyType>
+void radix_sort_reserve_internal(void* context, int n, void** mem_out, size_t* size_out)
 {
     ContextGuard guard(context);
 
-    cub::DoubleBuffer<int> d_keys;
+    cub::DoubleBuffer<KeyType> d_keys;
 	cub::DoubleBuffer<int> d_values;
 
     // compute temporary memory required
@@ -41,7 +51,7 @@ void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out)
         sort_temp_size,
         d_keys,
         d_values,
-        n, 0, 32,
+        n, 0, sizeof(KeyType)*8,
         (cudaStream_t)cuda_stream_get_current()));
 
     if (!context)
@@ -62,15 +72,21 @@ void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out)
         *size_out = temp.size;
 }
 
-void radix_sort_pairs_device(void* context, int* keys, int* values, int n)
+void radix_sort_reserve(void* context, int n, void** mem_out, size_t* size_out)
+{
+    radix_sort_reserve_internal<int>(context, n, mem_out, size_out);
+}
+
+template <typename KeyType>
+void radix_sort_pairs_device(void* context, KeyType* keys, int* values, int n)
 {
     ContextGuard guard(context);
 
-    cub::DoubleBuffer<int> d_keys(keys, keys + n);
+    cub::DoubleBuffer<KeyType> d_keys(keys, keys + n);
 	cub::DoubleBuffer<int> d_values(values, values + n);
 
     RadixSortTemp temp;
-    radix_sort_reserve(WP_CURRENT_CONTEXT, n, &temp.mem, &temp.size);
+    radix_sort_reserve_internal<KeyType>(WP_CURRENT_CONTEXT, n, &temp.mem, &temp.size);
 
     // sort
     check_cuda(cub::DeviceRadixSort::SortPairs(
@@ -78,16 +94,31 @@ void radix_sort_pairs_device(void* context, int* keys, int* values, int n)
         temp.size,
         d_keys, 
         d_values, 
-        n, 0, 32, 
+        n, 0, sizeof(KeyType)*8, 
         (cudaStream_t)cuda_stream_get_current()));
 
 	if (d_keys.Current() != keys)
-		memcpy_d2d(WP_CURRENT_CONTEXT, keys, d_keys.Current(), sizeof(int)*n);
+		memcpy_d2d(WP_CURRENT_CONTEXT, keys, d_keys.Current(), sizeof(KeyType)*n);
 
 	if (d_values.Current() != values)
 		memcpy_d2d(WP_CURRENT_CONTEXT, values, d_values.Current(), sizeof(int)*n);
 }
 
+void radix_sort_pairs_device(void* context, int* keys, int* values, int n)
+{
+    radix_sort_pairs_device<int>(context, keys, values, n);
+}
+
+void radix_sort_pairs_device(void* context, float* keys, int* values, int n)
+{
+    radix_sort_pairs_device<float>(context, keys, values, n);
+}
+
+void radix_sort_pairs_device(void* context, int64_t* keys, int* values, int n)
+{
+    radix_sort_pairs_device<int64_t>(context, keys, values, n);
+}
+
 void radix_sort_pairs_int_device(uint64_t keys, uint64_t values, int n)
 {
     radix_sort_pairs_device(
@@ -96,7 +127,69 @@ void radix_sort_pairs_int_device(uint64_t keys, uint64_t values, int n)
         reinterpret_cast<int *>(values), n);
 }
 
-void radix_sort_pairs_device(void* context, float* keys, int* values, int n)
+void radix_sort_pairs_float_device(uint64_t keys, uint64_t values, int n)
+{
+    radix_sort_pairs_device(
+        WP_CURRENT_CONTEXT,
+        reinterpret_cast<float *>(keys),
+        reinterpret_cast<int *>(values), n);
+}
+
+void radix_sort_pairs_int64_device(uint64_t keys, uint64_t values, int n)
+{
+    radix_sort_pairs_device(
+        WP_CURRENT_CONTEXT,
+        reinterpret_cast<int64_t *>(keys),
+        reinterpret_cast<int *>(values), n);
+}
+
+void segmented_sort_reserve(void* context, int n, int num_segments, void** mem_out, size_t* size_out)
+{
+    ContextGuard guard(context);
+
+    cub::DoubleBuffer<int> d_keys;
+	cub::DoubleBuffer<int> d_values;
+
+    int* start_indices = NULL;
+    int* end_indices = NULL;
+
+    // compute temporary memory required
+	size_t sort_temp_size;
+    check_cuda(cub::DeviceSegmentedRadixSort::SortPairs(
+        NULL,
+        sort_temp_size,
+        d_keys,
+        d_values,
+        n, 
+        num_segments,
+        start_indices,
+        end_indices,
+        0,
+        32,
+        (cudaStream_t)cuda_stream_get_current()));
+
+    if (!context)
+        context = cuda_context_get_current();
+
+    RadixSortTemp& temp = g_radix_sort_temp_map[context];
+
+    if (sort_temp_size > temp.size)
+    {
+	    free_device(WP_CURRENT_CONTEXT, temp.mem);
+        temp.mem = alloc_device(WP_CURRENT_CONTEXT, sort_temp_size);
+        temp.size = sort_temp_size;
+    }
+    
+    if (mem_out)
+        *mem_out = temp.mem;
+    if (size_out)
+        *size_out = temp.size;
+}
+
+// segment_start_indices and segment_end_indices are arrays of length num_segments, where segment_start_indices[i] is the index of the first element 
+// in the i-th segment and segment_end_indices[i] is the index after the last element in the i-th segment
+// https://nvidia.github.io/cccl/cub/api/structcub_1_1DeviceSegmentedRadixSort.html
+void segmented_sort_pairs_device(void* context, float* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments)
 {
     ContextGuard guard(context);
 
@@ -104,15 +197,20 @@ void radix_sort_pairs_device(void* context, float* keys, int* values, int n)
 	cub::DoubleBuffer<int> d_values(values, values + n);
 
     RadixSortTemp temp;
-    radix_sort_reserve(WP_CURRENT_CONTEXT, n, &temp.mem, &temp.size);
+    segmented_sort_reserve(WP_CURRENT_CONTEXT, n, num_segments, &temp.mem, &temp.size);
 
     // sort
-    check_cuda(cub::DeviceRadixSort::SortPairs(
+    check_cuda(cub::DeviceSegmentedRadixSort::SortPairs(
         temp.mem,
         temp.size,
         d_keys, 
         d_values, 
-        n, 0, 32, 
+        n,
+        num_segments,
+        segment_start_indices,
+        segment_end_indices,
+        0,
+        32,
         (cudaStream_t)cuda_stream_get_current()));
 
 	if (d_keys.Current() != keys)
@@ -122,10 +220,58 @@ void radix_sort_pairs_device(void* context, float* keys, int* values, int n)
 		memcpy_d2d(WP_CURRENT_CONTEXT, values, d_values.Current(), sizeof(int)*n);
 }
 
-void radix_sort_pairs_float_device(uint64_t keys, uint64_t values, int n)
+void segmented_sort_pairs_float_device(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments)
 {
-    radix_sort_pairs_device(
+    segmented_sort_pairs_device(
         WP_CURRENT_CONTEXT,
         reinterpret_cast<float *>(keys),
-        reinterpret_cast<int *>(values), n);
+        reinterpret_cast<int *>(values), n,
+        reinterpret_cast<int *>(segment_start_indices),
+        reinterpret_cast<int *>(segment_end_indices),
+        num_segments);
+}
+
+// segment_indices is an array of length num_segments + 1, where segment_indices[i] is the index of the first element in the i-th segment
+// The end of a segment is given by segment_indices[i+1]
+// https://nvidia.github.io/cccl/cub/api/structcub_1_1DeviceSegmentedSort.html#a-simple-example
+void segmented_sort_pairs_device(void* context, int* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments)
+{
+    ContextGuard guard(context);
+
+    cub::DoubleBuffer<int> d_keys(keys, keys + n);
+	cub::DoubleBuffer<int> d_values(values, values + n);
+
+    RadixSortTemp temp;
+    segmented_sort_reserve(WP_CURRENT_CONTEXT, n, num_segments, &temp.mem, &temp.size);
+
+    // sort
+    check_cuda(cub::DeviceSegmentedRadixSort::SortPairs(
+        temp.mem,
+        temp.size,
+        d_keys, 
+        d_values, 
+        n,
+        num_segments,
+        segment_start_indices,
+        segment_end_indices,
+        0,
+        32,
+        (cudaStream_t)cuda_stream_get_current()));
+
+	if (d_keys.Current() != keys)
+		memcpy_d2d(WP_CURRENT_CONTEXT, keys, d_keys.Current(), sizeof(float)*n);
+
+	if (d_values.Current() != values)
+		memcpy_d2d(WP_CURRENT_CONTEXT, values, d_values.Current(), sizeof(int)*n);
+}
+
+void segmented_sort_pairs_int_device(uint64_t keys, uint64_t values, int n, uint64_t segment_start_indices, uint64_t segment_end_indices, int num_segments)
+{
+    segmented_sort_pairs_device(
+        WP_CURRENT_CONTEXT,
+        reinterpret_cast<int *>(keys),
+        reinterpret_cast<int *>(values), n,
+        reinterpret_cast<int *>(segment_start_indices),
+        reinterpret_cast<int *>(segment_end_indices),
+        num_segments);
 }

warp/native/sort.h

@@ -1,9 +1,18 @@
-/** Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
- * NVIDIA CORPORATION and its licensors retain all intellectual property
- * and proprietary rights in and to this software, related documentation
- * and any modifications thereto.  Any use, reproduction, disclosure or
- * distribution of this software and related documentation without an express
- * license agreement from NVIDIA CORPORATION is strictly prohibited.
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
  */
 
 #pragma once
@@ -13,5 +22,12 @@
 void radix_sort_reserve(void* context, int n, void** mem_out=NULL, size_t* size_out=NULL);
 void radix_sort_pairs_host(int* keys, int* values, int n);
 void radix_sort_pairs_host(float* keys, int* values, int n);
+void radix_sort_pairs_host(int64_t* keys, int* values, int n);
 void radix_sort_pairs_device(void* context, int* keys, int* values, int n);
-void radix_sort_pairs_device(void* context, float* keys, int* values, int n);
+void radix_sort_pairs_device(void* context, float* keys, int* values, int n);
+void radix_sort_pairs_device(void* context, int64_t* keys, int* values, int n);
+
+void segmented_sort_pairs_host(float* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments);
+void segmented_sort_pairs_device(void* context, float* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments);
+void segmented_sort_pairs_host(void* context, int* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments);
+void segmented_sort_pairs_device(void* context, int* keys, int* values, int n, int* segment_start_indices, int* segment_end_indices, int num_segments);