From ead7929f55914161e25a1703ff9d80c013350d5b Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Sat, 28 Mar 2026 12:16:42 -0700 Subject: [PATCH 1/3] Add differentiable material support for channel tracing --- witwin/core/__init__.py | 2 + witwin/core/material.py | 94 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+) diff --git a/witwin/core/__init__.py b/witwin/core/__init__.py index d137d4d..0a08391 100644 --- a/witwin/core/__init__.py +++ b/witwin/core/__init__.py @@ -3,6 +3,7 @@ __version__ = "0.0.1" from .material import ( + DifferentiableMaterial, FrequencyMaterialSample, Material, MaterialCapabilities, @@ -46,6 +47,7 @@ "Box", "Sphere", "Cylinder", "Cone", "Ellipsoid", "Pyramid", "Prism", "Torus", "HollowBox", "Material", + "DifferentiableMaterial", "MaterialCapabilities", "MaterialSpec", "Mesh", diff --git a/witwin/core/material.py b/witwin/core/material.py index 2be856c..2d7d99f 100644 --- a/witwin/core/material.py +++ b/witwin/core/material.py @@ -20,6 +20,30 @@ def _coerce_nonnegative(value: float, *, name: str) -> float: return scalar +def _coerce_differentiable_scalar(value: Any, *, name: str) -> Any: + try: + return float(value) + except (TypeError, ValueError): + return value + + +def _coerce_differentiable_nonnegative(value: Any, *, name: str) -> Any: + try: + scalar = float(value) + except (TypeError, ValueError): + return value + if scalar < 0.0: + raise ValueError(f"{name} must be >= 0.") + return scalar + + +def _is_numeric_close(value: Any, target: float) -> bool: + try: + return bool(np.isclose(float(value), target)) + except (TypeError, ValueError): + return False + + @dataclass(frozen=True) class MaterialCapabilities: conductive: bool = False @@ -114,6 +138,76 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: ) +@dataclass(frozen=True, init=False) +class DifferentiableMaterial(MaterialSpec): + eps_r: Any + mu_r: Any + sigma_e: Any + name: str | None + + def __init__( + self, + eps_r: Any = 1.0, + mu_r: Any = 1.0, + sigma_e: Any = 0.0, + name: str | None = None, + ): + object.__setattr__(self, "eps_r", _coerce_differentiable_scalar(eps_r, name="eps_r")) + object.__setattr__(self, "mu_r", _coerce_differentiable_scalar(mu_r, name="mu_r")) + object.__setattr__(self, "sigma_e", _coerce_differentiable_nonnegative(sigma_e, name="sigma_e")) + object.__setattr__(self, "name", None if name is None else str(name)) + + def capabilities(self) -> MaterialCapabilities: + return MaterialCapabilities( + conductive=not _is_numeric_close(self.sigma_e, 0.0), + magnetic=not _is_numeric_close(self.mu_r, 1.0), + anisotropic=False, + dispersive=False, + ) + + def evaluate_static(self) -> StaticMaterialSample: + return StaticMaterialSample( + eps_r=self.eps_r, + mu_r=self.mu_r, + sigma_e=self.sigma_e, + ) + + def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: + resolved_frequency = float(frequency) + if resolved_frequency < 0.0: + raise ValueError("frequency must be >= 0.") + if resolved_frequency == 0.0: + if not _is_numeric_close(self.sigma_e, 0.0): + try: + scalar_sigma = float(self.sigma_e) + except (TypeError, ValueError): + scalar_sigma = None + if scalar_sigma is not None: + raise ValueError("Conductive materials require frequency > 0.") + return FrequencyMaterialSample( + eps_r=self.eps_r, + mu_r=self.mu_r, + sigma_e=self.sigma_e, + ) + + try: + eps_r = complex( + float(self.eps_r), + -float(self.sigma_e) / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), + ) + mu_r = float(self.mu_r) + sigma_e = float(self.sigma_e) + except (TypeError, ValueError): + eps_r = self.eps_r + mu_r = self.mu_r + sigma_e = self.sigma_e + return FrequencyMaterialSample( + eps_r=eps_r, + mu_r=mu_r, + sigma_e=sigma_e, + ) + + @dataclass(frozen=True, init=False) class Structure: geometry: Any From 24a8938dd38280f6a841b43db23d2fda1aed913f Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Wed, 8 Jul 2026 05:10:31 -0700 Subject: [PATCH 2/3] feat: PolySlab / ComplexPolySlab extruded-polygon primitives --- tests/test_geometry_sdf.py | 65 +++++++ tests/test_geometry_smoke.py | 52 ++++++ witwin/core/geometry/__init__.py | 6 +- witwin/core/geometry/primitives.py | 281 ++++++++++++++++++++++++++++- 4 files changed, 398 insertions(+), 6 deletions(-) diff --git a/tests/test_geometry_sdf.py b/tests/test_geometry_sdf.py index 8eeaf7f..d61f1a1 100644 --- a/tests/test_geometry_sdf.py +++ b/tests/test_geometry_sdf.py @@ -4,6 +4,7 @@ import torch from witwin.core import Box, Cone, Cylinder, Ellipsoid, HollowBox, Mesh, Prism, Pyramid, Sphere, Torus +from witwin.core.geometry import ComplexPolySlab, PolySlab from witwin.core.geometry.mesh_sdf import ( _slang_mesh_sdf_available, _triangle_mesh_smooth_signed_distance_torch, @@ -151,6 +152,39 @@ def _require_cuda_mesh_sdf(): (0.5, 0.0, 0.0), (0.7, 0.0, 0.0), ), + ( + PolySlab( + vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.5), (-0.5, 0.5)], + bounds=(-0.5, 0.5), + axis="z", + ), + (0.0, 0.0, 0.0), + (0.5, 0.0, 0.0), + (0.7, 0.0, 0.0), + ), + ( + # Non-convex L-shape: the notch quadrant must be outside. + PolySlab( + vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.0), (0.0, 0.0), (0.0, 0.5), (-0.5, 0.5)], + bounds=(-0.5, 0.5), + axis="z", + ), + (-0.25, -0.25, 0.0), + (0.5, -0.25, 0.0), + (0.25, 0.25, 0.0), + ), + ( + ComplexPolySlab( + loops=[ + [(-0.5, -0.5), (0.5, -0.5), (0.5, 0.5), (-0.5, 0.5)], + [(-0.2, -0.2), (0.2, -0.2), (0.2, 0.2), (-0.2, 0.2)], + ], + bounds=(-0.5, 0.5), + ), + (0.35, 0.0, 0.0), + (0.5, 0.0, 0.0), + (0.0, 0.0, 0.0), + ), ], ) def test_primitive_signed_distance_has_expected_sign(geometry, inside_point, surface_point, outside_point): @@ -180,6 +214,16 @@ def test_primitive_signed_distance_has_expected_sign(geometry, inside_point, sur (Cone(radius=0.4, height=1.0, axis="z"), (0.0, 0.0, 0.5), (0.2, 0.0, 0.5), (0.3, 0.0, 0.5)), (Pyramid(base_size=1.0, height=1.0, axis="z"), (0.0, 0.0, 0.5), (0.25, 0.0, 0.5), (0.35, 0.0, 0.5)), (Prism(radius=0.4, height=1.0, num_sides=5, axis="z"), (0.0, 0.0, 0.0), (0.4, 0.0, 0.0), (0.6, 0.0, 0.0)), + ( + PolySlab( + vertices=[(-0.4, -0.4), (0.4, -0.4), (0.4, 0.4), (-0.4, 0.4)], + bounds=(-0.5, 0.5), + axis="z", + ), + (0.0, 0.0, 0.0), + (0.4, 0.0, 0.0), + (0.7, 0.0, 0.0), + ), ], ) def test_primitive_occupancy_is_monotone_and_bounded(geometry, inside_point, boundary_point, outside_point): @@ -327,6 +371,27 @@ def test_prism_occupancy_has_gradients_for_position_radius_height_and_rotation() assert abs(grad.item()) > 1.0e-6 +def test_polyslab_occupancy_has_gradients_for_vertices_bounds_and_sidewall_angle(): + xx, yy, zz = _asymmetric_grid() + vertices = torch.nn.Parameter( + torch.tensor([[-0.31, -0.24], [0.29, -0.24], [0.29, 0.26], [-0.31, 0.26]], dtype=torch.float32) + ) + bounds = torch.nn.Parameter(torch.tensor([-0.28, 0.17], dtype=torch.float32)) + sidewall_angle = torch.nn.Parameter(torch.tensor(0.10, dtype=torch.float32)) + + geometry = PolySlab(vertices, bounds, axis="z", sidewall_angle=sidewall_angle) + loss = geometry.to_mask(xx, yy, zz, beta=0.06).sum() + loss.backward() + + for grad in (vertices.grad, bounds.grad): + assert grad is not None + assert torch.all(torch.isfinite(grad)) + assert torch.any(grad.abs() > 1.0e-6) + assert sidewall_angle.grad is not None + assert torch.isfinite(sidewall_angle.grad) + assert abs(sidewall_angle.grad.item()) > 1.0e-6 + + def test_mesh_signed_distance_matches_box_sign_for_closed_cube(): vertices, faces = _cube_mesh(size=1.0) mesh = Mesh(vertices, faces, recenter=False, fill_mode="solid") diff --git a/tests/test_geometry_smoke.py b/tests/test_geometry_smoke.py index 38b7945..d694e8d 100644 --- a/tests/test_geometry_smoke.py +++ b/tests/test_geometry_smoke.py @@ -4,6 +4,7 @@ import torch from witwin.core import Box, Cylinder, Mesh, Sphere +from witwin.core.geometry import PolySlab def _grid(): @@ -36,6 +37,17 @@ def _grid(): (0.4, 0.2, -0.1), 64, ), + ( + PolySlab( + vertices=[(-0.3, -0.3), (0.3, -0.3), (0.3, 0.3), (-0.3, 0.3)], + bounds=(-0.35, 0.35), + position=(0.1, -0.2, 0.3), + ), + 16, + (0.1, -0.2, 0.3), + (0.8, -0.2, 0.3), + 12, + ), ], ) def test_geometry_construction_to_mesh_and_to_mask(geometry, segments, inside_point, outside_point, expected_faces): @@ -71,6 +83,46 @@ def test_geometry_construction_to_mesh_and_to_mask(geometry, segments, inside_po assert torch.any(occupancy < 0.5) +def test_polyslab_nonconvex_to_mesh_produces_expected_topology(): + slab = PolySlab( + vertices=[(-0.5, -0.5), (0.5, -0.5), (0.5, 0.0), (0.0, 0.0), (0.0, 0.5), (-0.5, 0.5)], + bounds=(-0.25, 0.25), + sidewall_angle=0.1, + ) + vertices, faces = slab.to_mesh() + + # 6 polygon vertices: 12 side triangles + 2 * (6 - 2) ear-clipped cap triangles. + assert vertices.shape == (12, 3) + assert faces.shape == (20, 3) + assert faces.dtype == torch.int64 + + +def test_polyslab_to_mesh_matches_signed_distance_for_all_axes(): + # Asymmetric polygon so a u/v mirror in the axis remap cannot go unnoticed. + vertices_2d = [(0.1, 0.0), (0.9, 0.0), (0.1, 0.6)] + for axis in ("x", "y", "z"): + slab = PolySlab(vertices=vertices_2d, bounds=(-0.5, 0.5), axis=axis) + mesh_vertices, faces = slab.to_mesh() + + distance = slab.signed_distance(mesh_vertices[:, 0], mesh_vertices[:, 1], mesh_vertices[:, 2]) + assert torch.all(distance.abs() < 1.0e-5), axis + + # Positive divergence-theorem volume implies outward face normals. + tri = mesh_vertices[faces] + volume = torch.sum(torch.sum(tri[:, 0] * torch.linalg.cross(tri[:, 1], tri[:, 2], dim=1), dim=1)) / 6.0 + assert volume.item() == pytest.approx(0.5 * 0.8 * 0.6, rel=1.0e-4), axis + + +def test_polyslab_to_mesh_accepts_closed_ring_vertices(): + slab = PolySlab( + vertices=[(-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0), (-1.0, -1.0)], + bounds=(-0.5, 0.5), + ) + vertices, faces = slab.to_mesh() + assert vertices.shape == (8, 3) + assert faces.shape == (12, 3) + + def test_mesh_roundtrip_preserves_world_vertices_and_faces(): base = Box(position=(0.25, -0.15, 0.4), size=(0.5, 0.3, 0.7)) vertices, faces = base.to_mesh() diff --git a/witwin/core/geometry/__init__.py b/witwin/core/geometry/__init__.py index c8a052d..837b5de 100644 --- a/witwin/core/geometry/__init__.py +++ b/witwin/core/geometry/__init__.py @@ -2,10 +2,10 @@ from .base import GeometryBase from .mesh import Mesh -from .primitives import Box, Cone, Cylinder, Ellipsoid, HollowBox, Prism, Pyramid, Sphere, Torus +from .primitives import Box, ComplexPolySlab, Cone, Cylinder, Ellipsoid, HollowBox, PolySlab, Prism, Pyramid, Sphere, Torus from .smpl import SMPLBody -Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | Torus | HollowBox | Mesh | SMPLBody +Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | PolySlab | ComplexPolySlab | Torus | HollowBox | Mesh | SMPLBody __all__ = [ "GeometryBase", @@ -16,6 +16,8 @@ "Ellipsoid", "Pyramid", "Prism", + "PolySlab", + "ComplexPolySlab", "Torus", "HollowBox", "Mesh", diff --git a/witwin/core/geometry/primitives.py b/witwin/core/geometry/primitives.py index 836640c..782a30d 100644 --- a/witwin/core/geometry/primitives.py +++ b/witwin/core/geometry/primitives.py @@ -41,14 +41,32 @@ def _safe_signed_denom(value: torch.Tensor, eps: float) -> torch.Tensor: return torch.where(value >= 0.0, value.clamp_min(eps), value.clamp_max(-eps)) +def _polygon_edge_distance_sq_2d( + points: torch.Tensor, + polygon: torch.Tensor, + *, + edges: torch.Tensor | None = None, + rel: torch.Tensor | None = None, +) -> torch.Tensor: + """Squared unsigned distance from (P, 2) points to the closed edge loop of an (M, 2) polygon. + + ``edges`` and ``rel`` may be passed in when the caller has already materialized them. + """ + if edges is None: + edges = torch.roll(polygon, shifts=-1, dims=0) - polygon + if rel is None: + rel = points[:, None, :] - polygon[None, :, :] + edge_norm_sq = torch.sum(edges ** 2, dim=-1).clamp_min(torch.finfo(points.dtype).eps) + t = torch.clamp(torch.sum(rel * edges[None, :, :], dim=-1) / edge_norm_sq[None, :], 0.0, 1.0) + closest = polygon[None, :, :] + t[..., None] * edges[None, :, :] + return torch.sum((points[:, None, :] - closest) ** 2, dim=-1).min(dim=1).values.clamp_min(0.0) + + def _convex_polygon_signed_distance_2d(points: torch.Tensor, polygon: torch.Tensor) -> torch.Tensor: next_polygon = torch.roll(polygon, shifts=-1, dims=0) edges = next_polygon - polygon rel = points[:, None, :] - polygon[None, :, :] - edge_norm_sq = torch.sum(edges ** 2, dim=-1).clamp_min(torch.finfo(points.dtype).eps) - t = torch.clamp(torch.sum(rel * edges[None, :, :], dim=-1) / edge_norm_sq[None, :], 0.0, 1.0) - closest = polygon[None, :, :] + t[..., None] * edges[None, :, :] - dist = torch.sqrt(torch.sum((points[:, None, :] - closest) ** 2, dim=-1).min(dim=1).values.clamp_min(0.0)) + dist = torch.sqrt(_polygon_edge_distance_sq_2d(points, polygon, edges=edges, rel=rel)) polygon_area = 0.5 * torch.sum(_cross_2d(polygon, next_polygon)) crosses = _cross_2d(edges[None, :, :], rel) @@ -59,6 +77,34 @@ def _convex_polygon_signed_distance_2d(points: torch.Tensor, polygon: torch.Tens return torch.where(inside, -dist, dist) +def _polygon_crossing_count_2d(points: torch.Tensor, polygon: torch.Tensor) -> torch.Tensor: + """Count +u horizontal-ray crossings of a closed (M, 2) polygon loop for each (P, 2) point.""" + start = polygon + end = torch.roll(polygon, shifts=-1, dims=0) + pu = points[:, 0:1] + pv = points[:, 1:2] + straddles = (start[None, :, 1] <= pv) != (end[None, :, 1] <= pv) + denom = _safe_signed_denom(end[:, 1] - start[:, 1], torch.finfo(points.dtype).eps) + t = (pv - start[None, :, 1]) / denom[None, :] + crossing_u = start[None, :, 0] + t * (end[:, 0] - start[:, 0])[None, :] + return (straddles & (pu < crossing_u)).sum(dim=1) + + +def _polygon_loops_signed_distance_2d(points: torch.Tensor, loops) -> torch.Tensor: + """Even-odd signed distance from (P, 2) points to one or more closed 2D polygon loops. + + Negative inside; the interior is defined by the even-odd rule over the combined + crossing count of all loops, so winding direction is irrelevant and additional + loops carve holes. Valid for arbitrary (including non-convex and self-intersecting) + polygons. + """ + distance_sq = torch.stack([_polygon_edge_distance_sq_2d(points, loop) for loop in loops], dim=0) + distance = torch.sqrt(distance_sq.min(dim=0).values) + crossings = sum(_polygon_crossing_count_2d(points, loop) for loop in loops) + inside = crossings % 2 == 1 + return torch.where(inside, -distance, distance) + + def _canonical_coords(dx, dy, dz, axis: str): if axis == "z": return dx, dy, dz @@ -509,3 +555,230 @@ def to_mesh(self, segments=16): torch.cat([outer_vertices, inner_vertices], dim=0), torch.cat([outer_faces, inner_faces + outer_vertices.shape[0]], dim=0), ) + + +def _as_polygon_vertices(value, *, device=None) -> torch.Tensor: + if isinstance(value, torch.Tensor): + vertices = value.to(device=device, dtype=torch.float32) + else: + vertices = torch.tensor([[float(u), float(v)] for u, v in value], dtype=torch.float32, device=device) + if vertices.ndim != 2 or vertices.shape[0] < 3 or vertices.shape[1] != 2: + raise ValueError("vertices must have shape (N, 2) with N >= 3.") + return vertices + + +def _as_axis_bounds(value, *, device=None) -> torch.Tensor: + if isinstance(value, torch.Tensor): + bounds = value.to(device=device, dtype=torch.float32) + else: + bounds = torch.tensor([float(value[0]), float(value[1])], dtype=torch.float32, device=device) + if bounds.shape != (2,): + raise ValueError("bounds must contain (min, max) along the extrusion axis.") + return bounds + + +def _cross_2d_scalar(origin, first, second) -> float: + return float((first[0] - origin[0]) * (second[1] - origin[1]) - (first[1] - origin[1]) * (second[0] - origin[0])) + + +def _point_in_triangle_2d(point, a, b, c, eps=1.0e-12) -> bool: + """Inclusive point-in-triangle test for a CCW triangle (a, b, c).""" + return ( + _cross_2d_scalar(a, b, point) >= -eps + and _cross_2d_scalar(b, c, point) >= -eps + and _cross_2d_scalar(c, a, point) >= -eps + ) + + +def _ear_clip_faces_2d(polygon: np.ndarray) -> list[list[int]]: + """Triangulate a simple CCW (N, 2) polygon into CCW index triangles via ear clipping.""" + remaining = list(range(polygon.shape[0])) + faces: list[list[int]] = [] + while len(remaining) > 3: + clipped = False + for slot in range(len(remaining)): + i0 = remaining[slot - 1] + i1 = remaining[slot] + i2 = remaining[(slot + 1) % len(remaining)] + a, b, c = polygon[i0], polygon[i1], polygon[i2] + if _cross_2d_scalar(a, b, c) <= 1.0e-12: + continue + others = (polygon[j] for j in remaining if j not in (i0, i1, i2)) + if any(_point_in_triangle_2d(p, a, b, c) for p in others): + continue + faces.append([i0, i1, i2]) + remaining.pop(slot) + clipped = True + break + if not clipped: + raise ValueError("Failed to triangulate polygon; vertices must describe a simple polygon.") + faces.append(list(remaining)) + return faces + + +def _offset_polygon_2d(polygon: np.ndarray, offset: float) -> np.ndarray: + """Approximate outward miter offset of a CCW (N, 2) polygon (negative offset shrinks).""" + if offset == 0.0: + return polygon + previous_edges = polygon - np.roll(polygon, 1, axis=0) + next_edges = np.roll(polygon, -1, axis=0) - polygon + + def _outward_normals(edges: np.ndarray) -> np.ndarray: + normals = np.stack([edges[:, 1], -edges[:, 0]], axis=1) + return normals / np.maximum(np.linalg.norm(normals, axis=1, keepdims=True), 1.0e-12) + + normal_previous = _outward_normals(previous_edges) + normal_next = _outward_normals(next_edges) + denom = np.maximum(1.0 + np.sum(normal_previous * normal_next, axis=1, keepdims=True), 1.0e-3) + return polygon + offset * (normal_previous + normal_next) / denom + + +class PolySlab(GeometryBase): + """Extruded 2D polygon with optional linear sidewall taper (Tidy3D semantics). + + ``vertices`` is the (N, 2) cross-section in the plane perpendicular to ``axis`` + (plane coordinates follow ``_axial_split``: axis z -> (x, y), y -> (x, z), + x -> (y, z)), ``bounds`` is (min, max) along the axis, and a positive + ``sidewall_angle`` (radians) shrinks the cross-section as the axis coordinate + increases. ``reference_plane`` selects the axis position where the drawn + polygon is exact. The interior uses the even-odd rule, so non-convex polygons + are supported and vertex winding is irrelevant. + """ + + kind = "poly_slab" + + def __init__( + self, + vertices, + bounds, + axis="z", + sidewall_angle=0.0, + reference_plane="middle", + position=(0, 0, 0), + rotation=None, + *, + device=None, + ): + super().__init__(position=position, rotation=rotation, device=device) + self.vertices: torch.Tensor = _as_polygon_vertices(vertices, device=device) + self.bounds: torch.Tensor = _as_axis_bounds(bounds, device=device) + if isinstance(axis, int) and not isinstance(axis, bool): + if axis not in (0, 1, 2): + raise ValueError("integer axis must be 0, 1, or 2.") + axis = ("x", "y", "z")[axis] + self.axis: str = self._validate_axis(axis) + self.sidewall_angle: torch.Tensor = _as_scalar(sidewall_angle, device=device) + if reference_plane not in ("bottom", "middle", "top"): + raise ValueError("reference_plane must be 'bottom', 'middle', or 'top'.") + self.reference_plane: str = reference_plane + self.loop_sizes: tuple[int, ...] = (int(self.vertices.shape[0]),) + + def _reference_position(self, bounds: torch.Tensor) -> torch.Tensor: + if self.reference_plane == "bottom": + return bounds[0] + if self.reference_plane == "top": + return bounds[1] + return 0.5 * (bounds[0] + bounds[1]) + + def signed_distance(self, xx, yy, zz): + dx, dy, dz = self._local_coords(xx, yy, zz) + axial, pu, pv = _axial_split(dx, dy, dz, self.axis) + vertices = self.vertices.to(device=pu.device, dtype=pu.dtype) + loops = torch.split(vertices, list(self.loop_sizes), dim=0) + points = torch.stack([pu, pv], dim=-1).reshape(-1, 2) + polygon_distance = _polygon_loops_signed_distance_2d(points, loops).reshape(pu.shape) + bounds = self.bounds.to(device=axial.device, dtype=axial.dtype) + tan_angle = torch.tan(self.sidewall_angle.to(device=axial.device, dtype=axial.dtype)) + # Signed distance minus an offset is exact morphological erosion/dilation, so + # the taper is a z-dependent inward offset of the cross-section. + tapered_distance = polygon_distance + tan_angle * (axial - self._reference_position(bounds)) + axial_distance = torch.abs(axial - 0.5 * (bounds[0] + bounds[1])) - 0.5 * (bounds[1] - bounds[0]) + outside = torch.sqrt(torch.clamp(tapered_distance, min=0.0) ** 2 + torch.clamp(axial_distance, min=0.0) ** 2) + inside = torch.clamp(torch.maximum(tapered_distance, axial_distance), max=0.0) + return outside + inside + + def to_mesh(self, segments=16): + del segments + device = self.device + polygon = self.vertices.detach().cpu().to(torch.float64).numpy() + # Drop consecutive duplicate vertices (including a closed-ring wrap-around + # repeat); coincident points would block every candidate ear below. + keep = np.linalg.norm(polygon - np.roll(polygon, 1, axis=0), axis=1) > 1.0e-12 + polygon = polygon[keep] + if polygon.shape[0] < 3: + raise ValueError("vertices must contain at least 3 distinct points.") + signed_area = 0.5 * float(np.sum(polygon[:, 0] * np.roll(polygon[:, 1], -1) - np.roll(polygon[:, 0], -1) * polygon[:, 1])) + if signed_area < 0.0: + polygon = polygon[::-1].copy() + bounds = self.bounds.detach().cpu().to(torch.float64).numpy() + reference = {"bottom": bounds[0], "top": bounds[1], "middle": 0.5 * (bounds[0] + bounds[1])}[self.reference_plane] + tan_angle = float(np.tan(self.sidewall_angle.detach().cpu().to(torch.float64).item())) + bottom_ring = _offset_polygon_2d(polygon, -tan_angle * (bounds[0] - reference)) + top_ring = _offset_polygon_2d(polygon, -tan_angle * (bounds[1] - reference)) + count = polygon.shape[0] + bottom = np.concatenate([bottom_ring, np.full((count, 1), bounds[0])], axis=1) + top = np.concatenate([top_ring, np.full((count, 1), bounds[1])], axis=1) + vertices = _constant_tensor(np.concatenate([bottom, top], axis=0), device=device) + faces = [] + for index in range(count): + next_index = (index + 1) % count + faces.extend([ + [index, next_index, count + next_index], + [index, count + next_index, count + index], + ]) + for i0, i1, i2 in _ear_clip_faces_2d(polygon): + faces.append([i0, i2, i1]) + faces.append([count + i0, count + i1, count + i2]) + # Map local (u, v, axial) columns to world axes following _axial_split + # (axis z -> (x, y), y -> (x, z), x -> (y, z)); _remap_axis_torch would + # mirror u/v for axis "x" and disagree with signed_distance/to_mask. + vertices = vertices[:, {"z": [0, 1, 2], "y": [0, 2, 1], "x": [2, 0, 1]}[self.axis]] + faces_tensor = _faces_tensor(faces, device=device) + if self.axis == "y": + # Odd column permutation reflects orientation; flip winding to keep + # face normals outward. + faces_tensor = torch.flip(faces_tensor, dims=(1,)) + vertices = self._transform_mesh_verts(vertices) + return vertices, faces_tensor + + +class ComplexPolySlab(PolySlab): + """PolySlab variant for non-simple cross-sections. + + Accepts a list of vertex loops; self-intersecting loops (e.g. bowties) and + multiple loops are allowed. The interior is the even-odd rule over the combined + crossing count of all loops, so loop winding is irrelevant and extra loops carve + holes. The distance magnitude is the minimum distance over all edges of all loops. + """ + + kind = "complex_poly_slab" + + def __init__( + self, + loops, + bounds, + axis="z", + sidewall_angle=0.0, + reference_plane="middle", + position=(0, 0, 0), + rotation=None, + *, + device=None, + ): + loop_tensors = [_as_polygon_vertices(loop, device=device) for loop in loops] + if not loop_tensors: + raise ValueError("loops must contain at least one polygon loop.") + super().__init__( + torch.cat(loop_tensors, dim=0), + bounds, + axis=axis, + sidewall_angle=sidewall_angle, + reference_plane=reference_plane, + position=position, + rotation=rotation, + device=device, + ) + self.loop_sizes = tuple(int(loop.shape[0]) for loop in loop_tensors) + + def to_mesh(self, segments=16): + raise NotImplementedError("ComplexPolySlab does not support mesh export.") From 39a3e0bb06d6a31c6f8a57f1cd68b68d48883c1a Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Thu, 9 Jul 2026 00:32:54 -0700 Subject: [PATCH 3/3] feat: native CUDA mesh SDF backend with prebuilt CUDA wheels Replace the SlangTorch mesh_sdf module with a native CUDA extension (torch.utils.cpp_extension), mirroring the maxwell FDTD CUDA backend. This removes the intermittent SlangTorch load failures on non-UTF-8 Windows consoles and ships prebuilt per-platform wheels so installs never invoke nvcc/MSVC. - geometry/cuda/: six mesh SDF kernels (brute/BVH unsigned distance, distance+winding, BVH ray-parity sign, and hand-derived backward gradients) plus a pybind extension and a JIT/packaged-prebuilt loader. - mesh_sdf.py: load the native module instead of Slang; keep the Torch-native fallback when CUDA/toolchain is unavailable. - Packaging: hatch build hook marks the wheel platform-specific and ships prebuilt/*.pyd plus the CUDA sources; publish workflow builds a CUDA wheel matrix (ubuntu-22.04 + windows-2022, py3.10-3.12) and sdist. - Bump version to 0.0.2. --- .github/workflows/publish-witwin.yml | 219 +++++- .gitignore | 4 + hatch_build.py | 17 + pyproject.toml | 13 +- scripts/build_mesh_sdf_cuda_prebuilt.py | 68 ++ tests/test_geometry_sdf.py | 8 +- witwin/core/geometry/cuda/__init__.py | 9 + witwin/core/geometry/cuda/backend.py | 177 +++++ witwin/core/geometry/cuda/build.py | 247 +++++++ witwin/core/geometry/cuda/extension.cpp | 43 ++ witwin/core/geometry/cuda/mesh_sdf_kernels.cu | 688 ++++++++++++++++++ witwin/core/geometry/mesh_sdf.py | 122 ++-- witwin/core/mesh_sdf.slang | 609 ---------------- 13 files changed, 1537 insertions(+), 687 deletions(-) create mode 100644 hatch_build.py create mode 100644 scripts/build_mesh_sdf_cuda_prebuilt.py create mode 100644 witwin/core/geometry/cuda/__init__.py create mode 100644 witwin/core/geometry/cuda/backend.py create mode 100644 witwin/core/geometry/cuda/build.py create mode 100644 witwin/core/geometry/cuda/extension.cpp create mode 100644 witwin/core/geometry/cuda/mesh_sdf_kernels.cu delete mode 100644 witwin/core/mesh_sdf.slang diff --git a/.github/workflows/publish-witwin.yml b/.github/workflows/publish-witwin.yml index 690e2ca..8186784 100644 --- a/.github/workflows/publish-witwin.yml +++ b/.github/workflows/publish-witwin.yml @@ -1,19 +1,22 @@ -name: Publish witwin +name: Build and publish witwin on: + push: + branches: + - main + - "codex/**" + pull_request: release: types: [published] + workflow_dispatch: jobs: - publish: - if: startsWith(github.event.release.tag_name, 'witwin-v') && !github.event.release.prerelease + validate_release: + if: github.event_name != 'release' || (startsWith(github.event.release.tag_name, 'witwin-v') && !github.event.release.prerelease) + name: Validate release runs-on: ubuntu-latest - permissions: - contents: read - id-token: write - environment: - name: pypi - url: https://pypi.org/project/witwin/ + outputs: + package_dir: ${{ steps.package_dir.outputs.path }} steps: - name: Check out repository @@ -38,6 +41,7 @@ jobs: fi - name: Validate release tag and package version + if: github.event_name == 'release' shell: bash env: PACKAGE_DIR: ${{ steps.package_dir.outputs.path }} @@ -78,15 +82,206 @@ jobs: print(f"Validated {package_name} {package_version} from {pyproject_path}.") PY + build_cuda_wheels: + name: Build CUDA wheel / ${{ matrix.os }} / py${{ matrix.python_version }} + needs: validate_release + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-2022] + python_version: ["3.10", "3.11", "3.12"] + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python_version }} + + - name: Free disk space (Linux) + if: runner.os == 'Linux' + shell: bash + run: | + echo "Disk before cleanup:"; df -h / + sudo rm -rf /usr/share/dotnet /opt/ghc /usr/local/lib/android \ + /usr/local/share/boost /opt/hostedtoolcache/CodeQL \ + /usr/local/share/powershell /usr/share/swift || true + sudo docker image prune --all --force || true + echo "Disk after cleanup:"; df -h / + + - name: Install CUDA Toolkit + uses: Jimver/cuda-toolkit@v0.2.29 + with: + cuda: "12.5.0" + method: local + log-file-suffix: "${{ runner.os }}-py${{ matrix.python_version }}" + + - name: Set up MSVC + if: runner.os == 'Windows' + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + + - name: Export Windows CUDA_HOME + if: runner.os == 'Windows' + shell: pwsh + run: | + $cudaRoot = $env:CUDA_PATH + if (-not $cudaRoot) { + $nvcc = (Get-Command nvcc -ErrorAction Stop).Source + $cudaRoot = Split-Path (Split-Path $nvcc -Parent) -Parent + } + "CUDA_HOME=$cudaRoot" >> $env:GITHUB_ENV + "CUDA_PATH=$cudaRoot" >> $env:GITHUB_ENV + "CUDA_ROOT=$cudaRoot" >> $env:GITHUB_ENV + "CUDA_BIN_PATH=$cudaRoot" >> $env:GITHUB_ENV + + - name: Show toolchain + shell: bash + run: | + python --version + nvcc --version + command -v cl || true + + - name: Install CUDA extension build dependencies + run: | + python -m pip install --upgrade pip setuptools wheel ninja numpy + python -m pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu124 + python -m pip install build + + - name: Build packaged mesh SDF CUDA extension + shell: bash + env: + MAX_JOBS: "2" + TORCH_CUDA_ARCH_LIST: "7.0;8.0;8.6;8.9;9.0+PTX" + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/witwin_core_mesh_sdf_cuda + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + run: | + python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose + + - name: Build platform wheel + shell: bash + env: + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + run: | + python -m build --wheel "$PACKAGE_DIR" + + - name: Normalize Linux wheel platform tag + if: runner.os == 'Linux' + shell: bash + env: + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + run: | + # PyPI rejects the raw linux_x86_64 tag; the extension is built on the + # ubuntu-22.04 runner (glibc 2.35), so retag to the matching manylinux + # tag. torch/CUDA runtime libraries are runtime dependencies and stay + # unbundled, so no auditwheel repair is required. + shopt -s nullglob + for whl in "$PACKAGE_DIR"/dist/*-linux_x86_64.whl; do + echo "Retagging $whl" + python -m wheel tags --platform-tag manylinux_2_35_x86_64 --remove "$whl" + done + ls -l "$PACKAGE_DIR"/dist + + - name: Smoke install prebuilt wheel + shell: bash + env: + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + WITWIN_CORE_MESH_SDF_CUDA_PREBUILT: "1" + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/missing_prebuilt_build_dir + run: | + python -m pip uninstall -y witwin || true + python -m pip install --no-deps "$PACKAGE_DIR"/dist/*.whl + python - <<'PY' + import importlib.util + import site + from pathlib import Path + + candidates = [] + for root in site.getsitepackages() + [site.getusersitepackages()]: + candidates.append(Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py") + build_path = next((path for path in candidates if path.exists()), None) + assert build_path is not None, f"installed build.py not found in {candidates}" + + spec = importlib.util.spec_from_file_location("witwin_core_mesh_sdf_cuda_build", build_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + path = module.prebuilt_extension_path() + assert path.exists(), f"missing packaged prebuilt extension: {path}" + extension = module.build_extension() + print(f"Loaded packaged mesh SDF CUDA extension from {extension.__file__}") + PY + + - name: Upload wheel artifact + uses: actions/upload-artifact@v4 + with: + name: wheel-${{ matrix.os }}-py${{ matrix.python_version }} + path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.whl + if-no-files-found: error + + build_sdist: + name: Build source distribution + needs: validate_release + runs-on: ubuntu-latest + + steps: + - name: Check out repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + - name: Build distributions shell: bash env: - PACKAGE_DIR: ${{ steps.package_dir.outputs.path }} + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} run: | python -m pip install --upgrade pip build - python -m build "$PACKAGE_DIR" + python -m build --sdist "$PACKAGE_DIR" + + - name: Upload source distribution artifact + uses: actions/upload-artifact@v4 + with: + name: sdist + path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.tar.gz + if-no-files-found: error + + publish: + name: Publish to PyPI + if: github.event_name == 'release' + needs: + - build_cuda_wheels + - build_sdist + runs-on: ubuntu-latest + permissions: + contents: read + id-token: write + environment: + name: pypi + url: https://pypi.org/project/witwin/ + + steps: + - name: Download wheel artifacts + uses: actions/download-artifact@v4 + with: + pattern: wheel-* + path: dist + merge-multiple: true + + - name: Download source distribution artifact + uses: actions/download-artifact@v4 + with: + name: sdist + path: dist - name: Publish distributions to PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: - packages-dir: ${{ steps.package_dir.outputs.path }}/dist + packages-dir: dist + skip-existing: true diff --git a/.gitignore b/.gitignore index 0d710ff..618b9e1 100644 --- a/.gitignore +++ b/.gitignore @@ -228,3 +228,7 @@ __marimo__/ # Streamlit .streamlit/secrets.toml + +# Native CUDA mesh SDF prebuilt extensions (built per-platform in CI) +witwin/core/geometry/cuda/prebuilt/*.pyd +witwin/core/geometry/cuda/prebuilt/*.so diff --git a/hatch_build.py b/hatch_build.py new file mode 100644 index 0000000..b204e2f --- /dev/null +++ b/hatch_build.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from pathlib import Path + +from hatchling.builders.hooks.plugin.interface import BuildHookInterface + + +class CustomBuildHook(BuildHookInterface): + def initialize(self, version: str, build_data: dict) -> None: + del version + if self.target_name != "wheel": + return + prebuilt_dir = Path(self.root) / "witwin" / "core" / "geometry" / "cuda" / "prebuilt" + has_prebuilt_extension = any(prebuilt_dir.glob("witwin_core_mesh_sdf_cuda.*")) + if has_prebuilt_extension: + build_data["infer_tag"] = True + build_data["pure_python"] = False diff --git a/pyproject.toml b/pyproject.toml index 2379395..186fe70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "witwin" -version = "0.0.1" +version = "0.0.2" description = "Witwin - Multi-physics simulation platform" readme = "README.md" requires-python = ">=3.10" @@ -29,3 +29,14 @@ all = [ [tool.hatch.build.targets.wheel] packages = ["witwin"] +artifacts = [ + "witwin/core/geometry/cuda/**/*.cpp", + "witwin/core/geometry/cuda/**/*.cu", + "witwin/core/geometry/cuda/**/*.cuh", + "witwin/core/geometry/cuda/**/*.h", + "witwin/core/geometry/cuda/prebuilt/*.pyd", + "witwin/core/geometry/cuda/prebuilt/*.so", +] + +[tool.hatch.build.hooks.custom] +path = "hatch_build.py" diff --git a/scripts/build_mesh_sdf_cuda_prebuilt.py b/scripts/build_mesh_sdf_cuda_prebuilt.py new file mode 100644 index 0000000..86e3e43 --- /dev/null +++ b/scripts/build_mesh_sdf_cuda_prebuilt.py @@ -0,0 +1,68 @@ +from __future__ import annotations + +import argparse +import importlib.util +import os +import shutil +import tempfile +from pathlib import Path + +import torch + + +def _ensure_current_device_arch() -> None: + if not torch.cuda.is_available(): + return + major, minor = torch.cuda.get_device_capability() + current_arch = f"{major}.{minor}" + arch_list = os.environ.get("TORCH_CUDA_ARCH_LIST", "") + entries = [entry.strip() for entry in arch_list.split(";") if entry.strip()] + normalized = {entry.removesuffix("+PTX") for entry in entries} + if current_arch not in normalized: + entries.append(current_arch) + os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(entries) + + +def _load_cuda_build_module(): + repo_root = Path(__file__).resolve().parents[1] + build_path = repo_root / "witwin" / "core" / "geometry" / "cuda" / "build.py" + spec = importlib.util.spec_from_file_location("witwin_core_mesh_sdf_cuda_build", build_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to load mesh SDF CUDA build module from {build_path}.") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def main() -> None: + parser = argparse.ArgumentParser(description="Build the packaged mesh SDF CUDA extension.") + parser.add_argument("--verbose", action="store_true") + args = parser.parse_args() + + build_dir = Path( + os.environ.get( + "WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", + Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel", + ) + ) + os.environ["WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR"] = str(build_dir) + os.environ["WITWIN_CORE_MESH_SDF_CUDA_SKIP_PREBUILT"] = "1" + _ensure_current_device_arch() + + cuda_build = _load_cuda_build_module() + module = cuda_build.build_extension(verbose=args.verbose) + module_file = Path(module.__file__).resolve() + + target_dir = cuda_build.prebuilt_root() + target_dir.mkdir(parents=True, exist_ok=True) + for suffix in (".pyd", ".so"): + existing = target_dir / f"witwin_core_mesh_sdf_cuda{suffix}" + if existing.exists(): + existing.unlink() + target = cuda_build.prebuilt_extension_path() + shutil.copy2(module_file, target) + print(f"Built prebuilt mesh SDF CUDA extension: {target}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_geometry_sdf.py b/tests/test_geometry_sdf.py index d61f1a1..1cbc0a1 100644 --- a/tests/test_geometry_sdf.py +++ b/tests/test_geometry_sdf.py @@ -6,7 +6,7 @@ from witwin.core import Box, Cone, Cylinder, Ellipsoid, HollowBox, Mesh, Prism, Pyramid, Sphere, Torus from witwin.core.geometry import ComplexPolySlab, PolySlab from witwin.core.geometry.mesh_sdf import ( - _slang_mesh_sdf_available, + _mesh_sdf_available, _triangle_mesh_smooth_signed_distance_torch, _triangle_mesh_unsigned_distance_torch, triangle_mesh_smooth_signed_distance, @@ -90,9 +90,9 @@ def _grid_plane_mesh(*, resolution=24, size=1.0): def _require_cuda_mesh_sdf(): if not torch.cuda.is_available(): - pytest.skip("CUDA is required for the Slang mesh SDF path.") - if not _slang_mesh_sdf_available(): - pytest.skip("slangtorch is required for the Slang mesh SDF path.") + pytest.skip("CUDA is required for the native mesh SDF path.") + if not _mesh_sdf_available(): + pytest.skip("Native CUDA mesh SDF extension is required for this path.") @pytest.mark.parametrize( diff --git a/witwin/core/geometry/cuda/__init__.py b/witwin/core/geometry/cuda/__init__.py new file mode 100644 index 0000000..53f7368 --- /dev/null +++ b/witwin/core/geometry/cuda/__init__.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +from .backend import get_compiled_extension, get_native_mesh_sdf_module, is_available + +__all__ = [ + "get_compiled_extension", + "get_native_mesh_sdf_module", + "is_available", +] diff --git a/witwin/core/geometry/cuda/backend.py b/witwin/core/geometry/cuda/backend.py new file mode 100644 index 0000000..63afa2a --- /dev/null +++ b/witwin/core/geometry/cuda/backend.py @@ -0,0 +1,177 @@ +"""Loader for the native CUDA mesh-SDF extension. + +Exposes a module object that is call-compatible with the former SlangTorch +module: ``module.(**kwargs).launchRaw(blockSize=..., gridSize=...)``. +The kernels compute their own launch grid, so ``launchRaw`` ignores the block / +grid hints and simply dispatches the compiled function. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any + +import torch + + +_COMPILED_EXTENSION: Any | None = None + + +def is_available() -> bool: + return bool(torch.cuda.is_available()) + + +def get_compiled_extension(*, verbose: bool = False) -> Any: + global _COMPILED_EXTENSION + if _COMPILED_EXTENSION is None: + from .build import build_extension + + _COMPILED_EXTENSION = build_extension(verbose=verbose) + return _COMPILED_EXTENSION + + +def _query_mesh_unsigned_distance(*, triangles, points, unsignedDistance, closestTriangleIndex): + get_compiled_extension().query_mesh_unsigned_distance( + triangles, points, unsignedDistance, closestTriangleIndex + ) + + +def _query_mesh_distance_and_winding(*, triangles, points, unsignedDistance, windingAngle, closestTriangleIndex): + get_compiled_extension().query_mesh_distance_and_winding( + triangles, points, unsignedDistance, windingAngle, closestTriangleIndex + ) + + +def _query_mesh_unsigned_distance_bvh( + *, + triangles, + points, + nodeBBoxMin, + nodeBBoxMax, + nodeLeft, + nodeRight, + nodeStart, + nodeCount, + triangleIndices, + unsignedDistance, + closestTriangleIndex, +): + get_compiled_extension().query_mesh_unsigned_distance_bvh( + triangles, + points, + nodeBBoxMin, + nodeBBoxMax, + nodeLeft, + nodeRight, + nodeStart, + nodeCount, + triangleIndices, + unsignedDistance, + closestTriangleIndex, + ) + + +def _query_mesh_parity_sign_bvh( + *, + triangles, + points, + nodeBBoxMin, + nodeBBoxMax, + nodeLeft, + nodeRight, + nodeStart, + nodeCount, + triangleIndices, + jitterScale, + inside, +): + get_compiled_extension().query_mesh_parity_sign_bvh( + triangles, + points, + nodeBBoxMin, + nodeBBoxMax, + nodeLeft, + nodeRight, + nodeStart, + nodeCount, + triangleIndices, + float(jitterScale), + inside, + ) + + +def _backward_mesh_unsigned_distance( + *, triangles, points, closestTriangleIndex, gradUnsignedDistance, gradTriangles, gradPoints +): + get_compiled_extension().backward_mesh_unsigned_distance( + triangles, points, closestTriangleIndex, gradUnsignedDistance, gradTriangles, gradPoints + ) + + +def _backward_mesh_distance_and_winding( + *, + triangles, + points, + closestTriangleIndex, + gradUnsignedDistance, + gradWindingAngle, + gradTriangles, + gradPoints, +): + get_compiled_extension().backward_mesh_distance_and_winding( + triangles, + points, + closestTriangleIndex, + gradUnsignedDistance, + gradWindingAngle, + gradTriangles, + gradPoints, + ) + + +_KERNELS: dict[str, Callable[..., None]] = { + "queryMeshUnsignedDistance": _query_mesh_unsigned_distance, + "queryMeshDistanceAndWinding": _query_mesh_distance_and_winding, + "queryMeshUnsignedDistanceBVH": _query_mesh_unsigned_distance_bvh, + "queryMeshParitySignBVH": _query_mesh_parity_sign_bvh, + "backwardMeshUnsignedDistance": _backward_mesh_unsigned_distance, + "backwardMeshDistanceAndWinding": _backward_mesh_distance_and_winding, +} + + +class _Launch: + __slots__ = ("_kernel", "_kwargs") + + def __init__(self, kernel: Callable[..., None], kwargs: dict): + self._kernel = kernel + self._kwargs = kwargs + + def launchRaw(self, *, blockSize=None, gridSize=None): # noqa: N802 - kernel launch entry point + del blockSize, gridSize + self._kernel(**self._kwargs) + return None + + +class NativeMeshSDFModule: + def __getattr__(self, name: str): + kernel = _KERNELS.get(name) + if kernel is None: + raise AttributeError(f"Native CUDA mesh SDF backend does not implement kernel {name!r}.") + + def bind(**kwargs): + return _Launch(kernel, kwargs) + + bind.__name__ = name + object.__setattr__(self, name, bind) + return bind + + +_NATIVE_MODULE = NativeMeshSDFModule() + + +def get_native_mesh_sdf_module(*, verbose: bool = False) -> NativeMeshSDFModule: + """Compile the extension (once) and return the launch-compatible module.""" + if not is_available(): + raise RuntimeError("Native CUDA mesh SDF backend requires torch.cuda.is_available() to be True.") + get_compiled_extension(verbose=verbose) + return _NATIVE_MODULE diff --git a/witwin/core/geometry/cuda/build.py b/witwin/core/geometry/cuda/build.py new file mode 100644 index 0000000..c66e3ca --- /dev/null +++ b/witwin/core/geometry/cuda/build.py @@ -0,0 +1,247 @@ +"""JIT build of the native CUDA mesh-SDF extension. + +Self-contained ``torch.utils.cpp_extension.load`` wrapper with the Windows +MSVC + CUDA toolchain discovery the platform needs. ``witwin.core`` must not +depend on ``witwin.maxwell``, so the toolchain helpers are duplicated here +rather than imported from the maxwell FDTD CUDA build. +""" + +from __future__ import annotations + +import importlib.util +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +from torch.utils.cpp_extension import load + + +EXTENSION_NAME = "witwin_core_mesh_sdf_cuda" + + +def _candidate_vcvars64_paths() -> list[Path]: + paths: list[Path] = [] + program_files_x86 = os.environ.get("ProgramFiles(x86)") + if program_files_x86: + vswhere = Path(program_files_x86) / "Microsoft Visual Studio" / "Installer" / "vswhere.exe" + if vswhere.exists(): + try: + install_root = subprocess.check_output( + [ + str(vswhere), + "-latest", + "-products", + "*", + "-requires", + "Microsoft.VisualStudio.Component.VC.Tools.x86.x64", + "-property", + "installationPath", + ], + text=True, + encoding="mbcs", + errors="replace", + ).strip() + except (OSError, subprocess.CalledProcessError): + install_root = "" + if install_root: + paths.append(Path(install_root) / "VC" / "Auxiliary" / "Build" / "vcvars64.bat") + + root = os.environ.get("VSINSTALLDIR") + if root: + paths.append(Path(root) / "VC" / "Auxiliary" / "Build" / "vcvars64.bat") + + program_files = os.environ.get("ProgramFiles", r"C:\Program Files") + for edition in ("Community", "Professional", "Enterprise", "BuildTools"): + paths.append( + Path(program_files) + / "Microsoft Visual Studio" + / "2022" + / edition + / "VC" + / "Auxiliary" + / "Build" + / "vcvars64.bat" + ) + + seen: set[Path] = set() + unique_paths: list[Path] = [] + for path in paths: + resolved = path.resolve() if path.exists() else path + if resolved in seen: + continue + seen.add(resolved) + unique_paths.append(path) + return unique_paths + + +def _load_vcvars64_environment() -> bool: + for vcvars in _candidate_vcvars64_paths(): + if not vcvars.exists(): + continue + fd, probe_name = tempfile.mkstemp(prefix="witwin_core_mesh_sdf_vcvars_probe_", suffix=".cmd") + os.close(fd) + probe = Path(probe_name) + probe.write_text(f'@echo off\ncall "{vcvars}" >nul\nset\n', encoding="utf-8") + try: + output = subprocess.check_output( + ["cmd.exe", "/d", "/c", str(probe)], + text=True, + encoding="mbcs", + errors="replace", + ) + except (OSError, subprocess.CalledProcessError): + continue + finally: + try: + probe.unlink() + except OSError: + pass + updates: dict[str, str] = {} + for line in output.splitlines(): + if "=" not in line: + continue + key, value = line.split("=", 1) + updates[key] = value + for key, value in updates.items(): + os.environ[key] = value + path_value = updates.get("PATH", updates.get("Path")) + if path_value is not None: + os.environ["PATH"] = path_value + os.environ["Path"] = path_value + return True + return False + + +def _ensure_windows_build_tools_on_path() -> None: + if os.name != "nt": + return + if shutil.which("cl") is None: + _load_vcvars64_environment() + if shutil.which("cl") is None: + return + + current_path = os.environ.get("PATH") or os.environ.get("Path") or "" + prefixes: list[str] = [] + vc_tools = os.environ.get("VCToolsInstallDir") + if vc_tools: + prefixes.append(str(Path(vc_tools) / "bin" / "Hostx64" / "x64")) + vs_install = os.environ.get("VSINSTALLDIR") + if vs_install: + prefixes.append( + str(Path(vs_install) / "Common7" / "IDE" / "CommonExtensions" / "Microsoft" / "CMake" / "Ninja") + ) + if not prefixes: + return + merged_path = os.pathsep.join([*prefixes, current_path]) + os.environ["PATH"] = merged_path + os.environ["Path"] = merged_path + + +def _ensure_cuda_home_from_nvcc() -> None: + if os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH"): + return + nvcc = shutil.which("nvcc") + if nvcc is None: + return + cuda_home = Path(nvcc).resolve().parents[1] + os.environ["CUDA_HOME"] = str(cuda_home) + if os.name == "nt": + os.environ["CUDA_PATH"] = str(cuda_home) + + +def _conda_torch_ldflags() -> list[str]: + # Conda-distributed torch keeps c10.lib/torch_cuda.lib in \Library\lib + # instead of site-packages\torch\lib, which cpp_extension does not search. + if os.name != "nt": + return [] + library_lib = Path(sys.prefix) / "Library" / "lib" + if (library_lib / "c10.lib").exists(): + return [f"/LIBPATH:{library_lib}"] + return [] + + +def source_root() -> Path: + return Path(__file__).resolve().parent + + +def prebuilt_root() -> Path: + return source_root() / "prebuilt" + + +def extension_suffix() -> str: + return ".pyd" if os.name == "nt" else ".so" + + +def prebuilt_extension_path() -> Path: + return prebuilt_root() / f"{EXTENSION_NAME}{extension_suffix()}" + + +def extension_sources() -> list[Path]: + root = source_root() + return [root / "extension.cpp", root / "mesh_sdf_kernels.cu"] + + +def _load_extension_file(module_path: Path): + spec = importlib.util.spec_from_file_location(EXTENSION_NAME, module_path) + if spec is None or spec.loader is None: + raise ImportError(f"Unable to create import spec for {module_path}.") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_packaged_prebuilt_extension(): + # Wheels ship a compiled extension in ``prebuilt/`` so installs never invoke + # nvcc/MSVC. A build against a mismatched torch ABI raises on import; the + # caller falls back to a JIT rebuild in that case. + module_path = prebuilt_extension_path() + if not module_path.exists(): + return None + return _load_extension_file(module_path) + + +def _load_prebuilt_extension(build_directory: Path): + # Load the already-built module directly, bypassing the build toolchain. + module_path = build_directory / f"{EXTENSION_NAME}{extension_suffix()}" + if not module_path.exists(): + raise FileNotFoundError( + f"WITWIN_CORE_MESH_SDF_CUDA_PREBUILT=1 but {module_path} does not exist; " + "run a normal build first." + ) + return _load_extension_file(module_path) + + +def _jit_build(build_directory: Path, *, verbose: bool): + _ensure_windows_build_tools_on_path() + _ensure_cuda_home_from_nvcc() + build_directory.mkdir(parents=True, exist_ok=True) + return load( + name=EXTENSION_NAME, + sources=[str(path) for path in extension_sources()], + build_directory=str(build_directory), + extra_cflags=["/O2"] if os.name == "nt" else ["-O3"], + extra_cuda_cflags=["-O3"], + extra_ldflags=_conda_torch_ldflags(), + verbose=verbose, + ) + + +def build_extension(*, verbose: bool = False): + default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME + build_directory = Path(os.environ.get("WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", default_build_directory)) + + if os.environ.get("WITWIN_CORE_MESH_SDF_CUDA_SKIP_PREBUILT") != "1": + try: + module = _load_packaged_prebuilt_extension() + except Exception: # noqa: BLE001 - stale/ABI-mismatched prebuilt, rebuild instead + module = None + if module is not None: + return module + + if os.environ.get("WITWIN_CORE_MESH_SDF_CUDA_PREBUILT") == "1": + return _load_prebuilt_extension(build_directory) + + return _jit_build(build_directory, verbose=verbose) diff --git a/witwin/core/geometry/cuda/extension.cpp b/witwin/core/geometry/cuda/extension.cpp new file mode 100644 index 0000000..e1cad20 --- /dev/null +++ b/witwin/core/geometry/cuda/extension.cpp @@ -0,0 +1,43 @@ +#include +#include + +bool is_available() { + return torch::cuda::is_available(); +} + +void query_mesh_unsigned_distance_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, at::Tensor closest_triangle_index); +void query_mesh_distance_and_winding_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, + at::Tensor winding_angle, at::Tensor closest_triangle_index); +void query_mesh_unsigned_distance_bvh_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, + at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, + at::Tensor triangle_indices, at::Tensor unsigned_distance, at::Tensor closest_triangle_index); +void query_mesh_parity_sign_bvh_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, + at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, + at::Tensor triangle_indices, double jitter_scale, at::Tensor inside); +void backward_mesh_unsigned_distance_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, + at::Tensor grad_unsigned_distance, at::Tensor grad_triangles, at::Tensor grad_points); +void backward_mesh_distance_and_winding_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, + at::Tensor grad_unsigned_distance, at::Tensor grad_winding_angle, + at::Tensor grad_triangles, at::Tensor grad_points); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("is_available", &is_available, "Return whether CUDA is available to PyTorch."); + m.def("query_mesh_unsigned_distance", &query_mesh_unsigned_distance_cuda, + "Brute-force unsigned point-to-mesh distance and closest triangle index."); + m.def("query_mesh_distance_and_winding", &query_mesh_distance_and_winding_cuda, + "Brute-force unsigned distance plus solid-angle winding sum."); + m.def("query_mesh_unsigned_distance_bvh", &query_mesh_unsigned_distance_bvh_cuda, + "BVH-accelerated unsigned distance and closest triangle index."); + m.def("query_mesh_parity_sign_bvh", &query_mesh_parity_sign_bvh_cuda, + "BVH-accelerated ray-parity inside/outside test."); + m.def("backward_mesh_unsigned_distance", &backward_mesh_unsigned_distance_cuda, + "Backward pass for the unsigned distance query."); + m.def("backward_mesh_distance_and_winding", &backward_mesh_distance_and_winding_cuda, + "Backward pass for the distance-and-winding query."); +} diff --git a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu new file mode 100644 index 0000000..12a80df --- /dev/null +++ b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu @@ -0,0 +1,688 @@ +// Native CUDA triangle-mesh distance / winding / parity kernels. +// +// Direct port of the former ``mesh_sdf.slang`` module. The forward kernels +// reproduce the Slang closest-point (Ericson) and solid-angle math verbatim; +// the backward kernels replace Slang autodiff with hand-derived analytic +// gradients. Every closest-point feature (vertex / edge / face) contributes +// gradients to the triangle vertices through the barycentric weights of the +// closest point, so grad_p == -(grad_a + grad_b + grad_c) holds by +// construction (translation invariance of the distance). + +#include +#include +#include +#include + +namespace { + +constexpr float kSmoothEps = 1.0e-12f; +constexpr float kDenomEps = 1.0e-12f; +constexpr int kMaxStackDepth = 64; + +__host__ __device__ __forceinline__ float3 vsub(float3 a, float3 b) { + return make_float3(a.x - b.x, a.y - b.y, a.z - b.z); +} +__host__ __device__ __forceinline__ float3 vadd(float3 a, float3 b) { + return make_float3(a.x + b.x, a.y + b.y, a.z + b.z); +} +__host__ __device__ __forceinline__ float3 vscale(float3 a, float s) { + return make_float3(a.x * s, a.y * s, a.z * s); +} +__device__ __forceinline__ float vdot(float3 a, float3 b) { + return a.x * b.x + a.y * b.y + a.z * b.z; +} +__device__ __forceinline__ float3 vcross(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__ float vlen(float3 a) { + return sqrtf(vdot(a, a)); +} + +__device__ __forceinline__ float safe_signed_denom(float value, float eps) { + return value >= 0.0f ? fmaxf(value, eps) : fminf(value, -eps); +} + +__device__ __forceinline__ float3 read_point(const float* __restrict__ points, int index) { + const float* p = points + static_cast(index) * 3; + return make_float3(p[0], p[1], p[2]); +} + +__device__ __forceinline__ float3 read_vertex(const float* __restrict__ triangles, int triangle, int vertex) { + const float* t = triangles + (static_cast(triangle) * 3 + vertex) * 3; + return make_float3(t[0], t[1], t[2]); +} + +__device__ __forceinline__ float3 read_vec3(const float* __restrict__ data, int index) { + const float* v = data + static_cast(index) * 3; + return make_float3(v[0], v[1], v[2]); +} + +// Closest point on triangle (a, b, c) to p (Ericson, matching mesh_sdf.slang). +// Returns the squared distance and the barycentric weights (la, lb, lc) of the +// closest point so the backward pass can distribute vertex gradients. +__device__ __forceinline__ float closest_point_barycentric( + float3 p, float3 a, float3 b, float3 c, float& la, float& lb, float& lc) { + const float3 ab = vsub(b, a); + const float3 ac = vsub(c, a); + const float3 ap = vsub(p, a); + + const float d1 = vdot(ab, ap); + const float d2 = vdot(ac, ap); + if (d1 <= 0.0f && d2 <= 0.0f) { + la = 1.0f; lb = 0.0f; lc = 0.0f; + return vdot(ap, ap); + } + + const float3 bp = vsub(p, b); + const float d3 = vdot(ab, bp); + const float d4 = vdot(ac, bp); + if (d3 >= 0.0f && d4 <= d3) { + la = 0.0f; lb = 1.0f; lc = 0.0f; + return vdot(bp, bp); + } + + const float vc = d1 * d4 - d3 * d2; + if (vc <= 0.0f && d1 >= 0.0f && d3 <= 0.0f) { + const float v = d1 / (d1 - d3); + la = 1.0f - v; lb = v; lc = 0.0f; + const float3 delta = vsub(p, vadd(a, vscale(ab, v))); + return vdot(delta, delta); + } + + const float3 cp = vsub(p, c); + const float d5 = vdot(ab, cp); + const float d6 = vdot(ac, cp); + if (d6 >= 0.0f && d5 <= d6) { + la = 0.0f; lb = 0.0f; lc = 1.0f; + return vdot(cp, cp); + } + + const float vb = d5 * d2 - d1 * d6; + if (vb <= 0.0f && d2 >= 0.0f && d6 <= 0.0f) { + const float w = d2 / safe_signed_denom(d2 - d6, kDenomEps); + la = 1.0f - w; lb = 0.0f; lc = w; + const float3 delta = vsub(p, vadd(a, vscale(ac, w))); + return vdot(delta, delta); + } + + const float va = d3 * d6 - d5 * d4; + if (va <= 0.0f && (d4 - d3) >= 0.0f && (d5 - d6) >= 0.0f) { + const float w = (d4 - d3) / safe_signed_denom((d4 - d3) + (d5 - d6), kDenomEps); + la = 0.0f; lb = 1.0f - w; lc = w; + const float3 bc = vsub(c, b); + const float3 delta = vsub(p, vadd(b, vscale(bc, w))); + return vdot(delta, delta); + } + + const float denom = safe_signed_denom(va + vb + vc, kDenomEps); + const float v = vb / denom; + const float w = vc / denom; + la = 1.0f - v - w; lb = v; lc = w; + const float3 closest = vadd(vadd(a, vscale(ab, v)), vscale(ac, w)); + const float3 delta = vsub(p, closest); + return vdot(delta, delta); +} + +__device__ __forceinline__ float point_triangle_distance_squared(float3 p, float3 a, float3 b, float3 c) { + float la, lb, lc; + return closest_point_barycentric(p, a, b, c, la, lb, lc); +} + +__device__ __forceinline__ float triangle_solid_angle(float3 p, float3 a0, float3 b0, float3 c0) { + const float eps = 1.0e-12f; + const float3 a = vsub(a0, p); + const float3 b = vsub(b0, p); + const float3 c = vsub(c0, p); + const float la = fmaxf(vlen(a), eps); + const float lb = fmaxf(vlen(b), eps); + const float lc = fmaxf(vlen(c), eps); + const float numerator = vdot(a, vcross(b, c)); + const float denominator = + la * lb * lc + vdot(a, b) * lc + vdot(b, c) * la + vdot(c, a) * lb; + return 2.0f * atan2f(numerator, safe_signed_denom(denominator, eps)); +} + +__device__ __forceinline__ float point_aabb_distance_squared(float3 p, float3 lo, float3 hi) { + const float dx = p.x < lo.x ? (lo.x - p.x) : (p.x > hi.x ? p.x - hi.x : 0.0f); + const float dy = p.y < lo.y ? (lo.y - p.y) : (p.y > hi.y ? p.y - hi.y : 0.0f); + const float dz = p.z < lo.z ? (lo.z - p.z) : (p.z > hi.z ? p.z - hi.z : 0.0f); + return dx * dx + dy * dy + dz * dz; +} + +__device__ __forceinline__ bool ray_aabb_hit(float3 origin, float3 direction, float3 lo, float3 hi) { + const float eps = 1.0e-12f; + const float3 inv = make_float3( + fabsf(direction.x) > eps ? 1.0f / direction.x : 1.0e30f, + fabsf(direction.y) > eps ? 1.0f / direction.y : 1.0e30f, + fabsf(direction.z) > eps ? 1.0f / direction.z : 1.0e30f); + const float tx1 = (lo.x - origin.x) * inv.x; + const float tx2 = (hi.x - origin.x) * inv.x; + float tmin = fminf(tx1, tx2); + float tmax = fmaxf(tx1, tx2); + const float ty1 = (lo.y - origin.y) * inv.y; + const float ty2 = (hi.y - origin.y) * inv.y; + tmin = fmaxf(tmin, fminf(ty1, ty2)); + tmax = fminf(tmax, fmaxf(ty1, ty2)); + const float tz1 = (lo.z - origin.z) * inv.z; + const float tz2 = (hi.z - origin.z) * inv.z; + tmin = fmaxf(tmin, fminf(tz1, tz2)); + tmax = fminf(tmax, fmaxf(tz1, tz2)); + return tmax >= fmaxf(tmin, 0.0f); +} + +__device__ __forceinline__ bool ray_triangle_hit(float3 origin, float3 direction, float3 a, float3 b, float3 c) { + const float epsilon = 1.0e-6f; + const float3 edge1 = vsub(b, a); + const float3 edge2 = vsub(c, a); + const float3 h = vcross(direction, edge2); + const float det = vdot(edge1, h); + if (fabsf(det) < epsilon) { + return false; + } + const float inv_det = 1.0f / det; + const float3 s = vsub(origin, a); + const float u = inv_det * vdot(s, h); + if (u < 0.0f || u > 1.0f) { + return false; + } + const float3 q = vcross(s, edge1); + const float v = inv_det * vdot(direction, q); + if (v < 0.0f || (u + v) > 1.0f) { + return false; + } + const float t = inv_det * vdot(edge2, q); + return t > epsilon; +} + +// grad(u) where u = sqrt(d2 + kSmoothEps) - const distributes over the triangle +// vertices through the barycentric weights of the closest point. n scales the +// unit direction (p - q) / sqrt(d2 + eps) by the upstream gradient. +__device__ __forceinline__ void accumulate_closest_distance_gradient( + float3 p, float3 a, float3 b, float3 c, float grad_output, + float* __restrict__ grad_triangles, int triangle_index, float3& point_gradient) { + float la, lb, lc; + const float dist2 = closest_point_barycentric(p, a, b, c, la, lb, lc); + const float denom = sqrtf(dist2 + kSmoothEps); + // q = la*a + lb*b + lc*c, so p - q = grad_output * (p - q) / denom = n below. + const float3 q = vadd(vadd(vscale(a, la), vscale(b, lb)), vscale(c, lc)); + const float3 n = vscale(vsub(p, q), grad_output / denom); + point_gradient = vadd(point_gradient, n); + float* g = grad_triangles + static_cast(triangle_index) * 9; + atomicAdd(g + 0, -la * n.x); + atomicAdd(g + 1, -la * n.y); + atomicAdd(g + 2, -la * n.z); + atomicAdd(g + 3, -lb * n.x); + atomicAdd(g + 4, -lb * n.y); + atomicAdd(g + 5, -lb * n.z); + atomicAdd(g + 6, -lc * n.x); + atomicAdd(g + 7, -lc * n.y); + atomicAdd(g + 8, -lc * n.z); +} + +// grad of 2*atan2(N, D) w.r.t. the three (world) vertices, accumulated onto +// grad_triangles; the point gradient is -(gA + gB + gC). +__device__ __forceinline__ void accumulate_solid_angle_gradient( + float3 p, float3 a0, float3 b0, float3 c0, float grad_output, + float* __restrict__ grad_triangles, int triangle_index, float3& point_gradient) { + const float eps = 1.0e-12f; + const float3 a = vsub(a0, p); + const float3 b = vsub(b0, p); + const float3 c = vsub(c0, p); + const float la = fmaxf(vlen(a), eps); + const float lb = fmaxf(vlen(b), eps); + const float lc = fmaxf(vlen(c), eps); + const float ab = vdot(a, b); + const float bc = vdot(b, c); + const float ca = vdot(c, a); + const float N = vdot(a, vcross(b, c)); + const float D = la * lb * lc + ab * lc + bc * la + ca * lb; + const float scale = 2.0f * grad_output / fmaxf(N * N + D * D, eps); + + // dN/dA = b x c, dN/dB = c x a, dN/dC = a x b. + const float3 dN_da = vcross(b, c); + const float3 dN_db = vcross(c, a); + const float3 dN_dc = vcross(a, b); + const float3 a_hat = vscale(a, 1.0f / la); + const float3 b_hat = vscale(b, 1.0f / lb); + const float3 c_hat = vscale(c, 1.0f / lc); + // dD/dA = (lb*lc + b.c) a_hat + lc*b + lb*c, and cyclic variants. + const float3 dD_da = vadd(vadd(vscale(a_hat, lb * lc + bc), vscale(b, lc)), vscale(c, lb)); + const float3 dD_db = vadd(vadd(vscale(b_hat, la * lc + ca), vscale(a, lc)), vscale(c, la)); + const float3 dD_dc = vadd(vadd(vscale(c_hat, la * lb + ab), vscale(b, la)), vscale(a, lb)); + + // d(atan2(N, D)) = (D dN - N dD) / (N^2 + D^2); grad = scale * that. + const float3 gA = vscale(vsub(vscale(dN_da, D), vscale(dD_da, N)), scale); + const float3 gB = vscale(vsub(vscale(dN_db, D), vscale(dD_db, N)), scale); + const float3 gC = vscale(vsub(vscale(dN_dc, D), vscale(dD_dc, N)), scale); + + float* g = grad_triangles + static_cast(triangle_index) * 9; + atomicAdd(g + 0, gA.x); atomicAdd(g + 1, gA.y); atomicAdd(g + 2, gA.z); + atomicAdd(g + 3, gB.x); atomicAdd(g + 4, gB.y); atomicAdd(g + 5, gB.z); + atomicAdd(g + 6, gC.x); atomicAdd(g + 7, gC.y); atomicAdd(g + 8, gC.z); + point_gradient = vsub(point_gradient, vadd(vadd(gA, gB), gC)); +} + +__global__ void query_unsigned_distance_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + int triangle_count, + int point_count, + float* __restrict__ unsigned_distance, + int* __restrict__ closest_triangle_index) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + const float3 p = read_point(points, i); + float min_dist2 = 1.0e30f; + int closest = -1; + for (int t = 0; t < triangle_count; ++t) { + const float d2 = point_triangle_distance_squared( + p, read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2)); + if (d2 < min_dist2) { + min_dist2 = d2; + closest = t; + } + } + unsigned_distance[i] = sqrtf(fmaxf(min_dist2, 0.0f) + kSmoothEps) - sqrtf(kSmoothEps); + closest_triangle_index[i] = closest; +} + +__global__ void query_distance_and_winding_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + int triangle_count, + int point_count, + float* __restrict__ unsigned_distance, + float* __restrict__ winding_angle, + int* __restrict__ closest_triangle_index) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + const float3 p = read_point(points, i); + float min_dist2 = 1.0e30f; + float total_angle = 0.0f; + int closest = -1; + for (int t = 0; t < triangle_count; ++t) { + const float3 a = read_vertex(triangles, t, 0); + const float3 b = read_vertex(triangles, t, 1); + const float3 c = read_vertex(triangles, t, 2); + const float d2 = point_triangle_distance_squared(p, a, b, c); + if (d2 < min_dist2) { + min_dist2 = d2; + closest = t; + } + total_angle += triangle_solid_angle(p, a, b, c); + } + unsigned_distance[i] = sqrtf(fmaxf(min_dist2, 0.0f) + kSmoothEps) - sqrtf(kSmoothEps); + winding_angle[i] = total_angle; + closest_triangle_index[i] = closest; +} + +__global__ void query_unsigned_distance_bvh_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + const float* __restrict__ node_bbox_min, + const float* __restrict__ node_bbox_max, + const int* __restrict__ node_left, + const int* __restrict__ node_right, + const int* __restrict__ node_start, + const int* __restrict__ node_count, + const int* __restrict__ triangle_indices, + int point_count, + float* __restrict__ unsigned_distance, + int* __restrict__ closest_triangle_index) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + int stack[kMaxStackDepth]; + int stack_size = 0; + stack[stack_size++] = 0; + + const float3 p = read_point(points, i); + float min_dist2 = 1.0e30f; + int closest = -1; + + while (stack_size > 0) { + const int node = stack[--stack_size]; + const float node_dist2 = point_aabb_distance_squared( + p, read_vec3(node_bbox_min, node), read_vec3(node_bbox_max, node)); + if (node_dist2 > min_dist2) { + continue; + } + const int count = node_count[node]; + if (count > 0) { + const int start = node_start[node]; + for (int offset = 0; offset < count; ++offset) { + const int t = triangle_indices[start + offset]; + const float d2 = point_triangle_distance_squared( + p, read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2)); + if (d2 < min_dist2) { + min_dist2 = d2; + closest = t; + } + } + continue; + } + const int left = node_left[node]; + const int right = node_right[node]; + if (left < 0 || right < 0) { + continue; + } + const float left_dist2 = point_aabb_distance_squared( + p, read_vec3(node_bbox_min, left), read_vec3(node_bbox_max, left)); + const float right_dist2 = point_aabb_distance_squared( + p, read_vec3(node_bbox_min, right), read_vec3(node_bbox_max, right)); + if (left_dist2 < right_dist2) { + if (stack_size + 2 <= kMaxStackDepth) { + stack[stack_size++] = right; + stack[stack_size++] = left; + } + } else { + if (stack_size + 2 <= kMaxStackDepth) { + stack[stack_size++] = left; + stack[stack_size++] = right; + } + } + } + + unsigned_distance[i] = sqrtf(fmaxf(min_dist2, 0.0f) + kSmoothEps) - sqrtf(kSmoothEps); + closest_triangle_index[i] = closest; +} + +__global__ void query_parity_sign_bvh_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + const float* __restrict__ node_bbox_min, + const float* __restrict__ node_bbox_max, + const int* __restrict__ node_left, + const int* __restrict__ node_right, + const int* __restrict__ node_start, + const int* __restrict__ node_count, + const int* __restrict__ triangle_indices, + float jitter_scale, + int point_count, + int* __restrict__ inside) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + int stack[kMaxStackDepth]; + int stack_size = 0; + stack[stack_size++] = 0; + + const float3 p = read_point(points, i); + const float3 direction = make_float3(1.0f, 0.0f, 0.0f); + const float3 origin = make_float3(p.x, p.y + jitter_scale, p.z + jitter_scale * 1.618034f); + int intersections = 0; + + while (stack_size > 0) { + const int node = stack[--stack_size]; + if (!ray_aabb_hit(origin, direction, read_vec3(node_bbox_min, node), read_vec3(node_bbox_max, node))) { + continue; + } + const int count = node_count[node]; + if (count > 0) { + const int start = node_start[node]; + for (int offset = 0; offset < count; ++offset) { + const int t = triangle_indices[start + offset]; + if (ray_triangle_hit(origin, direction, + read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2))) { + intersections += 1; + } + } + continue; + } + const int left = node_left[node]; + const int right = node_right[node]; + if (left >= 0 && stack_size < kMaxStackDepth) { + stack[stack_size++] = left; + } + if (right >= 0 && stack_size < kMaxStackDepth) { + stack[stack_size++] = right; + } + } + + inside[i] = (intersections & 1) != 0 ? 1 : 0; +} + +__global__ void backward_unsigned_distance_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + const int* __restrict__ closest_triangle_index, + const float* __restrict__ grad_unsigned_distance, + int point_count, + float* __restrict__ grad_triangles, + float* __restrict__ grad_points) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + const float grad_output = grad_unsigned_distance[i]; + if (grad_output == 0.0f) { + return; + } + const int t = closest_triangle_index[i]; + if (t < 0) { + return; + } + const float3 p = read_point(points, i); + float3 point_gradient = make_float3(0.0f, 0.0f, 0.0f); + accumulate_closest_distance_gradient( + p, read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2), + grad_output, grad_triangles, t, point_gradient); + float* gp = grad_points + static_cast(i) * 3; + gp[0] = point_gradient.x; + gp[1] = point_gradient.y; + gp[2] = point_gradient.z; +} + +__global__ void backward_distance_and_winding_kernel( + const float* __restrict__ triangles, + const float* __restrict__ points, + const int* __restrict__ closest_triangle_index, + const float* __restrict__ grad_unsigned_distance, + const float* __restrict__ grad_winding_angle, + int triangle_count, + int point_count, + float* __restrict__ grad_triangles, + float* __restrict__ grad_points) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= point_count) { + return; + } + const float grad_unsigned = grad_unsigned_distance[i]; + const float grad_angle = grad_winding_angle[i]; + if (grad_unsigned == 0.0f && grad_angle == 0.0f) { + return; + } + const float3 p = read_point(points, i); + float3 point_gradient = make_float3(0.0f, 0.0f, 0.0f); + + if (grad_angle != 0.0f) { + for (int t = 0; t < triangle_count; ++t) { + accumulate_solid_angle_gradient( + p, read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2), + grad_angle, grad_triangles, t, point_gradient); + } + } + + if (grad_unsigned != 0.0f) { + const int t = closest_triangle_index[i]; + if (t >= 0) { + accumulate_closest_distance_gradient( + p, read_vertex(triangles, t, 0), read_vertex(triangles, t, 1), read_vertex(triangles, t, 2), + grad_unsigned, grad_triangles, t, point_gradient); + } + } + + float* gp = grad_points + static_cast(i) * 3; + gp[0] = point_gradient.x; + gp[1] = point_gradient.y; + gp[2] = point_gradient.z; +} + +dim3 linear_grid(int elements, int block_size) { + return dim3(static_cast((elements + block_size - 1) / block_size), 1, 1); +} + +void check_f32(const at::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == at::kFloat, name, " must be float32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +void check_i32(const at::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.scalar_type() == at::kInt, name, " must be int32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +constexpr int kBlock = 256; + +} // namespace + +void query_mesh_unsigned_distance_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, at::Tensor closest_triangle_index) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_f32(unsigned_distance, "unsigned_distance"); + check_i32(closest_triangle_index, "closest_triangle_index"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + query_unsigned_distance_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + static_cast(triangles.size(0)), point_count, + unsigned_distance.data_ptr(), closest_triangle_index.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void query_mesh_distance_and_winding_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor unsigned_distance, + at::Tensor winding_angle, at::Tensor closest_triangle_index) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_f32(unsigned_distance, "unsigned_distance"); + check_f32(winding_angle, "winding_angle"); + check_i32(closest_triangle_index, "closest_triangle_index"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + query_distance_and_winding_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + static_cast(triangles.size(0)), point_count, + unsigned_distance.data_ptr(), winding_angle.data_ptr(), + closest_triangle_index.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void query_mesh_unsigned_distance_bvh_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, + at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, + at::Tensor triangle_indices, at::Tensor unsigned_distance, at::Tensor closest_triangle_index) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_f32(node_bbox_min, "node_bbox_min"); + check_f32(node_bbox_max, "node_bbox_max"); + check_i32(node_left, "node_left"); + check_i32(node_right, "node_right"); + check_i32(node_start, "node_start"); + check_i32(node_count, "node_count"); + check_i32(triangle_indices, "triangle_indices"); + check_f32(unsigned_distance, "unsigned_distance"); + check_i32(closest_triangle_index, "closest_triangle_index"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + query_unsigned_distance_bvh_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + node_bbox_min.data_ptr(), node_bbox_max.data_ptr(), + node_left.data_ptr(), node_right.data_ptr(), + node_start.data_ptr(), node_count.data_ptr(), + triangle_indices.data_ptr(), point_count, + unsigned_distance.data_ptr(), closest_triangle_index.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void query_mesh_parity_sign_bvh_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor node_bbox_min, at::Tensor node_bbox_max, + at::Tensor node_left, at::Tensor node_right, at::Tensor node_start, at::Tensor node_count, + at::Tensor triangle_indices, double jitter_scale, at::Tensor inside) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_f32(node_bbox_min, "node_bbox_min"); + check_f32(node_bbox_max, "node_bbox_max"); + check_i32(node_left, "node_left"); + check_i32(node_right, "node_right"); + check_i32(node_start, "node_start"); + check_i32(node_count, "node_count"); + check_i32(triangle_indices, "triangle_indices"); + check_i32(inside, "inside"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + query_parity_sign_bvh_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + node_bbox_min.data_ptr(), node_bbox_max.data_ptr(), + node_left.data_ptr(), node_right.data_ptr(), + node_start.data_ptr(), node_count.data_ptr(), + triangle_indices.data_ptr(), static_cast(jitter_scale), point_count, + inside.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void backward_mesh_unsigned_distance_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, + at::Tensor grad_unsigned_distance, at::Tensor grad_triangles, at::Tensor grad_points) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_i32(closest_triangle_index, "closest_triangle_index"); + check_f32(grad_unsigned_distance, "grad_unsigned_distance"); + check_f32(grad_triangles, "grad_triangles"); + check_f32(grad_points, "grad_points"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + backward_unsigned_distance_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + closest_triangle_index.data_ptr(), grad_unsigned_distance.data_ptr(), + point_count, grad_triangles.data_ptr(), grad_points.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +void backward_mesh_distance_and_winding_cuda( + at::Tensor triangles, at::Tensor points, at::Tensor closest_triangle_index, + at::Tensor grad_unsigned_distance, at::Tensor grad_winding_angle, + at::Tensor grad_triangles, at::Tensor grad_points) { + check_f32(triangles, "triangles"); + check_f32(points, "points"); + check_i32(closest_triangle_index, "closest_triangle_index"); + check_f32(grad_unsigned_distance, "grad_unsigned_distance"); + check_f32(grad_winding_angle, "grad_winding_angle"); + check_f32(grad_triangles, "grad_triangles"); + check_f32(grad_points, "grad_points"); + const c10::cuda::CUDAGuard guard(points.device()); + const int point_count = static_cast(points.size(0)); + if (point_count == 0) { + return; + } + backward_distance_and_winding_kernel<<>>( + triangles.data_ptr(), points.data_ptr(), + closest_triangle_index.data_ptr(), grad_unsigned_distance.data_ptr(), + grad_winding_angle.data_ptr(), static_cast(triangles.size(0)), point_count, + grad_triangles.data_ptr(), grad_points.data_ptr()); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} diff --git a/witwin/core/geometry/mesh_sdf.py b/witwin/core/geometry/mesh_sdf.py index e5f4ffe..688991d 100644 --- a/witwin/core/geometry/mesh_sdf.py +++ b/witwin/core/geometry/mesh_sdf.py @@ -1,18 +1,16 @@ -"""Torch-native and Slang-backed triangle-mesh distance helpers.""" +"""Torch-native and native-CUDA triangle-mesh distance helpers.""" from __future__ import annotations import math -import os -import sys -from pathlib import Path -from typing import Any import numpy as np import torch -_MESH_SDF_MODULE_CACHE: dict[str, Any] = {} +_MESH_SDF_MODULE_UNSET = object() +_MESH_SDF_MODULE: object | None = _MESH_SDF_MODULE_UNSET +_MESH_SDF_BUILD_ERROR: str | None = None _MESH_SDF_BVH_LEAF_SIZE = 8 _MESH_SDF_BVH_MIN_TRIANGLES = 1024 @@ -25,42 +23,44 @@ def _safe_signed_denom(value: torch.Tensor, eps: float) -> torch.Tensor: return torch.where(value >= 0.0, value.clamp_min(eps), value.clamp_max(-eps)) -def _ensure_current_env_on_path() -> None: - scripts_dir = os.path.join(os.path.dirname(sys.executable), "Scripts") - if not os.path.isdir(scripts_dir): - return - current_path = os.environ.get("PATH", "") - path_entries = current_path.split(os.pathsep) if current_path else [] - if scripts_dir not in path_entries: - os.environ["PATH"] = scripts_dir + os.pathsep + current_path - - def _get_mesh_sdf_module(): - try: - import slangtorch - except ImportError: - return None - - _ensure_current_env_on_path() - slang_path = str(Path(__file__).resolve().parents[1] / "mesh_sdf.slang") - module = _MESH_SDF_MODULE_CACHE.get(slang_path) - if module is None: - module = slangtorch.loadModule(slang_path) - _MESH_SDF_MODULE_CACHE[slang_path] = module - return module + """Return the native CUDA mesh-SDF module, or None if it cannot be built. + + The module is compiled once via ``torch.utils.cpp_extension`` and cached. + A build failure (missing toolchain, no CUDA) is recorded and downgraded to + ``None`` so callers transparently fall back to the Torch-native path. + """ + global _MESH_SDF_MODULE, _MESH_SDF_BUILD_ERROR + if _MESH_SDF_MODULE is _MESH_SDF_MODULE_UNSET: + if not torch.cuda.is_available(): + _MESH_SDF_MODULE = None + else: + try: + from .cuda import get_native_mesh_sdf_module + + _MESH_SDF_MODULE = get_native_mesh_sdf_module() + except Exception as exc: # noqa: BLE001 - degrade to Torch fallback + _MESH_SDF_BUILD_ERROR = f"{type(exc).__name__}: {exc}" + _MESH_SDF_MODULE = None + return _MESH_SDF_MODULE + + +def _mesh_sdf_available() -> bool: + return _get_mesh_sdf_module() is not None -def _slang_mesh_sdf_available() -> bool: - return _get_mesh_sdf_module() is not None +def _mesh_sdf_build_error() -> str | None: + """Return the last native-CUDA build error message, if any.""" + return _MESH_SDF_BUILD_ERROR -def _should_use_slang(points: torch.Tensor, vertices: torch.Tensor) -> bool: +def _should_use_native(points: torch.Tensor, vertices: torch.Tensor) -> bool: return ( points.device.type == "cuda" and vertices.device.type == "cuda" and points.dtype == torch.float32 and vertices.dtype == torch.float32 - and _slang_mesh_sdf_available() + and _mesh_sdf_available() ) @@ -340,13 +340,13 @@ def _triangle_mesh_winding_angle_torch( ) -def _launch_slang_unsigned_distance( +def _launch_native_unsigned_distance( triangles: torch.Tensor, points: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: empty_f = torch.empty((0,), device=points.device, dtype=torch.float32) empty_i = torch.empty((0,), device=points.device, dtype=torch.int32) @@ -371,14 +371,14 @@ def _launch_slang_unsigned_distance( return distances, closest -def _launch_slang_unsigned_distance_bvh( +def _launch_native_unsigned_distance_bvh( triangles: torch.Tensor, points: torch.Tensor, bvh: dict[str, torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: empty_f = torch.empty((0,), device=points.device, dtype=torch.float32) empty_i = torch.empty((0,), device=points.device, dtype=torch.int32) @@ -410,7 +410,7 @@ def _launch_slang_unsigned_distance_bvh( return distances, closest -def _launch_slang_parity_sign_bvh( +def _launch_native_parity_sign_bvh( triangles: torch.Tensor, points: torch.Tensor, bvh: dict[str, torch.Tensor], @@ -419,7 +419,7 @@ def _launch_slang_parity_sign_bvh( ) -> torch.Tensor: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: return torch.empty((0,), device=points.device, dtype=torch.int32) @@ -444,13 +444,13 @@ def _launch_slang_parity_sign_bvh( return inside -def _launch_slang_distance_and_winding( +def _launch_native_distance_and_winding( triangles: torch.Tensor, points: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: empty_f = torch.empty((0,), device=points.device, dtype=torch.float32) empty_i = torch.empty((0,), device=points.device, dtype=torch.int32) @@ -483,9 +483,9 @@ def triangle_mesh_unsigned_distance_static_bvh( bvh: dict[str, torch.Tensor] | None, ) -> torch.Tensor: if bvh is None: - distances, _closest = _launch_slang_unsigned_distance(triangles, points) + distances, _closest = _launch_native_unsigned_distance(triangles, points) else: - distances, _closest = _launch_slang_unsigned_distance_bvh(triangles, points, bvh) + distances, _closest = _launch_native_unsigned_distance_bvh(triangles, points, bvh) return distances.to(dtype=points.dtype) @@ -495,15 +495,15 @@ def triangle_mesh_signed_distance_static_bvh( bvh: dict[str, torch.Tensor] | None, ) -> torch.Tensor: if bvh is None: - distances, _winding_angles, _closest = _launch_slang_distance_and_winding(triangles, points) + distances, _winding_angles, _closest = _launch_native_distance_and_winding(triangles, points) parity_inside = torch.signbit(_signed_distance_from_unsigned_and_winding( distances.to(dtype=points.dtype), _winding_angles.to(dtype=points.dtype), winding_beta=0.5, )).to(dtype=torch.int32) else: - distances, _closest = _launch_slang_unsigned_distance_bvh(triangles, points, bvh) - parity_inside = _launch_slang_parity_sign_bvh(triangles, points, bvh) + distances, _closest = _launch_native_unsigned_distance_bvh(triangles, points, bvh) + parity_inside = _launch_native_parity_sign_bvh(triangles, points, bvh) sign = torch.where( parity_inside.to(dtype=torch.bool), @@ -513,7 +513,7 @@ def triangle_mesh_signed_distance_static_bvh( return distances.to(dtype=points.dtype) * sign -def _launch_slang_unsigned_backward( +def _launch_native_unsigned_backward( triangles: torch.Tensor, points: torch.Tensor, closest_triangle_index: torch.Tensor, @@ -521,7 +521,7 @@ def _launch_slang_unsigned_backward( ) -> tuple[torch.Tensor, torch.Tensor]: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: return torch.zeros_like(triangles), torch.zeros_like(points) @@ -542,7 +542,7 @@ def _launch_slang_unsigned_backward( return grad_triangles, grad_points -def _launch_slang_distance_and_winding_backward( +def _launch_native_distance_and_winding_backward( triangles: torch.Tensor, points: torch.Tensor, closest_triangle_index: torch.Tensor, @@ -551,7 +551,7 @@ def _launch_slang_distance_and_winding_backward( ) -> tuple[torch.Tensor, torch.Tensor]: module = _get_mesh_sdf_module() if module is None: - raise RuntimeError("Mesh SDF Slang module is not available.") + raise RuntimeError("Native CUDA mesh SDF module is not available.") if points.shape[0] == 0: return torch.zeros_like(triangles), torch.zeros_like(points) @@ -576,14 +576,14 @@ def _launch_slang_distance_and_winding_backward( class _TriangleMeshUnsignedDistanceFunction(torch.autograd.Function): @staticmethod def forward(ctx, triangles: torch.Tensor, points: torch.Tensor) -> torch.Tensor: - distances, closest_triangle_index = _launch_slang_unsigned_distance(triangles, points) + distances, closest_triangle_index = _launch_native_unsigned_distance(triangles, points) ctx.save_for_backward(triangles.detach(), points.detach(), closest_triangle_index) return distances @staticmethod def backward(ctx, grad_output: torch.Tensor): triangles, points, closest_triangle_index = ctx.saved_tensors - grad_triangles, grad_points = _launch_slang_unsigned_backward( + grad_triangles, grad_points = _launch_native_unsigned_backward( triangles, points, closest_triangle_index, @@ -595,7 +595,7 @@ def backward(ctx, grad_output: torch.Tensor): class _TriangleMeshWindingAngleFunction(torch.autograd.Function): @staticmethod def forward(ctx, triangles: torch.Tensor, points: torch.Tensor) -> torch.Tensor: - _unsigned_distance, winding_angle, closest_triangle_index = _launch_slang_distance_and_winding(triangles, points) + _unsigned_distance, winding_angle, closest_triangle_index = _launch_native_distance_and_winding(triangles, points) ctx.save_for_backward(triangles.detach(), points.detach(), closest_triangle_index) return winding_angle @@ -603,7 +603,7 @@ def forward(ctx, triangles: torch.Tensor, points: torch.Tensor) -> torch.Tensor: def backward(ctx, grad_output: torch.Tensor): triangles, points, closest_triangle_index = ctx.saved_tensors zeros = torch.zeros_like(grad_output, device=triangles.device, dtype=torch.float32) - grad_triangles, grad_points = _launch_slang_distance_and_winding_backward( + grad_triangles, grad_points = _launch_native_distance_and_winding_backward( triangles, points, closest_triangle_index, @@ -616,7 +616,7 @@ def backward(ctx, grad_output: torch.Tensor): class _TriangleMeshSignedDistanceFunction(torch.autograd.Function): @staticmethod def forward(ctx, triangles: torch.Tensor, points: torch.Tensor, winding_beta: float) -> torch.Tensor: - unsigned_distance, winding_angle, closest_triangle_index = _launch_slang_distance_and_winding(triangles, points) + unsigned_distance, winding_angle, closest_triangle_index = _launch_native_distance_and_winding(triangles, points) signed_distance = _signed_distance_from_unsigned_and_winding(unsigned_distance, winding_angle, winding_beta) ctx.winding_beta = float(winding_beta) ctx.save_for_backward( @@ -642,7 +642,7 @@ def backward(ctx, grad_output: torch.Tensor): * (-(1.0 - tanh_shifted.square()) / beta) * torch.sign(winding_angle) ) - grad_triangles, grad_points = _launch_slang_distance_and_winding_backward( + grad_triangles, grad_points = _launch_native_distance_and_winding_backward( triangles, points, closest_triangle_index, @@ -662,10 +662,10 @@ def triangle_mesh_unsigned_distance( _triangles: torch.Tensor | None = None, ) -> torch.Tensor: triangles = _indexed_triangles(vertices, faces) if _triangles is None else _triangles - if _should_use_slang(points, vertices): + if _should_use_native(points, vertices): if torch.is_grad_enabled() and (triangles.requires_grad or points.requires_grad): return _TriangleMeshUnsignedDistanceFunction.apply(triangles, points.contiguous()) - primal, _closest_triangle_index = _launch_slang_unsigned_distance(triangles, points) + primal, _closest_triangle_index = _launch_native_unsigned_distance(triangles, points) return primal.to(dtype=points.dtype) return _triangle_mesh_unsigned_distance_torch_from_triangles( @@ -686,10 +686,10 @@ def triangle_mesh_winding_angle( _triangles: torch.Tensor | None = None, ) -> torch.Tensor: triangles = _indexed_triangles(vertices, faces) if _triangles is None else _triangles - if _should_use_slang(points, vertices): + if _should_use_native(points, vertices): if torch.is_grad_enabled() and (triangles.requires_grad or points.requires_grad): return _TriangleMeshWindingAngleFunction.apply(triangles, points.contiguous()) - _unsigned_distance, winding_angle, _closest_triangle_index = _launch_slang_distance_and_winding(triangles, points) + _unsigned_distance, winding_angle, _closest_triangle_index = _launch_native_distance_and_winding(triangles, points) return winding_angle.to(dtype=points.dtype) return _triangle_mesh_winding_angle_torch_from_triangles( @@ -737,10 +737,10 @@ def triangle_mesh_smooth_signed_distance( _triangles: torch.Tensor | None = None, ) -> torch.Tensor: triangles = _indexed_triangles(vertices, faces) if _triangles is None else _triangles - if _should_use_slang(points, vertices): + if _should_use_native(points, vertices): if torch.is_grad_enabled() and (triangles.requires_grad or points.requires_grad): return _TriangleMeshSignedDistanceFunction.apply(triangles, points.contiguous(), float(winding_beta)) - primal_unsigned_distance, primal_winding_angle, _closest_triangle_index = _launch_slang_distance_and_winding( + primal_unsigned_distance, primal_winding_angle, _closest_triangle_index = _launch_native_distance_and_winding( triangles, points, ) diff --git a/witwin/core/mesh_sdf.slang b/witwin/core/mesh_sdf.slang deleted file mode 100644 index b52d458..0000000 --- a/witwin/core/mesh_sdf.slang +++ /dev/null @@ -1,609 +0,0 @@ -float3 readPoint(TensorView points, int pointIndex) -{ - return float3( - points[pointIndex, 0], - points[pointIndex, 1], - points[pointIndex, 2] - ); -} - -float3 readTriangleVertex(TensorView triangles, int triangleIndex, int vertexIndex) -{ - return float3( - triangles[triangleIndex, vertexIndex, 0], - triangles[triangleIndex, vertexIndex, 1], - triangles[triangleIndex, vertexIndex, 2] - ); -} - -[Differentiable] -float safeSignedDenom(float value, float eps) -{ - return value >= 0.0 ? max(value, eps) : min(value, -eps); -} - -[Differentiable] -float pointTriangleDistanceSquared(float3 point, float3 a, float3 b, float3 c) -{ - float3 ab = b - a; - float3 ac = c - a; - float3 ap = point - a; - - float d1 = dot(ab, ap); - float d2 = dot(ac, ap); - if (d1 <= 0.0 && d2 <= 0.0) - { - return dot(ap, ap); - } - - float3 bp = point - b; - float d3 = dot(ab, bp); - float d4 = dot(ac, bp); - if (d3 >= 0.0 && d4 <= d3) - { - return dot(bp, bp); - } - - float vc = d1 * d4 - d3 * d2; - if (vc <= 0.0 && d1 >= 0.0 && d3 <= 0.0) - { - float v = d1 / (d1 - d3); - float3 closest = a + v * ab; - float3 delta = point - closest; - return dot(delta, delta); - } - - float3 cp = point - c; - float d5 = dot(ab, cp); - float d6 = dot(ac, cp); - if (d6 >= 0.0 && d5 <= d6) - { - return dot(cp, cp); - } - - float vb = d5 * d2 - d1 * d6; - if (vb <= 0.0 && d2 >= 0.0 && d6 <= 0.0) - { - float w = d2 / safeSignedDenom(d2 - d6, 1.0e-12); - float3 closest = a + w * ac; - float3 delta = point - closest; - return dot(delta, delta); - } - - float va = d3 * d6 - d5 * d4; - if (va <= 0.0 && (d4 - d3) >= 0.0 && (d5 - d6) >= 0.0) - { - float3 bc = c - b; - float w = (d4 - d3) / safeSignedDenom((d4 - d3) + (d5 - d6), 1.0e-12); - float3 closest = b + w * bc; - float3 delta = point - closest; - return dot(delta, delta); - } - - float denom = safeSignedDenom(va + vb + vc, 1.0e-12); - float vFace = vb / denom; - float wFace = vc / denom; - float3 closest = a + ab * vFace + ac * wFace; - float3 delta = point - closest; - return dot(delta, delta); -} - -[Differentiable] -float pointTriangleUnsignedDistance(float3 point, float3 a, float3 b, float3 c) -{ - const float smoothEps = 1.0e-12; - float distanceSquared = pointTriangleDistanceSquared(point, a, b, c); - return sqrt(max(distanceSquared, 0.0) + smoothEps) - sqrt(smoothEps); -} - -[Differentiable] -float triangleSolidAngle(float3 point, float3 a0, float3 b0, float3 c0) -{ - const float eps = 1.0e-12; - - float3 a = a0 - point; - float3 b = b0 - point; - float3 c = c0 - point; - - float la = max(length(a), eps); - float lb = max(length(b), eps); - float lc = max(length(c), eps); - float numerator = dot(a, cross(b, c)); - float denominator = - la * lb * lc - + dot(a, b) * lc - + dot(b, c) * la - + dot(c, a) * lb; - return 2.0 * atan2(numerator, safeSignedDenom(denominator, eps)); -} - -void addTriangleGradient( - TensorView gradTriangles, - int triangleIndex, - float3 gradA, - float3 gradB, - float3 gradC -) -{ - float oldValue = 0.0; - - gradTriangles.InterlockedAdd(uint3(triangleIndex, 0, 0), gradA.x, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 0, 1), gradA.y, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 0, 2), gradA.z, oldValue); - - gradTriangles.InterlockedAdd(uint3(triangleIndex, 1, 0), gradB.x, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 1, 1), gradB.y, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 1, 2), gradB.z, oldValue); - - gradTriangles.InterlockedAdd(uint3(triangleIndex, 2, 0), gradC.x, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 2, 1), gradC.y, oldValue); - gradTriangles.InterlockedAdd(uint3(triangleIndex, 2, 2), gradC.z, oldValue); -} - -float3 readNodeBBoxMin(TensorView bboxMin, int nodeIndex) -{ - return float3( - bboxMin[nodeIndex, 0], - bboxMin[nodeIndex, 1], - bboxMin[nodeIndex, 2] - ); -} - -float3 readNodeBBoxMax(TensorView bboxMax, int nodeIndex) -{ - return float3( - bboxMax[nodeIndex, 0], - bboxMax[nodeIndex, 1], - bboxMax[nodeIndex, 2] - ); -} - -float pointAabbDistanceSquared(float3 point, float3 bboxMin, float3 bboxMax) -{ - float dx = point.x < bboxMin.x ? (bboxMin.x - point.x) : (point.x > bboxMax.x ? point.x - bboxMax.x : 0.0); - float dy = point.y < bboxMin.y ? (bboxMin.y - point.y) : (point.y > bboxMax.y ? point.y - bboxMax.y : 0.0); - float dz = point.z < bboxMin.z ? (bboxMin.z - point.z) : (point.z > bboxMax.z ? point.z - bboxMax.z : 0.0); - return dx * dx + dy * dy + dz * dz; -} - -bool rayAabbHit(float3 origin, float3 direction, float3 bboxMin, float3 bboxMax) -{ - const float eps = 1.0e-12; - float3 invDirection = float3( - abs(direction.x) > eps ? 1.0 / direction.x : 1.0e30, - abs(direction.y) > eps ? 1.0 / direction.y : 1.0e30, - abs(direction.z) > eps ? 1.0 / direction.z : 1.0e30 - ); - - float tx1 = (bboxMin.x - origin.x) * invDirection.x; - float tx2 = (bboxMax.x - origin.x) * invDirection.x; - float tmin = min(tx1, tx2); - float tmax = max(tx1, tx2); - - float ty1 = (bboxMin.y - origin.y) * invDirection.y; - float ty2 = (bboxMax.y - origin.y) * invDirection.y; - tmin = max(tmin, min(ty1, ty2)); - tmax = min(tmax, max(ty1, ty2)); - - float tz1 = (bboxMin.z - origin.z) * invDirection.z; - float tz2 = (bboxMax.z - origin.z) * invDirection.z; - tmin = max(tmin, min(tz1, tz2)); - tmax = min(tmax, max(tz1, tz2)); - - return tmax >= max(tmin, 0.0); -} - -bool rayTriangleHit(float3 origin, float3 direction, float3 a, float3 b, float3 c, out float t) -{ - const float epsilon = 1.0e-6; - - float3 edge1 = b - a; - float3 edge2 = c - a; - float3 h = cross(direction, edge2); - float determinant = dot(edge1, h); - - if (abs(determinant) < epsilon) - { - t = 0.0; - return false; - } - - float invDeterminant = 1.0 / determinant; - float3 s = origin - a; - float u = invDeterminant * dot(s, h); - if (u < 0.0 || u > 1.0) - { - t = 0.0; - return false; - } - - float3 q = cross(s, edge1); - float v = invDeterminant * dot(direction, q); - if (v < 0.0 || (u + v) > 1.0) - { - t = 0.0; - return false; - } - - t = invDeterminant * dot(edge2, q); - return t > epsilon; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void queryMeshUnsignedDistance( - TensorView triangles, - TensorView points, - TensorView unsignedDistance, - TensorView closestTriangleIndex -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= unsignedDistance.size(0)) - { - return; - } - - const float smoothEps = 1.0e-12; - float3 point = readPoint(points, pointIndex); - float minDistanceSquared = 1.0e30; - int closestTriangle = -1; - - for (int triangleIndex = 0; triangleIndex < triangles.size(0); ++triangleIndex) - { - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - float distanceSquared = pointTriangleDistanceSquared(point, a, b, c); - if (distanceSquared < minDistanceSquared) - { - minDistanceSquared = distanceSquared; - closestTriangle = triangleIndex; - } - } - - unsignedDistance[pointIndex] = sqrt(max(minDistanceSquared, 0.0) + smoothEps) - sqrt(smoothEps); - closestTriangleIndex[pointIndex] = closestTriangle; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void queryMeshDistanceAndWinding( - TensorView triangles, - TensorView points, - TensorView unsignedDistance, - TensorView windingAngle, - TensorView closestTriangleIndex -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= unsignedDistance.size(0)) - { - return; - } - - const float smoothEps = 1.0e-12; - float3 point = readPoint(points, pointIndex); - float minDistanceSquared = 1.0e30; - float totalAngle = 0.0; - int closestTriangle = -1; - - for (int triangleIndex = 0; triangleIndex < triangles.size(0); ++triangleIndex) - { - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - float distanceSquared = pointTriangleDistanceSquared(point, a, b, c); - if (distanceSquared < minDistanceSquared) - { - minDistanceSquared = distanceSquared; - closestTriangle = triangleIndex; - } - totalAngle += triangleSolidAngle(point, a, b, c); - } - - unsignedDistance[pointIndex] = sqrt(max(minDistanceSquared, 0.0) + smoothEps) - sqrt(smoothEps); - windingAngle[pointIndex] = totalAngle; - closestTriangleIndex[pointIndex] = closestTriangle; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void queryMeshUnsignedDistanceBVH( - TensorView triangles, - TensorView points, - TensorView nodeBBoxMin, - TensorView nodeBBoxMax, - TensorView nodeLeft, - TensorView nodeRight, - TensorView nodeStart, - TensorView nodeCount, - TensorView triangleIndices, - TensorView unsignedDistance, - TensorView closestTriangleIndex -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= unsignedDistance.size(0)) - { - return; - } - - const float smoothEps = 1.0e-12; - const int maxStackDepth = 64; - int stack[maxStackDepth]; - int stackSize = 0; - stack[stackSize++] = 0; - - float3 point = readPoint(points, pointIndex); - float minDistanceSquared = 1.0e30; - int closestTriangle = -1; - - while (stackSize > 0) - { - int nodeIndex = stack[--stackSize]; - float3 bboxMin = readNodeBBoxMin(nodeBBoxMin, nodeIndex); - float3 bboxMax = readNodeBBoxMax(nodeBBoxMax, nodeIndex); - float nodeDistanceSquared = pointAabbDistanceSquared(point, bboxMin, bboxMax); - if (nodeDistanceSquared > minDistanceSquared) - { - continue; - } - - int count = nodeCount[nodeIndex]; - if (count > 0) - { - int start = nodeStart[nodeIndex]; - for (int offset = 0; offset < count; ++offset) - { - int triangleIndex = triangleIndices[start + offset]; - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - float distanceSquared = pointTriangleDistanceSquared(point, a, b, c); - if (distanceSquared < minDistanceSquared) - { - minDistanceSquared = distanceSquared; - closestTriangle = triangleIndex; - } - } - continue; - } - - int leftChild = nodeLeft[nodeIndex]; - int rightChild = nodeRight[nodeIndex]; - if (leftChild < 0 || rightChild < 0) - { - continue; - } - - float leftDistanceSquared = pointAabbDistanceSquared(point, readNodeBBoxMin(nodeBBoxMin, leftChild), readNodeBBoxMax(nodeBBoxMax, leftChild)); - float rightDistanceSquared = pointAabbDistanceSquared(point, readNodeBBoxMin(nodeBBoxMin, rightChild), readNodeBBoxMax(nodeBBoxMax, rightChild)); - - if (leftDistanceSquared < rightDistanceSquared) - { - if (stackSize + 2 <= maxStackDepth) - { - stack[stackSize++] = rightChild; - stack[stackSize++] = leftChild; - } - } - else - { - if (stackSize + 2 <= maxStackDepth) - { - stack[stackSize++] = leftChild; - stack[stackSize++] = rightChild; - } - } - } - - unsignedDistance[pointIndex] = sqrt(max(minDistanceSquared, 0.0) + smoothEps) - sqrt(smoothEps); - closestTriangleIndex[pointIndex] = closestTriangle; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void queryMeshParitySignBVH( - TensorView triangles, - TensorView points, - TensorView nodeBBoxMin, - TensorView nodeBBoxMax, - TensorView nodeLeft, - TensorView nodeRight, - TensorView nodeStart, - TensorView nodeCount, - TensorView triangleIndices, - float jitterScale, - TensorView inside -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= inside.size(0)) - { - return; - } - - const int maxStackDepth = 64; - int stack[maxStackDepth]; - int stackSize = 0; - stack[stackSize++] = 0; - - float3 point = readPoint(points, pointIndex); - float3 direction = float3(1.0, 0.0, 0.0); - float3 origin = point + float3(0.0, jitterScale, jitterScale * 1.618034); - int intersections = 0; - - while (stackSize > 0) - { - int nodeIndex = stack[--stackSize]; - float3 bboxMin = readNodeBBoxMin(nodeBBoxMin, nodeIndex); - float3 bboxMax = readNodeBBoxMax(nodeBBoxMax, nodeIndex); - if (!rayAabbHit(origin, direction, bboxMin, bboxMax)) - { - continue; - } - - int count = nodeCount[nodeIndex]; - if (count > 0) - { - int start = nodeStart[nodeIndex]; - for (int offset = 0; offset < count; ++offset) - { - int triangleIndex = triangleIndices[start + offset]; - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - float t; - if (rayTriangleHit(origin, direction, a, b, c, t)) - { - intersections += 1; - } - } - continue; - } - - int leftChild = nodeLeft[nodeIndex]; - int rightChild = nodeRight[nodeIndex]; - if (leftChild >= 0 && stackSize < maxStackDepth) - { - stack[stackSize++] = leftChild; - } - if (rightChild >= 0 && stackSize < maxStackDepth) - { - stack[stackSize++] = rightChild; - } - } - - inside[pointIndex] = (intersections & 1) != 0 ? 1 : 0; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void backwardMeshUnsignedDistance( - TensorView triangles, - TensorView points, - TensorView closestTriangleIndex, - TensorView gradUnsignedDistance, - TensorView gradTriangles, - TensorView gradPoints -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= gradUnsignedDistance.size(0)) - { - return; - } - - float gradOutput = gradUnsignedDistance[pointIndex]; - if (gradOutput == 0.0) - { - return; - } - - int triangleIndex = closestTriangleIndex[pointIndex]; - if (triangleIndex < 0) - { - return; - } - - float3 point = readPoint(points, pointIndex); - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - - DifferentialPair dpPoint = DifferentialPair(point, float3(0.0)); - DifferentialPair dpA = DifferentialPair(a, float3(0.0)); - DifferentialPair dpB = DifferentialPair(b, float3(0.0)); - DifferentialPair dpC = DifferentialPair(c, float3(0.0)); - bwd_diff(pointTriangleUnsignedDistance)(dpPoint, dpA, dpB, dpC, gradOutput); - - addTriangleGradient(gradTriangles, triangleIndex, dpA.d, dpB.d, dpC.d); - gradPoints[pointIndex, 0] = dpPoint.d.x; - gradPoints[pointIndex, 1] = dpPoint.d.y; - gradPoints[pointIndex, 2] = dpPoint.d.z; -} - -[AutoPyBindCUDA] -[CUDAKernel] -void backwardMeshDistanceAndWinding( - TensorView triangles, - TensorView points, - TensorView closestTriangleIndex, - TensorView gradUnsignedDistance, - TensorView gradWindingAngle, - TensorView gradTriangles, - TensorView gradPoints -) -{ - uint3 dispatchIdx = cudaThreadIdx() + cudaBlockIdx() * cudaBlockDim(); - int pointIndex = int(dispatchIdx.x); - - if (pointIndex >= gradUnsignedDistance.size(0)) - { - return; - } - - float gradUnsigned = gradUnsignedDistance[pointIndex]; - float gradAngle = gradWindingAngle[pointIndex]; - if (gradUnsigned == 0.0 && gradAngle == 0.0) - { - return; - } - - float3 point = readPoint(points, pointIndex); - float3 pointGradient = float3(0.0); - - if (gradAngle != 0.0) - { - for (int triangleIndex = 0; triangleIndex < triangles.size(0); ++triangleIndex) - { - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - - DifferentialPair dpPoint = DifferentialPair(point, float3(0.0)); - DifferentialPair dpA = DifferentialPair(a, float3(0.0)); - DifferentialPair dpB = DifferentialPair(b, float3(0.0)); - DifferentialPair dpC = DifferentialPair(c, float3(0.0)); - bwd_diff(triangleSolidAngle)(dpPoint, dpA, dpB, dpC, gradAngle); - - addTriangleGradient(gradTriangles, triangleIndex, dpA.d, dpB.d, dpC.d); - pointGradient += dpPoint.d; - } - } - - if (gradUnsigned != 0.0) - { - int triangleIndex = closestTriangleIndex[pointIndex]; - if (triangleIndex >= 0) - { - float3 a = readTriangleVertex(triangles, triangleIndex, 0); - float3 b = readTriangleVertex(triangles, triangleIndex, 1); - float3 c = readTriangleVertex(triangles, triangleIndex, 2); - - DifferentialPair dpPoint = DifferentialPair(point, float3(0.0)); - DifferentialPair dpA = DifferentialPair(a, float3(0.0)); - DifferentialPair dpB = DifferentialPair(b, float3(0.0)); - DifferentialPair dpC = DifferentialPair(c, float3(0.0)); - bwd_diff(pointTriangleUnsignedDistance)(dpPoint, dpA, dpB, dpC, gradUnsigned); - - addTriangleGradient(gradTriangles, triangleIndex, dpA.d, dpB.d, dpC.d); - pointGradient += dpPoint.d; - } - } - - gradPoints[pointIndex, 0] = pointGradient.x; - gradPoints[pointIndex, 1] = pointGradient.y; - gradPoints[pointIndex, 2] = pointGradient.z; -}