warp-lang 1.1.0__py3-none-macosx_10_13_universal2.whl → 1.2.1__py3-none-macosx_10_13_universal2.whl

warp/native/volume.cu

@@ -6,27 +6,53 @@
  * license agreement from NVIDIA CORPORATION is strictly prohibited.
  */
 
-#include "volume.h"
-#include "warp.h"
 #include "cuda_util.h"
+#include "volume_impl.h"
+#include "warp.h"
+
 
+__global__ void volume_get_leaf_coords(const uint32_t leaf_count, pnanovdb_coord_t *leaf_coords,
+                                       const pnanovdb_buf_t buf)
+{
+    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    if (tid < leaf_count)
+    {
+        pnanovdb_leaf_handle_t leaf = wp::volume::get_leaf(buf, tid);
+        leaf_coords[tid] = wp::volume::leaf_origin(buf, leaf);
+    }
+}
 
-__global__ void volume_get_leaf_coords(const uint32_t leaf_count, pnanovdb_coord_t *leaf_coords, const uint64_t first_leaf, const uint32_t leaf_stride)
-{ 
-    static constexpr uint32_t MASK = (1u << 3u) - 1u; // mask for bit operations
+__global__ void volume_get_voxel_coords(const uint32_t voxel_count, pnanovdb_coord_t *voxel_coords,
+                                        const pnanovdb_buf_t buf)
+{
+    const uint32_t leaf_index = blockIdx.x;
+    pnanovdb_leaf_handle_t leaf = wp::volume::get_leaf(buf, leaf_index);
+    pnanovdb_coord_t leaf_coords = wp::volume::leaf_origin(buf, leaf);
 
-    const int tid = blockIdx.x*blockDim.x + threadIdx.x;
+    pnanovdb_coord_t ijk = {
+        int32_t(threadIdx.x) + leaf_coords.x,
+        int32_t(threadIdx.y) + leaf_coords.y,
+        int32_t(threadIdx.z) + leaf_coords.z,
+    };
 
-    if (tid < leaf_count) {
-        leaf_coords[tid] = ((pnanovdb_leaf_t*)(first_leaf + leaf_stride * tid))->bbox_min;
-        leaf_coords[tid].x &= ~MASK;
-        leaf_coords[tid].y &= ~MASK;
-        leaf_coords[tid].z &= ~MASK;
+    const uint64_t index = wp::volume::leaf_voxel_index(buf, leaf_index, ijk);
+    if (index < voxel_count)
+    {
+        voxel_coords[index] = ijk;
     }
 }
 
-void launch_get_leaf_coords(void* context, const uint32_t leaf_count, pnanovdb_coord_t *leaf_coords, const uint64_t first_leaf, const uint32_t leaf_stride)
+void launch_get_leaf_coords(void *context, const uint32_t leaf_count, pnanovdb_coord_t *leaf_coords, pnanovdb_buf_t buf)
+{
+    ContextGuard guard(context);
+    wp_launch_device(WP_CURRENT_CONTEXT, volume_get_leaf_coords, leaf_count, (leaf_count, leaf_coords, buf));
+}
+
+void launch_get_voxel_coords(void *context, const uint32_t leaf_count, const uint32_t voxel_count,
+                             pnanovdb_coord_t *voxel_coords, pnanovdb_buf_t buf)
 {
     ContextGuard guard(context);
-    wp_launch_device(WP_CURRENT_CONTEXT, volume_get_leaf_coords, leaf_count, (leaf_count, leaf_coords, first_leaf, leaf_stride));
+    cudaStream_t stream = (cudaStream_t)cuda_stream_get_current();
+    volume_get_voxel_coords<<<leaf_count, dim3(8, 8, 8), 0, stream>>>(voxel_count, voxel_coords, buf);
 }