PyRXMesh 0.2.1rc0__tar.gz → 0.2.1rc2__tar.gz

{pyrxmesh-0.2.1rc0 → pyrxmesh-0.2.1rc2}/CMakeLists.txt

@@ -98,7 +98,10 @@ pybind11_add_module(_rxmesh MODULE
     src/bindings/handles.cu
     src/bindings/mesh.cu
     src/bindings/common.h
+    src/bindings/dispatch.h
+    src/bindings/plugin_launch.h
     src/bindings/dlpack_minimal.h
+    src/bindings/dlpack_utils.h
 
     src/bindings/py_attributes.cu
     src/bindings/py_attribute.h

{pyrxmesh-0.2.1rc0 → pyrxmesh-0.2.1rc2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyRXMesh
-Version: 0.2.1rc0
+Version: 0.2.1rc2
 Summary: Python bindings for RXMesh, a CUDA/C++ library for mesh processing on the GPU.
 Author: Ahmed Mahmoud
 License-Expression: BSD-2-Clause

{pyrxmesh-0.2.1rc0 → pyrxmesh-0.2.1rc2}/pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "scikit_build_core.build"
 
 [project]
 name = "PyRXMesh"
-version = "0.2.1rc0"
+version = "0.2.1rc2"
 description = "Python bindings for RXMesh, a CUDA/C++ library for mesh processing on the GPU."
 readme = "README.md"
 requires-python = ">=3.10"

pyrxmesh-0.2.1rc2/pyrxmesh/__init__.py

@@ -0,0 +1,168 @@
+"""PyRXMesh: Python bindings for the RXMesh GPU mesh-processing library.
+
+This package exposes:
+
+- The native classes (``RXMeshStatic``, ``DenseMatrix``, ``SparseMatrix``,
+  ``Attribute``, solvers, etc.) from the compiled ``_rxmesh`` extension.
+- Optional NumPy / SciPy / PyTorch interop methods attached to those classes
+  via the ``_numpy_interop``, ``_scipy_interop``, and ``_torch_interop``
+  submodules. The torch / scipy modules import their backing library lazily,
+  so attaching them is free if the user never calls them.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+
+_dll_handles = []
+
+
+def _add_windows_dll_dir(path: Path) -> None:
+    if os.name != "nt" or not path.exists():
+        return
+    try:
+        handle = os.add_dll_directory(str(path))
+    except (FileNotFoundError, OSError):
+        return
+    _dll_handles.append(handle)
+
+
+def _prepare_dll_search_path() -> None:
+    if os.name != "nt":
+        return
+
+    package_dir = Path(__file__).resolve().parent
+    _add_windows_dll_dir(package_dir)
+    _add_windows_dll_dir(package_dir / "bin")
+    _add_windows_dll_dir(package_dir / "lib")
+
+    for env_name in ("CUDA_PATH", "CUDA_HOME"):
+        cuda_root = os.environ.get(env_name)
+        if cuda_root:
+            _add_windows_dll_dir(Path(cuda_root) / "bin")
+
+    for env_name, cuda_root in os.environ.items():
+        if env_name.startswith("CUDA_PATH_V") and cuda_root:
+            _add_windows_dll_dir(Path(cuda_root) / "bin")
+
+    default_cuda_root = Path("C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA")
+    if default_cuda_root.exists():
+        for cuda_bin in sorted(default_cuda_root.glob("v*/bin"), reverse=True):
+            _add_windows_dll_dir(cuda_bin)
+
+
+_prepare_dll_search_path()
+
+_extension_import_error: ModuleNotFoundError | None = None
+
+try:
+    from ._rxmesh import (
+        Attribute,
+        CGSolver,
+        CholeskySolver,
+        DEdgeHandle,
+        DType,
+        DenseMatrix,
+        EdgeAttributeFloat32,
+        EdgeAttributeFloat64,
+        EdgeAttributeInt32,
+        EdgeAttributeInt8,
+        EdgeHandle,
+        ElementKind,
+        FaceAttributeFloat32,
+        FaceAttributeFloat64,
+        FaceAttributeInt32,
+        FaceAttributeInt8,
+        FaceHandle,
+        HessianSparseMatrix,
+        JacobianSparseMatrix,
+        Layout,
+        Location,
+        LogLevel,
+        Op,
+        PCGSolver,
+        QRSolver,
+        RXMeshStatic,
+        LUSolver,
+        SparseMatrix,
+        VertexAttributeFloat32,
+        VertexAttributeFloat64,
+        VertexAttributeInt32,
+        VertexAttributeInt8,
+        VertexHandle,
+        abi_version,
+        build_config_tag,
+        create_plane,
+        cuda_stream_synchronize,
+        cuDSSCholeskySolver,
+        has_cudss,
+        init,
+        show,
+    )
+except ModuleNotFoundError as exc:
+    if exc.name != f"{__name__}._rxmesh":
+        raise
+    _extension_import_error = exc
+    __all__ = []
+else:
+    __all__ = [
+        "Attribute",
+        "CGSolver",
+        "CholeskySolver",
+        "DEdgeHandle",
+        "DType",
+        "DenseMatrix",
+        "EdgeAttributeFloat32",
+        "EdgeAttributeFloat64",
+        "EdgeAttributeInt32",
+        "EdgeAttributeInt8",
+        "EdgeHandle",
+        "ElementKind",
+        "FaceAttributeFloat32",
+        "FaceAttributeFloat64",
+        "FaceAttributeInt32",
+        "FaceAttributeInt8",
+        "FaceHandle",
+        "HessianSparseMatrix",
+        "JacobianSparseMatrix",
+        "Layout",
+        "Location",
+        "LogLevel",
+        "Op",
+        "PCGSolver",
+        "QRSolver",
+        "RXMeshStatic",
+        "LUSolver",
+        "SparseMatrix",
+        "VertexAttributeFloat32",
+        "VertexAttributeFloat64",
+        "VertexAttributeInt32",
+        "VertexAttributeInt8",
+        "VertexHandle",
+        "abi_version",
+        "build_config_tag",
+        "create_plane",
+        "cuda_stream_synchronize",
+        "cuDSSCholeskySolver",
+        "has_cudss",
+        "init",
+        "show",
+    ]
+
+    # Attach optional interop methods to the native classes. Each submodule
+    # only does the patching; backing libraries (torch, scipy) are imported
+    # lazily inside the patched methods, so loading these modules is cheap.
+    from . import _numpy_interop  # noqa: F401  (side-effect: monkey-patches)
+    from . import _scipy_interop  # noqa: F401  (side-effect: monkey-patches)
+    from . import _torch_interop  # noqa: F401  (side-effect: monkey-patches)
+
+
+def __getattr__(name: str):
+    if _extension_import_error is not None:
+        raise ModuleNotFoundError(
+            "pyrxmesh._rxmesh is not installed yet. Build/install PyRXMesh "
+            "before using runtime mesh bindings."
+        ) from _extension_import_error
+    raise AttributeError(name)

pyrxmesh-0.2.1rc2/pyrxmesh/_numpy_interop.py

@@ -0,0 +1,46 @@
+"""NumPy interop helpers attached to ``SparseMatrix``.
+
+Overrides ``SparseMatrix.multiply_vector`` so it accepts a NumPy 1D/2D array
+in addition to a native ``DenseMatrix``. NumPy is already a hard PyRXMesh
+dependency, so this module is always importable.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from . import DenseMatrix, Location, SparseMatrix
+
+
+_native_multiply_vector = SparseMatrix.multiply_vector
+
+
+def _sparse_matrix_multiply_vector(self, vector, stream=None):
+    if isinstance(vector, DenseMatrix):
+        return _native_multiply_vector(self, vector, stream)
+
+    values = np.asarray(vector)
+    if values.ndim == 1:
+        values = values.reshape(-1, 1)
+    if values.ndim != 2 or values.shape[1] != 1:
+        raise ValueError(
+            "SparseMatrix.multiply_vector() expects a 1D array or a "
+            "DenseMatrix/array with shape (cols, 1)."
+        )
+    if values.shape[0] != self.cols:
+        raise ValueError(
+            "SparseMatrix.multiply_vector() vector length must match "
+            "matrix.cols."
+        )
+
+    dense = DenseMatrix(
+        self.cols,
+        1,
+        dtype=self.dtype,
+        location=Location.ALL,
+    )
+    dense.from_numpy_copy(values, target=Location.ALL, stream=stream)
+    return _native_multiply_vector(self, dense, stream)
+
+
+SparseMatrix.multiply_vector = _sparse_matrix_multiply_vector

pyrxmesh-0.2.1rc2/pyrxmesh/_scipy_interop.py

@@ -0,0 +1,41 @@
+"""SciPy interop helpers attached to ``SparseMatrix``.
+
+This module is imported automatically by ``pyrxmesh/__init__.py`` and patches
+``to_scipy_csr`` / ``to_scipy_csr_copy`` onto ``SparseMatrix``. ``scipy`` is
+imported lazily on first call, so importing this module is cheap.
+"""
+
+from __future__ import annotations
+
+from . import Location, SparseMatrix
+
+
+def _require_scipy_sparse(api):
+    try:
+        import scipy.sparse as scipy_sparse
+    except ImportError as exc:
+        raise ImportError(
+            f"SciPy is required for SparseMatrix.{api}(). "
+            "Install scipy or use to_numpy()/to_numpy_copy()."
+        ) from exc
+    return scipy_sparse
+
+
+def _sparse_matrix_to_scipy_csr(self):
+    scipy_sparse = _require_scipy_sparse("to_scipy_csr")
+    row_ptr, col_idx, values = self.to_numpy(Location.HOST)
+    return scipy_sparse.csr_matrix(
+        (values, col_idx, row_ptr), shape=self.shape
+    )
+
+
+def _sparse_matrix_to_scipy_csr_copy(self, source=Location.HOST, stream=None):
+    scipy_sparse = _require_scipy_sparse("to_scipy_csr_copy")
+    row_ptr, col_idx, values = self.to_numpy_copy(source=source, stream=stream)
+    return scipy_sparse.csr_matrix(
+        (values, col_idx, row_ptr), shape=self.shape
+    )
+
+
+SparseMatrix.to_scipy_csr = _sparse_matrix_to_scipy_csr
+SparseMatrix.to_scipy_csr_copy = _sparse_matrix_to_scipy_csr_copy

pyrxmesh-0.2.1rc2/pyrxmesh/_torch_interop.py

@@ -0,0 +1,206 @@
+"""PyTorch interop helpers attached to native PyRXMesh classes.
+
+This module is imported automatically by ``pyrxmesh/__init__.py`` and patches
+``to_torch`` / ``from_torch_copy`` (and friends) onto the relevant native
+classes. ``torch`` itself is imported lazily on first use, so importing this
+module does not pull torch into the process.
+"""
+
+from __future__ import annotations
+
+from . import (
+    Attribute,
+    DenseMatrix,
+    EdgeAttributeFloat32,
+    EdgeAttributeFloat64,
+    EdgeAttributeInt32,
+    EdgeAttributeInt8,
+    FaceAttributeFloat32,
+    FaceAttributeFloat64,
+    FaceAttributeInt32,
+    FaceAttributeInt8,
+    Location,
+    SparseMatrix,
+    VertexAttributeFloat32,
+    VertexAttributeFloat64,
+    VertexAttributeInt32,
+    VertexAttributeInt8,
+)
+
+
+def _require_torch():
+    try:
+        import torch
+    except ImportError as exc:
+        raise ImportError(
+            "PyTorch is required for PyRXMesh torch interop helpers. "
+            "Install torch or use the DLPack/NumPy APIs directly."
+        ) from exc
+    return torch
+
+
+def _torch_dtype_name(dtype):
+    torch = _require_torch()
+    if dtype == torch.float32:
+        return "float32"
+    if dtype == torch.float64:
+        return "float64"
+    if dtype == torch.int32:
+        return "int32"
+    raise TypeError(f"Unsupported torch dtype for PyRXMesh copy: {dtype}")
+
+
+class _DenseMatrixDlpackView:
+    """Wrap a ``DenseMatrix`` so ``torch.utils.dlpack.from_dlpack`` can
+    forward its ``stream`` argument into ``DenseMatrix.to_dlpack``."""
+
+    def __init__(self, matrix, location):
+        self._matrix = matrix
+        self._location = location
+        self._is_host = location == Location.HOST
+        self._is_device = location == Location.DEVICE
+
+    def __dlpack__(self, stream=None):
+        return self._matrix.to_dlpack(self._location, stream=stream)
+
+    def __dlpack_device__(self):
+        if self._is_host:
+            return (1, 0)
+        if self._is_device:
+            torch = _require_torch()
+            return (2, torch.cuda.current_device())
+        raise ValueError(
+            "DenseMatrix.to_torch() location must be Location.HOST or "
+            "Location.DEVICE."
+        )
+
+
+class _SparseMatrixDlpackView:
+    """Wrap a single ``SparseMatrix`` CSR component for
+    ``torch.utils.dlpack.from_dlpack``."""
+
+    def __init__(self, matrix, location, component):
+        self._matrix = matrix
+        self._location = location
+        self._component = component
+        self._is_host = location == Location.HOST
+        self._is_device = location == Location.DEVICE
+
+    def __dlpack__(self, stream=None):
+        if self._component == "row_ptr":
+            return self._matrix._row_ptr_dlpack(self._location, stream=stream)
+        if self._component == "col_idx":
+            return self._matrix._col_indices_dlpack(self._location, stream=stream)
+        if self._component == "values":
+            return self._matrix._values_dlpack(self._location, stream=stream)
+        raise ValueError("Unsupported SparseMatrix DLPack component.")
+
+    def __dlpack_device__(self):
+        if self._is_host:
+            return (1, 0)
+        if self._is_device:
+            torch = _require_torch()
+            return (2, torch.cuda.current_device())
+        raise ValueError(
+            "SparseMatrix.to_torch() location must be Location.HOST or "
+            "Location.DEVICE."
+        )
+
+
+def _dense_matrix_to_torch(self, location=Location.DEVICE):
+    torch = _require_torch()
+    return torch.utils.dlpack.from_dlpack(
+        _DenseMatrixDlpackView(self, location)
+    )
+
+
+def _dense_matrix_from_torch_copy(source):
+    tensor = source.detach() if hasattr(source, "detach") else source
+    return DenseMatrix.from_dlpack_copy(tensor)
+
+
+def _sparse_matrix_to_torch(self, location=Location.DEVICE):
+    torch = _require_torch()
+    crow = torch.utils.dlpack.from_dlpack(
+        _SparseMatrixDlpackView(self, location, "row_ptr")
+    )
+    col = torch.utils.dlpack.from_dlpack(
+        _SparseMatrixDlpackView(self, location, "col_idx")
+    )
+    val = torch.utils.dlpack.from_dlpack(
+        _SparseMatrixDlpackView(self, location, "values")
+    )
+    return torch.sparse_csr_tensor(
+        crow,
+        col,
+        val,
+        size=self.shape,
+        device=val.device,
+    )
+
+
+def _sparse_matrix_from_torch_values_copy(
+    self,
+    values,
+    target=Location.ALL,
+    stream=None,
+):
+    tensor = values.detach() if hasattr(values, "detach") else values
+    self.from_dlpack_values_copy(tensor, target=target, stream=stream)
+    return self
+
+
+def _sparse_matrix_from_torch_copy(source, dtype=None, stream=None):
+    torch = _require_torch()
+    if source.layout != torch.sparse_csr:
+        raise TypeError(
+            "SparseMatrix.from_torch_copy() expects a torch sparse CSR "
+            "tensor."
+        )
+    value_dtype = dtype or _torch_dtype_name(source.values().dtype)
+    return SparseMatrix.from_dlpack_copy(
+        source.crow_indices().detach(),
+        source.col_indices().detach(),
+        source.values().detach(),
+        tuple(source.shape),
+        dtype=value_dtype,
+        stream=stream,
+    )
+
+
+def _attribute_to_torch(self, location=Location.DEVICE):
+    torch = _require_torch()
+    return torch.utils.dlpack.from_dlpack(self.to_dlpack(location))
+
+
+def _attribute_from_torch_copy(self, values, target=Location.ALL):
+    tensor = values.detach() if hasattr(values, "detach") else values
+    self.from_dlpack_copy(tensor, target=target)
+    return self
+
+
+DenseMatrix.to_torch = _dense_matrix_to_torch
+DenseMatrix.from_torch_copy = staticmethod(_dense_matrix_from_torch_copy)
+SparseMatrix.to_torch = _sparse_matrix_to_torch
+SparseMatrix.from_torch_copy = staticmethod(_sparse_matrix_from_torch_copy)
+SparseMatrix.from_torch_values_copy = _sparse_matrix_from_torch_values_copy
+
+_attribute_types = (
+    Attribute,
+    VertexAttributeFloat32,
+    VertexAttributeFloat64,
+    VertexAttributeInt32,
+    VertexAttributeInt8,
+    EdgeAttributeFloat32,
+    EdgeAttributeFloat64,
+    EdgeAttributeInt32,
+    EdgeAttributeInt8,
+    FaceAttributeFloat32,
+    FaceAttributeFloat64,
+    FaceAttributeInt32,
+    FaceAttributeInt8,
+)
+for _attribute_type in _attribute_types:
+    _attribute_type.to_torch = _attribute_to_torch
+    _attribute_type.from_torch_copy = _attribute_from_torch_copy
+del _attribute_type, _attribute_types

pyrxmesh-0.2.1rc2/src/bindings/common.h

@@ -0,0 +1,169 @@
+#pragma once
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <cstdint>
+#include <functional>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <type_traits>
+#include <utility>
+#include <variant>
+#include <vector>
+
+#include <pybind11/functional.h>
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <cuda_runtime_api.h>
+
+#include "pyrxmesh/plugin_api.h"
+#include "rxmesh/handle.h"
+#include "rxmesh/matrix/dense_matrix.h"
+#include "rxmesh/rxmesh_static.h"
+#include "rxmesh/types.h"
+#include "rxmesh/util/log.h"
+#include "rxmesh/util/macros.h"
+
+#include "bindings/dispatch.h"
+
+#if USE_POLYSCOPE
+#include "polyscope/polyscope.h"
+#endif
+
+namespace py = pybind11;
+
+namespace pyrxmesh_py {
+
+inline constexpr uint32_t    plugin_abi_version = PYRXMESH_PLUGIN_ABI_VERSION;
+inline constexpr const char* build_config_tag   = PYRXMESH_BUILD_CONFIG;
+
+inline rxmesh::locationT parse_location(int location)
+{
+    return static_cast<rxmesh::locationT>(location);
+}
+
+inline int64_t cuda_stream_arg_value(py::object stream)
+{
+    if (stream.is_none()) {
+        return 1;
+    }
+    if (!py::isinstance<py::int_>(stream)) {
+        throw py::type_error("CUDA stream must be an integer or None.");
+    }
+    return stream.cast<int64_t>();
+}
+
+inline cudaStream_t parse_cuda_stream_arg(py::object stream)
+{
+    const int64_t value = cuda_stream_arg_value(std::move(stream));
+    if (value == 0) {
+        throw std::invalid_argument(
+            "CUDA stream value 0 is ambiguous and not supported.");
+    }
+    if (value == 1) {
+        return nullptr;
+    }
+    if (value == 2) {
+        return cudaStreamPerThread;
+    }
+    if (value > 2) {
+        return reinterpret_cast<cudaStream_t>(static_cast<uintptr_t>(value));
+    }
+    throw std::invalid_argument(
+        "CUDA stream value must be None, 1, 2, or a "
+        "positive raw stream pointer.");
+}
+
+inline void cuda_stream_synchronize_arg(py::object stream)
+{
+    using namespace rxmesh;
+    CUDA_ERROR(cudaStreamSynchronize(parse_cuda_stream_arg(std::move(stream))));
+}
+
+inline void make_stream_wait_for_legacy_default(cudaStream_t stream)
+{
+    using namespace rxmesh;
+    if (stream == nullptr) {
+        return;
+    }
+    cudaEvent_t event = nullptr;
+    CUDA_ERROR(cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
+    CUDA_ERROR(cudaEventRecord(event, nullptr));
+    CUDA_ERROR(cudaStreamWaitEvent(stream, event, 0));
+    CUDA_ERROR(cudaEventDestroy(event));
+}
+
+inline void synchronize_device_transfer(rxmesh::locationT source,
+                                        rxmesh::locationT target)
+{
+    using namespace rxmesh;
+    if (((source | target) & rxmesh::DEVICE) == rxmesh::DEVICE) {
+        CUDA_ERROR(cudaStreamSynchronize(nullptr));
+    }
+}
+
+inline rxmesh::layoutT parse_layout(int layout)
+{
+    return static_cast<rxmesh::layoutT>(layout);
+}
+
+inline void ensure_polyscope_available()
+{
+#if !USE_POLYSCOPE
+    throw std::runtime_error(
+        "RXMesh was built with RX_USE_POLYSCOPE=OFF; visualization is not "
+        "available.");
+#endif
+}
+
+inline void show_polyscope()
+{
+    ensure_polyscope_available();
+#if USE_POLYSCOPE
+    polyscope::show();
+#endif
+}
+
+inline glm::fvec3 sequence_to_fvec3(const py::sequence& values,
+                                    const char*         name)
+{
+    if (py::len(values) != 3) {
+        throw std::invalid_argument(std::string(name) +
+                                    " must contain exactly 3 values.");
+    }
+    return glm::fvec3(values[0].cast<float>(),
+                      values[1].cast<float>(),
+                      values[2].cast<float>());
+}
+
+inline py::array_t<float> vec3_to_numpy(const glm::vec3& value)
+{
+    py::array_t<float> out({static_cast<py::ssize_t>(3)});
+    auto               view = out.mutable_unchecked<1>();
+    view(0)                 = value[0];
+    view(1)                 = value[1];
+    view(2)                 = value[2];
+    return out;
+}
+
+void register_module_core(py::module_& m);
+void register_handles(py::module_& m);
+void register_attribute(py::module_& m);
+void register_dense_matrix(py::module_& m);
+void register_geometry(py::module_& m);
+void register_sparse_matrix(py::module_& m);
+void register_solvers(py::module_& m);
+void register_mesh(py::module_& m);
+
+
+py::object make_sparse_matrix_from_mesh(
+    std::shared_ptr<rxmesh::RXMeshStatic> mesh,
+    rxmesh::Op                            op,
+    std::string                           dtype);
+}  // namespace pyrxmesh_py