warp-lang 1.6.2__py3-none-win_amd64.whl → 1.7.1__py3-none-win_amd64.whl

warp/stubs.py

@@ -96,7 +96,12 @@ from warp.context import Stream, get_stream, set_stream, wait_stream, synchroniz
 from warp.context import Event, record_event, wait_event, synchronize_event, get_event_elapsed_time
 from warp.context import RegisteredGLBuffer
 from warp.context import is_mempool_supported, is_mempool_enabled, set_mempool_enabled
-from warp.context import set_mempool_release_threshold, get_mempool_release_threshold
+from warp.context import (
+    set_mempool_release_threshold,
+    get_mempool_release_threshold,
+    get_mempool_used_mem_current,
+    get_mempool_used_mem_high,
+)
 from warp.context import is_mempool_access_supported, is_mempool_access_enabled, set_mempool_access_enabled
 from warp.context import is_peer_access_supported, is_peer_access_enabled, set_peer_access_enabled
 
@@ -132,6 +137,7 @@ from warp.paddle import device_from_paddle, device_to_paddle
 from warp.paddle import stream_from_paddle
 
 from warp.build import clear_kernel_cache
+from warp.build import clear_lto_cache
 
 from warp.constants import *
 
@@ -648,6 +654,18 @@ def matrix(*args: Scalar, shape: Tuple[int, int], dtype: Scalar) -> Matrix[Any,
     ...
 
 
+@over
+def matrix_from_cols(*args: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
+    """Construct a matrix from column vectors."""
+    ...
+
+
+@over
+def matrix_from_rows(*args: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
+    """Construct a matrix from row vectors."""
+    ...
+
+
 @over
 def identity(n: int32, dtype: Scalar) -> Matrix[Any, Any, Scalar]:
     """Create an identity matrix with shape=(n,n) with the type given by ``dtype``."""
@@ -662,6 +680,14 @@ def svd3(A: Matrix[3, 3, Float], U: Matrix[3, 3, Float], sigma: Vector[3, Float]
     ...
 
 
+@over
+def svd2(A: Matrix[2, 2, Float], U: Matrix[2, 2, Float], sigma: Vector[2, Float], V: Matrix[2, 2, Scalar]):
+    """Compute the SVD of a 2x2 matrix ``A``. The singular values are returned in ``sigma``,
+    while the left and right basis vectors are returned in ``U`` and ``V``.
+    """
+    ...
+
+
 @over
 def qr3(A: Matrix[3, 3, Float], Q: Matrix[3, 3, Float], R: Matrix[3, 3, Float]):
     """Compute the QR decomposition of a 3x3 matrix ``A``. The orthogonal matrix is returned in ``Q``,
@@ -687,7 +713,7 @@ def quaternion(dtype: Float) -> Quaternion[Float]:
 
 
 @over
-def quaternion(x: Float, y: Float, z: Float, w: Float) -> Quaternion[Float]:
+def quaternion(x: Float, y: Float, z: Float, w: Float, dtype: Scalar) -> Quaternion[Float]:
     """Create a quaternion using the supplied components (type inferred from component type)."""
     ...
 
@@ -724,7 +750,19 @@ def quat_to_axis_angle(quat: Quaternion[Float], axis: Vector[3, Float], angle: F
 
 @over
 def quat_from_matrix(mat: Matrix[3, 3, Float]) -> Quaternion[Float]:
-    """Construct a quaternion from a 3x3 matrix."""
+    """Construct a quaternion from a 3x3 matrix.
+
+    If the matrix is not a pure rotation, but for example includes scaling or skewing, the result is undefined.
+    """
+    ...
+
+
+@over
+def quat_from_matrix(mat: Matrix[4, 4, Float]) -> Quaternion[Float]:
+    """Construct a quaternion from a 4x4 matrix.
+
+    If the top-left 3x3 block of the matrix is not a pure rotation, but for example includes scaling or skewing, the result is undefined.
+    """
     ...
 
 
@@ -1028,7 +1066,7 @@ def tile(x: Any) -> Tile:
 
     This function converts values computed using scalar kernel code to a tile representation for input into collective operations.
 
-    * If the input value is a scalar, then the resulting tile has ``shape=(1, block_dim)``
+    * If the input value is a scalar, then the resulting tile has ``shape=(block_dim,)``
     * If the input value is a vector, then the resulting tile has ``shape=(length(vector), block_dim)``
 
     :param x: A per-thread local value, e.g. scalar, vector, or matrix.
@@ -1121,13 +1159,12 @@ def tile_transpose(a: Tile) -> Tile:
 def tile_broadcast(a: Tile, shape: Tuple[int, ...]) -> Tile:
     """Broadcast a tile.
 
-    This function will attempt to broadcast the input tile ``a`` to the destination shape (m, n).
-
+    Broadcasts the input tile ``a`` to the destination shape.
     Broadcasting follows NumPy broadcast rules.
 
     :param a: Tile to broadcast
     :param shape: The shape to broadcast to
-    :returns: Tile with broadcast ``shape=(m, n)``
+    :returns: Tile with broadcast shape
     """
     ...
 
@@ -1810,6 +1847,18 @@ def randi(state: uint32, low: int32, high: int32) -> int:
     ...
 
 
+@over
+def randu(state: uint32) -> uint32:
+    """Return a random unsigned integer in the range [0, 2^32)."""
+    ...
+
+
+@over
+def randu(state: uint32, low: uint32, high: uint32) -> uint32:
+    """Return a random unsigned integer between [low, high)."""
+    ...
+
+
 @over
 def randf(state: uint32) -> float:
     """Return a random float between [0.0, 1.0)."""
@@ -2029,61 +2078,171 @@ def tid() -> Tuple[int, int, int, int]:
 
 @over
 def select(cond: bool, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: int8, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: uint8, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: int16, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: uint16, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: int32, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: uint32, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: int64, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(cond: uint64, value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(cond, value_if_true, value_if_false)``.
+    """
     ...
 
 
 @over
 def select(arr: Array[Any], value_if_false: Any, value_if_true: Any) -> Any:
-    """Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``"""
+    """Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``.
+
+    .. deprecated:: 1.7
+         Use :func:`where` instead, which has the more intuitive argument order:
+         ``where(arr, value_if_true, value_if_false)``.
+    """
+    ...
+
+
+@over
+def where(cond: bool, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: int8, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: uint8, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: int16, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: uint16, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: int32, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: uint32, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: int64, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(cond: uint64, value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
+    ...
+
+
+@over
+def where(arr: Array[Any], value_if_true: Any, value_if_false: Any) -> Any:
+    """Select between two arguments, if ``arr`` is not null then return ``value_if_true``, otherwise return ``value_if_false``."""
     ...
 
 
@@ -2492,7 +2651,19 @@ def expect_near(a: Float, b: Float, tolerance: Float):
 
 
 @over
-def expect_near(a: vec3f, b: vec3f, tolerance: float32):
+def expect_near(a: Vector[Any, Float], b: Vector[Any, Float], tolerance: Float):
+    """Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
+    ...
+
+
+@over
+def expect_near(a: Quaternion[Float], b: Quaternion[Float], tolerance: Float):
+    """Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
+    ...
+
+
+@over
+def expect_near(a: Matrix[Any, Any, Float], b: Matrix[Any, Any, Float], tolerance: Float):
     """Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
     ...
 
@@ -2901,7 +3072,7 @@ def tile_matmul(a: Tile, b: Tile, out: Tile) -> Tile:
         * fp16, fp32, fp64 (real)
         * vec2h, vec2f, vec2d (complex)
 
-    All input and output tiles must have the same datatype. Tile data will be automatically be migrated
+    All input and output tiles must have the same datatype. Tile data will automatically be migrated
     to shared memory if necessary and will use TensorCore operations when available.
 
     :param a: A tile with ``shape=(M, K)``
@@ -2920,7 +3091,7 @@ def tile_matmul(a: Tile, b: Tile) -> Tile:
         * fp16, fp32, fp64 (real)
         * vec2h, vec2f, vec2d (complex)
 
-    Both input tiles must have the same datatype. Tile data will be automatically be migrated
+    Both input tiles must have the same datatype. Tile data will automatically be migrated
     to shared memory if necessary and will use TensorCore operations when available.
 
     :param a: A tile with ``shape=(M, K)``
@@ -3134,3 +3305,29 @@ def smooth_normalize(v: Any, delta: float):
         Vector[Any,Float]: The normalized vector.
     """
     ...
+
+
+@over
+def transform_from_matrix(mat: Matrix[4, 4, float32]) -> Transformation[float32]:
+    """Construct a transformation from a 4x4 matrix.
+
+    Args:
+        mat (Matrix[4, 4, Float]): Matrix to convert.
+
+    Returns:
+        Transformation[Float]: The transformation.
+    """
+    ...
+
+
+@over
+def transform_to_matrix(xform: Transformation[float32]) -> Matrix[4, 4, float32]:
+    """Convert a transformation to a 4x4 matrix.
+
+    Args:
+        xform (Transformation[Float]): Transformation to convert.
+
+    Returns:
+        Matrix[4, 4, Float]: The matrix.
+    """
+    ...

warp/tests/__main__.py

@@ -1,18 +1,3 @@
-# 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 warp.thirdparty.unittest_parallel import main
 
 if __name__ == "__main__":

warp/tests/assets/torus.usda

@@ -59,9 +59,9 @@
     }
     defaultPrim = "World"
     endTimeCode = 100
+    framesPerSecond = 24
     metersPerUnit = 0.01
     startTimeCode = 0
-    timeCodesPerSecond = 24
     upAxis = "Y"
 )
 

warp/tests/{test_mempool.py → cuda/test_mempool.py}

@@ -71,6 +71,44 @@ def test_mempool_release_threshold(test, device):
     test.assertEqual(wp.get_mempool_release_threshold(device), saved_threshold)
 
 
+def test_mempool_usage_queries(test, device):
+    """Check API to query mempool memory usage."""
+
+    device = wp.get_device(device)
+    pre_alloc_mempool_usage_curr = wp.get_mempool_used_mem_current(device)
+    pre_alloc_mempool_usage_high = wp.get_mempool_used_mem_high(device)
+
+    # Allocate a 1 MiB array
+    test_data = wp.empty(262144, dtype=wp.float32, device=device)
+    wp.synchronize_device(device)
+
+    # Query memory usage again
+    post_alloc_mempool_usage_curr = wp.get_mempool_used_mem_current(device)
+    post_alloc_mempool_usage_high = wp.get_mempool_used_mem_high(device)
+
+    test.assertEqual(
+        post_alloc_mempool_usage_curr, pre_alloc_mempool_usage_curr + 1048576, "Memory usage did not increase by 1 MiB"
+    )
+    test.assertGreaterEqual(post_alloc_mempool_usage_high, 1048576, "High-water mark is not at least 1 MiB")
+
+    # Free the allocation
+    del test_data
+    wp.synchronize_device(device)
+
+    # Query memory usage
+    post_free_mempool_usage_curr = wp.get_mempool_used_mem_current(device)
+    post_free_mempool_usage_high = wp.get_mempool_used_mem_high(device)
+
+    test.assertEqual(
+        post_free_mempool_usage_curr,
+        pre_alloc_mempool_usage_curr,
+        "Test didn't end with the same amount of used memory as the test started with.",
+    )
+    test.assertEqual(
+        post_free_mempool_usage_high, post_alloc_mempool_usage_high, "High-water mark should not change after free"
+    )
+
+
 def test_mempool_exceptions(test, device):
     device = wp.get_device(device)
 
@@ -176,6 +214,7 @@ devices_without_mempools = [d for d in get_test_devices() if not d.is_mempool_su
 add_function_test(
     TestMempool, "test_mempool_release_threshold", test_mempool_release_threshold, devices=devices_with_mempools
 )
+add_function_test(TestMempool, "test_mempool_usage_queries", test_mempool_usage_queries, devices=devices_with_mempools)
 add_function_test(TestMempool, "test_mempool_access_self", test_mempool_access_self, devices=devices_with_mempools)
 
 # test devices without mempool support

warp/tests/{test_streams.py → cuda/test_streams.py}

@@ -342,6 +342,29 @@ def test_event_elapsed_time(test, device):
     test.assertGreater(elapsed, 0)
 
 
+def test_event_elapsed_time_graph(test, device):
+    stream = wp.get_stream(device)
+    e1 = wp.Event(device, enable_timing=True)
+    e2 = wp.Event(device, enable_timing=True)
+
+    a = wp.zeros(N, dtype=float, device=device)
+
+    wp.load_module(device=device)
+
+    with wp.ScopedCapture(device, force_module_load=False) as capture:
+        stream.record_event(e1)
+        wp.launch(inc, dim=N, inputs=[a], device=device)
+        stream.record_event(e2)
+
+    wp.capture_launch(capture.graph)
+
+    wp.synchronize_device(device)
+
+    elapsed = wp.get_event_elapsed_time(e1, e2)
+
+    test.assertGreater(elapsed, 0)
+
+
 def test_stream_priority_basics(test, device):
     standard_stream = wp.Stream(device)
     test.assertEqual(standard_stream.priority, 0, "Default priority of streams must be 0.")
@@ -401,6 +424,52 @@ def test_stream_priority_timings(test, device):
     test.assertLess(elapsed_hi, elapsed_lo, "Copies on higher-priority stream should be faster.")
 
 
+@wp.kernel
+def sum_threads(sum: wp.array(dtype=wp.uint64)):
+    i = wp.tid()
+    wp.atomic_add(sum, 0, wp.uint64(1))
+
+
+def test_stream_event_is_complete(test, device):
+    with wp.ScopedDevice(device):
+        stream = wp.Stream()
+        event = wp.Event()
+        # No operations on stream, should be complete
+        test.assertTrue(stream.is_complete)
+
+        # Event not recorded yet, should be complete
+        test.assertTrue(event.is_complete)
+
+        a = wp.zeros(1, dtype=wp.uint64)
+
+        threads = 1024 * 1024 * 8
+
+        with wp.ScopedStream(stream):
+            # Launch some work on the stream and reuse the event
+
+            for iter in range(5):
+                # Kernel takes about 1 ms to run on an RTX 3090
+                wp.launch(sum_threads, dim=threads, outputs=[a])
+
+                stream.record_event(event)
+
+                # Kernel should still be running
+                test.assertFalse(stream.is_complete)
+
+                # Event should not be finished
+                test.assertFalse(event.is_complete)
+
+                # Force the stream operations to complete
+                wp.synchronize_stream(stream)
+
+                # Now all operations are complete
+                test.assertTrue(stream.is_complete)
+                test.assertTrue(event.is_complete)
+
+                # Verify result
+                test.assertEqual(a.numpy()[0], (iter + 1) * threads)
+
+
 devices = get_selected_cuda_test_devices()
 
 
@@ -554,9 +623,11 @@ add_function_test(TestStreams, "test_stream_scope_wait_event", test_stream_scope
 add_function_test(TestStreams, "test_stream_scope_wait_stream", test_stream_scope_wait_stream, devices=devices)
 add_function_test(TestStreams, "test_stream_priority_basics", test_stream_priority_basics, devices=devices)
 add_function_test(TestStreams, "test_stream_priority_timings", test_stream_priority_timings, devices=devices)
+add_function_test(TestStreams, "test_stream_event_is_complete", test_stream_event_is_complete, devices=devices)
 
 add_function_test(TestStreams, "test_event_synchronize", test_event_synchronize, devices=devices)
 add_function_test(TestStreams, "test_event_elapsed_time", test_event_elapsed_time, devices=devices)
+add_function_test(TestStreams, "test_event_elapsed_time_graph", test_event_elapsed_time_graph, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()