Skip to content

mps-pointops 0.1.0

Choose a tag to compare

View all tags
@gamzerA gamzerA released this 30 Sep 19:40
· 222 commits to main since this release
v0.1.0
f94f94f

Correction: this release declares torch>=2.6, but it needs PyTorch 2.7 or later (torch.mps.compile_shader is new in 2.7). Use 0.1.1, which fixes the requirement.

First public release: point cloud ops for PyTorch on Apple Silicon GPUs (MPS), written as Metal kernels, plus drop-in stand-ins for the CUDA-only pointnet2_ops and knn_cuda packages.

Highlights

  • Metal kernels for farthest point sampling, k nearest neighbors and ball query (PyTorch3D-style first-K contract, with coordinate gradients).
  • Existing CUDA-only code runs unmodified: mps_pointops.compat.install() serves pointnet2_ops (furthest_point_sample, gather_operation, grouping_operation, ball_query) and knn_cuda (KNN).
  • Same results as on CUDA: MulSen-AD's Point-MAE 3D anomaly detector, re-fit on the Mac GPU through the stand-ins, gave the same object and 3D-label AUROC/AP as the original CUDA runs in all 45 runs (15 categories x 3 seeds). Per-sample scores differ by at most 9.85e-5 relative, with the same ranking in every run.

Performance

Apple M5 Pro, PyTorch 2.14.1, 100,000 points, batch 1, median of 5 runs. Timings move by a few ms between runs.

op mps-pointops (Metal) best CPU library plain PyTorch on MPS
FPS, 1024 samples 33.3 ms 174.0 ms (fpsample) 247.5 ms
kNN, 1024 queries, k=128 6.1 ms 17.6 ms (scipy cKDTree) 120.9 ms
Ball query, 1024 queries, K=64, r=0.1 13.2 ms 21.3 ms (scipy cKDTree) 286.4 ms

On 30 real MulSen-AD point clouds (21k to 117k points), MulSen-AD's own grouping step (FPS + kNN) took a median of 36.0 ms, against 163.4 ms for fpsample + scipy and 330.6 ms for plain PyTorch on MPS.

The Metal kernels returned the same indices as the references on every benchmark input. Full tables and the correctness checks behind them are in bench/results/.

Install

Needs an Apple Silicon Mac and PyTorch with MPS (torch>=2.6). Tested with PyTorch 2.14.1 on macOS 26.5, M5 Pro.

pip install "git+https://github.com/gamzerA/mps-pointops.git@v0.1.0"

Or install the wheel attached below. SHA256SUMS lists the checksums of both files.

Known limitations

  • Tested on one machine (M5 Pro). Results from other Apple GPUs are welcome.
  • The kernels assume 32-wide simdgroups; the first call checks this and raises an error otherwise.
  • On MPS: float32 input for FPS and kNN, float32 or float16 for ball query, and k <= 256 for kNN.
  • Ball query on spatially sorted input is slower than scipy's KD-tree (32.5 ms against 23.2 ms at 100k points), because the first-K-in-input-order contract rules out reordering the scan.
  • The ball query API is not full PyTorch3D compatibility yet: no return_nn, 3D coordinates only, and no lengths arguments in the public API.
  • The CUDA stand-ins can resolve near ties differently from the CUDA packages. As in pointnet2_ops, pointnet2_utils.ball_query returns index 0 in every slot for a query with no neighbor.
  • No torch_cluster / PyG style API (flat inputs with batch vectors) yet.

License

Apache-2.0. The ball query kernel and its Python wrapper are MIT; see LICENSES/MIT-ball-query.txt.

Developed with the help of AI coding assistants.

Assets 5
Loading