warp-lang 1.5.1__py3-none-manylinux2014_aarch64.whl → 1.6.1__py3-none-manylinux2014_aarch64.whl

warp/native/quat.h

@@ -487,6 +487,37 @@ inline CUDA_CALLABLE void adj_indexref(quat_t<Type>* q, int idx,
     // nop
 }
 
+
+template<typename Type>
+inline CUDA_CALLABLE void augassign_add(quat_t<Type>& q, int idx, Type value)
+{
+    q[idx] += value;
+}
+
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_augassign_add(quat_t<Type>& q, int idx, Type value,
+                                        quat_t<Type>& adj_q, int adj_idx, Type& adj_value)
+{
+    adj_value += adj_q[idx];
+}
+
+
+template<typename Type>
+inline CUDA_CALLABLE void augassign_sub(quat_t<Type>& q, int idx, Type value)
+{
+    q[idx] -= value;
+}
+
+
+template<typename Type>
+inline CUDA_CALLABLE void adj_augassign_sub(quat_t<Type>& q, int idx, Type value,
+                                        quat_t<Type>& adj_q, int adj_idx, Type& adj_value)
+{
+    adj_value -= adj_q[idx];
+}
+
+
 template<typename Type>
 inline CUDA_CALLABLE quat_t<Type> assign(quat_t<Type>& q, int idx, Type value)
 {
@@ -1229,6 +1260,15 @@ inline CUDA_CALLABLE quat_t<Type> quat_identity()
     return quat_t<Type>(Type(0), Type(0), Type(0), Type(1));
 }
 
+template<typename Type>
+CUDA_CALLABLE inline int len(const quat_t<Type>& x)
+{
+    return 4;
+}
 
+template<typename Type>
+CUDA_CALLABLE inline void adj_len(const quat_t<Type>& x, quat_t<Type>& adj_x, const int& adj_ret)
+{
+}
 
 } // namespace wp

warp/native/solid_angle.h

@@ -357,6 +357,13 @@ CUDA_CALLABLE inline void combine_precomputed_solid_angle_props(SolidAngleProps
 	my_data.max_p_dist_sq = length_sq(max(my_data.average_p - my_data.box.lower, my_data.box.upper - my_data.average_p));
 }
 
+CUDA_CALLABLE inline SolidAngleProps combine_precomputed_solid_angle_props(const SolidAngleProps* left_child_data, const SolidAngleProps* right_child_data)
+{
+	SolidAngleProps my_data;
+	combine_precomputed_solid_angle_props(my_data, left_child_data, right_child_data);
+	return my_data;
+}
+
 // Return whether need to
 CUDA_CALLABLE inline bool evaluate_node_solid_angle(const vec3 &query_point, SolidAngleProps *current_data, float &solid_angle, const float accuracy_scale_sq)
 {

warp/native/sort.cpp

@@ -77,12 +77,90 @@ void radix_sort_pairs_host(int* keys, int* values, int n)
 	}	
 }
 
+ //http://stereopsis.com/radix.html
+inline unsigned int radix_float_to_int(float f)
+{
+	unsigned int i = reinterpret_cast<unsigned int&>(f);
+	unsigned int mask = (unsigned int)(-(int)(i >> 31)) | 0x80000000;
+	return i ^ mask;
+}
+
+void radix_sort_pairs_host(float* keys, int* values, int n)
+{
+	static unsigned int tables[2][1 << 16];
+	memset(tables, 0, sizeof(tables));
+		
+	float* auxKeys = keys + n;
+	int* auxValues = values + n;
+
+	// build histograms
+	for (int i=0; i < n; ++i)
+	{
+		const unsigned int k = radix_float_to_int(keys[i]);
+		const unsigned short low = k & 0xffff;
+		const unsigned short high = k >> 16;
+		
+		++tables[0][low];
+		++tables[1][high];
+	}
+	
+	// convert histograms to offset tables in-place
+	unsigned int offlow = 0;
+	unsigned int offhigh = 0;
+	
+	for (int i=0; i < 65536; ++i)
+	{
+		const unsigned int newofflow = offlow + tables[0][i];
+		const unsigned int newoffhigh = offhigh + tables[1][i];
+		
+		tables[0][i] = offlow;
+		tables[1][i] = offhigh;
+		
+		offlow = newofflow;
+		offhigh = newoffhigh;
+	}
+		
+	// pass 1 - sort by low 16 bits
+	for (int i=0; i < n; ++i)
+	{
+		// lookup offset of input
+		const float f = keys[i];
+		const unsigned int k = radix_float_to_int(f);
+		const int v = values[i];
+		const unsigned int b = k & 0xffff;
+		
+		// find offset and increment
+		const unsigned int offset = tables[0][b]++;
+		
+		auxKeys[offset] = f;
+		auxValues[offset] = v;
+	}	
+		
+	// pass 2 - sort by high 16 bits
+	for (int i=0; i < n; ++i)
+	{
+		// lookup offset of input
+		const float f = auxKeys[i];
+		const unsigned int k = radix_float_to_int(f);
+		const int v = auxValues[i];
+
+		const unsigned int b = k >> 16;
+		
+		const unsigned int offset = tables[1][b]++;
+		
+		keys[offset] = f;
+		values[offset] = v;
+	}	
+}
+
 #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_float_device(uint64_t keys, uint64_t values, int n) {}
+
 #endif // !WP_ENABLE_CUDA
 
 
@@ -92,3 +170,10 @@ void radix_sort_pairs_int_host(uint64_t keys, uint64_t values, int n)
         reinterpret_cast<int *>(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);
+}

warp/native/sort.cu

@@ -95,3 +95,37 @@ void radix_sort_pairs_int_device(uint64_t keys, uint64_t values, int n)
         reinterpret_cast<int *>(keys),
         reinterpret_cast<int *>(values), n);
 }
+
+void radix_sort_pairs_device(void* context, float* keys, int* values, int n)
+{
+    ContextGuard guard(context);
+
+    cub::DoubleBuffer<float> 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);
+
+    // sort
+    check_cuda(cub::DeviceRadixSort::SortPairs(
+        temp.mem,
+        temp.size,
+        d_keys, 
+        d_values, 
+        n, 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 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);
+}

warp/native/sort.h

@@ -12,4 +12,6 @@
 
 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_device(void* context, int* keys, int* values, int n);
+void radix_sort_pairs_host(float* 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);

warp/native/spatial.h

@@ -400,6 +400,17 @@ CUDA_CALLABLE inline void adj_lerp(const transform_t<Type>& a, const transform_t
     adj_t += tensordot(b, adj_ret) - tensordot(a, adj_ret);
 }
 
+template<typename Type>
+CUDA_CALLABLE inline int len(const transform_t<Type>& t)
+{
+    return 7;
+}
+
+template<typename Type>
+CUDA_CALLABLE inline void adj_len(const transform_t<Type>& t, transform_t<Type>& adj_t, const int& adj_ret)
+{
+}
+
 template<typename Type>
 using spatial_matrix_t = mat_t<6,6,Type>;