Mesh: GPU-Accelerated Remesh

[mehdiataei]

PhysicsNeMo Pull Request

Description

This PR adds mesh remeshing to PhysicsNeMo using Warp.

The public API is available through physicsnemo.mesh.remeshing.remesh and Mesh.remesh. It supports CPU and CUDA meshes, preserves the input device, and provides controls for target vertex count, clustering, projection, topology reconstruction, and cleanup.

The implementation includes:

• A Warp-based remeshing functional and associated kernels.
• CPU and CUDA execution support.
• Public functional and object-oriented mesh APIs.
• Dynamic-shape and torch.compile support through FunctionSpec.
• Unit tests covering the public API, functional API, CPU and CUDA execution, validation, compilation, and topology.
• Documentation, examples, and ASV benchmarks.
• Updated comparison and scaling figures.
• PyACVD is removed as a dependency.

CUDA remeshing can be up to 300× faster than PyACVD on CPU.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards. (Not applicable; this PR does not add or modify a model.)

Dependencies

No new dependencies are required. The implementation uses the existing Warp dependency, and PyACVD is not required.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

This PR replaces the PyACVD-based remeshing backend with a Warp-accelerated implementation that runs on both CPU and CUDA (up to 300\u00d7 faster), exposing the feature through physicsnemo.mesh.remeshing.remesh and Mesh.remesh. The implementation uses area-weighted centroidal Voronoi clustering (Lloyd iterations), BVH-based centroid projection back onto the source surface, and a PyTorch-side topology reconstruction and cleanup pass.

• The core pipeline (launch_forward.py) is well-structured: two seeding strategies (FPS for small targets, voxel stratification for large ones), configurable Lloyd iterations, and thorough cleanup of degenerate and non-manifold faces.
• The custom-op registration, FakeTensor fake implementation, and torch.compile support are correctly set up with appropriate tags.
• The project_centroids_to_surface kernel manually recomputes barycentric interpolation with vertex weights that do not match Warp's documented convention (mesh_eval_position uses (1-u-v)*p0 + u*p1 + v*p2; the kernel uses the transposed form), so projected centroids land at a valid but non-closest surface point.

Important Files Changed

Filename	Overview
physicsnemo/nn/functional/geometry/remeshing/_warp_impl/_kernels/project_centroids_to_surface.py	Projection kernel manually computes barycentric interpolation with swapped vertex weights relative to Warp's documented convention, causing centroids to be placed at wrong positions on the source surface
physicsnemo/nn/functional/geometry/remeshing/_warp_impl/launch_forward.py	Core pipeline orchestration: FPS and voxel seeding, Lloyd iterations, surface projection, connectivity reconstruction; logic is sound apart from the barycentric issue in the kernel it calls
physicsnemo/nn/functional/geometry/remeshing/remeshing.py	FunctionSpec-based public tensor API: input validation, dispatch, fake-tensor registration, and benchmark cases are well structured and correct
physicsnemo/nn/functional/geometry/remeshing/_warp_impl/op.py	Custom-op registration with appropriate tags, correct sentinel encoding for max_iterations, and proper FakeTensor shape description
physicsnemo/nn/functional/geometry/remeshing/_config.py	Frozen dataclass with thorough type and range validation for all five tuning parameters; clean and correct
physicsnemo/mesh/remeshing/_remeshing.py	Mesh-level wrapper with solid validation, correct global_data preservation, and proper delegation to the tensor functional
physicsnemo/mesh/mesh.py	Adds Mesh.remesh() as a thin delegation to the functional; docstring and signature match the functional's contract
physicsnemo/utils/_index_tuple_ops.py	Fast packed-key integer-tuple deduplication utility with correct int64 overflow guard and fallback to torch.unique
test/mesh/remeshing/test_remesh_warp.py	Comprehensive CUDA test suite; one assertion uses exact vertex count equality where the range form should be used, and the sphere symmetry masks the barycentric projection bug
test/nn/functional/geometry/test_remeshing.py	Good coverage of input validation, FakeTensor propagation, torch.compile, opcheck, and internal helpers

_{Reviews (1): Last reviewed commit: "Enable Warp remeshing on CPU" | Re-trigger Greptile}

[greptile-apps]

Incorrect barycentric interpolation — centroids projected to wrong surface point

The manual barycentric formula swaps the vertex weights relative to Warp's convention. Warp's mesh_eval_position (and the underlying closest_point_to_triangle in intersect.h) uses the standard Möller-Trumbore parameterisation: closest = (1-u-v)*p0 + u*p1 + v*p2. The kernel instead computes u*p0 + v*p1 + (1-u-v)*p2, which places the centroid at a different point inside the same triangle — not the closest point on the surface. Warp examples universally use wp.mesh_eval_position(mesh_id, query.face, query.u, query.v) for exactly this reason.

The projected point is still geometrically on the surface (the coordinates are a valid convex combination), so the tests pass on symmetric geometry such as a sphere, but centroids on anisotropic or high-curvature surfaces will drift, producing worse mesh quality than intended. The simplest fix is to replace the manual computation with the built-in call: centroids[centroid_index] = wp.mesh_eval_position(mesh_id, query.face, query.u, query.v).

[peterdsharpe]

Submitting an initial partial review! This PR seems to have good bones, but a lot of new API surface here violates repo rules by not including jaxtyping on tensor inputs. Many functions are also missing complete (i.e., conforming to repo-wide style rules, with explanations for parameters) docstrings. Would you mind adding these, to aid in the review? Thanks!

[mehdiataei]

@peterdsharpe Regarding the use of bultins.int:

This is intentional. tensorclass generates methods like .int(), .bool(), and .float(). With Python 3.14’s lazy annotation lookup, a bare int | None inside the class can resolve to that generated method instead of the builtin type and break signature inspection. Using builtins.int avoids the collision. e.g.

@tensorclass
class Example:
    value: torch.Tensor

    # Bare `int | None` resolves to the generated `Example.int` method on Python 3.14.
    def method(self, count: builtins.int | None = None):
        pass

[mehdiataei]

For future reference based on our conversations:

I dug into this a bit more and the behavior is specific to Python 3.14, where annotations are evaluated lazily.

print(int) inside the class body still shows the builtin because it runs before @tensorclass. However, the method annotation is evaluated later, after tensorclass has added the .int() method. At that point, bare int resolves to the generated method.

Here’s a minimal repro:

import builtins
from typing import get_type_hints

import torch
from tensordict import tensorclass


@tensorclass
class Bare:
    value: torch.Tensor

    def method(self, count: int | None = None):
        pass


@tensorclass
class Qualified:
    value: torch.Tensor

    def method(self, count: builtins.int | None = None):
        pass


try:
    print(get_type_hints(Bare.method))
except TypeError as error:
    print("Bare annotation failed:", error)

print("Qualified annotation:", get_type_hints(Qualified.method))

On Python 3.14, this prints:

Bare annotation failed: unsupported operand type(s) for |: 'function' and 'NoneType'
Qualified annotation: {'count': int | None}

So I think keeping builtins.int is the safest option here. It explicitly identifies the builtin and avoids depending on when or where the annotation gets evaluated. While using string types is an alternative my understanding based on Python doc is that this is more future proof and safer.

[peterdsharpe]

This looks good — nice work. I checked out the branch and ran the test suites locally (109 passed, CPU + CUDA).

Key comments to fix below:

• The 2^24 torch.multinomial limit: remeshing currently crashes above ~16.7M input vertices (repro and a drop-in fix in the inline comment).
• The tuning knobs (search_radius_scale, hash_grid_resolution, etc.) might be best left out of the Mesh.remesh signature and kept on the tensor functional — if you agree? Once these ship in a release we are committed to them.
• Validation logic should be deduplicated between mesh.remeshing.remesh and the functional's _validate_inputs.
• bunny.pt should get regenerated so the checked-in asset matches the new remesher.

Everything else below is a suggestion / nit / optional.

[peterdsharpe]

Following up on the 2^24 torch.multinomial finding from my review: rather than patching the sampler in place, it might be cleaner to replace both seeding paths (_select_fps_centroids and _select_stratified_centroids) with the existing mesh_poisson_disk_sample functional — "K well-spaced, area-weighted points on the surface" is exactly its output contract, and it's on the same Warp stack.

That would:

• Fix the crash as a side effect: its candidate draw is an area-weighted triangle CDF in Warp (no torch.multinomial, so no 2^24 cap).
• Delete ~120 lines here, plus three of the five tuning knobs (voxel_width_scale, farthest_point_threshold, and farthest_point_oversampling exist only to configure the hand-rolled seeding) — which also largely resolves my comment on the Mesh.remesh signature.
• Probably improve seeding: blue-noise beats voxel-stratification-with-random-fill as a CVT initialization, and the seeds would lie on the surface rather than at input vertices — the input vertex distribution is the thing remeshing is trying to fix.

One mode note: weighted_sample_elimination gives exactly K but has a host-side elimination loop, so it may regress wall-clock at large K. Since n_clusters is a ceiling rather than a guarantee anyway, GPU-native dart throwing with max_points=n_clusters and radius ≈ 0.7 * sqrt(area / n_clusters) should preserve the output <= n_clusters contract — worth benchmarking both. (It also takes a random_seed arg, which would address the fixed-seed nit for free.)

[mehdiataei]

Thanks for the suggestion! I benchmarked both Poisson modes against the current initializer. The current approach generates exactly $K$ area-weighted vertex seeds: FPS for small targets and voxel stratification for larger ones, using a chunked sampler that avoids the torch.multinomial limit. The CVT/Lloyd remeshing pipeline itself is still Warp-based and runs on either CPU or CUDA.

Weighted elimination does produce an exact $K$, but it scales poorly because of the 5× candidate pool and the host-side elimination loop. Dart throwing is faster, but it only guarantees up to $K$ points, so it needs additional fill/retry logic. Both also duplicate mesh preprocessing. On a 10.5M-vertex mesh, that added about 2.24 GiB of memory and 25–27 ms just for seeding.

There's also a compatibility issue with this CVT formulation. Since Lloyd assignment integrates over the weighted input vertices, triangle-interior Poisson seeds can start out without owning any input vertex. With $K=N=642$, the current initializer retained all 642 clusters, while dart throwing retained 557 and weighted elimination retained 569.

The replacement sampler has been tested with 100 million candidate entries, and the full Warp remeshing pipeline has been validated on a 41.9M-vertex mesh, so I think the current initializer is the safer choice for this PR. That said, dart throwing is still worth revisiting later once preprocessing is shared and there's an exact-K fallback.

[mehdiataei]

/ok to test c141e8f

[peterdsharpe]

LGTM, thank you!

[mehdiataei]

/ok to test c141e8f

[mehdiataei]

/ok to test 972c432