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

warp/tests/test_atomic_cas.py

@@ -0,0 +1,299 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+def create_spinlock_test(dtype):
+    @wp.func
+    def spinlock_acquire(lock: wp.array(dtype=dtype)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, dtype(0), dtype(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release(lock: wp.array(dtype=dtype)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, dtype(0))
+
+    @wp.func
+    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
+        value = wp.atomic_exch(ptr, index, dtype(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    @wp.kernel
+    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype)):
+        # Try to acquire the lock
+        spinlock_acquire(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read(counter, 0)
+        counter[0] = value + dtype(1)
+
+        # Release the lock
+        spinlock_release(lock)
+
+    return test_spinlock_counter
+
+
+def test_atomic_cas(test, device, warp_type, numpy_type):
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.array([0], dtype=warp_type, device=device)
+
+    test_spinlock_counter = create_spinlock_test(warp_type)
+    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+
+    # Verify counter reached n
+    counter_np = counter.numpy()
+    expected = np.array([n], dtype=numpy_type)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+def create_spinlock_test_2d(dtype):
+    @wp.func
+    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=2)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, dtype(0), dtype(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release(lock: wp.array(dtype=dtype, ndim=2)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, dtype(0))
+
+    @wp.func
+    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
+        value = wp.atomic_exch(ptr, index, dtype(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    @wp.kernel
+    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=2)):
+        # Try to acquire the lock
+        spinlock_acquire(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read(counter, 0)
+        counter[0] = value + dtype(1)
+
+        # Release the lock
+        spinlock_release(lock)
+
+    return test_spinlock_counter
+
+
+def test_atomic_cas_2d(test, device, warp_type, numpy_type):
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1), dtype=warp_type, device=device)
+
+    test_spinlock_counter = create_spinlock_test_2d(warp_type)
+    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+
+    # Verify counter reached n
+    counter_np = counter.numpy()
+    expected = np.array([n], dtype=numpy_type)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+def create_spinlock_test_3d(dtype):
+    @wp.func
+    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=3)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, 0, dtype(0), dtype(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release(lock: wp.array(dtype=dtype, ndim=3)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, 0, dtype(0))
+
+    @wp.func
+    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
+        value = wp.atomic_exch(ptr, index, dtype(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    @wp.kernel
+    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=3)):
+        # Try to acquire the lock
+        spinlock_acquire(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read(counter, 0)
+        counter[0] = value + dtype(1)
+
+        # Release the lock
+        spinlock_release(lock)
+
+    return test_spinlock_counter
+
+
+def test_atomic_cas_3d(test, device, warp_type, numpy_type):
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1, 1), dtype=warp_type, device=device)
+
+    test_spinlock_counter = create_spinlock_test_3d(warp_type)
+    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+
+    # Verify counter reached n
+    counter_np = counter.numpy()
+    expected = np.array([n], dtype=numpy_type)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+def create_spinlock_test_4d(dtype):
+    @wp.func
+    def spinlock_acquire(lock: wp.array(dtype=dtype, ndim=4)):
+        # Try to acquire the lock by setting it to 1 if it's 0
+        while wp.atomic_cas(lock, 0, 0, 0, 0, dtype(0), dtype(1)) == 1:
+            pass
+
+    @wp.func
+    def spinlock_release(lock: wp.array(dtype=dtype, ndim=4)):
+        # Release the lock by setting it back to 0
+        wp.atomic_exch(lock, 0, 0, 0, 0, dtype(0))
+
+    @wp.func
+    def volatile_read(ptr: wp.array(dtype=dtype), index: int):
+        value = wp.atomic_exch(ptr, index, dtype(0))
+        wp.atomic_exch(ptr, index, value)
+        return value
+
+    @wp.kernel
+    def test_spinlock_counter(counter: wp.array(dtype=dtype), lock: wp.array(dtype=dtype, ndim=4)):
+        # Try to acquire the lock
+        spinlock_acquire(lock)
+
+        # Critical section - increment counter
+        # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+        # Work around since warp arrays cannot be marked as volatile
+        value = volatile_read(counter, 0)
+        counter[0] = value + dtype(1)
+
+        # Release the lock
+        spinlock_release(lock)
+
+    return test_spinlock_counter
+
+
+def test_atomic_cas_4d(test, device, warp_type, numpy_type):
+    n = 100
+    counter = wp.array([0], dtype=warp_type, device=device)
+    lock = wp.zeros(shape=(1, 1, 1, 1), dtype=warp_type, device=device)
+
+    test_spinlock_counter = create_spinlock_test_4d(warp_type)
+    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, lock], device=device)
+
+    # Verify counter reached n
+    counter_np = counter.numpy()
+    expected = np.array([n], dtype=numpy_type)
+
+    if not np.array_equal(counter_np, expected):
+        print(f"Counter mismatch: expected {expected}, got {counter_np}")
+
+    assert_np_equal(counter_np, expected)
+
+
+devices = get_test_devices()
+
+
+class TestAtomicCAS(unittest.TestCase):
+    pass
+
+
+# Test all supported types
+test_types = [
+    (wp.int32, np.int32),
+    (wp.uint32, np.uint32),
+    (wp.int64, np.int64),
+    (wp.uint64, np.uint64),
+    (wp.float32, np.float32),
+    (wp.float64, np.float64),
+]
+
+for warp_type, numpy_type in test_types:
+    type_name = warp_type.__name__
+    add_function_test(
+        TestAtomicCAS,
+        f"test_cas_{type_name}",
+        test_atomic_cas,
+        devices=devices,
+        warp_type=warp_type,
+        numpy_type=numpy_type,
+    )
+
+    # Add 2D test for each type
+    add_function_test(
+        TestAtomicCAS,
+        f"test_cas_2d_{type_name}",
+        test_atomic_cas_2d,
+        devices=devices,
+        warp_type=warp_type,
+        numpy_type=numpy_type,
+    )
+
+    # Add 3D test for each type
+    add_function_test(
+        TestAtomicCAS,
+        f"test_cas_3d_{type_name}",
+        test_atomic_cas_3d,
+        devices=devices,
+        warp_type=warp_type,
+        numpy_type=numpy_type,
+    )
+
+    # Add 4D test for each type
+    add_function_test(
+        TestAtomicCAS,
+        f"test_cas_4d_{type_name}",
+        test_atomic_cas_4d,
+        devices=devices,
+        warp_type=warp_type,
+        numpy_type=numpy_type,
+    )
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_codegen.py

@@ -435,14 +435,8 @@ def test_error_collection_construct(test, device):
         x = [1.0, 2.0, 3.0]
 
     def kernel_2_fn():
-        x = (1.0, 2.0, 3.0)
-
-    def kernel_3_fn():
         x = {"a": 1.0, "b": 2.0, "c": 3.0}
 
-    def kernel_4_fn():
-        wp.length((1.0, 2.0, 3.0))
-
     kernel = wp.Kernel(func=kernel_1_fn)
     with test.assertRaisesRegex(
         RuntimeError,
@@ -451,22 +445,9 @@ def test_error_collection_construct(test, device):
         wp.launch(kernel, dim=1, device=device)
 
     kernel = wp.Kernel(func=kernel_2_fn)
-    with test.assertRaisesRegex(
-        RuntimeError,
-        r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` for small collections instead.",
-    ):
-        wp.launch(kernel, dim=1, device=device)
-
-    kernel = wp.Kernel(func=kernel_3_fn)
     with test.assertRaisesRegex(RuntimeError, r"Construct `ast.Dict` not supported in kernels."):
         wp.launch(kernel, dim=1, device=device)
 
-    kernel = wp.Kernel(func=kernel_4_fn)
-    with test.assertRaisesRegex(
-        RuntimeError, r"Tuple constructs are not supported in kernels. Use vectors like `wp.vec3\(\)` instead."
-    ):
-        wp.launch(kernel, dim=1, device=device)
-
 
 def test_error_unmatched_arguments(test, device):
     def kernel_1_fn():
@@ -693,6 +674,138 @@ def test_codegen_return_in_kernel(test, device):
     test.assertEqual(result.numpy()[0], grid_size - 256)
 
 
+@wp.kernel
+def conditional_ifexp(x: float, result: wp.array(dtype=wp.int32)):
+    wp.atomic_add(result, 0, 1) if x > 0.0 else wp.atomic_add(result, 1, 1)
+
+
+def test_ifexp_only_executes_one_branch(test, device):
+    result = wp.zeros(2, dtype=wp.int32, device=device)
+
+    wp.launch(conditional_ifexp, dim=1, inputs=[1.0, result], device=device)
+
+    values = result.numpy()
+    # Only first branch is taken
+    test.assertEqual(values[0], 1)
+    test.assertEqual(values[1], 0)
+
+
+@wp.kernel
+def test_multiple_return_values_quat_to_axis_angle_kernel(
+    q: wp.quath,
+    expected_axis: wp.vec3h,
+    expected_angle: wp.float16,
+):
+    axis, angle = wp.quat_to_axis_angle(q)
+
+    wp.expect_near(axis[0], expected_axis[0], tolerance=wp.float16(1e-3))
+    wp.expect_near(axis[1], expected_axis[1], tolerance=wp.float16(1e-3))
+    wp.expect_near(axis[2], expected_axis[2], tolerance=wp.float16(1e-3))
+
+    wp.expect_near(angle, expected_angle, tolerance=wp.float16(1e-3))
+
+
+@wp.kernel
+def test_multiple_return_values_svd3_kernel(
+    A: wp.mat33f,
+    expected_U: wp.mat33f,
+    expected_sigma: wp.vec3f,
+    expected_V: wp.mat33f,
+):
+    U, sigma, V = wp.svd3(A)
+
+    wp.expect_near(U[0][0], expected_U[0][0], tolerance=1e-5)
+    wp.expect_near(U[0][1], expected_U[0][1], tolerance=1e-5)
+    wp.expect_near(U[0][2], expected_U[0][2], tolerance=1e-5)
+    wp.expect_near(U[1][0], expected_U[1][0], tolerance=1e-5)
+    wp.expect_near(U[1][1], expected_U[1][1], tolerance=1e-5)
+    wp.expect_near(U[1][2], expected_U[1][2], tolerance=1e-5)
+    wp.expect_near(U[2][0], expected_U[2][0], tolerance=1e-5)
+    wp.expect_near(U[2][1], expected_U[2][1], tolerance=1e-5)
+    wp.expect_near(U[2][2], expected_U[2][2], tolerance=1e-5)
+
+    wp.expect_near(sigma[0], expected_sigma[0], tolerance=1e-5)
+    wp.expect_near(sigma[1], expected_sigma[1], tolerance=1e-5)
+    wp.expect_near(sigma[2], expected_sigma[2], tolerance=1e-5)
+
+    wp.expect_near(V[0][0], expected_V[0][0], tolerance=1e-5)
+    wp.expect_near(V[0][1], expected_V[0][1], tolerance=1e-5)
+    wp.expect_near(V[0][2], expected_V[0][2], tolerance=1e-5)
+    wp.expect_near(V[1][0], expected_V[1][0], tolerance=1e-5)
+    wp.expect_near(V[1][1], expected_V[1][1], tolerance=1e-5)
+    wp.expect_near(V[1][2], expected_V[1][2], tolerance=1e-5)
+    wp.expect_near(V[2][0], expected_V[2][0], tolerance=1e-5)
+    wp.expect_near(V[2][1], expected_V[2][1], tolerance=1e-5)
+    wp.expect_near(V[2][2], expected_V[2][2], tolerance=1e-5)
+
+
+def test_multiple_return_values(test, device):
+    q = wp.quath(1.0, 2.0, 3.0, 4.0)
+    expected_axis = wp.vec3h(0.26726124, 0.53452247, 0.80178368)
+    expected_angle = 1.50408018
+
+    axis, angle = wp.quat_to_axis_angle(q)
+
+    test.assertAlmostEqual(axis[0], expected_axis[0], places=3)
+    test.assertAlmostEqual(axis[1], expected_axis[1], places=3)
+    test.assertAlmostEqual(axis[2], expected_axis[2], places=3)
+
+    test.assertAlmostEqual(angle, expected_angle, places=3)
+
+    wp.launch(
+        test_multiple_return_values_quat_to_axis_angle_kernel,
+        dim=1,
+        inputs=(q, expected_axis, expected_angle),
+    )
+
+    # fmt: off
+    A = wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)
+    expected_U = wp.mat33(
+        0.21483721, 0.88723058, -0.40824816,
+        0.52058744, 0.24964368, 0.81649637,
+        0.82633746, -0.38794267, -0.40824834,
+    )
+    expected_sigma = wp.vec3(16.84809875, 1.06836915, 0.00000019)
+    expected_V = wp.mat33(
+        0.47967088, -0.77669406, 0.40824246,
+        0.57236743, -0.07568054, -0.81649727,
+        0.66506463, 0.62531471, 0.40825251,
+    )
+    # fmt: on
+
+    U, sigma, V = wp.svd3(A)
+
+    test.assertAlmostEqual(U[0][0], expected_U[0][0], places=5)
+    test.assertAlmostEqual(U[0][1], expected_U[0][1], places=5)
+    test.assertAlmostEqual(U[0][2], expected_U[0][2], places=5)
+    test.assertAlmostEqual(U[1][0], expected_U[1][0], places=5)
+    test.assertAlmostEqual(U[1][1], expected_U[1][1], places=5)
+    test.assertAlmostEqual(U[1][2], expected_U[1][2], places=5)
+    test.assertAlmostEqual(U[2][0], expected_U[2][0], places=5)
+    test.assertAlmostEqual(U[2][1], expected_U[2][1], places=5)
+    test.assertAlmostEqual(U[2][2], expected_U[2][2], places=5)
+
+    test.assertAlmostEqual(sigma[0], expected_sigma[0], places=5)
+    test.assertAlmostEqual(sigma[1], expected_sigma[1], places=5)
+    test.assertAlmostEqual(sigma[2], expected_sigma[2], places=5)
+
+    test.assertAlmostEqual(V[0][0], expected_V[0][0], places=5)
+    test.assertAlmostEqual(V[0][1], expected_V[0][1], places=5)
+    test.assertAlmostEqual(V[0][2], expected_V[0][2], places=5)
+    test.assertAlmostEqual(V[1][0], expected_V[1][0], places=5)
+    test.assertAlmostEqual(V[1][1], expected_V[1][1], places=5)
+    test.assertAlmostEqual(V[1][2], expected_V[1][2], places=5)
+    test.assertAlmostEqual(V[2][0], expected_V[2][0], places=5)
+    test.assertAlmostEqual(V[2][1], expected_V[2][1], places=5)
+    test.assertAlmostEqual(V[2][2], expected_V[2][2], places=5)
+
+    wp.launch(
+        test_multiple_return_values_svd3_kernel,
+        dim=1,
+        inputs=(A, expected_U, expected_sigma, expected_V),
+    )
+
+
 class TestCodeGen(unittest.TestCase):
     pass
 
@@ -825,6 +938,16 @@ add_kernel_test(TestCodeGen, name="test_call_syntax", kernel=test_call_syntax, d
 add_kernel_test(TestCodeGen, name="test_shadow_builtin", kernel=test_shadow_builtin, dim=1, devices=devices)
 add_kernel_test(TestCodeGen, name="test_while_condition_eval", kernel=test_while_condition_eval, dim=1, devices=devices)
 add_function_test(TestCodeGen, "test_codegen_return_in_kernel", test_codegen_return_in_kernel, devices=devices)
+add_function_test(
+    TestCodeGen, "test_ifexp_only_executes_one_branch", test_ifexp_only_executes_one_branch, devices=devices
+)
+add_function_test(
+    TestCodeGen,
+    func=test_multiple_return_values,
+    name="test_multiple_return_values",
+    devices=devices,
+)
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_conditional.py

@@ -58,6 +58,48 @@ def test_conditional_if_else_nested():
     wp.expect_eq(e, -2.0)
 
 
+@wp.kernel
+def test_conditional_ifexp():
+    a = 0.5
+    b = 2.0
+
+    c = 1.0 if a > b else -1.0
+
+    wp.expect_eq(c, -1.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_nested():
+    a = 1.0
+    b = 2.0
+
+    c = 3.0 if a > b else 6.0
+    d = 4.0 if a > b else 7.0
+    e = 1.0 if (a > b and c > d) else (-1.0 if a > b else (2.0 if c > d else -2.0))
+
+    wp.expect_eq(e, -2.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_constant():
+    a = 1.0 if False else -1.0
+    b = 2.0 if 123 else -2.0
+
+    wp.expect_eq(a, -1.0)
+    wp.expect_eq(b, 2.0)
+
+
+@wp.kernel
+def test_conditional_ifexp_constant_nested():
+    a = 1.0 if False else (2.0 if True else 3.0)
+    b = 4.0 if 0 else (5.0 if 0 else (6.0 if False else 7.0))
+    c = 8.0 if False else (9.0 if False else (10.0 if 321 else 11.0))
+
+    wp.expect_eq(a, 2.0)
+    wp.expect_eq(b, 7.0)
+    wp.expect_eq(c, 10.0)
+
+
 @wp.kernel
 def test_boolean_and():
     a = 1.0
@@ -90,7 +132,7 @@ def test_boolean_compound():
 
     d = 1.0
 
-    if a > 0.0 and b > 0.0 or c > a:
+    if (a > 0.0 and b > 0.0) or c > a:
         d = -1.0
 
     wp.expect_eq(d, -1.0)
@@ -231,6 +273,10 @@ class TestConditional(unittest.TestCase):
 
 add_kernel_test(TestConditional, kernel=test_conditional_if_else, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_conditional_if_else_nested, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_nested, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_constant, dim=1, devices=devices)
+add_kernel_test(TestConditional, kernel=test_conditional_ifexp_constant_nested, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_and, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_or, dim=1, devices=devices)
 add_kernel_test(TestConditional, kernel=test_boolean_compound, dim=1, devices=devices)

warp/tests/test_ctypes.py

@@ -541,25 +541,6 @@ def test_scalar_array_types(test, device, load, store):
         )
 
 
-@wp.kernel
-def test_transform_matrix():
-    r = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), 0.5)
-    t = wp.vec3(0.25, 0.5, -0.75)
-    s = wp.vec3(2.0, 0.5, 0.75)
-
-    m = wp.mat44(t, r, s)
-
-    p = wp.vec3(1.0, 2.0, 3.0)
-
-    r_0 = wp.quat_rotate(r, wp.cw_mul(s, p)) + t
-    r_1 = wp.transform_point(m, p)
-
-    r_2 = wp.transform_vector(m, p)
-
-    wp.expect_near(r_0, r_1, 1.0e-4)
-    wp.expect_near(r_2, r_0 - t, 1.0e-4)
-
-
 devices = get_test_devices()
 
 
@@ -628,7 +609,6 @@ add_function_test(TestCTypes, "test_vec2_transform", test_vec2_transform, device
 add_function_test(TestCTypes, "test_vec3_arg", test_vec3_arg, devices=devices, n=8)
 add_function_test(TestCTypes, "test_vec3_transform", test_vec3_transform, devices=devices, n=8)
 add_function_test(TestCTypes, "test_transform_multiply", test_transform_multiply, devices=devices, n=8)
-add_kernel_test(TestCTypes, name="test_transform_matrix", kernel=test_transform_matrix, dim=1, devices=devices)
 add_function_test(TestCTypes, "test_scalar_array", test_scalar_array, devices=devices)
 add_function_test(TestCTypes, "test_vector_array", test_vector_array, devices=devices)
 

warp/tests/test_devices.py

@@ -70,6 +70,13 @@ def test_devices_can_access_self(test, device):
             test.assertNotEqual(device, device_str)
 
 
+def test_devices_sm_count(test, device):
+    if device.is_cuda:
+        test.assertTrue(device.sm_count > 0)
+    else:
+        test.assertEqual(device.sm_count, 0)
+
+
 devices = get_test_devices()
 
 
@@ -90,6 +97,7 @@ add_function_test(
 )
 add_function_test(TestDevices, "test_devices_verify_cuda_device", test_devices_verify_cuda_device, devices=devices)
 add_function_test(TestDevices, "test_devices_can_access_self", test_devices_can_access_self, devices=devices)
+add_function_test(TestDevices, "test_devices_sm_count", test_devices_sm_count, devices=devices)
 
 
 if __name__ == "__main__":

warp/tests/test_fabricarray.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import math
 import unittest
 from typing import Any
@@ -70,7 +72,7 @@ def _warp_type_to_fabric(dtype, is_array=False):
 
 
 # returns a fabric array interface constructed from a regular array
-def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list = None, copy=False):
+def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list[int] | None = None, copy=False):
     assert isinstance(data, wp.array)
     assert data.ndim == 1
 
@@ -138,7 +140,7 @@ def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: li
 
 
 # returns a fabric array array interface constructed from a list of regular arrays
-def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list = None):
+def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list[int] | None = None):
     # data should be a list of arrays
     assert isinstance(data, list)