warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.1__py3-none-macosx_10_13_universal2.whl

warp/native/quat.h

@@ -375,12 +375,14 @@ inline CUDA_CALLABLE mat_t<3,3,Type> quat_to_matrix(const quat_t<Type>& q)
     vec_t<3,Type> c2 = quat_rotate(q, vec_t<3,Type>(0.0, 1.0, 0.0));
     vec_t<3,Type> c3 = quat_rotate(q, vec_t<3,Type>(0.0, 0.0, 1.0));
 
-    return mat_t<3,3,Type>(c1, c2, c3);
+    return matrix_from_cols<Type>(c1, c2, c3);
 }
 
-template<typename Type>
-inline CUDA_CALLABLE quat_t<Type> quat_from_matrix(const mat_t<3,3,Type>& m)
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE quat_t<Type> quat_from_matrix(const mat_t<Rows,Cols,Type>& m)
 {
+    static_assert((Rows == 3 && Cols == 3) || (Rows == 4 && Cols == 4), "Non-square matrix");
+
     const Type tr = m.data[0][0] + m.data[1][1] + m.data[2][2];
     Type x, y, z, w, h = Type(0);
 
@@ -498,37 +500,98 @@ inline CUDA_CALLABLE void adj_indexref(quat_t<Type>* q, int idx,
 
 
 template<typename Type>
-inline CUDA_CALLABLE void augassign_add(quat_t<Type>& q, int idx, Type value)
+inline CUDA_CALLABLE void add_inplace(quat_t<Type>& q, int idx, Type value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     q[idx] += value;
 }
 
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_augassign_add(quat_t<Type>& q, int idx, Type value,
+inline CUDA_CALLABLE void adj_add_inplace(quat_t<Type>& q, int idx, Type value,
                                         quat_t<Type>& adj_q, int adj_idx, Type& adj_value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     adj_value += adj_q[idx];
 }
 
 
 template<typename Type>
-inline CUDA_CALLABLE void augassign_sub(quat_t<Type>& q, int idx, Type value)
+inline CUDA_CALLABLE void sub_inplace(quat_t<Type>& q, int idx, Type value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     q[idx] -= value;
 }
 
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_augassign_sub(quat_t<Type>& q, int idx, Type value,
+inline CUDA_CALLABLE void adj_sub_inplace(quat_t<Type>& q, int idx, Type value,
                                         quat_t<Type>& adj_q, int adj_idx, Type& adj_value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     adj_value -= adj_q[idx];
 }
 
 
 template<typename Type>
-inline CUDA_CALLABLE quat_t<Type> assign(quat_t<Type>& q, int idx, Type value)
+inline CUDA_CALLABLE void assign_inplace(quat_t<Type>& q, int idx, Type value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    q[idx] = value;
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_assign_inplace(quat_t<Type>& q, int idx, Type value, quat_t<Type>& adj_q, int& adj_idx, Type& adj_value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx > 3)
+    {
+        printf("quat index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    adj_value += adj_q[idx];
+}
+
+
+template<typename Type>
+inline CUDA_CALLABLE quat_t<Type> assign_copy(quat_t<Type>& q, int idx, Type value)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx > 3)
@@ -544,7 +607,7 @@ inline CUDA_CALLABLE quat_t<Type> assign(quat_t<Type>& q, int idx, Type value)
 }
 
 template<typename Type>
-inline CUDA_CALLABLE void adj_assign(quat_t<Type>& q, int idx, Type value, quat_t<Type>& adj_q, int& adj_idx, Type& adj_value, const quat_t<Type>& adj_ret)
+inline CUDA_CALLABLE void adj_assign_copy(quat_t<Type>& q, int idx, Type value, quat_t<Type>& adj_q, int& adj_idx, Type& adj_value, const quat_t<Type>& adj_ret)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx > 3)
@@ -562,6 +625,7 @@ inline CUDA_CALLABLE void adj_assign(quat_t<Type>& q, int idx, Type value, quat_
     }
 }
 
+
 template<typename Type>
 CUDA_CALLABLE inline quat_t<Type> lerp(const quat_t<Type>& a, const quat_t<Type>& b, Type t)
 {
@@ -1048,9 +1112,11 @@ inline CUDA_CALLABLE void adj_quat_to_matrix(const quat_t<Type>& q, quat_t<Type>
     adj_quat_rotate(q, vec_t<3,Type>(0.0, 0.0, 1.0), adj_q, t, adj_ret.get_col(2));
 }
 
-template<typename Type>
-inline CUDA_CALLABLE void adj_quat_from_matrix(const mat_t<3,3,Type>& m, mat_t<3,3,Type>& adj_m, const quat_t<Type>& adj_ret)
+template<unsigned Rows, unsigned Cols, typename Type>
+inline CUDA_CALLABLE void adj_quat_from_matrix(const mat_t<Rows,Cols,Type>& m, mat_t<Rows,Cols,Type>& adj_m, const quat_t<Type>& adj_ret)
 {
+    static_assert((Rows == 3 && Cols == 3) || (Rows == 4 && Cols == 4), "Non-square matrix");
+
     const Type tr = m.data[0][0] + m.data[1][1] + m.data[2][2];
     Type x, y, z, w, h = Type(0);
 
@@ -1280,4 +1346,26 @@ CUDA_CALLABLE inline void adj_len(const quat_t<Type>& x, quat_t<Type>& adj_x, co
 {
 }
 
+template<typename Type>
+inline CUDA_CALLABLE void expect_near(const quat_t<Type>& actual, const quat_t<Type>& expected, const Type& tolerance)
+{
+    Type diff(0);
+    for(size_t i = 0; i < 4; ++i)
+    {
+        diff = max(diff, abs(actual[i] - expected[i]));
+    }
+    if (diff > tolerance)
+    {
+        printf("Error, expect_near() failed with tolerance "); print(tolerance);
+        printf("\t Expected: "); print(expected);
+        printf("\t Actual: "); print(actual);
+    }
+}
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_expect_near(const quat_t<Type>& actual, const quat_t<Type>& expected, Type tolerance, quat_t<Type>& adj_actual, quat_t<Type>& adj_expected, Type adj_tolerance)
+{
+    // nop
+}
+
 } // namespace wp

warp/native/rand.h

@@ -53,6 +53,9 @@ inline CUDA_CALLABLE uint32 rand_init(int seed, int offset) { return rand_pcg(ui
 inline CUDA_CALLABLE int randi(uint32& state) { state = rand_pcg(state); return int(state); }
 inline CUDA_CALLABLE int randi(uint32& state, int min, int max) { state = rand_pcg(state); return state % (max - min) + min; }
 
+inline CUDA_CALLABLE uint32 randu(uint32& state) { state = rand_pcg(state); return state; }
+inline CUDA_CALLABLE uint32 randu(uint32& state, uint32 min, uint32 max) { state = rand_pcg(state); return state % (max - min) + min; }
+
 /*
  * We want to ensure randf adheres to a uniform distribution over [0,1). The set of all possible float32 (IEEE 754 standard) values is not uniformly distributed however.
  * On the other hand, for a given sign and exponent, the mantissa of the float32 representation is uniformly distributed.
@@ -74,6 +77,9 @@ inline CUDA_CALLABLE void adj_rand_init(int seed, int offset, int& adj_seed, int
 inline CUDA_CALLABLE void adj_randi(uint32& state, uint32& adj_state, float adj_ret) {}
 inline CUDA_CALLABLE void adj_randi(uint32& state, int min, int max, uint32& adj_state, int& adj_min, int& adj_max, float adj_ret) {}
 
+inline CUDA_CALLABLE void adj_randu(uint32& state, uint32& adj_state, float adj_ret) {}
+inline CUDA_CALLABLE void adj_randu(uint32& state, uint32 min, uint32 max, uint32& adj_state, uint32& adj_min, uint32& adj_max, float adj_ret) {}
+
 inline CUDA_CALLABLE void adj_randf(uint32& state, uint32& adj_state, float adj_ret) {}
 inline CUDA_CALLABLE void adj_randf(uint32& state, float min, float max, uint32& adj_state, float& adj_min, float& adj_max, float adj_ret) {}
 

warp/native/sort.cpp

@@ -21,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
@@ -94,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)
@@ -162,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
 
 
@@ -180,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

@@ -36,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
@@ -50,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)
@@ -71,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(
@@ -87,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(
@@ -105,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);
 
@@ -113,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)
@@ -131,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

@@ -22,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);