warp-lang 1.7.1__py3-none-macosx_10_13_universal2.whl → 1.7.2rc1__py3-none-macosx_10_13_universal2.whl

warp/stubs.py

@@ -2248,145 +2248,193 @@ def where(arr: Array[Any], value_if_true: Any, value_if_false: Any) -> Any:
 
 @over
 def atomic_add(arr: Array[Any], i: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: FabricArray[Any], i: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
+    """
     ...
 
 
 @over
 def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically add ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: Array[Any], i: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: FabricArray[Any], i: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
+    This function is automatically invoked when using the syntax ``arr[i] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
+    This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
+    """
     ...
 
 
 @over
 def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
-    """Atomically subtract ``value`` onto ``arr[i,j,k,l]`` and return the old value."""
+    """Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
+    This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
+    """
     ...
 
 

warp/tests/cuda/test_streams.py

@@ -442,7 +442,7 @@ def test_stream_event_is_complete(test, device):
 
         a = wp.zeros(1, dtype=wp.uint64)
 
-        threads = 1024 * 1024 * 8
+        threads = 1024 * 1024 * 64
 
         with wp.ScopedStream(stream):
             # Launch some work on the stream and reuse the event

warp/tests/sim/test_model.py

@@ -150,7 +150,9 @@ class TestModel(unittest.TestCase):
         assert builder.joint_count == 8
         assert builder.body_count == 8
         assert builder.articulation_count == 3
-        add_three_cubes(builder, parent_body=b4)
+        lb = add_three_cubes(builder, parent_body=b4)
+        # add shape with no mass (e.g. visual only shape)
+        builder.add_shape_sphere(lb, density=0.0)
 
         builder.collapse_fixed_joints()
 
@@ -158,8 +160,8 @@ class TestModel(unittest.TestCase):
         assert builder.articulation_count == 2
         assert builder.articulation_start == [0, 1]
         assert builder.joint_type == [wp.sim.JOINT_REVOLUTE, wp.sim.JOINT_FREE]
-        assert builder.shape_count == 11
-        assert builder.shape_body == [-1, -1, -1, -1, -1, -1, 0, 1, 1, 1, 1]
+        assert builder.shape_count == 12
+        assert builder.shape_body == [-1, -1, -1, -1, -1, -1, 0, 1, 1, 1, 1, 1]
         assert builder.body_count == 2
         assert builder.body_com[0] == wp.vec3(1.0, 2.0, 3.0)
         assert builder.body_com[1] == wp.vec3(0.25, 0.25, 0.25)

warp/tests/sim/test_sim_grad.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import platform
 import unittest
 
 import numpy as np
@@ -104,9 +103,6 @@ def test_sphere_pushing_on_rails(
     static_contacts=True,
     print_grad=False,
 ):
-    if platform.system() == "Darwin":
-        test.skipTest("Crashes on Mac runners")
-
     # Two spheres on a rail (prismatic or D6 joint), one is pushed, the other is passive.
     # The absolute distance to a target is measured and gradients are compared for
     # a push that is too far and too close.
@@ -161,9 +157,6 @@ def test_sphere_pushing_on_rails(
     model.joint_attach_ke = 32000.0 * 16
     model.joint_attach_kd = 500.0 * 4
 
-    model.shape_geo.scale.requires_grad = False
-    model.shape_geo.thickness.requires_grad = False
-
     if static_contacts:
         wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, model)
         model.rigid_contact_margin = 10.0
@@ -272,7 +265,7 @@ def test_sphere_pushing_on_rails(
     gradcheck(rollout, [action_too_close], device=device, eps=0.2, tol=tol, print_grad=print_grad)
 
 
-devices = get_test_devices()
+devices = get_test_devices(mode="basic")
 
 
 class TestSimGradients(unittest.TestCase):

warp/tests/test_array.py

@@ -2803,7 +2803,7 @@ def test_alloc_strides(test, device):
 
 def test_casting(test, device):
     idxs = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12)
-    idxs = wp.array(idxs, device=device).reshape((-1, 3))
+    idxs = wp.array(idxs, device=device, dtype=wp.int32).reshape((-1, 3))
     idxs = wp.array(idxs, shape=idxs.shape[0], dtype=wp.vec3i, device=device)
     assert idxs.dtype is wp.vec3i
     assert idxs.shape == (4,)
@@ -2846,25 +2846,25 @@ def test_array_len(test, device):
 
 def test_cuda_interface_conversion(test, device):
     class MyArrayInterface:
-        def __init__(self, data):
-            self.data = np.array(data)
+        def __init__(self, data, npdtype):
+            self.data = np.array(data, dtype=npdtype)
             self.__array_interface__ = self.data.__array_interface__
             self.__cuda_array_interface__ = self.data.__array_interface__
             self.__len__ = self.data.__len__
 
-    array = MyArrayInterface((1, 2, 3))
+    array = MyArrayInterface((1, 2, 3), np.int8)
     wp_array = wp.array(array, dtype=wp.int8, device=device)
     assert wp_array.ptr != 0
 
-    array = MyArrayInterface((1, 2, 3))
+    array = MyArrayInterface((1, 2, 3), np.float32)
     wp_array = wp.array(array, dtype=wp.float32, device=device)
     assert wp_array.ptr != 0
 
-    array = MyArrayInterface((1, 2, 3))
+    array = MyArrayInterface((1, 2, 3), np.float32)
     wp_array = wp.array(array, dtype=wp.vec3, device=device)
     assert wp_array.ptr != 0
 
-    array = MyArrayInterface((1, 2, 3, 4))
+    array = MyArrayInterface((1, 2, 3, 4), np.float32)
     wp_array = wp.array(array, dtype=wp.mat22, device=device)
     assert wp_array.ptr != 0
 
@@ -2883,6 +2883,7 @@ add_function_test(TestArray, "test_shape", test_shape, devices=devices)
 add_function_test(TestArray, "test_negative_shape", test_negative_shape, devices=devices)
 add_function_test(TestArray, "test_flatten", test_flatten, devices=devices)
 add_function_test(TestArray, "test_reshape", test_reshape, devices=devices)
+
 add_function_test(TestArray, "test_slicing", test_slicing, devices=devices)
 add_function_test(TestArray, "test_transpose", test_transpose, devices=devices)
 add_function_test(TestArray, "test_view", test_view, devices=devices)

warp/tests/test_atomic.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import re
 import unittest
 
 import numpy as np
@@ -32,8 +33,8 @@ def make_atomic_test(type):
         tid = wp.tid()
 
         wp.atomic_add(out_add, 0, val[tid])
-        wp.atomic_min(out_min, 0, val[tid])
-        wp.atomic_max(out_max, 0, val[tid])
+        wp.atomic_min(out_min, wp.uint32(0), val[tid])
+        wp.atomic_max(out_max, wp.int64(0), val[tid])
 
     # register a custom kernel (no decorator) function
     # this lets us register the same function definition
@@ -130,6 +131,102 @@ test_atomic_mat33 = make_atomic_test(wp.mat33)
 test_atomic_mat44 = make_atomic_test(wp.mat44)
 
 
+def test_atomic_add_supported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_add(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
+def test_atomic_min_supported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_min(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
+def test_atomic_max_supported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_max(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
+def test_atomic_add_unsupported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    dtype_str = re.escape(wp.types.type_repr(dtype))
+    scalar_type_str = wp.types.type_repr(scalar_type)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_add(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        (
+            r"atomic_add\(\) operations only work on arrays with \[u\]int32, \[u\]int64, float16, float32, or float64 "
+            rf"as the underlying scalar types, but got {dtype_str} \(with scalar type {scalar_type_str}\)$"
+        ),
+    ):
+        wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
+def test_atomic_min_unsupported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    dtype_str = re.escape(wp.types.type_repr(dtype))
+    scalar_type_str = wp.types.type_repr(scalar_type)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_min(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        (
+            r"atomic_min\(\) operations only work on arrays with \[u\]int32, \[u\]int64, float32, or float64 "
+            rf"as the underlying scalar types, but got {dtype_str} \(with scalar type {scalar_type_str}\)$"
+        ),
+    ):
+        wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
+def test_atomic_max_unsupported_dtypes(test, device, dtype):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    dtype_str = re.escape(wp.types.type_repr(dtype))
+    scalar_type_str = wp.types.type_repr(scalar_type)
+
+    @wp.kernel
+    def kernel(arr: wp.array(dtype=dtype)):
+        wp.atomic_max(arr, 0, dtype(scalar_type(0)))
+
+    arr = wp.zeros(1, dtype=dtype, device=device)
+    with test.assertRaisesRegex(
+        RuntimeError,
+        (
+            r"atomic_max\(\) operations only work on arrays with \[u\]int32, \[u\]int64, float32, or float64 "
+            rf"as the underlying scalar types, but got {dtype_str} \(with scalar type {scalar_type_str}\)$"
+        ),
+    ):
+        wp.launch(kernel, dim=1, outputs=(arr,), device=device)
+
+
 devices = get_test_devices()
 
 
@@ -147,6 +244,88 @@ add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=de
 add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
 add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
 
+for dtype in (
+    wp.int32,
+    wp.uint32,
+    wp.int64,
+    wp.uint64,
+    wp.float16,
+    wp.float32,
+    wp.float64,
+    wp.vec3i,
+    wp.vec3ui,
+    wp.vec3l,
+    wp.vec3ul,
+    wp.vec3h,
+    wp.vec3f,
+    wp.vec3d,
+):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    add_function_test(
+        TestAtomic,
+        f"test_atomic_add_supported_dtypes_{dtype.__name__}",
+        test_atomic_add_supported_dtypes,
+        devices=devices,
+        dtype=dtype,
+    )
+
+    if scalar_type is not wp.float16:
+        add_function_test(
+            TestAtomic,
+            f"test_atomic_min_supported_dtypes_{dtype.__name__}",
+            test_atomic_min_supported_dtypes,
+            devices=devices,
+            dtype=dtype,
+        )
+        add_function_test(
+            TestAtomic,
+            f"test_atomic_max_supported_dtypes_{dtype.__name__}",
+            test_atomic_max_supported_dtypes,
+            devices=devices,
+            dtype=dtype,
+        )
+
+
+for dtype in (
+    wp.int8,
+    wp.uint8,
+    wp.int16,
+    wp.uint16,
+    wp.float16,
+    wp.vec3b,
+    wp.vec3ub,
+    wp.vec3s,
+    wp.vec3us,
+    wp.vec3h,
+):
+    scalar_type = getattr(dtype, "_wp_scalar_type_", dtype)
+
+    if scalar_type is not wp.float16:
+        add_function_test(
+            TestAtomic,
+            f"test_atomic_add_unsupported_dtypes_{dtype.__name__}",
+            test_atomic_add_unsupported_dtypes,
+            devices=devices,
+            dtype=dtype,
+        )
+
+    add_function_test(
+        TestAtomic,
+        f"test_atomic_min_unsupported_dtypes_{dtype.__name__}",
+        test_atomic_min_unsupported_dtypes,
+        devices=devices,
+        dtype=dtype,
+    )
+
+    add_function_test(
+        TestAtomic,
+        f"test_atomic_max_unsupported_dtypes_{dtype.__name__}",
+        test_atomic_max_unsupported_dtypes,
+        devices=devices,
+        dtype=dtype,
+    )
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()