I implemented a cuda version of converting mesh into o_voxel

o-voxel/o_voxel/_C.pyi

@@ -0,0 +1,54 @@
+from typing import Any, List, Tuple
+
+import torch
+
+def hashmap_insert_cuda(*args: Any, **kwargs: Any) -> Any: ...
+def hashmap_lookup_cuda(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def hashmap_insert_3d_cuda(*args: Any, **kwargs: Any) -> Any: ...
+def hashmap_lookup_3d_cuda(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def hashmap_insert_3d_idx_as_val_cuda(*args: Any, **kwargs: Any) -> Any: ...
+def mesh_to_flexible_dual_grid_cpu(
+    vertices: torch.Tensor,
+    faces: torch.Tensor,
+    voxel_size: List[float],
+    grid_range: List[int],
+    face_weight: float,
+    boundary_weight: float,
+    regularization_weight: float,
+    timing: bool,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
+def mesh_to_flexible_dual_grid_cuda(
+    vertices: torch.Tensor,
+    faces: torch.Tensor,
+    voxel_size: List[float],
+    grid_range: List[int],
+    face_weight: float,
+    boundary_weight: float,
+    regularization_weight: float,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
+def intersect_occ_cpu(
+    triangles: torch.Tensor,
+    voxel_size: List[float],
+    grid_range: List[int],
+) -> torch.Tensor: ...
+def intersect_occ_cuda(
+    triangles: torch.Tensor,
+    voxel_size: List[float],
+    grid_range: List[int],
+) -> torch.Tensor: ...
+def textured_mesh_to_volumetric_attr_cpu(*args: Any, **kwargs: Any) -> Tuple[torch.Tensor, ...]: ...
+def z_order_encode_cpu(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ...
+def z_order_decode_cpu(code: torch.Tensor) -> torch.Tensor: ...
+def hilbert_encode_cpu(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ...
+def hilbert_decode_cpu(code: torch.Tensor) -> torch.Tensor: ...
+def z_order_encode_cuda(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ...
+def z_order_decode_cuda(code: torch.Tensor) -> torch.Tensor: ...
+def hilbert_encode_cuda(x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: ...
+def hilbert_decode_cuda(code: torch.Tensor) -> torch.Tensor: ...
+def encode_sparse_voxel_octree_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def decode_sparse_voxel_octree_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def encode_sparse_voxel_octree_attr_parent_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def decode_sparse_voxel_octree_attr_parent_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def encode_sparse_voxel_octree_attr_neighbor_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def decode_sparse_voxel_octree_attr_neighbor_cpu(*args: Any, **kwargs: Any) -> torch.Tensor: ...
+def rasterize_voxels_cuda(*args: Any, **kwargs: Any) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...

o-voxel/setup.py

@@ -4,6 +4,7 @@
 
 ROOT = os.path.dirname(os.path.abspath(__file__))
 BUILD_TARGET = os.environ.get("BUILD_TARGET", "auto")
+OVOXEL_VERSION = os.environ.get("OVOXEL_VERSION", "0.0.1")
 
 if BUILD_TARGET == "auto":
     if IS_HIP_EXTENSION:
@@ -24,10 +25,12 @@
 
 setup(
     name="o_voxel",
+    version=OVOXEL_VERSION,
+    python_requires=">=3.8",
     packages=[
-        'o_voxel',
-        'o_voxel.convert',
-        'o_voxel.io',
+        "o_voxel",
+        "o_voxel.convert",
+        "o_voxel.io",
     ],
     ext_modules=[
         CUDAExtension(
@@ -38,6 +41,11 @@
                 # Convert functions
                 "src/convert/flexible_dual_grid.cpp",
                 "src/convert/volumetic_attr.cpp",
+                "src/convert/mesh_to_flexible_dual_grid_gpu/intersect_qef.cu",
+                "src/convert/mesh_to_flexible_dual_grid_gpu/face_qef.cu",
+                "src/convert/mesh_to_flexible_dual_grid_gpu/boundary_qef.cu",
+                "src/convert/mesh_to_flexible_dual_grid_gpu/mesh_to_flexible_dual_grid.cu",
+                "src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_octree.cu",
                 ## Serialization functions
                 "src/serialize/api.cu",
                 "src/serialize/hilbert.cu",
@@ -48,7 +56,6 @@
                 "src/io/filter_neighbor.cpp",
                 # Rasterization functions
                 "src/rasterize/rasterize.cu",
-
                 # main
                 "src/ext.cpp",
             ],
@@ -57,11 +64,12 @@
             ],
             extra_compile_args={
                 "cxx": ["-O3", "-std=c++17"],
-                "nvcc": ["-O3","-std=c++17"] + cc_flag,
-            }
+                "nvcc": ["-O3", "-std=c++17"] + cc_flag,
+            },
         )
     ],
-    cmdclass={
-        'build_ext': BuildExtension
-    }
+    package_data={
+        "o_voxel": ["_C.pyi"],
+    },
+    cmdclass={"build_ext": BuildExtension},
 )

o-voxel/src/convert/api.h

@@ -11,37 +11,171 @@
 
 #pragma once
 #include <torch/extension.h>
-
+#include <cstdint>
+#include <vector>
 
 /**
  * Extract flexible dual grid from a triangle mesh.
  *
  * @param vertices: Tensor of shape (N, 3) containing vertex positions.
  * @param faces: Tensor of shape (M, 3) containing triangle vertex indices.
- * @param voxel_size: Tensor of shape (3,) containing the voxel size in each dimension.
- * @param grid_range: Tensor of shape (2, 3) containing the minimum and maximum coordinates of the grid range.
+ * @param voxel_size: Host vector of length 3 containing the voxel size in each dimension.
+ * @param grid_range: Host vector of length 6 containing the minimum and maximum coordinates of the grid range.
  * @param face_weight: Weight for the face edges in the QEM computation.
  * @param boundary_weight: Weight for the boundary edges in the QEM computation.
  * @param regularization_weight: Regularization factor to apply to the QEM matrices.
  * @param timing: Boolean flag to indicate whether to print timing information.
  *
- * @return a tuple ((x, y, z), vertices, intersected, faces) containing the remeshed vertices and the corresponding voxel grid.
+ * @return a tuple (voxels, dual_vertices, intersected) containing the sparse dual grid.
  */
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> mesh_to_flexible_dual_grid_cpu(
-    const torch::Tensor& vertices,
-    const torch::Tensor& faces,
-    const torch::Tensor& voxel_size,
-    const torch::Tensor& grid_range,
+    const torch::Tensor &vertices,
+    const torch::Tensor &faces,
+    const std::vector<float> &voxel_size,
+    const std::vector<int64_t> &grid_range,
     float face_weight,
     float boundary_weight,
     float regularization_weight,
-    bool timing
-);
+    bool timing);
+
+/**
+ * CPU-only occupancy pass for pre-gathered triangles.
+ *
+ * Input triangles are [T, 3, 3] float32 in grid-local coordinates. The function
+ * uses the same triangle/voxel intersection rules as the CPU flexible dual grid
+ * pipeline, but returns only occupied voxel coordinates [N, 3] int32.
+ */
+torch::Tensor intersect_occ_cpu(
+    const torch::Tensor &triangles,
+    const std::vector<float> &voxel_size,
+    const std::vector<int64_t> &grid_range);
+
+namespace o_voxel::fdg
+{
+
+    /**
+     * CUDA occupancy pass for pre-gathered triangles.
+     *
+     * This shares the same active-brick construction used by intersect_qef_cuda,
+     * but stops after compacting occupied voxels. It is useful when only
+     * occupancy is needed and mean/QEF/intersection flags would be wasted work.
+     */
+    torch::Tensor intersect_occ_cuda(
+        const torch::Tensor &triangles,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range);
 
+    /**
+     * CUDA triangle intersection and QEF pass.
+     *
+     * Input triangles are [T, 3, 3] float32 in grid-local coordinates. Large
+     * triangles are split into small scan tasks so many GPU threads can share
+     * their work. The return tuple is:
+     *   0 voxels          [N, 3]  int32
+     *   1 mean_sum        [N, 3]  float32
+     *   2 cnt             [N]     float32
+     *   3 intersected     [N, 3]  bool
+     *   4 qefs            [N, 10] float32, SymQEF10 layout
+     *   5 brick_hash_keys [H]     uint64
+     *   6 brick_hash_vals [H]     uint32
+     *   7 brick_bits      [B, 16] uint32
+     *   8 brick_base      [B]     int64
+     *
+     * The brick hash, bitset, and base tensors are lookup data for later
+     * face_qef_cuda and boundary_qef_cuda calls.
+     */
+    std::tuple<
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor,
+        torch::Tensor>
+    intersect_qef_cuda(
+        const torch::Tensor &triangles,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range);
+
+    /**
+     * In-place CUDA face QEF accumulation.
+     *
+     * Each triangle is paired with the active bricks overlapped by its bounding
+     * box. Threads then inspect only occupied voxels inside those bricks and add
+     * face_weight * face_qef directly into qefs [N, 10].
+     */
+    torch::Tensor face_qef_cuda(
+        const torch::Tensor &triangles,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range,
+        const torch::Tensor &voxels,
+        const torch::Tensor &qefs,
+        float face_weight,
+        const torch::Tensor &brick_hash_keys,
+        const torch::Tensor &brick_hash_vals,
+        const torch::Tensor &brick_bits,
+        const torch::Tensor &brick_base);
+
+    /**
+     * In-place CUDA boundary QEF accumulation.
+     *
+     * Boundaries are [E, 2, 3] float32 segments in grid-local coordinates. Each
+     * thread walks one segment through the voxel grid and adds boundary_weight *
+     * boundary_qef only to voxels found through the active brick lookup.
+     */
+    torch::Tensor boundary_qef_cuda(
+        const torch::Tensor &boundaries,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range,
+        float boundary_weight,
+        const torch::Tensor &voxels,
+        const torch::Tensor &qefs,
+        const torch::Tensor &brick_hash_keys,
+        const torch::Tensor &brick_hash_vals,
+        const torch::Tensor &brick_bits,
+        const torch::Tensor &brick_base);
+
+    /**
+     * Standalone CUDA octree voxelization.
+     *
+     * This is not part of mesh_to_flexible_dual_grid_cuda. It expands octree
+     * jobs from coarse to fine cells on the GPU and returns (prim_ids, voxels),
+     * where prim_ids is [K] int32 and voxels is [K, 3] int32.
+     */
+    std::tuple<torch::Tensor, torch::Tensor>
+    voxelize_mesh_octree_cuda(
+        const torch::Tensor &vertices,
+        const torch::Tensor &faces,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range);
+
+    /**
+     * Full CUDA flexible dual grid pipeline.
+     *
+     * This is a parallel implementation of the CPU pipeline semantics:
+     * gather triangles, compute intersection QEFs, accumulate face and boundary
+     * QEFs in-place, then solve the constrained QEF for each voxel.
+     *
+     * @return (voxels [N, 3] int32, dual_vertices [N, 3] float32,
+     *          intersected [N, 3] bool)
+     */
+    std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
+    mesh_to_flexible_dual_grid_cuda(
+        const torch::Tensor &vertices,
+        const torch::Tensor &faces,
+        const std::vector<float> &voxel_size,
+        const std::vector<int64_t> &grid_range,
+        float face_weight,
+        float boundary_weight,
+        float regularization_weight);
+
+} // namespace o_voxel::fdg
 
 /**
  * Voxelizes a triangle mesh with PBR materials
- * 
+ *
  * @param voxel_size                    [3] tensor containing the size of a voxel
  * @param grid_range                    [6] tensor containing the size of the grid
  * @param vertices                      [N_tri, 3, 3] array containing the triangle vertices
@@ -74,7 +208,7 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> mesh_to_flexible_dual_gr
  * @param normalTextureFilter           list of int indicating the normal texture filter (0: NEAREST, 1: LINEAR)
  * @param normalTextureWrap             list of int indicating the normal texture wrap (0: REPEAT, 1: CLAMP_TO_EDGE, 2: MIRRORED_REPEAT)
  * @param mipLevelOffset                float indicating the mip level offset for texture mipmap
- * 
+ *
  * @return tuple containing:
  *   - coords: tensor of shape [N, 3] containing the voxel coordinates
  *   - out_baseColor: tensor of shape [N, 3] containing the base color of each voxel
@@ -86,37 +220,36 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> mesh_to_flexible_dual_gr
  */
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 textured_mesh_to_volumetric_attr_cpu(
-    const torch::Tensor& voxel_size,
-    const torch::Tensor& grid_range,
-    const torch::Tensor& vertices,
-    const torch::Tensor& normals,
-    const torch::Tensor& uvs,
-    const torch::Tensor& materialIds,
-    const std::vector<torch::Tensor>& baseColorFactor,
-    const std::vector<torch::Tensor>& baseColorTexture,
-    const std::vector<int>& baseColorTextureFilter,
-    const std::vector<int>& baseColorTextureWrap,
-    const std::vector<float>& metallicFactor,
-    const std::vector<torch::Tensor>& metallicTexture,
-    const std::vector<int>& metallicTextureFilter,
-    const std::vector<int>& metallicTextureWrap,
-    const std::vector<float>& roughnessFactor,
-    const std::vector<torch::Tensor>& roughnessTexture,
-    const std::vector<int>& roughnessTextureFilter,
-    const std::vector<int>& roughnessTextureWrap,
-    const std::vector<torch::Tensor>& emissiveFactor,
-    const std::vector<torch::Tensor>& emissiveTexture,
-    const std::vector<int>& emissiveTextureFilter,
-    const std::vector<int>& emissiveTextureWrap,
-    const std::vector<int>& alphaMode,
-    const std::vector<float>& alphaCutoff,
-    const std::vector<float>& alphaFactor,
-    const std::vector<torch::Tensor>& alphaTexture,
-    const std::vector<int>& alphaTextureFilter,
-    const std::vector<int>& alphaTextureWrap,
-    const std::vector<torch::Tensor>& normalTexture,
-    const std::vector<int>& normalTextureFilter,
-    const std::vector<int>& normalTextureWrap,
+    const torch::Tensor &voxel_size,
+    const torch::Tensor &grid_range,
+    const torch::Tensor &vertices,
+    const torch::Tensor &normals,
+    const torch::Tensor &uvs,
+    const torch::Tensor &materialIds,
+    const std::vector<torch::Tensor> &baseColorFactor,
+    const std::vector<torch::Tensor> &baseColorTexture,
+    const std::vector<int> &baseColorTextureFilter,
+    const std::vector<int> &baseColorTextureWrap,
+    const std::vector<float> &metallicFactor,
+    const std::vector<torch::Tensor> &metallicTexture,
+    const std::vector<int> &metallicTextureFilter,
+    const std::vector<int> &metallicTextureWrap,
+    const std::vector<float> &roughnessFactor,
+    const std::vector<torch::Tensor> &roughnessTexture,
+    const std::vector<int> &roughnessTextureFilter,
+    const std::vector<int> &roughnessTextureWrap,
+    const std::vector<torch::Tensor> &emissiveFactor,
+    const std::vector<torch::Tensor> &emissiveTexture,
+    const std::vector<int> &emissiveTextureFilter,
+    const std::vector<int> &emissiveTextureWrap,
+    const std::vector<int> &alphaMode,
+    const std::vector<float> &alphaCutoff,
+    const std::vector<float> &alphaFactor,
+    const std::vector<torch::Tensor> &alphaTexture,
+    const std::vector<int> &alphaTextureFilter,
+    const std::vector<int> &alphaTextureWrap,
+    const std::vector<torch::Tensor> &normalTexture,
+    const std::vector<int> &normalTextureFilter,
+    const std::vector<int> &normalTextureWrap,
     const float mipLevelOffset,
-    const bool timing
-);
+    const bool timing);