warp-lang 1.0.2__py3-none-manylinux2014_x86_64.whl → 1.1.0__py3-none-manylinux2014_x86_64.whl

warp/native/hashgrid.cu

@@ -1,77 +1,78 @@
-/** 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.
- */
-
-#include "warp.h"
-#include "cuda_util.h"
-#include "hashgrid.h"
-#include "sort.h"
-
-namespace wp
-{
-
-__global__ void compute_cell_indices(HashGrid grid, const wp::vec3* points, int num_points)
-{
-    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
-
-    if (tid < num_points)
-    {
-        grid.point_cells[tid] = hash_grid_index(grid, points[tid]);
-        grid.point_ids[tid] = tid;
-    }
-}
-
-__global__ void compute_cell_offsets(int* cell_starts, int* cell_ends, const int* point_cells, int num_points)
-{
-    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
-
-    // compute cell start / end
-    if (tid < num_points)
-    {
-		// scan the particle-cell array to find the start and end
-		const int c = point_cells[tid];
-		
-		if (tid == 0)
-			cell_starts[c] = 0;
-		else
-		{
-			const int p = point_cells[tid-1];
-
-			if (c != p)
-			{
-				cell_starts[c] = tid;
-                cell_ends[p] = tid;
-			}
-		}
-
-        if (tid == num_points - 1)
-        {
-            cell_ends[c] = tid + 1;
-        }
-	}    
-}
-
-void hash_grid_rebuild_device(const wp::HashGrid& grid, const wp::vec3* points, int num_points)
-{
-    ContextGuard guard(grid.context);
-
-    wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_cell_indices, num_points, (grid, points, num_points));
-    
-    radix_sort_pairs_device(WP_CURRENT_CONTEXT, grid.point_cells, grid.point_ids, num_points);
-
-    const int num_cells = grid.dim_x * grid.dim_y * grid.dim_z;
-    
-    memset_device(WP_CURRENT_CONTEXT, grid.cell_starts, 0, sizeof(int) * num_cells);    
-    memset_device(WP_CURRENT_CONTEXT, grid.cell_ends, 0, sizeof(int) * num_cells);
-
-    wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_cell_offsets, num_points, (grid.cell_starts, grid.cell_ends, grid.point_cells, num_points));
-}
-
-
-} // namespace wp
-
-
+/** 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.
+ */
+
+#include "warp.h"
+#include "cuda_util.h"
+#include "hashgrid.h"
+#include "sort.h"
+
+namespace wp
+{
+
+__global__ void compute_cell_indices(HashGrid grid, wp::array_t<wp::vec3> points)
+{
+    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
+
+    if (tid < points.shape[0])
+    {
+        const vec3& point = wp::index(points, tid);
+        grid.point_cells[tid] = hash_grid_index(grid, point);
+        grid.point_ids[tid] = tid;
+    }
+}
+
+__global__ void compute_cell_offsets(int* cell_starts, int* cell_ends, const int* point_cells, int num_points)
+{
+    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
+
+    // compute cell start / end
+    if (tid < num_points)
+    {
+		// scan the particle-cell array to find the start and end
+		const int c = point_cells[tid];
+		
+		if (tid == 0)
+			cell_starts[c] = 0;
+		else
+		{
+			const int p = point_cells[tid-1];
+
+			if (c != p)
+			{
+				cell_starts[c] = tid;
+                cell_ends[p] = tid;
+			}
+		}
+
+        if (tid == num_points - 1)
+        {
+            cell_ends[c] = tid + 1;
+        }
+	}    
+}
+
+void hash_grid_rebuild_device(const wp::HashGrid& grid, const wp::array_t<wp::vec3>& points)
+{
+    ContextGuard guard(grid.context);
+
+    int num_points = points.shape[0];
+
+    wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_cell_indices, num_points, (grid, points));
+
+    radix_sort_pairs_device(WP_CURRENT_CONTEXT, grid.point_cells, grid.point_ids, num_points);
+
+    const int num_cells = grid.dim_x * grid.dim_y * grid.dim_z;
+    
+    memset_device(WP_CURRENT_CONTEXT, grid.cell_starts, 0, sizeof(int) * num_cells);    
+    memset_device(WP_CURRENT_CONTEXT, grid.cell_ends, 0, sizeof(int) * num_cells);
+
+    wp_launch_device(WP_CURRENT_CONTEXT, wp::compute_cell_offsets, num_points, (grid.cell_starts, grid.cell_ends, grid.point_cells, num_points));
+}
+
+
+} // namespace wp

warp/native/hashgrid.h

@@ -1,227 +1,227 @@
-/** 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.
- */
-
-#pragma once
-
-namespace wp
-{
-
-struct HashGrid
-{
-    float cell_width;
-    float cell_width_inv;
-
-    int* point_cells{nullptr};   // cell id of a point
-    int* point_ids{nullptr};     // index to original point
-    
-    int* cell_starts{nullptr};   // start index of a range of indices belonging to a cell, dim_x*dim_y*dim_z in length
-    int* cell_ends{nullptr};     // end index of a range of indices belonging to a cell, dim_x*dim_y*dim_z in length
-
-    int dim_x;
-    int dim_y;
-    int dim_z;
-
-    int num_points;
-    int max_points;
-
-    void* context;
-};
-
-// convert a virtual (world) cell coordinate to a physical one
-CUDA_CALLABLE inline int hash_grid_index(const HashGrid& grid, int x, int y, int z)
-{
-    // offset to ensure positive coordinates (means grid dim should be less than 4096^3)
-    const int origin = 1<<20;
-
-    x += origin;
-    y += origin;
-    z += origin;
-
-    assert(0 < x);
-    assert(0 < y);
-    assert(0 < z);
-
-    // clamp in case any particles fall outside the guard region (-10^20 cell index)
-    x = max(0, x);
-    y = max(0, y);
-    z = max(0, z);
-
-    // compute physical cell (assume pow2 grid dims)
-    // int cx = x & (grid.dim_x-1);
-    // int cy = y & (grid.dim_y-1);
-    // int cz = z & (grid.dim_z-1);
-
-    // compute physical cell (arbitrary grid dims)
-    int cx = x%grid.dim_x;
-    int cy = y%grid.dim_y;
-    int cz = z%grid.dim_z;
-
-    return cz*(grid.dim_x*grid.dim_y) + cy*grid.dim_x + cx;
-}
-
-CUDA_CALLABLE inline int hash_grid_index(const HashGrid& grid, const vec3& p)
-{
-    return hash_grid_index(grid, 
-                           int(p[0]*grid.cell_width_inv), 
-                           int(p[1]*grid.cell_width_inv),
-                           int(p[2]*grid.cell_width_inv));
-}
-
-// stores state required to traverse neighboring cells of a point
-struct hash_grid_query_t
-{
-    CUDA_CALLABLE hash_grid_query_t()
-        : x_start(0),
-          y_start(0),
-          z_start(0),
-          x_end(0),
-          y_end(0),
-          z_end(0),
-          x(0),
-          y(0),
-          z(0),
-          cell(0),
-          cell_index(0),
-          cell_end(0),
-          current(0),
-          grid()
-    {}
-
-    // Required for adjoint computations.
-    CUDA_CALLABLE inline hash_grid_query_t& operator+=(const hash_grid_query_t& other)
-    {
-        return *this;
-    }
-
-    int x_start;
-    int y_start;
-    int z_start;
-
-    int x_end;
-    int y_end;
-    int z_end;
-
-    int x;
-    int y;
-    int z;
-
-    int cell;
-    int cell_index;     // offset in the current cell (index into cell_indices)
-    int cell_end;       // index following the end of this cell 
-    
-    int current;        // index of the current iterator value
-
-    HashGrid grid;
-};
-
-
-CUDA_CALLABLE inline hash_grid_query_t hash_grid_query(uint64_t id, wp::vec3 pos, float radius)
-{
-    hash_grid_query_t query;
-
-    query.grid = *(const HashGrid*)(id);
-
-    // convert coordinate to grid
-    query.x_start = int((pos[0]-radius)*query.grid.cell_width_inv);
-    query.y_start = int((pos[1]-radius)*query.grid.cell_width_inv);
-    query.z_start = int((pos[2]-radius)*query.grid.cell_width_inv);
-
-    // do not want to visit any cells more than once, so limit large radius offset to one pass over each dimension
-    query.x_end = min(int((pos[0]+radius)*query.grid.cell_width_inv), query.x_start + query.grid.dim_x-1);
-    query.y_end = min(int((pos[1]+radius)*query.grid.cell_width_inv), query.y_start + query.grid.dim_y-1);
-    query.z_end = min(int((pos[2]+radius)*query.grid.cell_width_inv), query.z_start + query.grid.dim_z-1);
-
-    query.x = query.x_start;
-    query.y = query.y_start;
-    query.z = query.z_start;
-
-    const int cell = hash_grid_index(query.grid, query.x, query.y, query.z);
-    query.cell_index = query.grid.cell_starts[cell];
-    query.cell_end = query.grid.cell_ends[cell];
-
-    return query;
-}
-
-
-CUDA_CALLABLE inline bool hash_grid_query_next(hash_grid_query_t& query, int& index)
-{
-    const HashGrid& grid = query.grid;
-    if (!grid.point_cells)
-        return false;
-
-    while (1)
-    {
-        if (query.cell_index < query.cell_end)
-        {
-            // write output index
-            index = grid.point_ids[query.cell_index++];            
-            return true;
-        }
-        else
-        {
-            query.x++;
-            if (query.x > query.x_end)
-            {
-                query.x = query.x_start;
-                query.y++;
-            }
-
-            if (query.y > query.y_end)
-            {
-                query.y = query.y_start;
-                query.z++;
-            }
-
-            if (query.z > query.z_end)
-            {
-                // finished lookup grid
-                return false;
-            }
-
-            // update cell pointers
-            const int cell = hash_grid_index(grid, query.x, query.y, query.z);
-
-            query.cell_index = grid.cell_starts[cell];
-            query.cell_end = grid.cell_ends[cell];        
-        }
-    }
-}
-
-CUDA_CALLABLE inline int iter_next(hash_grid_query_t& query)
-{
-    return query.current;
-}
-
-CUDA_CALLABLE inline bool iter_cmp(hash_grid_query_t& query)
-{
-    bool finished = hash_grid_query_next(query, query.current);
-    return finished;
-}
-
-CUDA_CALLABLE inline hash_grid_query_t iter_reverse(const hash_grid_query_t& query)
-{
-    // can't reverse grid queries, users should not rely on neighbor ordering
-    return query;
-}
-
-
-
-CUDA_CALLABLE inline int hash_grid_point_id(uint64_t id, int& index)
-{
-    const HashGrid* grid = (const HashGrid*)(id);
-    if (grid->point_ids == nullptr)
-        return -1;
-    return grid->point_ids[index];
-}
-
-CUDA_CALLABLE inline void adj_hash_grid_query(uint64_t id, wp::vec3 pos, float radius, uint64_t& adj_id, wp::vec3& adj_pos, float& adj_radius, hash_grid_query_t& adj_res) {}
-CUDA_CALLABLE inline void adj_hash_grid_query_next(hash_grid_query_t& query, int& index, hash_grid_query_t& adj_query, int& adj_index, bool& adj_res) {}
-CUDA_CALLABLE inline void adj_hash_grid_point_id(uint64_t id, int& index, uint64_t & adj_id, int& adj_index, int& adj_res) {}
-
-
-} // namespace wp
+/** 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.
+ */
+
+#pragma once
+
+namespace wp
+{
+
+struct HashGrid
+{
+    float cell_width;
+    float cell_width_inv;
+
+    int* point_cells{nullptr};   // cell id of a point
+    int* point_ids{nullptr};     // index to original point
+    
+    int* cell_starts{nullptr};   // start index of a range of indices belonging to a cell, dim_x*dim_y*dim_z in length
+    int* cell_ends{nullptr};     // end index of a range of indices belonging to a cell, dim_x*dim_y*dim_z in length
+
+    int dim_x;
+    int dim_y;
+    int dim_z;
+
+    int num_points;
+    int max_points;
+
+    void* context{nullptr};
+};
+
+// convert a virtual (world) cell coordinate to a physical one
+CUDA_CALLABLE inline int hash_grid_index(const HashGrid& grid, int x, int y, int z)
+{
+    // offset to ensure positive coordinates (means grid dim should be less than 4096^3)
+    const int origin = 1<<20;
+
+    x += origin;
+    y += origin;
+    z += origin;
+
+    assert(0 < x);
+    assert(0 < y);
+    assert(0 < z);
+
+    // clamp in case any particles fall outside the guard region (-10^20 cell index)
+    x = max(0, x);
+    y = max(0, y);
+    z = max(0, z);
+
+    // compute physical cell (assume pow2 grid dims)
+    // int cx = x & (grid.dim_x-1);
+    // int cy = y & (grid.dim_y-1);
+    // int cz = z & (grid.dim_z-1);
+
+    // compute physical cell (arbitrary grid dims)
+    int cx = x%grid.dim_x;
+    int cy = y%grid.dim_y;
+    int cz = z%grid.dim_z;
+
+    return cz*(grid.dim_x*grid.dim_y) + cy*grid.dim_x + cx;
+}
+
+CUDA_CALLABLE inline int hash_grid_index(const HashGrid& grid, const vec3& p)
+{
+    return hash_grid_index(grid, 
+                           int(p[0]*grid.cell_width_inv), 
+                           int(p[1]*grid.cell_width_inv),
+                           int(p[2]*grid.cell_width_inv));
+}
+
+// stores state required to traverse neighboring cells of a point
+struct hash_grid_query_t
+{
+    CUDA_CALLABLE hash_grid_query_t()
+        : x_start(0),
+          y_start(0),
+          z_start(0),
+          x_end(0),
+          y_end(0),
+          z_end(0),
+          x(0),
+          y(0),
+          z(0),
+          cell(0),
+          cell_index(0),
+          cell_end(0),
+          current(0),
+          grid()
+    {}
+
+    // Required for adjoint computations.
+    CUDA_CALLABLE inline hash_grid_query_t& operator+=(const hash_grid_query_t& other)
+    {
+        return *this;
+    }
+
+    int x_start;
+    int y_start;
+    int z_start;
+
+    int x_end;
+    int y_end;
+    int z_end;
+
+    int x;
+    int y;
+    int z;
+
+    int cell;
+    int cell_index;     // offset in the current cell (index into cell_indices)
+    int cell_end;       // index following the end of this cell 
+    
+    int current;        // index of the current iterator value
+
+    HashGrid grid;
+};
+
+
+CUDA_CALLABLE inline hash_grid_query_t hash_grid_query(uint64_t id, wp::vec3 pos, float radius)
+{
+    hash_grid_query_t query;
+
+    query.grid = *(const HashGrid*)(id);
+
+    // convert coordinate to grid
+    query.x_start = int((pos[0]-radius)*query.grid.cell_width_inv);
+    query.y_start = int((pos[1]-radius)*query.grid.cell_width_inv);
+    query.z_start = int((pos[2]-radius)*query.grid.cell_width_inv);
+
+    // do not want to visit any cells more than once, so limit large radius offset to one pass over each dimension
+    query.x_end = min(int((pos[0]+radius)*query.grid.cell_width_inv), query.x_start + query.grid.dim_x-1);
+    query.y_end = min(int((pos[1]+radius)*query.grid.cell_width_inv), query.y_start + query.grid.dim_y-1);
+    query.z_end = min(int((pos[2]+radius)*query.grid.cell_width_inv), query.z_start + query.grid.dim_z-1);
+
+    query.x = query.x_start;
+    query.y = query.y_start;
+    query.z = query.z_start;
+
+    const int cell = hash_grid_index(query.grid, query.x, query.y, query.z);
+    query.cell_index = query.grid.cell_starts[cell];
+    query.cell_end = query.grid.cell_ends[cell];
+
+    return query;
+}
+
+
+CUDA_CALLABLE inline bool hash_grid_query_next(hash_grid_query_t& query, int& index)
+{
+    const HashGrid& grid = query.grid;
+    if (!grid.point_cells)
+        return false;
+
+    while (1)
+    {
+        if (query.cell_index < query.cell_end)
+        {
+            // write output index
+            index = grid.point_ids[query.cell_index++];            
+            return true;
+        }
+        else
+        {
+            query.x++;
+            if (query.x > query.x_end)
+            {
+                query.x = query.x_start;
+                query.y++;
+            }
+
+            if (query.y > query.y_end)
+            {
+                query.y = query.y_start;
+                query.z++;
+            }
+
+            if (query.z > query.z_end)
+            {
+                // finished lookup grid
+                return false;
+            }
+
+            // update cell pointers
+            const int cell = hash_grid_index(grid, query.x, query.y, query.z);
+
+            query.cell_index = grid.cell_starts[cell];
+            query.cell_end = grid.cell_ends[cell];        
+        }
+    }
+}
+
+CUDA_CALLABLE inline int iter_next(hash_grid_query_t& query)
+{
+    return query.current;
+}
+
+CUDA_CALLABLE inline bool iter_cmp(hash_grid_query_t& query)
+{
+    bool finished = hash_grid_query_next(query, query.current);
+    return finished;
+}
+
+CUDA_CALLABLE inline hash_grid_query_t iter_reverse(const hash_grid_query_t& query)
+{
+    // can't reverse grid queries, users should not rely on neighbor ordering
+    return query;
+}
+
+
+
+CUDA_CALLABLE inline int hash_grid_point_id(uint64_t id, int& index)
+{
+    const HashGrid* grid = (const HashGrid*)(id);
+    if (grid->point_ids == nullptr)
+        return -1;
+    return grid->point_ids[index];
+}
+
+CUDA_CALLABLE inline void adj_hash_grid_query(uint64_t id, wp::vec3 pos, float radius, uint64_t& adj_id, wp::vec3& adj_pos, float& adj_radius, hash_grid_query_t& adj_res) {}
+CUDA_CALLABLE inline void adj_hash_grid_query_next(hash_grid_query_t& query, int& index, hash_grid_query_t& adj_query, int& adj_index, bool& adj_res) {}
+CUDA_CALLABLE inline void adj_hash_grid_point_id(uint64_t id, int& index, uint64_t & adj_id, int& adj_index, int& adj_res) {}
+
+
+} // namespace wp

warp/native/initializer_array.h

@@ -1,32 +1,32 @@
-/** Copyright (c) 2023 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.
- */
-
-#pragma once
-
-namespace wp {
-
-// wp::initializer_array<> is a simple substitute for std::initializer_list<>
-// which doesn't depend on compiler implementation-specific support. It copies
-// elements by value and only supports array-style indexing.
-template<unsigned Length, typename Type>
-struct initializer_array
-{
-    const Type storage[Length];
-
-    CUDA_CALLABLE const Type operator[](unsigned i)
-    {
-        return storage[i];
-    }
-
-    CUDA_CALLABLE const Type operator[](unsigned i) const
-    {
-        return storage[i];
-    }
-};
-
-}  // namespace wp
+/** Copyright (c) 2023 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.
+ */
+
+#pragma once
+
+namespace wp {
+
+// wp::initializer_array<> is a simple substitute for std::initializer_list<>
+// which doesn't depend on compiler implementation-specific support. It copies
+// elements by value and only supports array-style indexing.
+template<unsigned Length, typename Type>
+struct initializer_array
+{
+    const Type storage[Length];
+
+    CUDA_CALLABLE const Type operator[](unsigned i)
+    {
+        return storage[i];
+    }
+
+    CUDA_CALLABLE const Type operator[](unsigned i) const
+    {
+        return storage[i];
+    }
+};
+
+}  // namespace wp