diff --git a/o-voxel/o_voxel/_C.pyi b/o-voxel/o_voxel/_C.pyi new file mode 100644 index 00000000..258ec369 --- /dev/null +++ b/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]: ... diff --git a/o-voxel/o_voxel/convert/flexible_dual_grid.py b/o-voxel/o_voxel/convert/flexible_dual_grid.py index 7cf1397e..3a39e84b 100644 --- a/o-voxel/o_voxel/convert/flexible_dual_grid.py +++ b/o-voxel/o_voxel/convert/flexible_dual_grid.py @@ -1,17 +1,27 @@ -from typing import * +from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import torch from .. import _C +FloatArrayLike = Union[float, List[Any], Tuple[Any, ...], np.ndarray, torch.Tensor] +IntArrayLike = Union[int, List[Any], Tuple[Any, ...], np.ndarray, torch.Tensor] +ArrayLike = Union[List[Any], Tuple[Any, ...], np.ndarray, torch.Tensor] +_EDGE_NEIGHBOR_VOXEL_OFFSET: Dict[torch.device, torch.Tensor] = {} +_QUAD_SPLIT_1: Dict[torch.device, torch.Tensor] = {} +_QUAD_SPLIT_2: Dict[torch.device, torch.Tensor] = {} +_QUAD_SPLIT_TRAIN: Dict[torch.device, torch.Tensor] = {} + __all__ = [ "mesh_to_flexible_dual_grid", + "intersect_occ", "flexible_dual_grid_to_mesh", ] def _init_hashmap(grid_size, capacity, device): + """Create the sparse voxel lookup table used when converting a dual grid to a mesh.""" VOL = (grid_size[0] * grid_size[1] * grid_size[2]).item() - + # If the number of elements in the tensor is less than 2^32, use uint32 as the hashmap type, otherwise use uint64. if VOL < 2**32: hashmap_keys = torch.full((capacity,), torch.iinfo(torch.uint32).max, dtype=torch.uint32, device=device) @@ -21,7 +31,7 @@ def _init_hashmap(grid_size, capacity, device): raise ValueError(f"The spatial size is too large to fit in a hashmap. Get volumn {VOL} > 2^64.") hashmap_vals = torch.empty((capacity,), dtype=torch.uint32, device=device) - + return hashmap_keys, hashmap_vals @@ -29,17 +39,23 @@ def _init_hashmap(grid_size, capacity, device): def mesh_to_flexible_dual_grid( vertices: torch.Tensor, faces: torch.Tensor, - voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None, - grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None, - aabb: Union[list, tuple, np.ndarray, torch.Tensor] = None, + voxel_size: Optional[FloatArrayLike] = None, + grid_size: Optional[IntArrayLike] = None, + aabb: Optional[ArrayLike] = None, face_weight: float = 1.0, boundary_weight: float = 1.0, regularization_weight: float = 0.1, timing: bool = False, -) -> Union[torch.Tensor, torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - Voxelize a mesh into a sparse voxel grid. - + Convert a triangle mesh into a sparse flexible dual grid. + + The mesh is first placed in a voxel grid, then intersection, face, and + boundary QEF terms are accumulated for the active voxels. The final dual + vertex in each voxel is found by solving the accumulated QEF. CUDA inputs use + the CUDA implementation; CPU inputs use the CPU implementation with the same + public semantics. + Args: vertices (torch.Tensor): The vertices of the mesh. faces (torch.Tensor): The faces of the mesh. @@ -52,29 +68,44 @@ def mesh_to_flexible_dual_grid( boundary_weight (float): The weight of the boundary term in the QEF when solving the dual vertices. regularization_weight (float): The weight of the regularization term in the QEF when solving the dual vertices. timing (bool): Whether to time the voxelization process. - + Returns: torch.Tensor: The indices of the voxels that are occupied by the mesh. The shape of the tensor is (N, 3), where N is the number of occupied voxels. torch.Tensor: The dual vertices of the mesh. torch.Tensor: The intersected flag of each voxel. """ - - # Load mesh - vertices = vertices.float() - faces = faces.int() - # Voxelize settings + assert isinstance(vertices, torch.Tensor), f"vertices must be a torch.Tensor, but got {type(vertices)}" + assert isinstance(faces, torch.Tensor), f"faces must be a torch.Tensor, but got {type(faces)}" + assert vertices.dim() == 2, f"vertices must be a 2D tensor, but got {vertices.shape}" + assert vertices.size(1) == 3, f"vertices must have 3 columns, but got {vertices.size(1)}" + assert faces.dim() == 2, f"faces must be a 2D tensor, but got {faces.shape}" + assert faces.size(1) == 3, f"faces must have 3 columns, but got {faces.size(1)}" + assert vertices.device == faces.device, "vertices and faces must be on the same device" assert voxel_size is not None or grid_size is not None, "Either voxel_size or grid_size must be provided" + device = vertices.device + vertices = vertices.to(device=device, dtype=torch.float32).contiguous() + faces = faces.to(device=device, dtype=torch.int32).contiguous() + voxel_size_cpu = None + grid_size_cpu = None + aabb_cpu = None + if voxel_size is not None: if isinstance(voxel_size, float): voxel_size = [voxel_size, voxel_size, voxel_size] if isinstance(voxel_size, (list, tuple)): voxel_size = np.array(voxel_size) if isinstance(voxel_size, np.ndarray): - voxel_size = torch.tensor(voxel_size, dtype=torch.float32) + voxel_size_cpu = torch.tensor(voxel_size, dtype=torch.float32) + voxel_size = voxel_size_cpu + else: + assert isinstance(voxel_size, torch.Tensor), f"voxel_size must be a float, list, tuple, np.ndarray, or torch.Tensor, but got {type(voxel_size)}" + assert not voxel_size.is_cuda, "voxel_size Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + voxel_size_cpu = voxel_size.to(dtype=torch.float32).contiguous() assert isinstance(voxel_size, torch.Tensor), f"voxel_size must be a float, list, tuple, np.ndarray, or torch.Tensor, but got {type(voxel_size)}" + voxel_size = voxel_size_cpu.to(device=device).contiguous() assert voxel_size.dim() == 1, f"voxel_size must be a 1D tensor, but got {voxel_size.shape}" assert voxel_size.size(0) == 3, f"voxel_size must have 3 elements, but got {voxel_size.size(0)}" @@ -84,8 +115,14 @@ def mesh_to_flexible_dual_grid( if isinstance(grid_size, (list, tuple)): grid_size = np.array(grid_size) if isinstance(grid_size, np.ndarray): - grid_size = torch.tensor(grid_size, dtype=torch.int32) + grid_size_cpu = torch.tensor(grid_size, dtype=torch.int32) + grid_size = grid_size_cpu + else: + assert isinstance(grid_size, torch.Tensor), f"grid_size must be an int, list, tuple, np.ndarray, or torch.Tensor, but got {type(grid_size)}" + assert not grid_size.is_cuda, "grid_size Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + grid_size_cpu = grid_size.to(dtype=torch.int32).contiguous() assert isinstance(grid_size, torch.Tensor), f"grid_size must be an int, list, tuple, np.ndarray, or torch.Tensor, but got {type(grid_size)}" + grid_size = grid_size_cpu.to(device=device).contiguous() assert grid_size.dim() == 1, f"grid_size must be a 1D tensor, but got {grid_size.shape}" assert grid_size.size(0) == 3, f"grid_size must have 3 elements, but got {grid_size.size(0)}" @@ -93,8 +130,14 @@ def mesh_to_flexible_dual_grid( if isinstance(aabb, (list, tuple)): aabb = np.array(aabb) if isinstance(aabb, np.ndarray): - aabb = torch.tensor(aabb, dtype=torch.float32) + aabb_cpu = torch.tensor(aabb, dtype=torch.float32) + aabb = aabb_cpu + else: + assert isinstance(aabb, torch.Tensor), f"aabb must be a list, tuple, np.ndarray, or torch.Tensor, but got {type(aabb)}" + assert not aabb.is_cuda, "aabb Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + aabb_cpu = aabb.to(dtype=torch.float32).contiguous() assert isinstance(aabb, torch.Tensor), f"aabb must be a list, tuple, np.ndarray, or torch.Tensor, but got {type(aabb)}" + aabb = aabb_cpu.to(device=device).contiguous() assert aabb.dim() == 2, f"aabb must be a 2D tensor, but got {aabb.shape}" assert aabb.size(0) == 2, f"aabb must have 2 rows, but got {aabb.size(0)}" assert aabb.size(1) == 3, f"aabb must have 3 columns, but got {aabb.size(1)}" @@ -103,7 +146,7 @@ def mesh_to_flexible_dual_grid( if aabb is None: min_xyz = vertices.min(dim=0).values max_xyz = vertices.max(dim=0).values - + if voxel_size is not None: padding = torch.ceil((max_xyz - min_xyz) / voxel_size) * voxel_size - (max_xyz - min_xyz) min_xyz -= padding * 0.5 @@ -113,45 +156,198 @@ def mesh_to_flexible_dual_grid( min_xyz -= padding * 0.5 max_xyz += padding * 0.5 - aabb = torch.stack([min_xyz, max_xyz], dim=0).float().cuda() + aabb = torch.stack([min_xyz, max_xyz], dim=0).float().contiguous() # Fill voxel size or grid size - if voxel_size is None: - voxel_size = (aabb[1] - aabb[0]) / grid_size - if grid_size is None: - grid_size = ((aabb[1] - aabb[0]) / voxel_size).round().int() - - # subdivide mesh + if voxel_size_cpu is None: + assert grid_size_cpu is not None + if aabb_cpu is None: + assert not aabb.is_cuda, "CUDA inputs require CPU aabb or explicit voxel_size to avoid implicit CUDA sync" + aabb_cpu = aabb.detach() + voxel_size_cpu = ((aabb_cpu[1] - aabb_cpu[0]) / grid_size_cpu).to(dtype=torch.float32).contiguous() + voxel_size = voxel_size_cpu.to(device=device).contiguous() + if grid_size_cpu is None: + assert voxel_size_cpu is not None + if aabb_cpu is None: + assert not aabb.is_cuda, "CUDA inputs require CPU aabb or explicit grid_size to avoid implicit CUDA sync" + aabb_cpu = aabb.detach() + grid_size_cpu = ((aabb_cpu[1] - aabb_cpu[0]) / voxel_size_cpu).round().to(dtype=torch.int32).contiguous() + grid_size = grid_size_cpu.to(device=device).contiguous() + grid_size = grid_size_cpu.to(device=device).contiguous() + voxel_size_arg = [float(x) for x in voxel_size_cpu.tolist()] + grid_range_arg = [0, 0, 0] + [int(x) for x in grid_size_cpu.tolist()] + + # Shift mesh vertices into grid-local coordinates before calling C++/CUDA. vertices = vertices - aabb[0].reshape(1, 3) - grid_range = torch.stack([torch.zeros_like(grid_size), grid_size], dim=0).int() - - ret = _C.mesh_to_flexible_dual_grid_cpu( - vertices, - faces, - voxel_size, - grid_range, - face_weight, - boundary_weight, - regularization_weight, - timing, - ) - + + if vertices.is_cuda: + ret = _C.mesh_to_flexible_dual_grid_cuda( + vertices, + faces, + voxel_size_arg, + grid_range_arg, + face_weight, + boundary_weight, + regularization_weight, + ) + else: + ret = _C.mesh_to_flexible_dual_grid_cpu( + vertices, + faces, + voxel_size_arg, + grid_range_arg, + face_weight, + boundary_weight, + regularization_weight, + timing, + ) + return ret +@torch.no_grad() +def intersect_occ( + vertices: torch.Tensor, + faces: torch.Tensor, + voxel_size: Optional[FloatArrayLike] = None, + grid_size: Optional[IntArrayLike] = None, + aabb: Optional[ArrayLike] = None, +) -> torch.Tensor: + """ + Return only the voxel coordinates intersected by a triangle mesh. + + This uses the same grid setup as mesh_to_flexible_dual_grid, but stops after + the occupancy stage and does not compute QEFs or dual vertices. It is the + matching public API for users who only need occupied voxels. + """ + + assert isinstance(vertices, torch.Tensor), f"vertices must be a torch.Tensor, but got {type(vertices)}" + assert isinstance(faces, torch.Tensor), f"faces must be a torch.Tensor, but got {type(faces)}" + assert vertices.dim() == 2, f"vertices must be a 2D tensor, but got {vertices.shape}" + assert vertices.size(1) == 3, f"vertices must have 3 columns, but got {vertices.size(1)}" + assert faces.dim() == 2, f"faces must be a 2D tensor, but got {faces.shape}" + assert faces.size(1) == 3, f"faces must have 3 columns, but got {faces.size(1)}" + assert vertices.device == faces.device, "vertices and faces must be on the same device" + assert voxel_size is not None or grid_size is not None, "Either voxel_size or grid_size must be provided" + + device = vertices.device + vertices = vertices.to(device=device, dtype=torch.float32).contiguous() + faces = faces.to(device=device, dtype=torch.int32).contiguous() + voxel_size_cpu = None + grid_size_cpu = None + aabb_cpu = None + + if voxel_size is not None: + if isinstance(voxel_size, float): + voxel_size = [voxel_size, voxel_size, voxel_size] + if isinstance(voxel_size, (list, tuple)): + voxel_size = np.array(voxel_size) + if isinstance(voxel_size, np.ndarray): + voxel_size_cpu = torch.tensor(voxel_size, dtype=torch.float32) + voxel_size = voxel_size_cpu + else: + assert isinstance(voxel_size, torch.Tensor), f"voxel_size must be a float, list, tuple, np.ndarray, or torch.Tensor, but got {type(voxel_size)}" + assert not voxel_size.is_cuda, "voxel_size Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + voxel_size_cpu = voxel_size.to(dtype=torch.float32).contiguous() + assert isinstance(voxel_size, torch.Tensor), f"voxel_size must be a float, list, tuple, np.ndarray, or torch.Tensor, but got {type(voxel_size)}" + voxel_size = voxel_size_cpu.to(device=device).contiguous() + assert voxel_size.dim() == 1, f"voxel_size must be a 1D tensor, but got {voxel_size.shape}" + assert voxel_size.size(0) == 3, f"voxel_size must have 3 elements, but got {voxel_size.size(0)}" + + if grid_size is not None: + if isinstance(grid_size, int): + grid_size = [grid_size, grid_size, grid_size] + if isinstance(grid_size, (list, tuple)): + grid_size = np.array(grid_size) + if isinstance(grid_size, np.ndarray): + grid_size_cpu = torch.tensor(grid_size, dtype=torch.int32) + grid_size = grid_size_cpu + else: + assert isinstance(grid_size, torch.Tensor), f"grid_size must be an int, list, tuple, np.ndarray, or torch.Tensor, but got {type(grid_size)}" + assert not grid_size.is_cuda, "grid_size Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + grid_size_cpu = grid_size.to(dtype=torch.int32).contiguous() + assert isinstance(grid_size, torch.Tensor), f"grid_size must be an int, list, tuple, np.ndarray, or torch.Tensor, but got {type(grid_size)}" + grid_size = grid_size_cpu.to(device=device).contiguous() + assert grid_size.dim() == 1, f"grid_size must be a 1D tensor, but got {grid_size.shape}" + assert grid_size.size(0) == 3, f"grid_size must have 3 elements, but got {grid_size.size(0)}" + + if aabb is not None: + if isinstance(aabb, (list, tuple)): + aabb = np.array(aabb) + if isinstance(aabb, np.ndarray): + aabb_cpu = torch.tensor(aabb, dtype=torch.float32) + aabb = aabb_cpu + else: + assert isinstance(aabb, torch.Tensor), f"aabb must be a list, tuple, np.ndarray, or torch.Tensor, but got {type(aabb)}" + assert not aabb.is_cuda, "aabb Tensor metadata must be on CPU; pass Python values to avoid implicit CUDA sync" + aabb_cpu = aabb.to(dtype=torch.float32).contiguous() + assert isinstance(aabb, torch.Tensor), f"aabb must be a list, tuple, np.ndarray, or torch.Tensor, but got {type(aabb)}" + aabb = aabb_cpu.to(device=device).contiguous() + assert aabb.dim() == 2, f"aabb must be a 2D tensor, but got {aabb.shape}" + assert aabb.size(0) == 2, f"aabb must have 2 rows, but got {aabb.size(0)}" + assert aabb.size(1) == 3, f"aabb must have 3 columns, but got {aabb.size(1)}" + + if aabb is None: + min_xyz = vertices.min(dim=0).values + max_xyz = vertices.max(dim=0).values + + if voxel_size is not None: + padding = torch.ceil((max_xyz - min_xyz) / voxel_size) * voxel_size - (max_xyz - min_xyz) + min_xyz -= padding * 0.5 + max_xyz += padding * 0.5 + if grid_size is not None: + padding = (max_xyz - min_xyz) / (grid_size - 1) + min_xyz -= padding * 0.5 + max_xyz += padding * 0.5 + + aabb = torch.stack([min_xyz, max_xyz], dim=0).float().contiguous() + + if voxel_size_cpu is None: + assert grid_size_cpu is not None + if aabb_cpu is None: + assert not aabb.is_cuda, "CUDA inputs require CPU aabb or explicit voxel_size to avoid implicit CUDA sync" + aabb_cpu = aabb.detach() + voxel_size_cpu = ((aabb_cpu[1] - aabb_cpu[0]) / grid_size_cpu).to(dtype=torch.float32).contiguous() + voxel_size = voxel_size_cpu.to(device=device).contiguous() + if grid_size_cpu is None: + assert voxel_size_cpu is not None + if aabb_cpu is None: + assert not aabb.is_cuda, "CUDA inputs require CPU aabb or explicit grid_size to avoid implicit CUDA sync" + aabb_cpu = aabb.detach() + grid_size_cpu = ((aabb_cpu[1] - aabb_cpu[0]) / voxel_size_cpu).round().to(dtype=torch.int32).contiguous() + grid_size = grid_size_cpu.to(device=device).contiguous() + grid_size = grid_size_cpu.to(device=device).contiguous() + voxel_size_arg = [float(x) for x in voxel_size_cpu.tolist()] + grid_range_arg = [0, 0, 0] + [int(x) for x in grid_size_cpu.tolist()] + + vertices = vertices - aabb[0].reshape(1, 3) + triangles = vertices[faces.to(dtype=torch.long)].contiguous() + + if vertices.is_cuda: + return _C.intersect_occ_cuda(triangles, voxel_size_arg, grid_range_arg) + else: + return _C.intersect_occ_cpu(triangles, voxel_size_arg, grid_range_arg) + + def flexible_dual_grid_to_mesh( coords: torch.Tensor, dual_vertices: torch.Tensor, intersected_flag: torch.Tensor, - split_weight: Union[torch.Tensor, None], - aabb: Union[list, tuple, np.ndarray, torch.Tensor], - voxel_size: Union[float, list, tuple, np.ndarray, torch.Tensor] = None, - grid_size: Union[int, list, tuple, np.ndarray, torch.Tensor] = None, + split_weight: Optional[torch.Tensor], + aabb: ArrayLike, + voxel_size: Optional[FloatArrayLike] = None, + grid_size: Optional[IntArrayLike] = None, train: bool = False, ): """ - Extract mesh from sparse voxel structures using flexible dual grid. - + Extract a triangle mesh from sparse flexible dual grid outputs. + + The function looks up neighboring active voxels around each intersected grid + edge, forms one quad from the four neighboring dual vertices, then splits + each quad into triangles. The sparse voxel lookup is built in PyTorch and the + returned vertices are moved back from grid-local coordinates into the input + AABB. + Args: coords (torch.Tensor): The coordinates of the voxels. dual_vertices (torch.Tensor): The dual vertices. @@ -163,24 +359,32 @@ def flexible_dual_grid_to_mesh( grid_size (int, list, tuple, np.ndarray, torch.Tensor): The size of the grid. NOTE: One of voxel_size and grid_size must be provided. train (bool): Whether to use training mode. - + Returns: vertices (torch.Tensor): The vertices of the mesh. faces (torch.Tensor): The faces of the mesh. """ - # Static variables - if not hasattr(flexible_dual_grid_to_mesh, "edge_neighbor_voxel_offset"): - flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset = torch.tensor([ - [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], # x-axis - [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], # y-axis - [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]], # z-axis - ], dtype=torch.int, device=coords.device).unsqueeze(0) - if not hasattr(flexible_dual_grid_to_mesh, "quad_split_1"): - flexible_dual_grid_to_mesh.quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=coords.device, requires_grad=False) - if not hasattr(flexible_dual_grid_to_mesh, "quad_split_2"): - flexible_dual_grid_to_mesh.quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=coords.device, requires_grad=False) - if not hasattr(flexible_dual_grid_to_mesh, "quad_split_train"): - flexible_dual_grid_to_mesh.quad_split_train = torch.tensor([0, 1, 4, 1, 2, 4, 2, 3, 4, 3, 0, 4], dtype=torch.long, device=coords.device, requires_grad=False) + device = coords.device + if device not in _EDGE_NEIGHBOR_VOXEL_OFFSET: + _EDGE_NEIGHBOR_VOXEL_OFFSET[device] = torch.tensor( + [ + [[0, 0, 0], [0, 0, 1], [0, 1, 1], [0, 1, 0]], # x-axis + [[0, 0, 0], [1, 0, 0], [1, 0, 1], [0, 0, 1]], # y-axis + [[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]], # z-axis + ], + dtype=torch.int, + device=device, + ).unsqueeze(0) + if device not in _QUAD_SPLIT_1: + _QUAD_SPLIT_1[device] = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device, requires_grad=False) + if device not in _QUAD_SPLIT_2: + _QUAD_SPLIT_2[device] = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device, requires_grad=False) + if device not in _QUAD_SPLIT_TRAIN: + _QUAD_SPLIT_TRAIN[device] = torch.tensor([0, 1, 4, 1, 2, 4, 2, 3, 4, 3, 0, 4], dtype=torch.long, device=device, requires_grad=False) + edge_neighbor_voxel_offset = _EDGE_NEIGHBOR_VOXEL_OFFSET[device] + quad_split_1 = _QUAD_SPLIT_1[device] + quad_split_2 = _QUAD_SPLIT_2[device] + quad_split_train = _QUAD_SPLIT_TRAIN[device] # AABB if isinstance(aabb, (list, tuple)): @@ -226,16 +430,13 @@ def flexible_dual_grid_to_mesh( _C.hashmap_insert_3d_idx_as_val_cuda(*hashmap, torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=-1), *grid_size.tolist()) # Find connected voxels - edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + flexible_dual_grid_to_mesh.edge_neighbor_voxel_offset # (N, 3, 4, 3) - connected_voxel = edge_neighbor_voxel[intersected_flag] # (M, 4, 3) + edge_neighbor_voxel = coords.reshape(N, 1, 1, 3) + edge_neighbor_voxel_offset # (N, 3, 4, 3) + connected_voxel = edge_neighbor_voxel[intersected_flag] # (M, 4, 3) M = connected_voxel.shape[0] - connected_voxel_hash_key = torch.cat([ - torch.zeros((M * 4, 1), dtype=torch.int, device=coords.device), - connected_voxel.reshape(-1, 3) - ], dim=1) + connected_voxel_hash_key = torch.cat([torch.zeros((M * 4, 1), dtype=torch.int, device=coords.device), connected_voxel.reshape(-1, 3)], dim=1) connected_voxel_indices = _C.hashmap_lookup_3d_cuda(*hashmap, connected_voxel_hash_key, *grid_size.tolist()).reshape(M, 4).int() - connected_voxel_valid = (connected_voxel_indices != 0xffffffff).all(dim=1) - quad_indices = connected_voxel_indices[connected_voxel_valid].int() # (L, 4) + connected_voxel_valid = (connected_voxel_indices != 0xFFFFFFFF).all(dim=1) + quad_indices = connected_voxel_indices[connected_voxel_valid].int() # (L, 4) L = quad_indices.shape[0] # Construct triangles @@ -243,12 +444,12 @@ def flexible_dual_grid_to_mesh( mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3) if split_weight is None: # if split 1 - atempt_triangles_0 = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_1] + atempt_triangles_0 = quad_indices[:, quad_split_1] normals0 = torch.cross(mesh_vertices[atempt_triangles_0[:, 1]] - mesh_vertices[atempt_triangles_0[:, 0]], mesh_vertices[atempt_triangles_0[:, 2]] - mesh_vertices[atempt_triangles_0[:, 0]]) normals1 = torch.cross(mesh_vertices[atempt_triangles_0[:, 2]] - mesh_vertices[atempt_triangles_0[:, 1]], mesh_vertices[atempt_triangles_0[:, 3]] - mesh_vertices[atempt_triangles_0[:, 1]]) align0 = (normals0 * normals1).sum(dim=1, keepdim=True).abs() # if split 2 - atempt_triangles_1 = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_2] + atempt_triangles_1 = quad_indices[:, quad_split_2] normals0 = torch.cross(mesh_vertices[atempt_triangles_1[:, 1]] - mesh_vertices[atempt_triangles_1[:, 0]], mesh_vertices[atempt_triangles_1[:, 2]] - mesh_vertices[atempt_triangles_1[:, 0]]) normals1 = torch.cross(mesh_vertices[atempt_triangles_1[:, 2]] - mesh_vertices[atempt_triangles_1[:, 1]], mesh_vertices[atempt_triangles_1[:, 3]] - mesh_vertices[atempt_triangles_1[:, 1]]) align1 = (normals0 * normals1).sum(dim=1, keepdim=True).abs() @@ -258,11 +459,7 @@ def flexible_dual_grid_to_mesh( split_weight_ws = split_weight[quad_indices] split_weight_ws_02 = split_weight_ws[:, 0] * split_weight_ws[:, 2] split_weight_ws_13 = split_weight_ws[:, 1] * split_weight_ws[:, 3] - mesh_triangles = torch.where( - split_weight_ws_02 > split_weight_ws_13, - quad_indices[:, flexible_dual_grid_to_mesh.quad_split_1], - quad_indices[:, flexible_dual_grid_to_mesh.quad_split_2] - ).reshape(-1, 3) + mesh_triangles = torch.where(split_weight_ws_02 > split_weight_ws_13, quad_indices[:, quad_split_1], quad_indices[:, quad_split_2]).reshape(-1, 3) else: assert split_weight is not None, "split_weight must be provided in training mode" mesh_vertices = (coords.float() + dual_vertices) * voxel_size + aabb[0].reshape(1, 3) @@ -272,12 +469,9 @@ def flexible_dual_grid_to_mesh( split_weight_ws = split_weight[quad_indices] split_weight_ws_02 = split_weight_ws[:, 0] * split_weight_ws[:, 2] split_weight_ws_13 = split_weight_ws[:, 1] * split_weight_ws[:, 3] - mid_vertices = ( - split_weight_ws_02 * mean_v02 + - split_weight_ws_13 * mean_v13 - ) / (split_weight_ws_02 + split_weight_ws_13) + mid_vertices = (split_weight_ws_02 * mean_v02 + split_weight_ws_13 * mean_v13) / (split_weight_ws_02 + split_weight_ws_13) mesh_vertices = torch.cat([mesh_vertices, mid_vertices], dim=0) - quad_indices = torch.cat([quad_indices, torch.arange(N, N + L, device='cuda').unsqueeze(1)], dim=1) - mesh_triangles = quad_indices[:, flexible_dual_grid_to_mesh.quad_split_train].reshape(-1, 3) - + quad_indices = torch.cat([quad_indices, torch.arange(N, N + L, device="cuda").unsqueeze(1)], dim=1) + mesh_triangles = quad_indices[:, quad_split_train].reshape(-1, 3) + return mesh_vertices, mesh_triangles diff --git a/o-voxel/setup.py b/o-voxel/setup.py index 91cb5cec..15e04e2a 100644 --- a/o-voxel/setup.py +++ b/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}, ) diff --git a/o-voxel/src/convert/api.h b/o-voxel/src/convert/api.h index b70551c8..248d475f 100644 --- a/o-voxel/src/convert/api.h +++ b/o-voxel/src/convert/api.h @@ -11,37 +11,171 @@ #pragma once #include - +#include +#include /** * 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 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 &voxel_size, + const std::vector &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 &voxel_size, + const std::vector &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 &voxel_size, + const std::vector &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 &voxel_size, + const std::vector &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 &voxel_size, + const std::vector &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 &voxel_size, + const std::vector &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 + voxelize_mesh_octree_cuda( + const torch::Tensor &vertices, + const torch::Tensor &faces, + const std::vector &voxel_size, + const std::vector &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 + mesh_to_flexible_dual_grid_cuda( + const torch::Tensor &vertices, + const torch::Tensor &faces, + const std::vector &voxel_size, + const std::vector &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 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 mesh_to_flexible_dual_gr */ std::tuple 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& baseColorFactor, - const std::vector& baseColorTexture, - const std::vector& baseColorTextureFilter, - const std::vector& baseColorTextureWrap, - const std::vector& metallicFactor, - const std::vector& metallicTexture, - const std::vector& metallicTextureFilter, - const std::vector& metallicTextureWrap, - const std::vector& roughnessFactor, - const std::vector& roughnessTexture, - const std::vector& roughnessTextureFilter, - const std::vector& roughnessTextureWrap, - const std::vector& emissiveFactor, - const std::vector& emissiveTexture, - const std::vector& emissiveTextureFilter, - const std::vector& emissiveTextureWrap, - const std::vector& alphaMode, - const std::vector& alphaCutoff, - const std::vector& alphaFactor, - const std::vector& alphaTexture, - const std::vector& alphaTextureFilter, - const std::vector& alphaTextureWrap, - const std::vector& normalTexture, - const std::vector& normalTextureFilter, - const std::vector& 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 &baseColorFactor, + const std::vector &baseColorTexture, + const std::vector &baseColorTextureFilter, + const std::vector &baseColorTextureWrap, + const std::vector &metallicFactor, + const std::vector &metallicTexture, + const std::vector &metallicTextureFilter, + const std::vector &metallicTextureWrap, + const std::vector &roughnessFactor, + const std::vector &roughnessTexture, + const std::vector &roughnessTextureFilter, + const std::vector &roughnessTextureWrap, + const std::vector &emissiveFactor, + const std::vector &emissiveTexture, + const std::vector &emissiveTextureFilter, + const std::vector &emissiveTextureWrap, + const std::vector &alphaMode, + const std::vector &alphaCutoff, + const std::vector &alphaFactor, + const std::vector &alphaTexture, + const std::vector &alphaTextureFilter, + const std::vector &alphaTextureWrap, + const std::vector &normalTexture, + const std::vector &normalTextureFilter, + const std::vector &normalTextureWrap, const float mipLevelOffset, - const bool timing -); + const bool timing); diff --git a/o-voxel/src/convert/flexible_dual_grid.cpp b/o-voxel/src/convert/flexible_dual_grid.cpp index ad89edc0..7e1bb564 100644 --- a/o-voxel/src/convert/flexible_dual_grid.cpp +++ b/o-voxel/src/convert/flexible_dual_grid.cpp @@ -6,76 +6,175 @@ #include "api.h" - constexpr size_t kInvalidIndex = std::numeric_limits::max(); - -struct float3 {float x, y, z; float& operator[](int i) {return (&x)[i];}}; -struct int3 {int x, y, z; int& operator[](int i) {return (&x)[i];}}; -struct int4 {int x, y, z, w; int& operator[](int i) {return (&x)[i];}}; -struct bool3 {bool x, y, z; bool& operator[](int i) {return (&x)[i];}}; - +struct float3 +{ + float x, y, z; + float &operator[](int i) { return (&x)[i]; } +}; +struct int3 +{ + int x, y, z; + int &operator[](int i) { return (&x)[i]; } +}; +struct int4 +{ + int x, y, z, w; + int &operator[](int i) { return (&x)[i]; } +}; +struct bool3 +{ + bool x, y, z; + bool &operator[](int i) { return (&x)[i]; } +}; template -static inline U lerp(const T& a, const T& b, const T& t, const U& val_a, const U& val_b) { - if (a == b) return val_a; // Avoid divide by zero +static inline U lerp(const T &a, const T &b, const T &t, const U &val_a, const U &val_b) +{ + if (a == b) + return val_a; // Avoid divide by zero T alpha = (t - a) / (b - a); return (1 - alpha) * val_a + alpha * val_b; } - template -static auto get_or_default(const Map& map, const Key& key, const Default& default_val) -> typename Map::mapped_type { +static auto get_or_default(const Map &map, const Key &key, const Default &default_val) -> typename Map::mapped_type +{ auto it = map.find(key); return (it != map.end()) ? it->second : default_val; } - // 3D voxel coordinate -struct VoxelCoord { +struct VoxelCoord +{ int x, y, z; - int& operator[](int i) { + int &operator[](int i) + { return (&x)[i]; } - bool operator==(const VoxelCoord& other) const { + bool operator==(const VoxelCoord &other) const + { return x == other.x && y == other.y && z == other.z; } }; // Hash function for VoxelCoord to use in unordered_map -namespace std { -template <> -struct hash { - size_t operator()(const VoxelCoord& v) const { - const std::size_t p1 = 73856093; - const std::size_t p2 = 19349663; - const std::size_t p3 = 83492791; - return (std::size_t)(v.x) * p1 ^ (std::size_t)(v.y) * p2 ^ (std::size_t)(v.z) * p3; - } -}; +namespace std +{ + template <> + struct hash + { + size_t operator()(const VoxelCoord &v) const + { + const std::size_t p1 = 73856093; + const std::size_t p2 = 19349663; + const std::size_t p3 = 83492791; + return (std::size_t)(v.x) * p1 ^ (std::size_t)(v.y) * p2 ^ (std::size_t)(v.z) * p3; + } + }; } +void intersect_occ( + const Eigen::Vector3f &voxel_size, + const Eigen::Vector3i &grid_min, + const Eigen::Vector3i &grid_max, + const std::vector &triangles, + std::unordered_map &hash_table, + std::vector &voxels) +{ + const size_t N_tri = triangles.size() / 3; + + for (size_t i = 0; i < N_tri; ++i) + { + const Eigen::Vector3f &v0 = triangles[i * 3 + 0]; + const Eigen::Vector3f &v1 = triangles[i * 3 + 1]; + const Eigen::Vector3f &v2 = triangles[i * 3 + 2]; + + auto scan_line_fill = [&](const int ax2) + { + int ax0 = (ax2 + 1) % 3; + int ax1 = (ax2 + 2) % 3; + + std::array t = { + Eigen::Vector3d(v0[ax0], v0[ax1], v0[ax2]), + Eigen::Vector3d(v1[ax0], v1[ax1], v1[ax2]), + Eigen::Vector3d(v2[ax0], v2[ax1], v2[ax2])}; + std::sort(t.begin(), t.end(), [](const Eigen::Vector3d &a, const Eigen::Vector3d &b) + { return a.y() < b.y(); }); + + int start = std::clamp(int(t[0].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + int mid = std::clamp(int(t[1].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + int end = std::clamp(int(t[2].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + + auto scan_line_half = [&](const int row_start, const int row_end, const Eigen::Vector3d t0, const Eigen::Vector3d t1, const Eigen::Vector3d t2) + { + for (int y_idx = row_start; y_idx < row_end; ++y_idx) + { + double y = (y_idx + 1) * voxel_size[ax1]; + Eigen::Vector2d t3 = lerp(t0.y(), t1.y(), y, Eigen::Vector2d(t0.x(), t0.z()), Eigen::Vector2d(t1.x(), t1.z())); + Eigen::Vector2d t4 = lerp(t0.y(), t2.y(), y, Eigen::Vector2d(t0.x(), t0.z()), Eigen::Vector2d(t2.x(), t2.z())); + if (t3.x() > t4.x()) + std::swap(t3, t4); + int line_start = std::clamp(int(t3.x() / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + int line_end = std::clamp(int(t4.x() / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + for (int x_idx = line_start; x_idx < line_end; ++x_idx) + { + double x = (x_idx + 1) * voxel_size[ax0]; + double z = lerp(t3.x(), t4.x(), x, t3.y(), t4.y()); + int z_idx = int(z / voxel_size[ax2]); + if (z_idx >= grid_min[ax2] && z_idx < grid_max[ax2]) + { + for (int dx = 0; dx < 2; ++dx) + { + for (int dy = 0; dy < 2; ++dy) + { + VoxelCoord coord; + coord[ax0] = x_idx + dx; + coord[ax1] = y_idx + dy; + coord[ax2] = z_idx; + if (hash_table.find(coord) == hash_table.end()) + { + hash_table[coord] = voxels.size(); + voxels.push_back({coord.x, coord.y, coord.z}); + } + } + } + } + } + } + }; + scan_line_half(start, mid, t[0], t[1], t[2]); + scan_line_half(mid, end, t[2], t[1], t[0]); + }; + scan_line_fill(0); + scan_line_fill(1); + scan_line_fill(2); + } +} void intersect_qef( - const Eigen::Vector3f& voxel_size, - const Eigen::Vector3i& grid_min, - const Eigen::Vector3i& grid_max, - const std::vector& triangles, // 3 vertices per triangle - std::unordered_map& hash_table, // Hash table for voxel lookup - std::vector& voxels, // Output: Voxel coordinates - std::vector& means, // Output: Mean vertex positions for each voxel - std::vector& cnt, // Output: Number of intersections for each voxel - std::vector& intersected, // Output: Whether edge of voxel intersects with triangle - std::vector& qefs // Output: QEF matrices for each voxel -) { + const Eigen::Vector3f &voxel_size, + const Eigen::Vector3i &grid_min, + const Eigen::Vector3i &grid_max, + const std::vector &triangles, // 3 vertices per triangle + std::unordered_map &hash_table, // Hash table for voxel lookup + std::vector &voxels, // Output: Voxel coordinates + std::vector &means, // Output: Mean vertex positions for each voxel + std::vector &cnt, // Output: Number of intersections for each voxel + std::vector &intersected, // Output: Whether edge of voxel intersects with triangle + std::vector &qefs // Output: QEF matrices for each voxel +) +{ const size_t N_tri = triangles.size() / 3; - for (size_t i = 0; i < N_tri; ++i) { - const Eigen::Vector3f& v0 = triangles[i * 3 + 0]; - const Eigen::Vector3f& v1 = triangles[i * 3 + 1]; - const Eigen::Vector3f& v2 = triangles[i * 3 + 2]; + for (size_t i = 0; i < N_tri; ++i) + { + const Eigen::Vector3f &v0 = triangles[i * 3 + 0]; + const Eigen::Vector3f &v1 = triangles[i * 3 + 1]; + const Eigen::Vector3f &v2 = triangles[i * 3 + 2]; // Compute edge vectors and face normal Eigen::Vector3f e0 = v1 - v0; @@ -93,7 +192,8 @@ void intersect_qef( | / t2 */ - auto scan_line_fill = [&] (const int ax2) { + auto scan_line_fill = [&](const int ax2) + { int ax0 = (ax2 + 1) % 3; int ax1 = (ax2 + 2) % 3; @@ -101,44 +201,56 @@ void intersect_qef( std::array t = { Eigen::Vector3d(v0[ax0], v0[ax1], v0[ax2]), Eigen::Vector3d(v1[ax0], v1[ax1], v1[ax2]), - Eigen::Vector3d(v2[ax0], v2[ax1], v2[ax2]) - }; - std::sort(t.begin(), t.end(), [](const Eigen::Vector3d& a, const Eigen::Vector3d& b) { return a.y() < b.y(); }); + Eigen::Vector3d(v2[ax0], v2[ax1], v2[ax2])}; + std::sort(t.begin(), t.end(), [](const Eigen::Vector3d &a, const Eigen::Vector3d &b) + { return a.y() < b.y(); }); // Scan-line algorithm int start = std::clamp(int(t[0].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); int mid = std::clamp(int(t[1].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); int end = std::clamp(int(t[2].y() / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); - auto scan_line_half = [&] (const int row_start, const int row_end, const Eigen::Vector3d t0, const Eigen::Vector3d t1, const Eigen::Vector3d t2) { - /* - t0 - | \ - t3-t4 - | \ - t1---t2 - */ - for (int y_idx = row_start; y_idx < row_end; ++y_idx) { + auto scan_line_half = [&](const int row_start, const int row_end, const Eigen::Vector3d t0, const Eigen::Vector3d t1, const Eigen::Vector3d t2) + { + /* + t0 + | \ + t3-t4 + | \ + t1---t2 + */ + for (int y_idx = row_start; y_idx < row_end; ++y_idx) + { double y = (y_idx + 1) * voxel_size[ax1]; Eigen::Vector2d t3 = lerp(t0.y(), t1.y(), y, Eigen::Vector2d(t0.x(), t0.z()), Eigen::Vector2d(t1.x(), t1.z())); Eigen::Vector2d t4 = lerp(t0.y(), t2.y(), y, Eigen::Vector2d(t0.x(), t0.z()), Eigen::Vector2d(t2.x(), t2.z())); - if (t3.x() > t4.x()) std::swap(t3, t4); + if (t3.x() > t4.x()) + std::swap(t3, t4); int line_start = std::clamp(int(t3.x() / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); int line_end = std::clamp(int(t4.x() / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); - for (int x_idx = line_start; x_idx < line_end; ++x_idx) { + for (int x_idx = line_start; x_idx < line_end; ++x_idx) + { double x = (x_idx + 1) * voxel_size[ax0]; double z = lerp(t3.x(), t4.x(), x, t3.y(), t4.y()); int z_idx = int(z / voxel_size[ax2]); - if (z_idx >= grid_min[ax2] && z_idx < grid_max[ax2]) { + if (z_idx >= grid_min[ax2] && z_idx < grid_max[ax2]) + { // For 4-connected voxels - for (int dx = 0; dx < 2; ++dx) { - for (int dy = 0; dy < 2; ++dy) { + for (int dx = 0; dx < 2; ++dx) + { + for (int dy = 0; dy < 2; ++dy) + { VoxelCoord coord; - coord[ax0] = x_idx + dx; coord[ax1] = y_idx + dy; coord[ax2] = z_idx; + coord[ax0] = x_idx + dx; + coord[ax1] = y_idx + dy; + coord[ax2] = z_idx; Eigen::Vector3d intersect; - intersect[ax0] = x; intersect[ax1] = y; intersect[ax2] = z; + intersect[ax0] = x; + intersect[ax1] = y; + intersect[ax2] = z; auto kv = hash_table.find(coord); - if (kv == hash_table.end()) { + if (kv == hash_table.end()) + { hash_table[coord] = voxels.size(); voxels.push_back({coord.x, coord.y, coord.z}); means.push_back(intersect.cast()); @@ -148,7 +260,8 @@ void intersect_qef( if (dx == 0 && dy == 0) intersected.back()[ax2] = true; } - else { + else + { auto i = kv->second; means[i] += intersect.cast(); cnt[i] += 1; @@ -163,7 +276,7 @@ void intersect_qef( } }; scan_line_half(start, mid, t[0], t[1], t[2]); - scan_line_half(mid, end, t[2], t[1], t[0]); + scan_line_half(mid, end, t[2], t[1], t[0]); }; scan_line_fill(0); scan_line_fill(1); @@ -171,21 +284,22 @@ void intersect_qef( } } - void face_qef( - const Eigen::Vector3f& voxel_size, - const Eigen::Vector3i& grid_min, - const Eigen::Vector3i& grid_max, - const std::vector& triangles, // 3 vertices per triangle - std::unordered_map& hash_table, // Hash table for voxel lookup - std::vector& qefs // Output: QEF matrices for each voxel -) { + const Eigen::Vector3f &voxel_size, + const Eigen::Vector3i &grid_min, + const Eigen::Vector3i &grid_max, + const std::vector &triangles, // 3 vertices per triangle + std::unordered_map &hash_table, // Hash table for voxel lookup + std::vector &qefs // Output: QEF matrices for each voxel +) +{ const size_t N_tri = triangles.size() / 3; - for (size_t i = 0; i < N_tri; ++i) { - const Eigen::Vector3f& v0 = triangles[i * 3 + 0]; - const Eigen::Vector3f& v1 = triangles[i * 3 + 1]; - const Eigen::Vector3f& v2 = triangles[i * 3 + 2]; + for (size_t i = 0; i < N_tri; ++i) + { + const Eigen::Vector3f &v0 = triangles[i * 3 + 0]; + const Eigen::Vector3f &v1 = triangles[i * 3 + 1]; + const Eigen::Vector3f &v2 = triangles[i * 3 + 2]; // Compute edge vectors and face normal Eigen::Vector3f e0 = v1 - v0; @@ -211,8 +325,7 @@ void face_qef( Eigen::Vector3f c( n.x() > 0.0f ? voxel_size.x() : 0.0f, n.y() > 0.0f ? voxel_size.y() : 0.0f, - n.z() > 0.0f ? voxel_size.z() : 0.0f - ); + n.z() > 0.0f ? voxel_size.z() : 0.0f); float d1 = n.dot(c - v0); float d2 = n.dot(voxel_size - c - v0); @@ -247,38 +360,52 @@ void face_qef( float d_zx_e2 = -n_zx_e2.dot(Eigen::Vector2f(v2.z(), v2.x())) + n_zx_e2.cwiseMax(0.0f).dot(Eigen::Vector2f(voxel_size.z(), voxel_size.x())); // Loop over candidate voxels inside bounding box - for (int z = bb_min.z(); z < bb_max.z(); ++z) { - for (int y = bb_min.y(); y < bb_max.y(); ++y) { - for (int x = bb_min.x(); x < bb_max.x(); ++x) { + for (int z = bb_min.z(); z < bb_max.z(); ++z) + { + for (int y = bb_min.y(); y < bb_max.y(); ++y) + { + for (int x = bb_min.x(); x < bb_max.x(); ++x) + { // Voxel center Eigen::Vector3f p = voxel_size.cwiseProduct(Eigen::Vector3f(x, y, z)); // Plane through box test float nDOTp = n.dot(p); - if (((nDOTp + d1) * (nDOTp + d2)) > 0.0f) continue; + if (((nDOTp + d1) * (nDOTp + d2)) > 0.0f) + continue; // XY projection test Eigen::Vector2f p_xy(p.x(), p.y()); - if (n_xy_e0.dot(p_xy) + d_xy_e0 < 0) continue; - if (n_xy_e1.dot(p_xy) + d_xy_e1 < 0) continue; - if (n_xy_e2.dot(p_xy) + d_xy_e2 < 0) continue; + if (n_xy_e0.dot(p_xy) + d_xy_e0 < 0) + continue; + if (n_xy_e1.dot(p_xy) + d_xy_e1 < 0) + continue; + if (n_xy_e2.dot(p_xy) + d_xy_e2 < 0) + continue; // YZ projection test Eigen::Vector2f p_yz(p.y(), p.z()); - if (n_yz_e0.dot(p_yz) + d_yz_e0 < 0) continue; - if (n_yz_e1.dot(p_yz) + d_yz_e1 < 0) continue; - if (n_yz_e2.dot(p_yz) + d_yz_e2 < 0) continue; + if (n_yz_e0.dot(p_yz) + d_yz_e0 < 0) + continue; + if (n_yz_e1.dot(p_yz) + d_yz_e1 < 0) + continue; + if (n_yz_e2.dot(p_yz) + d_yz_e2 < 0) + continue; // ZX projection test Eigen::Vector2f p_zx(p.z(), p.x()); - if (n_zx_e0.dot(p_zx) + d_zx_e0 < 0) continue; - if (n_zx_e1.dot(p_zx) + d_zx_e1 < 0) continue; - if (n_zx_e2.dot(p_zx) + d_zx_e2 < 0) continue; + if (n_zx_e0.dot(p_zx) + d_zx_e0 < 0) + continue; + if (n_zx_e1.dot(p_zx) + d_zx_e1 < 0) + continue; + if (n_zx_e2.dot(p_zx) + d_zx_e2 < 0) + continue; // Passed all tests — mark voxel auto coord = VoxelCoord{x, y, z}; auto kv = hash_table.find(coord); - if (kv != hash_table.end()) { + if (kv != hash_table.end()) + { qefs[kv->second] += Q; } } @@ -287,25 +414,27 @@ void face_qef( } } - void boundry_qef( - const Eigen::Vector3f& voxel_size, - const Eigen::Vector3i& grid_min, - const Eigen::Vector3i& grid_max, - const std::vector& boundries, // 2 vertices per segment - const float boundary_weight, // Weight for boundary edges - std::unordered_map& hash_table, // Hash table for voxel lookup - std::vector& qefs // Output: QEF matrices for each voxel -) { - for (size_t i = 0; i < boundries.size() / 2; ++i) { - const Eigen::Vector3f& v0 = boundries[i * 2 + 0]; - const Eigen::Vector3f& v1 = boundries[i * 2 + 1]; + const Eigen::Vector3f &voxel_size, + const Eigen::Vector3i &grid_min, + const Eigen::Vector3i &grid_max, + const std::vector &boundries, // 2 vertices per segment + const float boundary_weight, // Weight for boundary edges + std::unordered_map &hash_table, // Hash table for voxel lookup + std::vector &qefs // Output: QEF matrices for each voxel +) +{ + for (size_t i = 0; i < boundries.size() / 2; ++i) + { + const Eigen::Vector3f &v0 = boundries[i * 2 + 0]; + const Eigen::Vector3f &v1 = boundries[i * 2 + 1]; // Calculate the QEF for the edge (boundary) defined by v0 and v1 Eigen::Vector3d dir(v1.x() - v0.x(), v1.y() - v0.y(), v1.z() - v0.z()); double segment_length = dir.norm(); - if (segment_length < 1e-6d) continue; // Skip degenerate edges (zero-length) - dir.normalize(); // unit direction vector + if (segment_length < 1e-6d) + continue; // Skip degenerate edges (zero-length) + dir.normalize(); // unit direction vector // Projection matrix orthogonal to the direction: I - d d^T Eigen::Matrix3f A = Eigen::Matrix3f::Identity() - (dir * dir.transpose()).cast(); @@ -333,11 +462,15 @@ void boundry_qef( Eigen::Vector3i step = (dir.array() > 0).select(Eigen::Vector3i(1, 1, 1), Eigen::Vector3i(-1, -1, -1)); Eigen::Vector3d tMax, tDelta; - for (int axis = 0; axis < 3; ++axis) { - if (dir[axis] == 0.0d) { + for (int axis = 0; axis < 3; ++axis) + { + if (dir[axis] == 0.0d) + { tMax[axis] = std::numeric_limits::infinity(); tDelta[axis] = std::numeric_limits::infinity(); - } else { + } + else + { float voxel_border = voxel_size[axis] * (v0_voxel[axis] + (step[axis] > 0 ? 1 : 0)); tMax[axis] = (voxel_border - v0[axis]) / dir[axis]; tDelta[axis] = voxel_size[axis] / std::abs(dir[axis]); @@ -352,15 +485,20 @@ void boundry_qef( voxels.push_back({current.x(), current.y(), current.z()}); // Traverse the voxels - while (true) { + while (true) + { int axis; - if (tMax.x() < tMax.y()) { + if (tMax.x() < tMax.y()) + { axis = (tMax.x() < tMax.z()) ? 0 : 2; - } else { + } + else + { axis = (tMax.y() < tMax.z()) ? 1 : 2; } - if (tMax[axis] > segment_length) break; + if (tMax[axis] > segment_length) + break; current[axis] += step[axis]; tMax[axis] += tDelta[axis]; @@ -369,12 +507,15 @@ void boundry_qef( } // Accumulate QEF for each voxel passed through - for (const auto& coord : voxels) { + for (const auto &coord : voxels) + { // Make sure the voxel is within bounds if ((coord.x < grid_min.x() || coord.x >= grid_max.x()) || (coord.y < grid_min.y() || coord.y >= grid_max.y()) || - (coord.z < grid_min.z() || coord.z >= grid_max.z())) continue; - if (!hash_table.count(coord)) continue; // Skip if voxel not in hash table + (coord.z < grid_min.z() || coord.z >= grid_max.z())) + continue; + if (!hash_table.count(coord)) + continue; // Skip if voxel not in hash table // Accumulate the QEF for this voxel qefs[hash_table[coord]] += boundary_weight * Q; // Scale by boundary weight @@ -382,11 +523,10 @@ void boundry_qef( } } - std::array quad_to_2tri( - const std::vector& vertices, - const int4& quad_indices -) { + const std::vector &vertices, + const int4 &quad_indices) +{ int ia = quad_indices.x; int ib = quad_indices.y; int ic = quad_indices.z; @@ -407,22 +547,25 @@ std::array quad_to_2tri( Eigen::Vector3f n_bcd = (c - b).cross(d - b).normalized(); float angle_bd = std::acos(std::clamp(n_abd.dot(n_bcd), -1.0f, 1.0f)); - if (angle_ac <= angle_bd) { + if (angle_ac <= angle_bd) + { return {int3{ia, ib, ic}, int3{ia, ic, id}}; - } else { + } + else + { return {int3{ia, ib, id}, int3{ib, ic, id}}; } } - void face_from_dual_vertices( - const std::unordered_map& hash_table, - const std::vector& voxels, - const std::vector& dual_vertices, - const std::vector& intersected, - std::vector& face_indices -) { - for (int i = 0; i < dual_vertices.size(); ++i) { + const std::unordered_map &hash_table, + const std::vector &voxels, + const std::vector &dual_vertices, + const std::vector &intersected, + std::vector &face_indices) +{ + for (int i = 0; i < dual_vertices.size(); ++i) + { int3 coord = voxels[i]; bool3 is_intersected = intersected[i]; @@ -433,23 +576,25 @@ void face_from_dual_vertices( get_or_default(hash_table, VoxelCoord{coord.x + 1, coord.y + 1, coord.z}, kInvalidIndex), get_or_default(hash_table, VoxelCoord{coord.x, coord.y, coord.z + 1}, kInvalidIndex), get_or_default(hash_table, VoxelCoord{coord.x + 1, coord.y, coord.z + 1}, kInvalidIndex), - get_or_default(hash_table, VoxelCoord{coord.x, coord.y + 1, coord.z + 1}, kInvalidIndex) - }; + get_or_default(hash_table, VoxelCoord{coord.x, coord.y + 1, coord.z + 1}, kInvalidIndex)}; // xy-plane - if (is_intersected[2] && neigh_indices[0] != kInvalidIndex && neigh_indices[1] != kInvalidIndex && neigh_indices[2] != kInvalidIndex) { + if (is_intersected[2] && neigh_indices[0] != kInvalidIndex && neigh_indices[1] != kInvalidIndex && neigh_indices[2] != kInvalidIndex) + { int4 quad_indices{i, neigh_indices[0], neigh_indices[2], neigh_indices[1]}; auto tri_indices = quad_to_2tri(dual_vertices, quad_indices); face_indices.insert(face_indices.end(), tri_indices.begin(), tri_indices.end()); } // yz-plane - if (is_intersected[0] && neigh_indices[1] != kInvalidIndex && neigh_indices[3] != kInvalidIndex && neigh_indices[5] != kInvalidIndex) { + if (is_intersected[0] && neigh_indices[1] != kInvalidIndex && neigh_indices[3] != kInvalidIndex && neigh_indices[5] != kInvalidIndex) + { int4 quad_indices{i, neigh_indices[1], neigh_indices[5], neigh_indices[3]}; auto tri_indices = quad_to_2tri(dual_vertices, quad_indices); face_indices.insert(face_indices.end(), tri_indices.begin(), tri_indices.end()); } // xz-plane - if (is_intersected[1] && neigh_indices[0] != kInvalidIndex && neigh_indices[3] != kInvalidIndex && neigh_indices[4] != kInvalidIndex) { + if (is_intersected[1] && neigh_indices[0] != kInvalidIndex && neigh_indices[3] != kInvalidIndex && neigh_indices[4] != kInvalidIndex) + { int4 quad_indices{i, neigh_indices[0], neigh_indices[4], neigh_indices[3]}; auto tri_indices = quad_to_2tri(dual_vertices, quad_indices); face_indices.insert(face_indices.end(), tri_indices.begin(), tri_indices.end()); @@ -457,13 +602,54 @@ void face_from_dual_vertices( } } +torch::Tensor intersect_occ_cpu( + const torch::Tensor &triangles, + const std::vector &voxel_size, + const std::vector &grid_range) +{ + TORCH_CHECK(!triangles.is_cuda(), "triangles must be a CPU tensor"); + + const int64_t N_tri = triangles.size(0); + const float *t_ptr = triangles.data_ptr(); + + Eigen::Vector3f e_voxel_size(voxel_size[0], voxel_size[1], voxel_size[2]); + Eigen::Vector3i e_grid_min( + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])); + Eigen::Vector3i e_grid_max( + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])); + + std::vector e_triangles; + e_triangles.reserve(N_tri * 3); + for (int64_t i = 0; i < N_tri * 3; ++i) + { + e_triangles.push_back(Eigen::Vector3f( + t_ptr[i * 3 + 0], + t_ptr[i * 3 + 1], + t_ptr[i * 3 + 2])); + } + + std::unordered_map hash_table; + std::vector voxels; + intersect_occ(e_voxel_size, e_grid_min, e_grid_max, e_triangles, hash_table, voxels); + + if (voxels.empty()) + { + return torch::empty({0, 3}, torch::kInt32); + } + return torch::from_blob(voxels.data(), {static_cast(voxels.size()), 3}, torch::kInt32).clone(); +} + /** * 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 QEF computation. * @param boundary_weight: Weight for the boundary edges in the QEF computation. * @param regularization_weight: Regularization factor to apply to the QEF matrices. @@ -472,119 +658,135 @@ void face_from_dual_vertices( * @return a tuple ((x, y, z), vertices, intersected, faces) containing the remeshed vertices and the corresponding voxel grid. */ std::tuple 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 &voxel_size, + const std::vector &grid_range, float face_weight, float boundary_weight, float regularization_weight, - bool timing -) { + bool timing) +{ + TORCH_CHECK(!vertices.is_cuda(), "vertices must be a CPU tensor"); + TORCH_CHECK(!faces.is_cuda(), "faces must be a CPU tensor"); + const int F = faces.size(0); - const float* v_ptr = vertices.data_ptr(); - const int* f_ptr = faces.data_ptr(); - const float* voxel_size_ptr = voxel_size.data_ptr(); - const int* grid_range_ptr = grid_range.data_ptr(); + const float *v_ptr = vertices.data_ptr(); + const int *f_ptr = faces.data_ptr(); clock_t start, end; std::unordered_map hash_table; - std::vector voxels; // Voxel coordinates + std::vector voxels; // Voxel coordinates std::vector means; // Mean vertex positions for each voxel - std::vector cnt; // Number of intersections for each voxel - std::vector intersected; // Indicate whether edges of voxels intersect with surface - std::vector qefs; // QEF matrices for each voxel + std::vector cnt; // Number of intersections for each voxel + std::vector intersected; // Indicate whether edges of voxels intersect with surface + std::vector qefs; // QEF matrices for each voxel // Convert tensors to Eigen types - Eigen::Vector3f e_voxel_size(voxel_size_ptr[0], voxel_size_ptr[1], voxel_size_ptr[2]); - Eigen::Vector3i e_grid_min(grid_range_ptr[0], grid_range_ptr[1], grid_range_ptr[2]); - Eigen::Vector3i e_grid_max(grid_range_ptr[3], grid_range_ptr[4], grid_range_ptr[5]); - + Eigen::Vector3f e_voxel_size(voxel_size[0], voxel_size[1], voxel_size[2]); + Eigen::Vector3i e_grid_min( + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])); + Eigen::Vector3i e_grid_max( + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])); + // Intersect QEF computation start = clock(); std::vector triangles; triangles.reserve(F * 3); - for (int f = 0; f < F; ++f) { - for (int v = 0; v < 3; ++v) { + for (int f = 0; f < F; ++f) + { + for (int v = 0; v < 3; ++v) + { triangles.push_back(Eigen::Vector3f( v_ptr[f_ptr[f * 3 + v] * 3 + 0], v_ptr[f_ptr[f * 3 + v] * 3 + 1], - v_ptr[f_ptr[f * 3 + v] * 3 + 2] - )); + v_ptr[f_ptr[f * 3 + v] * 3 + 2])); } } intersect_qef(e_voxel_size, e_grid_min, e_grid_max, triangles, hash_table, voxels, means, cnt, intersected, qefs); end = clock(); - if (timing) std::cout << "Intersect QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; + if (timing) + std::cout << "Intersect QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; // Face QEF computation - if (face_weight > 0.0f) { + if (face_weight > 0.0f) + { start = clock(); face_qef(e_voxel_size, e_grid_min, e_grid_max, triangles, hash_table, qefs); end = clock(); - if (timing) std::cout << "Face QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; + if (timing) + std::cout << "Face QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; } // Boundary QEF computation - if (boundary_weight > 0.0f) { + if (boundary_weight > 0.0f) + { start = clock(); std::map, int> edge_count; - for (int f = 0; f < F; ++f) { - for (int v0 = 0; v0 < 3; ++v0) { + for (int f = 0; f < F; ++f) + { + for (int v0 = 0; v0 < 3; ++v0) + { int e0 = f_ptr[f * 3 + v0]; int e1 = f_ptr[f * 3 + (v0 + 1) % 3]; - if (e0 > e1) std::swap(e0, e1); + if (e0 > e1) + std::swap(e0, e1); edge_count[std::make_pair(e0, e1)]++; } } std::vector boundries; - for (const auto& e : edge_count) { - if (e.second == 1) { + for (const auto &e : edge_count) + { + if (e.second == 1) + { int v0 = e.first.first; int v1 = e.first.second; boundries.push_back(Eigen::Vector3f( v_ptr[v0 * 3 + 0], v_ptr[v0 * 3 + 1], - v_ptr[v0 * 3 + 2] - )); + v_ptr[v0 * 3 + 2])); boundries.push_back(Eigen::Vector3f( v_ptr[v1 * 3 + 0], v_ptr[v1 * 3 + 1], - v_ptr[v1 * 3 + 2] - )); + v_ptr[v1 * 3 + 2])); } } boundry_qef(e_voxel_size, e_grid_min, e_grid_max, boundries, boundary_weight, hash_table, qefs); end = clock(); - if (timing) std::cout << "Boundary QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; + if (timing) + std::cout << "Boundary QEF computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; } // Solve the QEF system to obtain final dual vertices start = clock(); std::vector dual_vertices(voxels.size()); - for (int i = 0; i < voxels.size(); ++i) { + for (int i = 0; i < voxels.size(); ++i) + { int3 coord = voxels[i]; Eigen::Matrix4f Q = qefs[i]; float min_corner[3] = { coord.x * e_voxel_size.x(), coord.y * e_voxel_size.y(), - coord.z * e_voxel_size.z() - }; + coord.z * e_voxel_size.z()}; float max_corner[3] = { (coord.x + 1) * e_voxel_size.x(), (coord.y + 1) * e_voxel_size.y(), - (coord.z + 1) * e_voxel_size.z() - }; + (coord.z + 1) * e_voxel_size.z()}; // Add regularization term - if (regularization_weight > 0.0f) { + if (regularization_weight > 0.0f) + { Eigen::Vector3f p = means[i] / cnt[i]; // Construct the QEF matrix for this vertex Eigen::Matrix4f Qreg = Eigen::Matrix4f::Zero(); - Qreg.topLeftCorner<3,3>() = Eigen::Matrix3f::Identity(); - Qreg.block<3,1>(0,3) = -p; - Qreg.block<1,3>(3,0) = -p.transpose(); - Qreg(3,3) = p.dot(p); + Qreg.topLeftCorner<3, 3>() = Eigen::Matrix3f::Identity(); + Qreg.block<3, 1>(0, 3) = -p; + Qreg.block<1, 3>(3, 0) = -p.transpose(); + Qreg(3, 3) = p.dot(p); Q += regularization_weight * cnt[i] * Qreg; // Scale by regularization weight } @@ -595,15 +797,16 @@ std::tuple mesh_to_flexible_dual_gr Eigen::Vector3f v_new = A.colPivHouseholderQr().solve(b); if (!( - v_new.x() >= min_corner[0] && v_new.x() <= max_corner[0] && - v_new.y() >= min_corner[1] && v_new.y() <= max_corner[1] && - v_new.z() >= min_corner[2] && v_new.z() <= max_corner[2] - )) { + v_new.x() >= min_corner[0] && v_new.x() <= max_corner[0] && + v_new.y() >= min_corner[1] && v_new.y() <= max_corner[1] && + v_new.z() >= min_corner[2] && v_new.z() <= max_corner[2])) + { // Starting enumeration of constraints float best = std::numeric_limits::infinity(); // Solve single-constraint - auto solve_single_constraint = [&](int fixed_axis) { + auto solve_single_constraint = [&](int fixed_axis) + { int ax1 = (fixed_axis + 1) % 3; int ax2 = (fixed_axis + 2) % 3; @@ -612,9 +815,9 @@ std::tuple mesh_to_flexible_dual_gr Eigen::Vector2f q, b, x; A << Q(ax1, ax1), Q(ax1, ax2), - Q(ax2, ax1), Q(ax2, ax2); + Q(ax2, ax1), Q(ax2, ax2); B << Q(ax1, fixed_axis), Q(ax1, 3), - Q(ax2, fixed_axis), Q(ax2, 3); + Q(ax2, fixed_axis), Q(ax2, 3); auto Asol = A.colPivHouseholderQr(); // if lower bound @@ -623,15 +826,16 @@ std::tuple mesh_to_flexible_dual_gr x = Asol.solve(b); if ( x.x() >= min_corner[ax1] && x.x() <= max_corner[ax1] && - x.y() >= min_corner[ax2] && x.y() <= max_corner[ax2] - ) { + x.y() >= min_corner[ax2] && x.y() <= max_corner[ax2]) + { Eigen::Vector4f p; p[fixed_axis] = min_corner[fixed_axis]; p[ax1] = x.x(); p[ax2] = x.y(); p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -643,15 +847,16 @@ std::tuple mesh_to_flexible_dual_gr x = Asol.solve(b); if ( x.x() >= min_corner[ax1] && x.x() <= max_corner[ax1] && - x.y() >= min_corner[ax2] && x.y() <= max_corner[ax2] - ) { + x.y() >= min_corner[ax2] && x.y() <= max_corner[ax2]) + { Eigen::Vector4f p; p[fixed_axis] = max_corner[fixed_axis]; p[ax1] = x.x(); p[ax2] = x.y(); p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -662,7 +867,8 @@ std::tuple mesh_to_flexible_dual_gr solve_single_constraint(2); // fix z // Solve two-constraint - auto solve_two_constraint = [&](int free_axis) { + auto solve_two_constraint = [&](int free_axis) + { int ax1 = (free_axis + 1) % 3; int ax2 = (free_axis + 2) % 3; @@ -675,14 +881,16 @@ std::tuple mesh_to_flexible_dual_gr // if lower-lower bound q << min_corner[ax1], min_corner[ax2], 1; x = -(b.dot(q)) / a; - if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) + { Eigen::Vector4f p; p[free_axis] = x; p[ax1] = min_corner[ax1]; p[ax2] = min_corner[ax2]; p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -691,14 +899,16 @@ std::tuple mesh_to_flexible_dual_gr // if lower-upper bound q << min_corner[ax1], max_corner[ax2], 1; x = -(b.dot(q)) / a; - if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) + { Eigen::Vector4f p; p[free_axis] = x; p[ax1] = min_corner[ax1]; p[ax2] = max_corner[ax2]; p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -707,14 +917,16 @@ std::tuple mesh_to_flexible_dual_gr // if upper-lower bound q << max_corner[ax1], min_corner[ax2], 1; x = -(b.dot(q)) / a; - if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) + { Eigen::Vector4f p; p[free_axis] = x; p[ax1] = max_corner[ax1]; p[ax2] = min_corner[ax2]; p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -723,14 +935,16 @@ std::tuple mesh_to_flexible_dual_gr // if upper-upper bound q << max_corner[ax1], max_corner[ax2], 1; x = -(b.dot(q)) / a; - if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) + { Eigen::Vector4f p; p[free_axis] = x; p[ax1] = max_corner[ax1]; p[ax2] = max_corner[ax2]; p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -741,9 +955,12 @@ std::tuple mesh_to_flexible_dual_gr solve_two_constraint(2); // free z // Solve three-constraint - for (int x_constraint = 0; x_constraint < 2; ++x_constraint) { - for (int y_constraint = 0; y_constraint < 2; ++y_constraint) { - for (int z_constraint = 0; z_constraint < 2; ++z_constraint) { + for (int x_constraint = 0; x_constraint < 2; ++x_constraint) + { + for (int y_constraint = 0; y_constraint < 2; ++y_constraint) + { + for (int z_constraint = 0; z_constraint < 2; ++z_constraint) + { Eigen::Vector4f p; p[0] = x_constraint ? min_corner[0] : max_corner[0]; p[1] = y_constraint ? min_corner[1] : max_corner[1]; @@ -751,7 +968,8 @@ std::tuple mesh_to_flexible_dual_gr p[3] = 1.0f; float err = p.transpose() * Q * p; - if (err < best) { + if (err < best) + { best = err; v_new << p[0], p[1], p[2]; } @@ -764,12 +982,11 @@ std::tuple mesh_to_flexible_dual_gr dual_vertices[i] = float3{v_new.x(), v_new.y(), v_new.z()}; } end = clock(); - if (timing) std::cout << "Dual vertices computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; + if (timing) + std::cout << "Dual vertices computation took " << double(end - start) / CLOCKS_PER_SEC << " seconds." << std::endl; return std::make_tuple( - torch::from_blob(voxels.data(), {int(voxels .size()), 3}, torch::kInt32).clone(), + torch::from_blob(voxels.data(), {int(voxels.size()), 3}, torch::kInt32).clone(), torch::from_blob(dual_vertices.data(), {int(dual_vertices.size()), 3}, torch::kFloat32).clone(), - torch::from_blob(intersected.data(), {int(intersected.size()), 3}, torch::kBool).clone() - ); + torch::from_blob(intersected.data(), {int(intersected.size()), 3}, torch::kBool).clone()); } - diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/boundary_qef.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/boundary_qef.cu new file mode 100644 index 00000000..44abc346 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/boundary_qef.cu @@ -0,0 +1,336 @@ +#include "../api.h" + +#include "types.cuh" + +#include +#include +#include +#include +#include + +#include + +// Boundary QEFs are accumulated by walking each boundary segment through the +// voxel grid. A thread visits the voxels crossed by one segment, looks them up +// in the active-brick structure built by intersect_qef_cuda, and adds directly +// into the running qefs tensor. +namespace o_voxel::fdg +{ + namespace + { + + constexpr int kThreads = 256; + + __device__ __forceinline__ void add_boundary_voxel( + int x, + int y, + int z, + GridSpec grid, + BrickLookup lookup, + SymQEF10 qef, + int &cached_bx, + int &cached_by, + int &cached_bz, + const uint32_t *&cached_bits, + int64_t &cached_base, + bool &cached_found, + float *out_qefs) + { + // Map a voxel coordinate to its compact row and add one QEF. The + // brick cache avoids repeating the hash lookup while a segment stays + // within the same 8x8x8 brick. + if (x < grid.grid_min.x || x >= grid.grid_max.x) + return; + if (y < grid.grid_min.y || y >= grid.grid_max.y) + return; + if (z < grid.grid_min.z || z >= grid.grid_max.z) + return; + + const int rx = x - grid.grid_min.x; + const int ry = y - grid.grid_min.y; + const int rz = z - grid.grid_min.z; + const int bx = rx / kBrickSize; + const int by = ry / kBrickSize; + const int bz = rz / kBrickSize; + const int lx = rx - bx * kBrickSize; + const int ly = ry - by * kBrickSize; + const int lz = rz - bz * kBrickSize; + const int local_id = lx + kBrickSize * (ly + kBrickSize * lz); + + // Boundary DDA often visits many consecutive voxels in the same + // brick. Reuse the previous lookup until the brick coordinate + // changes. + if (bx != cached_bx || by != cached_by || bz != cached_bz) + { + cached_bx = bx; + cached_by = by; + cached_bz = bz; + cached_found = lookup_brick_bits_and_base(bx, by, bz, grid, lookup, &cached_bits, &cached_base); + } + if (!cached_found) + return; + + const int word = local_id / 32; + const int bit = local_id - word * 32; + if ((cached_bits[word] & (1u << bit)) == 0) + return; + + int rank = 0; + for (int i = 0; i < word; ++i) + rank += __popc(cached_bits[i]); + const uint32_t mask = bit == 0 ? 0u : ((1u << bit) - 1u); + rank += __popc(cached_bits[word] & mask); + + // Compact row inside qefs is the brick's base row plus the number + // of active local bits before this voxel. + float *dst = out_qefs + 10 * (cached_base + rank); + atomicAdd(dst + 0, qef.q00); + atomicAdd(dst + 1, qef.q01); + atomicAdd(dst + 2, qef.q02); + atomicAdd(dst + 3, qef.q03); + atomicAdd(dst + 4, qef.q11); + atomicAdd(dst + 5, qef.q12); + atomicAdd(dst + 6, qef.q13); + atomicAdd(dst + 7, qef.q22); + atomicAdd(dst + 8, qef.q23); + atomicAdd(dst + 9, qef.q33); + } + + __global__ void accumulate_boundary_qef_kernel( + const float *__restrict__ boundaries, + int64_t num_boundaries, + GridSpec grid, + float boundary_weight, + BrickLookup lookup, + float *__restrict__ out_qefs) + { + // One thread handles one boundary segment and advances through grid + // cells in DDA order until the segment length is reached. + const int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_boundaries) + return; + + const float *seg = boundaries + 6 * eid; + const float3 p0 = make_float3(seg[0], seg[1], seg[2]); + const float3 p1 = make_float3(seg[3], seg[4], seg[5]); + + const double dx = static_cast(p1.x) - static_cast(p0.x); + const double dy = static_cast(p1.y) - static_cast(p0.y); + const double dz = static_cast(p1.z) - static_cast(p0.z); + const double segment_length = sqrt(dx * dx + dy * dy + dz * dz); + if (segment_length < 1e-6) + return; + + const double dir_x = dx / segment_length; + const double dir_y = dy / segment_length; + const double dir_z = dz / segment_length; + // Squared distance to a line through p0 with unit direction d: + // ||(I - d d^T)(x - p0)||^2. The symmetric matrix below stores + // A = I - d d^T, b = -A*p0, c = p0^T*A*p0 in QEF form. + const float a00 = 1.0f - static_cast(dir_x * dir_x); + const float a01 = -static_cast(dir_x * dir_y); + const float a02 = -static_cast(dir_x * dir_z); + const float a11 = 1.0f - static_cast(dir_y * dir_y); + const float a12 = -static_cast(dir_y * dir_z); + const float a22 = 1.0f - static_cast(dir_z * dir_z); + const float b0 = -(a00 * p0.x + a01 * p0.y + a02 * p0.z); + const float b1 = -(a01 * p0.x + a11 * p0.y + a12 * p0.z); + const float b2 = -(a02 * p0.x + a12 * p0.y + a22 * p0.z); + const float av0_x = a00 * p0.x + a01 * p0.y + a02 * p0.z; + const float av0_y = a01 * p0.x + a11 * p0.y + a12 * p0.z; + const float av0_z = a02 * p0.x + a12 * p0.y + a22 * p0.z; + const float c = p0.x * av0_x + p0.y * av0_y + p0.z * av0_z; + const SymQEF10 qef{ + boundary_weight * a00, + boundary_weight * a01, + boundary_weight * a02, + boundary_weight * b0, + boundary_weight * a11, + boundary_weight * a12, + boundary_weight * b1, + boundary_weight * a22, + boundary_weight * b2, + boundary_weight * c, + }; + const int step_x = dir_x > 0.0 ? 1 : -1; + const int step_y = dir_y > 0.0 ? 1 : -1; + const int step_z = dir_z > 0.0 ? 1 : -1; + int cur_x = static_cast(floorf(p0.x / grid.voxel_size.x)); + int cur_y = static_cast(floorf(p0.y / grid.voxel_size.y)); + int cur_z = static_cast(floorf(p0.z / grid.voxel_size.z)); + + // tmax_* is the distance along the segment to the next grid plane on + // that axis. tdelta_* is the distance between later grid-plane + // crossings on the same axis. + double tmax_x; + double tmax_y; + double tmax_z; + double tdelta_x; + double tdelta_y; + double tdelta_z; + + if (dir_x == 0.0) + { + tmax_x = CUDART_INF; + tdelta_x = CUDART_INF; + } + else + { + const float border = grid.voxel_size.x * static_cast(cur_x + (step_x > 0 ? 1 : 0)); + tmax_x = static_cast(border - p0.x) / dir_x; + tdelta_x = static_cast(grid.voxel_size.x) / fabs(dir_x); + } + if (dir_y == 0.0) + { + tmax_y = CUDART_INF; + tdelta_y = CUDART_INF; + } + else + { + const float border = grid.voxel_size.y * static_cast(cur_y + (step_y > 0 ? 1 : 0)); + tmax_y = static_cast(border - p0.y) / dir_y; + tdelta_y = static_cast(grid.voxel_size.y) / fabs(dir_y); + } + if (dir_z == 0.0) + { + tmax_z = CUDART_INF; + tdelta_z = CUDART_INF; + } + else + { + const float border = grid.voxel_size.z * static_cast(cur_z + (step_z > 0 ? 1 : 0)); + tmax_z = static_cast(border - p0.z) / dir_z; + tdelta_z = static_cast(grid.voxel_size.z) / fabs(dir_z); + } + + int cached_bx = -1; + int cached_by = -1; + int cached_bz = -1; + const uint32_t *cached_bits = nullptr; + int64_t cached_base = 0; + bool cached_found = false; + + add_boundary_voxel( + cur_x, + cur_y, + cur_z, + grid, + lookup, + qef, + cached_bx, + cached_by, + cached_bz, + cached_bits, + cached_base, + cached_found, + out_qefs); + + while (true) + { + int axis; + // Advance to the nearest next grid plane, visit the voxel on the + // other side, and stop once that crossing would be past p1. + if (tmax_x < tmax_y) + axis = (tmax_x < tmax_z) ? 0 : 2; + else + axis = (tmax_y < tmax_z) ? 1 : 2; + + if (axis == 0 && tmax_x > segment_length) + break; + if (axis == 1 && tmax_y > segment_length) + break; + if (axis == 2 && tmax_z > segment_length) + break; + + if (axis == 0) + { + cur_x += step_x; + tmax_x += tdelta_x; + } + else if (axis == 1) + { + cur_y += step_y; + tmax_y += tdelta_y; + } + else + { + cur_z += step_z; + tmax_z += tdelta_z; + } + + add_boundary_voxel( + cur_x, + cur_y, + cur_z, + grid, + lookup, + qef, + cached_bx, + cached_by, + cached_bz, + cached_bits, + cached_base, + cached_found, + out_qefs); + } + } + + } // namespace + + torch::Tensor boundary_qef_cuda( + const torch::Tensor &boundaries, + const std::vector &voxel_size, + const std::vector &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) + { + // qefs is an in-place accumulator. boundary_weight is already folded + // into each segment QEF before the atomic adds. + TORCH_CHECK(boundaries.is_cuda(), "boundaries must be a CUDA tensor"); + TORCH_CHECK(voxels.is_cuda(), "voxels must be a CUDA tensor"); + static_assert(sizeof(SymQEF10) == sizeof(float) * 10, "Unexpected SymQEF10 layout"); + + const c10::cuda::CUDAGuard guard(boundaries.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(boundaries.get_device()).stream(); + const int64_t num_boundaries = boundaries.size(0); + const int64_t num_voxels = voxels.size(0); + if (num_boundaries == 0 || num_voxels == 0 || boundary_weight <= 0.0f) + return qefs; + + const GridSpec grid{ + float3{voxel_size[0], voxel_size[1], voxel_size[2]}, + Int3{ + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])}, + Int3{ + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])}, + }; + const BrickLookup lookup{ + brick_hash_keys.data_ptr(), + brick_hash_vals.data_ptr(), + brick_bits.data_ptr(), + brick_base.data_ptr(), + static_cast(brick_hash_keys.numel()), + }; + + const int blocks = static_cast((num_boundaries + kThreads - 1) / kThreads); + accumulate_boundary_qef_kernel<<>>( + boundaries.data_ptr(), + num_boundaries, + grid, + boundary_weight, + lookup, + qefs.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return qefs; + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/face_qef.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/face_qef.cu new file mode 100644 index 00000000..321750f9 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/face_qef.cu @@ -0,0 +1,471 @@ +#include "../api.h" + +#include "qef.cuh" + +#include +#include +#include +#include +#include + +#include + +// Face QEFs are accumulated as a triangle-brick task stream. Each task checks +// one triangle against one active brick-sized region, scans only voxels that are +// already active in that brick, and adds face_weight * plane_qef directly into +// the running qefs tensor. +namespace o_voxel::fdg +{ + namespace + { + + constexpr int kThreads = 256; + + // A compact work item: triangle tri_id tested against brick (bx, by, bz). + struct FaceBrickTask + { + int32_t tri_id; + int32_t bx; + int32_t by; + int32_t bz; + }; + + // Matches CPU e0.cross(e1).normalized(); explicit RN ops prevent FMA residuals. + __device__ __forceinline__ float3 aligned_cross_normalize3( + float e0x, + float e0y, + float e0z, + float e1x, + float e1y, + float e1z) + { + const float nx = __fsub_rn(__fmul_rn(e0y, e1z), __fmul_rn(e0z, e1y)); + const float ny = __fsub_rn(__fmul_rn(e0z, e1x), __fmul_rn(e0x, e1z)); + const float nz = __fsub_rn(__fmul_rn(e0x, e1y), __fmul_rn(e0y, e1x)); + const float n2 = __fadd_rn( + __fadd_rn(__fmul_rn(nx, nx), __fmul_rn(ny, ny)), + __fmul_rn(nz, nz)); + if (n2 > 0.0f) + { + const float len = sqrtf(n2); + return make_float3(nx / len, ny / len, nz / len); + } + return make_float3(nx, ny, nz); + } + + __global__ void count_face_brick_tasks_kernel( + const float *__restrict__ triangles, + int64_t num_triangles, + GridSpec grid, + int64_t *__restrict__ task_counts) + { + const int64_t tri_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tri_id >= num_triangles) + return; + + const float *tri = triangles + tri_id * 9; + const float vs[3] = {grid.voxel_size.x, grid.voxel_size.y, grid.voxel_size.z}; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + + // The first pass only counts how many active-brick-sized regions a + // triangle's voxel bbox can overlap. CUB scan then gives each + // triangle a compact output range for task emission. + float min_x = tri[0] < tri[3] ? tri[0] : tri[3]; + min_x = min_x < tri[6] ? min_x : tri[6]; + float min_y = tri[1] < tri[4] ? tri[1] : tri[4]; + min_y = min_y < tri[7] ? min_y : tri[7]; + float min_z = tri[2] < tri[5] ? tri[2] : tri[5]; + min_z = min_z < tri[8] ? min_z : tri[8]; + float max_x = tri[0] > tri[3] ? tri[0] : tri[3]; + max_x = max_x > tri[6] ? max_x : tri[6]; + float max_y = tri[1] > tri[4] ? tri[1] : tri[4]; + max_y = max_y > tri[7] ? max_y : tri[7]; + float max_z = tri[2] > tri[5] ? tri[2] : tri[5]; + max_z = max_z > tri[8] ? max_z : tri[8]; + + const int bb_min_x = max(static_cast(min_x / vs[0]), grid_min.x); + const int bb_min_y = max(static_cast(min_y / vs[1]), grid_min.y); + const int bb_min_z = max(static_cast(min_z / vs[2]), grid_min.z); + const int bb_max_x = min(static_cast(max_x / vs[0] + 1.0f), grid_max.x); + const int bb_max_y = min(static_cast(max_y / vs[1] + 1.0f), grid_max.y); + const int bb_max_z = min(static_cast(max_z / vs[2] + 1.0f), grid_max.z); + if (bb_max_x <= bb_min_x || bb_max_y <= bb_min_y || bb_max_z <= bb_min_z) + { + task_counts[tri_id] = 0; + return; + } + + const int bx0 = (bb_min_x - grid_min.x) / kBrickSize; + const int by0 = (bb_min_y - grid_min.y) / kBrickSize; + const int bz0 = (bb_min_z - grid_min.z) / kBrickSize; + const int bx1 = (bb_max_x - 1 - grid_min.x) / kBrickSize; + const int by1 = (bb_max_y - 1 - grid_min.y) / kBrickSize; + const int bz1 = (bb_max_z - 1 - grid_min.z) / kBrickSize; + // One FaceBrickTask is cheaper than asking every active voxel to + // test every triangle. Later, the task only scans active bits inside + // that brick. + task_counts[tri_id] = + static_cast(bx1 - bx0 + 1) * + static_cast(by1 - by0 + 1) * + static_cast(bz1 - bz0 + 1); + } + + __global__ void emit_face_brick_tasks_kernel( + const float *__restrict__ triangles, + int64_t num_triangles, + GridSpec grid, + const int64_t *__restrict__ task_offsets, + FaceBrickTask *__restrict__ tasks) + { + const int64_t tri_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tri_id >= num_triangles) + return; + + const float *tri = triangles + tri_id * 9; + const float vs[3] = {grid.voxel_size.x, grid.voxel_size.y, grid.voxel_size.z}; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + + float min_x = tri[0] < tri[3] ? tri[0] : tri[3]; + min_x = min_x < tri[6] ? min_x : tri[6]; + float min_y = tri[1] < tri[4] ? tri[1] : tri[4]; + min_y = min_y < tri[7] ? min_y : tri[7]; + float min_z = tri[2] < tri[5] ? tri[2] : tri[5]; + min_z = min_z < tri[8] ? min_z : tri[8]; + float max_x = tri[0] > tri[3] ? tri[0] : tri[3]; + max_x = max_x > tri[6] ? max_x : tri[6]; + float max_y = tri[1] > tri[4] ? tri[1] : tri[4]; + max_y = max_y > tri[7] ? max_y : tri[7]; + float max_z = tri[2] > tri[5] ? tri[2] : tri[5]; + max_z = max_z > tri[8] ? max_z : tri[8]; + + const int bb_min_x = max(static_cast(min_x / vs[0]), grid_min.x); + const int bb_min_y = max(static_cast(min_y / vs[1]), grid_min.y); + const int bb_min_z = max(static_cast(min_z / vs[2]), grid_min.z); + const int bb_max_x = min(static_cast(max_x / vs[0] + 1.0f), grid_max.x); + const int bb_max_y = min(static_cast(max_y / vs[1] + 1.0f), grid_max.y); + const int bb_max_z = min(static_cast(max_z / vs[2] + 1.0f), grid_max.z); + if (bb_max_x <= bb_min_x || bb_max_y <= bb_min_y || bb_max_z <= bb_min_z) + return; + + const int bx0 = (bb_min_x - grid_min.x) / kBrickSize; + const int by0 = (bb_min_y - grid_min.y) / kBrickSize; + const int bz0 = (bb_min_z - grid_min.z) / kBrickSize; + const int bx1 = (bb_max_x - 1 - grid_min.x) / kBrickSize; + const int by1 = (bb_max_y - 1 - grid_min.y) / kBrickSize; + const int bz1 = (bb_max_z - 1 - grid_min.z) / kBrickSize; + + int64_t out = task_offsets[tri_id]; + // The task grid is triangle bbox clipped to brick coordinates. It is + // intentionally conservative; accumulate_face_qef_kernel performs + // the triangle/voxel-box overlap test. + for (int bz = bz0; bz <= bz1; ++bz) + for (int by = by0; by <= by1; ++by) + for (int bx = bx0; bx <= bx1; ++bx) + tasks[out++] = FaceBrickTask{static_cast(tri_id), bx, by, bz}; + } + + __global__ void accumulate_face_qef_kernel( + const FaceBrickTask *__restrict__ tasks, + int64_t num_tasks, + const float *__restrict__ triangles, + GridSpec grid, + BrickLookup lookup, + float face_weight, + float *__restrict__ out_qefs) + { + const int64_t task_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (task_id >= num_tasks) + return; + + const FaceBrickTask task = tasks[task_id]; + const uint32_t *bits; + int64_t base; + // If the triangle bbox brick was never activated by intersect_qef, + // it has no output rows and the whole task can be skipped. + if (!lookup_brick_bits_and_base(task.bx, task.by, task.bz, grid, lookup, &bits, &base)) + return; + + const float *tri = triangles + static_cast(task.tri_id) * 9; + const float3 v0 = make_float3(tri[0], tri[1], tri[2]); + const float3 v1 = make_float3(tri[3], tri[4], tri[5]); + const float3 v2 = make_float3(tri[6], tri[7], tri[8]); + const float vs[3] = {grid.voxel_size.x, grid.voxel_size.y, grid.voxel_size.z}; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + + const float e0x = v1.x - v0.x; + const float e0y = v1.y - v0.y; + const float e0z = v1.z - v0.z; + const float e1x = v2.x - v1.x; + const float e1y = v2.y - v1.y; + const float e1z = v2.z - v1.z; + const float e2x = v0.x - v2.x; + const float e2y = v0.y - v2.y; + const float e2z = v0.z - v2.z; + + const float3 n = aligned_cross_normalize3(e0x, e0y, e0z, e1x, e1y, e1z); + const float nx = n.x; + const float ny = n.y; + const float nz = n.z; + // Face QEF is the squared distance to the triangle plane, scaled + // once here before all matching voxels receive it. + const SymQEF10 qef = qef_scale( + qef_from_plane(make_float4(nx, ny, nz, -(nx * v0.x + ny * v0.y + nz * v0.z))), + face_weight); + + float min_x = v0.x < v1.x ? v0.x : v1.x; + min_x = min_x < v2.x ? min_x : v2.x; + float min_y = v0.y < v1.y ? v0.y : v1.y; + min_y = min_y < v2.y ? min_y : v2.y; + float min_z = v0.z < v1.z ? v0.z : v1.z; + min_z = min_z < v2.z ? min_z : v2.z; + float max_x = v0.x > v1.x ? v0.x : v1.x; + max_x = max_x > v2.x ? max_x : v2.x; + float max_y = v0.y > v1.y ? v0.y : v1.y; + max_y = max_y > v2.y ? max_y : v2.y; + float max_z = v0.z > v1.z ? v0.z : v1.z; + max_z = max_z > v2.z ? max_z : v2.z; + + const int bb_min_x = max(static_cast(min_x / vs[0]), grid_min.x); + const int bb_min_y = max(static_cast(min_y / vs[1]), grid_min.y); + const int bb_min_z = max(static_cast(min_z / vs[2]), grid_min.z); + const int bb_max_x = min(static_cast(max_x / vs[0] + 1.0f), grid_max.x); + const int bb_max_y = min(static_cast(max_y / vs[1] + 1.0f), grid_max.y); + const int bb_max_z = min(static_cast(max_z / vs[2] + 1.0f), grid_max.z); + + const float c_x = nx > 0.0f ? vs[0] : 0.0f; + const float c_y = ny > 0.0f ? vs[1] : 0.0f; + const float c_z = nz > 0.0f ? vs[2] : 0.0f; + const float d1 = nx * (c_x - v0.x) + ny * (c_y - v0.y) + nz * (c_z - v0.z); + const float d2 = nx * (vs[0] - c_x - v0.x) + ny * (vs[1] - c_y - v0.y) + nz * (vs[2] - c_z - v0.z); + + // Plane slab test: choose the two voxel-box corners that are most + // separated along the triangle normal. If both signed distances have + // the same sign, the box cannot cross the triangle plane. + const int mul_xy = nz < 0.0f ? -1 : 1; + const float n_xy_e0_x = -mul_xy * e0y; + const float n_xy_e0_y = mul_xy * e0x; + const float n_xy_e1_x = -mul_xy * e1y; + const float n_xy_e1_y = mul_xy * e1x; + const float n_xy_e2_x = -mul_xy * e2y; + const float n_xy_e2_y = mul_xy * e2x; + const float d_xy_e0 = -(n_xy_e0_x * v0.x + n_xy_e0_y * v0.y) + fmaxf(n_xy_e0_x, 0.0f) * vs[0] + fmaxf(n_xy_e0_y, 0.0f) * vs[1]; + const float d_xy_e1 = -(n_xy_e1_x * v1.x + n_xy_e1_y * v1.y) + fmaxf(n_xy_e1_x, 0.0f) * vs[0] + fmaxf(n_xy_e1_y, 0.0f) * vs[1]; + const float d_xy_e2 = -(n_xy_e2_x * v2.x + n_xy_e2_y * v2.y) + fmaxf(n_xy_e2_x, 0.0f) * vs[0] + fmaxf(n_xy_e2_y, 0.0f) * vs[1]; + + const int mul_yz = nx < 0.0f ? -1 : 1; + const float n_yz_e0_x = -mul_yz * e0z; + const float n_yz_e0_y = mul_yz * e0y; + const float n_yz_e1_x = -mul_yz * e1z; + const float n_yz_e1_y = mul_yz * e1y; + const float n_yz_e2_x = -mul_yz * e2z; + const float n_yz_e2_y = mul_yz * e2y; + const float d_yz_e0 = -(n_yz_e0_x * v0.y + n_yz_e0_y * v0.z) + fmaxf(n_yz_e0_x, 0.0f) * vs[1] + fmaxf(n_yz_e0_y, 0.0f) * vs[2]; + const float d_yz_e1 = -(n_yz_e1_x * v1.y + n_yz_e1_y * v1.z) + fmaxf(n_yz_e1_x, 0.0f) * vs[1] + fmaxf(n_yz_e1_y, 0.0f) * vs[2]; + const float d_yz_e2 = -(n_yz_e2_x * v2.y + n_yz_e2_y * v2.z) + fmaxf(n_yz_e2_x, 0.0f) * vs[1] + fmaxf(n_yz_e2_y, 0.0f) * vs[2]; + + const int mul_zx = ny < 0.0f ? -1 : 1; + const float n_zx_e0_x = -mul_zx * e0x; + const float n_zx_e0_y = mul_zx * e0z; + const float n_zx_e1_x = -mul_zx * e1x; + const float n_zx_e1_y = mul_zx * e1z; + const float n_zx_e2_x = -mul_zx * e2x; + const float n_zx_e2_y = mul_zx * e2z; + const float d_zx_e0 = -(n_zx_e0_x * v0.z + n_zx_e0_y * v0.x) + fmaxf(n_zx_e0_x, 0.0f) * vs[2] + fmaxf(n_zx_e0_y, 0.0f) * vs[0]; + const float d_zx_e1 = -(n_zx_e1_x * v1.z + n_zx_e1_y * v1.x) + fmaxf(n_zx_e1_x, 0.0f) * vs[2] + fmaxf(n_zx_e1_y, 0.0f) * vs[0]; + const float d_zx_e2 = -(n_zx_e2_x * v2.z + n_zx_e2_y * v2.x) + fmaxf(n_zx_e2_x, 0.0f) * vs[2] + fmaxf(n_zx_e2_y, 0.0f) * vs[0]; + + // The xy/yz/zx edge functions are projected half-space tests. The + // fmax terms move the sampled box corner to the side most favorable + // to overlap, making this a voxel-box vs triangle test rather than + // a point-in-triangle test. + int rank_before_word = 0; + for (int word = 0; word < kBrickBitWords; ++word) + { + const uint32_t word_bits = bits[word]; + uint32_t active = word_bits; + int rank_in_word = 0; + while (active != 0) + { + const int bit = __ffs(active) - 1; + const int local_id = word * 32 + bit; + const int local_rank = rank_before_word + rank_in_word; + // local_rank is the number of active bits before this voxel + // in the brick, so base + local_rank is the compact qef row. + const int lz = local_id / (kBrickSize * kBrickSize); + const int rem = local_id - lz * kBrickSize * kBrickSize; + const int ly = rem / kBrickSize; + const int lx = rem - ly * kBrickSize; + const int x = grid_min.x + task.bx * kBrickSize + lx; + const int y = grid_min.y + task.by * kBrickSize + ly; + const int z = grid_min.z + task.bz * kBrickSize + lz; + + if (x >= bb_min_x && x < bb_max_x && + y >= bb_min_y && y < bb_max_y && + z >= bb_min_z && z < bb_max_z) + { + const float px = x * vs[0]; + const float py = y * vs[1]; + const float pz = z * vs[2]; + const float n_dot_p = nx * px + ny * py + nz * pz; + bool hit = true; + // Reject if the voxel box is entirely on one side of + // the triangle plane, or outside any projected edge + // half-space. + if ((n_dot_p + d1) * (n_dot_p + d2) > 0.0f) + hit = false; + if (n_xy_e0_x * px + n_xy_e0_y * py + d_xy_e0 < 0.0f) + hit = false; + if (n_xy_e1_x * px + n_xy_e1_y * py + d_xy_e1 < 0.0f) + hit = false; + if (n_xy_e2_x * px + n_xy_e2_y * py + d_xy_e2 < 0.0f) + hit = false; + if (n_yz_e0_x * py + n_yz_e0_y * pz + d_yz_e0 < 0.0f) + hit = false; + if (n_yz_e1_x * py + n_yz_e1_y * pz + d_yz_e1 < 0.0f) + hit = false; + if (n_yz_e2_x * py + n_yz_e2_y * pz + d_yz_e2 < 0.0f) + hit = false; + if (n_zx_e0_x * pz + n_zx_e0_y * px + d_zx_e0 < 0.0f) + hit = false; + if (n_zx_e1_x * pz + n_zx_e1_y * px + d_zx_e1 < 0.0f) + hit = false; + if (n_zx_e2_x * pz + n_zx_e2_y * px + d_zx_e2 < 0.0f) + hit = false; + if (hit) + { + // Different triangle-brick tasks can contribute to + // the same voxel, so the total QEF is accumulated + // with atomic adds. + float *dst = out_qefs + 10 * (base + local_rank); + atomicAdd(dst + 0, qef.q00); + atomicAdd(dst + 1, qef.q01); + atomicAdd(dst + 2, qef.q02); + atomicAdd(dst + 3, qef.q03); + atomicAdd(dst + 4, qef.q11); + atomicAdd(dst + 5, qef.q12); + atomicAdd(dst + 6, qef.q13); + atomicAdd(dst + 7, qef.q22); + atomicAdd(dst + 8, qef.q23); + atomicAdd(dst + 9, qef.q33); + } + } + + active &= active - 1u; + ++rank_in_word; + } + rank_before_word += __popc(word_bits); + } + } + + } // namespace + + torch::Tensor face_qef_cuda( + const torch::Tensor &triangles, + const std::vector &voxel_size, + const std::vector &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) + { + // qefs is an in-place accumulator. No separate face-QEF buffer is + // allocated; the caller should pass the running total QEF tensor. + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + TORCH_CHECK(voxels.is_cuda(), "voxels must be a CUDA tensor"); + + const c10::cuda::CUDAGuard guard(triangles.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(triangles.get_device()).stream(); + const torch::Device device = triangles.device(); + const int64_t num_triangles = triangles.size(0); + const int64_t num_voxels = voxels.size(0); + if (num_triangles == 0 || num_voxels == 0) + return qefs; + + const auto opts_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto opts_u8 = torch::TensorOptions().dtype(torch::kUInt8).device(device); + const GridSpec grid{ + float3{voxel_size[0], voxel_size[1], voxel_size[2]}, + Int3{ + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])}, + Int3{ + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])}, + }; + const BrickLookup lookup{ + brick_hash_keys.data_ptr(), + brick_hash_vals.data_ptr(), + brick_bits.data_ptr(), + brick_base.data_ptr(), + static_cast(brick_hash_keys.numel()), + }; + + auto task_counts = torch::empty({num_triangles}, opts_i64); + auto task_offsets = torch::empty({num_triangles}, opts_i64); + int blocks = static_cast((num_triangles + kThreads - 1) / kThreads); + // count -> exclusive scan -> emit keeps task storage exact while still + // letting every triangle count its brick work independently. + count_face_brick_tasks_kernel<<>>( + triangles.data_ptr(), + num_triangles, + grid, + task_counts.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + size_t temp_bytes = 0; + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + nullptr, + temp_bytes, + task_counts.data_ptr(), + task_offsets.data_ptr(), + static_cast(num_triangles), + stream)); + auto temp = torch::empty({static_cast(temp_bytes)}, opts_u8); + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + temp.data_ptr(), + temp_bytes, + task_counts.data_ptr(), + task_offsets.data_ptr(), + static_cast(num_triangles), + stream)); + + int64_t tail[2] = {0, 0}; + C10_CUDA_CHECK(cudaMemcpyAsync(tail, task_counts.data_ptr() + num_triangles - 1, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaMemcpyAsync(tail + 1, task_offsets.data_ptr() + num_triangles - 1, sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + const int64_t num_tasks = tail[0] + tail[1]; + if (num_tasks == 0) + return qefs; + + auto tasks = torch::empty({num_tasks, 4}, opts_i32); + emit_face_brick_tasks_kernel<<>>( + triangles.data_ptr(), + num_triangles, + grid, + task_offsets.data_ptr(), + reinterpret_cast(tasks.data_ptr())); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + blocks = static_cast((num_tasks + kThreads - 1) / kThreads); + accumulate_face_qef_kernel<<>>( + reinterpret_cast(tasks.data_ptr()), + num_tasks, + triangles.data_ptr(), + grid, + lookup, + face_weight, + qefs.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return qefs; + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_common.h b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_common.h new file mode 100644 index 00000000..247ffef3 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_common.h @@ -0,0 +1,152 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace fdg_gpu { + +inline void throw_cuda_error(cudaError_t error, const char* context) { + if (error == cudaSuccess) return; + throw std::runtime_error(std::string(context) + ": " + cudaGetErrorString(error)); +} + +struct int2_ { + int x; + int y; +}; + +struct int3_ { + int x; + int y; + int z; + + __host__ __device__ int& operator[](int i) { return (&x)[i]; } + __host__ __device__ int operator[](int i) const { return (&x)[i]; } +}; + +struct bool3_ { + bool x; + bool y; + bool z; + + __host__ __device__ bool& operator[](int i) { return (&x)[i]; } + __host__ __device__ bool operator[](int i) const { return (&x)[i]; } +}; + +template +class DeviceBuffer { +public: + DeviceBuffer() = default; + explicit DeviceBuffer(int64_t count) { allocate(count); } + ~DeviceBuffer() { release(); } + + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + DeviceBuffer(DeviceBuffer&& other) noexcept + : ptr_(other.ptr_), size_(other.size_), owns_(other.owns_) { + other.ptr_ = nullptr; + other.size_ = 0; + other.owns_ = true; + } + + DeviceBuffer& operator=(DeviceBuffer&& other) noexcept { + if (this != &other) { + release(); + ptr_ = other.ptr_; + size_ = other.size_; + owns_ = other.owns_; + other.ptr_ = nullptr; + other.size_ = 0; + other.owns_ = true; + } + return *this; + } + + void allocate(int64_t count) { + if (count < 0) { + throw std::invalid_argument("DeviceBuffer::allocate count must be non-negative"); + } + release(); + size_ = count; + owns_ = true; + if (count == 0) return; + throw_cuda_error(cudaMalloc(reinterpret_cast(&ptr_), static_cast(count) * sizeof(T)), + "cudaMalloc failed in DeviceBuffer::allocate"); + } + + void adopt(T* ptr, int64_t count) { + release(); + ptr_ = ptr; + size_ = count; + owns_ = true; + } + + void clear_async(cudaStream_t stream = nullptr) { + if (size_ == 0) return; + throw_cuda_error(cudaMemsetAsync(ptr_, 0, static_cast(size_) * sizeof(T), stream), + "cudaMemsetAsync failed in DeviceBuffer::clear_async"); + } + + T* data() noexcept { return ptr_; } + const T* data() const noexcept { return ptr_; } + int64_t size() const noexcept { return size_; } + bool empty() const noexcept { return size_ == 0; } + + T* release_ownership() noexcept { + T* out = ptr_; + ptr_ = nullptr; + size_ = 0; + owns_ = true; + return out; + } + +private: + void release() noexcept { + if (ptr_ != nullptr && owns_) { + cudaFree(ptr_); + } + ptr_ = nullptr; + size_ = 0; + owns_ = true; + } + + T* ptr_ = nullptr; + int64_t size_ = 0; + bool owns_ = true; +}; + +struct SymQEF10 { + float q00, q01, q02, q03; + float q11, q12, q13; + float q22, q23; + float q33; +}; + +struct PrimitivePairResult { + int64_t size = 0; + DeviceBuffer prim_id; + DeviceBuffer voxel_i; + DeviceBuffer voxel_j; + DeviceBuffer voxel_k; +}; + +__host__ __device__ __forceinline__ int ceil_div_i64(int64_t n, int block) { + return static_cast((n + block - 1) / block); +} + +__host__ __device__ __forceinline__ uint64_t pack_voxel_key( + int x, int y, int z, int3_ grid_min, int3_ grid_max) { + const uint64_t sx = static_cast(grid_max.x - grid_min.x); + const uint64_t sy = static_cast(grid_max.y - grid_min.y); + const uint64_t ux = static_cast(x - grid_min.x); + const uint64_t uy = static_cast(y - grid_min.y); + const uint64_t uz = static_cast(z - grid_min.z); + return ux + sx * (uy + sy * uz); +} + +} // namespace fdg_gpu diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_small_cpqr_device.cuh b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_small_cpqr_device.cuh new file mode 100644 index 00000000..e0394a3d --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/fdg_gpu_small_cpqr_device.cuh @@ -0,0 +1,316 @@ +#pragma once + +#include +#include + +namespace fdg_gpu::small_cpqr { +namespace detail { + +template +__device__ __forceinline__ float absf(float x) { + return x < 0.0f ? -x : x; +} + +template +__device__ __forceinline__ void swap_cols( + float* qr, + float* col_norms_updated, + float* col_norms_direct, + int* perm, + int c0, + int c1) { + if (c0 == c1) return; + #pragma unroll + for (int r = 0; r < N; ++r) { + const float tmp = qr[r * N + c0]; + qr[r * N + c0] = qr[r * N + c1]; + qr[r * N + c1] = tmp; + } + const float tmp_u = col_norms_updated[c0]; + col_norms_updated[c0] = col_norms_updated[c1]; + col_norms_updated[c1] = tmp_u; + + const float tmp_d = col_norms_direct[c0]; + col_norms_direct[c0] = col_norms_direct[c1]; + col_norms_direct[c1] = tmp_d; + + const int tmp_p = perm[c0]; + perm[c0] = perm[c1]; + perm[c1] = tmp_p; +} + +template +__device__ __forceinline__ void make_householder_real( + float x0, + const float* tail_in, + int tail_len, + float* beta, + float* tau, + float* essential_out) { + float tail_sq_norm = 0.0f; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + tail_sq_norm += tail_in[i] * tail_in[i]; + } + } + + const float tol = FLT_MIN; + if (tail_sq_norm <= tol) { + *beta = x0; + *tau = 0.0f; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + essential_out[i] = 0.0f; + } + return; + } + + float b = sqrtf(x0 * x0 + tail_sq_norm); + if (x0 >= 0.0f) { + b = -b; + } + const float denom = x0 - b; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + essential_out[i] = (i < tail_len) ? (tail_in[i] / denom) : 0.0f; + } + *beta = b; + *tau = (b - x0) / b; +} + +template +__device__ __forceinline__ void apply_householder_left_matrix( + float* qr, + int row0, + int col0, + const float* essential, + int tail_len, + float tau) { + if (tau == 0.0f) return; + #pragma unroll + for (int j = 0; j < N; ++j) { + if (j < col0) continue; + float tmp = qr[row0 * N + j]; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + tmp += essential[i] * qr[(row0 + 1 + i) * N + j]; + } + } + qr[row0 * N + j] -= tau * tmp; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + qr[(row0 + 1 + i) * N + j] -= tau * essential[i] * tmp; + } + } + } +} + +template +__device__ __forceinline__ void apply_householder_left_vector( + float* c, + int row0, + const float* essential, + int tail_len, + float tau) { + if (tau == 0.0f) return; + float tmp = c[row0]; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + tmp += essential[i] * c[row0 + 1 + i]; + } + } + c[row0] -= tau * tmp; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + c[row0 + 1 + i] -= tau * essential[i] * tmp; + } + } +} + +template +__device__ __forceinline__ void backsolve_upper_ranked( + const float* qr, + int rank, + const float* c, + const int* perm, + float* x_out) { + float y[N]; + #pragma unroll + for (int i = 0; i < N; ++i) { + y[i] = 0.0f; + x_out[i] = 0.0f; + } + for (int i = rank - 1; i >= 0; --i) { + float s = c[i]; + #pragma unroll + for (int j = 0; j < N; ++j) { + if (j > i && j < rank) { + s -= qr[i * N + j] * y[j]; + } + } + y[i] = s / qr[i * N + i]; + } + #pragma unroll + for (int i = 0; i < N; ++i) { + if (i < rank) { + x_out[perm[i]] = y[i]; + } else { + x_out[perm[i]] = 0.0f; + } + } +} + +template +__device__ __forceinline__ void cpqr_solve_small_impl( + const float* A_in, + const float* b_in, + float* x_out) { + float qr[N * N]; + float c[N]; + int perm[N]; + float col_norms_direct[N]; + float col_norms_updated[N]; + float essential[N > 1 ? N - 1 : 1]; + + #pragma unroll + for (int i = 0; i < N * N; ++i) { + qr[i] = A_in[i]; + } + #pragma unroll + for (int i = 0; i < N; ++i) { + c[i] = b_in[i]; + perm[i] = i; + x_out[i] = 0.0f; + } + + #pragma unroll + for (int j = 0; j < N; ++j) { + float norm_sq = 0.0f; + #pragma unroll + for (int r = 0; r < N; ++r) { + const float v = qr[r * N + j]; + norm_sq += v * v; + } + const float norm = sqrtf(norm_sq); + col_norms_direct[j] = norm; + col_norms_updated[j] = norm; + } + + float max_norm_updated = col_norms_updated[0]; + #pragma unroll + for (int j = 1; j < N; ++j) { + if (col_norms_updated[j] > max_norm_updated) { + max_norm_updated = col_norms_updated[j]; + } + } + + const float threshold_helper = (max_norm_updated * FLT_EPSILON) * (max_norm_updated * FLT_EPSILON) / float(N); + const float norm_downdate_threshold = sqrtf(FLT_EPSILON); + int nonzero_pivots = N; + float maxpivot = 0.0f; + + #pragma unroll + for (int k = 0; k < N; ++k) { + int biggest_col_index = k; + float best_updated = col_norms_updated[k]; + #pragma unroll + for (int j = 0; j < N; ++j) { + if (j > k && col_norms_updated[j] > best_updated) { + best_updated = col_norms_updated[j]; + biggest_col_index = j; + } + } + const float biggest_col_sq_norm = best_updated * best_updated; + if (nonzero_pivots == N && biggest_col_sq_norm < threshold_helper * float(N - k)) { + nonzero_pivots = k; + } + + swap_cols(qr, col_norms_updated, col_norms_direct, perm, k, biggest_col_index); + + const int tail_len = N - k - 1; + float tail_local[N > 1 ? N - 1 : 1]; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + tail_local[i] = (i < tail_len) ? qr[(k + 1 + i) * N + k] : 0.0f; + } + + float beta = 0.0f; + float tau = 0.0f; + make_householder_real(qr[k * N + k], tail_local, tail_len, &beta, &tau, essential); + + qr[k * N + k] = beta; + #pragma unroll + for (int i = 0; i < N - 1; ++i) { + if (i < tail_len) { + qr[(k + 1 + i) * N + k] = essential[i]; + } + } + const float abs_beta = absf(beta); + if (abs_beta > maxpivot) { + maxpivot = abs_beta; + } + + apply_householder_left_matrix(qr, k, k + 1, essential, tail_len, tau); + if (k < nonzero_pivots) { + apply_householder_left_vector(c, k, essential, tail_len, tau); + } + + #pragma unroll + for (int j = 0; j < N; ++j) { + if (j <= k) continue; + if (col_norms_updated[j] != 0.0f) { + float temp = absf(qr[k * N + j]) / col_norms_updated[j]; + temp = (1.0f + temp) * (1.0f - temp); + if (temp < 0.0f) temp = 0.0f; + const float ratio = col_norms_updated[j] / col_norms_direct[j]; + const float temp2 = temp * ratio * ratio; + if (temp2 <= norm_downdate_threshold) { + float norm_sq = 0.0f; + #pragma unroll + for (int r = 0; r < N; ++r) { + if (r > k) { + const float v = qr[r * N + j]; + norm_sq += v * v; + } + } + const float norm = sqrtf(norm_sq); + col_norms_direct[j] = norm; + col_norms_updated[j] = norm; + } else { + col_norms_updated[j] *= sqrtf(temp); + } + } + } + } + + if (nonzero_pivots == 0) { + #pragma unroll + for (int i = 0; i < N; ++i) { + x_out[i] = 0.0f; + } + return; + } + + backsolve_upper_ranked(qr, nonzero_pivots, c, perm, x_out); +} + +} // namespace detail + +__device__ __forceinline__ void cpqr_solve_3x3(const float A[9], const float b[3], float x[3]) { + detail::cpqr_solve_small_impl<3>(A, b, x); +} + +__device__ __forceinline__ void cpqr_solve_2x2(const float A[4], const float b[2], float x[2]) { + detail::cpqr_solve_small_impl<2>(A, b, x); +} + +__device__ __forceinline__ float solve_1x1_unchecked(float a, float rhs) { + return rhs / a; +} + +} // namespace fdg_gpu::small_cpqr diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.cu new file mode 100644 index 00000000..71549fab --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.cu @@ -0,0 +1,766 @@ +#include "flexible_dual_grid_gpu.h" + +#include "intersection_qef.h" +#include "voxel_traverse_edge_dda.h" +#include "voxelize_mesh_oct.h" + +#include +#include +#include "fdg_gpu_small_cpqr_device.cuh" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace fdg_gpu { +namespace { + +constexpr int kBlockSize = 128; + +struct FlatTriangles { + int64_t num_triangles = 0; + DeviceBuffer triangles; // [3 * num_triangles, 3] +}; + +struct BoundaryEdgeIndexResult { + int64_t size = 0; + DeviceBuffer edge_vertex_ids; // [size, 2] +}; + +struct BoundarySegments { + int64_t size = 0; + DeviceBuffer segments; // [2 * size, 3] +}; + +struct IsOne { + __host__ __device__ bool operator()(int v) const { return v == 1; } +}; + +__host__ __device__ __forceinline__ uint64_t pack_edge_key(int a, int b) { + return (static_cast(static_cast(a)) << 32) | + static_cast(b); +} + +__host__ __device__ __forceinline__ int edge_key_v0(uint64_t key) { + return static_cast(key >> 32); +} + +__host__ __device__ __forceinline__ int edge_key_v1(uint64_t key) { + return static_cast(key & 0xffffffffu); +} + +__host__ __device__ __forceinline__ SymQEF10 sym10_zero() { + return SymQEF10{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; +} + +__host__ __device__ __forceinline__ SymQEF10 sym10_add(const SymQEF10& a, const SymQEF10& b) { + return SymQEF10{ + a.q00 + b.q00, + a.q01 + b.q01, + a.q02 + b.q02, + a.q03 + b.q03, + a.q11 + b.q11, + a.q12 + b.q12, + a.q13 + b.q13, + a.q22 + b.q22, + a.q23 + b.q23, + a.q33 + b.q33, + }; +} + +__host__ __device__ __forceinline__ SymQEF10 sym10_scale(const SymQEF10& a, float s) { + return SymQEF10{ + a.q00 * s, + a.q01 * s, + a.q02 * s, + a.q03 * s, + a.q11 * s, + a.q12 * s, + a.q13 * s, + a.q22 * s, + a.q23 * s, + a.q33 * s, + }; +} + +__global__ void gather_flat_triangles_kernel( + const float* vertices, + const int32_t* faces, + int64_t num_faces, + float* triangles) { + const int64_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid >= 3 * num_faces) return; + + const int64_t f = tid / 3; + const int lv = static_cast(tid % 3); + const int32_t vid = faces[3 * f + lv]; + + triangles[3 * tid + 0] = vertices[3 * vid + 0]; + triangles[3 * tid + 1] = vertices[3 * vid + 1]; + triangles[3 * tid + 2] = vertices[3 * vid + 2]; +} + +__global__ void emit_face_edges_kernel( + const int32_t* faces, + int64_t num_faces, + uint64_t* edge_keys) { + const int64_t f = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (f >= num_faces) return; + + int32_t e00 = faces[3 * f + 0]; + int32_t e01 = faces[3 * f + 1]; + int32_t e10 = faces[3 * f + 1]; + int32_t e11 = faces[3 * f + 2]; + int32_t e20 = faces[3 * f + 2]; + int32_t e21 = faces[3 * f + 0]; + + if (e00 > e01) { const int32_t t = e00; e00 = e01; e01 = t; } + if (e10 > e11) { const int32_t t = e10; e10 = e11; e11 = t; } + if (e20 > e21) { const int32_t t = e20; e20 = e21; e21 = t; } + + edge_keys[3 * f + 0] = pack_edge_key(e00, e01); + edge_keys[3 * f + 1] = pack_edge_key(e10, e11); + edge_keys[3 * f + 2] = pack_edge_key(e20, e21); +} + +__global__ void unpack_boundary_keys_kernel( + const uint64_t* boundary_keys, + int64_t size, + int32_t* edge_vertex_ids) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= size) return; + + const uint64_t key = boundary_keys[i]; + edge_vertex_ids[2 * i + 0] = edge_key_v0(key); + edge_vertex_ids[2 * i + 1] = edge_key_v1(key); +} + +__global__ void gather_boundary_segments_kernel( + const float* vertices, + const int32_t* edge_vertex_ids, + int64_t num_boundary_edges, + float* segments) { + const int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_boundary_edges) return; + + const int32_t v0 = edge_vertex_ids[2 * eid + 0]; + const int32_t v1 = edge_vertex_ids[2 * eid + 1]; + + segments[6 * eid + 0] = vertices[3 * v0 + 0]; + segments[6 * eid + 1] = vertices[3 * v0 + 1]; + segments[6 * eid + 2] = vertices[3 * v0 + 2]; + segments[6 * eid + 3] = vertices[3 * v1 + 0]; + segments[6 * eid + 4] = vertices[3 * v1 + 1]; + segments[6 * eid + 5] = vertices[3 * v1 + 2]; +} + +__global__ void zero_qef_kernel(SymQEF10* qefs, int64_t n) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + qefs[i] = sym10_zero(); +} + +__global__ void sum_qef_kernel( + const SymQEF10* qef_init, + const SymQEF10* qef_face, + const SymQEF10* qef_boundary, + int64_t n, + float face_weight, + SymQEF10* qef_total) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + + SymQEF10 q = qef_init[i]; + q = sym10_add(q, sym10_scale(qef_face[i], face_weight)); + q = sym10_add(q, qef_boundary[i]); + qef_total[i] = q; +} + +__global__ void unpack_intersected_kernel( + const uint8_t* intersected_mask, + int64_t n, + bool* intersected_bool) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + + const uint8_t m = intersected_mask[i]; + intersected_bool[3 * i + 0] = (m & (1u << 0)) != 0; + intersected_bool[3 * i + 1] = (m & (1u << 1)) != 0; + intersected_bool[3 * i + 2] = (m & (1u << 2)) != 0; +} + + +__device__ __forceinline__ int idx4(int r, int c) { return r * 4 + c; } +__device__ __forceinline__ int idx3(int r, int c) { return r * 3 + c; } +__device__ __forceinline__ int idx2(int r, int c) { return r * 2 + c; } + +__device__ __forceinline__ void sym10_to_dense4x4(const SymQEF10& q, float Q[16]) { + Q[idx4(0,0)] = q.q00; Q[idx4(0,1)] = q.q01; Q[idx4(0,2)] = q.q02; Q[idx4(0,3)] = q.q03; + Q[idx4(1,0)] = q.q01; Q[idx4(1,1)] = q.q11; Q[idx4(1,2)] = q.q12; Q[idx4(1,3)] = q.q13; + Q[idx4(2,0)] = q.q02; Q[idx4(2,1)] = q.q12; Q[idx4(2,2)] = q.q22; Q[idx4(2,3)] = q.q23; + Q[idx4(3,0)] = q.q03; Q[idx4(3,1)] = q.q13; Q[idx4(3,2)] = q.q23; Q[idx4(3,3)] = q.q33; +} + +__device__ __forceinline__ bool point_inside_box3( + const float v[3], + const float min_corner[3], + const float max_corner[3]) { + return ( + v[0] >= min_corner[0] && v[0] <= max_corner[0] && + v[1] >= min_corner[1] && v[1] <= max_corner[1] && + v[2] >= min_corner[2] && v[2] <= max_corner[2]); +} + +__device__ __forceinline__ float qef_error4(const float Q[16], const float p[4]) { + const float y0 = Q[idx4(0,0)] * p[0] + Q[idx4(0,1)] * p[1] + Q[idx4(0,2)] * p[2] + Q[idx4(0,3)] * p[3]; + const float y1 = Q[idx4(1,0)] * p[0] + Q[idx4(1,1)] * p[1] + Q[idx4(1,2)] * p[2] + Q[idx4(1,3)] * p[3]; + const float y2 = Q[idx4(2,0)] * p[0] + Q[idx4(2,1)] * p[1] + Q[idx4(2,2)] * p[2] + Q[idx4(2,3)] * p[3]; + const float y3 = Q[idx4(3,0)] * p[0] + Q[idx4(3,1)] * p[1] + Q[idx4(3,2)] * p[2] + Q[idx4(3,3)] * p[3]; + return p[0] * y0 + p[1] * y1 + p[2] * y2 + p[3] * y3; +} + +__device__ __forceinline__ void add_qef_regularization_inplace( + float Q[16], + const float mean_sum[3], + float cnt, + float regularization_weight) { + if (regularization_weight <= 0.0f || cnt <= 0.0f) { + return; + } + + const float px = mean_sum[0] / cnt; + const float py = mean_sum[1] / cnt; + const float pz = mean_sum[2] / cnt; + const float w = regularization_weight * cnt; + + Q[idx4(0,0)] += w; + Q[idx4(1,1)] += w; + Q[idx4(2,2)] += w; + + Q[idx4(0,3)] += -w * px; + Q[idx4(1,3)] += -w * py; + Q[idx4(2,3)] += -w * pz; + + Q[idx4(3,0)] += -w * px; + Q[idx4(3,1)] += -w * py; + Q[idx4(3,2)] += -w * pz; + + Q[idx4(3,3)] += w * (px * px + py * py + pz * pz); +} + +__device__ __forceinline__ void try_single_constraint( + const float Q[16], + int fixed_axis, + const float min_corner[3], + const float max_corner[3], + float& best, + float v_new[3]) { + const int ax1 = (fixed_axis + 1) % 3; + const int ax2 = (fixed_axis + 2) % 3; + + float A2[4]; + float B2[4]; + float q2[2]; + float rhs2[2]; + float x2[2]; + + A2[idx2(0,0)] = Q[idx4(ax1, ax1)]; + A2[idx2(0,1)] = Q[idx4(ax1, ax2)]; + A2[idx2(1,0)] = Q[idx4(ax2, ax1)]; + A2[idx2(1,1)] = Q[idx4(ax2, ax2)]; + + B2[idx2(0,0)] = Q[idx4(ax1, fixed_axis)]; + B2[idx2(0,1)] = Q[idx4(ax1, 3)]; + B2[idx2(1,0)] = Q[idx4(ax2, fixed_axis)]; + B2[idx2(1,1)] = Q[idx4(ax2, 3)]; + + q2[0] = min_corner[fixed_axis]; + q2[1] = 1.0f; + rhs2[0] = -(B2[idx2(0,0)] * q2[0] + B2[idx2(0,1)] * q2[1]); + rhs2[1] = -(B2[idx2(1,0)] * q2[0] + B2[idx2(1,1)] * q2[1]); + fdg_gpu::small_cpqr::cpqr_solve_2x2(A2, rhs2, x2); + if (x2[0] >= min_corner[ax1] && x2[0] <= max_corner[ax1] && + x2[1] >= min_corner[ax2] && x2[1] <= max_corner[ax2]) { + float p4[4]; + p4[fixed_axis] = min_corner[fixed_axis]; + p4[ax1] = x2[0]; + p4[ax2] = x2[1]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + + q2[0] = max_corner[fixed_axis]; + q2[1] = 1.0f; + rhs2[0] = -(B2[idx2(0,0)] * q2[0] + B2[idx2(0,1)] * q2[1]); + rhs2[1] = -(B2[idx2(1,0)] * q2[0] + B2[idx2(1,1)] * q2[1]); + fdg_gpu::small_cpqr::cpqr_solve_2x2(A2, rhs2, x2); + if (x2[0] >= min_corner[ax1] && x2[0] <= max_corner[ax1] && + x2[1] >= min_corner[ax2] && x2[1] <= max_corner[ax2]) { + float p4[4]; + p4[fixed_axis] = max_corner[fixed_axis]; + p4[ax1] = x2[0]; + p4[ax2] = x2[1]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } +} + +__device__ __forceinline__ void try_two_constraint( + const float Q[16], + int free_axis, + const float min_corner[3], + const float max_corner[3], + float& best, + float v_new[3]) { + const int ax1 = (free_axis + 1) % 3; + const int ax2 = (free_axis + 2) % 3; + + const float a = Q[idx4(free_axis, free_axis)]; + const float b0 = Q[idx4(free_axis, ax1)]; + const float b1 = Q[idx4(free_axis, ax2)]; + const float b2 = Q[idx4(free_axis, 3)]; + + float rhs = -(b0 * min_corner[ax1] + b1 * min_corner[ax2] + b2); + float x = fdg_gpu::small_cpqr::solve_1x1_unchecked(a, rhs); + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + float p4[4]; + p4[free_axis] = x; + p4[ax1] = min_corner[ax1]; + p4[ax2] = min_corner[ax2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + + rhs = -(b0 * min_corner[ax1] + b1 * max_corner[ax2] + b2); + x = fdg_gpu::small_cpqr::solve_1x1_unchecked(a, rhs); + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + float p4[4]; + p4[free_axis] = x; + p4[ax1] = min_corner[ax1]; + p4[ax2] = max_corner[ax2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + + rhs = -(b0 * max_corner[ax1] + b1 * min_corner[ax2] + b2); + x = fdg_gpu::small_cpqr::solve_1x1_unchecked(a, rhs); + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + float p4[4]; + p4[free_axis] = x; + p4[ax1] = max_corner[ax1]; + p4[ax2] = min_corner[ax2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + + rhs = -(b0 * max_corner[ax1] + b1 * max_corner[ax2] + b2); + x = fdg_gpu::small_cpqr::solve_1x1_unchecked(a, rhs); + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) { + float p4[4]; + p4[free_axis] = x; + p4[ax1] = max_corner[ax1]; + p4[ax2] = max_corner[ax2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } +} + +__device__ __forceinline__ void try_three_constraint( + const float Q[16], + const float min_corner[3], + const float max_corner[3], + float& best, + float v_new[3]) { + for (int x_constraint = 0; x_constraint < 2; ++x_constraint) { + for (int y_constraint = 0; y_constraint < 2; ++y_constraint) { + for (int z_constraint = 0; z_constraint < 2; ++z_constraint) { + float p4[4]; + p4[0] = x_constraint ? min_corner[0] : max_corner[0]; + p4[1] = y_constraint ? min_corner[1] : max_corner[1]; + p4[2] = z_constraint ? min_corner[2] : max_corner[2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + } + } +} + +__global__ void solve_qef_full_kernel( + const int* voxel_coords, + const float* mean_sum, + const float* cnt, + const SymQEF10* qef_total, + int64_t n, + float3 voxel_size, + float regularization_weight, + float* dual_vertices) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + + const int x = voxel_coords[3 * i + 0]; + const int y = voxel_coords[3 * i + 1]; + const int z = voxel_coords[3 * i + 2]; + + float min_corner[3] = { + x * voxel_size.x, + y * voxel_size.y, + z * voxel_size.z, + }; + float max_corner[3] = { + (x + 1) * voxel_size.x, + (y + 1) * voxel_size.y, + (z + 1) * voxel_size.z, + }; + + float Q[16]; + sym10_to_dense4x4(qef_total[i], Q); + + const float mean_i[3] = { + mean_sum[3 * i + 0], + mean_sum[3 * i + 1], + mean_sum[3 * i + 2], + }; + add_qef_regularization_inplace(Q, mean_i, cnt[i], regularization_weight); + + float A3[9]; + float b3[3]; + float v_new[3]; + + A3[idx3(0,0)] = Q[idx4(0,0)]; + A3[idx3(0,1)] = Q[idx4(0,1)]; + A3[idx3(0,2)] = Q[idx4(0,2)]; + A3[idx3(1,0)] = Q[idx4(1,0)]; + A3[idx3(1,1)] = Q[idx4(1,1)]; + A3[idx3(1,2)] = Q[idx4(1,2)]; + A3[idx3(2,0)] = Q[idx4(2,0)]; + A3[idx3(2,1)] = Q[idx4(2,1)]; + A3[idx3(2,2)] = Q[idx4(2,2)]; + + b3[0] = -Q[idx4(0,3)]; + b3[1] = -Q[idx4(1,3)]; + b3[2] = -Q[idx4(2,3)]; + + fdg_gpu::small_cpqr::cpqr_solve_3x3(A3, b3, v_new); + + if (!point_inside_box3(v_new, min_corner, max_corner)) { + float best = CUDART_INF_F; + try_single_constraint(Q, 0, min_corner, max_corner, best, v_new); + try_single_constraint(Q, 1, min_corner, max_corner, best, v_new); + try_single_constraint(Q, 2, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 0, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 1, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 2, min_corner, max_corner, best, v_new); + try_three_constraint(Q, min_corner, max_corner, best, v_new); + } + + dual_vertices[3 * i + 0] = v_new[0]; + dual_vertices[3 * i + 1] = v_new[1]; + dual_vertices[3 * i + 2] = v_new[2]; +} + +inline FlatTriangles build_flat_triangles_gpu( + const float* vertices, + const int32_t* faces, + int64_t num_faces, + cudaStream_t stream) { + FlatTriangles out; + out.num_triangles = num_faces; + out.triangles.allocate(9 * num_faces); + if (num_faces == 0) return out; + + gather_flat_triangles_kernel<<>>( + vertices, + faces, + num_faces, + out.triangles.data()); + throw_cuda_error(cudaGetLastError(), "gather_flat_triangles_kernel"); + return out; +} + +inline BoundaryEdgeIndexResult detect_boundary_edges_gpu( + const int32_t* faces, + int64_t num_faces, + cudaStream_t stream) { + BoundaryEdgeIndexResult out; + if (num_faces == 0) return out; + + DeviceBuffer edge_keys(3 * num_faces); + emit_face_edges_kernel<<>>( + faces, + num_faces, + edge_keys.data()); + throw_cuda_error(cudaGetLastError(), "emit_face_edges_kernel"); + + auto policy = thrust::cuda::par.on(stream); + auto edge_keys_begin = thrust::device_pointer_cast(edge_keys.data()); + thrust::sort(policy, edge_keys_begin, edge_keys_begin + 3 * num_faces); + + DeviceBuffer unique_keys(3 * num_faces); + DeviceBuffer counts(3 * num_faces); + auto reduce_end = thrust::reduce_by_key( + policy, + edge_keys_begin, + edge_keys_begin + 3 * num_faces, + thrust::make_constant_iterator(1), + thrust::device_pointer_cast(unique_keys.data()), + thrust::device_pointer_cast(counts.data())); + const int64_t unique_size = reduce_end.first - thrust::device_pointer_cast(unique_keys.data()); + if (unique_size == 0) return out; + + const int64_t boundary_count = thrust::count_if( + policy, + thrust::device_pointer_cast(counts.data()), + thrust::device_pointer_cast(counts.data()) + unique_size, + IsOne{}); + out.size = boundary_count; + out.edge_vertex_ids.allocate(2 * boundary_count); + if (boundary_count == 0) return out; + + DeviceBuffer boundary_keys(boundary_count); + auto copied_end = thrust::copy_if( + policy, + thrust::device_pointer_cast(unique_keys.data()), + thrust::device_pointer_cast(unique_keys.data()) + unique_size, + thrust::device_pointer_cast(counts.data()), + thrust::device_pointer_cast(boundary_keys.data()), + IsOne{}); + const int64_t copied = copied_end - thrust::device_pointer_cast(boundary_keys.data()); + if (copied != boundary_count) { + throw std::runtime_error("boundary edge count mismatch"); + } + + unpack_boundary_keys_kernel<<>>( + boundary_keys.data(), + boundary_count, + out.edge_vertex_ids.data()); + throw_cuda_error(cudaGetLastError(), "unpack_boundary_keys_kernel"); + return out; +} + +inline BoundarySegments gather_boundary_segments_gpu( + const float* vertices, + const BoundaryEdgeIndexResult& boundary_edges, + cudaStream_t stream) { + BoundarySegments out; + out.size = boundary_edges.size; + out.segments.allocate(6 * out.size); + if (out.size == 0) return out; + + gather_boundary_segments_kernel<<>>( + vertices, + boundary_edges.edge_vertex_ids.data(), + out.size, + out.segments.data()); + throw_cuda_error(cudaGetLastError(), "gather_boundary_segments_kernel"); + return out; +} + +inline DeviceBuffer make_zero_qef_buffer(int64_t n, cudaStream_t stream) { + DeviceBuffer out(n); + if (n > 0) { + zero_qef_kernel<<>>(out.data(), n); + throw_cuda_error(cudaGetLastError(), "zero_qef_kernel"); + } + return out; +} + +} // namespace + +cudaError_t mesh_to_flexible_dual_grid_gpu( + const float* vertices, + int64_t num_vertices, + const int32_t* faces, + int64_t num_faces, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + float face_weight, + float boundary_weight, + float regularization_weight, + int64_t intersect_chunk_triangles, + int boundary_chunk_steps, + cudaStream_t stream, + FlexibleDualGridGPUOutput* out) { + if (out == nullptr) { + return cudaErrorInvalidValue; + } + out->size = 0; + out->voxel_coords = nullptr; + out->dual_vertices = nullptr; + out->intersected = nullptr; + + if (num_vertices < 0 || num_faces < 0) { + return cudaErrorInvalidValue; + } + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + return cudaErrorInvalidValue; + } + if (grid_max.x <= grid_min.x || grid_max.y <= grid_min.y || grid_max.z <= grid_min.z) { + return cudaErrorInvalidValue; + } + if (num_vertices > 0 && vertices == nullptr) { + return cudaErrorInvalidValue; + } + if (num_faces > 0 && faces == nullptr) { + return cudaErrorInvalidValue; + } + if (intersect_chunk_triangles <= 0 || boundary_chunk_steps <= 0) { + return cudaErrorInvalidValue; + } + + try { + FlatTriangles flat = build_flat_triangles_gpu(vertices, faces, num_faces, stream); + + auto surface = intersection_qef::intersect_qef_gpu( + flat.triangles.data(), + flat.num_triangles, + voxel_size, + grid_min, + grid_max, + intersect_chunk_triangles, + stream); + + auto face_qefs = make_zero_qef_buffer(surface.size, stream); + if (surface.size > 0 && face_weight > 0.0f) { + auto face_result = oct_pairs::face_qef_gpu( + voxel_size, + grid_min, + grid_max, + flat.triangles.data(), + flat.num_triangles, + surface.voxels.data(), + surface.size, + stream); + face_qefs = std::move(face_result.qefs); + } + + auto boundary_qefs = make_zero_qef_buffer(surface.size, stream); + if (surface.size > 0 && boundary_weight > 0.0f) { + BoundaryEdgeIndexResult boundary_edges = detect_boundary_edges_gpu(faces, num_faces, stream); + BoundarySegments boundary_segments = gather_boundary_segments_gpu(vertices, boundary_edges, stream); + auto boundary_result = edge_dda::boundary_qef_gpu( + voxel_size, + grid_min, + grid_max, + boundary_segments.segments.data(), + boundary_segments.size, + boundary_weight, + surface.voxels.data(), + surface.size, + boundary_chunk_steps, + stream); + boundary_qefs = std::move(boundary_result.qefs); + } + + DeviceBuffer qef_total(surface.size); + if (surface.size > 0) { + sum_qef_kernel<<>>( + surface.qefs.data(), + face_qefs.data(), + boundary_qefs.data(), + surface.size, + face_weight, + qef_total.data()); + throw_cuda_error(cudaGetLastError(), "sum_qef_kernel"); + } + + DeviceBuffer dual_vertices(3 * surface.size); + if (surface.size > 0) { + solve_qef_full_kernel<<>>( + surface.voxels.data(), + surface.mean_sum.data(), + surface.cnt.data(), + qef_total.data(), + surface.size, + voxel_size, + regularization_weight, + dual_vertices.data()); + throw_cuda_error(cudaGetLastError(), "solve_qef_full_kernel"); + } + + DeviceBuffer intersected_bool(3 * surface.size); + if (surface.size > 0) { + unpack_intersected_kernel<<>>( + surface.intersected.data(), + surface.size, + intersected_bool.data()); + throw_cuda_error(cudaGetLastError(), "unpack_intersected_kernel"); + } + + out->size = surface.size; + out->voxel_coords = surface.voxels.release_ownership(); + out->dual_vertices = dual_vertices.release_ownership(); + out->intersected = intersected_bool.release_ownership(); + return cudaSuccess; + } catch (const std::bad_alloc&) { + return cudaErrorMemoryAllocation; + } catch (const std::invalid_argument&) { + return cudaErrorInvalidValue; + } catch (const std::exception&) { + return cudaErrorUnknown; + } +} + +void free_flexible_dual_grid_gpu_output(FlexibleDualGridGPUOutput* out) noexcept { + if (out == nullptr) return; + if (out->voxel_coords) cudaFree(out->voxel_coords); + if (out->dual_vertices) cudaFree(out->dual_vertices); + if (out->intersected) cudaFree(out->intersected); + out->size = 0; + out->voxel_coords = nullptr; + out->dual_vertices = nullptr; + out->intersected = nullptr; +} + +} // namespace fdg_gpu diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.h b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.h new file mode 100644 index 00000000..04b6729c --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/flexible_dual_grid_gpu.h @@ -0,0 +1,35 @@ +#pragma once + +#include "fdg_gpu_common.h" + +#include +#include + +namespace fdg_gpu { + +struct FlexibleDualGridGPUOutput { + int64_t size = 0; + int32_t* voxel_coords = nullptr; // [size, 3] + float* dual_vertices = nullptr; // [size, 3] + bool* intersected = nullptr; // [size, 3] +}; + +cudaError_t mesh_to_flexible_dual_grid_gpu( + const float* vertices, + int64_t num_vertices, + const int32_t* faces, + int64_t num_faces, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + float face_weight, + float boundary_weight, + float regularization_weight, + int64_t intersect_chunk_triangles, + int boundary_chunk_steps, + cudaStream_t stream, + FlexibleDualGridGPUOutput* out); + +void free_flexible_dual_grid_gpu_output(FlexibleDualGridGPUOutput* out) noexcept; + +} // namespace fdg_gpu diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersect_qef.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersect_qef.cu new file mode 100644 index 00000000..769f81d3 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersect_qef.cu @@ -0,0 +1,1013 @@ +#include "../api.h" + +#include "qef.cuh" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +// Triangle/voxel intersection is expressed as a stream of small scan tasks. +// A large triangle can cover many grid cells, so each axis projection is split +// into 16x16 tiles; GPU threads then process tiles rather than whole triangles. +// The occupancy pass first marks active voxels in brick bitsets and only later +// compacts those bits into the final voxel rows used by QEF arrays. +namespace o_voxel::fdg +{ + namespace + { + + constexpr int kThreads = 256; + constexpr int kTileU = 16; + constexpr int kTileV = 16; + + // One triangle-axis tile. u/v are the two scan axes for the chosen axis, + // and the half-open range [u0, u1) x [v0, v1) bounds the work item. + struct ScanTask + { + int32_t tri_id; + int32_t axis; + int32_t u0; + int32_t u1; + int32_t v0; + int32_t v1; + }; + + // Tensor-owned result of the shared occupancy build. intersect_occ_cuda + // returns only voxels; intersect_qef_cuda also reuses tasks and brick + // lookup tensors to accumulate intersection QEFs. + struct IntersectionOccupancy + { + torch::Tensor tasks; + torch::Tensor hash_keys; + torch::Tensor hash_vals; + torch::Tensor brick_coords; + torch::Tensor brick_bits; + torch::Tensor brick_base; + torch::Tensor voxels; + torch::Tensor overflow_flag; + int64_t num_tasks = 0; + int64_t num_bricks = 0; + int64_t num_voxels = 0; + uint64_t hash_capacity = 0; + }; + + __host__ __device__ __forceinline__ int div_up_i32(int n, int d) + { + return (n + d - 1) / d; + } + + __host__ __device__ __forceinline__ int64_t div_up_i64(int64_t n, int64_t d) + { + return (n + d - 1) / d; + } + + __device__ __forceinline__ int clamp_int(int v, int lo, int hi) + { + return max(min(v, hi), lo); + } + + __device__ __forceinline__ uint64_t mix64(uint64_t x) + { + x ^= x >> 33; + x *= 0xff51afd7ed558ccdULL; + x ^= x >> 33; + x *= 0xc4ceb9fe1a85ec53ULL; + x ^= x >> 33; + return x; + } + + __device__ bool compute_scan_bbox( + const float *tri, + int axis, + GridSpec grid, + int &u0, + int &u1, + int &v0, + int &v1) + { + // For one depth axis, scan the triangle in the other two axes. The + // returned u/v box is conservative because each scan event later + // touches a 2x2 voxel neighborhood in the projection plane. + const int ax0 = (axis + 1) % 3; + const int ax1 = (axis + 2) % 3; + const float3 voxel_size = grid.voxel_size; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const float vs[3] = {voxel_size.x, voxel_size.y, voxel_size.z}; + float min_u = tri[ax0]; + float max_u = tri[ax0]; + float min_v = tri[ax1]; + float max_v = tri[ax1]; + for (int i = 1; i < 3; ++i) + { + const float u = tri[3 * i + ax0]; + const float v = tri[3 * i + ax1]; + min_u = fminf(min_u, u); + max_u = fmaxf(max_u, u); + min_v = fminf(min_v, v); + max_v = fmaxf(max_v, v); + } + + // u is expanded by one cell on the low side and two on the high + // side to preserve scanline event coverage near triangle edges and + // voxel boundaries. + u0 = clamp_int(static_cast(min_u / vs[ax0]) - 1, grid_min[ax0], grid_max[ax0] - 1); + u1 = clamp_int(static_cast(max_u / vs[ax0]) + 2, grid_min[ax0], grid_max[ax0] - 1); + v0 = clamp_int(static_cast(min_v / vs[ax1]), grid_min[ax1], grid_max[ax1] - 1); + v1 = clamp_int(static_cast(max_v / vs[ax1]), grid_min[ax1], grid_max[ax1] - 1); + return u1 > u0 && v1 > v0; + } + + __device__ int64_t task_brick_bound( + const float *tri, + const ScanTask &task, + GridSpec grid) + { + // This is an allocation bound, not the exact output count. For the + // tile's u/v range, estimate the depth interval covered by the + // triangle plane, then count all bricks overlapped by that 3D box. + const int axis = task.axis; + const int ax0 = (axis + 1) % 3; + const int ax1 = (axis + 2) % 3; + const int ax2 = axis; + const float3 voxel_size = grid.voxel_size; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const float vs[3] = {voxel_size.x, voxel_size.y, voxel_size.z}; + int lo[3]; + int hi[3]; + lo[ax0] = max(task.u0, grid_min[ax0]); + hi[ax0] = min(task.u1 + 1, grid_max[ax0]); + lo[ax1] = max(task.v0, grid_min[ax1]); + hi[ax1] = min(task.v1 + 1, grid_max[ax1]); + + const double u0 = tri[ax0]; + const double v0 = tri[ax1]; + const double z0 = tri[ax2]; + const double u1 = tri[3 + ax0]; + const double v1 = tri[3 + ax1]; + const double z1 = tri[3 + ax2]; + const double u2 = tri[6 + ax0]; + const double v2 = tri[6 + ax1]; + const double z2 = tri[6 + ax2]; + const double denom = (u1 - u0) * (v2 - v0) - (u2 - u0) * (v1 - v0); + double z_min = fmin(fmin(z0, z1), z2); + double z_max = fmax(fmax(z0, z1), z2); + if (fabs(denom) > 1e-20) + { + // z = a*u + b*v + c is the triangle plane written over the + // current projection. Evaluating the four tile corners gives a + // conservative depth interval for all scan events in this tile. + const double a = ((z1 - z0) * (v2 - v0) - (z2 - z0) * (v1 - v0)) / denom; + const double b = ((u1 - u0) * (z2 - z0) - (u2 - u0) * (z1 - z0)) / denom; + const double c = z0 - a * u0 - b * v0; + const double ru0 = static_cast(task.u0) * vs[ax0]; + const double ru1 = static_cast(task.u1 + 1) * vs[ax0]; + const double rv0 = static_cast(task.v0) * vs[ax1]; + const double rv1 = static_cast(task.v1 + 1) * vs[ax1]; + const double d0 = a * ru0 + b * rv0 + c; + const double d1 = a * ru0 + b * rv1 + c; + const double d2 = a * ru1 + b * rv0 + c; + const double d3 = a * ru1 + b * rv1 + c; + z_min = fmin(fmin(d0, d1), fmin(d2, d3)); + z_max = fmax(fmax(d0, d1), fmax(d2, d3)); + } + + lo[ax2] = clamp_int(static_cast(z_min / vs[ax2]) - 2, grid_min[ax2], grid_max[ax2]); + hi[ax2] = clamp_int(static_cast(z_max / vs[ax2]) + 3, grid_min[ax2], grid_max[ax2]); + if (hi[0] <= lo[0] || hi[1] <= lo[1] || hi[2] <= lo[2]) + return 0; + + // Convert the conservative voxel box into a conservative brick box. + // Summing these bounds lets the host allocate hash/bitset storage + // once, with overflow treated as a bug guard. + const int64_t bx0 = (lo[0] - grid_min.x) / kBrickSize; + const int64_t by0 = (lo[1] - grid_min.y) / kBrickSize; + const int64_t bz0 = (lo[2] - grid_min.z) / kBrickSize; + const int64_t bx1 = div_up_i64(hi[0] - grid_min.x, kBrickSize); + const int64_t by1 = div_up_i64(hi[1] - grid_min.y, kBrickSize); + const int64_t bz1 = div_up_i64(hi[2] - grid_min.z, kBrickSize); + return (bx1 - bx0) * (by1 - by0) * (bz1 - bz0); + } + + template + __device__ void scan_triangle_events_tiled( + const float *tri, + const ScanTask &task, + GridSpec grid, + Emit emit) + { + // Shared scanline generator used by occupancy and QEF passes. The + // template callback keeps both passes on the exact same event stream. + const int ax2 = task.axis; + const int ax0 = (ax2 + 1) % 3; + const int ax1 = (ax2 + 2) % 3; + const float3 voxel_size = grid.voxel_size; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + double t[3][3] = { + {static_cast(tri[ax0]), static_cast(tri[ax1]), static_cast(tri[ax2])}, + {static_cast(tri[3 + ax0]), static_cast(tri[3 + ax1]), static_cast(tri[3 + ax2])}, + {static_cast(tri[6 + ax0]), static_cast(tri[6 + ax1]), static_cast(tri[6 + ax2])}, + }; + int order[3] = {0, 1, 2}; + // Sort vertices by the scan row axis so the triangle can be scanned + // as two monotonic halves: top->middle and middle->bottom. + if (t[order[0]][1] > t[order[1]][1]) + { + const int tmp = order[0]; + order[0] = order[1]; + order[1] = tmp; + } + if (t[order[1]][1] > t[order[2]][1]) + { + const int tmp = order[1]; + order[1] = order[2]; + order[2] = tmp; + } + if (t[order[0]][1] > t[order[1]][1]) + { + const int tmp = order[0]; + order[0] = order[1]; + order[1] = tmp; + } + + const double *t0 = t[order[0]]; + const double *t1 = t[order[1]]; + const double *t2 = t[order[2]]; + const float vs[3] = {voxel_size.x, voxel_size.y, voxel_size.z}; + const int start = max(min(static_cast(t0[1] / vs[ax1]), grid_max[ax1] - 1), grid_min[ax1]); + const int mid = max(min(static_cast(t1[1] / vs[ax1]), grid_max[ax1] - 1), grid_min[ax1]); + const int end = max(min(static_cast(t2[1] / vs[ax1]), grid_max[ax1] - 1), grid_min[ax1]); + + auto scan_half = [&](int row_start, int row_end, const double *a, const double *b, const double *c) + { + // For each scan row, intersect the row with two triangle edges, + // then interpolate along the horizontal span to recover depth. + row_start = max(row_start, task.v0); + row_end = min(row_end, task.v1); + for (int y_idx = row_start; y_idx < row_end; ++y_idx) + { + // y and x use the high cell boundary, matching the original + // event placement for voxel face crossings. + const double y = (static_cast(y_idx) + 1.0) * vs[ax1]; + const double ab = fabs(a[1] - b[1]) < 1e-12 ? 0.0 : (y - a[1]) / (b[1] - a[1]); + const double ac = fabs(a[1] - c[1]) < 1e-12 ? 0.0 : (y - a[1]) / (c[1] - a[1]); + double t3x = (1.0 - ab) * a[0] + ab * b[0]; + double t3z = (1.0 - ab) * a[2] + ab * b[2]; + double t4x = (1.0 - ac) * a[0] + ac * c[0]; + double t4z = (1.0 - ac) * a[2] + ac * c[2]; + if (t3x > t4x) + { + double tmp = t3x; + t3x = t4x; + t4x = tmp; + tmp = t3z; + t3z = t4z; + t4z = tmp; + } + + int line_start = max(min(static_cast(t3x / vs[ax0]), grid_max[ax0] - 1), grid_min[ax0]); + int line_end = max(min(static_cast(t4x / vs[ax0]), grid_max[ax0] - 1), grid_min[ax0]); + line_start = max(line_start, task.u0); + line_end = min(line_end, task.u1); + for (int x_idx = line_start; x_idx < line_end; ++x_idx) + { + const double x = (static_cast(x_idx) + 1.0) * vs[ax0]; + // alpha moves across the row segment; z is the point + // where the triangle plane crosses this projected event. + const double alpha = fabs(t4x - t3x) < 1e-12 ? 0.0 : (x - t3x) / (t4x - t3x); + const double z = (1.0 - alpha) * t3z + alpha * t4z; + const int z_idx = static_cast(z / vs[ax2]); + if (z_idx >= grid_min[ax2] && z_idx < grid_max[ax2]) + emit(ax0, ax1, ax2, x_idx, y_idx, z_idx, x, y, z); + } + } + }; + scan_half(start, mid, t0, t1, t2); + scan_half(mid, end, t2, t1, t0); + } + + __device__ __forceinline__ uint64_t voxel_to_brick( + int x, + int y, + int z, + GridSpec grid, + Int3 &brick, + int &local_id) + { + // Split one voxel coordinate into a brick key plus a local bit id. + // The key addresses the hash table; local_id addresses the 512-bit + // occupancy mask inside that brick. + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const int rx = x - grid_min.x; + const int ry = y - grid_min.y; + const int rz = z - grid_min.z; + brick.x = rx / kBrickSize; + brick.y = ry / kBrickSize; + brick.z = rz / kBrickSize; + const int lx = rx - brick.x * kBrickSize; + const int ly = ry - brick.y * kBrickSize; + const int lz = rz - brick.z * kBrickSize; + local_id = lx + kBrickSize * (ly + kBrickSize * lz); + const uint64_t nbx = div_up_i64(grid_max.x - grid_min.x, kBrickSize); + const uint64_t nby = div_up_i64(grid_max.y - grid_min.y, kBrickSize); + return static_cast(brick.x) + nbx * (static_cast(brick.y) + nby * static_cast(brick.z)); + } + + __device__ uint32_t get_or_create_brick( + uint64_t key, + Int3 brick, + uint64_t *hash_keys, + uint32_t *hash_vals, + uint32_t *brick_count, + int32_t *brick_coords, + int32_t *overflow_flag, + uint64_t hash_capacity, + uint32_t max_bricks) + { + uint64_t slot = mix64(key) & (hash_capacity - 1); + for (uint64_t probe = 0; probe < hash_capacity; ++probe) + { + // The first thread that installs the key owns brick allocation. + // Other threads finding the same key wait until hash_vals is + // published, then reuse the existing compact brick index. + const uint64_t prev = atomicCAS( + reinterpret_cast(hash_keys + slot), + static_cast(kEmptyBrickKey), + static_cast(key)); + if (prev == kEmptyBrickKey) + { + const uint32_t idx = atomicAdd(brick_count, 1u); + if (idx >= max_bricks) + { + atomicExch(overflow_flag, 1); + __threadfence(); + hash_vals[slot] = kOverflowBrickVal; + return kEmptyBrickVal; + } + brick_coords[3 * idx + 0] = brick.x; + brick_coords[3 * idx + 1] = brick.y; + brick_coords[3 * idx + 2] = brick.z; + __threadfence(); + hash_vals[slot] = idx; + return idx; + } + if (prev == key) + { + volatile uint32_t *val_ptr = hash_vals + slot; + uint32_t val = *val_ptr; + // key becomes visible before value; spin only on that slot + // until the creating thread publishes the brick index. + while (val == kEmptyBrickVal) + val = *val_ptr; + return val == kOverflowBrickVal ? kEmptyBrickVal : val; + } + slot = (slot + 1u) & (hash_capacity - 1); + } + atomicExch(overflow_flag, 1); + return kEmptyBrickVal; + } + + __device__ __forceinline__ SymQEF10 triangle_qef(const float *tri) + { + // Plane QEF for the triangle itself. The normal uses the same edge + // order as the CPU path so the plane sign and d term stay aligned. + const float e0x = tri[3] - tri[0]; + const float e0y = tri[4] - tri[1]; + const float e0z = tri[5] - tri[2]; + const float e1x = tri[6] - tri[3]; + const float e1y = tri[7] - tri[4]; + const float e1z = tri[8] - tri[5]; + float nx = e0y * e1z - e0z * e1y; + float ny = e0z * e1x - e0x * e1z; + float nz = e0x * e1y - e0y * e1x; + const float inv_len = rsqrtf(nx * nx + ny * ny + nz * nz + 1e-30f); + nx *= inv_len; + ny *= inv_len; + nz *= inv_len; + return qef_from_plane(float4{nx, ny, nz, -(nx * tri[0] + ny * tri[1] + nz * tri[2])}); + } + + __device__ __forceinline__ void atomic_add_qef(float *dst, const SymQEF10 &q) + { + // Multiple scan tasks can hit the same compact voxel row, so every + // matrix coefficient is accumulated atomically. + atomicAdd(dst + 0, q.q00); + atomicAdd(dst + 1, q.q01); + atomicAdd(dst + 2, q.q02); + atomicAdd(dst + 3, q.q03); + atomicAdd(dst + 4, q.q11); + atomicAdd(dst + 5, q.q12); + atomicAdd(dst + 6, q.q13); + atomicAdd(dst + 7, q.q22); + atomicAdd(dst + 8, q.q23); + atomicAdd(dst + 9, q.q33); + } + + __global__ void count_scan_tasks_kernel( + const float *triangles, + int64_t num_triangles, + GridSpec grid, + int64_t *task_counts) + { + const int64_t pair_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (pair_id >= num_triangles * 3) + return; + const int64_t tri_id = pair_id / 3; + const int axis = static_cast(pair_id - tri_id * 3); + // Parallel unit is (triangle, depth axis). Large projected boxes + // become many fixed-size tile tasks instead of one long thread. + int u0, u1, v0, v1; + if (!compute_scan_bbox(triangles + tri_id * 9, axis, grid, u0, u1, v0, v1)) + { + task_counts[pair_id] = 0; + return; + } + task_counts[pair_id] = static_cast(div_up_i32(u1 - u0, kTileU)) * + static_cast(div_up_i32(v1 - v0, kTileV)); + } + + __global__ void emit_scan_tasks_and_brick_bounds_kernel( + const float *triangles, + int64_t num_triangles, + GridSpec grid, + const int64_t *task_offsets, + ScanTask *tasks, + int64_t *task_brick_bounds) + { + const int64_t pair_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (pair_id >= num_triangles * 3) + return; + const int64_t tri_id = pair_id / 3; + const int axis = static_cast(pair_id - tri_id * 3); + int u0, u1, v0, v1; + if (!compute_scan_bbox(triangles + tri_id * 9, axis, grid, u0, u1, v0, v1)) + return; + + int64_t out = task_offsets[pair_id]; + // task_offsets is the CUB exclusive scan of per-pair task counts. + // Each thread owns a disjoint output range and writes all tiles for + // its (triangle, axis) pair. + for (int v = v0; v < v1; v += kTileV) + { + for (int u = u0; u < u1; u += kTileU) + { + ScanTask task{ + static_cast(tri_id), + static_cast(axis), + static_cast(u), + static_cast(min(u + kTileU, u1)), + static_cast(v), + static_cast(min(v + kTileV, v1)), + }; + tasks[out] = task; + task_brick_bounds[out] = task_brick_bound(triangles + tri_id * 9, task, grid); + ++out; + } + } + } + + __global__ void mark_occupied_voxel_bits_kernel( + const ScanTask *tasks, + int64_t num_tasks, + const float *triangles, + GridSpec grid, + uint64_t *hash_keys, + uint32_t *hash_vals, + uint32_t *brick_count, + int32_t *brick_coords, + uint32_t *brick_bits, + int32_t *overflow_flag, + uint64_t hash_capacity, + uint32_t max_bricks) + { + const int64_t task_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (task_id >= num_tasks) + return; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const ScanTask task = tasks[task_id]; + const float *tri = triangles + static_cast(task.tri_id) * 9; + auto emit = [&](int ax0, int ax1, int ax2, int x_idx, int y_idx, int z_idx, double, double, double) + { + // One geometric event activates the four voxels around the + // crossed face in the projection plane. This is the occupancy + // subset later reused by face and boundary QEF stages. + for (int dx = 0; dx < 2; ++dx) + { + for (int dy = 0; dy < 2; ++dy) + { + int coord[3]; + coord[ax0] = x_idx + dx; + coord[ax1] = y_idx + dy; + coord[ax2] = z_idx; + if (coord[0] < grid_min.x || coord[0] >= grid_max.x) + continue; + if (coord[1] < grid_min.y || coord[1] >= grid_max.y) + continue; + if (coord[2] < grid_min.z || coord[2] >= grid_max.z) + continue; + Int3 brick; + int local_id; + const uint64_t key = voxel_to_brick(coord[0], coord[1], coord[2], grid, brick, local_id); + const uint32_t brick_idx = get_or_create_brick( + key, brick, hash_keys, hash_vals, brick_count, brick_coords, overflow_flag, hash_capacity, max_bricks); + if (brick_idx == kEmptyBrickVal) + continue; + atomicOr(brick_bits + static_cast(brick_idx) * kBrickBitWords + local_id / 32, 1u << (local_id & 31)); + } + } + }; + scan_triangle_events_tiled(tri, task, grid, emit); + } + + __global__ void count_brick_voxels_kernel( + const uint32_t *brick_bits, + int64_t num_bricks, + int64_t *brick_counts) + { + const int64_t brick_idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (brick_idx >= num_bricks) + return; + // Popcount turns each brick bitset into the number of compact voxel + // rows owned by that brick. + int64_t count = 0; + const uint32_t *bits = brick_bits + brick_idx * kBrickBitWords; + for (int i = 0; i < kBrickBitWords; ++i) + count += __popc(bits[i]); + brick_counts[brick_idx] = count; + } + + __global__ void emit_occupied_voxels_kernel( + const int32_t *brick_coords, + const uint32_t *brick_bits, + const int64_t *brick_base, + int64_t num_bricks, + GridSpec grid, + int32_t *voxels) + { + const int64_t brick_idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (brick_idx >= num_bricks) + return; + const Int3 grid_min = grid.grid_min; + const int bx = brick_coords[3 * brick_idx + 0]; + const int by = brick_coords[3 * brick_idx + 1]; + const int bz = brick_coords[3 * brick_idx + 2]; + const uint32_t *bits = brick_bits + brick_idx * kBrickBitWords; + int64_t out = brick_base[brick_idx]; + for (int local_id = 0; local_id < kBrickLocalCells; ++local_id) + { + if ((bits[local_id / 32] & (1u << (local_id & 31))) == 0) + continue; + // Enumerating local_id in increasing order defines the local + // row order inside this compact active brick. + const int lz = local_id / (kBrickSize * kBrickSize); + const int rem = local_id - lz * kBrickSize * kBrickSize; + const int ly = rem / kBrickSize; + const int lx = rem - ly * kBrickSize; + voxels[3 * out + 0] = grid_min.x + bx * kBrickSize + lx; + voxels[3 * out + 1] = grid_min.y + by * kBrickSize + ly; + voxels[3 * out + 2] = grid_min.z + bz * kBrickSize + lz; + ++out; + } + } + + __global__ void accumulate_intersection_qef_kernel( + const ScanTask *tasks, + int64_t num_tasks, + const float *triangles, + GridSpec grid, + const uint64_t *hash_keys, + const uint32_t *hash_vals, + const uint32_t *brick_bits, + const int64_t *brick_base, + float *mean_sum, + float *cnt, + uint32_t *intersected_mask, + float *qefs, + uint64_t hash_capacity) + { + const int64_t task_id = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (task_id >= num_tasks) + return; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const BrickLookup lookup{hash_keys, hash_vals, brick_bits, brick_base, hash_capacity}; + const ScanTask task = tasks[task_id]; + const float *tri = triangles + static_cast(task.tri_id) * 9; + const SymQEF10 qef = triangle_qef(tri); + auto emit = [&](int ax0, int ax1, int ax2, int x_idx, int y_idx, int z_idx, double x, double y, double z) + { + // Replay the same event stream used for occupancy. This pass now + // maps each active voxel to its compact row and accumulates the + // intersection point and triangle plane QEF. + for (int dx = 0; dx < 2; ++dx) + { + for (int dy = 0; dy < 2; ++dy) + { + int coord[3]; + coord[ax0] = x_idx + dx; + coord[ax1] = y_idx + dy; + coord[ax2] = z_idx; + if (coord[0] < grid_min.x || coord[0] >= grid_max.x) + continue; + if (coord[1] < grid_min.y || coord[1] >= grid_max.y) + continue; + if (coord[2] < grid_min.z || coord[2] >= grid_max.z) + continue; + const int64_t out_idx = lookup_voxel_row_in_bricks( + coord[0], + coord[1], + coord[2], + grid, + lookup); + if (out_idx < 0) + continue; + float p[3]; + p[ax0] = static_cast(x); + p[ax1] = static_cast(y); + p[ax2] = static_cast(z); + atomicAdd(mean_sum + 3 * out_idx + 0, p[0]); + atomicAdd(mean_sum + 3 * out_idx + 1, p[1]); + atomicAdd(mean_sum + 3 * out_idx + 2, p[2]); + atomicAdd(cnt + out_idx, 1.0f); + // The base event marks which grid edge direction was + // crossed. Neighbor voxels receive QEF/mean updates but + // should not duplicate the axis flag. + if (dx == 0 && dy == 0) + atomicOr(intersected_mask + out_idx, 1u << ax2); + atomic_add_qef(qefs + 10 * out_idx, qef); + } + } + }; + scan_triangle_events_tiled(tri, task, grid, emit); + } + + __global__ void decode_intersection_masks_kernel( + const uint32_t *mask, + int64_t n, + bool *intersected) + { + // Convert one uint32 bitfield per voxel into the public [N, 3] bool + // layout: bit 0/1/2 means the voxel was intersected along x/y/z. + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) + return; + const uint32_t m = mask[i]; + intersected[3 * i + 0] = (m & 1u) != 0; + intersected[3 * i + 1] = (m & 2u) != 0; + intersected[3 * i + 2] = (m & 4u) != 0; + } + + int64_t next_power_of_two_i64(int64_t x) + { + int64_t p = 1; + while (p < x) + p <<= 1; + return p; + } + + int64_t read_i64(const torch::Tensor &t, cudaStream_t stream) + { + int64_t value = 0; + C10_CUDA_CHECK(cudaMemcpyAsync(&value, t.data_ptr(), sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + return value; + } + + int64_t read_scan_total(const torch::Tensor &counts, const torch::Tensor &offsets, int64_t n, cudaStream_t stream) + { + // After an exclusive scan, total = last_count + last_offset. This + // small host readback gives the exact tensor size for the next pass. + int64_t tail[2] = {0, 0}; + C10_CUDA_CHECK(cudaMemcpyAsync( + tail, + counts.data_ptr() + n - 1, + sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaMemcpyAsync( + tail + 1, + offsets.data_ptr() + n - 1, + sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + return tail[0] + tail[1]; + } + + void cub_exclusive_sum_i64(const torch::Tensor &in, const torch::Tensor &out, int64_t n, cudaStream_t stream) + { + if (n == 0) + return; + TORCH_CHECK(n <= std::numeric_limits::max(), "CUB item count exceeds int"); + size_t temp_bytes = 0; + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + nullptr, + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + auto temp = torch::empty( + {static_cast(temp_bytes)}, + torch::TensorOptions().dtype(torch::kUInt8).device(in.device())); + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + temp.data_ptr(), + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + } + + int64_t cub_sum_i64(const torch::Tensor &in, int64_t n, cudaStream_t stream) + { + if (n == 0) + return 0; + TORCH_CHECK(n <= std::numeric_limits::max(), "CUB item count exceeds int"); + auto out = torch::empty({1}, in.options()); + size_t temp_bytes = 0; + C10_CUDA_CHECK(cub::DeviceReduce::Sum( + nullptr, + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + auto temp = torch::empty( + {static_cast(temp_bytes)}, + torch::TensorOptions().dtype(torch::kUInt8).device(in.device())); + C10_CUDA_CHECK(cub::DeviceReduce::Sum( + temp.data_ptr(), + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + return read_i64(out, stream); + } + + IntersectionOccupancy build_intersection_occupancy( + const torch::Tensor &triangles, + GridSpec grid, + const torch::Device &device, + cudaStream_t stream) + { + // Pipeline: + // 1. Count and emit scan tasks for triangle-axis tiles. + // 2. Estimate a strict active-brick bound from those tasks. + // 3. Mark occupied voxel bits in active bricks through a hash table. + // 4. Prefix-sum brick popcounts and emit compact voxel coordinates. + IntersectionOccupancy out; + const Int3 grid_min = grid.grid_min; + const Int3 grid_max = grid.grid_max; + const int64_t num_triangles = triangles.size(0); + const auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto opts_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto opts_u32 = torch::TensorOptions().dtype(torch::kUInt32).device(device); + const auto opts_u64 = torch::TensorOptions().dtype(torch::kUInt64).device(device); + out.tasks = torch::empty({0, 6}, opts_i32); + out.hash_keys = torch::empty({0}, opts_u64); + out.hash_vals = torch::empty({0}, opts_u32); + out.brick_coords = torch::empty({0, 3}, opts_i32); + out.brick_bits = torch::empty({0, kBrickBitWords}, opts_u32); + out.brick_base = torch::empty({0}, opts_i64); + out.overflow_flag = torch::empty({0}, opts_i32); + out.voxels = torch::empty({0, 3}, opts_i32); + if (num_triangles == 0) + return out; + + const int64_t pair_count = num_triangles * 3; + auto task_counts = torch::empty({pair_count}, opts_i64); + auto task_offsets = torch::empty({pair_count}, opts_i64); + int blocks = static_cast((pair_count + kThreads - 1) / kThreads); + count_scan_tasks_kernel<<>>( + triangles.data_ptr(), + num_triangles, + grid, + task_counts.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + cub_exclusive_sum_i64(task_counts, task_offsets, pair_count, stream); + const int64_t num_tasks = read_scan_total(task_counts, task_offsets, pair_count, stream); + out.num_tasks = num_tasks; + if (num_tasks == 0) + return out; + + out.tasks = torch::empty({num_tasks, 6}, opts_i32); + auto task_brick_bounds = torch::empty({num_tasks}, opts_i64); + emit_scan_tasks_and_brick_bounds_kernel<<>>( + triangles.data_ptr(), + num_triangles, + grid, + task_offsets.data_ptr(), + reinterpret_cast(out.tasks.data_ptr()), + task_brick_bounds.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const int64_t sum_brick_bounds = cub_sum_i64(task_brick_bounds, num_tasks, stream); + const int64_t nbx = div_up_i64(grid_max.x - grid_min.x, kBrickSize); + const int64_t nby = div_up_i64(grid_max.y - grid_min.y, kBrickSize); + const int64_t nbz = div_up_i64(grid_max.z - grid_min.z, kBrickSize); + const int64_t total_grid_bricks = nbx * nby * nbz; + const int64_t max_bricks = std::min(sum_brick_bounds, total_grid_bricks); + if (max_bricks == 0) + return out; + TORCH_CHECK(max_bricks <= std::numeric_limits::max(), "active brick bound exceeds uint32_t"); + const int64_t hash_capacity_i64 = next_power_of_two_i64(std::max(2, max_bricks * 2)); + out.hash_capacity = static_cast(hash_capacity_i64); + + out.hash_keys = torch::empty({hash_capacity_i64}, opts_u64); + out.hash_vals = torch::empty({hash_capacity_i64}, opts_u32); + auto brick_count = torch::zeros({1}, opts_u32); + out.brick_coords = torch::empty({max_bricks, 3}, opts_i32); + out.brick_bits = torch::zeros({max_bricks, kBrickBitWords}, opts_u32); + out.overflow_flag = torch::zeros({1}, opts_i32); + C10_CUDA_CHECK(cudaMemsetAsync(out.hash_keys.data_ptr(), 0xff, hash_capacity_i64 * sizeof(uint64_t), stream)); + C10_CUDA_CHECK(cudaMemsetAsync(out.hash_vals.data_ptr(), 0xff, hash_capacity_i64 * sizeof(uint32_t), stream)); + + blocks = static_cast((num_tasks + kThreads - 1) / kThreads); + mark_occupied_voxel_bits_kernel<<>>( + reinterpret_cast(out.tasks.data_ptr()), + num_tasks, + triangles.data_ptr(), + grid, + out.hash_keys.data_ptr(), + out.hash_vals.data_ptr(), + brick_count.data_ptr(), + out.brick_coords.data_ptr(), + out.brick_bits.data_ptr(), + out.overflow_flag.data_ptr(), + out.hash_capacity, + static_cast(max_bricks)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + int32_t overflow = 0; + uint32_t brick_count_h = 0; + C10_CUDA_CHECK(cudaMemcpyAsync(&overflow, out.overflow_flag.data_ptr(), sizeof(int32_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaMemcpyAsync(&brick_count_h, brick_count.data_ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + TORCH_CHECK(overflow == 0, "brick occupancy hash overflow"); + + const int64_t num_bricks = static_cast(brick_count_h); + out.num_bricks = num_bricks; + TORCH_CHECK(num_bricks <= max_bricks, "active brick count exceeds bound"); + if (num_bricks == 0) + return out; + + auto brick_counts = torch::empty({max_bricks}, opts_i64); + out.brick_base = torch::empty({max_bricks}, opts_i64); + blocks = static_cast((num_bricks + kThreads - 1) / kThreads); + count_brick_voxels_kernel<<>>( + out.brick_bits.data_ptr(), + num_bricks, + brick_counts.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + cub_exclusive_sum_i64(brick_counts, out.brick_base, num_bricks, stream); + const int64_t num_voxels = read_scan_total(brick_counts, out.brick_base, num_bricks, stream); + out.num_voxels = num_voxels; + if (num_voxels == 0) + return out; + + out.voxels = torch::empty({num_voxels, 3}, opts_i32); + emit_occupied_voxels_kernel<<>>( + out.brick_coords.data_ptr(), + out.brick_bits.data_ptr(), + out.brick_base.data_ptr(), + num_bricks, + grid, + out.voxels.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; + } + + } // namespace + + torch::Tensor intersect_occ_cuda( + const torch::Tensor &triangles, + const std::vector &voxel_size, + const std::vector &grid_range) + { + // Same occupancy construction as intersect_qef_cuda, but no QEF, + // mean/cnt, or intersected mask work is performed. + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + + const c10::cuda::CUDAGuard guard(triangles.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(triangles.get_device()).stream(); + const torch::Device device = triangles.device(); + const GridSpec grid{ + float3{voxel_size[0], voxel_size[1], voxel_size[2]}, + Int3{ + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])}, + Int3{ + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])}, + }; + return build_intersection_occupancy(triangles, grid, device, stream).voxels; + } + + 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 &voxel_size, + const std::vector &grid_range) + { + // Start from the shared occupancy pass, then accumulate the intersection + // planes and per-axis flags needed by the full flexible dual grid solve. + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + + const c10::cuda::CUDAGuard guard(triangles.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(triangles.get_device()).stream(); + const torch::Device device = triangles.device(); + const auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(device); + const auto opts_bool = torch::TensorOptions().dtype(torch::kBool).device(device); + const auto opts_u32 = torch::TensorOptions().dtype(torch::kUInt32).device(device); + const GridSpec grid{ + float3{voxel_size[0], voxel_size[1], voxel_size[2]}, + Int3{ + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])}, + Int3{ + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])}, + }; + + IntersectionOccupancy occ = build_intersection_occupancy(triangles, grid, device, stream); + const int64_t n = occ.num_voxels; + auto mean_sum = torch::zeros({n, 3}, opts_f32); + auto cnt = torch::zeros({n}, opts_f32); + auto intersected = torch::empty({n, 3}, opts_bool); + auto qefs = torch::zeros({n, 10}, opts_f32); + if (n == 0) + return std::make_tuple( + occ.voxels, + mean_sum, + cnt, + intersected, + qefs, + occ.hash_keys, + occ.hash_vals, + occ.brick_bits, + occ.brick_base); + + auto intersected_mask = torch::zeros({n}, opts_u32); + const int blocks_tasks = static_cast((occ.num_tasks + kThreads - 1) / kThreads); + accumulate_intersection_qef_kernel<<>>( + reinterpret_cast(occ.tasks.data_ptr()), + occ.num_tasks, + triangles.data_ptr(), + grid, + occ.hash_keys.data_ptr(), + occ.hash_vals.data_ptr(), + occ.brick_bits.data_ptr(), + occ.brick_base.data_ptr(), + mean_sum.data_ptr(), + cnt.data_ptr(), + intersected_mask.data_ptr(), + qefs.data_ptr(), + occ.hash_capacity); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + const int blocks_voxels = static_cast((n + kThreads - 1) / kThreads); + decode_intersection_masks_kernel<<>>( + intersected_mask.data_ptr(), + n, + intersected.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return std::make_tuple( + occ.voxels, + mean_sum, + cnt, + intersected, + qefs, + occ.hash_keys, + occ.hash_vals, + occ.brick_bits, + occ.brick_base); + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.cu new file mode 100644 index 00000000..692fd522 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.cu @@ -0,0 +1,823 @@ +#include "intersection_qef.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace intersection_qef { +namespace { + +using fdg_gpu::DeviceBuffer; +using fdg_gpu::SymQEF10; +using fdg_gpu::int3_; +using fdg_gpu::throw_cuda_error; + +#define IQ_CUDA_CHECK(expr) ::fdg_gpu::throw_cuda_error((expr), #expr) + +struct D2 { + double x; + double z; +}; + +struct D3 { + double x; + double y; + double z; + + __host__ __device__ double operator[](int i) const { return (&x)[i]; } +}; + +struct QEFEventValue { + float mean_sum_x; + float mean_sum_y; + float mean_sum_z; + float cnt; + uint8_t intersected; + SymQEF10 qef; +}; + +struct OccChunk { + DeviceBuffer keys; + int64_t size = 0; +}; + +struct QEFChunk { + DeviceBuffer keys; + DeviceBuffer values; + int64_t size = 0; +}; + +struct AddQEFEventValue { + __host__ __device__ QEFEventValue operator()(const QEFEventValue& a, const QEFEventValue& b) const { + QEFEventValue out; + out.mean_sum_x = a.mean_sum_x + b.mean_sum_x; + out.mean_sum_y = a.mean_sum_y + b.mean_sum_y; + out.mean_sum_z = a.mean_sum_z + b.mean_sum_z; + out.cnt = a.cnt + b.cnt; + out.intersected = static_cast(a.intersected | b.intersected); + out.qef.q00 = a.qef.q00 + b.qef.q00; + out.qef.q01 = a.qef.q01 + b.qef.q01; + out.qef.q02 = a.qef.q02 + b.qef.q02; + out.qef.q03 = a.qef.q03 + b.qef.q03; + out.qef.q11 = a.qef.q11 + b.qef.q11; + out.qef.q12 = a.qef.q12 + b.qef.q12; + out.qef.q13 = a.qef.q13 + b.qef.q13; + out.qef.q22 = a.qef.q22 + b.qef.q22; + out.qef.q23 = a.qef.q23 + b.qef.q23; + out.qef.q33 = a.qef.q33 + b.qef.q33; + return out; + } +}; + +__host__ __device__ inline double lerp_scalar(double a, double b, double t, double va, double vb) { + if (a == b) return va; + const double alpha = (t - a) / (b - a); + return (1.0 - alpha) * va + alpha * vb; +} + +__host__ __device__ inline D2 lerp_vec2(double a, double b, double t, D2 va, D2 vb) { + if (a == b) return va; + const double alpha = (t - a) / (b - a); + return D2{(1.0 - alpha) * va.x + alpha * vb.x, (1.0 - alpha) * va.z + alpha * vb.z}; +} + +__host__ __device__ inline int clamp_int(int x, int lo, int hi) { + return x < lo ? lo : (x > hi ? hi : x); +} + +__host__ __device__ inline void swap_d3(D3& a, D3& b) { + D3 t = a; + a = b; + b = t; +} + +__host__ __device__ inline void sort_by_y(D3& t0, D3& t1, D3& t2) { + if (t0.y > t1.y) swap_d3(t0, t1); + if (t1.y > t2.y) swap_d3(t1, t2); + if (t0.y > t1.y) swap_d3(t0, t1); +} + +__device__ inline void normalize3(double& x, double& y, double& z) { + const double n = sqrt(x * x + y * y + z * z); + if (n > 0.0) { + x /= n; + y /= n; + z /= n; + } +} + +__device__ inline SymQEF10 make_plane_qef_from_triangle(const double v0[3], const double v1[3], const double v2[3]) { + const double e0x = v1[0] - v0[0]; + const double e0y = v1[1] - v0[1]; + const double e0z = v1[2] - v0[2]; + + const double e1x = v2[0] - v1[0]; + const double e1y = v2[1] - v1[1]; + const double e1z = v2[2] - v1[2]; + + double nx = e0y * e1z - e0z * e1y; + double ny = e0z * e1x - e0x * e1z; + double nz = e0x * e1y - e0y * e1x; + normalize3(nx, ny, nz); + + const double d = -(nx * v0[0] + ny * v0[1] + nz * v0[2]); + + SymQEF10 q; + q.q00 = static_cast(nx * nx); + q.q01 = static_cast(nx * ny); + q.q02 = static_cast(nx * nz); + q.q03 = static_cast(nx * d); + q.q11 = static_cast(ny * ny); + q.q12 = static_cast(ny * nz); + q.q13 = static_cast(ny * d); + q.q22 = static_cast(nz * nz); + q.q23 = static_cast(nz * d); + q.q33 = static_cast(d * d); + return q; +} + +__device__ inline int64_t count_triangle_axis_surface_voxels( + const float* tri, + int ax2, + const float voxel_size[3], + int3_ grid_min, + int3_ grid_max) { + const double v0[3] = {static_cast(tri[0]), static_cast(tri[1]), static_cast(tri[2])}; + const double v1[3] = {static_cast(tri[3]), static_cast(tri[4]), static_cast(tri[5])}; + const double v2[3] = {static_cast(tri[6]), static_cast(tri[7]), static_cast(tri[8])}; + + const int ax0 = (ax2 + 1) % 3; + const int ax1 = (ax2 + 2) % 3; + + D3 t0{v0[ax0], v0[ax1], v0[ax2]}; + D3 t1{v1[ax0], v1[ax1], v1[ax2]}; + D3 t2{v2[ax0], v2[ax1], v2[ax2]}; + sort_by_y(t0, t1, t2); + + const int start = clamp_int(static_cast(t0.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int mid = clamp_int(static_cast(t1.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int end = clamp_int(static_cast(t2.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + + int64_t total = 0; + auto scan_half = [&](int row_start, int row_end, D3 a, D3 b, D3 c) { + for (int y_idx = row_start; y_idx < row_end; ++y_idx) { + const double y = (static_cast(y_idx) + 1.0) * voxel_size[ax1]; + D2 t3 = lerp_vec2(a.y, b.y, y, D2{a.x, a.z}, D2{b.x, b.z}); + D2 t4 = lerp_vec2(a.y, c.y, y, D2{a.x, a.z}, D2{c.x, c.z}); + if (t3.x > t4.x) { + D2 tmp = t3; + t3 = t4; + t4 = tmp; + } + + const int line_start = clamp_int(static_cast(t3.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + const int line_end = clamp_int(static_cast(t4.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + for (int x_idx = line_start; x_idx < line_end; ++x_idx) { + const double x = (static_cast(x_idx) + 1.0) * voxel_size[ax0]; + const double z = lerp_scalar(t3.x, t4.x, x, t3.z, t4.z); + const int z_idx = static_cast(z / voxel_size[ax2]); + if (z_idx < grid_min[ax2] || z_idx >= grid_max[ax2]) continue; + total += 4; + } + } + }; + + scan_half(start, mid, t0, t1, t2); + scan_half(mid, end, t2, t1, t0); + return total; +} + +__global__ void intersection_count_kernel( + const float* triangles, + int64_t tri_begin, + int64_t tri_count, + float vx, + float vy, + float vz, + int3_ grid_min, + int3_ grid_max, + int64_t* counts) { + const int64_t local_t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local_t >= tri_count) return; + + const int64_t t = tri_begin + local_t; + const float* tri = triangles + t * 9; + const float voxel_size[3] = {vx, vy, vz}; + + int64_t total = 0; + total += count_triangle_axis_surface_voxels(tri, 0, voxel_size, grid_min, grid_max); + total += count_triangle_axis_surface_voxels(tri, 1, voxel_size, grid_min, grid_max); + total += count_triangle_axis_surface_voxels(tri, 2, voxel_size, grid_min, grid_max); + counts[local_t] = total; +} + +__global__ void intersection_occ_emit_kernel( + const float* triangles, + int64_t tri_begin, + int64_t tri_count, + float vx, + float vy, + float vz, + int3_ grid_min, + int3_ grid_max, + const int64_t* offsets, + uint64_t* event_keys) { + const int64_t local_t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local_t >= tri_count) return; + + const int64_t t = tri_begin + local_t; + const float* tri = triangles + t * 9; + const double v0[3] = {static_cast(tri[0]), static_cast(tri[1]), static_cast(tri[2])}; + const double v1[3] = {static_cast(tri[3]), static_cast(tri[4]), static_cast(tri[5])}; + const double v2[3] = {static_cast(tri[6]), static_cast(tri[7]), static_cast(tri[8])}; + const float voxel_size[3] = {vx, vy, vz}; + + int64_t out = offsets[local_t]; + + for (int ax2 = 0; ax2 < 3; ++ax2) { + const int ax0 = (ax2 + 1) % 3; + const int ax1 = (ax2 + 2) % 3; + + D3 t0{v0[ax0], v0[ax1], v0[ax2]}; + D3 t1{v1[ax0], v1[ax1], v1[ax2]}; + D3 t2{v2[ax0], v2[ax1], v2[ax2]}; + sort_by_y(t0, t1, t2); + + const int start = clamp_int(static_cast(t0.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int mid = clamp_int(static_cast(t1.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int end = clamp_int(static_cast(t2.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + + auto emit_one = [&](int x_idx, int y_idx, int z_idx) { + int coord[3]; + coord[ax0] = x_idx; + coord[ax1] = y_idx; + coord[ax2] = z_idx; + event_keys[out++] = fdg_gpu::pack_voxel_key(coord[0], coord[1], coord[2], grid_min, grid_max); + }; + + auto scan_half = [&](int row_start, int row_end, D3 a, D3 b, D3 c) { + for (int y_idx = row_start; y_idx < row_end; ++y_idx) { + const double y = (static_cast(y_idx) + 1.0) * voxel_size[ax1]; + D2 t3 = lerp_vec2(a.y, b.y, y, D2{a.x, a.z}, D2{b.x, b.z}); + D2 t4 = lerp_vec2(a.y, c.y, y, D2{a.x, a.z}, D2{c.x, c.z}); + if (t3.x > t4.x) { + D2 tmp = t3; + t3 = t4; + t4 = tmp; + } + + const int line_start = clamp_int(static_cast(t3.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + const int line_end = clamp_int(static_cast(t4.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + + for (int x_idx = line_start; x_idx < line_end; ++x_idx) { + const double x = (static_cast(x_idx) + 1.0) * voxel_size[ax0]; + const double z = lerp_scalar(t3.x, t4.x, x, t3.z, t4.z); + const int z_idx = static_cast(z / voxel_size[ax2]); + if (z_idx < grid_min[ax2] || z_idx >= grid_max[ax2]) continue; + + emit_one(x_idx + 0, y_idx + 0, z_idx); + emit_one(x_idx + 1, y_idx + 0, z_idx); + emit_one(x_idx + 0, y_idx + 1, z_idx); + emit_one(x_idx + 1, y_idx + 1, z_idx); + } + } + }; + + scan_half(start, mid, t0, t1, t2); + scan_half(mid, end, t2, t1, t0); + } +} + +__global__ void intersect_qef_emit_kernel( + const float* triangles, + int64_t tri_begin, + int64_t tri_count, + float vx, + float vy, + float vz, + int3_ grid_min, + int3_ grid_max, + const int64_t* offsets, + uint64_t* event_keys, + QEFEventValue* event_values) { + const int64_t local_t = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local_t >= tri_count) return; + + const int64_t t = tri_begin + local_t; + const float* tri = triangles + t * 9; + const double v0[3] = {static_cast(tri[0]), static_cast(tri[1]), static_cast(tri[2])}; + const double v1[3] = {static_cast(tri[3]), static_cast(tri[4]), static_cast(tri[5])}; + const double v2[3] = {static_cast(tri[6]), static_cast(tri[7]), static_cast(tri[8])}; + const float voxel_size[3] = {vx, vy, vz}; + const SymQEF10 qef = make_plane_qef_from_triangle(v0, v1, v2); + + int64_t out = offsets[local_t]; + + for (int ax2 = 0; ax2 < 3; ++ax2) { + const int ax0 = (ax2 + 1) % 3; + const int ax1 = (ax2 + 2) % 3; + + D3 t0{v0[ax0], v0[ax1], v0[ax2]}; + D3 t1{v1[ax0], v1[ax1], v1[ax2]}; + D3 t2{v2[ax0], v2[ax1], v2[ax2]}; + sort_by_y(t0, t1, t2); + + const int start = clamp_int(static_cast(t0.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int mid = clamp_int(static_cast(t1.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + const int end = clamp_int(static_cast(t2.y / voxel_size[ax1]), grid_min[ax1], grid_max[ax1] - 1); + + auto emit_one = [&](int x_idx, int y_idx, int z_idx, double x, double y, double z, uint8_t mask) { + int coord[3]; + coord[ax0] = x_idx; + coord[ax1] = y_idx; + coord[ax2] = z_idx; + event_keys[out] = fdg_gpu::pack_voxel_key(coord[0], coord[1], coord[2], grid_min, grid_max); + event_values[out].mean_sum_x = static_cast(ax0 == 0 ? x : (ax1 == 0 ? y : z)); + event_values[out].mean_sum_y = static_cast(ax0 == 1 ? x : (ax1 == 1 ? y : z)); + event_values[out].mean_sum_z = static_cast(ax0 == 2 ? x : (ax1 == 2 ? y : z)); + event_values[out].cnt = 1.0f; + event_values[out].intersected = mask; + event_values[out].qef = qef; + ++out; + }; + + auto scan_half = [&](int row_start, int row_end, D3 a, D3 b, D3 c) { + for (int y_idx = row_start; y_idx < row_end; ++y_idx) { + const double y = (static_cast(y_idx) + 1.0) * voxel_size[ax1]; + D2 t3 = lerp_vec2(a.y, b.y, y, D2{a.x, a.z}, D2{b.x, b.z}); + D2 t4 = lerp_vec2(a.y, c.y, y, D2{a.x, a.z}, D2{c.x, c.z}); + if (t3.x > t4.x) { + D2 tmp = t3; + t3 = t4; + t4 = tmp; + } + + const int line_start = clamp_int(static_cast(t3.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + const int line_end = clamp_int(static_cast(t4.x / voxel_size[ax0]), grid_min[ax0], grid_max[ax0] - 1); + + for (int x_idx = line_start; x_idx < line_end; ++x_idx) { + const double x = (static_cast(x_idx) + 1.0) * voxel_size[ax0]; + const double z = lerp_scalar(t3.x, t4.x, x, t3.z, t4.z); + const int z_idx = static_cast(z / voxel_size[ax2]); + if (z_idx < grid_min[ax2] || z_idx >= grid_max[ax2]) continue; + + emit_one(x_idx + 0, y_idx + 0, z_idx, x, y, z, static_cast(1u << ax2)); + emit_one(x_idx + 1, y_idx + 0, z_idx, x, y, z, static_cast(0u)); + emit_one(x_idx + 0, y_idx + 1, z_idx, x, y, z, static_cast(0u)); + emit_one(x_idx + 1, y_idx + 1, z_idx, x, y, z, static_cast(0u)); + } + } + }; + + scan_half(start, mid, t0, t1, t2); + scan_half(mid, end, t2, t1, t0); + } +} + +__global__ void decode_occ_output_kernel( + const uint64_t* keys, + int64_t size, + int3_ grid_min, + int3_ grid_max, + int* out_voxels) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= size) return; + int x, y, z; + const uint64_t key = keys[i]; + const uint64_t sx = static_cast(grid_max.x - grid_min.x); + const uint64_t sy = static_cast(grid_max.y - grid_min.y); + const uint64_t yz = sx * sy; + const uint64_t zz = key / yz; + const uint64_t rem = key - zz * yz; + const uint64_t yy = rem / sx; + const uint64_t xx = rem - yy * sx; + x = static_cast(xx) + grid_min.x; + y = static_cast(yy) + grid_min.y; + z = static_cast(zz) + grid_min.z; + out_voxels[3 * i + 0] = x; + out_voxels[3 * i + 1] = y; + out_voxels[3 * i + 2] = z; +} + +__global__ void decode_qef_output_kernel( + const uint64_t* keys, + const QEFEventValue* values, + int64_t size, + int3_ grid_min, + int3_ grid_max, + int* out_voxels, + float* out_mean_sum, + float* out_cnt, + uint8_t* out_intersected, + SymQEF10* out_qefs) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= size) return; + + const uint64_t key = keys[i]; + const uint64_t sx = static_cast(grid_max.x - grid_min.x); + const uint64_t sy = static_cast(grid_max.y - grid_min.y); + const uint64_t yz = sx * sy; + const uint64_t zz = key / yz; + const uint64_t rem = key - zz * yz; + const uint64_t yy = rem / sx; + const uint64_t xx = rem - yy * sx; + + out_voxels[3 * i + 0] = static_cast(xx) + grid_min.x; + out_voxels[3 * i + 1] = static_cast(yy) + grid_min.y; + out_voxels[3 * i + 2] = static_cast(zz) + grid_min.z; + + out_mean_sum[3 * i + 0] = values[i].mean_sum_x; + out_mean_sum[3 * i + 1] = values[i].mean_sum_y; + out_mean_sum[3 * i + 2] = values[i].mean_sum_z; + out_cnt[i] = values[i].cnt; + out_intersected[i] = values[i].intersected; + out_qefs[i] = values[i].qef; +} + +inline int64_t copy_last_i64(const int64_t* ptr, int64_t count, cudaStream_t stream) { + if (count <= 0) return 0; + int64_t value = 0; + IQ_CUDA_CHECK(cudaMemcpyAsync(&value, ptr + (count - 1), sizeof(int64_t), cudaMemcpyDeviceToHost, stream)); + IQ_CUDA_CHECK(cudaStreamSynchronize(stream)); + return value; +} + +OccChunk make_occ_chunk_exact(DeviceBuffer& src_keys, int64_t size, cudaStream_t stream) { + OccChunk out; + out.size = size; + if (size <= 0) return out; + out.keys.allocate(size); + thrust::copy_n( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(src_keys.data()), + size, + thrust::device_pointer_cast(out.keys.data())); + return out; +} + +QEFChunk make_qef_chunk_exact( + DeviceBuffer& src_keys, + DeviceBuffer& src_values, + int64_t size, + cudaStream_t stream) { + QEFChunk out; + out.size = size; + if (size <= 0) return out; + out.keys.allocate(size); + out.values.allocate(size); + thrust::copy_n( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(src_keys.data()), + size, + thrust::device_pointer_cast(out.keys.data())); + thrust::copy_n( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(src_values.data()), + size, + thrust::device_pointer_cast(out.values.data())); + return out; +} + +OccChunk merge_occ_two_chunks(OccChunk a, OccChunk b, cudaStream_t stream) { + if (a.size == 0) return std::move(b); + if (b.size == 0) return std::move(a); + + DeviceBuffer merged(a.size + b.size); + auto merged_end = thrust::merge( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(a.keys.data()), + thrust::device_pointer_cast(a.keys.data()) + a.size, + thrust::device_pointer_cast(b.keys.data()), + thrust::device_pointer_cast(b.keys.data()) + b.size, + thrust::device_pointer_cast(merged.data())); + const int64_t merged_size = merged_end - thrust::device_pointer_cast(merged.data()); + + auto unique_end = thrust::unique( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(merged.data()), + thrust::device_pointer_cast(merged.data()) + merged_size); + + OccChunk out; + out.size = unique_end - thrust::device_pointer_cast(merged.data()); + out.keys = std::move(merged); + return out; +} + +QEFChunk merge_qef_two_chunks(QEFChunk a, QEFChunk b, cudaStream_t stream) { + if (a.size == 0) return std::move(b); + if (b.size == 0) return std::move(a); + + DeviceBuffer merged_keys(a.size + b.size); + DeviceBuffer merged_values(a.size + b.size); + + auto merged_end = thrust::merge_by_key( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(a.keys.data()), + thrust::device_pointer_cast(a.keys.data()) + a.size, + thrust::device_pointer_cast(b.keys.data()), + thrust::device_pointer_cast(b.keys.data()) + b.size, + thrust::device_pointer_cast(a.values.data()), + thrust::device_pointer_cast(b.values.data()), + thrust::device_pointer_cast(merged_keys.data()), + thrust::device_pointer_cast(merged_values.data())); + const int64_t merged_size = merged_end.first - thrust::device_pointer_cast(merged_keys.data()); + + DeviceBuffer next_keys(merged_size); + DeviceBuffer next_values(merged_size); + auto next_end = thrust::reduce_by_key( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(merged_keys.data()), + thrust::device_pointer_cast(merged_keys.data()) + merged_size, + thrust::device_pointer_cast(merged_values.data()), + thrust::device_pointer_cast(next_keys.data()), + thrust::device_pointer_cast(next_values.data()), + thrust::equal_to(), + AddQEFEventValue()); + + QEFChunk out; + out.size = next_end.first - thrust::device_pointer_cast(next_keys.data()); + out.keys = std::move(next_keys); + out.values = std::move(next_values); + return out; +} + +OccChunk final_merge_occ_chunks(std::vector chunks, cudaStream_t stream) { + if (chunks.empty()) return OccChunk{}; + while (chunks.size() > 1) { + std::vector next_level; + next_level.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + if (i + 1 >= chunks.size()) { + next_level.push_back(std::move(chunks[i])); + } else { + next_level.push_back(merge_occ_two_chunks(std::move(chunks[i]), std::move(chunks[i + 1]), stream)); + } + } + chunks = std::move(next_level); + } + return std::move(chunks[0]); +} + +QEFChunk final_merge_qef_chunks(std::vector chunks, cudaStream_t stream) { + if (chunks.empty()) return QEFChunk{}; + while (chunks.size() > 1) { + std::vector next_level; + next_level.reserve((chunks.size() + 1) / 2); + for (size_t i = 0; i < chunks.size(); i += 2) { + if (i + 1 >= chunks.size()) { + next_level.push_back(std::move(chunks[i])); + } else { + next_level.push_back(merge_qef_two_chunks(std::move(chunks[i]), std::move(chunks[i + 1]), stream)); + } + } + chunks = std::move(next_level); + } + return std::move(chunks[0]); +} + +IntersectionOccResult run_occ_impl( + const float* triangles, + int64_t num_triangles, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + int64_t chunk_triangles, + cudaStream_t stream) { + if (num_triangles < 0) throw std::invalid_argument("num_triangles must be non-negative"); + if (chunk_triangles <= 0) throw std::invalid_argument("chunk_triangles must be positive"); + + constexpr int threads = 256; + std::vector chunks; + chunks.reserve(static_cast((num_triangles + chunk_triangles - 1) / chunk_triangles)); + + for (int64_t tri_begin = 0; tri_begin < num_triangles; tri_begin += chunk_triangles) { + const int64_t tri_count = std::min(chunk_triangles, num_triangles - tri_begin); + if (tri_count == 0) continue; + + DeviceBuffer counts(tri_count); + const int blocks = fdg_gpu::ceil_div_i64(tri_count, threads); + intersection_count_kernel<<>>( + triangles, + tri_begin, + tri_count, + voxel_size.x, + voxel_size.y, + voxel_size.z, + grid_min, + grid_max, + counts.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + + DeviceBuffer offsets(tri_count); + thrust::exclusive_scan( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(counts.data()), + thrust::device_pointer_cast(counts.data()) + tri_count, + thrust::device_pointer_cast(offsets.data())); + + const int64_t last_count = copy_last_i64(counts.data(), tri_count, stream); + const int64_t last_offset = copy_last_i64(offsets.data(), tri_count, stream); + const int64_t raw_size = last_offset + last_count; + if (raw_size == 0) continue; + + DeviceBuffer partial_keys(raw_size); + intersection_occ_emit_kernel<<>>( + triangles, + tri_begin, + tri_count, + voxel_size.x, + voxel_size.y, + voxel_size.z, + grid_min, + grid_max, + offsets.data(), + partial_keys.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + + thrust::sort( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(partial_keys.data()), + thrust::device_pointer_cast(partial_keys.data()) + raw_size); + + auto partial_end = thrust::unique( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(partial_keys.data()), + thrust::device_pointer_cast(partial_keys.data()) + raw_size); + const int64_t partial_size = partial_end - thrust::device_pointer_cast(partial_keys.data()); + if (partial_size == 0) continue; + + chunks.push_back(make_occ_chunk_exact(partial_keys, partial_size, stream)); + } + + OccChunk final_chunk = final_merge_occ_chunks(std::move(chunks), stream); + + IntersectionOccResult out; + out.size = final_chunk.size; + out.voxels.allocate(out.size * 3); + if (out.size == 0) return out; + + const int blocks = fdg_gpu::ceil_div_i64(out.size, threads); + decode_occ_output_kernel<<>>( + final_chunk.keys.data(), out.size, grid_min, grid_max, out.voxels.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + return out; +} + +IntersectQEFResult run_qef_impl( + const float* triangles, + int64_t num_triangles, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + int64_t chunk_triangles, + cudaStream_t stream) { + if (num_triangles < 0) throw std::invalid_argument("num_triangles must be non-negative"); + if (chunk_triangles <= 0) throw std::invalid_argument("chunk_triangles must be positive"); + + constexpr int threads = 256; + std::vector chunks; + chunks.reserve(static_cast((num_triangles + chunk_triangles - 1) / chunk_triangles)); + + for (int64_t tri_begin = 0; tri_begin < num_triangles; tri_begin += chunk_triangles) { + const int64_t tri_count = std::min(chunk_triangles, num_triangles - tri_begin); + if (tri_count == 0) continue; + + DeviceBuffer counts(tri_count); + const int blocks = fdg_gpu::ceil_div_i64(tri_count, threads); + intersection_count_kernel<<>>( + triangles, + tri_begin, + tri_count, + voxel_size.x, + voxel_size.y, + voxel_size.z, + grid_min, + grid_max, + counts.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + + DeviceBuffer offsets(tri_count); + thrust::exclusive_scan( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(counts.data()), + thrust::device_pointer_cast(counts.data()) + tri_count, + thrust::device_pointer_cast(offsets.data())); + + const int64_t last_count = copy_last_i64(counts.data(), tri_count, stream); + const int64_t last_offset = copy_last_i64(offsets.data(), tri_count, stream); + const int64_t raw_size = last_offset + last_count; + if (raw_size == 0) continue; + + DeviceBuffer partial_keys(raw_size); + DeviceBuffer partial_values(raw_size); + intersect_qef_emit_kernel<<>>( + triangles, + tri_begin, + tri_count, + voxel_size.x, + voxel_size.y, + voxel_size.z, + grid_min, + grid_max, + offsets.data(), + partial_keys.data(), + partial_values.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + + thrust::sort_by_key( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(partial_keys.data()), + thrust::device_pointer_cast(partial_keys.data()) + raw_size, + thrust::device_pointer_cast(partial_values.data())); + + DeviceBuffer reduced_keys(raw_size); + DeviceBuffer reduced_values(raw_size); + auto reduce_end = thrust::reduce_by_key( + thrust::cuda::par.on(stream), + thrust::device_pointer_cast(partial_keys.data()), + thrust::device_pointer_cast(partial_keys.data()) + raw_size, + thrust::device_pointer_cast(partial_values.data()), + thrust::device_pointer_cast(reduced_keys.data()), + thrust::device_pointer_cast(reduced_values.data()), + thrust::equal_to(), + AddQEFEventValue()); + const int64_t reduced_size = reduce_end.first - thrust::device_pointer_cast(reduced_keys.data()); + if (reduced_size == 0) continue; + + chunks.push_back(make_qef_chunk_exact(reduced_keys, reduced_values, reduced_size, stream)); + } + + QEFChunk final_chunk = final_merge_qef_chunks(std::move(chunks), stream); + + IntersectQEFResult out; + out.size = final_chunk.size; + out.voxels.allocate(out.size * 3); + out.mean_sum.allocate(out.size * 3); + out.cnt.allocate(out.size); + out.intersected.allocate(out.size); + out.qefs.allocate(out.size); + if (out.size == 0) return out; + + const int blocks = fdg_gpu::ceil_div_i64(out.size, threads); + decode_qef_output_kernel<<>>( + final_chunk.keys.data(), + final_chunk.values.data(), + out.size, + grid_min, + grid_max, + out.voxels.data(), + out.mean_sum.data(), + out.cnt.data(), + out.intersected.data(), + out.qefs.data()); + IQ_CUDA_CHECK(cudaGetLastError()); + return out; +} + +} // namespace + +IntersectionOccResult intersection_occ_gpu( + const float* triangles, + int64_t num_triangles, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + int64_t chunk_triangles, + cudaStream_t stream) { + if (triangles == nullptr && num_triangles > 0) throw std::invalid_argument("triangles is null"); + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + throw std::invalid_argument("voxel_size must be positive"); + } + return run_occ_impl(triangles, num_triangles, voxel_size, grid_min, grid_max, chunk_triangles, stream); +} + +IntersectQEFResult intersect_qef_gpu( + const float* triangles, + int64_t num_triangles, + float3 voxel_size, + int3_ grid_min, + int3_ grid_max, + int64_t chunk_triangles, + cudaStream_t stream) { + if (triangles == nullptr && num_triangles > 0) throw std::invalid_argument("triangles is null"); + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + throw std::invalid_argument("voxel_size must be positive"); + } + return run_qef_impl(triangles, num_triangles, voxel_size, grid_min, grid_max, chunk_triangles, stream); +} + +} // namespace intersection_qef diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.h b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.h new file mode 100644 index 00000000..c5bebcc5 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/intersection_qef.h @@ -0,0 +1,39 @@ +#pragma once + +#include "fdg_gpu_common.h" + +namespace intersection_qef { + +struct IntersectionOccResult { + int64_t size = 0; + fdg_gpu::DeviceBuffer voxels; // [size,3] flattened +}; + +struct IntersectQEFResult { + int64_t size = 0; + fdg_gpu::DeviceBuffer voxels; // [size,3] flattened + fdg_gpu::DeviceBuffer mean_sum; // [size,3] flattened + fdg_gpu::DeviceBuffer cnt; // [size] + fdg_gpu::DeviceBuffer intersected; // [size], bitmask for bool3 + fdg_gpu::DeviceBuffer qefs;// [size] +}; + +IntersectionOccResult intersection_occ_gpu( + const float* triangles, // [num_triangles, 3, 3] flattened + int64_t num_triangles, + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int64_t chunk_triangles = 4096, + cudaStream_t stream = nullptr); + +IntersectQEFResult intersect_qef_gpu( + const float* triangles, // [num_triangles, 3, 3] flattened + int64_t num_triangles, + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int64_t chunk_triangles = 4096, + cudaStream_t stream = nullptr); + +} // namespace intersection_qef diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/mesh_to_flexible_dual_grid.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/mesh_to_flexible_dual_grid.cu new file mode 100644 index 00000000..543bdf13 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/mesh_to_flexible_dual_grid.cu @@ -0,0 +1,897 @@ +#include "../api.h" + +#include "qef.cuh" + +#include +#include +#include +#include +#include + +#include +#include +#include + +// Full CUDA flexible dual grid pipeline. It follows the CPU algorithm's +// semantics, but each stage is expressed as GPU work: gather triangles, build +// intersection QEFs, add face and boundary QEFs in-place, then solve one QEF per +// active voxel. +namespace o_voxel::fdg +{ + namespace + { + + constexpr int kThreads = 256; + constexpr uint64_t kEmptyEdgeKey = UINT64_MAX; + + __host__ __device__ __forceinline__ int64_t div_up_i64(int64_t n, int64_t d) + { + return (n + d - 1) / d; + } + + int64_t next_power_of_two_i64(int64_t x) + { + int64_t out = 1; + while (out < x) + out <<= 1; + return out; + } + + __host__ __device__ __forceinline__ uint64_t pack_edge_key(int32_t a, int32_t b) + { + // Callers pass endpoints sorted as (min, max), so the same + // undirected mesh edge from two faces lands in the same hash slot. + return (static_cast(static_cast(a)) << 32) | + static_cast(b); + } + + __host__ __device__ __forceinline__ int32_t edge_key_v0(uint64_t key) + { + return static_cast(key >> 32); + } + + __host__ __device__ __forceinline__ int32_t edge_key_v1(uint64_t key) + { + return static_cast(key & 0xffffffffu); + } + + __device__ __forceinline__ uint64_t mix64(uint64_t x) + { + x ^= x >> 33; + x *= 0xff51afd7ed558ccdULL; + x ^= x >> 33; + x *= 0xc4ceb9fe1a85ec53ULL; + x ^= x >> 33; + return x; + } + + namespace small_cpqr + { + // Small column-pivoted QR routines used by the constrained QEF + // solver. This keeps the CUDA solve aligned with the CPU path + // instead of using a closed-form solver with different numerical behavior. + namespace detail + { + + template + __device__ __forceinline__ float absf(float x) + { + return x < 0.0f ? -x : x; + } + + template + __device__ __forceinline__ void swap_cols( + float *qr, + float *col_norms_updated, + float *col_norms_direct, + int *perm, + int c0, + int c1) + { + if (c0 == c1) + return; + for (int r = 0; r < N; ++r) + { + const float tmp = qr[r * N + c0]; + qr[r * N + c0] = qr[r * N + c1]; + qr[r * N + c1] = tmp; + } + const float tmp_u = col_norms_updated[c0]; + col_norms_updated[c0] = col_norms_updated[c1]; + col_norms_updated[c1] = tmp_u; + const float tmp_d = col_norms_direct[c0]; + col_norms_direct[c0] = col_norms_direct[c1]; + col_norms_direct[c1] = tmp_d; + const int tmp_p = perm[c0]; + perm[c0] = perm[c1]; + perm[c1] = tmp_p; + } + + template + __device__ __forceinline__ void make_householder_real( + float x0, + const float *tail_in, + int tail_len, + float *beta, + float *tau, + float *essential_out) + { + // Build a Householder reflector that zeroes the tail of a + // column. This is the small-matrix version of QR + // factorization used by the voxel-local QEF solve. + float tail_sq_norm = 0.0f; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + tail_sq_norm += tail_in[i] * tail_in[i]; + + if (tail_sq_norm <= FLT_MIN) + { + *beta = x0; + *tau = 0.0f; + for (int i = 0; i < N - 1; ++i) + essential_out[i] = 0.0f; + return; + } + + float b = sqrtf(x0 * x0 + tail_sq_norm); + if (x0 >= 0.0f) + b = -b; + const float denom = x0 - b; + for (int i = 0; i < N - 1; ++i) + essential_out[i] = (i < tail_len) ? (tail_in[i] / denom) : 0.0f; + *beta = b; + *tau = (b - x0) / b; + } + + template + __device__ __forceinline__ void apply_householder_left_matrix( + float *qr, + int row0, + int col0, + const float *essential, + int tail_len, + float tau) + { + if (tau == 0.0f) + return; + for (int j = 0; j < N; ++j) + { + if (j < col0) + continue; + float tmp = qr[row0 * N + j]; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + tmp += essential[i] * qr[(row0 + 1 + i) * N + j]; + qr[row0 * N + j] -= tau * tmp; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + qr[(row0 + 1 + i) * N + j] -= tau * essential[i] * tmp; + } + } + + template + __device__ __forceinline__ void apply_householder_left_vector( + float *c, + int row0, + const float *essential, + int tail_len, + float tau) + { + if (tau == 0.0f) + return; + float tmp = c[row0]; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + tmp += essential[i] * c[row0 + 1 + i]; + c[row0] -= tau * tmp; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + c[row0 + 1 + i] -= tau * essential[i] * tmp; + } + + template + __device__ __forceinline__ void backsolve_upper_ranked( + const float *qr, + int rank, + const float *c, + const int *perm, + float *x_out) + { + float y[N]; + for (int i = 0; i < N; ++i) + { + y[i] = 0.0f; + x_out[i] = 0.0f; + } + for (int i = rank - 1; i >= 0; --i) + { + float s = c[i]; + for (int j = 0; j < N; ++j) + if (j > i && j < rank) + s -= qr[i * N + j] * y[j]; + y[i] = s / qr[i * N + i]; + } + for (int i = 0; i < N; ++i) + x_out[perm[i]] = (i < rank) ? y[i] : 0.0f; + } + + template + __device__ __forceinline__ void cpqr_solve_small_impl( + const float *A_in, + const float *b_in, + float *x_out) + { + // Column-pivoted QR handles rank-deficient QEF systems: if + // the plane constraints do not determine all coordinates, + // low-norm columns are moved to the end and dropped. + float qr[N * N]; + float c[N]; + int perm[N]; + float col_norms_direct[N]; + float col_norms_updated[N]; + float essential[N > 1 ? N - 1 : 1]; + + for (int i = 0; i < N * N; ++i) + qr[i] = A_in[i]; + for (int i = 0; i < N; ++i) + { + c[i] = b_in[i]; + perm[i] = i; + x_out[i] = 0.0f; + } + + // Track both exact and downdated column norms. Pivoting on + // the largest remaining norm improves stability for nearly + // degenerate local systems. + for (int j = 0; j < N; ++j) + { + float norm_sq = 0.0f; + for (int r = 0; r < N; ++r) + norm_sq += qr[r * N + j] * qr[r * N + j]; + const float norm = sqrtf(norm_sq); + col_norms_direct[j] = norm; + col_norms_updated[j] = norm; + } + + float max_norm_updated = col_norms_updated[0]; + for (int j = 1; j < N; ++j) + if (col_norms_updated[j] > max_norm_updated) + max_norm_updated = col_norms_updated[j]; + const float threshold_helper = (max_norm_updated * FLT_EPSILON) * (max_norm_updated * FLT_EPSILON) / float(N); + const float norm_downdate_threshold = sqrtf(FLT_EPSILON); + int nonzero_pivots = N; + + for (int k = 0; k < N; ++k) + { + // Choose the strongest remaining column as the next + // pivot, then apply a Householder reflector to form R. + int biggest_col_index = k; + float best_updated = col_norms_updated[k]; + for (int j = 0; j < N; ++j) + if (j > k && col_norms_updated[j] > best_updated) + { + best_updated = col_norms_updated[j]; + biggest_col_index = j; + } + if (nonzero_pivots == N && best_updated * best_updated < threshold_helper * float(N - k)) + nonzero_pivots = k; + + swap_cols(qr, col_norms_updated, col_norms_direct, perm, k, biggest_col_index); + const int tail_len = N - k - 1; + float tail_local[N > 1 ? N - 1 : 1]; + for (int i = 0; i < N - 1; ++i) + tail_local[i] = (i < tail_len) ? qr[(k + 1 + i) * N + k] : 0.0f; + + float beta = 0.0f; + float tau = 0.0f; + make_householder_real(qr[k * N + k], tail_local, tail_len, &beta, &tau, essential); + qr[k * N + k] = beta; + for (int i = 0; i < N - 1; ++i) + if (i < tail_len) + qr[(k + 1 + i) * N + k] = essential[i]; + + apply_householder_left_matrix(qr, k, k + 1, essential, tail_len, tau); + if (k < nonzero_pivots) + apply_householder_left_vector(c, k, essential, tail_len, tau); + + // Cheap norm downdates avoid recomputing every column + // norm after each reflector; recompute only when the + // downdate becomes unreliable. + for (int j = 0; j < N; ++j) + { + if (j <= k || col_norms_updated[j] == 0.0f) + continue; + float temp = absf(qr[k * N + j]) / col_norms_updated[j]; + temp = (1.0f + temp) * (1.0f - temp); + if (temp < 0.0f) + temp = 0.0f; + const float ratio = col_norms_updated[j] / col_norms_direct[j]; + const float temp2 = temp * ratio * ratio; + if (temp2 <= norm_downdate_threshold) + { + float norm_sq = 0.0f; + for (int r = 0; r < N; ++r) + if (r > k) + norm_sq += qr[r * N + j] * qr[r * N + j]; + const float norm = sqrtf(norm_sq); + col_norms_direct[j] = norm; + col_norms_updated[j] = norm; + } + else + { + col_norms_updated[j] *= sqrtf(temp); + } + } + } + + if (nonzero_pivots == 0) + return; + // Solve the retained upper-triangular part and undo the + // pivot permutation. Dropped columns stay zero. + backsolve_upper_ranked(qr, nonzero_pivots, c, perm, x_out); + } + + } // namespace detail + + __device__ __forceinline__ void cpqr_solve_3x3(const float A[9], const float b[3], float x[3]) + { + detail::cpqr_solve_small_impl<3>(A, b, x); + } + + __device__ __forceinline__ void cpqr_solve_2x2(const float A[4], const float b[2], float x[2]) + { + detail::cpqr_solve_small_impl<2>(A, b, x); + } + + __device__ __forceinline__ float solve_1x1_unchecked(float a, float rhs) + { + return rhs / a; + } + + } // namespace small_cpqr + + __host__ __device__ __forceinline__ int idx4(int r, int c) { return r * 4 + c; } + __host__ __device__ __forceinline__ int idx2(int r, int c) { return r * 2 + c; } + + __device__ __forceinline__ void sym10_to_dense4x4(const SymQEF10 &q, float Q[16]) + { + // Expand compact symmetric storage so the solver can evaluate + // p^T Q p and extract 3x3/2x2 constrained systems directly. + Q[idx4(0, 0)] = q.q00; + Q[idx4(0, 1)] = q.q01; + Q[idx4(0, 2)] = q.q02; + Q[idx4(0, 3)] = q.q03; + Q[idx4(1, 0)] = q.q01; + Q[idx4(1, 1)] = q.q11; + Q[idx4(1, 2)] = q.q12; + Q[idx4(1, 3)] = q.q13; + Q[idx4(2, 0)] = q.q02; + Q[idx4(2, 1)] = q.q12; + Q[idx4(2, 2)] = q.q22; + Q[idx4(2, 3)] = q.q23; + Q[idx4(3, 0)] = q.q03; + Q[idx4(3, 1)] = q.q13; + Q[idx4(3, 2)] = q.q23; + Q[idx4(3, 3)] = q.q33; + } + + __device__ __forceinline__ bool point_inside_box3(const float v[3], const float min_corner[3], const float max_corner[3]) + { + return v[0] >= min_corner[0] && v[0] <= max_corner[0] && + v[1] >= min_corner[1] && v[1] <= max_corner[1] && + v[2] >= min_corner[2] && v[2] <= max_corner[2]; + } + + __device__ __forceinline__ float qef_error4(const float Q[16], const float p[4]) + { + // Homogeneous point p=(x,y,z,1). The scalar p^T Q p is the total + // squared plane/line distance represented by the accumulated QEF. + const float y0 = Q[idx4(0, 0)] * p[0] + Q[idx4(0, 1)] * p[1] + Q[idx4(0, 2)] * p[2] + Q[idx4(0, 3)] * p[3]; + const float y1 = Q[idx4(1, 0)] * p[0] + Q[idx4(1, 1)] * p[1] + Q[idx4(1, 2)] * p[2] + Q[idx4(1, 3)] * p[3]; + const float y2 = Q[idx4(2, 0)] * p[0] + Q[idx4(2, 1)] * p[1] + Q[idx4(2, 2)] * p[2] + Q[idx4(2, 3)] * p[3]; + const float y3 = Q[idx4(3, 0)] * p[0] + Q[idx4(3, 1)] * p[1] + Q[idx4(3, 2)] * p[2] + Q[idx4(3, 3)] * p[3]; + return p[0] * y0 + p[1] * y1 + p[2] * y2 + p[3] * y3; + } + + __device__ __forceinline__ void add_qef_regularization_inplace( + float Q[16], + const float mean_sum[3], + float cnt, + float regularization_weight) + { + if (regularization_weight <= 0.0f || cnt <= 0.0f) + return; + + // Regularization adds w * ||x - mean_intersection||^2. It keeps the + // solve near the observed intersection points when QEF planes alone + // are under-constrained. + const float px = mean_sum[0] / cnt; + const float py = mean_sum[1] / cnt; + const float pz = mean_sum[2] / cnt; + const float w = regularization_weight * cnt; + + Q[idx4(0, 0)] += w; + Q[idx4(1, 1)] += w; + Q[idx4(2, 2)] += w; + Q[idx4(0, 3)] += -w * px; + Q[idx4(1, 3)] += -w * py; + Q[idx4(2, 3)] += -w * pz; + Q[idx4(3, 0)] += -w * px; + Q[idx4(3, 1)] += -w * py; + Q[idx4(3, 2)] += -w * pz; + Q[idx4(3, 3)] += w * (px * px + py * py + pz * pz); + } + + __device__ __forceinline__ void try_single_constraint( + const float Q[16], + int fixed_axis, + const float min_corner[3], + const float max_corner[3], + float &best, + float v_new[3]) + { + const int ax1 = (fixed_axis + 1) % 3; + const int ax2 = (fixed_axis + 2) % 3; + // Candidate on one voxel face: fix one coordinate to min or max and + // solve the remaining 2D QEF on that face. + float A2[4] = { + Q[idx4(ax1, ax1)], Q[idx4(ax1, ax2)], + Q[idx4(ax2, ax1)], Q[idx4(ax2, ax2)]}; + float B2[4] = { + Q[idx4(ax1, fixed_axis)], Q[idx4(ax1, 3)], + Q[idx4(ax2, fixed_axis)], Q[idx4(ax2, 3)]}; + float x2[2]; + + for (int bound = 0; bound < 2; ++bound) + { + const float fixed = bound == 0 ? min_corner[fixed_axis] : max_corner[fixed_axis]; + const float rhs2[2] = { + -(B2[idx2(0, 0)] * fixed + B2[idx2(0, 1)]), + -(B2[idx2(1, 0)] * fixed + B2[idx2(1, 1)])}; + small_cpqr::cpqr_solve_2x2(A2, rhs2, x2); + if (x2[0] >= min_corner[ax1] && x2[0] <= max_corner[ax1] && + x2[1] >= min_corner[ax2] && x2[1] <= max_corner[ax2]) + { + float p4[4]; + p4[fixed_axis] = fixed; + p4[ax1] = x2[0]; + p4[ax2] = x2[1]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) + { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + } + } + + __device__ __forceinline__ void try_two_constraint( + const float Q[16], + int free_axis, + const float min_corner[3], + const float max_corner[3], + float &best, + float v_new[3]) + { + const int ax1 = (free_axis + 1) % 3; + const int ax2 = (free_axis + 2) % 3; + // Candidate on one voxel edge: fix two coordinates to box bounds and + // solve the remaining 1D minimizer along the free axis. + const float a = Q[idx4(free_axis, free_axis)]; + const float b0 = Q[idx4(free_axis, ax1)]; + const float b1 = Q[idx4(free_axis, ax2)]; + const float b2 = Q[idx4(free_axis, 3)]; + + for (int c0 = 0; c0 < 2; ++c0) + { + for (int c1 = 0; c1 < 2; ++c1) + { + const float v0 = c0 == 0 ? min_corner[ax1] : max_corner[ax1]; + const float v1 = c1 == 0 ? min_corner[ax2] : max_corner[ax2]; + const float x = small_cpqr::solve_1x1_unchecked(a, -(b0 * v0 + b1 * v1 + b2)); + if (x >= min_corner[free_axis] && x <= max_corner[free_axis]) + { + float p4[4]; + p4[free_axis] = x; + p4[ax1] = v0; + p4[ax2] = v1; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) + { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + } + } + } + + __device__ __forceinline__ void try_three_constraint( + const float Q[16], + const float min_corner[3], + const float max_corner[3], + float &best, + float v_new[3]) + { + // Final fallback candidates are the eight box corners. This makes + // the constrained search complete over the voxel box boundary. + for (int cx = 0; cx < 2; ++cx) + { + for (int cy = 0; cy < 2; ++cy) + { + for (int cz = 0; cz < 2; ++cz) + { + float p4[4]; + p4[0] = cx ? min_corner[0] : max_corner[0]; + p4[1] = cy ? min_corner[1] : max_corner[1]; + p4[2] = cz ? min_corner[2] : max_corner[2]; + p4[3] = 1.0f; + const float err = qef_error4(Q, p4); + if (err < best) + { + best = err; + v_new[0] = p4[0]; + v_new[1] = p4[1]; + v_new[2] = p4[2]; + } + } + } + } + } + + __global__ void gather_triangles_kernel( + const float *__restrict__ vertices, + const int32_t *__restrict__ faces, + int64_t num_faces, + float *__restrict__ triangles) + { + const int64_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid >= 3 * num_faces) + return; + + const int64_t f = tid / 3; + const int lv = static_cast(tid - 3 * f); + const int32_t vid = faces[3 * f + lv]; + // One thread copies one face vertex, producing a dense [F,3,3] + // triangle tensor consumed by the three QEF stages. + triangles[3 * tid + 0] = vertices[3 * static_cast(vid) + 0]; + triangles[3 * tid + 1] = vertices[3 * static_cast(vid) + 1]; + triangles[3 * tid + 2] = vertices[3 * static_cast(vid) + 2]; + } + + __global__ void count_edges_kernel( + int64_t num_faces, + const int32_t *__restrict__ faces, + uint64_t *__restrict__ hash_keys, + uint32_t *__restrict__ edge_counts, + int32_t *__restrict__ overflow, + uint64_t hash_capacity) + { + const int64_t fid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (fid >= num_faces) + return; + + const int32_t v[3] = { + faces[3 * fid + 0], + faces[3 * fid + 1], + faces[3 * fid + 2], + }; + for (int e = 0; e < 3; ++e) + { + int32_t a = v[e]; + int32_t b = v[(e + 1) % 3]; + if (a > b) + { + const int32_t tmp = a; + a = b; + b = tmp; + } + // Count each undirected edge. Interior manifold edges are seen + // twice; boundary edges are seen once. + const uint64_t key = pack_edge_key(a, b); + uint64_t slot = mix64(key) & (hash_capacity - 1); + bool inserted = false; + for (uint64_t probe = 0; probe < hash_capacity; ++probe) + { + const uint64_t old = atomicCAS( + reinterpret_cast(hash_keys + slot), + static_cast(kEmptyEdgeKey), + static_cast(key)); + if (old == kEmptyEdgeKey || old == key) + { + atomicAdd(edge_counts + slot, 1u); + inserted = true; + break; + } + slot = (slot + 1u) & (hash_capacity - 1); + } + if (!inserted) + atomicExch(overflow, 1); + } + } + + __global__ void count_boundary_edges_kernel( + uint64_t hash_capacity, + const uint64_t *__restrict__ hash_keys, + const uint32_t *__restrict__ edge_counts, + uint32_t *__restrict__ boundary_count) + { + const uint64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= hash_capacity) + return; + // A valid hash slot with count one is exactly one boundary edge. + if (hash_keys[i] != kEmptyEdgeKey && edge_counts[i] == 1u) + atomicAdd(boundary_count, 1u); + } + + __global__ void emit_boundaries_kernel( + uint64_t hash_capacity, + const uint64_t *__restrict__ hash_keys, + const uint32_t *__restrict__ edge_counts, + const float *__restrict__ vertices, + uint32_t *__restrict__ boundary_count, + float *__restrict__ boundaries) + { + const uint64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= hash_capacity) + return; + if (hash_keys[i] == kEmptyEdgeKey || edge_counts[i] != 1u) + return; + + const uint32_t out = atomicAdd(boundary_count, 1u); + const int32_t v0 = edge_key_v0(hash_keys[i]); + const int32_t v1 = edge_key_v1(hash_keys[i]); + // Emit the two endpoint positions in grid-local coordinates for the + // boundary DDA kernel. + boundaries[6 * static_cast(out) + 0] = vertices[3 * static_cast(v0) + 0]; + boundaries[6 * static_cast(out) + 1] = vertices[3 * static_cast(v0) + 1]; + boundaries[6 * static_cast(out) + 2] = vertices[3 * static_cast(v0) + 2]; + boundaries[6 * static_cast(out) + 3] = vertices[3 * static_cast(v1) + 0]; + boundaries[6 * static_cast(out) + 4] = vertices[3 * static_cast(v1) + 1]; + boundaries[6 * static_cast(out) + 5] = vertices[3 * static_cast(v1) + 2]; + } + + __global__ void solve_qef_kernel( + const int32_t *__restrict__ voxels, + const float *__restrict__ mean_sum, + const float *__restrict__ cnt, + const SymQEF10 *__restrict__ qefs, + int64_t n, + float3 voxel_size, + float regularization_weight, + float *__restrict__ dual_vertices) + { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) + return; + + const int x = voxels[3 * i + 0]; + const int y = voxels[3 * i + 1]; + const int z = voxels[3 * i + 2]; + // The dual vertex for this voxel must remain inside this voxel box. + const float min_corner[3] = { + x * voxel_size.x, + y * voxel_size.y, + z * voxel_size.z, + }; + const float max_corner[3] = { + (x + 1) * voxel_size.x, + (y + 1) * voxel_size.y, + (z + 1) * voxel_size.z, + }; + + float Q[16]; + sym10_to_dense4x4(qefs[i], Q); + const float mean_i[3] = { + mean_sum[3 * i + 0], + mean_sum[3 * i + 1], + mean_sum[3 * i + 2], + }; + add_qef_regularization_inplace(Q, mean_i, cnt[i], regularization_weight); + + // First solve the unconstrained 3D minimizer of p^T Q p with p.w=1. + const float A3[9] = { + Q[idx4(0, 0)], Q[idx4(0, 1)], Q[idx4(0, 2)], + Q[idx4(1, 0)], Q[idx4(1, 1)], Q[idx4(1, 2)], + Q[idx4(2, 0)], Q[idx4(2, 1)], Q[idx4(2, 2)]}; + const float b3[3] = {-Q[idx4(0, 3)], -Q[idx4(1, 3)], -Q[idx4(2, 3)]}; + float v_new[3]; + small_cpqr::cpqr_solve_3x3(A3, b3, v_new); + + if (!point_inside_box3(v_new, min_corner, max_corner)) + { + // If the best unconstrained point leaves the voxel, search the + // voxel boundary: first faces, then edges, then corners. Each + // candidate is scored with the full 4x4 QEF. + float best = CUDART_INF_F; + try_single_constraint(Q, 0, min_corner, max_corner, best, v_new); + try_single_constraint(Q, 1, min_corner, max_corner, best, v_new); + try_single_constraint(Q, 2, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 0, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 1, min_corner, max_corner, best, v_new); + try_two_constraint(Q, 2, min_corner, max_corner, best, v_new); + try_three_constraint(Q, min_corner, max_corner, best, v_new); + } + + dual_vertices[3 * i + 0] = v_new[0]; + dual_vertices[3 * i + 1] = v_new[1]; + dual_vertices[3 * i + 2] = v_new[2]; + } + + torch::Tensor extract_boundaries_cuda( + const torch::Tensor &vertices, + const torch::Tensor &faces) + { + // Boundary extraction uses a hash table over undirected mesh edges. + // An edge seen exactly once is a boundary edge and is emitted as a + // [2, 3] segment for boundary_qef_cuda. + const c10::cuda::CUDAGuard guard(vertices.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(vertices.get_device()).stream(); + const torch::Device device = vertices.device(); + const int64_t num_faces = faces.size(0); + const auto opts_u64 = torch::TensorOptions().dtype(torch::kUInt64).device(device); + const auto opts_u32 = torch::TensorOptions().dtype(torch::kUInt32).device(device); + const auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(device); + if (num_faces == 0) + return torch::empty({0, 2, 3}, opts_f32); + + const int64_t num_edges = num_faces * 3; + const int64_t hash_capacity_i64 = next_power_of_two_i64(num_edges * 2 > 2 ? num_edges * 2 : 2); + auto hash_keys = torch::empty({hash_capacity_i64}, opts_u64); + auto edge_counts = torch::zeros({hash_capacity_i64}, opts_u32); + auto boundary_count_t = torch::zeros({1}, opts_u32); + auto overflow = torch::zeros({1}, opts_i32); + C10_CUDA_CHECK(cudaMemsetAsync(hash_keys.data_ptr(), 0xff, hash_capacity_i64 * sizeof(uint64_t), stream)); + + int blocks = static_cast(div_up_i64(num_faces, kThreads)); + count_edges_kernel<<>>( + num_faces, + faces.data_ptr(), + hash_keys.data_ptr(), + edge_counts.data_ptr(), + overflow.data_ptr(), + static_cast(hash_capacity_i64)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + blocks = static_cast(div_up_i64(hash_capacity_i64, kThreads)); + count_boundary_edges_kernel<<>>( + static_cast(hash_capacity_i64), + hash_keys.data_ptr(), + edge_counts.data_ptr(), + boundary_count_t.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + int32_t overflow_h = 0; + uint32_t boundary_count = 0; + C10_CUDA_CHECK(cudaMemcpyAsync(&overflow_h, overflow.data_ptr(), sizeof(int32_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaMemcpyAsync(&boundary_count, boundary_count_t.data_ptr(), sizeof(uint32_t), cudaMemcpyDeviceToHost, stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + TORCH_CHECK(overflow_h == 0, "edge hash overflow while extracting boundaries"); + + auto boundaries = torch::empty({static_cast(boundary_count), 2, 3}, opts_f32); + if (boundary_count == 0) + return boundaries; + + C10_CUDA_CHECK(cudaMemsetAsync(boundary_count_t.data_ptr(), 0, sizeof(uint32_t), stream)); + emit_boundaries_kernel<<>>( + static_cast(hash_capacity_i64), + hash_keys.data_ptr(), + edge_counts.data_ptr(), + vertices.data_ptr(), + boundary_count_t.data_ptr(), + boundaries.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return boundaries; + } + + } // namespace + + std::tuple + mesh_to_flexible_dual_grid_cuda( + const torch::Tensor &vertices, + const torch::Tensor &faces, + const std::vector &voxel_size, + const std::vector &grid_range, + float face_weight, + float boundary_weight, + float regularization_weight) + { + // CUDA version of the Python/CPU flexible dual grid entry point: + // vertices/faces -> triangles -> intersection QEF -> face QEF -> + // boundary QEF -> constrained solve. + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(faces.is_cuda(), "faces must be a CUDA tensor"); + static_assert(sizeof(SymQEF10) == sizeof(float) * 10, "Unexpected SymQEF10 layout"); + + const c10::cuda::CUDAGuard guard(vertices.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(vertices.get_device()).stream(); + const torch::Device device = vertices.device(); + const auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(device); + + const int64_t num_faces = faces.size(0); + auto triangles = torch::empty({num_faces, 3, 3}, opts_f32); + if (num_faces > 0) + { + const int blocks = static_cast(div_up_i64(num_faces * 3, kThreads)); + gather_triangles_kernel<<>>( + vertices.data_ptr(), + faces.data_ptr(), + num_faces, + triangles.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + auto intersect = intersect_qef_cuda(triangles, voxel_size, grid_range); + torch::Tensor voxels = std::get<0>(intersect); + torch::Tensor mean_sum = std::get<1>(intersect); + torch::Tensor cnt = std::get<2>(intersect); + torch::Tensor intersected = std::get<3>(intersect); + torch::Tensor total_qefs = std::get<4>(intersect); + torch::Tensor brick_hash_keys = std::get<5>(intersect); + torch::Tensor brick_hash_vals = std::get<6>(intersect); + torch::Tensor brick_bits = std::get<7>(intersect); + torch::Tensor brick_base = std::get<8>(intersect); + + const int64_t num_voxels = voxels.size(0); + if (num_voxels == 0) + { + auto dual_vertices = torch::empty({0, 3}, opts_f32); + return std::make_tuple(voxels, dual_vertices, intersected); + } + + if (face_weight > 0.0f) + face_qef_cuda( + triangles, + voxel_size, + grid_range, + voxels, + total_qefs, + face_weight, + brick_hash_keys, + brick_hash_vals, + brick_bits, + brick_base); + + if (boundary_weight > 0.0f) + { + torch::Tensor boundaries = extract_boundaries_cuda(vertices, faces); + if (boundaries.size(0) > 0) + boundary_qef_cuda( + boundaries, + voxel_size, + grid_range, + boundary_weight, + voxels, + total_qefs, + brick_hash_keys, + brick_hash_vals, + brick_bits, + brick_base); + } + + int blocks = static_cast(div_up_i64(num_voxels, kThreads)); + const float3 voxel_size_h = make_float3(voxel_size[0], voxel_size[1], voxel_size[2]); + auto dual_vertices = torch::empty({num_voxels, 3}, opts_f32); + solve_qef_kernel<<>>( + voxels.data_ptr(), + mean_sum.data_ptr(), + cnt.data_ptr(), + reinterpret_cast(total_qefs.data_ptr()), + num_voxels, + voxel_size_h, + regularization_weight, + dual_vertices.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return std::make_tuple(voxels, dual_vertices, intersected); + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/qef.cuh b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/qef.cuh new file mode 100644 index 00000000..45fd1952 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/qef.cuh @@ -0,0 +1,74 @@ +#pragma once + +#include "types.cuh" + +namespace o_voxel::fdg +{ + + __host__ __device__ __forceinline__ SymQEF10 qef_zero() + { + return SymQEF10{0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + } + + // Small value helpers kept host/device so the same QEF math can be used in + // kernels and in any light host-side setup code. + __host__ __device__ __forceinline__ SymQEF10 qef_add( + const SymQEF10 &a, + const SymQEF10 &b) + { + return SymQEF10{ + a.q00 + b.q00, + a.q01 + b.q01, + a.q02 + b.q02, + a.q03 + b.q03, + a.q11 + b.q11, + a.q12 + b.q12, + a.q13 + b.q13, + a.q22 + b.q22, + a.q23 + b.q23, + a.q33 + b.q33, + }; + } + + // Scaling folds a stage weight into a QEF before it is accumulated into the + // running per-voxel total. + __host__ __device__ __forceinline__ SymQEF10 qef_scale( + const SymQEF10 &q, + float s) + { + return SymQEF10{ + q.q00 * s, + q.q01 * s, + q.q02 * s, + q.q03 * s, + q.q11 * s, + q.q12 * s, + q.q13 * s, + q.q22 * s, + q.q23 * s, + q.q33 * s, + }; + } + + __host__ __device__ __forceinline__ SymQEF10 qef_from_plane(float4 p) + { + // Plane ax + by + cz + d = 0 contributes p * p^T in SymQEF10 layout. + const float a = p.x; + const float b = p.y; + const float c = p.z; + const float d = p.w; + return SymQEF10{ + a * a, + a * b, + a * c, + a * d, + b * b, + b * c, + b * d, + c * c, + c * d, + d * d, + }; + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/torch_bindings.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/torch_bindings.cu new file mode 100644 index 00000000..2c35c4ed --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/torch_bindings.cu @@ -0,0 +1,936 @@ +#include +#include + +#include + +#include +#include +#include +#include +#include + +#include "../api.h" +#include "flexible_dual_grid_gpu.h" +#include "intersection_qef.h" +#include "voxelize_mesh_oct.h" +#include "voxel_traverse_edge_dda.h" + +struct bool3 { bool x, y, z; bool& operator[](int i) { return (&x)[i]; } }; + +struct VoxelCoord { + int x, y, z; + + int& operator[](int i) { return (&x)[i]; } + + bool operator==(const VoxelCoord& other) const { + return x == other.x && y == other.y && z == other.z; + } +}; + +namespace std { +template <> +struct hash { + size_t operator()(const VoxelCoord& v) const { + const std::size_t p1 = 73856093; + const std::size_t p2 = 19349663; + const std::size_t p3 = 83492791; + return static_cast(v.x) * p1 ^ + static_cast(v.y) * p2 ^ + static_cast(v.z) * p3; + } +}; +} // namespace std + +void intersect_qef( + const Eigen::Vector3f& voxel_size, + const Eigen::Vector3i& grid_min, + const Eigen::Vector3i& grid_max, + const std::vector& triangles, + std::unordered_map& hash_table, + std::vector& voxels, + std::vector& means, + std::vector& cnt, + std::vector& intersected, + std::vector& qefs +); + +void face_qef( + const Eigen::Vector3f& voxel_size, + const Eigen::Vector3i& grid_min, + const Eigen::Vector3i& grid_max, + const std::vector& triangles, + std::unordered_map& hash_table, + std::vector& qefs +); + +void boundry_qef( + const Eigen::Vector3f& voxel_size, + const Eigen::Vector3i& grid_min, + const Eigen::Vector3i& grid_max, + const std::vector& boundries, + float boundary_weight, + std::unordered_map& hash_table, + std::vector& qefs +); + +namespace { + +inline void check_cuda_success(cudaError_t err, const char* context) { + TORCH_CHECK(err == cudaSuccess, context, ": ", cudaGetErrorString(err)); +} + +inline float3 tensor_to_float3_cpu(const torch::Tensor& t) { + auto tc = t.to(torch::kFloat32).contiguous().cpu(); + TORCH_CHECK(tc.dim() == 1 && tc.size(0) == 3, "voxel_size must have shape [3]"); + const float* p = tc.data_ptr(); + return float3{p[0], p[1], p[2]}; +} + +inline void tensor_to_grid_min_max_cpu( + const torch::Tensor& t, + fdg_gpu::int3_& grid_min, + fdg_gpu::int3_& grid_max +) { + auto tc = t.to(torch::kInt32).contiguous().cpu(); + TORCH_CHECK(tc.dim() == 2 && tc.size(0) == 2 && tc.size(1) == 3, "grid_range must have shape [2, 3]"); + const int32_t* p = tc.data_ptr(); + grid_min = fdg_gpu::int3_{p[0], p[1], p[2]}; + grid_max = fdg_gpu::int3_{p[3], p[4], p[5]}; +} + +inline fdg_gpu::int3_ grid_size_from_min_max( + const fdg_gpu::int3_& grid_min, + const fdg_gpu::int3_& grid_max +) { + return fdg_gpu::int3_{ + grid_max.x - grid_min.x, + grid_max.y - grid_min.y, + grid_max.z - grid_min.z, + }; +} + +__global__ void unpack_intersected_mask_kernel( + const uint8_t* mask, + int64_t n, + bool* out_bool3 +) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= n) return; + const uint8_t m = mask[i]; + out_bool3[3 * i + 0] = (m & (1u << 0)) != 0; + out_bool3[3 * i + 1] = (m & (1u << 1)) != 0; + out_bool3[3 * i + 2] = (m & (1u << 2)) != 0; +} + +inline void check_triangles_tensor(const torch::Tensor& triangles_c) { + TORCH_CHECK( + triangles_c.dim() == 3 && triangles_c.size(1) == 3 && triangles_c.size(2) == 3, + "triangles must have shape [T, 3, 3]" + ); +} + +inline void check_voxels_tensor(const torch::Tensor& voxels_c) { + TORCH_CHECK( + voxels_c.dim() == 2 && voxels_c.size(1) == 3, + "voxels must have shape [N, 3]" + ); +} + +inline void check_edges_tensor(const torch::Tensor& edges_c) { + TORCH_CHECK(edges_c.dim() == 2 && edges_c.size(1) == 2, "edges must have shape [E, 2]"); +} + +inline void check_boundaries_tensor(const torch::Tensor& boundaries_c) { + TORCH_CHECK( + boundaries_c.dim() == 3 && boundaries_c.size(1) == 2 && boundaries_c.size(2) == 3, + "boundaries must have shape [B, 2, 3]" + ); +} + +inline void check_cpu_tensor(const torch::Tensor& t, const char* name) { + TORCH_CHECK(!t.is_cuda(), name, " must be a CPU tensor"); +} + +inline Eigen::Vector3f tensor_to_eigen_vec3_cpu(const torch::Tensor& t) { + auto tc = t.to(torch::kFloat32).contiguous().cpu(); + TORCH_CHECK(tc.dim() == 1 && tc.size(0) == 3, "voxel_size must have shape [3]"); + const float* p = tc.data_ptr(); + return Eigen::Vector3f(p[0], p[1], p[2]); +} + +inline void tensor_to_eigen_grid_min_max_cpu( + const torch::Tensor& t, + Eigen::Vector3i& grid_min, + Eigen::Vector3i& grid_max +) { + auto tc = t.to(torch::kInt32).contiguous().cpu(); + TORCH_CHECK(tc.dim() == 2 && tc.size(0) == 2 && tc.size(1) == 3, "grid_range must have shape [2, 3]"); + const int32_t* p = tc.data_ptr(); + grid_min = Eigen::Vector3i(p[0], p[1], p[2]); + grid_max = Eigen::Vector3i(p[3], p[4], p[5]); +} + +inline std::vector triangles_tensor_to_vector_cpu(const torch::Tensor& triangles) { + auto triangles_c = triangles.to(torch::kFloat32).contiguous().cpu(); + check_triangles_tensor(triangles_c); + const float* p = triangles_c.data_ptr(); + const int64_t n = triangles_c.size(0); + std::vector out; + out.reserve(static_cast(n) * 3); + for (int64_t i = 0; i < n; ++i) { + for (int v = 0; v < 3; ++v) { + const int64_t base = (i * 3 + v) * 3; + out.emplace_back(p[base + 0], p[base + 1], p[base + 2]); + } + } + return out; +} + +inline std::vector boundaries_tensor_to_vector_cpu(const torch::Tensor& boundaries) { + auto boundaries_c = boundaries.to(torch::kFloat32).contiguous().cpu(); + check_boundaries_tensor(boundaries_c); + const float* p = boundaries_c.data_ptr(); + const int64_t n = boundaries_c.size(0); + std::vector out; + out.reserve(static_cast(n) * 2); + for (int64_t i = 0; i < n; ++i) { + for (int v = 0; v < 2; ++v) { + const int64_t base = (i * 2 + v) * 3; + out.emplace_back(p[base + 0], p[base + 1], p[base + 2]); + } + } + return out; +} + +inline std::vector voxels_tensor_to_vector_cpu( + const torch::Tensor& voxels, + std::unordered_map& hash_table +) { + auto voxels_c = voxels.to(torch::kInt32).contiguous().cpu(); + check_voxels_tensor(voxels_c); + const int32_t* p = voxels_c.data_ptr(); + const int64_t n = voxels_c.size(0); + std::vector out; + out.reserve(static_cast(n)); + hash_table.reserve(static_cast(n)); + for (int64_t i = 0; i < n; ++i) { + const VoxelCoord coord{p[3 * i + 0], p[3 * i + 1], p[3 * i + 2]}; + hash_table[coord] = static_cast(i); + out.push_back(int3{coord.x, coord.y, coord.z}); + } + return out; +} + +inline torch::Tensor int3_vector_to_tensor_cpu(const std::vector& values) { + auto out = torch::empty({static_cast(values.size()), 3}, torch::TensorOptions().dtype(torch::kInt32).device(torch::kCPU)); + int32_t* p = out.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) { + p[3 * i + 0] = values[i].x; + p[3 * i + 1] = values[i].y; + p[3 * i + 2] = values[i].z; + } + return out; +} + +inline torch::Tensor vec3f_vector_to_tensor_cpu(const std::vector& values) { + auto out = torch::empty({static_cast(values.size()), 3}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU)); + float* p = out.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) { + p[3 * i + 0] = values[i].x(); + p[3 * i + 1] = values[i].y(); + p[3 * i + 2] = values[i].z(); + } + return out; +} + +inline torch::Tensor float_vector_to_tensor_cpu(const std::vector& values) { + auto out = torch::empty({static_cast(values.size())}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU)); + float* p = out.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) { + p[i] = values[i]; + } + return out; +} + +inline torch::Tensor bool3_vector_to_tensor_cpu(const std::vector& values) { + auto out = torch::empty({static_cast(values.size()), 3}, torch::TensorOptions().dtype(torch::kBool).device(torch::kCPU)); + bool* p = out.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) { + p[3 * i + 0] = values[i].x; + p[3 * i + 1] = values[i].y; + p[3 * i + 2] = values[i].z; + } + return out; +} + +inline torch::Tensor matrix4f_vector_to_tensor_cpu(const std::vector& values) { + auto out = torch::empty({static_cast(values.size()), 4, 4}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCPU)); + float* p = out.data_ptr(); + for (size_t i = 0; i < values.size(); ++i) { + for (int r = 0; r < 4; ++r) { + for (int c = 0; c < 4; ++c) { + p[i * 16 + r * 4 + c] = values[i](r, c); + } + } + } + return out; +} + +inline std::tuple primitive_pair_to_tensors( + const fdg_gpu::PrimitivePairResult& pairs, + const torch::Device& device, + cudaStream_t stream +) { + auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + torch::Tensor prim_id = torch::empty({pairs.size}, opts_i32); + torch::Tensor voxels_axis_major = torch::empty({3, pairs.size}, opts_i32); + + if (pairs.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + prim_id.data_ptr(), + pairs.prim_id.data(), + static_cast(pairs.size) * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync primitive prim_id" + ); + + check_cuda_success( + cudaMemcpyAsync( + voxels_axis_major.data_ptr() + pairs.size * 0, + pairs.voxel_i.data(), + static_cast(pairs.size) * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync primitive voxel_i" + ); + check_cuda_success( + cudaMemcpyAsync( + voxels_axis_major.data_ptr() + pairs.size * 1, + pairs.voxel_j.data(), + static_cast(pairs.size) * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync primitive voxel_j" + ); + check_cuda_success( + cudaMemcpyAsync( + voxels_axis_major.data_ptr() + pairs.size * 2, + pairs.voxel_k.data(), + static_cast(pairs.size) * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync primitive voxel_k" + ); + + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize primitive_pair_to_tensors"); + } + + torch::Tensor voxels = voxels_axis_major.transpose(0, 1).contiguous(); + return std::make_tuple(prim_id, voxels); +} + +} // namespace + + +std::tuple mesh_to_flexible_dual_grid_gpu( + const torch::Tensor& vertices, + const torch::Tensor& faces, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + float face_weight, + float boundary_weight, + float regularization_weight, + int64_t intersect_chunk_triangles, + int boundary_chunk_steps +) { + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(faces.is_cuda(), "faces must be a CUDA tensor"); + TORCH_CHECK(vertices.device() == faces.device(), "vertices and faces must be on the same CUDA device"); + + auto vertices_c = vertices.to(torch::kFloat32).contiguous(); + auto faces_c = faces.to(torch::kInt32).contiguous(); + + TORCH_CHECK(vertices_c.dim() == 2 && vertices_c.size(1) == 3, "vertices must have shape [V, 3]"); + TORCH_CHECK(faces_c.dim() == 2 && faces_c.size(1) == 3, "faces must have shape [F, 3]"); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + fdg_gpu::FlexibleDualGridGPUOutput out{}; + cudaStream_t stream = nullptr; + + cudaError_t status = fdg_gpu::mesh_to_flexible_dual_grid_gpu( + vertices_c.data_ptr(), + vertices_c.size(0), + faces_c.data_ptr(), + faces_c.size(0), + voxel_size_h, + grid_min, + grid_max, + face_weight, + boundary_weight, + regularization_weight, + intersect_chunk_triangles, + boundary_chunk_steps, + stream, + &out + ); + + if (status != cudaSuccess) { + fdg_gpu::free_flexible_dual_grid_gpu_output(&out); + TORCH_CHECK(false, "mesh_to_flexible_dual_grid_gpu failed: ", cudaGetErrorString(status)); + } + + auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(vertices_c.device()); + auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(vertices_c.device()); + auto opts_b = torch::TensorOptions().dtype(torch::kBool).device(vertices_c.device()); + + torch::Tensor voxel_coords = torch::empty({out.size, 3}, opts_i32); + torch::Tensor dual_vertices = torch::empty({out.size, 3}, opts_f32); + torch::Tensor intersected = torch::empty({out.size, 3}, opts_b); + + if (out.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + voxel_coords.data_ptr(), + out.voxel_coords, + static_cast(out.size) * 3 * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync voxel_coords" + ); + + check_cuda_success( + cudaMemcpyAsync( + dual_vertices.data_ptr(), + out.dual_vertices, + static_cast(out.size) * 3 * sizeof(float), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync dual_vertices" + ); + + check_cuda_success( + cudaMemcpyAsync( + intersected.data_ptr(), + out.intersected, + static_cast(out.size) * 3 * sizeof(bool), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersected" + ); + + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize"); + } + + fdg_gpu::free_flexible_dual_grid_gpu_output(&out); + return std::make_tuple(voxel_coords, dual_vertices, intersected); +} + + +std::tuple intersect_qef_cpu( + const torch::Tensor& triangles, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range +) { + check_cpu_tensor(triangles, "triangles"); + auto triangles_c = triangles.to(torch::kFloat32).contiguous(); + check_triangles_tensor(triangles_c); + + Eigen::Vector3f voxel_size_h = tensor_to_eigen_vec3_cpu(voxel_size); + Eigen::Vector3i grid_min, grid_max; + tensor_to_eigen_grid_min_max_cpu(grid_range, grid_min, grid_max); + + std::vector triangles_vec = triangles_tensor_to_vector_cpu(triangles_c); + std::unordered_map hash_table; + std::vector voxels_vec; + std::vector mean_sum; + std::vector cnt; + std::vector intersected_vec; + std::vector qefs; + + intersect_qef( + voxel_size_h, + grid_min, + grid_max, + triangles_vec, + hash_table, + voxels_vec, + mean_sum, + cnt, + intersected_vec, + qefs + ); + + return std::make_tuple( + int3_vector_to_tensor_cpu(voxels_vec), + vec3f_vector_to_tensor_cpu(mean_sum), + float_vector_to_tensor_cpu(cnt), + bool3_vector_to_tensor_cpu(intersected_vec), + matrix4f_vector_to_tensor_cpu(qefs) + ); +} + + +torch::Tensor face_qef_cpu( + const torch::Tensor& triangles, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + const torch::Tensor& voxels +) { + check_cpu_tensor(triangles, "triangles"); + check_cpu_tensor(voxels, "voxels"); + auto triangles_c = triangles.to(torch::kFloat32).contiguous(); + auto voxels_c = voxels.to(torch::kInt32).contiguous(); + check_triangles_tensor(triangles_c); + check_voxels_tensor(voxels_c); + + Eigen::Vector3f voxel_size_h = tensor_to_eigen_vec3_cpu(voxel_size); + Eigen::Vector3i grid_min, grid_max; + tensor_to_eigen_grid_min_max_cpu(grid_range, grid_min, grid_max); + + std::vector triangles_vec = triangles_tensor_to_vector_cpu(triangles_c); + std::unordered_map hash_table; + std::vector voxels_vec = voxels_tensor_to_vector_cpu(voxels_c, hash_table); + std::vector qefs(voxels_vec.size(), Eigen::Matrix4f::Zero()); + + face_qef( + voxel_size_h, + grid_min, + grid_max, + triangles_vec, + hash_table, + qefs + ); + + return matrix4f_vector_to_tensor_cpu(qefs); +} + + +torch::Tensor boundary_qef_cpu( + const torch::Tensor& boundaries, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + float boundary_weight, + const torch::Tensor& voxels +) { + check_cpu_tensor(boundaries, "boundaries"); + check_cpu_tensor(voxels, "voxels"); + auto boundaries_c = boundaries.to(torch::kFloat32).contiguous(); + auto voxels_c = voxels.to(torch::kInt32).contiguous(); + check_boundaries_tensor(boundaries_c); + check_voxels_tensor(voxels_c); + + Eigen::Vector3f voxel_size_h = tensor_to_eigen_vec3_cpu(voxel_size); + Eigen::Vector3i grid_min, grid_max; + tensor_to_eigen_grid_min_max_cpu(grid_range, grid_min, grid_max); + + std::vector boundaries_vec = boundaries_tensor_to_vector_cpu(boundaries_c); + std::unordered_map hash_table; + std::vector voxels_vec = voxels_tensor_to_vector_cpu(voxels_c, hash_table); + std::vector qefs(voxels_vec.size(), Eigen::Matrix4f::Zero()); + + boundry_qef( + voxel_size_h, + grid_min, + grid_max, + boundaries_vec, + boundary_weight, + hash_table, + qefs + ); + + return matrix4f_vector_to_tensor_cpu(qefs); +} + + +torch::Tensor intersection_occ_gpu( + const torch::Tensor& triangles, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + int64_t chunk_triangles +) { + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + TORCH_CHECK(chunk_triangles > 0, "chunk_triangles must be > 0"); + + auto triangles_c = triangles.to(torch::kFloat32).contiguous(); + check_triangles_tensor(triangles_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = intersection_qef::intersection_occ_gpu( + triangles_c.data_ptr(), + triangles_c.size(0), + voxel_size_h, + grid_min, + grid_max, + chunk_triangles, + stream + ); + + auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(triangles_c.device()); + torch::Tensor voxels = torch::empty({out.size, 3}, opts_i32); + if (out.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + voxels.data_ptr(), + out.voxels.data(), + static_cast(out.size) * 3 * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersection_occ voxels" + ); + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize intersection_occ"); + } + return voxels; +} + + +std::tuple intersect_qef_gpu( + const torch::Tensor& triangles, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + int64_t chunk_triangles +) { + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + TORCH_CHECK(chunk_triangles > 0, "chunk_triangles must be > 0"); + + auto triangles_c = triangles.to(torch::kFloat32).contiguous(); + check_triangles_tensor(triangles_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = intersection_qef::intersect_qef_gpu( + triangles_c.data_ptr(), + triangles_c.size(0), + voxel_size_h, + grid_min, + grid_max, + chunk_triangles, + stream + ); + + auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(triangles_c.device()); + auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(triangles_c.device()); + auto opts_u8 = torch::TensorOptions().dtype(torch::kUInt8).device(triangles_c.device()); + auto opts_b = torch::TensorOptions().dtype(torch::kBool).device(triangles_c.device()); + + static_assert(sizeof(fdg_gpu::SymQEF10) == sizeof(float) * 10, "Unexpected SymQEF10 layout"); + + torch::Tensor voxels = torch::empty({out.size, 3}, opts_i32); + torch::Tensor mean_sum = torch::empty({out.size, 3}, opts_f32); + torch::Tensor cnt = torch::empty({out.size}, opts_f32); + torch::Tensor intersected_mask = torch::empty({out.size}, opts_u8); + torch::Tensor qefs = torch::empty({out.size, 10}, opts_f32); + + if (out.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + voxels.data_ptr(), + out.voxels.data(), + static_cast(out.size) * 3 * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersect_qef voxels" + ); + check_cuda_success( + cudaMemcpyAsync( + mean_sum.data_ptr(), + out.mean_sum.data(), + static_cast(out.size) * 3 * sizeof(float), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersect_qef mean_sum" + ); + check_cuda_success( + cudaMemcpyAsync( + cnt.data_ptr(), + out.cnt.data(), + static_cast(out.size) * sizeof(float), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersect_qef cnt" + ); + check_cuda_success( + cudaMemcpyAsync( + intersected_mask.data_ptr(), + out.intersected.data(), + static_cast(out.size) * sizeof(uint8_t), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersect_qef intersected" + ); + check_cuda_success( + cudaMemcpyAsync( + qefs.data_ptr(), + out.qefs.data(), + static_cast(out.size) * sizeof(fdg_gpu::SymQEF10), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync intersect_qef qefs" + ); + } + + torch::Tensor intersected = torch::empty({out.size, 3}, opts_b); + if (out.size > 0) { + const int kBlock = 256; + const int grid = static_cast((out.size + kBlock - 1) / kBlock); + unpack_intersected_mask_kernel<<>>( + intersected_mask.data_ptr(), + out.size, + intersected.data_ptr() + ); + check_cuda_success(cudaGetLastError(), "unpack_intersected_mask_kernel"); + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize intersect_qef"); + } + + return std::make_tuple(voxels, mean_sum, cnt, intersected, qefs); +} + + +std::tuple voxelize_mesh_oct_gpu( + const torch::Tensor& vertices, + const torch::Tensor& faces, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range +) { + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(faces.is_cuda(), "faces must be a CUDA tensor"); + TORCH_CHECK(vertices.device() == faces.device(), "vertices and faces must be on the same CUDA device"); + + auto vertices_c = vertices.to(torch::kFloat32).contiguous(); + auto faces_c = faces.to(torch::kInt32).contiguous(); + TORCH_CHECK(vertices_c.dim() == 2 && vertices_c.size(1) == 3, "vertices must have shape [V, 3]"); + TORCH_CHECK(faces_c.dim() == 2 && faces_c.size(1) == 3, "faces must have shape [F, 3]"); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + fdg_gpu::int3_ grid_size = grid_size_from_min_max(grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = oct_pairs::voxelize_mesh_oct_gpu( + vertices_c.data_ptr(), + vertices_c.size(0), + faces_c.data_ptr(), + faces_c.size(0), + grid_min, + grid_size, + voxel_size_h, + stream + ); + + return primitive_pair_to_tensors(out, vertices_c.device(), stream); +} + + +std::tuple voxelize_edge_oct_gpu( + const torch::Tensor& vertices, + const torch::Tensor& edges, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range +) { + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(edges.is_cuda(), "edges must be a CUDA tensor"); + TORCH_CHECK(vertices.device() == edges.device(), "vertices and edges must be on the same CUDA device"); + + auto vertices_c = vertices.to(torch::kFloat32).contiguous(); + auto edges_c = edges.to(torch::kInt32).contiguous(); + TORCH_CHECK(vertices_c.dim() == 2 && vertices_c.size(1) == 3, "vertices must have shape [V, 3]"); + check_edges_tensor(edges_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + fdg_gpu::int3_ grid_size = grid_size_from_min_max(grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = oct_pairs::voxelize_edge_oct_gpu( + vertices_c.data_ptr(), + vertices_c.size(0), + edges_c.data_ptr(), + edges_c.size(0), + grid_min, + grid_size, + voxel_size_h, + stream + ); + + return primitive_pair_to_tensors(out, vertices_c.device(), stream); +} + + +torch::Tensor face_qef_gpu( + const torch::Tensor& triangles, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + const torch::Tensor& voxels +) { + TORCH_CHECK(triangles.is_cuda(), "triangles must be a CUDA tensor"); + TORCH_CHECK(voxels.is_cuda(), "voxels must be a CUDA tensor"); + TORCH_CHECK(triangles.device() == voxels.device(), "triangles and voxels must be on the same CUDA device"); + + auto triangles_c = triangles.to(torch::kFloat32).contiguous(); + auto voxels_c = voxels.to(torch::kInt32).contiguous(); + check_triangles_tensor(triangles_c); + check_voxels_tensor(voxels_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = oct_pairs::face_qef_gpu( + voxel_size_h, + grid_min, + grid_max, + triangles_c.data_ptr(), + triangles_c.size(0), + voxels_c.data_ptr(), + voxels_c.size(0), + stream + ); + + static_assert(sizeof(fdg_gpu::SymQEF10) == sizeof(float) * 10, "Unexpected SymQEF10 layout"); + auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(triangles_c.device()); + torch::Tensor qefs = torch::empty({out.size, 10}, opts_f32); + if (out.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + qefs.data_ptr(), + out.qefs.data(), + static_cast(out.size) * sizeof(fdg_gpu::SymQEF10), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync face_qef qefs" + ); + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize face_qef"); + } + return qefs; +} + + +std::tuple voxel_traverse_edge_dda_gpu( + const torch::Tensor& vertices, + const torch::Tensor& edges, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + int chunk_steps +) { + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(edges.is_cuda(), "edges must be a CUDA tensor"); + TORCH_CHECK(vertices.device() == edges.device(), "vertices and edges must be on the same CUDA device"); + TORCH_CHECK(chunk_steps > 0, "chunk_steps must be > 0"); + + auto vertices_c = vertices.to(torch::kFloat32).contiguous(); + auto edges_c = edges.to(torch::kInt32).contiguous(); + TORCH_CHECK(vertices_c.dim() == 2 && vertices_c.size(1) == 3, "vertices must have shape [V, 3]"); + check_edges_tensor(edges_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = edge_dda::voxel_traverse_edge_dda_gpu( + vertices_c.data_ptr(), + vertices_c.size(0), + edges_c.data_ptr(), + edges_c.size(0), + voxel_size_h, + grid_min, + grid_max, + chunk_steps, + stream + ); + + return primitive_pair_to_tensors(out, vertices_c.device(), stream); +} + + +torch::Tensor boundary_qef_gpu( + const torch::Tensor& boundaries, + const torch::Tensor& voxel_size, + const torch::Tensor& grid_range, + float boundary_weight, + const torch::Tensor& voxels, + int chunk_steps +) { + TORCH_CHECK(boundaries.is_cuda(), "boundaries must be a CUDA tensor"); + TORCH_CHECK(voxels.is_cuda(), "voxels must be a CUDA tensor"); + TORCH_CHECK(boundaries.device() == voxels.device(), "boundaries and voxels must be on the same CUDA device"); + TORCH_CHECK(chunk_steps > 0, "chunk_steps must be > 0"); + + auto boundaries_c = boundaries.to(torch::kFloat32).contiguous(); + auto voxels_c = voxels.to(torch::kInt32).contiguous(); + check_boundaries_tensor(boundaries_c); + check_voxels_tensor(voxels_c); + + float3 voxel_size_h = tensor_to_float3_cpu(voxel_size); + fdg_gpu::int3_ grid_min{}; + fdg_gpu::int3_ grid_max{}; + tensor_to_grid_min_max_cpu(grid_range, grid_min, grid_max); + + cudaStream_t stream = nullptr; + auto out = edge_dda::boundary_qef_gpu( + voxel_size_h, + grid_min, + grid_max, + boundaries_c.data_ptr(), + boundaries_c.size(0), + boundary_weight, + voxels_c.data_ptr(), + voxels_c.size(0), + chunk_steps, + stream + ); + + static_assert(sizeof(fdg_gpu::SymQEF10) == sizeof(float) * 10, "Unexpected SymQEF10 layout"); + auto opts_f32 = torch::TensorOptions().dtype(torch::kFloat32).device(boundaries_c.device()); + torch::Tensor qefs = torch::empty({out.size, 10}, opts_f32); + if (out.size > 0) { + check_cuda_success( + cudaMemcpyAsync( + qefs.data_ptr(), + out.qefs.data(), + static_cast(out.size) * sizeof(fdg_gpu::SymQEF10), + cudaMemcpyDeviceToDevice, + stream + ), + "cudaMemcpyAsync boundary_qef qefs" + ); + check_cuda_success(cudaStreamSynchronize(stream), "cudaStreamSynchronize boundary_qef"); + } + + return qefs; +} diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/types.cuh b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/types.cuh new file mode 100644 index 00000000..27717e67 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/types.cuh @@ -0,0 +1,201 @@ +#pragma once + +#include + +#include + +namespace o_voxel::fdg +{ + + struct Int3 + { + int x; + int y; + int z; + + __host__ __device__ int &operator[](int i) { return (&x)[i]; } + __host__ __device__ int operator[](int i) const { return (&x)[i]; } + }; + + struct GridSpec + { + // Grid-local coordinates use voxel indices in [grid_min, grid_max). + float3 voxel_size; + Int3 grid_min; + Int3 grid_max; + }; + + // Upper-triangle storage for a symmetric 4x4 QEF matrix: + // q00 q01 q02 q03 + // q11 q12 q13 + // q22 q23 + // q33 + struct SymQEF10 + { + float q00; + float q01; + float q02; + float q03; + float q11; + float q12; + float q13; + float q22; + float q23; + float q33; + }; + + // Active voxels are grouped into 8x8x8 bricks. Each brick stores 512 + // occupancy bits as 16 uint32 words, and active bricks are addressed by an + // open-addressing hash table. + inline constexpr int kBrickSize = 8; + inline constexpr int kBrickLocalCells = kBrickSize * kBrickSize * kBrickSize; + inline constexpr int kBrickBitWords = kBrickLocalCells / 32; + inline constexpr uint64_t kEmptyBrickKey = UINT64_MAX; + inline constexpr uint32_t kEmptyBrickVal = UINT32_MAX; + inline constexpr uint32_t kOverflowBrickVal = UINT32_MAX - 1u; + + // Lookup data returned by intersect_qef_cuda. hash_keys/hash_vals map a + // brick key to a compact brick index, brick_bits stores local occupancy, and + // brick_base gives the first row for that brick in compacted voxel/QEF arrays. + struct BrickLookup + { + const uint64_t *hash_keys; + const uint32_t *hash_vals; + const uint32_t *brick_bits; + const int64_t *brick_base; + uint64_t hash_capacity; + }; + + __host__ __device__ __forceinline__ uint64_t pack_voxel_key( + int x, + int y, + int z, + Int3 grid_min, + Int3 grid_max) + { + // Dense row-major key for a voxel coordinate inside the current grid. + const uint64_t sx = static_cast(grid_max.x - grid_min.x); + const uint64_t sy = static_cast(grid_max.y - grid_min.y); + const uint64_t ux = static_cast(x - grid_min.x); + const uint64_t uy = static_cast(y - grid_min.y); + const uint64_t uz = static_cast(z - grid_min.z); + return ux + sx * (uy + sy * uz); + } + + __host__ __device__ __forceinline__ Int3 unpack_voxel_key( + uint64_t key, + Int3 grid_min, + Int3 grid_max) + { + // Inverse of pack_voxel_key, used when compact voxel rows are emitted. + const uint64_t sx = static_cast(grid_max.x - grid_min.x); + const uint64_t sy = static_cast(grid_max.y - grid_min.y); + const uint64_t yz = sx * sy; + const uint64_t z = key / yz; + const uint64_t rem = key - z * yz; + const uint64_t y = rem / sx; + const uint64_t x = rem - y * sx; + return Int3{ + static_cast(x) + grid_min.x, + static_cast(y) + grid_min.y, + static_cast(z) + grid_min.z, + }; + } + + __device__ __forceinline__ bool lookup_brick_bits_and_base( + int bx, + int by, + int bz, + GridSpec grid, + BrickLookup lookup, + const uint32_t **bits, + int64_t *base) + { + // Find one active brick and return both its occupancy bitset and its + // first compact voxel row. Callers can then test local bits cheaply. + if (lookup.hash_capacity == 0) + return false; + + const uint64_t nbx = static_cast( + (grid.grid_max.x - grid.grid_min.x + kBrickSize - 1) / kBrickSize); + const uint64_t nby = static_cast( + (grid.grid_max.y - grid.grid_min.y + kBrickSize - 1) / kBrickSize); + const uint64_t key = + static_cast(bx) + nbx * (static_cast(by) + nby * static_cast(bz)); + + uint64_t slot_key = key; + slot_key ^= slot_key >> 33; + slot_key *= 0xff51afd7ed558ccdULL; + slot_key ^= slot_key >> 33; + slot_key *= 0xc4ceb9fe1a85ec53ULL; + slot_key ^= slot_key >> 33; + + uint64_t slot = slot_key & (lookup.hash_capacity - 1); + for (uint64_t probe = 0; probe < lookup.hash_capacity; ++probe) + { + const uint64_t found = lookup.hash_keys[slot]; + if (found == kEmptyBrickKey) + return false; + if (found == key) + { + const uint32_t brick_idx = lookup.hash_vals[slot]; + if (brick_idx == kEmptyBrickVal || brick_idx == kOverflowBrickVal) + return false; + *bits = lookup.brick_bits + static_cast(brick_idx) * kBrickBitWords; + *base = lookup.brick_base[brick_idx]; + return true; + } + slot = (slot + 1u) & (lookup.hash_capacity - 1); + } + return false; + } + + __device__ __forceinline__ int64_t lookup_voxel_row_in_bricks( + int x, + int y, + int z, + GridSpec grid, + BrickLookup lookup) + { + // Returns the row in compacted voxels/qefs, or -1 when the voxel is not active. + if (lookup.hash_capacity == 0) + return -1; + if (x < grid.grid_min.x || x >= grid.grid_max.x) + return -1; + if (y < grid.grid_min.y || y >= grid.grid_max.y) + return -1; + if (z < grid.grid_min.z || z >= grid.grid_max.z) + return -1; + + const int rx = x - grid.grid_min.x; + const int ry = y - grid.grid_min.y; + const int rz = z - grid.grid_min.z; + const int bx = rx / kBrickSize; + const int by = ry / kBrickSize; + const int bz = rz / kBrickSize; + const int lx = rx - bx * kBrickSize; + const int ly = ry - by * kBrickSize; + const int lz = rz - bz * kBrickSize; + const int local_id = lx + kBrickSize * (ly + kBrickSize * lz); + + const uint32_t *bits; + int64_t base; + if (!lookup_brick_bits_and_base(bx, by, bz, grid, lookup, &bits, &base)) + return -1; + const int word = local_id / 32; + const int bit = local_id - word * 32; + if ((bits[word] & (1u << bit)) == 0) + return -1; + + // Rank is the number of active local voxels before local_id. Since + // emit_occupied_voxels_kernel writes voxels in the same bit order, this + // reproduces the compact row without a per-voxel hash table. + int rank = 0; + for (int i = 0; i < word; ++i) + rank += __popc(bits[i]); + const uint32_t mask = bit == 0 ? 0u : ((1u << bit) - 1u); + rank += __popc(bits[word] & mask); + return base + rank; + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.cu new file mode 100644 index 00000000..dbc0d88c --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.cu @@ -0,0 +1,1029 @@ +#include "voxel_traverse_edge_dda.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include + + +#include +#include +#include +#include +#include +#include + +namespace voxel_traverse_edge_dda_impl { + +#define VOX_CUDA_CHECK(expr) \ + do { \ + cudaError_t _err = (expr); \ + if (_err != cudaSuccess) { \ + throw std::runtime_error(std::string("CUDA error: ") + \ + cudaGetErrorString(_err) + \ + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +constexpr int kDefaultBlockSize = 128; + +struct EdgeDesc { + float3 v0_ws; + float3 v1_ws; + + double3 dir_unit; + double segment_length; + + int32_t start_x; + int32_t start_y; + int32_t start_z; + + int8_t step_x; + int8_t step_y; + int8_t step_z; + + double tmax0_x; + double tmax0_y; + double tmax0_z; + + double tdelta_x; + double tdelta_y; + double tdelta_z; +}; + +struct DDAJobQueue { + int32_t* edge_id = nullptr; + + int32_t* cur_x = nullptr; + int32_t* cur_y = nullptr; + int32_t* cur_z = nullptr; + + double* tmax_x = nullptr; + double* tmax_y = nullptr; + double* tmax_z = nullptr; + + int64_t size = 0; + int64_t capacity = 0; +}; + +struct RoundBuffers { + int32_t* pair_count = nullptr; + int32_t* next_job_count = nullptr; + + int32_t* pair_offset = nullptr; + int32_t* next_job_offset = nullptr; + + void* cub_temp_storage = nullptr; + size_t cub_temp_bytes = 0; + int64_t capacity = 0; +}; + +struct ResultBuffer { + int32_t* edge_id = nullptr; + int32_t* vi = nullptr; + int32_t* vj = nullptr; + int32_t* vk = nullptr; + int64_t size = 0; +}; + +struct DeviceResult { + int32_t* edge_id = nullptr; + int32_t* voxel_i = nullptr; + int32_t* voxel_j = nullptr; + int32_t* voxel_k = nullptr; + int64_t size = 0; +}; + +struct Workspace { + EdgeDesc* edge_desc = nullptr; + + uint8_t* edge_valid = nullptr; + int32_t* init_count = nullptr; + int32_t* init_offset = nullptr; + + DDAJobQueue queue_a; + DDAJobQueue queue_b; + + RoundBuffers round; + std::vector result_rounds; +}; + +inline int ceil_div_i64(int64_t n, int block) { + return static_cast((n + block - 1) / block); +} + +inline void free_ptr(void* ptr) { + if (ptr != nullptr) { + cudaFree(ptr); + } +} + +inline void alloc_i32(int32_t** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(int32_t) * n)); + } +} + +inline void alloc_u8(uint8_t** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(uint8_t) * n)); + } +} + +inline void alloc_double(double** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(double) * n)); + } +} + +inline void alloc_edge_desc(EdgeDesc** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(EdgeDesc) * n)); + } +} + +inline void release_dda_job_queue(DDAJobQueue& q) { + free_ptr(q.edge_id); + free_ptr(q.cur_x); + free_ptr(q.cur_y); + free_ptr(q.cur_z); + free_ptr(q.tmax_x); + free_ptr(q.tmax_y); + free_ptr(q.tmax_z); + q = {}; +} + +inline void release_round_buffers(RoundBuffers& b) { + free_ptr(b.pair_count); + free_ptr(b.next_job_count); + free_ptr(b.pair_offset); + free_ptr(b.next_job_offset); + free_ptr(b.cub_temp_storage); + b = {}; +} + +inline void release_result_buffer(ResultBuffer& r) { + free_ptr(r.edge_id); + free_ptr(r.vi); + free_ptr(r.vj); + free_ptr(r.vk); + r = {}; +} + +inline void release_workspace(Workspace& ws) { + free_ptr(ws.edge_desc); + free_ptr(ws.edge_valid); + free_ptr(ws.init_count); + free_ptr(ws.init_offset); + release_dda_job_queue(ws.queue_a); + release_dda_job_queue(ws.queue_b); + release_round_buffers(ws.round); + for (auto& r : ws.result_rounds) { + release_result_buffer(r); + } + ws.result_rounds.clear(); +} + +inline void ensure_dda_job_queue_capacity(DDAJobQueue& q, int64_t capacity) { + if (capacity <= q.capacity) { + return; + } + release_dda_job_queue(q); + alloc_i32(&q.edge_id, capacity); + alloc_i32(&q.cur_x, capacity); + alloc_i32(&q.cur_y, capacity); + alloc_i32(&q.cur_z, capacity); + alloc_double(&q.tmax_x, capacity); + alloc_double(&q.tmax_y, capacity); + alloc_double(&q.tmax_z, capacity); + q.capacity = capacity; + q.size = 0; +} + +inline void ensure_round_capacity(RoundBuffers& b, int64_t capacity) { + if (capacity <= b.capacity) { + return; + } + free_ptr(b.pair_count); + free_ptr(b.next_job_count); + free_ptr(b.pair_offset); + free_ptr(b.next_job_offset); + + alloc_i32(&b.pair_count, capacity); + alloc_i32(&b.next_job_count, capacity); + alloc_i32(&b.pair_offset, capacity); + alloc_i32(&b.next_job_offset, capacity); + b.capacity = capacity; +} + +inline void ensure_scan_temp_storage( + RoundBuffers& b, + int32_t* d_in, + int32_t* d_out, + int64_t count, + cudaStream_t stream) { + if (count <= 0) { + return; + } + if (count > INT32_MAX) { + throw std::runtime_error("CUB scan count exceeds int32 range"); + } + size_t bytes = 0; + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + nullptr, + bytes, + d_in, + d_out, + static_cast(count), + stream)); + if (bytes > b.cub_temp_bytes) { + free_ptr(b.cub_temp_storage); + VOX_CUDA_CHECK(cudaMalloc(&b.cub_temp_storage, bytes)); + b.cub_temp_bytes = bytes; + } +} + +inline void exclusive_scan_i32( + RoundBuffers& b, + int32_t* d_in, + int32_t* d_out, + int64_t count, + cudaStream_t stream) { + if (count <= 0) { + return; + } + ensure_scan_temp_storage(b, d_in, d_out, count, stream); + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + b.cub_temp_storage, + b.cub_temp_bytes, + d_in, + d_out, + static_cast(count), + stream)); +} + +inline int32_t copy_last_i32(const int32_t* ptr, int64_t count, cudaStream_t stream) { + if (count <= 0) { + return 0; + } + int32_t value = 0; + VOX_CUDA_CHECK(cudaMemcpyAsync( + &value, + ptr + (count - 1), + sizeof(int32_t), + cudaMemcpyDeviceToHost, + stream)); + VOX_CUDA_CHECK(cudaStreamSynchronize(stream)); + return value; +} + +inline ResultBuffer make_result_buffer(int64_t count) { + ResultBuffer r; + if (count <= 0) { + return r; + } + alloc_i32(&r.edge_id, count); + alloc_i32(&r.vi, count); + alloc_i32(&r.vj, count); + alloc_i32(&r.vk, count); + r.size = count; + return r; +} + +inline DeviceResult gather_result_rounds( + const std::vector& rounds, + cudaStream_t stream) { + DeviceResult out; + int64_t total = 0; + for (const auto& r : rounds) { + total += r.size; + } + out.size = total; + if (total == 0) { + return out; + } + + alloc_i32(&out.edge_id, total); + alloc_i32(&out.voxel_i, total); + alloc_i32(&out.voxel_j, total); + alloc_i32(&out.voxel_k, total); + + int64_t cursor = 0; + for (const auto& r : rounds) { + if (r.size == 0) { + continue; + } + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.edge_id + cursor, + r.edge_id, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_i + cursor, + r.vi, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_j + cursor, + r.vj, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_k + cursor, + r.vk, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + cursor += r.size; + } + + VOX_CUDA_CHECK(cudaStreamSynchronize(stream)); + return out; +} + +__device__ inline int argmin_axis_strict(double tx, double ty, double tz) { + if (tx < ty) { + return (tx < tz) ? 0 : 2; + } + return (ty < tz) ? 1 : 2; +} + +__device__ inline bool in_bounds_voxel_abs( + int x, + int y, + int z, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max) { + return (grid_min.x <= x && x < grid_max.x) && + (grid_min.y <= y && y < grid_max.y) && + (grid_min.z <= z && z < grid_max.z); +} + +__global__ void kernel_build_edge_desc( + const float* __restrict__ vertices, + const int32_t* __restrict__ edges, + int64_t num_edges, + float3 voxel_size, + EdgeDesc* __restrict__ edge_desc, + uint8_t* __restrict__ edge_valid) { + int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_edges) { + return; + } + + int v0_id = edges[2 * eid + 0]; + int v1_id = edges[2 * eid + 1]; + + float3 v0 = make_float3(vertices[3 * v0_id + 0], vertices[3 * v0_id + 1], vertices[3 * v0_id + 2]); + float3 v1 = make_float3(vertices[3 * v1_id + 0], vertices[3 * v1_id + 1], vertices[3 * v1_id + 2]); + + double dx = static_cast(v1.x) - static_cast(v0.x); + double dy = static_cast(v1.y) - static_cast(v0.y); + double dz = static_cast(v1.z) - static_cast(v0.z); + double segment_length = sqrt(dx * dx + dy * dy + dz * dz); + + if (segment_length < 1e-6) { + edge_valid[eid] = 0; + return; + } + + double3 dir_unit = make_double3(dx / segment_length, dy / segment_length, dz / segment_length); + + int32_t sx = static_cast(floor(static_cast(v0.x) / static_cast(voxel_size.x))); + int32_t sy = static_cast(floor(static_cast(v0.y) / static_cast(voxel_size.y))); + int32_t sz = static_cast(floor(static_cast(v0.z) / static_cast(voxel_size.z))); + + int8_t step_x = (dir_unit.x > 0.0) ? 1 : -1; + int8_t step_y = (dir_unit.y > 0.0) ? 1 : -1; + int8_t step_z = (dir_unit.z > 0.0) ? 1 : -1; + + double tmax_x, tmax_y, tmax_z; + double tdelta_x, tdelta_y, tdelta_z; + + if (dir_unit.x == 0.0) { + tmax_x = CUDART_INF; + tdelta_x = CUDART_INF; + } else { + double voxel_border = static_cast(voxel_size.x) * static_cast(sx + (step_x > 0 ? 1 : 0)); + tmax_x = (voxel_border - static_cast(v0.x)) / dir_unit.x; + tdelta_x = static_cast(voxel_size.x) / fabs(dir_unit.x); + } + + if (dir_unit.y == 0.0) { + tmax_y = CUDART_INF; + tdelta_y = CUDART_INF; + } else { + double voxel_border = static_cast(voxel_size.y) * static_cast(sy + (step_y > 0 ? 1 : 0)); + tmax_y = (voxel_border - static_cast(v0.y)) / dir_unit.y; + tdelta_y = static_cast(voxel_size.y) / fabs(dir_unit.y); + } + + if (dir_unit.z == 0.0) { + tmax_z = CUDART_INF; + tdelta_z = CUDART_INF; + } else { + double voxel_border = static_cast(voxel_size.z) * static_cast(sz + (step_z > 0 ? 1 : 0)); + tmax_z = (voxel_border - static_cast(v0.z)) / dir_unit.z; + tdelta_z = static_cast(voxel_size.z) / fabs(dir_unit.z); + } + + EdgeDesc desc; + desc.v0_ws = v0; + desc.v1_ws = v1; + desc.dir_unit = dir_unit; + desc.segment_length = segment_length; + desc.start_x = sx; + desc.start_y = sy; + desc.start_z = sz; + desc.step_x = step_x; + desc.step_y = step_y; + desc.step_z = step_z; + desc.tmax0_x = tmax_x; + desc.tmax0_y = tmax_y; + desc.tmax0_z = tmax_z; + desc.tdelta_x = tdelta_x; + desc.tdelta_y = tdelta_y; + desc.tdelta_z = tdelta_z; + + edge_desc[eid] = desc; + edge_valid[eid] = 1; +} + +__global__ void kernel_count_init_jobs( + const uint8_t* __restrict__ edge_valid, + int64_t num_edges, + int32_t* __restrict__ init_count) { + int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_edges) { + return; + } + init_count[eid] = edge_valid[eid] ? 1 : 0; +} + +__global__ void kernel_emit_init_jobs( + const uint8_t* __restrict__ edge_valid, + const EdgeDesc* __restrict__ edge_desc, + const int32_t* __restrict__ init_offset, + int64_t num_edges, + DDAJobQueue out_q) { + int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_edges || !edge_valid[eid]) { + return; + } + + int32_t out = init_offset[eid]; + const EdgeDesc& desc = edge_desc[eid]; + out_q.edge_id[out] = static_cast(eid); + out_q.cur_x[out] = desc.start_x; + out_q.cur_y[out] = desc.start_y; + out_q.cur_z[out] = desc.start_z; + out_q.tmax_x[out] = desc.tmax0_x; + out_q.tmax_y[out] = desc.tmax0_y; + out_q.tmax_z[out] = desc.tmax0_z; +} + +__global__ void kernel_count_dda_jobs( + DDAJobQueue curr_q, + const EdgeDesc* __restrict__ edge_desc, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int chunk_steps, + int32_t* __restrict__ pair_count, + int32_t* __restrict__ next_job_count) { + int64_t jid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (jid >= curr_q.size) { + return; + } + + int32_t eid = curr_q.edge_id[jid]; + const EdgeDesc& desc = edge_desc[eid]; + + int32_t cx = curr_q.cur_x[jid]; + int32_t cy = curr_q.cur_y[jid]; + int32_t cz = curr_q.cur_z[jid]; + double tx = curr_q.tmax_x[jid]; + double ty = curr_q.tmax_y[jid]; + double tz = curr_q.tmax_z[jid]; + + int32_t local_pairs = 0; + bool alive = true; + + if (in_bounds_voxel_abs(cx, cy, cz, grid_min, grid_max)) { + local_pairs += 1; + } + + for (int step_idx = 0; step_idx < chunk_steps; ++step_idx) { + int axis = argmin_axis_strict(tx, ty, tz); + double t_axis = (axis == 0) ? tx : (axis == 1 ? ty : tz); + if (t_axis > desc.segment_length) { + alive = false; + break; + } + + if (axis == 0) { + cx += static_cast(desc.step_x); + tx += desc.tdelta_x; + } else if (axis == 1) { + cy += static_cast(desc.step_y); + ty += desc.tdelta_y; + } else { + cz += static_cast(desc.step_z); + tz += desc.tdelta_z; + } + + if (in_bounds_voxel_abs(cx, cy, cz, grid_min, grid_max)) { + local_pairs += 1; + } + } + + pair_count[jid] = local_pairs; + next_job_count[jid] = alive ? 1 : 0; +} + +__global__ void kernel_emit_dda_jobs( + DDAJobQueue curr_q, + const EdgeDesc* __restrict__ edge_desc, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int chunk_steps, + const int32_t* __restrict__ pair_offset, + const int32_t* __restrict__ next_job_offset, + ResultBuffer out_res, + DDAJobQueue next_q) { + int64_t jid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (jid >= curr_q.size) { + return; + } + + int32_t eid = curr_q.edge_id[jid]; + const EdgeDesc& desc = edge_desc[eid]; + + int32_t cx = curr_q.cur_x[jid]; + int32_t cy = curr_q.cur_y[jid]; + int32_t cz = curr_q.cur_z[jid]; + double tx = curr_q.tmax_x[jid]; + double ty = curr_q.tmax_y[jid]; + double tz = curr_q.tmax_z[jid]; + + int32_t out_pair = pair_offset[jid]; + bool alive = true; + + if (in_bounds_voxel_abs(cx, cy, cz, grid_min, grid_max)) { + out_res.edge_id[out_pair] = eid; + out_res.vi[out_pair] = cx; + out_res.vj[out_pair] = cy; + out_res.vk[out_pair] = cz; + out_pair += 1; + } + + for (int step_idx = 0; step_idx < chunk_steps; ++step_idx) { + int axis = argmin_axis_strict(tx, ty, tz); + double t_axis = (axis == 0) ? tx : (axis == 1 ? ty : tz); + if (t_axis > desc.segment_length) { + alive = false; + break; + } + + if (axis == 0) { + cx += static_cast(desc.step_x); + tx += desc.tdelta_x; + } else if (axis == 1) { + cy += static_cast(desc.step_y); + ty += desc.tdelta_y; + } else { + cz += static_cast(desc.step_z); + tz += desc.tdelta_z; + } + + if (in_bounds_voxel_abs(cx, cy, cz, grid_min, grid_max)) { + out_res.edge_id[out_pair] = eid; + out_res.vi[out_pair] = cx; + out_res.vj[out_pair] = cy; + out_res.vk[out_pair] = cz; + out_pair += 1; + } + } + + if (alive) { + int32_t out_job = next_job_offset[jid]; + next_q.edge_id[out_job] = eid; + next_q.cur_x[out_job] = cx; + next_q.cur_y[out_job] = cy; + next_q.cur_z[out_job] = cz; + next_q.tmax_x[out_job] = tx; + next_q.tmax_y[out_job] = ty; + next_q.tmax_z[out_job] = tz; + } +} + +inline void release_device_result(DeviceResult& out) { + free_ptr(out.edge_id); + free_ptr(out.voxel_i); + free_ptr(out.voxel_j); + free_ptr(out.voxel_k); + out = {}; +} + +} // namespace voxel_traverse_edge_dda_impl + +namespace { + +inline fdg_gpu::PrimitivePairResult to_primitive_pair(voxel_traverse_edge_dda_impl::DeviceResult&& r) { + fdg_gpu::PrimitivePairResult out; + out.size = r.size; + out.prim_id.adopt(r.edge_id, r.size); + out.voxel_i.adopt(r.voxel_i, r.size); + out.voxel_j.adopt(r.voxel_j, r.size); + out.voxel_k.adopt(r.voxel_k, r.size); + r.edge_id = nullptr; + r.voxel_i = nullptr; + r.voxel_j = nullptr; + r.voxel_k = nullptr; + r.size = 0; + return out; +} + +__host__ __device__ inline fdg_gpu::SymQEF10 symqef10_zero() { + return fdg_gpu::SymQEF10{0,0,0,0,0,0,0,0,0,0}; +} + +struct SymQEF10Add { + __host__ __device__ fdg_gpu::SymQEF10 operator()(const fdg_gpu::SymQEF10& a, const fdg_gpu::SymQEF10& b) const { + return fdg_gpu::SymQEF10{ + a.q00 + b.q00, a.q01 + b.q01, a.q02 + b.q02, a.q03 + b.q03, + a.q11 + b.q11, a.q12 + b.q12, a.q13 + b.q13, + a.q22 + b.q22, a.q23 + b.q23, + a.q33 + b.q33}; + } +}; + +struct SurfaceLookup { + int64_t size = 0; + fdg_gpu::DeviceBuffer keys_sorted; + fdg_gpu::DeviceBuffer ids_sorted; +}; + +struct EdgePairKeys { + int64_t size = 0; + fdg_gpu::DeviceBuffer pair_keys; +}; + +struct BoundaryContribStream { + int64_t size = 0; + fdg_gpu::DeviceBuffer voxel_id; + fdg_gpu::DeviceBuffer qef; +}; + + +__global__ void copy_boundaries_to_vertices_kernel(const float* boundaries, int64_t num_boundaries, float* vertices_out) { + int64_t tid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (tid >= 2 * num_boundaries) return; + vertices_out[3 * tid + 0] = boundaries[3 * tid + 0]; + vertices_out[3 * tid + 1] = boundaries[3 * tid + 1]; + vertices_out[3 * tid + 2] = boundaries[3 * tid + 2]; +} + +__global__ void build_synth_edges_kernel(int64_t num_boundaries, int32_t* edges_out) { + int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_boundaries) return; + edges_out[2 * eid + 0] = static_cast(2 * eid + 0); + edges_out[2 * eid + 1] = static_cast(2 * eid + 1); +} + +__global__ void build_surface_keys_kernel(const int* voxels, int64_t num_voxels, fdg_gpu::int3_ grid_min, fdg_gpu::int3_ grid_max, uint64_t* keys, int32_t* ids) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_voxels) return; + int x = voxels[3 * i + 0]; + int y = voxels[3 * i + 1]; + int z = voxels[3 * i + 2]; + keys[i] = fdg_gpu::pack_voxel_key(x, y, z, grid_min, grid_max); + ids[i] = static_cast(i); +} + +__global__ void build_raw_pair_voxel_keys_kernel(const int32_t* voxel_i, const int32_t* voxel_j, const int32_t* voxel_k, int64_t num_pairs, fdg_gpu::int3_ grid_min, fdg_gpu::int3_ grid_max, uint64_t* pair_voxel_keys) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + pair_voxel_keys[i] = fdg_gpu::pack_voxel_key(voxel_i[i], voxel_j[i], voxel_k[i], grid_min, grid_max); +} + +__device__ inline int lower_bound_u64(const uint64_t* arr, int64_t n, uint64_t key) { + int64_t lo = 0; + int64_t hi = n; + while (lo < hi) { + int64_t mid = (lo + hi) >> 1; + uint64_t v = arr[mid]; + if (v < key) lo = mid + 1; + else hi = mid; + } + return static_cast(lo); +} + +__global__ void map_pair_to_voxel_id_kernel(const uint64_t* pair_voxel_keys, const int32_t* edge_id, int64_t num_pairs, const uint64_t* surface_keys_sorted, const int32_t* surface_ids_sorted, int64_t num_voxels, int32_t* mapped_voxel_id, int32_t* mapped_edge_id, int32_t* valid) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + uint64_t key = pair_voxel_keys[i]; + int pos = lower_bound_u64(surface_keys_sorted, num_voxels, key); + if (pos < num_voxels && surface_keys_sorted[pos] == key) { + mapped_voxel_id[i] = surface_ids_sorted[pos]; + mapped_edge_id[i] = edge_id[i]; + valid[i] = 1; + } else { + mapped_voxel_id[i] = -1; + mapped_edge_id[i] = -1; + valid[i] = 0; + } +} + +__host__ __device__ inline uint64_t pack_edge_voxel_pair_key(int32_t edge_id, uint64_t voxel_key) { + return (static_cast(static_cast(edge_id)) << 32) ^ voxel_key; +} + +__global__ void compact_valid_pairs_kernel(const int32_t* mapped_voxel_id, const int32_t* mapped_edge_id, const int32_t* valid, const int32_t* offsets, int64_t num_pairs, uint64_t* pair_keys_out) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs || valid[i] == 0) return; + int32_t out = offsets[i]; + uint64_t voxel_key = static_cast(mapped_voxel_id[i]); + pair_keys_out[out] = (static_cast(static_cast(mapped_edge_id[i])) << 32) | voxel_key; +} + +__global__ void decode_pair_keys_kernel(const uint64_t* pair_keys, int64_t num_pairs, int32_t* voxel_id, int32_t* edge_id) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + edge_id[i] = static_cast(pair_keys[i] >> 32); + voxel_id[i] = static_cast(pair_keys[i] & 0xffffffffu); +} + +__device__ inline fdg_gpu::SymQEF10 symqef10_from_boundary(float3 p0, float3 p1, float boundary_weight) { + double dx = static_cast(p1.x) - static_cast(p0.x); + double dy = static_cast(p1.y) - static_cast(p0.y); + double dz = static_cast(p1.z) - static_cast(p0.z); + double L = sqrt(dx * dx + dy * dy + dz * dz); + if (L < 1e-6) return symqef10_zero(); + double ux = dx / L; + double uy = dy / L; + double uz = dz / L; + double A00 = 1.0 - ux * ux; + double A01 = -ux * uy; + double A02 = -ux * uz; + double A11 = 1.0 - uy * uy; + double A12 = -uy * uz; + double A22 = 1.0 - uz * uz; + double bx = -(A00 * p0.x + A01 * p0.y + A02 * p0.z); + double by = -(A01 * p0.x + A11 * p0.y + A12 * p0.z); + double bz = -(A02 * p0.x + A12 * p0.y + A22 * p0.z); + double c = p0.x * (A00 * p0.x + A01 * p0.y + A02 * p0.z) + + p0.y * (A01 * p0.x + A11 * p0.y + A12 * p0.z) + + p0.z * (A02 * p0.x + A12 * p0.y + A22 * p0.z); + float w = boundary_weight; + return fdg_gpu::SymQEF10{ + static_cast(w * A00), static_cast(w * A01), static_cast(w * A02), static_cast(w * bx), + static_cast(w * A11), static_cast(w * A12), static_cast(w * by), + static_cast(w * A22), static_cast(w * bz), + static_cast(w * c) + }; +} + +__global__ void build_boundary_qef_contrib_kernel(const int32_t* voxel_id, const int32_t* edge_id, int64_t num_pairs, const float* boundary_vertices, const int32_t* boundary_edges, float boundary_weight, int32_t* out_voxel_id, fdg_gpu::SymQEF10* out_qef) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + int32_t eid = edge_id[i]; + int32_t i0 = boundary_edges[2 * eid + 0]; + int32_t i1 = boundary_edges[2 * eid + 1]; + float3 p0 = make_float3(boundary_vertices[3 * i0 + 0], boundary_vertices[3 * i0 + 1], boundary_vertices[3 * i0 + 2]); + float3 p1 = make_float3(boundary_vertices[3 * i1 + 0], boundary_vertices[3 * i1 + 1], boundary_vertices[3 * i1 + 2]); + out_voxel_id[i] = voxel_id[i]; + out_qef[i] = symqef10_from_boundary(p0, p1, boundary_weight); +} + +__global__ void scatter_reduced_qef_kernel(const int32_t* reduced_voxel_id, const fdg_gpu::SymQEF10* reduced_qef, int64_t M, fdg_gpu::SymQEF10* full_qef) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= M) return; + full_qef[reduced_voxel_id[i]] = reduced_qef[i]; +} + +inline SurfaceLookup build_surface_lookup(const int* voxels, int64_t num_voxels, fdg_gpu::int3_ grid_min, fdg_gpu::int3_ grid_max, cudaStream_t stream) { + SurfaceLookup out; + out.size = num_voxels; + out.keys_sorted.allocate(num_voxels); + out.ids_sorted.allocate(num_voxels); + constexpr int kBlock = 128; + build_surface_keys_kernel<<>>(voxels, num_voxels, grid_min, grid_max, out.keys_sorted.data(), out.ids_sorted.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_surface_keys_kernel"); + thrust::sort_by_key(thrust::cuda::par.on(stream), thrust::device_pointer_cast(out.keys_sorted.data()), thrust::device_pointer_cast(out.keys_sorted.data()) + num_voxels, thrust::device_pointer_cast(out.ids_sorted.data())); + return out; +} + +inline EdgePairKeys map_and_unique_edge_pairs(const fdg_gpu::PrimitivePairResult& raw_pairs, const SurfaceLookup& lookup, fdg_gpu::int3_ grid_min, fdg_gpu::int3_ grid_max, cudaStream_t stream) { + EdgePairKeys out; + const int64_t N = raw_pairs.size; + if (N == 0) return out; + fdg_gpu::DeviceBuffer pair_voxel_keys(N); + fdg_gpu::DeviceBuffer mapped_voxel_id(N); + fdg_gpu::DeviceBuffer mapped_edge_id(N); + fdg_gpu::DeviceBuffer valid(N); + fdg_gpu::DeviceBuffer offsets(N); + constexpr int kBlock = 128; + build_raw_pair_voxel_keys_kernel<<>>(raw_pairs.voxel_i.data(), raw_pairs.voxel_j.data(), raw_pairs.voxel_k.data(), N, grid_min, grid_max, pair_voxel_keys.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_raw_pair_voxel_keys_kernel"); + map_pair_to_voxel_id_kernel<<>>(pair_voxel_keys.data(), raw_pairs.prim_id.data(), N, lookup.keys_sorted.data(), lookup.ids_sorted.data(), lookup.size, mapped_voxel_id.data(), mapped_edge_id.data(), valid.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "map_pair_to_voxel_id_kernel"); + size_t temp_bytes = 0; + void* temp = nullptr; + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum(nullptr, temp_bytes, valid.data(), offsets.data(), static_cast(N), stream)); + VOX_CUDA_CHECK(cudaMalloc(&temp, temp_bytes)); + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum(temp, temp_bytes, valid.data(), offsets.data(), static_cast(N), stream)); + int32_t last_off = voxel_traverse_edge_dda_impl::copy_last_i32(offsets.data(), N, stream); + int32_t last_valid = voxel_traverse_edge_dda_impl::copy_last_i32(valid.data(), N, stream); + cudaFree(temp); + int64_t M = static_cast(last_off) + static_cast(last_valid); + out.size = M; + out.pair_keys.allocate(M); + compact_valid_pairs_kernel<<>>(mapped_voxel_id.data(), mapped_edge_id.data(), valid.data(), offsets.data(), N, out.pair_keys.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "compact_valid_pairs_kernel"); + auto ptr = thrust::device_pointer_cast(out.pair_keys.data()); + thrust::sort(thrust::cuda::par.on(stream), ptr, ptr + M); + auto new_end = thrust::unique(thrust::cuda::par.on(stream), ptr, ptr + M); + out.size = static_cast(new_end - ptr); + return out; +} + +inline BoundaryContribStream build_boundary_contrib_stream(const EdgePairKeys& pair_keys, const float* boundary_vertices, const int32_t* boundary_edges, float boundary_weight, cudaStream_t stream) { + BoundaryContribStream out; + out.size = pair_keys.size; + out.voxel_id.allocate(out.size); + out.qef.allocate(out.size); + fdg_gpu::DeviceBuffer edge_id(out.size); + constexpr int kBlock = 128; + decode_pair_keys_kernel<<>>(pair_keys.pair_keys.data(), out.size, out.voxel_id.data(), edge_id.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "decode_pair_keys_kernel"); + build_boundary_qef_contrib_kernel<<>>(out.voxel_id.data(), edge_id.data(), out.size, boundary_vertices, boundary_edges, boundary_weight, out.voxel_id.data(), out.qef.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_boundary_qef_contrib_kernel"); + return out; +} + +inline edge_dda::BoundaryQEFResult reduce_boundary_contribs(BoundaryContribStream&& contrib, int64_t num_voxels, cudaStream_t stream) { + edge_dda::BoundaryQEFResult out; + out.size = num_voxels; + out.qefs.allocate(num_voxels); + out.qefs.clear_async(stream); + if (contrib.size == 0) return out; + auto kptr = thrust::device_pointer_cast(contrib.voxel_id.data()); + auto vptr = thrust::device_pointer_cast(contrib.qef.data()); + thrust::sort_by_key(thrust::cuda::par.on(stream), kptr, kptr + contrib.size, vptr); + fdg_gpu::DeviceBuffer reduced_ids(contrib.size); + fdg_gpu::DeviceBuffer reduced_qefs(contrib.size); + auto end_pair = thrust::reduce_by_key(thrust::cuda::par.on(stream), kptr, kptr + contrib.size, vptr, thrust::device_pointer_cast(reduced_ids.data()), thrust::device_pointer_cast(reduced_qefs.data()), thrust::equal_to(), SymQEF10Add{}); + int64_t M = end_pair.first - thrust::device_pointer_cast(reduced_ids.data()); + constexpr int kBlock = 128; + scatter_reduced_qef_kernel<<>>(reduced_ids.data(), reduced_qefs.data(), M, out.qefs.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "scatter_reduced_qef_kernel"); + return out; +} + +inline fdg_gpu::PrimitivePairResult dedup_pairs(fdg_gpu::PrimitivePairResult&& in, cudaStream_t stream) { + auto pid = thrust::device_pointer_cast(in.prim_id.data()); + auto vi = thrust::device_pointer_cast(in.voxel_i.data()); + auto vj = thrust::device_pointer_cast(in.voxel_j.data()); + auto vk = thrust::device_pointer_cast(in.voxel_k.data()); + auto begin = thrust::make_zip_iterator(thrust::make_tuple(pid, vi, vj, vk)); + auto end = thrust::make_zip_iterator(thrust::make_tuple(pid + in.size, vi + in.size, vj + in.size, vk + in.size)); + thrust::sort(thrust::cuda::par.on(stream), begin, end); + auto new_end = thrust::unique(thrust::cuda::par.on(stream), begin, end); + in.size = static_cast(new_end - begin); + return std::move(in); +} + +} // anonymous namespace + +namespace edge_dda { + +fdg_gpu::PrimitivePairResult voxel_traverse_edge_dda_gpu( + const float* d_vertices, + int64_t num_vertices, + const int32_t* d_edges, + int64_t num_edges, + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int chunk_steps, + cudaStream_t stream) { + using namespace voxel_traverse_edge_dda_impl; + if (d_vertices == nullptr || d_edges == nullptr || num_vertices < 0 || num_edges < 0) { + throw std::invalid_argument("invalid edge inputs"); + } + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + throw std::invalid_argument("invalid voxel_size"); + } + if (grid_max.x <= grid_min.x || grid_max.y <= grid_min.y || grid_max.z <= grid_min.z) { + throw std::invalid_argument("invalid grid range"); + } + if (chunk_steps <= 0) { + throw std::invalid_argument("chunk_steps must be positive"); + } + if (num_vertices == 0 || num_edges == 0) return {}; + + Workspace ws; + DeviceResult gathered{}; + gathered.edge_id = nullptr; + gathered.voxel_i = nullptr; + gathered.voxel_j = nullptr; + gathered.voxel_k = nullptr; + gathered.size = 0; + + VOX_CUDA_CHECK(cudaGetLastError()); + try { + alloc_edge_desc(&ws.edge_desc, num_edges); + alloc_u8(&ws.edge_valid, num_edges); + alloc_i32(&ws.init_count, num_edges); + alloc_i32(&ws.init_offset, num_edges); + kernel_build_edge_desc<<>>(d_vertices, d_edges, num_edges, voxel_size, ws.edge_desc, ws.edge_valid); + VOX_CUDA_CHECK(cudaGetLastError()); + kernel_count_init_jobs<<>>(ws.edge_valid, num_edges, ws.init_count); + VOX_CUDA_CHECK(cudaGetLastError()); + ensure_round_capacity(ws.round, num_edges); + exclusive_scan_i32(ws.round, ws.init_count, ws.init_offset, num_edges, stream); + int32_t last_init_offset = copy_last_i32(ws.init_offset, num_edges, stream); + int32_t last_init_count = copy_last_i32(ws.init_count, num_edges, stream); + int64_t num_init_jobs = static_cast(last_init_offset) + static_cast(last_init_count); + ensure_dda_job_queue_capacity(ws.queue_a, num_init_jobs); + ws.queue_a.size = num_init_jobs; + kernel_emit_init_jobs<<>>(ws.edge_valid, ws.edge_desc, ws.init_offset, num_edges, ws.queue_a); + VOX_CUDA_CHECK(cudaGetLastError()); + DDAJobQueue* curr = &ws.queue_a; + DDAJobQueue* next = &ws.queue_b; + while (curr->size > 0) { + int64_t nj = curr->size; + ensure_round_capacity(ws.round, nj); + kernel_count_dda_jobs<<>>(*curr, ws.edge_desc, grid_min, grid_max, chunk_steps, ws.round.pair_count, ws.round.next_job_count); + VOX_CUDA_CHECK(cudaGetLastError()); + exclusive_scan_i32(ws.round, ws.round.pair_count, ws.round.pair_offset, nj, stream); + exclusive_scan_i32(ws.round, ws.round.next_job_count, ws.round.next_job_offset, nj, stream); + int32_t last_pair_offset = copy_last_i32(ws.round.pair_offset, nj, stream); + int32_t last_pair_count = copy_last_i32(ws.round.pair_count, nj, stream); + int64_t num_pairs = static_cast(last_pair_offset) + static_cast(last_pair_count); + int32_t last_next_offset = copy_last_i32(ws.round.next_job_offset, nj, stream); + int32_t last_next_count = copy_last_i32(ws.round.next_job_count, nj, stream); + int64_t num_next_jobs = static_cast(last_next_offset) + static_cast(last_next_count); + ensure_dda_job_queue_capacity(*next, num_next_jobs); + next->size = num_next_jobs; + ResultBuffer round_result = make_result_buffer(num_pairs); + kernel_emit_dda_jobs<<>>(*curr, ws.edge_desc, grid_min, grid_max, chunk_steps, ws.round.pair_offset, ws.round.next_job_offset, round_result, *next); + VOX_CUDA_CHECK(cudaGetLastError()); + if (num_pairs > 0) ws.result_rounds.push_back(round_result); + else release_result_buffer(round_result); + std::swap(curr, next); + } + gathered = gather_result_rounds(ws.result_rounds, stream); + release_workspace(ws); + return dedup_pairs(to_primitive_pair(std::move(gathered)), stream); + } catch (...) { + + release_workspace(ws); + throw; + } +} + +BoundaryQEFResult boundary_qef_gpu( + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + const float* boundaries, + int64_t num_boundaries, + float boundary_weight, + const int* voxels, + int64_t num_voxels, + int chunk_steps, + cudaStream_t stream) { + BoundaryQEFResult out; + out.size = num_voxels; + out.qefs.allocate(num_voxels); + out.qefs.clear_async(stream); + if (num_voxels == 0 || num_boundaries == 0) return out; + if (boundaries == nullptr || voxels == nullptr) throw std::invalid_argument("null boundary_qef inputs"); + fdg_gpu::DeviceBuffer boundary_vertices(2 * num_boundaries * 3); + fdg_gpu::DeviceBuffer boundary_edges(num_boundaries * 2); + constexpr int kBlock = 128; + copy_boundaries_to_vertices_kernel<<>>(boundaries, num_boundaries, boundary_vertices.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "copy_boundaries_to_vertices_kernel"); + build_synth_edges_kernel<<>>(num_boundaries, boundary_edges.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_synth_edges_kernel"); + auto raw_pairs = voxel_traverse_edge_dda_gpu(boundary_vertices.data(), 2 * num_boundaries, boundary_edges.data(), num_boundaries, voxel_size, grid_min, grid_max, chunk_steps, stream); + auto lookup = build_surface_lookup(voxels, num_voxels, grid_min, grid_max, stream); + auto pair_keys = map_and_unique_edge_pairs(raw_pairs, lookup, grid_min, grid_max, stream); + auto contrib = build_boundary_contrib_stream(pair_keys, boundary_vertices.data(), boundary_edges.data(), boundary_weight, stream); + return reduce_boundary_contribs(std::move(contrib), num_voxels, stream); +} + +} // namespace edge_dda diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.h b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.h new file mode 100644 index 00000000..a1e54e0e --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxel_traverse_edge_dda.h @@ -0,0 +1,37 @@ +#pragma once + +#include "fdg_gpu_common.h" +#include +#include + +namespace edge_dda { + +fdg_gpu::PrimitivePairResult voxel_traverse_edge_dda_gpu( + const float* vertices, + int64_t num_vertices, + const int32_t* edges, + int64_t num_edges, + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + int chunk_steps, + cudaStream_t stream = nullptr); + +struct BoundaryQEFResult { + int64_t size = 0; + fdg_gpu::DeviceBuffer qefs; +}; + +BoundaryQEFResult boundary_qef_gpu( + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + const float* boundaries, + int64_t num_boundaries, + float boundary_weight, + const int* voxels, + int64_t num_voxels, + int chunk_steps, + cudaStream_t stream = nullptr); + +} // namespace edge_dda diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.cu new file mode 100644 index 00000000..68c16f56 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.cu @@ -0,0 +1,1502 @@ +#include "voxelize_mesh_oct.h" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace voxelize_oct_impl { + +#define VOX_CUDA_CHECK(expr) \ + do { \ + cudaError_t _err = (expr); \ + if (_err != cudaSuccess) { \ + throw std::runtime_error(std::string("CUDA error: ") + \ + cudaGetErrorString(_err) + \ + " at " + __FILE__ + ":" + \ + std::to_string(__LINE__)); \ + } \ + } while (0) + +constexpr int kRootNeighborCount = 27; +constexpr int kDefaultBlockSize = 128; + +struct FaceDesc { + float3 v0; + float3 v1; + float3 v2; + + float3 e0; + float3 e1; + float3 e2; + + float3 n_unit; + + float3 tri_bmin; + float3 tri_bmax; +}; + +struct EdgeDesc { + float3 p0; + float3 p1; + + float3 seg; + float seg_len; + float3 dir_unit; + + float3 seg_bmin; + float3 seg_bmax; +}; + +struct JobQueue { + int32_t* prim_id = nullptr; + uint8_t* level = nullptr; + int32_t* i = nullptr; + int32_t* j = nullptr; + int32_t* k = nullptr; + int64_t size = 0; + int64_t capacity = 0; +}; + +struct RoundBuffers { + uint8_t* job_hit = nullptr; + int32_t* child_count = nullptr; + int32_t* result_count = nullptr; + int32_t* child_offset = nullptr; + int32_t* result_offset = nullptr; + + void* cub_temp_storage = nullptr; + size_t cub_temp_bytes = 0; + int64_t capacity = 0; +}; + +struct ResultBuffer { + int32_t* prim_id = nullptr; + int32_t* vi = nullptr; + int32_t* vj = nullptr; + int32_t* vk = nullptr; + int64_t size = 0; +}; + +struct DeviceResult { + int32_t* prim_id = nullptr; + int32_t* voxel_i = nullptr; + int32_t* voxel_j = nullptr; + int32_t* voxel_k = nullptr; + int64_t size = 0; +}; + +struct VoxelizeWorkspace { + int32_t* leaf_ix = nullptr; + int32_t* leaf_iy = nullptr; + int32_t* leaf_iz = nullptr; + + JobQueue queue_a; + JobQueue queue_b; + RoundBuffers round; + std::vector result_rounds; +}; + +inline int ceil_div_i64(int64_t n, int block) { + return static_cast((n + block - 1) / block); +} + +inline int max3_int(int a, int b, int c) { + return a > b ? (a > c ? a : c) : (b > c ? b : c); +} + +inline int ceil_log2_pos_int(int x) { + int d = 0; + int v = 1; + while (v < x) { + v <<= 1; + ++d; + } + return d; +} + +inline int compute_grid_depth_from_grid_size(int3 grid_size) { + const int max_dim = max3_int(grid_size.x, grid_size.y, grid_size.z); + return ceil_log2_pos_int(max_dim); +} + +inline float3 reciprocal_voxel_size(float3 voxel_size) { + return make_float3(1.0f / voxel_size.x, 1.0f / voxel_size.y, 1.0f / voxel_size.z); +} + +inline void free_ptr(void* ptr) { + if (ptr != nullptr) { + cudaFree(ptr); + } +} + +inline void alloc_i32(int32_t** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(int32_t) * n)); + } +} + +inline void alloc_u8(uint8_t** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(uint8_t) * n)); + } +} + +inline void alloc_face_desc(FaceDesc** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(FaceDesc) * n)); + } +} + +inline void alloc_edge_desc(EdgeDesc** ptr, int64_t n) { + *ptr = nullptr; + if (n > 0) { + VOX_CUDA_CHECK(cudaMalloc(reinterpret_cast(ptr), sizeof(EdgeDesc) * n)); + } +} + +inline void release_job_queue(JobQueue& q) { + free_ptr(q.prim_id); + free_ptr(q.level); + free_ptr(q.i); + free_ptr(q.j); + free_ptr(q.k); + q = {}; +} + +inline void release_round_buffers(RoundBuffers& b) { + free_ptr(b.job_hit); + free_ptr(b.child_count); + free_ptr(b.result_count); + free_ptr(b.child_offset); + free_ptr(b.result_offset); + free_ptr(b.cub_temp_storage); + b = {}; +} + +inline void release_result_buffer(ResultBuffer& r) { + free_ptr(r.prim_id); + free_ptr(r.vi); + free_ptr(r.vj); + free_ptr(r.vk); + r = {}; +} + +inline void release_workspace(VoxelizeWorkspace& ws) { + free_ptr(ws.leaf_ix); + free_ptr(ws.leaf_iy); + free_ptr(ws.leaf_iz); + release_job_queue(ws.queue_a); + release_job_queue(ws.queue_b); + release_round_buffers(ws.round); + for (auto& r : ws.result_rounds) { + release_result_buffer(r); + } + ws.result_rounds.clear(); +} + +inline void ensure_job_queue_capacity(JobQueue& q, int64_t capacity) { + if (capacity <= q.capacity) { + return; + } + release_job_queue(q); + alloc_i32(&q.prim_id, capacity); + alloc_u8(&q.level, capacity); + alloc_i32(&q.i, capacity); + alloc_i32(&q.j, capacity); + alloc_i32(&q.k, capacity); + q.capacity = capacity; + q.size = 0; +} + +inline void ensure_round_capacity(RoundBuffers& b, int64_t capacity) { + if (capacity <= b.capacity) { + return; + } + + free_ptr(b.job_hit); + free_ptr(b.child_count); + free_ptr(b.result_count); + free_ptr(b.child_offset); + free_ptr(b.result_offset); + + alloc_u8(&b.job_hit, capacity); + alloc_i32(&b.child_count, capacity); + alloc_i32(&b.result_count, capacity); + alloc_i32(&b.child_offset, capacity); + alloc_i32(&b.result_offset, capacity); + b.capacity = capacity; +} + +inline void ensure_scan_temp_storage( + RoundBuffers& b, + int32_t* d_in, + int32_t* d_out, + int64_t count, + cudaStream_t stream) { + if (count <= 0) { + return; + } + if (count > INT32_MAX) { + throw std::runtime_error("CUB scan count exceeds int32 range"); + } + size_t bytes = 0; + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + nullptr, + bytes, + d_in, + d_out, + static_cast(count), + stream)); + if (bytes > b.cub_temp_bytes) { + free_ptr(b.cub_temp_storage); + VOX_CUDA_CHECK(cudaMalloc(&b.cub_temp_storage, bytes)); + b.cub_temp_bytes = bytes; + } +} + +inline void exclusive_scan_i32( + RoundBuffers& b, + int32_t* d_in, + int32_t* d_out, + int64_t count, + cudaStream_t stream) { + if (count <= 0) { + return; + } + ensure_scan_temp_storage(b, d_in, d_out, count, stream); + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + b.cub_temp_storage, + b.cub_temp_bytes, + d_in, + d_out, + static_cast(count), + stream)); +} + +inline int32_t copy_last_i32(const int32_t* ptr, int64_t count, cudaStream_t stream) { + if (count <= 0) { + return 0; + } + int32_t value = 0; + VOX_CUDA_CHECK(cudaMemcpyAsync( + &value, + ptr + (count - 1), + sizeof(int32_t), + cudaMemcpyDeviceToHost, + stream)); + VOX_CUDA_CHECK(cudaStreamSynchronize(stream)); + return value; +} + +inline ResultBuffer make_result_buffer(int64_t count) { + ResultBuffer r; + if (count <= 0) { + return r; + } + alloc_i32(&r.prim_id, count); + alloc_i32(&r.vi, count); + alloc_i32(&r.vj, count); + alloc_i32(&r.vk, count); + r.size = count; + return r; +} + +inline DeviceResult gather_result_rounds( + const std::vector& rounds, + cudaStream_t stream) { + DeviceResult out; + int64_t total = 0; + for (const auto& r : rounds) { + total += r.size; + } + out.size = total; + if (total == 0) { + return out; + } + + alloc_i32(&out.prim_id, total); + alloc_i32(&out.voxel_i, total); + alloc_i32(&out.voxel_j, total); + alloc_i32(&out.voxel_k, total); + + int64_t cursor = 0; + for (const auto& r : rounds) { + if (r.size == 0) { + continue; + } + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.prim_id + cursor, + r.prim_id, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_i + cursor, + r.vi, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_j + cursor, + r.vj, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + VOX_CUDA_CHECK(cudaMemcpyAsync( + out.voxel_k + cursor, + r.vk, + sizeof(int32_t) * r.size, + cudaMemcpyDeviceToDevice, + stream)); + cursor += r.size; + } + + VOX_CUDA_CHECK(cudaStreamSynchronize(stream)); + return out; +} + +__device__ inline float3 add3(const float3& a, const float3& b) { + return make_float3(a.x + b.x, a.y + b.y, a.z + b.z); +} + +__device__ inline float3 sub3(const float3& a, const float3& b) { + return make_float3(a.x - b.x, a.y - b.y, a.z - b.z); +} + +__device__ inline float3 mul3(const float3& a, float s) { + return make_float3(a.x * s, a.y * s, a.z * s); +} + +__device__ inline float3 mul3_comp(const float3& a, const float3& b) { + return make_float3(a.x * b.x, a.y * b.y, a.z * b.z); +} + +__device__ inline float dot3(const float3& a, const float3& b) { + return a.x * b.x + a.y * b.y + a.z * b.z; +} + +__device__ inline float2 dot2_pair(const float2& a, const float2& b) { + return make_float2(a.x * b.x, a.y * b.y); +} + +__device__ inline float dot2(const float2& a, const float2& b) { + return a.x * b.x + a.y * b.y; +} + +__device__ inline float3 cross3(const float3& a, const float3& b) { + return make_float3( + a.y * b.z - a.z * b.y, + a.z * b.x - a.x * b.z, + a.x * b.y - a.y * b.x); +} + +__device__ inline float3 min3(const float3& a, const float3& b) { + return make_float3(fminf(a.x, b.x), fminf(a.y, b.y), fminf(a.z, b.z)); +} + +__device__ inline float3 max3(const float3& a, const float3& b) { + return make_float3(fmaxf(a.x, b.x), fmaxf(a.y, b.y), fmaxf(a.z, b.z)); +} + +__device__ inline float3 normalize3(const float3& a) { + const float z = dot3(a, a); + if (z > 0.0f) { + const float n = sqrtf(z); + return make_float3(a.x / n, a.y / n, a.z / n); + } else { + return a; + } +} + +__device__ inline bool bbox_overlap_closed( + const float3& a_min, + const float3& a_max, + const float3& b_min, + const float3& b_max) { + return !(a_max.x < b_min.x || b_max.x < a_min.x || + a_max.y < b_min.y || b_max.y < a_min.y || + a_max.z < b_min.z || b_max.z < a_min.z); +} + +__device__ inline float2 max2_zero(const float2& a) { + return make_float2(fmaxf(a.x, 0.0f), fmaxf(a.y, 0.0f)); +} + +__device__ inline void compute_face_min_level_and_root( + int ix0, int iy0, int iz0, + int ix1, int iy1, int iz1, + int ix2, int iy2, int iz2, + int d, + uint8_t& level, + int& root_i, + int& root_j, + int& root_k) { + uint32_t diff = + static_cast(ix0 ^ ix1) | static_cast(ix0 ^ ix2) | + static_cast(iy0 ^ iy1) | static_cast(iy0 ^ iy2) | + static_cast(iz0 ^ iz1) | static_cast(iz0 ^ iz2); + + if (diff == 0) { + level = static_cast(d); + root_i = ix0; + root_j = iy0; + root_k = iz0; + return; + } + + int msb = 31 - __clz(diff); + int l = d - 1 - msb; + level = static_cast(l); + + int shift = d - l; + root_i = ix0 >> shift; + root_j = iy0 >> shift; + root_k = iz0 >> shift; +} + +__device__ inline void compute_edge_min_level_and_root( + int ix0, int iy0, int iz0, + int ix1, int iy1, int iz1, + int d, + uint8_t& level, + int& root_i, + int& root_j, + int& root_k) { + uint32_t diff = + static_cast(ix0 ^ ix1) | + static_cast(iy0 ^ iy1) | + static_cast(iz0 ^ iz1); + + if (diff == 0) { + level = static_cast(d); + root_i = ix0; + root_j = iy0; + root_k = iz0; + return; + } + + int msb = 31 - __clz(diff); + int l = d - 1 - msb; + level = static_cast(l); + + int shift = d - l; + root_i = ix0 >> shift; + root_j = iy0 >> shift; + root_k = iz0 >> shift; +} + +__device__ inline bool node_intersects_valid_domain( + int d, + int level, + int i, + int j, + int k, + int3 grid_size) { + int cells = 1 << level; + if (i < 0 || i >= cells || j < 0 || j >= cells || k < 0 || k >= cells) { + return false; + } + + int node_span = 1 << (d - level); + int x0 = i * node_span; + int y0 = j * node_span; + int z0 = k * node_span; + + return (x0 < grid_size.x) && (y0 < grid_size.y) && (z0 < grid_size.z); +} + +__device__ inline void compute_node_box_world( + int d, + int level, + int i, + int j, + int k, + int3 grid_min, + float3 voxel_size, + float3& box_min, + float3& box_size, + float3& box_max) { + int node_span = 1 << (d - level); + + int global_base_x = grid_min.x + i * node_span; + int global_base_y = grid_min.y + j * node_span; + int global_base_z = grid_min.z + k * node_span; + + int global_end_x = global_base_x + node_span; + int global_end_y = global_base_y + node_span; + int global_end_z = global_base_z + node_span; + + box_min = make_float3( + static_cast(global_base_x) * voxel_size.x, + static_cast(global_base_y) * voxel_size.y, + static_cast(global_base_z) * voxel_size.z); + box_max = make_float3( + static_cast(global_end_x) * voxel_size.x, + static_cast(global_end_y) * voxel_size.y, + static_cast(global_end_z) * voxel_size.z); + box_size = sub3(box_max, box_min); +} + +__device__ inline bool face_qef_style_triangle_box_hit( + const FaceDesc& f, + const float3& box_min, + const float3& box_size, + const float3& box_max) { + if (!bbox_overlap_closed(f.tri_bmin, f.tri_bmax, box_min, box_max)) { + return false; + } + + const float3& n = f.n_unit; + + float3 c = make_float3( + n.x > 0.0f ? box_size.x : 0.0f, + n.y > 0.0f ? box_size.y : 0.0f, + n.z > 0.0f ? box_size.z : 0.0f); + + float d1 = dot3(n, sub3(c, f.v0)); + float d2 = dot3(n, sub3(sub3(box_size, c), f.v0)); + + int mul_xy = (n.z < 0.0f) ? -1 : 1; + float2 n_xy_e0 = make_float2(-mul_xy * f.e0.y, mul_xy * f.e0.x); + float2 n_xy_e1 = make_float2(-mul_xy * f.e1.y, mul_xy * f.e1.x); + float2 n_xy_e2 = make_float2(-mul_xy * f.e2.y, mul_xy * f.e2.x); + + float d_xy_e0 = -dot2(n_xy_e0, make_float2(f.v0.x, f.v0.y)) + + dot2(max2_zero(n_xy_e0), make_float2(box_size.x, box_size.y)); + float d_xy_e1 = -dot2(n_xy_e1, make_float2(f.v1.x, f.v1.y)) + + dot2(max2_zero(n_xy_e1), make_float2(box_size.x, box_size.y)); + float d_xy_e2 = -dot2(n_xy_e2, make_float2(f.v2.x, f.v2.y)) + + dot2(max2_zero(n_xy_e2), make_float2(box_size.x, box_size.y)); + + int mul_yz = (n.x < 0.0f) ? -1 : 1; + float2 n_yz_e0 = make_float2(-mul_yz * f.e0.z, mul_yz * f.e0.y); + float2 n_yz_e1 = make_float2(-mul_yz * f.e1.z, mul_yz * f.e1.y); + float2 n_yz_e2 = make_float2(-mul_yz * f.e2.z, mul_yz * f.e2.y); + + float d_yz_e0 = -dot2(n_yz_e0, make_float2(f.v0.y, f.v0.z)) + + dot2(max2_zero(n_yz_e0), make_float2(box_size.y, box_size.z)); + float d_yz_e1 = -dot2(n_yz_e1, make_float2(f.v1.y, f.v1.z)) + + dot2(max2_zero(n_yz_e1), make_float2(box_size.y, box_size.z)); + float d_yz_e2 = -dot2(n_yz_e2, make_float2(f.v2.y, f.v2.z)) + + dot2(max2_zero(n_yz_e2), make_float2(box_size.y, box_size.z)); + + int mul_zx = (n.y < 0.0f) ? -1 : 1; + float2 n_zx_e0 = make_float2(-mul_zx * f.e0.x, mul_zx * f.e0.z); + float2 n_zx_e1 = make_float2(-mul_zx * f.e1.x, mul_zx * f.e1.z); + float2 n_zx_e2 = make_float2(-mul_zx * f.e2.x, mul_zx * f.e2.z); + + float d_zx_e0 = -dot2(n_zx_e0, make_float2(f.v0.z, f.v0.x)) + + dot2(max2_zero(n_zx_e0), make_float2(box_size.z, box_size.x)); + float d_zx_e1 = -dot2(n_zx_e1, make_float2(f.v1.z, f.v1.x)) + + dot2(max2_zero(n_zx_e1), make_float2(box_size.z, box_size.x)); + float d_zx_e2 = -dot2(n_zx_e2, make_float2(f.v2.z, f.v2.x)) + + dot2(max2_zero(n_zx_e2), make_float2(box_size.z, box_size.x)); + + float n_dot_p = dot3(n, box_min); + if (((n_dot_p + d1) * (n_dot_p + d2)) > 0.0f) { + return false; + } + + float2 p_xy = make_float2(box_min.x, box_min.y); + if (dot2(n_xy_e0, p_xy) + d_xy_e0 < 0.0f) return false; + if (dot2(n_xy_e1, p_xy) + d_xy_e1 < 0.0f) return false; + if (dot2(n_xy_e2, p_xy) + d_xy_e2 < 0.0f) return false; + + float2 p_yz = make_float2(box_min.y, box_min.z); + if (dot2(n_yz_e0, p_yz) + d_yz_e0 < 0.0f) return false; + if (dot2(n_yz_e1, p_yz) + d_yz_e1 < 0.0f) return false; + if (dot2(n_yz_e2, p_yz) + d_yz_e2 < 0.0f) return false; + + float2 p_zx = make_float2(box_min.z, box_min.x); + if (dot2(n_zx_e0, p_zx) + d_zx_e0 < 0.0f) return false; + if (dot2(n_zx_e1, p_zx) + d_zx_e1 < 0.0f) return false; + if (dot2(n_zx_e2, p_zx) + d_zx_e2 < 0.0f) return false; + + return true; +} + +__device__ inline bool segment_box_overlap_world( + const EdgeDesc& e, + const float3& box_min, + const float3& box_max) { + if (e.seg_len < 1.0e-6f) { + return false; + } + + if (!bbox_overlap_closed(e.seg_bmin, e.seg_bmax, box_min, box_max)) { + return false; + } + + float tmin = 0.0f; + float tmax = 1.0f; + + if (e.seg.x == 0.0f) { + if (!(box_min.x <= e.p0.x && e.p0.x <= box_max.x)) return false; + } else { + float inv_d = 1.0f / e.seg.x; + float t1 = (box_min.x - e.p0.x) * inv_d; + float t2 = (box_max.x - e.p0.x) * inv_d; + if (t1 > t2) { float tmp = t1; t1 = t2; t2 = tmp; } + tmin = fmaxf(tmin, t1); + tmax = fminf(tmax, t2); + if (tmin > tmax) return false; + } + + if (e.seg.y == 0.0f) { + if (!(box_min.y <= e.p0.y && e.p0.y <= box_max.y)) return false; + } else { + float inv_d = 1.0f / e.seg.y; + float t1 = (box_min.y - e.p0.y) * inv_d; + float t2 = (box_max.y - e.p0.y) * inv_d; + if (t1 > t2) { float tmp = t1; t1 = t2; t2 = tmp; } + tmin = fmaxf(tmin, t1); + tmax = fminf(tmax, t2); + if (tmin > tmax) return false; + } + + if (e.seg.z == 0.0f) { + if (!(box_min.z <= e.p0.z && e.p0.z <= box_max.z)) return false; + } else { + float inv_d = 1.0f / e.seg.z; + float t1 = (box_min.z - e.p0.z) * inv_d; + float t2 = (box_max.z - e.p0.z) * inv_d; + if (t1 > t2) { float tmp = t1; t1 = t2; t2 = tmp; } + tmin = fmaxf(tmin, t1); + tmax = fminf(tmax, t2); + if (tmin > tmax) return false; + } + + return true; +} + +__global__ void kernel_build_leaf_coords( + const float* __restrict__ vertices, + int64_t num_vertices, + float3 inv_voxel_size, + int3 grid_min, + int3 grid_size, + int32_t* __restrict__ leaf_ix, + int32_t* __restrict__ leaf_iy, + int32_t* __restrict__ leaf_iz) { + int64_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (vid >= num_vertices) { + return; + } + + float x = vertices[3 * vid + 0] * inv_voxel_size.x; + float y = vertices[3 * vid + 1] * inv_voxel_size.y; + float z = vertices[3 * vid + 2] * inv_voxel_size.z; + + int ix = static_cast(floorf(x)) - grid_min.x; + int iy = static_cast(floorf(y)) - grid_min.y; + int iz = static_cast(floorf(z)) - grid_min.z; + + ix = max(0, min(ix, grid_size.x - 1)); + iy = max(0, min(iy, grid_size.y - 1)); + iz = max(0, min(iz, grid_size.z - 1)); + + leaf_ix[vid] = ix; + leaf_iy[vid] = iy; + leaf_iz[vid] = iz; +} + +__global__ void kernel_init_faces_and_emit_root27( + const float* __restrict__ vertices, + const int32_t* __restrict__ faces, + const int32_t* __restrict__ leaf_ix, + const int32_t* __restrict__ leaf_iy, + const int32_t* __restrict__ leaf_iz, + int64_t num_faces, + int d, + FaceDesc* __restrict__ face_desc, + JobQueue out_q) { + int64_t fid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (fid >= num_faces) { + return; + } + + int v0_id = faces[3 * fid + 0]; + int v1_id = faces[3 * fid + 1]; + int v2_id = faces[3 * fid + 2]; + + float3 v0 = make_float3(vertices[3 * v0_id + 0], vertices[3 * v0_id + 1], vertices[3 * v0_id + 2]); + float3 v1 = make_float3(vertices[3 * v1_id + 0], vertices[3 * v1_id + 1], vertices[3 * v1_id + 2]); + float3 v2 = make_float3(vertices[3 * v2_id + 0], vertices[3 * v2_id + 1], vertices[3 * v2_id + 2]); + + FaceDesc fd; + fd.v0 = v0; + fd.v1 = v1; + fd.v2 = v2; + fd.e0 = sub3(v1, v0); + fd.e1 = sub3(v2, v1); + fd.e2 = sub3(v0, v2); + fd.n_unit = normalize3(cross3(fd.e0, fd.e1)); + fd.tri_bmin = min3(v0, min3(v1, v2)); + fd.tri_bmax = max3(v0, max3(v1, v2)); + face_desc[fid] = fd; + + uint8_t level; + int root_i, root_j, root_k; + compute_face_min_level_and_root( + leaf_ix[v0_id], leaf_iy[v0_id], leaf_iz[v0_id], + leaf_ix[v1_id], leaf_iy[v1_id], leaf_iz[v1_id], + leaf_ix[v2_id], leaf_iy[v2_id], leaf_iz[v2_id], + d, + level, + root_i, + root_j, + root_k); + + int64_t base = static_cast(kRootNeighborCount) * fid; + int slot = 0; + for (int dz = -1; dz <= 1; ++dz) { + for (int dy = -1; dy <= 1; ++dy) { + for (int dx = -1; dx <= 1; ++dx) { + int64_t out = base + slot++; + out_q.prim_id[out] = static_cast(fid); + out_q.level[out] = level; + out_q.i[out] = root_i + dx; + out_q.j[out] = root_j + dy; + out_q.k[out] = root_k + dz; + } + } + } +} + +__global__ void kernel_init_edges_and_emit_root27( + const float* __restrict__ vertices, + const int32_t* __restrict__ edges, + const int32_t* __restrict__ leaf_ix, + const int32_t* __restrict__ leaf_iy, + const int32_t* __restrict__ leaf_iz, + int64_t num_edges, + int d, + EdgeDesc* __restrict__ edge_desc, + JobQueue out_q) { + int64_t eid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (eid >= num_edges) { + return; + } + + int v0_id = edges[2 * eid + 0]; + int v1_id = edges[2 * eid + 1]; + + float3 p0 = make_float3(vertices[3 * v0_id + 0], vertices[3 * v0_id + 1], vertices[3 * v0_id + 2]); + float3 p1 = make_float3(vertices[3 * v1_id + 0], vertices[3 * v1_id + 1], vertices[3 * v1_id + 2]); + + EdgeDesc ed; + ed.p0 = p0; + ed.p1 = p1; + ed.seg = sub3(p1, p0); + ed.seg_len = sqrtf(dot3(ed.seg, ed.seg)); + ed.dir_unit = (ed.seg_len >= 1.0e-6f) ? mul3(ed.seg, 1.0f / ed.seg_len) : make_float3(0.0f, 0.0f, 0.0f); + ed.seg_bmin = min3(p0, p1); + ed.seg_bmax = max3(p0, p1); + edge_desc[eid] = ed; + + uint8_t level; + int root_i, root_j, root_k; + compute_edge_min_level_and_root( + leaf_ix[v0_id], leaf_iy[v0_id], leaf_iz[v0_id], + leaf_ix[v1_id], leaf_iy[v1_id], leaf_iz[v1_id], + d, + level, + root_i, + root_j, + root_k); + + int64_t base = static_cast(kRootNeighborCount) * eid; + int slot = 0; + for (int dz = -1; dz <= 1; ++dz) { + for (int dy = -1; dy <= 1; ++dy) { + for (int dx = -1; dx <= 1; ++dx) { + int64_t out = base + slot++; + out_q.prim_id[out] = static_cast(eid); + out_q.level[out] = level; + out_q.i[out] = root_i + dx; + out_q.j[out] = root_j + dy; + out_q.k[out] = root_k + dz; + } + } + } +} + +__global__ void kernel_count_face_jobs( + JobQueue curr_q, + const FaceDesc* __restrict__ face_desc, + int d, + int3 grid_min, + int3 grid_size, + float3 voxel_size, + uint8_t* __restrict__ job_hit, + int32_t* __restrict__ child_count, + int32_t* __restrict__ result_count) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= curr_q.size) { + return; + } + + int fid = curr_q.prim_id[idx]; + int level = static_cast(curr_q.level[idx]); + int i = curr_q.i[idx]; + int j = curr_q.j[idx]; + int k = curr_q.k[idx]; + + if (!node_intersects_valid_domain(d, level, i, j, k, grid_size)) { + job_hit[idx] = 0; + child_count[idx] = 0; + result_count[idx] = 0; + return; + } + + float3 box_min, box_size, box_max; + compute_node_box_world(d, level, i, j, k, grid_min, voxel_size, box_min, box_size, box_max); + + bool hit = face_qef_style_triangle_box_hit(face_desc[fid], box_min, box_size, box_max); + job_hit[idx] = static_cast(hit ? 1 : 0); + + if (!hit) { + child_count[idx] = 0; + result_count[idx] = 0; + } else if (level < d) { + child_count[idx] = 8; + result_count[idx] = 0; + } else { + child_count[idx] = 0; + result_count[idx] = 1; + } +} + +__global__ void kernel_count_edge_jobs( + JobQueue curr_q, + const EdgeDesc* __restrict__ edge_desc, + int d, + int3 grid_min, + int3 grid_size, + float3 voxel_size, + uint8_t* __restrict__ job_hit, + int32_t* __restrict__ child_count, + int32_t* __restrict__ result_count) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= curr_q.size) { + return; + } + + int eid = curr_q.prim_id[idx]; + int level = static_cast(curr_q.level[idx]); + int i = curr_q.i[idx]; + int j = curr_q.j[idx]; + int k = curr_q.k[idx]; + + if (!node_intersects_valid_domain(d, level, i, j, k, grid_size)) { + job_hit[idx] = 0; + child_count[idx] = 0; + result_count[idx] = 0; + return; + } + + float3 box_min, box_size, box_max; + compute_node_box_world(d, level, i, j, k, grid_min, voxel_size, box_min, box_size, box_max); + + bool hit = segment_box_overlap_world(edge_desc[eid], box_min, box_max); + job_hit[idx] = static_cast(hit ? 1 : 0); + + if (!hit) { + child_count[idx] = 0; + result_count[idx] = 0; + } else if (level < d) { + child_count[idx] = 8; + result_count[idx] = 0; + } else { + child_count[idx] = 0; + result_count[idx] = 1; + } +} + +__global__ void kernel_emit_jobs( + JobQueue curr_q, + const uint8_t* __restrict__ job_hit, + const int32_t* __restrict__ child_offset, + const int32_t* __restrict__ result_offset, + int d, + JobQueue next_q, + ResultBuffer out_res) { + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= curr_q.size) { + return; + } + if (job_hit[idx] == 0) { + return; + } + + int prim_id = curr_q.prim_id[idx]; + int level = static_cast(curr_q.level[idx]); + int i = curr_q.i[idx]; + int j = curr_q.j[idx]; + int k = curr_q.k[idx]; + + if (level < d) { + int32_t base = child_offset[idx]; + int child_level = level + 1; + int slot = 0; + for (int bz = 0; bz < 2; ++bz) { + for (int by = 0; by < 2; ++by) { + for (int bx = 0; bx < 2; ++bx) { + int32_t out = base + slot++; + next_q.prim_id[out] = prim_id; + next_q.level[out] = static_cast(child_level); + next_q.i[out] = 2 * i + bx; + next_q.j[out] = 2 * j + by; + next_q.k[out] = 2 * k + bz; + } + } + } + } else { + int32_t out = result_offset[idx]; + out_res.prim_id[out] = prim_id; + out_res.vi[out] = i; + out_res.vj[out] = j; + out_res.vk[out] = k; + } +} + +inline void reset_output(DeviceResult& out) { + out.prim_id = nullptr; + out.voxel_i = nullptr; + out.voxel_j = nullptr; + out.voxel_k = nullptr; + out.size = 0; +} + +inline void release_device_result(DeviceResult& out) { + free_ptr(out.prim_id); + free_ptr(out.voxel_i); + free_ptr(out.voxel_j); + free_ptr(out.voxel_k); + out = {}; +} + +} // namespace voxelize_oct_impl + + +namespace { + +inline fdg_gpu::PrimitivePairResult to_primitive_pair(voxelize_oct_impl::DeviceResult&& r) { + fdg_gpu::PrimitivePairResult out; + out.size = r.size; + out.prim_id.adopt(r.prim_id, r.size); + out.voxel_i.adopt(r.voxel_i, r.size); + out.voxel_j.adopt(r.voxel_j, r.size); + out.voxel_k.adopt(r.voxel_k, r.size); + r.prim_id = nullptr; + r.voxel_i = nullptr; + r.voxel_j = nullptr; + r.voxel_k = nullptr; + r.size = 0; + return out; +} + +struct SurfaceLookup { + int64_t size = 0; + fdg_gpu::DeviceBuffer keys_sorted; + fdg_gpu::DeviceBuffer ids_sorted; +}; + +struct FacePairKeys { + int64_t size = 0; + fdg_gpu::DeviceBuffer keys; +}; + +struct FaceContribStream { + int64_t size = 0; + fdg_gpu::DeviceBuffer voxel_id; + fdg_gpu::DeviceBuffer qef; +}; + +__host__ __device__ inline fdg_gpu::SymQEF10 symqef10_zero() { + return fdg_gpu::SymQEF10{0,0,0,0,0,0,0,0,0,0}; +} + +struct SymQEF10Add { + __host__ __device__ fdg_gpu::SymQEF10 operator()(const fdg_gpu::SymQEF10& a, const fdg_gpu::SymQEF10& b) const { + return fdg_gpu::SymQEF10{ + a.q00 + b.q00, a.q01 + b.q01, a.q02 + b.q02, a.q03 + b.q03, + a.q11 + b.q11, a.q12 + b.q12, a.q13 + b.q13, + a.q22 + b.q22, a.q23 + b.q23, + a.q33 + b.q33}; + } +}; + +__host__ __device__ inline uint64_t pack_pair_key(int32_t voxel_id, int32_t face_id) { + return (static_cast(static_cast(voxel_id)) << 32) | + static_cast(face_id); +} + +__host__ __device__ inline int32_t unpack_pair_voxel_id(uint64_t k) { + return static_cast(k >> 32); +} + +__host__ __device__ inline int32_t unpack_pair_face_id(uint64_t k) { + return static_cast(k & 0xffffffffu); +} + +__host__ __device__ inline fdg_gpu::SymQEF10 symqef10_from_plane(float4 p) { + const float a = p.x, b = p.y, c = p.z, d = p.w; + return fdg_gpu::SymQEF10{ + a*a, a*b, a*c, a*d, + b*b, b*c, b*d, + c*c, c*d, + d*d + }; +} + +__global__ void build_synth_faces_kernel(int64_t num_triangles, int32_t* __restrict__ faces) { + int64_t fid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (fid >= num_triangles) return; + faces[3 * fid + 0] = static_cast(3 * fid + 0); + faces[3 * fid + 1] = static_cast(3 * fid + 1); + faces[3 * fid + 2] = static_cast(3 * fid + 2); +} + +__global__ void build_surface_keys_kernel( + const int* __restrict__ voxels, + int64_t num_voxels, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + uint64_t* __restrict__ keys, + int32_t* __restrict__ ids) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_voxels) return; + int x = voxels[3 * i + 0]; + int y = voxels[3 * i + 1]; + int z = voxels[3 * i + 2]; + keys[i] = fdg_gpu::pack_voxel_key(x, y, z, grid_min, grid_max); + ids[i] = static_cast(i); +} + +__global__ void build_raw_pair_keys_kernel( + const int32_t* __restrict__ voxel_i, + const int32_t* __restrict__ voxel_j, + const int32_t* __restrict__ voxel_k, + const int32_t* __restrict__ face_id, + int64_t num_pairs, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + uint64_t* __restrict__ voxel_keys, + int32_t* __restrict__ pair_face_ids) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + const int32_t gx = voxel_i[i] + grid_min.x; + const int32_t gy = voxel_j[i] + grid_min.y; + const int32_t gz = voxel_k[i] + grid_min.z; + voxel_keys[i] = fdg_gpu::pack_voxel_key(gx, gy, gz, grid_min, grid_max); + pair_face_ids[i] = face_id[i]; +} + +__device__ inline int lower_bound_u64(const uint64_t* arr, int64_t n, uint64_t key) { + int64_t lo = 0; + int64_t hi = n; + while (lo < hi) { + int64_t mid = (lo + hi) >> 1; + uint64_t v = arr[mid]; + if (v < key) lo = mid + 1; + else hi = mid; + } + return static_cast(lo); +} + +__global__ void map_pair_to_voxel_id_kernel( + const uint64_t* __restrict__ pair_keys, + const int32_t* __restrict__ pair_face_ids, + int64_t num_pairs, + const uint64_t* __restrict__ surface_keys_sorted, + const int32_t* __restrict__ surface_ids_sorted, + int64_t num_voxels, + int32_t* __restrict__ mapped_voxel_id, + int32_t* __restrict__ mapped_face_id, + int32_t* __restrict__ valid) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + uint64_t key = pair_keys[i]; + int pos = lower_bound_u64(surface_keys_sorted, num_voxels, key); + if (pos < num_voxels && surface_keys_sorted[pos] == key) { + mapped_voxel_id[i] = surface_ids_sorted[pos]; + mapped_face_id[i] = pair_face_ids[i]; + valid[i] = 1; + } else { + mapped_voxel_id[i] = -1; + mapped_face_id[i] = -1; + valid[i] = 0; + } +} + +__global__ void compact_valid_pairs_kernel( + const int32_t* __restrict__ mapped_voxel_id, + const int32_t* __restrict__ mapped_face_id, + const int32_t* __restrict__ valid, + const int32_t* __restrict__ offsets, + int64_t num_pairs, + uint64_t* __restrict__ pair_keys_out) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs || valid[i] == 0) return; + int32_t out = offsets[i]; + pair_keys_out[out] = pack_pair_key(mapped_voxel_id[i], mapped_face_id[i]); +} + +__global__ void build_face_qef_contrib_kernel( + const uint64_t* __restrict__ pair_keys, + int64_t num_pairs, + const float* __restrict__ vertices, + const int32_t* __restrict__ faces, + int32_t* __restrict__ voxel_id_out, + fdg_gpu::SymQEF10* __restrict__ qef_out) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_pairs) return; + + int32_t voxel_id = unpack_pair_voxel_id(pair_keys[i]); + int32_t fid = unpack_pair_face_id(pair_keys[i]); + + int32_t i0 = faces[3 * fid + 0]; + int32_t i1 = faces[3 * fid + 1]; + int32_t i2 = faces[3 * fid + 2]; + + float3 v0 = make_float3(vertices[3 * i0 + 0], vertices[3 * i0 + 1], vertices[3 * i0 + 2]); + float3 v1 = make_float3(vertices[3 * i1 + 0], vertices[3 * i1 + 1], vertices[3 * i1 + 2]); + float3 v2 = make_float3(vertices[3 * i2 + 0], vertices[3 * i2 + 1], vertices[3 * i2 + 2]); + + float3 e0 = voxelize_oct_impl::sub3(v1, v0); + float3 e1 = voxelize_oct_impl::sub3(v2, v1); + float3 n = voxelize_oct_impl::normalize3(voxelize_oct_impl::cross3(e0, e1)); + float4 plane = make_float4(n.x, n.y, n.z, -voxelize_oct_impl::dot3(n, v0)); + + voxel_id_out[i] = voxel_id; + qef_out[i] = symqef10_from_plane(plane); +} + +__global__ void scatter_reduced_face_qef_kernel( + const int32_t* __restrict__ reduced_voxel_id, + const fdg_gpu::SymQEF10* __restrict__ reduced_qef, + int64_t num_reduced, + fdg_gpu::SymQEF10* __restrict__ full_qefs) { + int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= num_reduced) return; + full_qefs[reduced_voxel_id[i]] = reduced_qef[i]; +} + +inline SurfaceLookup build_surface_lookup( + const int* voxels, + int64_t num_voxels, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + cudaStream_t stream) { + SurfaceLookup out; + out.size = num_voxels; + out.keys_sorted.allocate(num_voxels); + out.ids_sorted.allocate(num_voxels); + constexpr int kBlock = 256; + build_surface_keys_kernel<<>>( + voxels, num_voxels, grid_min, grid_max, out.keys_sorted.data(), out.ids_sorted.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_surface_keys_kernel"); + thrust::device_ptr kptr(out.keys_sorted.data()); + thrust::device_ptr iptr(out.ids_sorted.data()); + thrust::sort_by_key(thrust::cuda::par.on(stream), kptr, kptr + num_voxels, iptr); + return out; +} + +inline FacePairKeys map_and_unique_face_pairs( + const fdg_gpu::PrimitivePairResult& raw_pairs, + const SurfaceLookup& lookup, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + cudaStream_t stream) { + FacePairKeys out; + const int64_t N = raw_pairs.size; + if (N == 0) return out; + fdg_gpu::DeviceBuffer pair_voxel_keys(N); + fdg_gpu::DeviceBuffer pair_face_ids(N); + fdg_gpu::DeviceBuffer mapped_voxel_id(N); + fdg_gpu::DeviceBuffer mapped_face_id(N); + fdg_gpu::DeviceBuffer valid(N); + fdg_gpu::DeviceBuffer offsets(N); + constexpr int kBlock = 256; + + build_raw_pair_keys_kernel<<>>( + raw_pairs.voxel_i.data(), raw_pairs.voxel_j.data(), raw_pairs.voxel_k.data(), raw_pairs.prim_id.data(), + N, grid_min, grid_max, pair_voxel_keys.data(), pair_face_ids.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_raw_pair_keys_kernel"); + + map_pair_to_voxel_id_kernel<<>>( + pair_voxel_keys.data(), pair_face_ids.data(), N, + lookup.keys_sorted.data(), lookup.ids_sorted.data(), lookup.size, + mapped_voxel_id.data(), mapped_face_id.data(), valid.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "map_pair_to_voxel_id_kernel"); + + void* temp = nullptr; + size_t temp_bytes = 0; + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum(nullptr, temp_bytes, valid.data(), offsets.data(), static_cast(N), stream)); + VOX_CUDA_CHECK(cudaMalloc(&temp, temp_bytes)); + VOX_CUDA_CHECK(cub::DeviceScan::ExclusiveSum(temp, temp_bytes, valid.data(), offsets.data(), static_cast(N), stream)); + int32_t last_off = voxelize_oct_impl::copy_last_i32(offsets.data(), N, stream); + int32_t last_valid = voxelize_oct_impl::copy_last_i32(valid.data(), N, stream); + int64_t M = static_cast(last_off + last_valid); + cudaFree(temp); + + out.size = M; + out.keys.allocate(M); + compact_valid_pairs_kernel<<>>( + mapped_voxel_id.data(), mapped_face_id.data(), valid.data(), offsets.data(), N, out.keys.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "compact_valid_pairs_kernel"); + + thrust::device_ptr kptr(out.keys.data()); + thrust::sort(thrust::cuda::par.on(stream), kptr, kptr + M); + auto new_end = thrust::unique(thrust::cuda::par.on(stream), kptr, kptr + M); + out.size = static_cast(new_end - kptr); + return out; +} + +inline FaceContribStream build_face_contrib_stream( + const FacePairKeys& pair_keys, + const float* triangles_world, + const int32_t* faces_synth, + cudaStream_t stream) { + FaceContribStream out; + out.size = pair_keys.size; + if (out.size == 0) return out; + out.voxel_id.allocate(out.size); + out.qef.allocate(out.size); + constexpr int kBlock = 256; + build_face_qef_contrib_kernel<<>>( + pair_keys.keys.data(), out.size, triangles_world, faces_synth, out.voxel_id.data(), out.qef.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_face_qef_contrib_kernel"); + return out; +} + +inline oct_pairs::FaceQEFResult reduce_face_contribs( + FaceContribStream&& contrib, + int64_t num_voxels, + cudaStream_t stream) { + oct_pairs::FaceQEFResult out; + out.size = num_voxels; + out.qefs.allocate(num_voxels); + out.qefs.clear_async(stream); + if (contrib.size == 0) return out; + + thrust::device_ptr kptr(contrib.voxel_id.data()); + thrust::device_ptr vptr(contrib.qef.data()); + thrust::sort_by_key(thrust::cuda::par.on(stream), kptr, kptr + contrib.size, vptr); + + fdg_gpu::DeviceBuffer reduced_ids(contrib.size); + fdg_gpu::DeviceBuffer reduced_qefs(contrib.size); + auto end_pair = thrust::reduce_by_key( + thrust::cuda::par.on(stream), + kptr, kptr + contrib.size, + vptr, + thrust::device_pointer_cast(reduced_ids.data()), + thrust::device_pointer_cast(reduced_qefs.data()), + thrust::equal_to(), + SymQEF10Add()); + int64_t M = end_pair.first - thrust::device_pointer_cast(reduced_ids.data()); + + constexpr int kBlock = 256; + scatter_reduced_face_qef_kernel<<>>( + reduced_ids.data(), reduced_qefs.data(), M, out.qefs.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "scatter_reduced_face_qef_kernel"); + return out; +} + +} // anonymous namespace + +namespace oct_pairs { + +fdg_gpu::PrimitivePairResult voxelize_mesh_oct_gpu( + const float* d_vertices, + int64_t num_vertices, + const int32_t* d_faces, + int64_t num_faces, + fdg_gpu::int3_ grid_min_, + fdg_gpu::int3_ grid_size_, + float3 voxel_size, + cudaStream_t stream) { + using namespace voxelize_oct_impl; + int3 grid_min{grid_min_.x, grid_min_.y, grid_min_.z}; + int3 grid_size{grid_size_.x, grid_size_.y, grid_size_.z}; + if (d_vertices == nullptr || d_faces == nullptr || num_vertices < 0 || num_faces < 0) { + throw std::invalid_argument("invalid mesh inputs"); + } + if (grid_size.x <= 0 || grid_size.y <= 0 || grid_size.z <= 0) { + throw std::invalid_argument("invalid grid_size"); + } + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + throw std::invalid_argument("invalid voxel_size"); + } + if (num_vertices == 0 || num_faces == 0) { + return {}; + } + const int d = compute_grid_depth_from_grid_size(grid_size); + if (d < 0 || d > 21) throw std::invalid_argument("grid depth exceeds 21"); + const float3 inv_voxel_size = reciprocal_voxel_size(voxel_size); + + VoxelizeWorkspace ws; + FaceDesc* face_desc = nullptr; + DeviceResult gathered; reset_output(gathered); + + alloc_i32(&ws.leaf_ix, num_vertices); + alloc_i32(&ws.leaf_iy, num_vertices); + alloc_i32(&ws.leaf_iz, num_vertices); + alloc_face_desc(&face_desc, num_faces); + + kernel_build_leaf_coords<<>>( + d_vertices, num_vertices, inv_voxel_size, grid_min, grid_size, ws.leaf_ix, ws.leaf_iy, ws.leaf_iz); + VOX_CUDA_CHECK(cudaGetLastError()); + + ensure_job_queue_capacity(ws.queue_a, static_cast(kRootNeighborCount) * num_faces); + ws.queue_a.size = static_cast(kRootNeighborCount) * num_faces; + + kernel_init_faces_and_emit_root27<<>>( + d_vertices, d_faces, ws.leaf_ix, ws.leaf_iy, ws.leaf_iz, num_faces, d, face_desc, ws.queue_a); + VOX_CUDA_CHECK(cudaGetLastError()); + + JobQueue* curr = &ws.queue_a; JobQueue* next = &ws.queue_b; + while (curr->size > 0) { + int64_t nj = curr->size; + ensure_round_capacity(ws.round, nj); + kernel_count_face_jobs<<>>( + *curr, face_desc, d, grid_min, grid_size, voxel_size, ws.round.job_hit, ws.round.child_count, ws.round.result_count); + VOX_CUDA_CHECK(cudaGetLastError()); + exclusive_scan_i32(ws.round, ws.round.child_count, ws.round.child_offset, nj, stream); + exclusive_scan_i32(ws.round, ws.round.result_count, ws.round.result_offset, nj, stream); + const int32_t num_children_total = copy_last_i32(ws.round.child_offset, nj, stream) + copy_last_i32(ws.round.child_count, nj, stream); + const int32_t num_results_total = copy_last_i32(ws.round.result_offset, nj, stream) + copy_last_i32(ws.round.result_count, nj, stream); + ensure_job_queue_capacity(*next, num_children_total); + next->size = num_children_total; + ResultBuffer round_res = make_result_buffer(num_results_total); + kernel_emit_jobs<<>>( + *curr, ws.round.job_hit, ws.round.child_offset, ws.round.result_offset, d, *next, round_res); + VOX_CUDA_CHECK(cudaGetLastError()); + if (num_results_total > 0) ws.result_rounds.push_back(round_res); + std::swap(curr, next); + } + gathered = gather_result_rounds(ws.result_rounds, stream); + free_ptr(face_desc); + release_workspace(ws); + return to_primitive_pair(std::move(gathered)); +} + +fdg_gpu::PrimitivePairResult voxelize_edge_oct_gpu( + const float* d_vertices, + int64_t num_vertices, + const int32_t* d_edges, + int64_t num_edges, + fdg_gpu::int3_ grid_min_, + fdg_gpu::int3_ grid_size_, + float3 voxel_size, + cudaStream_t stream) { + using namespace voxelize_oct_impl; + int3 grid_min{grid_min_.x, grid_min_.y, grid_min_.z}; + int3 grid_size{grid_size_.x, grid_size_.y, grid_size_.z}; + if (d_vertices == nullptr || d_edges == nullptr || num_vertices < 0 || num_edges < 0) { + throw std::invalid_argument("invalid edge inputs"); + } + if (grid_size.x <= 0 || grid_size.y <= 0 || grid_size.z <= 0) { + throw std::invalid_argument("invalid grid_size"); + } + if (!(voxel_size.x > 0.0f && voxel_size.y > 0.0f && voxel_size.z > 0.0f)) { + throw std::invalid_argument("invalid voxel_size"); + } + if (num_vertices == 0 || num_edges == 0) { + return {}; + } + const int d = compute_grid_depth_from_grid_size(grid_size); + if (d < 0 || d > 21) throw std::invalid_argument("grid depth exceeds 21"); + const float3 inv_voxel_size = reciprocal_voxel_size(voxel_size); + + VoxelizeWorkspace ws; + EdgeDesc* edge_desc = nullptr; + DeviceResult gathered; reset_output(gathered); + + alloc_i32(&ws.leaf_ix, num_vertices); + alloc_i32(&ws.leaf_iy, num_vertices); + alloc_i32(&ws.leaf_iz, num_vertices); + alloc_edge_desc(&edge_desc, num_edges); + kernel_build_leaf_coords<<>>( + d_vertices, num_vertices, inv_voxel_size, grid_min, grid_size, ws.leaf_ix, ws.leaf_iy, ws.leaf_iz); + VOX_CUDA_CHECK(cudaGetLastError()); + ensure_job_queue_capacity(ws.queue_a, static_cast(kRootNeighborCount) * num_edges); + ws.queue_a.size = static_cast(kRootNeighborCount) * num_edges; + kernel_init_edges_and_emit_root27<<>>( + d_vertices, d_edges, ws.leaf_ix, ws.leaf_iy, ws.leaf_iz, num_edges, d, edge_desc, ws.queue_a); + VOX_CUDA_CHECK(cudaGetLastError()); + JobQueue* curr = &ws.queue_a; JobQueue* next = &ws.queue_b; + while (curr->size > 0) { + int64_t nj = curr->size; + ensure_round_capacity(ws.round, nj); + kernel_count_edge_jobs<<>>( + *curr, edge_desc, d, grid_min, grid_size, voxel_size, ws.round.job_hit, ws.round.child_count, ws.round.result_count); + VOX_CUDA_CHECK(cudaGetLastError()); + exclusive_scan_i32(ws.round, ws.round.child_count, ws.round.child_offset, nj, stream); + exclusive_scan_i32(ws.round, ws.round.result_count, ws.round.result_offset, nj, stream); + const int32_t num_children_total = copy_last_i32(ws.round.child_offset, nj, stream) + copy_last_i32(ws.round.child_count, nj, stream); + const int32_t num_results_total = copy_last_i32(ws.round.result_offset, nj, stream) + copy_last_i32(ws.round.result_count, nj, stream); + ensure_job_queue_capacity(*next, num_children_total); + next->size = num_children_total; + ResultBuffer round_res = make_result_buffer(num_results_total); + kernel_emit_jobs<<>>( + *curr, ws.round.job_hit, ws.round.child_offset, ws.round.result_offset, d, *next, round_res); + VOX_CUDA_CHECK(cudaGetLastError()); + if (num_results_total > 0) ws.result_rounds.push_back(round_res); + std::swap(curr, next); + } + gathered = gather_result_rounds(ws.result_rounds, stream); + free_ptr(edge_desc); + release_workspace(ws); + return to_primitive_pair(std::move(gathered)); +} + +FaceQEFResult face_qef_gpu( + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + const float* triangles, + int64_t num_triangles, + const int* voxels, + int64_t num_voxels, + cudaStream_t stream) { + FaceQEFResult out; + out.size = num_voxels; + out.qefs.allocate(num_voxels); + out.qefs.clear_async(stream); + if (num_voxels == 0 || num_triangles == 0) return out; + if (triangles == nullptr || voxels == nullptr) throw std::invalid_argument("null face_qef inputs"); + + const int64_t num_tri_vertices = num_triangles * 3; + fdg_gpu::DeviceBuffer faces_synth(num_triangles * 3); + + constexpr int kBlock = 256; + build_synth_faces_kernel<<>>( + num_triangles, faces_synth.data()); + fdg_gpu::throw_cuda_error(cudaGetLastError(), "build_synth_faces_kernel"); + + fdg_gpu::int3_ grid_size{grid_max.x - grid_min.x, grid_max.y - grid_min.y, grid_max.z - grid_min.z}; + auto raw_pairs = voxelize_mesh_oct_gpu( + triangles, num_tri_vertices, faces_synth.data(), num_triangles, grid_min, grid_size, voxel_size, stream); + + auto lookup = build_surface_lookup(voxels, num_voxels, grid_min, grid_max, stream); + auto face_pair_keys = map_and_unique_face_pairs(raw_pairs, lookup, grid_min, grid_max, stream); + auto contrib = build_face_contrib_stream(face_pair_keys, triangles, faces_synth.data(), stream); + return reduce_face_contribs(std::move(contrib), num_voxels, stream); +} + +} // namespace oct_pairs diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.h b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.h new file mode 100644 index 00000000..20195901 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_oct.h @@ -0,0 +1,44 @@ +#pragma once + +#include "fdg_gpu_common.h" +#include +#include + +namespace oct_pairs { + +fdg_gpu::PrimitivePairResult voxelize_mesh_oct_gpu( + const float* vertices, + int64_t num_vertices, + const int32_t* faces, + int64_t num_faces, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_size, + float3 voxel_size, + cudaStream_t stream = nullptr); + +fdg_gpu::PrimitivePairResult voxelize_edge_oct_gpu( + const float* vertices, + int64_t num_vertices, + const int32_t* edges, + int64_t num_edges, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_size, + float3 voxel_size, + cudaStream_t stream = nullptr); + +struct FaceQEFResult { + int64_t size = 0; + fdg_gpu::DeviceBuffer qefs; +}; + +FaceQEFResult face_qef_gpu( + float3 voxel_size, + fdg_gpu::int3_ grid_min, + fdg_gpu::int3_ grid_max, + const float* triangles, + int64_t num_triangles, + const int* voxels, + int64_t num_voxels, + cudaStream_t stream = nullptr); + +} // namespace oct_pairs diff --git a/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_octree.cu b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_octree.cu new file mode 100644 index 00000000..b636b338 --- /dev/null +++ b/o-voxel/src/convert/mesh_to_flexible_dual_grid_gpu/voxelize_mesh_octree.cu @@ -0,0 +1,677 @@ +#include "../api.h" + +#include "types.cuh" + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +// Standalone octree voxelization. It is kept as a public CUDA module, but the +// main flexible dual grid CUDA path does not call it. Work is represented as an +// octree job stream that expands from coarse cells to fine cells on the GPU. +namespace o_voxel::fdg +{ + namespace + { + + constexpr int kThreads = 256; + constexpr int kRootNeighborCount = 27; + + // Per-face data reused while octree jobs for that face are refined. + struct FaceDesc + { + float3 v0; + float3 v1; + float3 v2; + float3 e0; + float3 e1; + float3 e2; + float3 n; + float3 bmin; + float3 bmax; + }; + + __host__ __device__ __forceinline__ int64_t div_up_i64(int64_t n, int64_t d) + { + return (n + d - 1) / d; + } + + int ceil_log2_pos(int x) + { + int d = 0; + int v = 1; + while (v < x) + { + v <<= 1; + ++d; + } + return d; + } + + int grid_depth(Int3 grid_size) + { + return ceil_log2_pos(std::max(grid_size.x, std::max(grid_size.y, grid_size.z))); + } + + int64_t read_scan_total(const torch::Tensor &counts, const torch::Tensor &offsets, int64_t n, cudaStream_t stream) + { + // After an exclusive scan, total = last_count + last_offset. The + // host loop needs this exact size before allocating next_jobs/results. + int64_t tail[2] = {0, 0}; + C10_CUDA_CHECK(cudaMemcpyAsync( + tail, + counts.data_ptr() + n - 1, + sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaMemcpyAsync( + tail + 1, + offsets.data_ptr() + n - 1, + sizeof(int64_t), + cudaMemcpyDeviceToHost, + stream)); + C10_CUDA_CHECK(cudaStreamSynchronize(stream)); + return tail[0] + tail[1]; + } + + void cub_exclusive_sum_i64(const torch::Tensor &in, const torch::Tensor &out, int64_t n, cudaStream_t stream) + { + if (n == 0) + return; + TORCH_CHECK(n <= std::numeric_limits::max(), "CUB item count exceeds int"); + size_t temp_bytes = 0; + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + nullptr, + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + auto temp = torch::empty( + {static_cast(temp_bytes)}, + torch::TensorOptions().dtype(torch::kUInt8).device(in.device())); + C10_CUDA_CHECK(cub::DeviceScan::ExclusiveSum( + temp.data_ptr(), + temp_bytes, + in.data_ptr(), + out.data_ptr(), + static_cast(n), + stream)); + } + + __device__ __forceinline__ float3 sub3(float3 a, float3 b) + { + return make_float3(a.x - b.x, a.y - b.y, a.z - b.z); + } + + __device__ __forceinline__ float dot3(float3 a, float3 b) + { + return a.x * b.x + a.y * b.y + a.z * b.z; + } + + __device__ __forceinline__ float dot2(float2 a, float2 b) + { + return a.x * b.x + a.y * b.y; + } + + __device__ __forceinline__ float3 cross3(float3 a, float3 b) + { + return make_float3( + a.y * b.z - a.z * b.y, + a.z * b.x - a.x * b.z, + a.x * b.y - a.y * b.x); + } + + __device__ __forceinline__ float3 min3(float3 a, float3 b) + { + return make_float3(fminf(a.x, b.x), fminf(a.y, b.y), fminf(a.z, b.z)); + } + + __device__ __forceinline__ float3 max3(float3 a, float3 b) + { + return make_float3(fmaxf(a.x, b.x), fmaxf(a.y, b.y), fmaxf(a.z, b.z)); + } + + __device__ __forceinline__ float2 max2_zero(float2 a) + { + return make_float2(fmaxf(a.x, 0.0f), fmaxf(a.y, 0.0f)); + } + + __device__ __forceinline__ float3 normalize3(float3 a) + { + const float n2 = dot3(a, a); + if (n2 <= 0.0f) + return a; + const float n = sqrtf(n2); + return make_float3(a.x / n, a.y / n, a.z / n); + } + + __device__ __forceinline__ bool bbox_overlap_closed(float3 a_min, float3 a_max, float3 b_min, float3 b_max) + { + return !(a_max.x < b_min.x || b_max.x < a_min.x || + a_max.y < b_min.y || b_max.y < a_min.y || + a_max.z < b_min.z || b_max.z < a_min.z); + } + + __device__ bool triangle_box_hit(const FaceDesc &f, float3 box_min, float3 box_size, float3 box_max) + { + // Triangle-box overlap test. AABB rejects obvious misses first; the + // remaining tests check whether the triangle plane crosses the box + // and whether the box overlaps the triangle in all three projections. + if (!bbox_overlap_closed(f.bmin, f.bmax, box_min, box_max)) + return false; + + const float3 n = f.n; + const float3 c = make_float3( + n.x > 0.0f ? box_size.x : 0.0f, + n.y > 0.0f ? box_size.y : 0.0f, + n.z > 0.0f ? box_size.z : 0.0f); + const float d1 = dot3(n, sub3(c, f.v0)); + const float d2 = dot3(n, sub3(sub3(box_size, c), f.v0)); + + // Projected edge half-spaces. The sign flip chooses inward-facing + // edge normals for each projection, and max2_zero shifts the tested + // box corner to the side most favorable to overlap. + const int mul_xy = n.z < 0.0f ? -1 : 1; + const float2 n_xy_e0 = make_float2(-mul_xy * f.e0.y, mul_xy * f.e0.x); + const float2 n_xy_e1 = make_float2(-mul_xy * f.e1.y, mul_xy * f.e1.x); + const float2 n_xy_e2 = make_float2(-mul_xy * f.e2.y, mul_xy * f.e2.x); + const float d_xy_e0 = -dot2(n_xy_e0, make_float2(f.v0.x, f.v0.y)) + + dot2(max2_zero(n_xy_e0), make_float2(box_size.x, box_size.y)); + const float d_xy_e1 = -dot2(n_xy_e1, make_float2(f.v1.x, f.v1.y)) + + dot2(max2_zero(n_xy_e1), make_float2(box_size.x, box_size.y)); + const float d_xy_e2 = -dot2(n_xy_e2, make_float2(f.v2.x, f.v2.y)) + + dot2(max2_zero(n_xy_e2), make_float2(box_size.x, box_size.y)); + + const int mul_yz = n.x < 0.0f ? -1 : 1; + const float2 n_yz_e0 = make_float2(-mul_yz * f.e0.z, mul_yz * f.e0.y); + const float2 n_yz_e1 = make_float2(-mul_yz * f.e1.z, mul_yz * f.e1.y); + const float2 n_yz_e2 = make_float2(-mul_yz * f.e2.z, mul_yz * f.e2.y); + const float d_yz_e0 = -dot2(n_yz_e0, make_float2(f.v0.y, f.v0.z)) + + dot2(max2_zero(n_yz_e0), make_float2(box_size.y, box_size.z)); + const float d_yz_e1 = -dot2(n_yz_e1, make_float2(f.v1.y, f.v1.z)) + + dot2(max2_zero(n_yz_e1), make_float2(box_size.y, box_size.z)); + const float d_yz_e2 = -dot2(n_yz_e2, make_float2(f.v2.y, f.v2.z)) + + dot2(max2_zero(n_yz_e2), make_float2(box_size.y, box_size.z)); + + const int mul_zx = n.y < 0.0f ? -1 : 1; + const float2 n_zx_e0 = make_float2(-mul_zx * f.e0.x, mul_zx * f.e0.z); + const float2 n_zx_e1 = make_float2(-mul_zx * f.e1.x, mul_zx * f.e1.z); + const float2 n_zx_e2 = make_float2(-mul_zx * f.e2.x, mul_zx * f.e2.z); + const float d_zx_e0 = -dot2(n_zx_e0, make_float2(f.v0.z, f.v0.x)) + + dot2(max2_zero(n_zx_e0), make_float2(box_size.z, box_size.x)); + const float d_zx_e1 = -dot2(n_zx_e1, make_float2(f.v1.z, f.v1.x)) + + dot2(max2_zero(n_zx_e1), make_float2(box_size.z, box_size.x)); + const float d_zx_e2 = -dot2(n_zx_e2, make_float2(f.v2.z, f.v2.x)) + + dot2(max2_zero(n_zx_e2), make_float2(box_size.z, box_size.x)); + + const float n_dot_p = dot3(n, box_min); + // Plane slab test: the two extreme box corners along the triangle + // normal must lie on opposite sides, or one side exactly on plane. + if (((n_dot_p + d1) * (n_dot_p + d2)) > 0.0f) + return false; + + const float2 p_xy = make_float2(box_min.x, box_min.y); + if (dot2(n_xy_e0, p_xy) + d_xy_e0 < 0.0f) + return false; + if (dot2(n_xy_e1, p_xy) + d_xy_e1 < 0.0f) + return false; + if (dot2(n_xy_e2, p_xy) + d_xy_e2 < 0.0f) + return false; + + const float2 p_yz = make_float2(box_min.y, box_min.z); + if (dot2(n_yz_e0, p_yz) + d_yz_e0 < 0.0f) + return false; + if (dot2(n_yz_e1, p_yz) + d_yz_e1 < 0.0f) + return false; + if (dot2(n_yz_e2, p_yz) + d_yz_e2 < 0.0f) + return false; + + const float2 p_zx = make_float2(box_min.z, box_min.x); + if (dot2(n_zx_e0, p_zx) + d_zx_e0 < 0.0f) + return false; + if (dot2(n_zx_e1, p_zx) + d_zx_e1 < 0.0f) + return false; + if (dot2(n_zx_e2, p_zx) + d_zx_e2 < 0.0f) + return false; + + return true; + } + + __device__ void compute_face_root( + int ix0, + int iy0, + int iz0, + int ix1, + int iy1, + int iz1, + int ix2, + int iy2, + int iz2, + int d, + int &level, + int &root_i, + int &root_j, + int &root_k) + { + // XOR of vertex leaf coordinates tells which octree bits differ + // across the triangle vertices. The highest differing bit selects + // the smallest common ancestor node containing all three vertex cells. + const uint32_t diff = + static_cast(ix0 ^ ix1) | static_cast(ix0 ^ ix2) | + static_cast(iy0 ^ iy1) | static_cast(iy0 ^ iy2) | + static_cast(iz0 ^ iz1) | static_cast(iz0 ^ iz2); + if (diff == 0) + { + level = d; + root_i = ix0; + root_j = iy0; + root_k = iz0; + return; + } + const int msb = 31 - __clz(diff); + level = d - 1 - msb; + const int shift = d - level; + root_i = ix0 >> shift; + root_j = iy0 >> shift; + root_k = iz0 >> shift; + } + + __device__ bool node_in_domain(int d, int level, int i, int j, int k, Int3 grid_size) + { + const int cells = 1 << level; + if (i < 0 || i >= cells || j < 0 || j >= cells || k < 0 || k >= cells) + return false; + const int span = 1 << (d - level); + return i * span < grid_size.x && j * span < grid_size.y && k * span < grid_size.z; + } + + __device__ void node_box(int d, int level, int i, int j, int k, Int3 grid_min, float3 voxel_size, float3 &box_min, float3 &box_size, float3 &box_max) + { + const int span = 1 << (d - level); + const int gx0 = grid_min.x + i * span; + const int gy0 = grid_min.y + j * span; + const int gz0 = grid_min.z + k * span; + const int gx1 = gx0 + span; + const int gy1 = gy0 + span; + const int gz1 = gz0 + span; + box_min = make_float3(gx0 * voxel_size.x, gy0 * voxel_size.y, gz0 * voxel_size.z); + box_max = make_float3(gx1 * voxel_size.x, gy1 * voxel_size.y, gz1 * voxel_size.z); + box_size = sub3(box_max, box_min); + } + + __global__ void build_leaf_coords_kernel( + const float *__restrict__ vertices, + int64_t num_vertices, + float3 inv_voxel_size, + Int3 grid_min, + Int3 grid_size, + int32_t *__restrict__ leaf_coords) + { + const int64_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (vid >= num_vertices) + return; + // Convert each vertex to a leaf voxel coordinate in grid-local + // space. These coordinates are only used to seed each face's octree + // root; geometric overlap is still checked later by triangle_box_hit. + int ix = static_cast(floorf(vertices[3 * vid + 0] * inv_voxel_size.x)) - grid_min.x; + int iy = static_cast(floorf(vertices[3 * vid + 1] * inv_voxel_size.y)) - grid_min.y; + int iz = static_cast(floorf(vertices[3 * vid + 2] * inv_voxel_size.z)) - grid_min.z; + ix = max(0, min(ix, grid_size.x - 1)); + iy = max(0, min(iy, grid_size.y - 1)); + iz = max(0, min(iz, grid_size.z - 1)); + leaf_coords[3 * vid + 0] = ix; + leaf_coords[3 * vid + 1] = iy; + leaf_coords[3 * vid + 2] = iz; + } + + __global__ void init_faces_kernel( + const float *__restrict__ vertices, + const int32_t *__restrict__ faces, + const int32_t *__restrict__ leaf_coords, + int64_t num_faces, + int d, + FaceDesc *__restrict__ face_desc, + int32_t *__restrict__ jobs) + { + const int64_t fid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (fid >= num_faces) + return; + + const int v0_id = faces[3 * fid + 0]; + const int v1_id = faces[3 * fid + 1]; + const int v2_id = faces[3 * fid + 2]; + const float3 v0 = make_float3(vertices[3 * v0_id + 0], vertices[3 * v0_id + 1], vertices[3 * v0_id + 2]); + const float3 v1 = make_float3(vertices[3 * v1_id + 0], vertices[3 * v1_id + 1], vertices[3 * v1_id + 2]); + const float3 v2 = make_float3(vertices[3 * v2_id + 0], vertices[3 * v2_id + 1], vertices[3 * v2_id + 2]); + + FaceDesc f; + f.v0 = v0; + f.v1 = v1; + f.v2 = v2; + f.e0 = sub3(v1, v0); + f.e1 = sub3(v2, v1); + f.e2 = sub3(v0, v2); + f.n = normalize3(cross3(f.e0, f.e1)); + f.bmin = min3(v0, min3(v1, v2)); + f.bmax = max3(v0, max3(v1, v2)); + face_desc[fid] = f; + + // Start from the smallest octree node that contains the triangle's + // vertex leaf cells. The 3x3x3 neighbor seed around that root keeps + // conservative coverage when the triangle crosses nearby cells. + int level; + int root_i; + int root_j; + int root_k; + compute_face_root( + leaf_coords[3 * v0_id + 0], + leaf_coords[3 * v0_id + 1], + leaf_coords[3 * v0_id + 2], + leaf_coords[3 * v1_id + 0], + leaf_coords[3 * v1_id + 1], + leaf_coords[3 * v1_id + 2], + leaf_coords[3 * v2_id + 0], + leaf_coords[3 * v2_id + 1], + leaf_coords[3 * v2_id + 2], + d, + level, + root_i, + root_j, + root_k); + + const int64_t base = fid * kRootNeighborCount; + int slot = 0; + for (int dz = -1; dz <= 1; ++dz) + { + for (int dy = -1; dy <= 1; ++dy) + { + for (int dx = -1; dx <= 1; ++dx) + { + const int64_t out = base + slot++; + jobs[5 * out + 0] = static_cast(fid); + jobs[5 * out + 1] = level; + jobs[5 * out + 2] = root_i + dx; + jobs[5 * out + 3] = root_j + dy; + jobs[5 * out + 4] = root_k + dz; + } + } + } + } + + __global__ void count_jobs_kernel( + const int32_t *__restrict__ jobs, + int64_t num_jobs, + const FaceDesc *__restrict__ face_desc, + int d, + Int3 grid_min, + Int3 grid_size, + float3 voxel_size, + int64_t *__restrict__ child_count, + int64_t *__restrict__ result_count) + { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_jobs) + return; + const int fid = jobs[5 * idx + 0]; + const int level = jobs[5 * idx + 1]; + const int i = jobs[5 * idx + 2]; + const int j = jobs[5 * idx + 3]; + const int k = jobs[5 * idx + 4]; + + if (!node_in_domain(d, level, i, j, k, grid_size)) + { + child_count[idx] = 0; + result_count[idx] = 0; + return; + } + + float3 box_min; + float3 box_size; + float3 box_max; + node_box(d, level, i, j, k, grid_min, voxel_size, box_min, box_size, box_max); + if (!triangle_box_hit(face_desc[fid], box_min, box_size, box_max)) + { + // Miss: no children and no leaf result. + child_count[idx] = 0; + result_count[idx] = 0; + } + else if (level < d) + { + // Hit at an internal octree node: split into eight children for + // the next breadth-first refinement level. + child_count[idx] = 8; + result_count[idx] = 0; + } + else + { + // Hit at leaf level: this face intersects this voxel. + child_count[idx] = 0; + result_count[idx] = 1; + } + } + + __global__ void emit_jobs_kernel( + const int32_t *__restrict__ jobs, + int64_t num_jobs, + const int64_t *__restrict__ child_count, + const int64_t *__restrict__ result_count, + const int64_t *__restrict__ child_offsets, + const int64_t *__restrict__ result_offsets, + int32_t *__restrict__ next_jobs, + int32_t *__restrict__ result_prim, + int32_t *__restrict__ result_voxels) + { + const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (idx >= num_jobs) + return; + const int prim_id = jobs[5 * idx + 0]; + const int level = jobs[5 * idx + 1]; + const int i = jobs[5 * idx + 2]; + const int j = jobs[5 * idx + 3]; + const int k = jobs[5 * idx + 4]; + + if (child_count[idx] != 0) + { + const int64_t base = child_offsets[idx]; + const int child_level = level + 1; + int slot = 0; + // child_offsets is the scan of child_count, so every surviving + // job writes its eight children into a disjoint range. + for (int bz = 0; bz < 2; ++bz) + { + for (int by = 0; by < 2; ++by) + { + for (int bx = 0; bx < 2; ++bx) + { + const int64_t out = base + slot++; + next_jobs[5 * out + 0] = prim_id; + next_jobs[5 * out + 1] = child_level; + next_jobs[5 * out + 2] = 2 * i + bx; + next_jobs[5 * out + 3] = 2 * j + by; + next_jobs[5 * out + 4] = 2 * k + bz; + } + } + } + } + else if (result_count[idx] != 0) + { + const int64_t out = result_offsets[idx]; + // result_offsets compacts leaf hits from this level into the + // chunk returned to the host loop. + result_prim[out] = prim_id; + result_voxels[3 * out + 0] = i; + result_voxels[3 * out + 1] = j; + result_voxels[3 * out + 2] = k; + } + } + + void copy_chunk(torch::Tensor &dst, int64_t dst_offset, const torch::Tensor &src, int64_t values, cudaStream_t stream) + { + if (values == 0) + return; + C10_CUDA_CHECK(cudaMemcpyAsync( + dst.data_ptr() + dst_offset, + src.data_ptr(), + values * sizeof(int32_t), + cudaMemcpyDeviceToDevice, + stream)); + } + + } // namespace + + std::tuple + voxelize_mesh_octree_cuda( + const torch::Tensor &vertices, + const torch::Tensor &faces, + const std::vector &voxel_size, + const std::vector &grid_range) + { + // Returns one row per primitive/voxel hit: + // prim_ids [K] int32 and voxels [K, 3] int32. + TORCH_CHECK(vertices.is_cuda(), "vertices must be a CUDA tensor"); + TORCH_CHECK(faces.is_cuda(), "faces must be a CUDA tensor"); + + const c10::cuda::CUDAGuard guard(vertices.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(vertices.get_device()).stream(); + const torch::Device device = vertices.device(); + const auto opts_i32 = torch::TensorOptions().dtype(torch::kInt32).device(device); + const auto opts_i64 = torch::TensorOptions().dtype(torch::kInt64).device(device); + const auto opts_u8 = torch::TensorOptions().dtype(torch::kUInt8).device(device); + const int64_t num_vertices = vertices.size(0); + const int64_t num_faces = faces.size(0); + auto empty_prim = torch::empty({0}, opts_i32); + auto empty_voxels = torch::empty({0, 3}, opts_i32); + if (num_vertices == 0 || num_faces == 0) + return std::make_tuple(empty_prim, empty_voxels); + + const float3 voxel_size_h{voxel_size[0], voxel_size[1], voxel_size[2]}; + const float3 inv_voxel_size{1.0f / voxel_size_h.x, 1.0f / voxel_size_h.y, 1.0f / voxel_size_h.z}; + const Int3 grid_min{ + static_cast(grid_range[0]), + static_cast(grid_range[1]), + static_cast(grid_range[2])}; + const Int3 grid_max{ + static_cast(grid_range[3]), + static_cast(grid_range[4]), + static_cast(grid_range[5])}; + const Int3 grid_size{grid_max.x - grid_min.x, grid_max.y - grid_min.y, grid_max.z - grid_min.z}; + const int d = grid_depth(grid_size); + // d is the full leaf depth of the cubic octree that covers the grid. + // A leaf node corresponds to one voxel-sized cell. + TORCH_CHECK(d <= 21, "grid depth exceeds 21"); + + auto leaf_coords = torch::empty({num_vertices, 3}, opts_i32); + int blocks = static_cast(div_up_i64(num_vertices, kThreads)); + build_leaf_coords_kernel<<>>( + vertices.data_ptr(), + num_vertices, + inv_voxel_size, + grid_min, + grid_size, + leaf_coords.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto face_desc_storage = torch::empty( + {num_faces * static_cast(sizeof(FaceDesc))}, + opts_u8); + FaceDesc *face_desc = reinterpret_cast(face_desc_storage.data_ptr()); + int64_t num_jobs = num_faces * kRootNeighborCount; + torch::Tensor jobs = torch::empty({num_jobs, 5}, opts_i32); + blocks = static_cast(div_up_i64(num_faces, kThreads)); + init_faces_kernel<<>>( + vertices.data_ptr(), + faces.data_ptr(), + leaf_coords.data_ptr(), + num_faces, + d, + face_desc, + jobs.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + std::vector prim_chunks; + std::vector voxel_chunks; + std::vector chunk_sizes; + int64_t total_results = 0; + + while (num_jobs > 0) + { + // Breadth-first refinement. Each iteration classifies all current + // jobs in parallel, scans counts with CUB, then emits the next level + // jobs and any leaf hits. + auto child_count = torch::empty({num_jobs}, opts_i64); + auto result_count = torch::empty({num_jobs}, opts_i64); + auto child_offsets = torch::empty({num_jobs}, opts_i64); + auto result_offsets = torch::empty({num_jobs}, opts_i64); + blocks = static_cast(div_up_i64(num_jobs, kThreads)); + count_jobs_kernel<<>>( + jobs.data_ptr(), + num_jobs, + face_desc, + d, + grid_min, + grid_size, + voxel_size_h, + child_count.data_ptr(), + result_count.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + cub_exclusive_sum_i64(child_count, child_offsets, num_jobs, stream); + cub_exclusive_sum_i64(result_count, result_offsets, num_jobs, stream); + const int64_t next_count = read_scan_total(child_count, child_offsets, num_jobs, stream); + const int64_t result_count_h = read_scan_total(result_count, result_offsets, num_jobs, stream); + + auto next_jobs = torch::empty({next_count, 5}, opts_i32); + auto result_prim = torch::empty({result_count_h}, opts_i32); + auto result_voxels = torch::empty({result_count_h, 3}, opts_i32); + emit_jobs_kernel<<>>( + jobs.data_ptr(), + num_jobs, + child_count.data_ptr(), + result_count.data_ptr(), + child_offsets.data_ptr(), + result_offsets.data_ptr(), + next_jobs.data_ptr(), + result_prim.data_ptr(), + result_voxels.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + if (result_count_h > 0) + { + total_results += result_count_h; + prim_chunks.push_back(result_prim); + voxel_chunks.push_back(result_voxels); + chunk_sizes.push_back(result_count_h); + } + + jobs = next_jobs; + num_jobs = next_count; + } + + if (total_results == 0) + return std::make_tuple(empty_prim, empty_voxels); + + auto prim_ids = torch::empty({total_results}, opts_i32); + auto voxels = torch::empty({total_results, 3}, opts_i32); + int64_t cursor = 0; + // Leaf hits are produced level by level, so collect chunks first and + // concatenate once the total result count is known. + for (size_t i = 0; i < chunk_sizes.size(); ++i) + { + const int64_t n = chunk_sizes[i]; + copy_chunk(prim_ids, cursor, prim_chunks[i], n, stream); + copy_chunk(voxels, cursor * 3, voxel_chunks[i], n * 3, stream); + cursor += n; + } + return std::make_tuple(prim_ids, voxels); + } + +} // namespace o_voxel::fdg diff --git a/o-voxel/src/convert/volumetic_attr.cpp b/o-voxel/src/convert/volumetic_attr.cpp index 64c3d098..f5a996d4 100644 --- a/o-voxel/src/convert/volumetic_attr.cpp +++ b/o-voxel/src/convert/volumetic_attr.cpp @@ -711,6 +711,12 @@ textured_mesh_to_volumetric_attr_cpu( const float mipLevelOffset, const bool timing ) { + TORCH_CHECK(!voxel_size.is_cuda(), "voxel_size must be a CPU tensor"); + TORCH_CHECK(!grid_range.is_cuda(), "grid_range must be a CPU tensor"); + TORCH_CHECK(!vertices.is_cuda(), "vertices must be a CPU tensor"); + TORCH_CHECK(!normals.is_cuda(), "normals must be a CPU tensor"); + TORCH_CHECK(!uvs.is_cuda(), "uvs must be a CPU tensor"); + TORCH_CHECK(!materialIds.is_cuda(), "materialIds must be a CPU tensor"); auto N_mat = baseColorFactor.size(); int N_tri = vertices.size(0); @@ -736,6 +742,14 @@ textured_mesh_to_volumetric_attr_cpu( std::vector H_nTex(N_mat), W_nTex(N_mat); for (int i = 0; i < N_mat; ++i) { + TORCH_CHECK(!baseColorFactor[i].is_cuda(), "baseColorFactor tensors must be CPU tensors"); + TORCH_CHECK(!baseColorTexture[i].is_cuda(), "baseColorTexture tensors must be CPU tensors"); + TORCH_CHECK(!metallicTexture[i].is_cuda(), "metallicTexture tensors must be CPU tensors"); + TORCH_CHECK(!roughnessTexture[i].is_cuda(), "roughnessTexture tensors must be CPU tensors"); + TORCH_CHECK(!emissiveFactor[i].is_cuda(), "emissiveFactor tensors must be CPU tensors"); + TORCH_CHECK(!emissiveTexture[i].is_cuda(), "emissiveTexture tensors must be CPU tensors"); + TORCH_CHECK(!alphaTexture[i].is_cuda(), "alphaTexture tensors must be CPU tensors"); + TORCH_CHECK(!normalTexture[i].is_cuda(), "normalTexture tensors must be CPU tensors"); baseColorFactor_ptrs[i] = baseColorFactor[i].contiguous().data_ptr(); if (baseColorTexture[i].numel() > 0) { baseColorTexture_ptrs[i] = baseColorTexture[i].contiguous().data_ptr(); @@ -869,4 +883,3 @@ textured_mesh_to_volumetric_attr_cpu( out_normals ); } - diff --git a/o-voxel/src/ext.cpp b/o-voxel/src/ext.cpp index e2ac946d..ed848632 100644 --- a/o-voxel/src/ext.cpp +++ b/o-voxel/src/ext.cpp @@ -1,12 +1,13 @@ #include +#include #include "hash/api.h" #include "convert/api.h" #include "io/api.h" #include "serialize/api.h" #include "rasterize/api.h" - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ // Hash functions m.def("hashmap_insert_cuda", &hashmap_insert_cuda); m.def("hashmap_lookup_cuda", &hashmap_lookup_cuda); @@ -15,16 +16,19 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("hashmap_insert_3d_idx_as_val_cuda", &hashmap_insert_3d_idx_as_val_cuda); // Convert functions m.def("mesh_to_flexible_dual_grid_cpu", &mesh_to_flexible_dual_grid_cpu, py::call_guard()); + m.def("mesh_to_flexible_dual_grid_cuda", &o_voxel::fdg::mesh_to_flexible_dual_grid_cuda, py::call_guard()); + m.def("intersect_occ_cpu", &intersect_occ_cpu, py::call_guard()); + m.def("intersect_occ_cuda", &o_voxel::fdg::intersect_occ_cuda, py::call_guard()); m.def("textured_mesh_to_volumetric_attr_cpu", &textured_mesh_to_volumetric_attr_cpu, py::call_guard()); // Serialization functions m.def("z_order_encode_cpu", &z_order_encode_cpu, py::call_guard()); - m.def("z_order_decode_cpu", &z_order_decode_cpu, py::call_guard()); - m.def("hilbert_encode_cpu", &hilbert_encode_cpu, py::call_guard()); - m.def("hilbert_decode_cpu", &hilbert_decode_cpu, py::call_guard()); + m.def("z_order_decode_cpu", &z_order_decode_cpu, py::call_guard()); + m.def("hilbert_encode_cpu", &hilbert_encode_cpu, py::call_guard()); + m.def("hilbert_decode_cpu", &hilbert_decode_cpu, py::call_guard()); m.def("z_order_encode_cuda", &z_order_encode_cuda, py::call_guard()); - m.def("z_order_decode_cuda", &z_order_decode_cuda, py::call_guard()); - m.def("hilbert_encode_cuda", &hilbert_encode_cuda, py::call_guard()); - m.def("hilbert_decode_cuda", &hilbert_decode_cuda, py::call_guard()); + m.def("z_order_decode_cuda", &z_order_decode_cuda, py::call_guard()); + m.def("hilbert_encode_cuda", &hilbert_encode_cuda, py::call_guard()); + m.def("hilbert_decode_cuda", &hilbert_decode_cuda, py::call_guard()); // IO functions m.def("encode_sparse_voxel_octree_cpu", &encode_sparse_voxel_octree_cpu, py::call_guard()); m.def("decode_sparse_voxel_octree_cpu", &decode_sparse_voxel_octree_cpu, py::call_guard()); @@ -34,4 +38,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("decode_sparse_voxel_octree_attr_neighbor_cpu", &decode_sparse_voxel_octree_attr_neighbor_cpu, py::call_guard()); // Rasterization functions m.def("rasterize_voxels_cuda", &rasterize_voxels_cuda); -} \ No newline at end of file +}