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

warp/utils.py

@@ -1,9 +1,19 @@
-# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
 
 import cProfile
 import ctypes
@@ -11,7 +21,7 @@ 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
 
@@ -133,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:
@@ -140,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")
 
@@ -665,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
@@ -791,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
 
@@ -805,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
@@ -820,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
@@ -834,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
@@ -860,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.
 
@@ -919,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:
@@ -952,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:
@@ -998,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.1.dist-info → warp_lang-1.7.0.dist-info}/METADATA

@@ -1,9 +1,9 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: warp-lang
-Version: 1.6.1
+Version: 1.7.0
 Summary: A Python framework for high-performance simulation and graphics programming
-Author-email: NVIDIA Corporation <mmacklin@nvidia.com>
-License: NVIDIA Software License
+Author-email: NVIDIA Corporation <warp-python@nvidia.com>
+License: Apache-2.0
 Project-URL: GitHub, https://github.com/NVIDIA/warp
 Project-URL: Documentation, https://nvidia.github.io/warp
 Project-URL: Changelog, https://github.com/NVIDIA/warp/blob/main/CHANGELOG.md
@@ -12,8 +12,9 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: License :: Other/Proprietary License
+Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
@@ -23,15 +24,18 @@ Provides-Extra: dev
 Requires-Dist: pre-commit; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: nvtx; extra == "dev"
-Requires-Dist: nvidia-sphinx-theme; extra == "dev"
+Requires-Dist: nvidia-sphinx-theme; python_version >= "3.9" and extra == "dev"
 Requires-Dist: sphinx-copybutton; extra == "dev"
 Requires-Dist: coverage[toml]; extra == "dev"
 Provides-Extra: extras
 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)
 ![GitHub commit activity](https://img.shields.io/github/commit-activity/m/NVIDIA/warp?link=https%3A%2F%2Fgithub.com%2FNVIDIA%2Fwarp%2Fcommits%2Fmain)
 [![Downloads](https://static.pepy.tech/badge/warp-lang/month)](https://pepy.tech/project/warp-lang)
 [![codecov](https://codecov.io/github/NVIDIA/warp/graph/badge.svg?token=7O1KSM79FG)](https://codecov.io/github/NVIDIA/warp)
@@ -78,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.1/warp_lang-1.6.1+cu11-py3-none-manylinux2014_aarch64.whl` |
-| Linux x86-64    | `pip install https://github.com/NVIDIA/warp/releases/download/v1.6.1/warp_lang-1.6.1+cu11-py3-none-manylinux2014_x86_64.whl`  |
-| Windows x86-64  | `pip install https://github.com/NVIDIA/warp/releases/download/v1.6.1/warp_lang-1.6.1+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.
 
@@ -96,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.
@@ -247,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>
@@ -320,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>
 
@@ -449,6 +474,8 @@ Problems, questions, and feature requests can be opened on [GitHub Issues](https
 
 The Warp team also monitors the **#warp** forum on the public [Omniverse Discord](https://discord.com/invite/nvidiaomniverse) server, come chat with us!
 
+For inquiries not suited for GitHub Issues or Discord, please email warp-python@nvidia.com.
+
 ## Versioning
 
 Versions take the format X.Y.Z, similar to [Python itself](https://devguide.python.org/developer-workflow/development-cycle/#devcycle):
@@ -464,7 +491,7 @@ Note that prior to 0.11.0, this schema was not strictly adhered to.
 
 ## License
 
-Warp is provided under the NVIDIA Software License, please see [LICENSE.md](./LICENSE.md) for full license text.
+Warp is provided under the Apache License, Version 2.0. Please see [LICENSE.md](./LICENSE.md) for full license text.
 
 ## Contributing