From a111e6eb18ea1d119bd110f5dcc30d8fad8e6b45 Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Thu, 9 Jul 2026 04:00:49 -0700 Subject: [PATCH 1/6] smpl: vectorized differentiable LBS with finite gradient at zero pose Replace smplpytorch's per-call forward with a batched LBS (~10x faster, ~2ms vs ~21ms) that is numerically equivalent, plus a Rodrigues that offsets the axis-angle vector rather than its norm so the gradient at exactly-zero pose stays finite. Cache SMPL faces per (gender, root, device). Enables stable pose optimization starting from T-pose. --- witwin/core/geometry/smpl.py | 71 ++++++++++++++++++++++++++++++++++-- 1 file changed, 68 insertions(+), 3 deletions(-) diff --git a/witwin/core/geometry/smpl.py b/witwin/core/geometry/smpl.py index 79d65aa..c3c2b24 100644 --- a/witwin/core/geometry/smpl.py +++ b/witwin/core/geometry/smpl.py @@ -132,6 +132,7 @@ def _setup_smpl_compat(): _SMPL_LAYER_CACHE: dict[tuple[str, str, str], Any] = {} +_SMPL_FACES_CACHE: dict[tuple[str, str, str], np.ndarray] = {} def _resolve_scene_device(device: str | None) -> str: @@ -166,6 +167,66 @@ def _get_smpl_layer(*, gender: str, model_root: str, device: str): return layer +def _axis_angle_matrices(rvecs: torch.Tensor) -> torch.Tensor: + """Differentiable Rodrigues for (J, 3) axis-angle vectors -> (J, 3, 3).""" + # smplpytorch's zero-rotation guard: offset the vector, not the norm, so + # the gradient at exactly zero pose stays finite. + safe = rvecs + 1e-8 + angle = safe.norm(dim=1, keepdim=True) + axis = safe / angle + x, y, z = axis.unbind(-1) + zero = torch.zeros_like(x) + skew = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).reshape(-1, 3, 3) + eye = torch.eye(3, dtype=rvecs.dtype, device=rvecs.device).expand_as(skew) + sin = torch.sin(angle).unsqueeze(-1) + cos = torch.cos(angle).unsqueeze(-1) + return eye + sin * skew + (1.0 - cos) * (skew @ skew) + + +def _fast_smpl_forward(layer, pose: torch.Tensor, betas: torch.Tensor): + """Vectorized SMPL LBS equivalent to smplpytorch's forward (batch 1). + + smplpytorch's Python implementation costs ~21 ms per call; this batched + version is ~2 ms and numerically matches it (verified by diag_fast_lbs). + """ + device = pose.device + v_template = layer.th_v_template[0] + joints_rest = layer.th_J_regressor @ ( + v_template + torch.einsum("vdk,k->vd", layer.th_shapedirs, betas.reshape(-1)) + ) + v_shaped = v_template + torch.einsum("vdk,k->vd", layer.th_shapedirs, betas.reshape(-1)) + + rotations = _axis_angle_matrices(pose.reshape(24, 3)) + eye = torch.eye(3, dtype=pose.dtype, device=device) + pose_feature = (rotations[1:] - eye).reshape(-1) + v_posed = v_shaped + torch.einsum("vdk,k->vd", layer.th_posedirs, pose_feature) + + parents = layer.kintree_parents + relative = [joints_rest[0]] + for j in range(1, 24): + relative.append(joints_rest[j] - joints_rest[parents[j]]) + + bottom = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=pose.dtype, device=device) + transforms: list[torch.Tensor] = [] + for j in range(24): + local = torch.cat( + [torch.cat([rotations[j], relative[j].reshape(3, 1)], dim=1), bottom], dim=0 + ) + transforms.append(local if j == 0 else transforms[parents[j]] @ local) + global_transforms = torch.stack(transforms) + + posed_joints = global_transforms[:, :3, 3] + corrected_t = posed_joints - torch.einsum( + "jab,jb->ja", global_transforms[:, :3, :3], joints_rest + ) + skin_rot = torch.einsum("vj,jab->vab", layer.th_weights, global_transforms[:, :3, :3]) + skin_t = layer.th_weights @ corrected_t + vertices = torch.einsum("vab,vb->va", skin_rot, v_posed) + skin_t + + center = posed_joints[0] + return (vertices - center).unsqueeze(0), (posed_joints - center).unsqueeze(0) + + class SMPLBody(GeometryBase): """Differentiable SMPL geometry with position and rotation.""" @@ -210,11 +271,15 @@ def _evaluate(self, *, device: str): shape_tensor = self.shape.to(device=device, dtype=torch.float32).view(1, -1) if shape_tensor.requires_grad: shape_tensor = shape_tensor + 1e-8 - vertices, joints = layer(pose_tensor, th_betas=shape_tensor) + vertices, joints = _fast_smpl_forward(layer, pose_tensor.reshape(-1), shape_tensor.reshape(-1)) vertices = self._transform_mesh_verts(vertices[0]) joints = self._transform_mesh_verts(joints[0]) - faces = layer.th_faces.detach().cpu().numpy().astype(np.int32) - return vertices.contiguous(), np.ascontiguousarray(faces), joints.contiguous() + cache_key = (self.gender, self.model_root, device) + faces = _SMPL_FACES_CACHE.get(cache_key) + if faces is None: + faces = np.ascontiguousarray(layer.th_faces.detach().cpu().numpy().astype(np.int32)) + _SMPL_FACES_CACHE[cache_key] = faces + return vertices.contiguous(), faces, joints.contiguous() def to_mesh(self, segments=16, *, device=None): del segments From 0c158dc4b6dd5d036681f1e37c89c8bdf2fec34c Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Thu, 9 Jul 2026 17:23:46 -0700 Subject: [PATCH 2/6] core: narrow shared geometry and material contracts --- .gitignore | 3 - README.md | 10 +- pyproject.toml | 20 +- tests/test_contracts.py | 56 ++++ tests/test_dependency_contract.py | 15 + tests/test_geometry_sdf.py | 65 ---- tests/test_geometry_smoke.py | 52 --- tests/test_scene_base.py | 18 -- witwin/core/__init__.py | 25 +- witwin/core/geometry/__init__.py | 11 +- witwin/core/geometry/base.py | 12 + witwin/core/geometry/cuda/backend.py | 2 +- witwin/core/geometry/cuda/mesh_sdf_kernels.cu | 8 +- witwin/core/geometry/polygon.py | 91 ++++++ witwin/core/geometry/primitives.py | 302 +----------------- witwin/core/geometry/smpl.py | 294 ----------------- witwin/core/material.py | 157 ++++----- witwin/core/scene.py | 38 --- witwin/core/scene_to_mitsuba.py | 138 -------- 19 files changed, 260 insertions(+), 1057 deletions(-) create mode 100644 tests/test_contracts.py create mode 100644 tests/test_dependency_contract.py delete mode 100644 tests/test_scene_base.py create mode 100644 witwin/core/geometry/polygon.py delete mode 100644 witwin/core/geometry/smpl.py delete mode 100644 witwin/core/scene.py delete mode 100644 witwin/core/scene_to_mitsuba.py diff --git a/.gitignore b/.gitignore index 618b9e1..f3a7b5c 100644 --- a/.gitignore +++ b/.gitignore @@ -171,9 +171,6 @@ venv.bak/ # mkdocs documentation /site -.slangtorch_cache/ -.fdtd_runtime_*.slang -.adjoint_kernels_runtime_*.slang *.png *.npy *.npz diff --git a/README.md b/README.md index 3dcf49c..9a886d3 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # WiTwin Core -WiTwin Core is the core package of the WiTwin stack. It provides the shared foundations used across WiTwin projects, with a focus on common scene representation, geometry primitives, mesh utilities, and state management that can be reused by different simulation modules. +WiTwin Core is the shared data and geometry package of the WiTwin stack. It defines the geometry and material contracts consumed by solver packages without imposing a common solver-specific scene lifecycle. ## Get Started @@ -12,9 +12,11 @@ pip install witwin ## What It Provides -- Shared scene and geometry primitives in `witwin.core` -- Common state management for simulation-facing scene objects -- Reusable mesh, structure, and material building blocks for WiTwin solvers +- Runtime-checkable `GeometrySpec` and `MaterialSpec` contracts +- Shared analytic geometry primitives and differentiable mesh/SDF utilities +- Reusable `Material` and `Structure` values for solver-owned scenes + +Each solver owns its public `Scene` implementation. Radar-specific SMPL geometry lives in `witwin.radar.geometry`; Maxwell-specific `PolySlab` geometry lives in `witwin.maxwell.geometry`. ## Related Solvers diff --git a/pyproject.toml b/pyproject.toml index 186fe70..0ab3826 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "hatchling.build" [project] name = "witwin" -version = "0.0.2" -description = "Witwin - Multi-physics simulation platform" +version = "0.1.0" +description = "Witwin shared geometry, material, and structure contracts" readme = "README.md" requires-python = ">=3.10" license = "MIT" @@ -16,15 +16,15 @@ dependencies = [ ] [project.optional-dependencies] -maxwell = ["witwin-maxwell>=0.0.1"] -radar = ["witwin-radar>=0.0.1"] -channel = ["witwin-channel>=0.0.1"] -genesis = ["witwin-genesis>=0.0.1"] +maxwell = ["witwin-maxwell>=0.2,<0.3"] +radar = ["witwin-radar>=0.2,<0.3"] +channel = ["witwin-channel>=0.1,<0.2"] +genesis = ["witwin-genesis>=0.1,<0.2"] all = [ - "witwin-maxwell>=0.0.1", - "witwin-radar>=0.0.1", - "witwin-channel>=0.0.1", - "witwin-genesis>=0.0.1", + "witwin-maxwell>=0.2,<0.3", + "witwin-radar>=0.2,<0.3", + "witwin-channel>=0.1,<0.2", + "witwin-genesis>=0.1,<0.2", ] [tool.hatch.build.targets.wheel] diff --git a/tests/test_contracts.py b/tests/test_contracts.py new file mode 100644 index 0000000..692f86a --- /dev/null +++ b/tests/test_contracts.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import pytest +import torch + +from witwin.core import ( + FrequencyMaterialSample, + GeometrySpec, + MaterialCapabilities, + MaterialSpec, + StaticMaterialSample, + Structure, +) + + +class DuckGeometry: + kind = "duck" + + def to_mesh(self, segments: int = 16): + del segments + return torch.zeros((3, 3)), torch.tensor([[0, 1, 2]]) + + +class DuckMaterial: + name = "duck" + + def capabilities(self): + return MaterialCapabilities() + + def evaluate_static(self): + return StaticMaterialSample(eps_r=2.0) + + def evaluate_at_frequency(self, frequency: float): + del frequency + return FrequencyMaterialSample(eps_r=2.0) + + +def test_structure_accepts_protocol_conforming_values(): + geometry = DuckGeometry() + material = DuckMaterial() + + assert isinstance(geometry, GeometrySpec) + assert isinstance(material, MaterialSpec) + assert Structure(geometry=geometry, material=material).geometry is geometry + + +@pytest.mark.parametrize( + ("geometry", "material", "message"), + [ + (object(), DuckMaterial(), "GeometrySpec"), + (DuckGeometry(), object(), "MaterialSpec"), + ], +) +def test_structure_rejects_values_outside_shared_contracts(geometry, material, message): + with pytest.raises(TypeError, match=message): + Structure(geometry=geometry, material=material) diff --git a/tests/test_dependency_contract.py b/tests/test_dependency_contract.py new file mode 100644 index 0000000..96c4365 --- /dev/null +++ b/tests/test_dependency_contract.py @@ -0,0 +1,15 @@ +from __future__ import annotations + +import tomllib +from pathlib import Path + + +def test_runtime_and_optional_dependencies_do_not_include_slangtorch(): + pyproject = tomllib.loads(Path("pyproject.toml").read_text(encoding="utf-8")) + groups = [pyproject["project"]["dependencies"]] + groups.extend(pyproject["project"].get("optional-dependencies", {}).values()) + assert all( + not dependency.lower().startswith("slangtorch") + for group in groups + for dependency in group + ) diff --git a/tests/test_geometry_sdf.py b/tests/test_geometry_sdf.py index 1cbc0a1..49a255b 100644 --- a/tests/test_geometry_sdf.py +++ b/tests/test_geometry_sdf.py @@ -4,7 +4,6 @@ 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 ( _mesh_sdf_available, _triangle_mesh_smooth_signed_distance_torch, @@ -152,39 +151,6 @@ 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): @@ -214,16 +180,6 @@ 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): @@ -371,27 +327,6 @@ 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 d694e8d..38b7945 100644 --- a/tests/test_geometry_smoke.py +++ b/tests/test_geometry_smoke.py @@ -4,7 +4,6 @@ import torch from witwin.core import Box, Cylinder, Mesh, Sphere -from witwin.core.geometry import PolySlab def _grid(): @@ -37,17 +36,6 @@ 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): @@ -83,46 +71,6 @@ 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/tests/test_scene_base.py b/tests/test_scene_base.py deleted file mode 100644 index 628d496..0000000 --- a/tests/test_scene_base.py +++ /dev/null @@ -1,18 +0,0 @@ -from witwin.core.scene import SceneBase - - -def test_scene_base_uses_add_mutators_without_with_aliases(): - scene = SceneBase(device="cpu") - structure = object() - source = object() - monitor = object() - - returned = scene.add_structure(structure).add_source(source).add_monitor(monitor) - - assert returned is scene - assert scene.structures == [structure] - assert scene.sources == [source] - assert scene.monitors == [monitor] - assert not hasattr(scene, "with_structure") - assert not hasattr(scene, "with_source") - assert not hasattr(scene, "with_monitor") diff --git a/witwin/core/__init__.py b/witwin/core/__init__.py index 0a08391..c2e426d 100644 --- a/witwin/core/__init__.py +++ b/witwin/core/__init__.py @@ -1,9 +1,8 @@ -"""Witwin Core - Shared geometry, materials, and scene utilities.""" +"""Witwin Core - Shared geometry, material, and structure contracts.""" -__version__ = "0.0.1" +__version__ = "0.1.0" from .material import ( - DifferentiableMaterial, FrequencyMaterialSample, Material, MaterialCapabilities, @@ -18,11 +17,11 @@ Ellipsoid, Geometry, GeometryBase, + GeometrySpec, HollowBox, Mesh, Prism, Pyramid, - SMPLBody, Sphere, Torus, ) @@ -33,38 +32,22 @@ quat_to_rotation_matrix, quat_to_rotation_matrix_np, ) -from .scene import SceneBase -from .scene_to_mitsuba import ( - MitsubaRenderable, - MitsubaSceneHandle, - build_mitsuba_scene, - create_mitsuba_mesh, - update_mitsuba_scene_vertices, -) - __all__ = [ "GeometryBase", + "GeometrySpec", "Box", "Sphere", "Cylinder", "Cone", "Ellipsoid", "Pyramid", "Prism", "Torus", "HollowBox", "Material", - "DifferentiableMaterial", "MaterialCapabilities", "MaterialSpec", "Mesh", - "SMPLBody", "FrequencyMaterialSample", "Geometry", - "SceneBase", "StaticMaterialSample", "Structure", - "MitsubaRenderable", - "MitsubaSceneHandle", "quat_from_euler", "quat_identity", "quat_multiply", "quat_to_rotation_matrix", "quat_to_rotation_matrix_np", - "create_mitsuba_mesh", - "build_mitsuba_scene", - "update_mitsuba_scene_vertices", ] diff --git a/witwin/core/geometry/__init__.py b/witwin/core/geometry/__init__.py index 837b5de..e7152f9 100644 --- a/witwin/core/geometry/__init__.py +++ b/witwin/core/geometry/__init__.py @@ -1,14 +1,14 @@ """Shared differentiable geometry package.""" -from .base import GeometryBase +from .base import GeometryBase, GeometrySpec from .mesh import Mesh -from .primitives import Box, ComplexPolySlab, Cone, Cylinder, Ellipsoid, HollowBox, PolySlab, Prism, Pyramid, Sphere, Torus -from .smpl import SMPLBody +from .primitives import Box, Cone, Cylinder, Ellipsoid, HollowBox, Prism, Pyramid, Sphere, Torus -Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | PolySlab | ComplexPolySlab | Torus | HollowBox | Mesh | SMPLBody +Geometry = Box | Sphere | Cylinder | Cone | Ellipsoid | Pyramid | Prism | Torus | HollowBox | Mesh __all__ = [ "GeometryBase", + "GeometrySpec", "Box", "Sphere", "Cylinder", @@ -16,11 +16,8 @@ "Ellipsoid", "Pyramid", "Prism", - "PolySlab", - "ComplexPolySlab", "Torus", "HollowBox", "Mesh", - "SMPLBody", "Geometry", ] diff --git a/witwin/core/geometry/base.py b/witwin/core/geometry/base.py index b2cd58a..36a2f75 100644 --- a/witwin/core/geometry/base.py +++ b/witwin/core/geometry/base.py @@ -2,6 +2,8 @@ from __future__ import annotations +from typing import Protocol, runtime_checkable + import numpy as np import torch @@ -136,6 +138,16 @@ def occupancy_from_signed_distance(signed_distance: torch.Tensor, *, xx, yy, zz, return occupancy.clamp(0.0, 1.0) +@runtime_checkable +class GeometrySpec(Protocol): + """Minimal geometry contract consumed by solver-specific scenes.""" + + kind: str + + def to_mesh(self, segments: int = 16) -> tuple: + ... + + class GeometryBase: """Base class for all geometry primitives.""" diff --git a/witwin/core/geometry/cuda/backend.py b/witwin/core/geometry/cuda/backend.py index 63afa2a..c0cff8f 100644 --- a/witwin/core/geometry/cuda/backend.py +++ b/witwin/core/geometry/cuda/backend.py @@ -1,6 +1,6 @@ """Loader for the native CUDA mesh-SDF extension. -Exposes a module object that is call-compatible with the former SlangTorch +Exposes a module object that is call-compatible with the former JIT-kernel 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. diff --git a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu index 12a80df..5e413c8 100644 --- a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu +++ b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu @@ -1,8 +1,8 @@ // 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 +// Direct port of the former mesh-SDF JIT module. The forward kernels +// reproduce its closest-point (Ericson) and solid-angle math verbatim; +// the backward kernels replace generated 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 @@ -57,7 +57,7 @@ __device__ __forceinline__ float3 read_vec3(const float* __restrict__ data, int return make_float3(v[0], v[1], v[2]); } -// Closest point on triangle (a, b, c) to p (Ericson, matching mesh_sdf.slang). +// Closest point on triangle (a, b, c) to p (Ericson, matching the reference implementation). // 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( diff --git a/witwin/core/geometry/polygon.py b/witwin/core/geometry/polygon.py new file mode 100644 index 0000000..8b56fb8 --- /dev/null +++ b/witwin/core/geometry/polygon.py @@ -0,0 +1,91 @@ +"""Shared differentiable polygon distance helpers.""" + +from __future__ import annotations + +import torch + + +def _cross_2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] + + +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, :, :] + 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) + if polygon_area.detach().item() >= 0.0: + inside = torch.all(crosses >= -1.0e-7, dim=1) + else: + inside = torch.all(crosses <= 1.0e-7, dim=1) + 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) diff --git a/witwin/core/geometry/primitives.py b/witwin/core/geometry/primitives.py index 782a30d..4968935 100644 --- a/witwin/core/geometry/primitives.py +++ b/witwin/core/geometry/primitives.py @@ -13,6 +13,7 @@ _constant_tensor, ) from .mesh_sdf import triangle_mesh_unsigned_distance +from .polygon import _safe_signed_denom, convex_polygon_signed_distance_2d def _faces_tensor(data, *, device) -> torch.Tensor: @@ -33,76 +34,6 @@ def _box_signed_distance(dx, dy, dz, half_size: torch.Tensor) -> torch.Tensor: return outside + inside -def _cross_2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: - return a[..., 0] * b[..., 1] - a[..., 1] * b[..., 0] - - -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, :, :] - 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) - if polygon_area.detach().item() >= 0.0: - inside = torch.all(crosses >= -1.0e-7, dim=1) - else: - inside = torch.all(crosses <= 1.0e-7, dim=1) - 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): @@ -350,7 +281,7 @@ def signed_distance(self, xx, yy, zz): torch.stack([height, radius]), torch.stack([height, -radius]), ]) - return _convex_polygon_signed_distance_2d(profile_points, triangle).reshape(axial.shape) + return convex_polygon_signed_distance_2d(profile_points, triangle).reshape(axial.shape) def to_mesh(self, segments=16): device = self.device @@ -430,7 +361,7 @@ def signed_distance(self, xx, yy, zz): radius = torch.clamp(self.radius.to(device=px.device, dtype=px.dtype), min=torch.finfo(px.dtype).eps) angles = torch.linspace(0, 2 * np.pi, self.num_sides + 1, device=px.device)[:-1] polygon = torch.stack([radius * torch.cos(angles), radius * torch.sin(angles)], dim=1) - polygon_distance = _convex_polygon_signed_distance_2d(torch.stack([px, py], dim=-1).reshape(-1, 2), polygon).reshape(px.shape) + polygon_distance = convex_polygon_signed_distance_2d(torch.stack([px, py], dim=-1).reshape(-1, 2), polygon).reshape(px.shape) axial_distance = torch.abs(axial) - self.height.to(device=axial.device, dtype=axial.dtype) / 2 outside = torch.sqrt(torch.clamp(polygon_distance, min=0.0) ** 2 + torch.clamp(axial_distance, min=0.0) ** 2) inside = torch.clamp(torch.maximum(polygon_distance, axial_distance), max=0.0) @@ -555,230 +486,3 @@ 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.") diff --git a/witwin/core/geometry/smpl.py b/witwin/core/geometry/smpl.py deleted file mode 100644 index c3c2b24..0000000 --- a/witwin/core/geometry/smpl.py +++ /dev/null @@ -1,294 +0,0 @@ -"""Differentiable SMPL geometry with optional smplpytorch dependency.""" - -from __future__ import annotations - -import pickle -from pathlib import Path -from typing import Any - -import numpy as np -import torch - -from .base import GeometryBase - - -class _Arr(np.ndarray): - @property - def r(self): - return np.asarray(self) - - -class _ChRecon: - def __init__(self, *args, **kwargs): - self._data = np.array([]) - for value in args: - if isinstance(value, np.ndarray): - self._data = value - return - if isinstance(value, _ChRecon): - self._data = value._data - return - - def __setstate__(self, state): - if isinstance(state, dict): - for value in state.values(): - if isinstance(value, np.ndarray): - self._data = value - return - if isinstance(value, _ChRecon): - self._data = value._data - return - elif isinstance(state, np.ndarray): - self._data = state - elif isinstance(state, (list, tuple)): - for value in state: - if isinstance(value, np.ndarray): - self._data = value - return - - -class _Unpickler(pickle.Unpickler): - def find_class(self, module, name): - if module.startswith("chumpy"): - return _ChRecon - return super().find_class(module, name) - - -def _ready_arguments_numpy(fname_or_dict): - import cv2 - - if not isinstance(fname_or_dict, dict): - with open(fname_or_dict, "rb") as handle: - data = _Unpickler(handle, encoding="latin1").load() - else: - data = fname_or_dict - - for key, value in list(data.items()): - if isinstance(value, _ChRecon): - data[key] = value._data - - want_shape_model = "shapedirs" in data - num_pose_params = data["kintree_table"].shape[1] * 3 - - if "trans" not in data: - data["trans"] = np.zeros(3) - if "pose" not in data: - data["pose"] = np.zeros(num_pose_params) - if "shapedirs" in data and "betas" not in data: - data["betas"] = np.zeros(data["shapedirs"].shape[-1]) - - if want_shape_model: - data["v_shaped"] = data["shapedirs"].dot(data["betas"]) + data["v_template"] - v_shaped = data["v_shaped"] - joint_regressor = data["J_regressor"] - data["J"] = np.column_stack([ - joint_regressor.dot(v_shaped[:, 0]), - joint_regressor.dot(v_shaped[:, 1]), - joint_regressor.dot(v_shaped[:, 2]), - ]) - pose = data["pose"].ravel()[3:] - rotations = np.concatenate([ - (cv2.Rodrigues(np.array(pp, dtype=np.float64))[0] - np.eye(3)).ravel() - for pp in pose.reshape((-1, 3)) - ]).ravel() - data["v_posed"] = v_shaped + data["posedirs"].dot(rotations) - else: - pose = data["pose"].ravel()[3:] - rotations = np.concatenate([ - (cv2.Rodrigues(np.array(pp, dtype=np.float64))[0] - np.eye(3)).ravel() - for pp in pose.reshape((-1, 3)) - ]).ravel() - data["v_posed"] = data["v_template"] + data["posedirs"].dot(rotations) - - for key, value in data.items(): - if isinstance(value, np.ndarray): - data[key] = value.view(_Arr) - return data - - -def _setup_smpl_compat(): - try: - import chumpy # noqa: F401 - - return - except ImportError: - pass - - try: - import smplpytorch.native.webuser.serialization as serialization - except ImportError: - return - - serialization.ready_arguments = _ready_arguments_numpy - - -try: - _setup_smpl_compat() - from smplpytorch.pytorch.smpl_layer import SMPL_Layer - - _SMPL_AVAILABLE = True -except ImportError: - _SMPL_AVAILABLE = False - - -_SMPL_LAYER_CACHE: dict[tuple[str, str, str], Any] = {} -_SMPL_FACES_CACHE: dict[tuple[str, str, str], np.ndarray] = {} - - -def _resolve_scene_device(device: str | None) -> str: - requested = "cuda" if device is None else device - resolved = torch.device(requested) - if resolved.type == "cuda" and not torch.cuda.is_available(): - raise RuntimeError( - "SMPLBody defaults to CUDA, but torch.cuda.is_available() is False. " - "Pass device='cpu' only for scene construction or non-rendering workflows." - ) - return str(resolved) - - -def _default_smpl_model_root() -> str: - return str(Path(__file__).resolve().parents[4] / "radar" / "models" / "smpl_models") - - -def _to_vertex_tensor(value, *, device: str) -> torch.Tensor: - if isinstance(value, torch.Tensor): - return value.to(device=device, dtype=torch.float32) - return torch.as_tensor(value, device=device, dtype=torch.float32) - - -def _get_smpl_layer(*, gender: str, model_root: str, device: str): - if not _SMPL_AVAILABLE: - raise ImportError("smplpytorch is required to instantiate or evaluate SMPLBody.") - key = (str(gender), str(model_root), str(device)) - layer = _SMPL_LAYER_CACHE.get(key) - if layer is None: - layer = SMPL_Layer(center_idx=0, gender=gender, model_root=model_root).to(device) - _SMPL_LAYER_CACHE[key] = layer - return layer - - -def _axis_angle_matrices(rvecs: torch.Tensor) -> torch.Tensor: - """Differentiable Rodrigues for (J, 3) axis-angle vectors -> (J, 3, 3).""" - # smplpytorch's zero-rotation guard: offset the vector, not the norm, so - # the gradient at exactly zero pose stays finite. - safe = rvecs + 1e-8 - angle = safe.norm(dim=1, keepdim=True) - axis = safe / angle - x, y, z = axis.unbind(-1) - zero = torch.zeros_like(x) - skew = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).reshape(-1, 3, 3) - eye = torch.eye(3, dtype=rvecs.dtype, device=rvecs.device).expand_as(skew) - sin = torch.sin(angle).unsqueeze(-1) - cos = torch.cos(angle).unsqueeze(-1) - return eye + sin * skew + (1.0 - cos) * (skew @ skew) - - -def _fast_smpl_forward(layer, pose: torch.Tensor, betas: torch.Tensor): - """Vectorized SMPL LBS equivalent to smplpytorch's forward (batch 1). - - smplpytorch's Python implementation costs ~21 ms per call; this batched - version is ~2 ms and numerically matches it (verified by diag_fast_lbs). - """ - device = pose.device - v_template = layer.th_v_template[0] - joints_rest = layer.th_J_regressor @ ( - v_template + torch.einsum("vdk,k->vd", layer.th_shapedirs, betas.reshape(-1)) - ) - v_shaped = v_template + torch.einsum("vdk,k->vd", layer.th_shapedirs, betas.reshape(-1)) - - rotations = _axis_angle_matrices(pose.reshape(24, 3)) - eye = torch.eye(3, dtype=pose.dtype, device=device) - pose_feature = (rotations[1:] - eye).reshape(-1) - v_posed = v_shaped + torch.einsum("vdk,k->vd", layer.th_posedirs, pose_feature) - - parents = layer.kintree_parents - relative = [joints_rest[0]] - for j in range(1, 24): - relative.append(joints_rest[j] - joints_rest[parents[j]]) - - bottom = torch.tensor([[0.0, 0.0, 0.0, 1.0]], dtype=pose.dtype, device=device) - transforms: list[torch.Tensor] = [] - for j in range(24): - local = torch.cat( - [torch.cat([rotations[j], relative[j].reshape(3, 1)], dim=1), bottom], dim=0 - ) - transforms.append(local if j == 0 else transforms[parents[j]] @ local) - global_transforms = torch.stack(transforms) - - posed_joints = global_transforms[:, :3, 3] - corrected_t = posed_joints - torch.einsum( - "jab,jb->ja", global_transforms[:, :3, :3], joints_rest - ) - skin_rot = torch.einsum("vj,jab->vab", layer.th_weights, global_transforms[:, :3, :3]) - skin_t = layer.th_weights @ corrected_t - vertices = torch.einsum("vab,vb->va", skin_rot, v_posed) + skin_t - - center = posed_joints[0] - return (vertices - center).unsqueeze(0), (posed_joints - center).unsqueeze(0) - - -class SMPLBody(GeometryBase): - """Differentiable SMPL geometry with position and rotation.""" - - kind = "smpl" - - def __init__( - self, - pose, - shape, - *, - position=(0.0, 0.0, 0.0), - gender: str = "male", - model_root: str | None = None, - rotation=None, - device=None, - ): - super().__init__(position=position, rotation=rotation, device=device) - tensor_device = str(self.position.device) - self.pose = _to_vertex_tensor(pose, device=tensor_device).reshape(-1) - self.shape = _to_vertex_tensor(shape, device=tensor_device).reshape(-1) - self.gender = str(gender) - self.model_root = _default_smpl_model_root() if model_root is None else str(model_root) - - def updated(self, **changes) -> "SMPLBody": - updated = SMPLBody( - pose=changes.pop("pose", self.pose), - shape=changes.pop("shape", self.shape), - position=changes.pop("position", self.position), - gender=changes.pop("gender", self.gender), - model_root=changes.pop("model_root", self.model_root), - rotation=changes.pop("rotation", self.rotation), - device=changes.pop("device", self.position.device), - ) - if changes: - unsupported = ", ".join(sorted(changes)) - raise TypeError(f"Unsupported SMPLBody updates: {unsupported}") - return updated - - def _evaluate(self, *, device: str): - layer = _get_smpl_layer(gender=self.gender, model_root=self.model_root, device=device) - pose_tensor = self.pose.to(device=device, dtype=torch.float32).view(1, -1) - shape_tensor = self.shape.to(device=device, dtype=torch.float32).view(1, -1) - if shape_tensor.requires_grad: - shape_tensor = shape_tensor + 1e-8 - vertices, joints = _fast_smpl_forward(layer, pose_tensor.reshape(-1), shape_tensor.reshape(-1)) - vertices = self._transform_mesh_verts(vertices[0]) - joints = self._transform_mesh_verts(joints[0]) - cache_key = (self.gender, self.model_root, device) - faces = _SMPL_FACES_CACHE.get(cache_key) - if faces is None: - faces = np.ascontiguousarray(layer.th_faces.detach().cpu().numpy().astype(np.int32)) - _SMPL_FACES_CACHE[cache_key] = faces - return vertices.contiguous(), faces, joints.contiguous() - - def to_mesh(self, segments=16, *, device=None): - del segments - resolved_device = _resolve_scene_device(device or self.position.device) - vertices, faces, _ = self._evaluate(device=resolved_device) - face_tensor = torch.as_tensor(faces, device=vertices.device, dtype=torch.int64) - return vertices, face_tensor - - def joints(self, *, device=None) -> torch.Tensor: - resolved_device = _resolve_scene_device(device or self.position.device) - _, _, joints = self._evaluate(device=resolved_device) - return joints diff --git a/witwin/core/material.py b/witwin/core/material.py index 2d7d99f..1324bbf 100644 --- a/witwin/core/material.py +++ b/witwin/core/material.py @@ -6,42 +6,46 @@ import numpy as np -VACUUM_PERMITTIVITY = 8.8541878128e-12 - - -def _coerce_scalar(value: float, *, name: str) -> float: - return float(value) +from .geometry.base import GeometrySpec - -def _coerce_nonnegative(value: float, *, name: str) -> float: - scalar = float(value) - if scalar < 0.0: - raise ValueError(f"{name} must be >= 0.") - return scalar +VACUUM_PERMITTIVITY = 8.8541878128e-12 -def _coerce_differentiable_scalar(value: Any, *, name: str) -> Any: +def _coerce_scalar(value: Any, *, name: str) -> Any: try: return float(value) except (TypeError, ValueError): return value -def _coerce_differentiable_nonnegative(value: Any, *, name: str) -> Any: +def _numeric_scalar(value: Any) -> float | None: + for candidate in ( + lambda item: float(item), + lambda item: float(item.item()), + lambda item: float(item[0]), + ): + try: + return candidate(value) + except (TypeError, ValueError, IndexError, KeyError, AttributeError): + continue + return None + + +def _coerce_nonnegative(value: Any, *, name: str) -> Any: + scalar = _numeric_scalar(value) + if scalar is not None and scalar < 0.0: + raise ValueError(f"{name} must be >= 0.") try: - scalar = float(value) + return 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): + scalar = _numeric_scalar(value) + if scalar is None: return False + return bool(np.isclose(scalar, target)) @dataclass(frozen=True) @@ -82,64 +86,6 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: @dataclass(frozen=True, init=False) class Material(MaterialSpec): - eps_r: float - mu_r: float - sigma_e: float - name: str | None - - def __init__( - self, - eps_r: float = 1.0, - mu_r: float = 1.0, - sigma_e: float = 0.0, - name: str | None = None, - ): - object.__setattr__(self, "eps_r", _coerce_scalar(eps_r, name="eps_r")) - object.__setattr__(self, "mu_r", _coerce_scalar(mu_r, name="mu_r")) - object.__setattr__(self, "sigma_e", _coerce_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 np.isclose(self.sigma_e, 0.0), - magnetic=not np.isclose(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 np.isclose(self.sigma_e, 0.0): - raise ValueError("Conductive materials require frequency > 0.") - return FrequencyMaterialSample( - eps_r=complex(self.eps_r, 0.0), - mu_r=self.mu_r, - sigma_e=self.sigma_e, - ) - - eps_r = complex( - self.eps_r, - -self.sigma_e / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), - ) - return FrequencyMaterialSample( - eps_r=eps_r, - mu_r=self.mu_r, - sigma_e=self.sigma_e, - ) - - -@dataclass(frozen=True, init=False) -class DifferentiableMaterial(MaterialSpec): eps_r: Any mu_r: Any sigma_e: Any @@ -152,9 +98,9 @@ def __init__( 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, "eps_r", _coerce_scalar(eps_r, name="eps_r")) + object.__setattr__(self, "mu_r", _coerce_scalar(mu_r, name="mu_r")) + object.__setattr__(self, "sigma_e", _coerce_nonnegative(sigma_e, name="sigma_e")) object.__setattr__(self, "name", None if name is None else str(name)) def capabilities(self) -> MaterialCapabilities: @@ -176,31 +122,31 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: resolved_frequency = float(frequency) if resolved_frequency < 0.0: raise ValueError("frequency must be >= 0.") + sigma_scalar = _numeric_scalar(self.sigma_e) 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.") + if sigma_scalar is not None and not np.isclose(sigma_scalar, 0.0): + raise ValueError("Conductive materials require frequency > 0.") + eps_scalar = _numeric_scalar(self.eps_r) + mu_scalar = _numeric_scalar(self.mu_r) return FrequencyMaterialSample( - eps_r=self.eps_r, - mu_r=self.mu_r, - sigma_e=self.sigma_e, + eps_r=self.eps_r if eps_scalar is None else complex(eps_scalar, 0.0), + mu_r=self.mu_r if mu_scalar is None else mu_scalar, + sigma_e=self.sigma_e if sigma_scalar is None else sigma_scalar, ) - 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_scalar = _numeric_scalar(self.eps_r) + mu_scalar = _numeric_scalar(self.mu_r) + if eps_scalar is None or sigma_scalar is None: eps_r = self.eps_r mu_r = self.mu_r sigma_e = self.sigma_e + else: + eps_r = complex( + eps_scalar, + -sigma_scalar / (2.0 * np.pi * resolved_frequency * VACUUM_PERMITTIVITY), + ) + mu_r = self.mu_r if mu_scalar is None else mu_scalar + sigma_e = self.sigma_e if sigma_scalar is None else sigma_scalar return FrequencyMaterialSample( eps_r=eps_r, mu_r=mu_r, @@ -210,7 +156,7 @@ def evaluate_at_frequency(self, frequency: float) -> FrequencyMaterialSample: @dataclass(frozen=True, init=False) class Structure: - geometry: Any + geometry: GeometrySpec material: MaterialSpec name: str | None priority: int @@ -220,7 +166,7 @@ class Structure: def __init__( self, - geometry: Any, + geometry: GeometrySpec, material: MaterialSpec, name: str | None = None, priority: int = 0, @@ -228,8 +174,13 @@ def __init__( tags=(), metadata: Mapping[str, Any] | None = None, ): - if material is None: - raise ValueError("Structure requires material=.") + if not isinstance(geometry, GeometrySpec): + raise TypeError("Structure geometry must implement GeometrySpec (kind and to_mesh()).") + if not isinstance(material, MaterialSpec): + raise TypeError( + "Structure material must implement MaterialSpec " + "(name, capabilities(), evaluate_static(), and evaluate_at_frequency())." + ) object.__setattr__(self, "geometry", geometry) object.__setattr__(self, "material", material) object.__setattr__(self, "name", None if name is None else str(name)) diff --git a/witwin/core/scene.py b/witwin/core/scene.py deleted file mode 100644 index 867c5c6..0000000 --- a/witwin/core/scene.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Base Scene class for all solver-specific scenes.""" - - -class SceneBase: - """Base scene for all solver-specific scenes. - - Solver-specific scenes (e.g. maxwell.Scene, radar.Scene) extend this - with domain-specific fields and methods. - """ - - def __init__( - self, - *, - structures=None, - sources=None, - monitors=None, - metadata=None, - device="cuda", - verbose=False, - ): - self.structures = list(structures or []) - self.sources = list(sources or []) - self.monitors = list(monitors or []) - self.metadata = dict(metadata or {}) - self.device = device - self.verbose = verbose - - def add_structure(self, structure): - self.structures.append(structure) - return self - - def add_source(self, source): - self.sources.append(source) - return self - - def add_monitor(self, monitor): - self.monitors.append(monitor) - return self diff --git a/witwin/core/scene_to_mitsuba.py b/witwin/core/scene_to_mitsuba.py deleted file mode 100644 index 90f5b94..0000000 --- a/witwin/core/scene_to_mitsuba.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Shared helpers for lowering mesh-based scenes into Mitsuba.""" - -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Mapping - -import numpy as np -import torch - - -def _load_mitsuba(variant: str): - import drjit as dr - import mitsuba as mi - - mi.set_variant(variant) - return mi, dr - - -def _as_vertices_array(vertices) -> np.ndarray: - if isinstance(vertices, torch.Tensor): - vertices = vertices.detach().cpu().numpy() - array = np.asarray(vertices, dtype=np.float32) - if array.ndim != 2 or array.shape[1] != 3: - raise ValueError("Mitsuba vertices must have shape (V, 3).") - return np.ascontiguousarray(array) - - -def _as_faces_array(faces) -> np.ndarray: - if isinstance(faces, torch.Tensor): - faces = faces.detach().cpu().numpy() - array = np.asarray(faces, dtype=np.uint32) - if array.ndim != 2 or array.shape[1] != 3: - raise ValueError("Mitsuba faces must have shape (F, 3).") - return np.ascontiguousarray(array) - - -def _materialize_plugin(mi, plugin): - if isinstance(plugin, Mapping): - return mi.load_dict(dict(plugin)) - return plugin - - -@dataclass(frozen=True) -class MitsubaRenderable: - name: str - vertices: torch.Tensor | np.ndarray - faces: torch.Tensor | np.ndarray - bsdf: dict[str, Any] | None = None - - -@dataclass(frozen=True) -class MitsubaSceneHandle: - scene: Any - params: Any - - -def create_mitsuba_mesh( - renderable: MitsubaRenderable, - *, - default_bsdf: dict[str, Any] | None = None, - variant: str = "cuda_ad_rgb", - has_vertex_texcoords: bool = False, -): - mi, dr = _load_mitsuba(variant) - vertices = _as_vertices_array(renderable.vertices) - faces = _as_faces_array(renderable.faces) - bsdf = default_bsdf or { - "type": "diffuse", - "reflectance": {"type": "rgb", "value": (0.8, 0.8, 0.8)}, - } - - mesh = mi.Mesh( - renderable.name, - vertex_count=vertices.shape[0], - face_count=faces.shape[0], - has_vertex_texcoords=has_vertex_texcoords, - ) - mesh_params = mi.traverse(mesh) - mesh_params["vertex_positions"] = dr.ravel(dr.cuda.ad.TensorXf(vertices)) - mesh_params["faces"] = dr.ravel(dr.cuda.ad.TensorXu(faces)) - mesh_params.update() - mesh.set_bsdf(mi.load_dict(renderable.bsdf or bsdf)) - return mesh - - -def build_mitsuba_scene( - *, - sensor, - renderables: Mapping[str, MitsubaRenderable | Any], - integrator: dict[str, Any] | Any | None = None, - default_bsdf: dict[str, Any] | None = None, - variant: str = "cuda_ad_rgb", -) -> MitsubaSceneHandle: - mi, _ = _load_mitsuba(variant) - scene_dict = { - "type": "scene", - "integrator": _materialize_plugin(mi, integrator or {"type": "direct"}), - "sensor": _materialize_plugin(mi, sensor), - "default_bsdf": default_bsdf - or { - "type": "diffuse", - "reflectance": {"type": "rgb", "value": (0.8, 0.8, 0.8)}, - }, - } - - for name, renderable in renderables.items(): - if not isinstance(renderable, MitsubaRenderable): - renderable = MitsubaRenderable( - name=name, - vertices=renderable.vertices, - faces=renderable.faces, - bsdf=getattr(renderable, "bsdf", None), - ) - scene_dict[name] = create_mitsuba_mesh( - renderable, - default_bsdf=scene_dict["default_bsdf"], - variant=variant, - ) - - scene = mi.load_dict(scene_dict) - return MitsubaSceneHandle(scene=scene, params=mi.traverse(scene)) - - -def update_mitsuba_scene_vertices( - params, - renderables: Mapping[str, MitsubaRenderable | Any], - *, - variant: str = "cuda_ad_rgb", -) -> None: - mi, dr = _load_mitsuba(variant) - for name, renderable in renderables.items(): - key = f"{name}.vertex_positions" - if key not in params: - continue - vertices = renderable.vertices if not isinstance(renderable, MitsubaRenderable) else renderable.vertices - params[key] = dr.ravel(dr.cuda.ad.TensorXf(_as_vertices_array(vertices))) - params.update() From bb829581b30a1af4f730c9f120227a70639731bb Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Thu, 9 Jul 2026 19:32:28 -0700 Subject: [PATCH 3/6] core: expand CUDA wheel compatibility --- .github/workflows/publish-witwin.yml | 56 +++++++++++++------ README.md | 10 +++- hatch_build.py | 2 +- pyproject.toml | 8 +-- scripts/build_mesh_sdf_cuda_prebuilt.py | 6 +- scripts/verify_cuda_binary_arches.py | 73 +++++++++++++++++++++++++ witwin/core/geometry/cuda/build.py | 41 +++++++++++++- 7 files changed, 168 insertions(+), 28 deletions(-) create mode 100644 scripts/verify_cuda_binary_arches.py diff --git a/.github/workflows/publish-witwin.yml b/.github/workflows/publish-witwin.yml index 8186784..f2f3121 100644 --- a/.github/workflows/publish-witwin.yml +++ b/.github/workflows/publish-witwin.yml @@ -90,7 +90,7 @@ jobs: fail-fast: false matrix: os: [ubuntu-22.04, windows-2022] - python_version: ["3.10", "3.11", "3.12"] + python_version: ["3.10", "3.11", "3.12", "3.13", "3.14"] steps: - name: Check out repository @@ -113,9 +113,9 @@ jobs: echo "Disk after cleanup:"; df -h / - name: Install CUDA Toolkit - uses: Jimver/cuda-toolkit@v0.2.29 + uses: Jimver/cuda-toolkit@v0.2.35 with: - cuda: "12.5.0" + cuda: "12.8.1" method: local log-file-suffix: "${{ runner.os }}-py${{ matrix.python_version }}" @@ -149,18 +149,30 @@ jobs: - 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 + - name: Build packaged mesh SDF CUDA extensions 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 + WITWIN_CUDA_GENCODE_ARCHES: "7.0;7.5;8.0;8.6;8.9;9.0;10.0;10.1;12.0+PTX" PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} run: | - python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose + for torch_version in 2.10.0 2.11.0; do + python -m pip install --force-reinstall --no-cache-dir "torch==$torch_version" \ + --index-url https://download.pytorch.org/whl/cu128 + export WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR="${RUNNER_TEMP}/witwin_core_mesh_sdf_cuda/${torch_version}" + python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose + done + + - name: Verify compiled CUDA architectures + shell: bash + env: + PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + run: | + python "$PACKAGE_DIR"/scripts/verify_cuda_binary_arches.py \ + --stem witwin_core_mesh_sdf_cuda \ + "$PACKAGE_DIR"/witwin/core/geometry/cuda/prebuilt - name: Build platform wheel shell: bash @@ -169,21 +181,25 @@ jobs: run: | python -m build --wheel "$PACKAGE_DIR" - - name: Normalize Linux wheel platform tag + - name: Repair 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 + python -m pip install auditwheel patchelf + mkdir -p "$PACKAGE_DIR"/wheelhouse 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" + python -m auditwheel repair "$whl" \ + --plat manylinux_2_35_x86_64 \ + --wheel-dir "$PACKAGE_DIR"/wheelhouse \ + --exclude libc10.so --exclude libc10_cuda.so \ + --exclude libtorch.so --exclude libtorch_cpu.so \ + --exclude libtorch_cuda.so --exclude libtorch_python.so \ + --exclude libcudart.so.12 --exclude libcuda.so.1 done + rm -f "$PACKAGE_DIR"/dist/*-linux_x86_64.whl + mv "$PACKAGE_DIR"/wheelhouse/*.whl "$PACKAGE_DIR"/dist/ ls -l "$PACKAGE_DIR"/dist - name: Smoke install prebuilt wheel @@ -197,9 +213,16 @@ jobs: python -m pip install --no-deps "$PACKAGE_DIR"/dist/*.whl python - <<'PY' import importlib.util + import os import site + import zipfile from pathlib import Path + wheel = next((Path(os.environ["PACKAGE_DIR"]) / "dist").glob("*.whl")) + names = zipfile.ZipFile(wheel).namelist() + for abi_tag in ("torch_2_10_cu128", "torch_2_11_cu128"): + assert any(f"/{abi_tag}/witwin_core_mesh_sdf_cuda" in name for name in names), abi_tag + candidates = [] for root in site.getsitepackages() + [site.getusersitepackages()]: candidates.append(Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py") @@ -212,6 +235,7 @@ jobs: path = module.prebuilt_extension_path() assert path.exists(), f"missing packaged prebuilt extension: {path}" + assert path.parent.name == module.torch_abi_tag() extension = module.build_extension() print(f"Loaded packaged mesh SDF CUDA extension from {extension.__file__}") PY diff --git a/README.md b/README.md index 9a886d3..273c9dd 100644 --- a/README.md +++ b/README.md @@ -4,12 +4,20 @@ WiTwin Core is the shared data and geometry package of the WiTwin stack. It defi ## Get Started -Python 3.10+ and an NVIDIA GPU are required for the main WiTwin simulation workflows. +CPython 3.10-3.14, PyTorch 2.10-2.11, and an NVIDIA GPU are supported for the main WiTwin simulation workflows. ```bash pip install witwin ``` +## Prebuilt CUDA Support + +Release wheels are built for Linux x86_64 and Windows x86_64 with CUDA 12.8. Each wheel carries separate PyTorch 2.10/cu128 and 2.11/cu128 native-extension variants, selected at runtime by the active PyTorch ABI. The fat binaries contain native code for compute capabilities 7.0, 7.5, 8.0, 8.6, 8.9, 9.0, 10.0, 10.1, and 12.0, plus compute 12.0 PTX for forward-compatible Blackwell execution. This includes native coverage for RTX 2080-class Turing GPUs and current data-center and RTX/RTX PRO Blackwell families. + +Linux wheels target `manylinux_2_35_x86_64`. The installed NVIDIA driver must support the CUDA 12.x runtime supplied by PyTorch; the CUDA toolkit is only needed for source/JIT builds. + +For full CUDA 12.8 and Blackwell support, use at least driver 570.26 on Linux or 570.65 on Windows. Pre-Blackwell systems can use NVIDIA's CUDA 12.x minor-version compatibility floor (525.60.13 on Linux or 528.33 on Windows), subject to NVIDIA's compatibility-mode feature limits. + ## What It Provides - Runtime-checkable `GeometrySpec` and `MaterialSpec` contracts diff --git a/hatch_build.py b/hatch_build.py index b204e2f..c5eb095 100644 --- a/hatch_build.py +++ b/hatch_build.py @@ -11,7 +11,7 @@ def initialize(self, version: str, build_data: dict) -> None: 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.*")) + 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 0ab3826..2ae9ec7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,12 +7,12 @@ name = "witwin" version = "0.1.0" description = "Witwin shared geometry, material, and structure contracts" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.10,<3.15" license = "MIT" dependencies = [ "numpy>=1.24", - "torch>=2.0", + "torch>=2.10,<2.12", ] [project.optional-dependencies] @@ -34,8 +34,8 @@ artifacts = [ "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", + "witwin/core/geometry/cuda/prebuilt/**/*.pyd", + "witwin/core/geometry/cuda/prebuilt/**/*.so", ] [tool.hatch.build.hooks.custom] diff --git a/scripts/build_mesh_sdf_cuda_prebuilt.py b/scripts/build_mesh_sdf_cuda_prebuilt.py index 86e3e43..fbf5f4c 100644 --- a/scripts/build_mesh_sdf_cuda_prebuilt.py +++ b/scripts/build_mesh_sdf_cuda_prebuilt.py @@ -39,21 +39,21 @@ def main() -> None: parser.add_argument("--verbose", action="store_true") args = parser.parse_args() + cuda_build = _load_cuda_build_module() build_dir = Path( os.environ.get( "WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", - Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel", + Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel" / cuda_build.torch_abi_tag(), ) ) 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 = cuda_build.prebuilt_variant_root() target_dir.mkdir(parents=True, exist_ok=True) for suffix in (".pyd", ".so"): existing = target_dir / f"witwin_core_mesh_sdf_cuda{suffix}" diff --git a/scripts/verify_cuda_binary_arches.py b/scripts/verify_cuda_binary_arches.py new file mode 100644 index 0000000..b7f4bee --- /dev/null +++ b/scripts/verify_cuda_binary_arches.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import argparse +import subprocess +import tempfile +import zipfile +from pathlib import Path + + +EXPECTED_SASS = ("70", "75", "80", "86", "89", "90", "100", "101", "120") +EXPECTED_PTX_TARGET = "sm_120" + + +def _matches(path: Path, stems: tuple[str, ...]) -> bool: + return path.suffix in {".pyd", ".so"} and any(path.name.startswith(stem) for stem in stems) + + +def _collect_binaries(inputs: list[Path], stems: tuple[str, ...], extract_root: Path) -> list[Path]: + binaries: list[Path] = [] + for input_path in inputs: + if input_path.suffix == ".whl": + with zipfile.ZipFile(input_path) as wheel: + for name in wheel.namelist(): + member = Path(name) + if not _matches(member, stems): + continue + wheel.extract(name, extract_root) + binaries.append(extract_root / member) + continue + if input_path.is_dir(): + binaries.extend(path for path in input_path.rglob("*") if _matches(path, stems)) + continue + if _matches(input_path, stems): + binaries.append(input_path) + return sorted(set(path.resolve() for path in binaries)) + + +def _cuobjdump(flag: str, binary: Path) -> str: + result = subprocess.run( + ["cuobjdump", flag, str(binary)], + check=True, + capture_output=True, + text=True, + ) + return f"{result.stdout}\n{result.stderr}" + + +def main() -> None: + parser = argparse.ArgumentParser(description="Verify CUDA SASS and PTX targets in release binaries.") + parser.add_argument("inputs", nargs="+", type=Path) + parser.add_argument("--stem", action="append", required=True) + args = parser.parse_args() + + with tempfile.TemporaryDirectory(prefix="witwin_cuda_arch_verify_") as temp_dir: + binaries = _collect_binaries(args.inputs, tuple(args.stem), Path(temp_dir)) + if not binaries: + raise SystemExit(f"No native binaries matching {args.stem!r} were found.") + + for binary in binaries: + elf_listing = _cuobjdump("--list-elf", binary) + missing_sass = [arch for arch in EXPECTED_SASS if f"sm_{arch}" not in elf_listing] + if missing_sass: + raise SystemExit(f"{binary} is missing SASS targets: {', '.join(missing_sass)}") + + ptx_dump = _cuobjdump("--dump-ptx", binary) + if f".target {EXPECTED_PTX_TARGET}" not in ptx_dump: + raise SystemExit(f"{binary} is missing compute 12.0 PTX.") + + print(f"Verified CUDA architectures in {binary}") + + +if __name__ == "__main__": + main() diff --git a/witwin/core/geometry/cuda/build.py b/witwin/core/geometry/cuda/build.py index c66e3ca..d231730 100644 --- a/witwin/core/geometry/cuda/build.py +++ b/witwin/core/geometry/cuda/build.py @@ -16,6 +16,7 @@ import tempfile from pathlib import Path +import torch from torch.utils.cpp_extension import load @@ -171,12 +172,26 @@ def prebuilt_root() -> Path: return source_root() / "prebuilt" +def torch_abi_tag() -> str: + """Return the directory tag for the active PyTorch CUDA ABI.""" + release = torch.__version__.split("+", 1)[0].split(".") + if len(release) < 2: + raise RuntimeError(f"Unable to determine the PyTorch ABI from {torch.__version__!r}.") + cuda_version = torch.version.cuda + cuda_tag = "cpu" if cuda_version is None else f"cu{cuda_version.replace('.', '')}" + return f"torch_{release[0]}_{release[1]}_{cuda_tag}" + + +def prebuilt_variant_root() -> Path: + return prebuilt_root() / torch_abi_tag() + + 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()}" + return prebuilt_variant_root() / f"{EXTENSION_NAME}{extension_suffix()}" def extension_sources() -> list[Path]: @@ -184,6 +199,26 @@ def extension_sources() -> list[Path]: return [root / "extension.cpp", root / "mesh_sdf_kernels.cu"] +def _cuda_gencode_flags() -> list[str]: + """Translate the release architecture list directly into nvcc flags.""" + value = os.environ.get("WITWIN_CUDA_GENCODE_ARCHES") + if not value: + return [] + flags: list[str] = [] + for entry in value.split(";"): + entry = entry.strip() + if not entry: + continue + include_ptx = entry.endswith("+PTX") + number = entry.removesuffix("+PTX").replace(".", "") + if not number.isdigit(): + raise ValueError(f"Invalid CUDA architecture {entry!r} in WITWIN_CUDA_GENCODE_ARCHES.") + flags.append(f"-gencode=arch=compute_{number},code=sm_{number}") + if include_ptx: + flags.append(f"-gencode=arch=compute_{number},code=compute_{number}") + return flags + + 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: @@ -223,14 +258,14 @@ def _jit_build(build_directory: Path, *, verbose: bool): 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_cuda_cflags=["-O3", *_cuda_gencode_flags()], extra_ldflags=_conda_torch_ldflags(), verbose=verbose, ) def build_extension(*, verbose: bool = False): - default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME + default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME / torch_abi_tag() 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": From 6e75274131bafc0ba2ab94ecb68546418a869e6c Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Sat, 11 Jul 2026 00:47:47 -0700 Subject: [PATCH 4/6] Release witwin 0.3.0 with Stable ABI wheels --- .github/workflows/publish-witwin.yml | 102 +++++++-- README.md | 4 +- hatch_build.py | 6 +- pyproject.toml | 12 +- scripts/build_mesh_sdf_cuda_prebuilt.py | 6 +- witwin/core/__init__.py | 2 +- witwin/core/geometry/cuda/build.py | 66 +++--- witwin/core/geometry/cuda/extension.cpp | 68 +++--- witwin/core/geometry/cuda/mesh_sdf_kernels.cu | 193 ++++++++++++------ 9 files changed, 291 insertions(+), 168 deletions(-) diff --git a/.github/workflows/publish-witwin.yml b/.github/workflows/publish-witwin.yml index f2f3121..c9215e7 100644 --- a/.github/workflows/publish-witwin.yml +++ b/.github/workflows/publish-witwin.yml @@ -83,14 +83,13 @@ jobs: PY build_cuda_wheels: - name: Build CUDA wheel / ${{ matrix.os }} / py${{ matrix.python_version }} + name: Build stable CUDA wheel / ${{ matrix.os }} 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", "3.13", "3.14"] steps: - name: Check out repository @@ -99,7 +98,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v5 with: - python-version: ${{ matrix.python_version }} + python-version: "3.10" - name: Free disk space (Linux) if: runner.os == 'Linux' @@ -117,7 +116,7 @@ jobs: with: cuda: "12.8.1" method: local - log-file-suffix: "${{ runner.os }}-py${{ matrix.python_version }}" + log-file-suffix: "${{ runner.os }}-stable-abi" - name: Set up MSVC if: runner.os == 'Windows' @@ -149,21 +148,18 @@ jobs: - name: Install CUDA extension build dependencies run: | python -m pip install --upgrade pip setuptools wheel ninja numpy + python -m pip install torch==2.10.0 --index-url https://download.pytorch.org/whl/cu128 python -m pip install build - - name: Build packaged mesh SDF CUDA extensions + - name: Build packaged stable mesh SDF CUDA extension shell: bash env: MAX_JOBS: "2" WITWIN_CUDA_GENCODE_ARCHES: "7.0;7.5;8.0;8.6;8.9;9.0;10.0;10.1;12.0+PTX" PACKAGE_DIR: ${{ needs.validate_release.outputs.package_dir }} + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/witwin_core_mesh_sdf_cuda/stable_abi_v1 run: | - for torch_version in 2.10.0 2.11.0; do - python -m pip install --force-reinstall --no-cache-dir "torch==$torch_version" \ - --index-url https://download.pytorch.org/whl/cu128 - export WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR="${RUNNER_TEMP}/witwin_core_mesh_sdf_cuda/${torch_version}" - python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose - done + python "$PACKAGE_DIR"/scripts/build_mesh_sdf_cuda_prebuilt.py --verbose - name: Verify compiled CUDA architectures shell: bash @@ -220,8 +216,10 @@ jobs: wheel = next((Path(os.environ["PACKAGE_DIR"]) / "dist").glob("*.whl")) names = zipfile.ZipFile(wheel).namelist() - for abi_tag in ("torch_2_10_cu128", "torch_2_11_cu128"): - assert any(f"/{abi_tag}/witwin_core_mesh_sdf_cuda" in name for name in names), abi_tag + native = [name for name in names if "witwin_core_mesh_sdf_cuda" in name and name.endswith((".so", ".pyd"))] + assert len(native) == 1, native + assert "/prebuilt/" in f"/{native[0]}" + assert "torch_" not in native[0] candidates = [] for root in site.getsitepackages() + [site.getusersitepackages()]: @@ -235,7 +233,7 @@ jobs: path = module.prebuilt_extension_path() assert path.exists(), f"missing packaged prebuilt extension: {path}" - assert path.parent.name == module.torch_abi_tag() + assert path.parent.name == "prebuilt" extension = module.build_extension() print(f"Loaded packaged mesh SDF CUDA extension from {extension.__file__}") PY @@ -243,10 +241,83 @@ jobs: - name: Upload wheel artifact uses: actions/upload-artifact@v4 with: - name: wheel-${{ matrix.os }}-py${{ matrix.python_version }} + name: wheel-${{ matrix.os }} path: ${{ needs.validate_release.outputs.package_dir }}/dist/*.whl if-no-files-found: error + test_torch_compatibility: + name: Stable ABI / ${{ matrix.os }} / py${{ matrix.compatibility.python_version }} / torch${{ matrix.compatibility.torch_version }} + needs: + - validate_release + - build_cuda_wheels + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-2022] + compatibility: + - {python_version: "3.10", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.11", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.12", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.13", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.10.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.11.0", cuda_index: "cu128"} + - {python_version: "3.14", torch_version: "2.12.0", cuda_index: "cu126"} + + steps: + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.compatibility.python_version }} + + - name: Download stable wheel + uses: actions/download-artifact@v4 + with: + name: wheel-${{ matrix.os }} + path: dist + + - name: Install compatibility environment + shell: bash + run: | + python -m pip install --upgrade pip + python -m pip install --no-cache-dir "torch==${{ matrix.compatibility.torch_version }}" \ + --index-url "https://download.pytorch.org/whl/${{ matrix.compatibility.cuda_index }}" + python -m pip install --no-deps dist/*.whl + + - name: Load the same stable binary + shell: bash + env: + EXPECTED_TORCH: ${{ matrix.compatibility.torch_version }} + WITWIN_CORE_MESH_SDF_CUDA_PREBUILT: "1" + WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR: ${{ runner.temp }}/must-not-jit-build + run: | + python - <<'PY' + import importlib.util + import os + import site + from pathlib import Path + + import torch + + assert torch.__version__.split("+", 1)[0] == os.environ["EXPECTED_TORCH"] + candidates = [ + Path(root) / "witwin" / "core" / "geometry" / "cuda" / "build.py" + for root in site.getsitepackages() + [site.getusersitepackages()] + ] + build_path = next((path for path in candidates if path.exists()), None) + assert build_path is not None, 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) + + native_path = module.prebuilt_extension_path() + assert native_path.exists(), native_path + extension = module.build_extension() + assert Path(extension.__file__).resolve() == native_path.resolve() + assert extension.query_mesh_unsigned_distance is not None + print(f"Loaded {native_path.name} with torch {torch.__version__} on Python {os.sys.version.split()[0]}") + PY + build_sdist: name: Build source distribution needs: validate_release @@ -281,6 +352,7 @@ jobs: if: github.event_name == 'release' needs: - build_cuda_wheels + - test_torch_compatibility - build_sdist runs-on: ubuntu-latest permissions: diff --git a/README.md b/README.md index 273c9dd..3a98233 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ WiTwin Core is the shared data and geometry package of the WiTwin stack. It defi ## Get Started -CPython 3.10-3.14, PyTorch 2.10-2.11, and an NVIDIA GPU are supported for the main WiTwin simulation workflows. +CPython 3.10-3.14, PyTorch 2.10 or newer, and an NVIDIA GPU are supported for the main WiTwin simulation workflows. ```bash pip install witwin @@ -12,7 +12,7 @@ pip install witwin ## Prebuilt CUDA Support -Release wheels are built for Linux x86_64 and Windows x86_64 with CUDA 12.8. Each wheel carries separate PyTorch 2.10/cu128 and 2.11/cu128 native-extension variants, selected at runtime by the active PyTorch ABI. The fat binaries contain native code for compute capabilities 7.0, 7.5, 8.0, 8.6, 8.9, 9.0, 10.0, 10.1, and 12.0, plus compute 12.0 PTX for forward-compatible Blackwell execution. This includes native coverage for RTX 2080-class Turing GPUs and current data-center and RTX/RTX PRO Blackwell families. +Release wheels are built for Linux x86_64 and Windows x86_64 with CUDA 12.8. Each platform wheel carries one PyTorch Stable ABI native library built against PyTorch 2.10 and reused across supported Python and PyTorch versions; CI currently verifies PyTorch 2.10-2.12. The fat binaries contain native code for compute capabilities 7.0, 7.5, 8.0, 8.6, 8.9, 9.0, 10.0, 10.1, and 12.0, plus compute 12.0 PTX for forward-compatible Blackwell execution. This includes native coverage for RTX 2080-class Turing GPUs and current data-center and RTX/RTX PRO Blackwell families. Linux wheels target `manylinux_2_35_x86_64`. The installed NVIDIA driver must support the CUDA 12.x runtime supplied by PyTorch; the CUDA toolkit is only needed for source/JIT builds. diff --git a/hatch_build.py b/hatch_build.py index c5eb095..cfe0486 100644 --- a/hatch_build.py +++ b/hatch_build.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sysconfig from pathlib import Path from hatchling.builders.hooks.plugin.interface import BuildHookInterface @@ -11,7 +12,8 @@ def initialize(self, version: str, build_data: dict) -> None: 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.*")) + has_prebuilt_extension = any(prebuilt_dir.glob("witwin_core_mesh_sdf_cuda.*")) if has_prebuilt_extension: - build_data["infer_tag"] = True + platform_tag = sysconfig.get_platform().replace("-", "_").replace(".", "_") + build_data["tag"] = f"py3-none-{platform_tag}" build_data["pure_python"] = False diff --git a/pyproject.toml b/pyproject.toml index 2ae9ec7..9375354 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "witwin" -version = "0.1.0" +version = "0.3.0" description = "Witwin shared geometry, material, and structure contracts" readme = "README.md" requires-python = ">=3.10,<3.15" @@ -12,18 +12,18 @@ license = "MIT" dependencies = [ "numpy>=1.24", - "torch>=2.10,<2.12", + "torch>=2.10", ] [project.optional-dependencies] maxwell = ["witwin-maxwell>=0.2,<0.3"] -radar = ["witwin-radar>=0.2,<0.3"] -channel = ["witwin-channel>=0.1,<0.2"] +radar = ["witwin-radar>=0.3,<0.4"] +channel = ["witwin-channel>=0.3,<0.4"] genesis = ["witwin-genesis>=0.1,<0.2"] all = [ "witwin-maxwell>=0.2,<0.3", - "witwin-radar>=0.2,<0.3", - "witwin-channel>=0.1,<0.2", + "witwin-radar>=0.3,<0.4", + "witwin-channel>=0.3,<0.4", "witwin-genesis>=0.1,<0.2", ] diff --git a/scripts/build_mesh_sdf_cuda_prebuilt.py b/scripts/build_mesh_sdf_cuda_prebuilt.py index fbf5f4c..1d2b226 100644 --- a/scripts/build_mesh_sdf_cuda_prebuilt.py +++ b/scripts/build_mesh_sdf_cuda_prebuilt.py @@ -43,7 +43,7 @@ def main() -> None: build_dir = Path( os.environ.get( "WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR", - Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel" / cuda_build.torch_abi_tag(), + Path(tempfile.gettempdir()) / "witwin_core_mesh_sdf_cuda_wheel" / "stable_abi_v1", ) ) os.environ["WITWIN_CORE_MESH_SDF_CUDA_BUILD_DIR"] = str(build_dir) @@ -53,10 +53,10 @@ def main() -> None: module = cuda_build.build_extension(verbose=args.verbose) module_file = Path(module.__file__).resolve() - target_dir = cuda_build.prebuilt_variant_root() + 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}" + existing = target_dir / f"{cuda_build.EXTENSION_NAME}{suffix}" if existing.exists(): existing.unlink() target = cuda_build.prebuilt_extension_path() diff --git a/witwin/core/__init__.py b/witwin/core/__init__.py index c2e426d..0b748ed 100644 --- a/witwin/core/__init__.py +++ b/witwin/core/__init__.py @@ -1,6 +1,6 @@ """Witwin Core - Shared geometry, material, and structure contracts.""" -__version__ = "0.1.0" +__version__ = "0.3.0" from .material import ( FrequencyMaterialSample, diff --git a/witwin/core/geometry/cuda/build.py b/witwin/core/geometry/cuda/build.py index d231730..44aba49 100644 --- a/witwin/core/geometry/cuda/build.py +++ b/witwin/core/geometry/cuda/build.py @@ -8,7 +8,6 @@ from __future__ import annotations -import importlib.util import os import shutil import subprocess @@ -172,26 +171,12 @@ def prebuilt_root() -> Path: return source_root() / "prebuilt" -def torch_abi_tag() -> str: - """Return the directory tag for the active PyTorch CUDA ABI.""" - release = torch.__version__.split("+", 1)[0].split(".") - if len(release) < 2: - raise RuntimeError(f"Unable to determine the PyTorch ABI from {torch.__version__!r}.") - cuda_version = torch.version.cuda - cuda_tag = "cpu" if cuda_version is None else f"cu{cuda_version.replace('.', '')}" - return f"torch_{release[0]}_{release[1]}_{cuda_tag}" - - -def prebuilt_variant_root() -> Path: - return prebuilt_root() / torch_abi_tag() - - def extension_suffix() -> str: return ".pyd" if os.name == "nt" else ".so" def prebuilt_extension_path() -> Path: - return prebuilt_variant_root() / f"{EXTENSION_NAME}{extension_suffix()}" + return prebuilt_root() / f"{EXTENSION_NAME}{extension_suffix()}" def extension_sources() -> list[Path]: @@ -219,19 +204,29 @@ def _cuda_gencode_flags() -> list[str]: return flags -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 +class _StableOpsModule: + """Attribute-compatible view of the Stable ABI dispatcher operators.""" + + def __init__(self, library_path: Path) -> None: + self.__file__ = str(library_path) + + def is_available(self) -> bool: + return bool(torch.cuda.is_available()) + + def __getattr__(self, name: str): + return getattr(torch.ops.witwin_core_mesh_sdf_cuda, name) + + +def _load_extension_file(library_path: Path) -> _StableOpsModule: + torch.ops.load_library(str(library_path)) + if not hasattr(torch.ops.witwin_core_mesh_sdf_cuda, "query_mesh_unsigned_distance"): + raise ImportError(f"{library_path} does not register the Stable ABI mesh SDF operators.") + return _StableOpsModule(library_path) 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. + # Wheels ship a Stable ABI library in ``prebuilt/`` so installs never invoke + # nvcc/MSVC and the same binary loads across supported Python/Torch versions. module_path = prebuilt_extension_path() if not module_path.exists(): return None @@ -253,19 +248,30 @@ 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( + library_path = 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", *_cuda_gencode_flags()], + extra_cflags=( + ["/O2", "/DTORCH_TARGET_VERSION=0x020a000000000000"] + if os.name == "nt" + else ["-O3", "-DTORCH_TARGET_VERSION=0x020a000000000000"] + ), + extra_cuda_cflags=[ + "-O3", + "-DTORCH_TARGET_VERSION=0x020a000000000000", + "-DUSE_CUDA", + *_cuda_gencode_flags(), + ], extra_ldflags=_conda_torch_ldflags(), + is_python_module=False, verbose=verbose, ) + return _StableOpsModule(Path(library_path)) def build_extension(*, verbose: bool = False): - default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME / torch_abi_tag() + default_build_directory = Path(tempfile.gettempdir()) / EXTENSION_NAME / "stable_abi_v1" 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": diff --git a/witwin/core/geometry/cuda/extension.cpp b/witwin/core/geometry/cuda/extension.cpp index e1cad20..802d8f6 100644 --- a/witwin/core/geometry/cuda/extension.cpp +++ b/witwin/core/geometry/cuda/extension.cpp @@ -1,43 +1,29 @@ -#include -#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."); +STABLE_TORCH_LIBRARY(witwin_core_mesh_sdf_cuda, m) { + m.def( + "query_mesh_unsigned_distance(Tensor triangles, Tensor points, " + "Tensor(a!) unsigned_distance, Tensor(b!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_distance_and_winding(Tensor triangles, Tensor points, " + "Tensor(a!) unsigned_distance, Tensor(b!) winding_angle, " + "Tensor(c!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_unsigned_distance_bvh(Tensor triangles, Tensor points, " + "Tensor node_bbox_min, Tensor node_bbox_max, Tensor node_left, Tensor node_right, " + "Tensor node_start, Tensor node_count, Tensor triangle_indices, " + "Tensor(a!) unsigned_distance, Tensor(b!) closest_triangle_index) -> ()"); + m.def( + "query_mesh_parity_sign_bvh(Tensor triangles, Tensor points, " + "Tensor node_bbox_min, Tensor node_bbox_max, Tensor node_left, Tensor node_right, " + "Tensor node_start, Tensor node_count, Tensor triangle_indices, float jitter_scale, " + "Tensor(a!) inside) -> ()"); + m.def( + "backward_mesh_unsigned_distance(Tensor triangles, Tensor points, " + "Tensor closest_triangle_index, Tensor grad_unsigned_distance, " + "Tensor(a!) grad_triangles, Tensor(b!) grad_points) -> ()"); + m.def( + "backward_mesh_distance_and_winding(Tensor triangles, Tensor points, " + "Tensor closest_triangle_index, Tensor grad_unsigned_distance, Tensor grad_winding_angle, " + "Tensor(a!) grad_triangles, Tensor(b!) grad_points) -> ()"); } diff --git a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu index 5e413c8..a558c61 100644 --- a/witwin/core/geometry/cuda/mesh_sdf_kernels.cu +++ b/witwin/core/geometry/cuda/mesh_sdf_kernels.cu @@ -8,10 +8,15 @@ // closest point, so grad_p == -(grad_a + grad_b + grad_c) holds by // construction (translation invariance of the distance). -#include -#include -#include -#include +#include +#include +#include +#include +#include +#include +#include + +#include namespace { @@ -528,16 +533,29 @@ 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_f32(const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + STD_TORCH_CHECK( + tensor.scalar_type() == torch::headeronly::ScalarType::Float, + name, + " must be float32"); + STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +void check_i32(const torch::stable::Tensor& tensor, const char* name) { + STD_TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + STD_TORCH_CHECK( + tensor.scalar_type() == torch::headeronly::ScalarType::Int, + name, + " must be int32"); + STD_TORCH_CHECK(tensor.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"); +cudaStream_t current_cuda_stream(const torch::stable::Tensor& tensor) { + void* stream_ptr = nullptr; + TORCH_ERROR_CODE_CHECK( + aoti_torch_get_current_cuda_stream(tensor.get_device_index(), &stream_ptr)); + return static_cast(stream_ptr); } constexpr int kBlock = 256; @@ -545,48 +563,62 @@ 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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + torch::stable::Tensor& unsigned_distance, + torch::stable::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 torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); const int point_count = static_cast(points.size(0)); if (point_count == 0) { return; } - query_unsigned_distance_kernel<<>>( - triangles.data_ptr(), points.data_ptr(), + query_unsigned_distance_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), static_cast(triangles.size(0)), point_count, - unsigned_distance.data_ptr(), closest_triangle_index.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); + unsigned_distance.mutable_data_ptr(), closest_triangle_index.mutable_data_ptr()); + STD_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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + torch::stable::Tensor& unsigned_distance, + torch::stable::Tensor& winding_angle, + torch::stable::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 torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); 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(), + query_distance_and_winding_kernel<<>>( + triangles.const_data_ptr(), points.const_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(); + unsigned_distance.mutable_data_ptr(), winding_angle.mutable_data_ptr(), + closest_triangle_index.mutable_data_ptr()); + STD_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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& node_bbox_min, + const torch::stable::Tensor& node_bbox_max, + const torch::stable::Tensor& node_left, + const torch::stable::Tensor& node_right, + const torch::stable::Tensor& node_start, + const torch::stable::Tensor& node_count, + const torch::stable::Tensor& triangle_indices, + torch::stable::Tensor& unsigned_distance, + torch::stable::Tensor& closest_triangle_index) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(node_bbox_min, "node_bbox_min"); @@ -598,25 +630,33 @@ void query_mesh_unsigned_distance_bvh_cuda( 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 torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); 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(); + query_unsigned_distance_bvh_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + node_bbox_min.const_data_ptr(), node_bbox_max.const_data_ptr(), + node_left.const_data_ptr(), node_right.const_data_ptr(), + node_start.const_data_ptr(), node_count.const_data_ptr(), + triangle_indices.const_data_ptr(), point_count, + unsigned_distance.mutable_data_ptr(), closest_triangle_index.mutable_data_ptr()); + STD_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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& node_bbox_min, + const torch::stable::Tensor& node_bbox_max, + const torch::stable::Tensor& node_left, + const torch::stable::Tensor& node_right, + const torch::stable::Tensor& node_start, + const torch::stable::Tensor& node_count, + const torch::stable::Tensor& triangle_indices, + double jitter_scale, + torch::stable::Tensor& inside) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_f32(node_bbox_min, "node_bbox_min"); @@ -627,46 +667,54 @@ void query_mesh_parity_sign_bvh_cuda( check_i32(node_count, "node_count"); check_i32(triangle_indices, "triangle_indices"); check_i32(inside, "inside"); - const c10::cuda::CUDAGuard guard(points.device()); + const torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); 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(); + query_parity_sign_bvh_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + node_bbox_min.const_data_ptr(), node_bbox_max.const_data_ptr(), + node_left.const_data_ptr(), node_right.const_data_ptr(), + node_start.const_data_ptr(), node_count.const_data_ptr(), + triangle_indices.const_data_ptr(), static_cast(jitter_scale), point_count, + inside.mutable_data_ptr()); + STD_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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& closest_triangle_index, + const torch::stable::Tensor& grad_unsigned_distance, + torch::stable::Tensor& grad_triangles, + torch::stable::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 torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); 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(); + backward_unsigned_distance_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + closest_triangle_index.const_data_ptr(), grad_unsigned_distance.const_data_ptr(), + point_count, grad_triangles.mutable_data_ptr(), grad_points.mutable_data_ptr()); + STD_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) { + const torch::stable::Tensor& triangles, + const torch::stable::Tensor& points, + const torch::stable::Tensor& closest_triangle_index, + const torch::stable::Tensor& grad_unsigned_distance, + const torch::stable::Tensor& grad_winding_angle, + torch::stable::Tensor& grad_triangles, + torch::stable::Tensor& grad_points) { check_f32(triangles, "triangles"); check_f32(points, "points"); check_i32(closest_triangle_index, "closest_triangle_index"); @@ -674,15 +722,24 @@ void backward_mesh_distance_and_winding_cuda( 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 torch::stable::accelerator::DeviceGuard device_guard(points.get_device_index()); 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(); + backward_distance_and_winding_kernel<<>>( + triangles.const_data_ptr(), points.const_data_ptr(), + closest_triangle_index.const_data_ptr(), grad_unsigned_distance.const_data_ptr(), + grad_winding_angle.const_data_ptr(), static_cast(triangles.size(0)), point_count, + grad_triangles.mutable_data_ptr(), grad_points.mutable_data_ptr()); + STD_CUDA_KERNEL_LAUNCH_CHECK(); +} + +STABLE_TORCH_LIBRARY_IMPL(witwin_core_mesh_sdf_cuda, CUDA, m) { + m.impl("query_mesh_unsigned_distance", TORCH_BOX(&query_mesh_unsigned_distance_cuda)); + m.impl("query_mesh_distance_and_winding", TORCH_BOX(&query_mesh_distance_and_winding_cuda)); + m.impl("query_mesh_unsigned_distance_bvh", TORCH_BOX(&query_mesh_unsigned_distance_bvh_cuda)); + m.impl("query_mesh_parity_sign_bvh", TORCH_BOX(&query_mesh_parity_sign_bvh_cuda)); + m.impl("backward_mesh_unsigned_distance", TORCH_BOX(&backward_mesh_unsigned_distance_cuda)); + m.impl("backward_mesh_distance_and_winding", TORCH_BOX(&backward_mesh_distance_and_winding_cuda)); } From 9ee6655656d4483cc7ad035e75fe49be0823f0f9 Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Sat, 11 Jul 2026 18:46:09 -0700 Subject: [PATCH 5/6] Route channel extras to channel native --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9375354..a6d507f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,12 +18,12 @@ dependencies = [ [project.optional-dependencies] maxwell = ["witwin-maxwell>=0.2,<0.3"] radar = ["witwin-radar>=0.3,<0.4"] -channel = ["witwin-channel>=0.3,<0.4"] +channel = ["witwin-channel-native>=0.1,<0.2"] genesis = ["witwin-genesis>=0.1,<0.2"] all = [ "witwin-maxwell>=0.2,<0.3", "witwin-radar>=0.3,<0.4", - "witwin-channel>=0.3,<0.4", + "witwin-channel-native>=0.1,<0.2", "witwin-genesis>=0.1,<0.2", ] From 6b6866e057746b242bc65f0624c6ba2809f6ba01 Mon Sep 17 00:00:00 2001 From: Xingyu Chen Date: Mon, 13 Jul 2026 02:55:49 -0700 Subject: [PATCH 6/6] Fix torus mesh parameterization --- tests/test_geometry_smoke.py | 24 +++++++++++++++++++++++- witwin/core/geometry/primitives.py | 2 +- 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tests/test_geometry_smoke.py b/tests/test_geometry_smoke.py index 38b7945..63afb37 100644 --- a/tests/test_geometry_smoke.py +++ b/tests/test_geometry_smoke.py @@ -1,9 +1,11 @@ from __future__ import annotations +import math + import pytest import torch -from witwin.core import Box, Cylinder, Mesh, Sphere +from witwin.core import Box, Cylinder, Mesh, Sphere, Torus def _grid(): @@ -71,6 +73,26 @@ def test_geometry_construction_to_mesh_and_to_mask(geometry, segments, inside_po assert torch.any(occupancy < 0.5) +def test_torus_mesh_matches_analytic_extents_and_volume(): + major_radius = 0.8 + minor_radius = 0.2 + torus = Torus(major_radius=major_radius, minor_radius=minor_radius, axis="z") + vertices, faces = torus.to_mesh(segments=48) + + expected_extents = torch.tensor( + [major_radius + minor_radius, major_radius + minor_radius, minor_radius], + dtype=vertices.dtype, + ) + torch.testing.assert_close(vertices.abs().amax(dim=0), expected_extents, rtol=1e-6, atol=1e-6) + + triangles = vertices[faces].to(torch.float64) + signed_volume = torch.sum( + torch.sum(triangles[:, 0] * torch.cross(triangles[:, 1], triangles[:, 2], dim=1), dim=1) + ) / 6.0 + expected_volume = 2.0 * math.pi**2 * major_radius * minor_radius**2 + assert abs(float(signed_volume)) == pytest.approx(expected_volume, rel=1e-2) + + 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/primitives.py b/witwin/core/geometry/primitives.py index 4968935..fa48b8b 100644 --- a/witwin/core/geometry/primitives.py +++ b/witwin/core/geometry/primitives.py @@ -427,7 +427,7 @@ def to_mesh(self, segments=16): for minor in range(minor_segments): phi = 2 * np.pi * minor / minor_segments vertices.append([ - np.cos(phi) * np.cos(theta), + np.cos(phi), np.cos(phi) * np.sin(theta), np.sin(phi), np.cos(theta),