warp-lang 1.3.3__py3-none-manylinux2014_x86_64.whl → 1.4.1__py3-none-manylinux2014_x86_64.whl

warp/tests/test_vec.py

@@ -1035,6 +1035,142 @@ def test_casting_constructors(test, device, dtype, register_kernels=False):
     assert_np_equal(out, a_grad.numpy())
 
 
+def test_vec_assign(test, device, dtype, register_kernels=False):
+    np_type = np.dtype(dtype)
+    wp_type = wp.types.np_dtype_to_warp_type[np_type]
+
+    vec2 = wp.types.vector(length=2, dtype=wp_type)
+    vec3 = wp.types.vector(length=3, dtype=wp_type)
+    vec4 = wp.types.vector(length=4, dtype=wp_type)
+
+    def vectest_read_write_store(
+        x: wp.array(dtype=wp_type), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3), c: wp.array(dtype=vec4)
+    ):
+        tid = wp.tid()
+
+        t = a[tid]
+        t[0] = x[tid]
+        a[tid] = t
+
+        u = b[tid]
+        u[1] = x[tid]
+        b[tid] = u
+
+        v = c[tid]
+        v[2] = x[tid]
+        c[tid] = v
+
+    def vectest_in_register(
+        x: wp.array(dtype=wp_type), y: wp.array(dtype=vec3), a: wp.array(dtype=vec2), b: wp.array(dtype=vec3)
+    ):
+        tid = wp.tid()
+
+        f = vec3(wp_type(0.0))
+        b_vec = b[tid]
+        f[0] = b_vec[1]
+        f[2] = b_vec[0] * b_vec[1]
+        y[tid] = f
+
+        g = wp_type(0.0)
+        a_vec = a[tid]
+        g = a_vec[0] + a_vec[1]
+        x[tid] = g
+
+    def vectest_in_register_overwrite(x: wp.array(dtype=vec3), a: wp.array(dtype=vec3)):
+        tid = wp.tid()
+
+        f = vec3(wp_type(0.0))
+        a_vec = a[tid]
+        f = a_vec
+        f[1] = wp_type(3.0)
+
+        x[tid] = f
+
+    def vectest_component(x: wp.array(dtype=vec3), y: wp.array(dtype=wp_type)):
+        i = wp.tid()
+
+        a = vec3(wp_type(0.0))
+        a.x = wp_type(1.0) * y[i]
+        a.y = wp_type(2.0) * y[i]
+        a.z = wp_type(3.0) * y[i]
+        x[i] = a
+
+    kernel_read_write_store = getkernel(vectest_read_write_store, suffix=dtype.__name__)
+    kernel_in_register = getkernel(vectest_in_register, suffix=dtype.__name__)
+    kernel_in_register_overwrite = getkernel(vectest_in_register_overwrite, suffix=dtype.__name__)
+    kernel_component = getkernel(vectest_component, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+    b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+    c = wp.ones(1, dtype=vec4, device=device, requires_grad=True)
+    x = wp.full(1, value=2.0, dtype=wp_type, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_read_write_store, dim=1, inputs=[x, a, b, c], device=device)
+
+    tape.backward(
+        grads={
+            a: wp.ones_like(a, requires_grad=False),
+            b: wp.ones_like(b, requires_grad=False),
+            c: wp.ones_like(c, requires_grad=False),
+        }
+    )
+
+    assert_np_equal(a.numpy(), np.array([[2.0, 1.0]], dtype=np_type))
+    assert_np_equal(b.numpy(), np.array([[1.0, 2.0, 1.0]], dtype=np_type))
+    assert_np_equal(c.numpy(), np.array([[1.0, 1.0, 2.0, 1.0]], dtype=np_type))
+    assert_np_equal(x.grad.numpy(), np.array([3.0], dtype=np_type))
+
+    tape.reset()
+
+    a = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+    b = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+    x = wp.zeros(1, dtype=wp_type, device=device, requires_grad=True)
+    y = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+
+    with tape:
+        wp.launch(kernel_in_register, dim=1, inputs=[x, y, a, b], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False), y: wp.ones_like(y, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([2.0], dtype=np_type))
+    assert_np_equal(y.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 1.0]], dtype=np_type))
+    assert_np_equal(b.grad.numpy(), np.array([[1.0, 2.0, 0.0]], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+    y = wp.ones(1, dtype=wp_type, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_component, dim=1, inputs=[x, y], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 2.0, 3.0]], dtype=np_type))
+    assert_np_equal(y.grad.numpy(), np.array([6.0], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros(1, dtype=vec3, device=device, requires_grad=True)
+    a = wp.ones(1, dtype=vec3, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(kernel_in_register_overwrite, dim=1, inputs=[x, a], device=device)
+
+    tape.backward(grads={x: wp.ones_like(x, requires_grad=False)})
+
+    assert_np_equal(x.numpy(), np.array([[1.0, 3.0, 1.0]], dtype=np_type))
+    assert_np_equal(a.grad.numpy(), np.array([[1.0, 0.0, 1.0]], dtype=np_type))
+
+
 @wp.kernel
 def test_vector_constructor_value_func():
     a = wp.vec2()
@@ -1170,6 +1306,13 @@ for dtype in np_float_types:
         devices=devices,
         dtype=dtype,
     )
+    add_function_test_register_kernel(
+        TestVec,
+        f"test_vec_assign_{dtype.__name__}",
+        test_vec_assign,
+        devices=devices,
+        dtype=dtype,
+    )
 
 add_function_test(
     TestVec,

warp/tests/test_vec_scalar_ops.py

@@ -192,12 +192,14 @@ def test_py_arithmetic_ops(test, device, dtype):
     test.assertSequenceEqual(-v, make_vec(-1, 2, -3))
     test.assertSequenceEqual(v + vec_cls(5, 5, 5), make_vec(6, 3, 8))
     test.assertSequenceEqual(v - vec_cls(5, 5, 5), make_vec(-4, -7, -2))
+    test.assertSequenceEqual(v % vec_cls(2, 2, 2), make_vec(1, 0, 1))
 
     v = vec_cls(2, 4, 6)
     test.assertSequenceEqual(v * wptype(2), make_vec(4, 8, 12))
     test.assertSequenceEqual(wptype(2) * v, make_vec(4, 8, 12))
     test.assertSequenceEqual(v / wptype(2), make_vec(1, 2, 3))
     test.assertSequenceEqual(wptype(24) / v, make_vec(12, 6, 4))
+    test.assertSequenceEqual(v % vec_cls(3, 3, 3), make_vec(2, 1, 0))
 
 
 def test_constructors(test, device, dtype, register_kernels=False):
@@ -1797,6 +1799,140 @@ def test_dotproduct(test, device, dtype, register_kernels=False):
         tape.zero()
 
 
+def test_modulo(test, device, dtype, register_kernels=False):
+    rng = np.random.default_rng(123)
+
+    tol = {
+        np.float16: 1.0e-2,
+        np.float32: 1.0e-6,
+        np.float64: 1.0e-8,
+    }.get(dtype, 0)
+
+    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec2 = wp.types.vector(length=2, dtype=wptype)
+    vec3 = wp.types.vector(length=3, dtype=wptype)
+    vec4 = wp.types.vector(length=4, dtype=wptype)
+    vec5 = wp.types.vector(length=5, dtype=wptype)
+
+    def check_mod(
+        s2: wp.array(dtype=vec2),
+        s3: wp.array(dtype=vec3),
+        s4: wp.array(dtype=vec4),
+        s5: wp.array(dtype=vec5),
+        v2: wp.array(dtype=vec2),
+        v3: wp.array(dtype=vec3),
+        v4: wp.array(dtype=vec4),
+        v5: wp.array(dtype=vec5),
+        v20: wp.array(dtype=wptype),
+        v21: wp.array(dtype=wptype),
+        v30: wp.array(dtype=wptype),
+        v31: wp.array(dtype=wptype),
+        v32: wp.array(dtype=wptype),
+        v40: wp.array(dtype=wptype),
+        v41: wp.array(dtype=wptype),
+        v42: wp.array(dtype=wptype),
+        v43: wp.array(dtype=wptype),
+        v50: wp.array(dtype=wptype),
+        v51: wp.array(dtype=wptype),
+        v52: wp.array(dtype=wptype),
+        v53: wp.array(dtype=wptype),
+        v54: wp.array(dtype=wptype),
+    ):
+        v20[0] = (wptype(2) * wp.mod(v2[0], s2[0]))[0]
+        v21[0] = (wptype(2) * wp.mod(v2[0], s2[0]))[1]
+
+        v30[0] = (wptype(2) * wp.mod(v3[0], s3[0]))[0]
+        v31[0] = (wptype(2) * wp.mod(v3[0], s3[0]))[1]
+        v32[0] = (wptype(2) * wp.mod(v3[0], s3[0]))[2]
+
+        v40[0] = (wptype(2) * wp.mod(v4[0], s4[0]))[0]
+        v41[0] = (wptype(2) * wp.mod(v4[0], s4[0]))[1]
+        v42[0] = (wptype(2) * wp.mod(v4[0], s4[0]))[2]
+        v43[0] = (wptype(2) * wp.mod(v4[0], s4[0]))[3]
+
+        v50[0] = (wptype(2) * wp.mod(v5[0], s5[0]))[0]
+        v51[0] = (wptype(2) * wp.mod(v5[0], s5[0]))[1]
+        v52[0] = (wptype(2) * wp.mod(v5[0], s5[0]))[2]
+        v53[0] = (wptype(2) * wp.mod(v5[0], s5[0]))[3]
+        v54[0] = (wptype(2) * wp.mod(v5[0], s5[0]))[4]
+
+    kernel = getkernel(check_mod, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    s2 = wp.array(randvals(rng, (1, 2), dtype), dtype=vec2, requires_grad=True, device=device)
+    s3 = wp.array(randvals(rng, (1, 3), dtype), dtype=vec3, requires_grad=True, device=device)
+    s4 = wp.array(randvals(rng, (1, 4), dtype), dtype=vec4, requires_grad=True, device=device)
+    s5 = wp.array(randvals(rng, (1, 5), dtype), dtype=vec5, requires_grad=True, device=device)
+    v2 = wp.array(randvals(rng, (1, 2), dtype), dtype=vec2, requires_grad=True, device=device)
+    v3 = wp.array(randvals(rng, (1, 3), dtype), dtype=vec3, requires_grad=True, device=device)
+    v4 = wp.array(randvals(rng, (1, 4), dtype), dtype=vec4, requires_grad=True, device=device)
+    v5 = wp.array(randvals(rng, (1, 5), dtype), dtype=vec5, requires_grad=True, device=device)
+    v20 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v21 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v30 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v31 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v32 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v40 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v41 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v42 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v43 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v50 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v51 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v52 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v53 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    v54 = wp.zeros(1, dtype=wptype, requires_grad=True, device=device)
+    tape = wp.Tape()
+    with tape:
+        wp.launch(
+            kernel,
+            dim=1,
+            inputs=[
+                s2,
+                s3,
+                s4,
+                s5,
+                v2,
+                v3,
+                v4,
+                v5,
+            ],
+            outputs=[
+                v20,
+                v21,
+                v30,
+                v31,
+                v32,
+                v40,
+                v41,
+                v42,
+                v43,
+                v50,
+                v51,
+                v52,
+                v53,
+                v54,
+            ],
+            device=device,
+        )
+
+    assert_np_equal(v20.numpy()[0], 2.0 * np.fmod(v2.numpy(), s2.numpy())[0, 0], tol=10 * tol)
+    assert_np_equal(v21.numpy()[0], 2.0 * np.fmod(v2.numpy(), s2.numpy())[0, 1], tol=10 * tol)
+    assert_np_equal(v30.numpy()[0], 2.0 * np.fmod(v3.numpy(), s3.numpy())[0, 0], tol=10 * tol)
+    assert_np_equal(v31.numpy()[0], 2.0 * np.fmod(v3.numpy(), s3.numpy())[0, 1], tol=10 * tol)
+    assert_np_equal(v32.numpy()[0], 2.0 * np.fmod(v3.numpy(), s3.numpy())[0, 2], tol=10 * tol)
+    assert_np_equal(v40.numpy()[0], 2.0 * np.fmod(v4.numpy(), s4.numpy())[0, 0], tol=10 * tol)
+    assert_np_equal(v41.numpy()[0], 2.0 * np.fmod(v4.numpy(), s4.numpy())[0, 1], tol=10 * tol)
+    assert_np_equal(v42.numpy()[0], 2.0 * np.fmod(v4.numpy(), s4.numpy())[0, 2], tol=10 * tol)
+    assert_np_equal(v43.numpy()[0], 2.0 * np.fmod(v4.numpy(), s4.numpy())[0, 3], tol=10 * tol)
+    assert_np_equal(v50.numpy()[0], 2.0 * np.fmod(v5.numpy(), s5.numpy())[0, 0], tol=10 * tol)
+    assert_np_equal(v51.numpy()[0], 2.0 * np.fmod(v5.numpy(), s5.numpy())[0, 1], tol=10 * tol)
+    assert_np_equal(v52.numpy()[0], 2.0 * np.fmod(v5.numpy(), s5.numpy())[0, 2], tol=10 * tol)
+    assert_np_equal(v53.numpy()[0], 2.0 * np.fmod(v5.numpy(), s5.numpy())[0, 3], tol=10 * tol)
+    assert_np_equal(v54.numpy()[0], 2.0 * np.fmod(v5.numpy(), s5.numpy())[0, 4], tol=10 * tol)
+
+
 def test_equivalent_types(test, device, dtype, register_kernels=False):
     wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
@@ -2150,6 +2286,9 @@ for dtype in np_scalar_types:
     add_function_test_register_kernel(
         TestVecScalarOps, f"test_addition_{dtype.__name__}", test_addition, devices=devices, dtype=dtype
     )
+    add_function_test_register_kernel(
+        TestVecScalarOps, f"test_modulo_{dtype.__name__}", test_modulo, devices=devices, dtype=dtype
+    )
     add_function_test_register_kernel(
         TestVecScalarOps, f"test_dotproduct_{dtype.__name__}", test_dotproduct, devices=devices, dtype=dtype
     )

warp/tests/unittest_suites.py

@@ -98,6 +98,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_bvh import TestBvh
     from warp.tests.test_closest_point_edge_edge import TestClosestPointEdgeEdgeMethods
     from warp.tests.test_codegen import TestCodeGen
+    from warp.tests.test_codegen_instancing import TestCodeGenInstancing
     from warp.tests.test_compile_consts import TestConstants
     from warp.tests.test_conditional import TestConditional
     from warp.tests.test_copy import TestCopy
@@ -117,6 +118,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_fem import TestFem, TestFemShapeFunctions
     from warp.tests.test_fp16 import TestFp16
     from warp.tests.test_func import TestFunc
+    from warp.tests.test_future_annotations import TestFutureAnnotations
     from warp.tests.test_generics import TestGenerics
     from warp.tests.test_grad import TestGrad
     from warp.tests.test_grad_customs import TestGradCustoms
@@ -160,6 +162,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_reload import TestReload
     from warp.tests.test_rounding import TestRounding
     from warp.tests.test_runlength_encode import TestRunlengthEncode
+    from warp.tests.test_scalar_ops import TestScalarOps
     from warp.tests.test_sim_grad import TestSimGradients
     from warp.tests.test_sim_kinematics import TestSimKinematics
     from warp.tests.test_smoothstep import TestSmoothstep
@@ -172,8 +175,10 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
     from warp.tests.test_tape import TestTape
     from warp.tests.test_torch import TestTorch
     from warp.tests.test_transient_module import TestTransientModule
+    from warp.tests.test_triangle_closest_point import TestTriangleClosestPoint
     from warp.tests.test_types import TestTypes
     from warp.tests.test_utils import TestUtils
+    from warp.tests.test_vbd import TestVBD
     from warp.tests.test_vec import TestVec
     from warp.tests.test_vec_lite import TestVecLite
     from warp.tests.test_vec_scalar_ops import TestVecScalarOps
@@ -193,6 +198,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestBvh,
         TestClosestPointEdgeEdgeMethods,
         TestCodeGen,
+        TestCodeGenInstancing,
         TestConstants,
         TestConditional,
         TestCopy,
@@ -211,6 +217,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestFemShapeFunctions,
         TestFp16,
         TestFunc,
+        TestFutureAnnotations,
         TestGenerics,
         TestGrad,
         TestGradCustoms,
@@ -254,6 +261,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestReload,
         TestRounding,
         TestRunlengthEncode,
+        TestScalarOps,
         TestSimGradients,
         TestSimKinematics,
         TestSmoothstep,
@@ -266,8 +274,10 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
         TestTape,
         TestTorch,
         TestTransientModule,
+        TestTriangleClosestPoint,
         TestTypes,
         TestUtils,
+        TestVBD,
         TestVec,
         TestVecLite,
         TestVecScalarOps,
@@ -288,6 +298,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
     from warp.tests.test_array_reduce import TestArrayReduce
     from warp.tests.test_bvh import TestBvh
     from warp.tests.test_codegen import TestCodeGen
+    from warp.tests.test_codegen_instancing import TestCodeGenInstancing
     from warp.tests.test_compile_consts import TestConstants
     from warp.tests.test_conditional import TestConditional
     from warp.tests.test_ctypes import TestCTypes
@@ -332,6 +343,7 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
         TestArrayReduce,
         TestBvh,
         TestCodeGen,
+        TestCodeGenInstancing,
         TestConstants,
         TestConditional,
         TestCTypes,

warp/tests/unittest_utils.py

@@ -198,11 +198,15 @@ class CheckOutput:
         if s != "":
             print(s.rstrip())
 
-        # fail if test produces unexpected output (e.g.: from wp.expect_eq() builtins)
-        # we allow strings starting of the form "Module xxx load on device xxx"
-        # for lazy loaded modules
-        if s != "" and not s.startswith("Module"):
-            self.test.fail(f"Unexpected output:\n'{s.rstrip()}'")
+            # fail if test produces unexpected output (e.g.: from wp.expect_eq() builtins)
+            # we allow strings starting of the form "Module xxx load on device xxx"
+            # for lazy loaded modules
+            filtered_s = "\n".join(
+                [line for line in s.splitlines() if not (line.startswith("Module") and "load on device" in line)]
+            )
+
+            if filtered_s.strip():
+                self.test.fail(f"Unexpected output:\n'{s.rstrip()}'")
 
 
 def assert_array_equal(result: wp.array, expect: wp.array):

warp/thirdparty/dlpack.py

@@ -58,6 +58,7 @@ class DLDataTypeCode(ctypes.c_uint8):
     kDLOpaquePointer = 3
     kDLBfloat = 4
     kDLComplex = 5
+    kDLBool = 6
 
     def __str__(self):
         return {
@@ -66,6 +67,7 @@ class DLDataTypeCode(ctypes.c_uint8):
             self.kDLFloat: "float",
             self.kDLBfloat: "bfloat",
             self.kDLComplex: "complex",
+            self.kDLBool: "bool",
             self.kDLOpaquePointer: "void_p",
         }[self.value]
 
@@ -85,7 +87,7 @@ class DLDataType(ctypes.Structure):
         ("lanes", ctypes.c_uint16),
     ]
     TYPE_MAP = {
-        "bool": (DLDataTypeCode.kDLUInt, 1, 1),
+        "bool": (DLDataTypeCode.kDLBool, 8, 1),
         "int8": (DLDataTypeCode.kDLInt, 8, 1),
         "int16": (DLDataTypeCode.kDLInt, 16, 1),
         "int32": (DLDataTypeCode.kDLInt, 32, 1),