warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/utils.py

@@ -13,13 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import cProfile
 import ctypes
 import os
 import sys
 import time
 import warnings
-from typing import Any, Optional
+from typing import Any, Callable, Dict, List, Optional, Union
 
 import numpy as np
 
@@ -141,6 +143,8 @@ def radix_sort_pairs(keys, values, count: int):
             runtime.core.radix_sort_pairs_int_host(keys.ptr, values.ptr, count)
         elif keys.dtype == wp.float32 and values.dtype == wp.int32:
             runtime.core.radix_sort_pairs_float_host(keys.ptr, values.ptr, count)
+        elif keys.dtype == wp.int64 and values.dtype == wp.int32:
+            runtime.core.radix_sort_pairs_int64_host(keys.ptr, values.ptr, count)
         else:
             raise RuntimeError("Unsupported data type")
     elif keys.device.is_cuda:
@@ -148,6 +152,90 @@ def radix_sort_pairs(keys, values, count: int):
             runtime.core.radix_sort_pairs_int_device(keys.ptr, values.ptr, count)
         elif keys.dtype == wp.float32 and values.dtype == wp.int32:
             runtime.core.radix_sort_pairs_float_device(keys.ptr, values.ptr, count)
+        elif keys.dtype == wp.int64 and values.dtype == wp.int32:
+            runtime.core.radix_sort_pairs_int64_device(keys.ptr, values.ptr, count)
+        else:
+            raise RuntimeError("Unsupported data type")
+
+
+def segmented_sort_pairs(
+    keys,
+    values,
+    count: int,
+    segment_start_indices: wp.array(dtype=wp.int32),
+    segment_end_indices: wp.array(dtype=wp.int32) = None,
+):
+    """Sort key-value pairs within segments.
+
+    This function performs a segmented sort of key-value pairs, where the sorting is done independently within each segment.
+    The segments are defined by their start and optionally end indices.
+
+    Args:
+        keys: Array of keys to sort. Must be of type int32 or float32.
+        values: Array of values to sort along with keys. Must be of type int32.
+        count: Number of elements to sort.
+        segment_start_indices: Array containing start index of each segment. Must be of type int32.
+            If segment_end_indices is None, this array must have length at least num_segments + 1,
+            and segment_end_indices will be inferred as segment_start_indices[1:].
+            If segment_end_indices is provided, this array must have length at least num_segments.
+        segment_end_indices: Optional array containing end index of each segment. Must be of type int32 if provided.
+            If None, segment_end_indices will be inferred from segment_start_indices[1:].
+            If provided, must have length at least num_segments.
+
+    Raises:
+        RuntimeError: If array storage devices don't match, if storage size is insufficient,
+                     if segment_start_indices is not of type int32, or if data types are unsupported.
+    """
+    if keys.device != values.device:
+        raise RuntimeError("Array storage devices do not match")
+
+    if count == 0:
+        return
+
+    if keys.size < 2 * count or values.size < 2 * count:
+        raise RuntimeError("Array storage must be large enough to contain 2*count elements")
+
+    from warp.context import runtime
+
+    if segment_start_indices.dtype != wp.int32:
+        raise RuntimeError("segment_start_indices array must be of type int32")
+
+    # Handle case where segment_end_indices is not provided
+    if segment_end_indices is None:
+        num_segments = max(0, segment_start_indices.size - 1)
+
+        segment_end_indices = segment_start_indices[1:]
+        segment_end_indices_ptr = segment_end_indices.ptr
+        segment_start_indices_ptr = segment_start_indices.ptr
+    else:
+        if segment_end_indices.dtype != wp.int32:
+            raise RuntimeError("segment_end_indices array must be of type int32")
+
+        num_segments = segment_start_indices.size
+
+        segment_end_indices_ptr = segment_end_indices.ptr
+        segment_start_indices_ptr = segment_start_indices.ptr
+
+    if keys.device.is_cpu:
+        if keys.dtype == wp.int32 and values.dtype == wp.int32:
+            runtime.core.segmented_sort_pairs_int_host(
+                keys.ptr, values.ptr, count, segment_start_indices_ptr, segment_end_indices_ptr, num_segments
+            )
+        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
+            runtime.core.segmented_sort_pairs_float_host(
+                keys.ptr, values.ptr, count, segment_start_indices_ptr, segment_end_indices_ptr, num_segments
+            )
+        else:
+            raise RuntimeError("Unsupported data type")
+    elif keys.device.is_cuda:
+        if keys.dtype == wp.int32 and values.dtype == wp.int32:
+            runtime.core.segmented_sort_pairs_int_device(
+                keys.ptr, values.ptr, count, segment_start_indices_ptr, segment_end_indices_ptr, num_segments
+            )
+        elif keys.dtype == wp.float32 and values.dtype == wp.int32:
+            runtime.core.segmented_sort_pairs_float_device(
+                keys.ptr, values.ptr, count, segment_start_indices_ptr, segment_end_indices_ptr, num_segments
+            )
         else:
             raise RuntimeError("Unsupported data type")
 
@@ -673,37 +761,38 @@ class ScopedTimer:
 
     def __init__(
         self,
-        name,
-        active=True,
-        print=True,
-        detailed=False,
-        dict=None,
-        use_nvtx=False,
-        color="rapids",
-        synchronize=False,
-        cuda_filter=0,
-        report_func=None,
-        skip_tape=False,
+        name: str,
+        active: bool = True,
+        print: bool = True,
+        detailed: bool = False,
+        dict: Optional[Dict[str, List[float]]] = None,
+        use_nvtx: bool = False,
+        color: Union[int, str] = "rapids",
+        synchronize: bool = False,
+        cuda_filter: int = 0,
+        report_func: Optional[Callable[[List[TimingResult], str], None]] = None,
+        skip_tape: bool = False,
     ):
         """Context manager object for a timer
 
         Parameters:
-            name (str): Name of timer
-            active (bool): Enables this timer
-            print (bool): At context manager exit, print elapsed time to sys.stdout
-            detailed (bool): Collects additional profiling data using cProfile and calls ``print_stats()`` at context exit
-            dict (dict): A dictionary of lists to which the elapsed time will be appended using ``name`` as a key
-            use_nvtx (bool): If true, timing functionality is replaced by an NVTX range
-            color (int or str): ARGB value (e.g. 0x00FFFF) or color name (e.g. 'cyan') associated with the NVTX range
-            synchronize (bool): Synchronize the CPU thread with any outstanding CUDA work to return accurate GPU timings
-            cuda_filter (int): Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
-            report_func (Callable): A callback function to print the activity report (``wp.timing_print()`` is used by default)
-            skip_tape (bool): If true, the timer will not be recorded in the tape
+            name: Name of timer
+            active: Enables this timer
+            print: At context manager exit, print elapsed time to ``sys.stdout``
+            detailed: Collects additional profiling data using cProfile and calls ``print_stats()`` at context exit
+            dict: A dictionary of lists to which the elapsed time will be appended using ``name`` as a key
+            use_nvtx: If true, timing functionality is replaced by an NVTX range
+            color: ARGB value (e.g. 0x00FFFF) or color name (e.g. 'cyan') associated with the NVTX range
+            synchronize: Synchronize the CPU thread with any outstanding CUDA work to return accurate GPU timings
+            cuda_filter: Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
+            report_func: A callback function to print the activity report.
+              If ``None``,  :func:`wp.timing_print() <timing_print>` will be used.
+            skip_tape: If true, the timer will not be recorded in the tape
 
         Attributes:
             extra_msg (str): Can be set to a string that will be added to the printout at context exit.
             elapsed (float): The duration of the ``with`` block used with this object
-            timing_results (list[TimingResult]): The list of activity timing results, if collection was requested using ``cuda_filter``
+            timing_results (List[TimingResult]): The list of activity timing results, if collection was requested using ``cuda_filter``
         """
         self.name = name
         self.active = active and self.enabled
@@ -799,7 +888,7 @@ class ScopedTimer:
 
 # Allow temporarily enabling/disabling mempool allocators
 class ScopedMempool:
-    def __init__(self, device, enable: bool):
+    def __init__(self, device: Devicelike, enable: bool):
         self.device = wp.get_device(device)
         self.enable = enable
 
@@ -813,7 +902,7 @@ class ScopedMempool:
 
 # Allow temporarily enabling/disabling mempool access
 class ScopedMempoolAccess:
-    def __init__(self, target_device, peer_device, enable: bool):
+    def __init__(self, target_device: Devicelike, peer_device: Devicelike, enable: bool):
         self.target_device = target_device
         self.peer_device = peer_device
         self.enable = enable
@@ -828,7 +917,7 @@ class ScopedMempoolAccess:
 
 # Allow temporarily enabling/disabling peer access
 class ScopedPeerAccess:
-    def __init__(self, target_device, peer_device, enable: bool):
+    def __init__(self, target_device: Devicelike, peer_device: Devicelike, enable: bool):
         self.target_device = target_device
         self.peer_device = peer_device
         self.enable = enable
@@ -842,7 +931,7 @@ class ScopedPeerAccess:
 
 
 class ScopedCapture:
-    def __init__(self, device=None, stream=None, force_module_load=None, external=False):
+    def __init__(self, device: Devicelike = None, stream=None, force_module_load=None, external=False):
         self.device = device
         self.stream = stream
         self.force_module_load = force_module_load
@@ -868,27 +957,6 @@ class ScopedCapture:
                 self.active = False
 
 
-# helper kernels for adj_matmul
-@wp.kernel
-def add_kernel_2d(x: wp.array2d(dtype=Any), acc: wp.array2d(dtype=Any), beta: Any):
-    i, j = wp.tid()
-
-    x[i, j] = x[i, j] + beta * acc[i, j]
-
-
-@wp.kernel
-def add_kernel_3d(x: wp.array3d(dtype=Any), acc: wp.array3d(dtype=Any), beta: Any):
-    i, j, k = wp.tid()
-
-    x[i, j, k] = x[i, j, k] + beta * acc[i, j, k]
-
-
-# explicit instantiations of generic kernels for adj_matmul
-for T in [wp.float16, wp.float32, wp.float64]:
-    wp.overload(add_kernel_2d, [wp.array2d(dtype=T), wp.array2d(dtype=T), T])
-    wp.overload(add_kernel_3d, [wp.array3d(dtype=T), wp.array3d(dtype=T), T])
-
-
 def check_p2p():
     """Check if the machine is configured properly for peer-to-peer transfers.
 
@@ -927,31 +995,28 @@ class timing_result_t(ctypes.Structure):
 
 
 class TimingResult:
-    """Timing result for a single activity.
+    """Timing result for a single activity."""
 
-    Parameters:
-        raw_result (warp.utils.timing_result_t): The result structure obtained from C++ (internal use only)
+    def __init__(self, device, name, filter, elapsed):
+        self.device: warp.context.Device = device
+        """The device where the activity was recorded."""
 
-    Attributes:
-        device (warp.Device): The device where the activity was recorded.
-        name (str): The activity name.
-        filter (int): The type of activity (e.g., ``warp.TIMING_KERNEL``).
-        elapsed (float): The elapsed time in milliseconds.
-    """
+        self.name: str = name
+        """The activity name."""
 
-    def __init__(self, device, name, filter, elapsed):
-        self.device = device
-        self.name = name
-        self.filter = filter
-        self.elapsed = elapsed
+        self.filter: int = filter
+        """The type of activity (e.g., ``warp.TIMING_KERNEL``)."""
+
+        self.elapsed: float = elapsed
+        """The elapsed time in milliseconds."""
 
 
-def timing_begin(cuda_filter=TIMING_ALL, synchronize=True):
+def timing_begin(cuda_filter: int = TIMING_ALL, synchronize: bool = True) -> None:
     """Begin detailed activity timing.
 
     Parameters:
-        cuda_filter (int): Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
-        synchronize (bool): Whether to synchronize all CUDA devices before timing starts
+        cuda_filter: Filter flags for CUDA activity timing, e.g. ``warp.TIMING_KERNEL`` or ``warp.TIMING_ALL``
+        synchronize: Whether to synchronize all CUDA devices before timing starts
     """
 
     if synchronize:
@@ -960,14 +1025,14 @@ def timing_begin(cuda_filter=TIMING_ALL, synchronize=True):
     warp.context.runtime.core.cuda_timing_begin(cuda_filter)
 
 
-def timing_end(synchronize=True):
+def timing_end(synchronize: bool = True) -> List[TimingResult]:
     """End detailed activity timing.
 
     Parameters:
-        synchronize (bool): Whether to synchronize all CUDA devices before timing ends
+        synchronize: Whether to synchronize all CUDA devices before timing ends
 
     Returns:
-        list[TimingResult]: A list of ``TimingResult`` objects for all recorded activities.
+        A list of :class:`TimingResult` objects for all recorded activities.
     """
 
     if synchronize:
@@ -1006,12 +1071,12 @@ def timing_end(synchronize=True):
     return results
 
 
-def timing_print(results, indent=""):
+def timing_print(results: List[TimingResult], indent: str = "") -> None:
     """Print timing results.
 
     Parameters:
-        results (list[TimingResult]): List of ``TimingResult`` objects.
-        indent (str): Optional indentation for the output.
+        results: List of :class:`TimingResult` objects to print.
+        indent: Optional indentation to prepend to all output lines.
     """
 
     if not results:

{warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: warp-lang
-Version: 1.6.2
+Version: 1.7.0
 Summary: A Python framework for high-performance simulation and graphics programming
 Author-email: NVIDIA Corporation <warp-python@nvidia.com>
 License: Apache-2.0
@@ -32,6 +32,7 @@ Requires-Dist: usd-core; extra == "extras"
 Requires-Dist: matplotlib; extra == "extras"
 Requires-Dist: pillow; extra == "extras"
 Requires-Dist: pyglet; extra == "extras"
+Dynamic: license-file
 
 [![PyPI version](https://badge.fury.io/py/warp-lang.svg)](https://badge.fury.io/py/warp-lang)
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
@@ -81,9 +82,9 @@ the `pip install` command, e.g.
 
 | Platform        | Install Command                                                                                                               |
 | --------------- | ----------------------------------------------------------------------------------------------------------------------------- |
-| Linux aarch64   | `pip install https://github.com/NVIDIA/warp/releases/download/v1.6.2/warp_lang-1.6.2+cu11-py3-none-manylinux2014_aarch64.whl` |
-| Linux x86-64    | `pip install https://github.com/NVIDIA/warp/releases/download/v1.6.2/warp_lang-1.6.2+cu11-py3-none-manylinux2014_x86_64.whl`  |
-| Windows x86-64  | `pip install https://github.com/NVIDIA/warp/releases/download/v1.6.2/warp_lang-1.6.2+cu11-py3-none-win_amd64.whl`             |
+| Linux aarch64   | `pip install https://github.com/NVIDIA/warp/releases/download/v1.7.0/warp_lang-1.7.0+cu11-py3-none-manylinux2014_aarch64.whl` |
+| Linux x86-64    | `pip install https://github.com/NVIDIA/warp/releases/download/v1.7.0/warp_lang-1.7.0+cu11-py3-none-manylinux2014_x86_64.whl`  |
+| Windows x86-64  | `pip install https://github.com/NVIDIA/warp/releases/download/v1.7.0/warp_lang-1.7.0+cu11-py3-none-win_amd64.whl`             |
 
 The `--force-reinstall` option may need to be used to overwrite a previous installation.
 
@@ -99,6 +100,15 @@ pip install -U --pre warp-lang --extra-index-url=https://pypi.nvidia.com/
 
 Note that the nightly builds are built with the CUDA 12 runtime and are not published for macOS.
 
+If you plan to install nightly builds regularly, you can simplify future installations by adding NVIDIA's package
+repository as an extra index via the `PIP_EXTRA_INDEX_URL` environment variable. For example:
+
+```text
+export PIP_EXTRA_INDEX_URL="https://pypi.nvidia.com"
+```
+
+This ensures the index is automatically used for `pip` commands, avoiding the need to specify it explicitly.
+
 ### CUDA Requirements
 
 * Warp packages built with CUDA Toolkit 11.x require NVIDIA driver 470 or newer.
@@ -250,16 +260,16 @@ python -m warp.tests
             <td align="center">raymarch</td>
         </tr>
         <tr>
+            <td><a href="https://github.com/NVIDIA/warp/tree/main/warp/examples/core/example_sample_mesh.py"><img src="https://media.githubusercontent.com/media/NVIDIA/warp/refs/heads/main/docs/img/examples/core_sample_mesh.png"></a></td>
             <td><a href="https://github.com/NVIDIA/warp/tree/main/warp/examples/core/example_sph.py"><img src="https://media.githubusercontent.com/media/NVIDIA/warp/refs/heads/main/docs/img/examples/core_sph.png"></a></td>
             <td><a href="https://github.com/NVIDIA/warp/tree/main/warp/examples/core/example_torch.py"><img src="https://media.githubusercontent.com/media/NVIDIA/warp/refs/heads/main/docs/img/examples/core_torch.png"></a></td>
             <td><a href="https://github.com/NVIDIA/warp/tree/main/warp/examples/core/example_wave.py"><img src="https://media.githubusercontent.com/media/NVIDIA/warp/refs/heads/main/docs/img/examples/core_wave.png"></a></td>
-            <td></td>
         </tr>
         <tr>
+            <td align="center">sample mesh</td>
             <td align="center">sph</td>
             <td align="center">torch</td>
             <td align="center">wave</td>
-            <td align="center"></td>
         </tr>
     </tbody>
 </table>
@@ -323,6 +333,18 @@ python -m warp.tests
             <td align="center">trajectory</td>
             <td align="center">soft body properties</td>
         </tr>
+        <tr>
+            <td><a href="https://github.com/NVIDIA/warp/tree/main/warp/examples/optim/example_fluid_checkpoint.py"><img src="https://media.githubusercontent.com/media/NVIDIA/warp/refs/heads/main/docs/img/examples/optim_fluid_checkpoint.png"></a></td>
+            <td></td>
+            <td></td>
+            <td></td>
+        </tr>
+        <tr>
+            <td align="center">fluid checkpoint</td>
+            <td align="center"></td>
+            <td align="center"></td>
+            <td align="center"></td>
+        </tr>
     </tbody>
 </table>