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

warp/tests/test_mat.py

@@ -1559,6 +1559,83 @@ def test_transform_vector(test, device, dtype, register_kernels=False):
             tape.zero()
 
 
+def test_mat_array_type_indexing(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)
+    mat22 = wp.types.matrix(shape=(2, 2), dtype=wp_type)
+    mat33 = wp.types.matrix(shape=(3, 3), dtype=wp_type)
+
+    def mattest_read_write_store(x: wp.array(dtype=wp_type), a: wp.array(dtype=mat22)):
+        tid = wp.tid()
+
+        t = a[tid]
+        t[0, 0] = x[tid]
+        a[tid] = t
+
+    def mattest_in_register(x: wp.array2d(dtype=mat22), y: wp.array(dtype=vec2)):
+        i, j = wp.tid()
+
+        a = mat22(wp_type(0.0))
+        a[0] = y[i]
+        a[1, 1] = wp_type(3.0)
+        x[i, j] = a
+
+    def mattest_in_register_overwrite(x: wp.array2d(dtype=mat22), y: wp.array(dtype=vec2)):
+        i, j = wp.tid()
+
+        a = mat22(wp_type(0.0))
+        a[0] = y[i]
+        a[0, 1] = wp_type(3.0)
+        x[i, j] = a
+
+    kernel_read_write_store = getkernel(mattest_read_write_store, suffix=dtype.__name__)
+    kernel_in_register = getkernel(mattest_in_register, suffix=dtype.__name__)
+    kernel_in_register_overwrite = getkernel(mattest_in_register_overwrite, suffix=dtype.__name__)
+
+    if register_kernels:
+        return
+
+    a = wp.ones(1, dtype=mat22, 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], device=device)
+
+    tape.backward(grads={a: wp.ones_like(a, requires_grad=False)})
+
+    assert_np_equal(a.numpy(), np.array([[[2.0, 1.0], [1.0, 1.0]]], dtype=np_type))
+    assert_np_equal(x.grad.numpy(), np.array([1.0], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros((1, 1), dtype=mat22, device=device, requires_grad=True)
+    y = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+
+    with tape:
+        wp.launch(kernel_in_register, dim=(1, 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, 1.0], [0.0, 3.0]]]], dtype=np_type))
+    assert_np_equal(y.grad.numpy(), np.array([[1.0, 1.0]], dtype=np_type))
+
+    tape.reset()
+
+    x = wp.zeros((1, 1), dtype=mat22, device=device, requires_grad=True)
+    y = wp.ones(1, dtype=vec2, device=device, requires_grad=True)
+
+    with tape:
+        wp.launch(kernel_in_register_overwrite, dim=(1, 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, 3.0], [0.0, 0.0]]]], dtype=np_type))
+    assert_np_equal(y.grad.numpy(), np.array([[1.0, 0.0]], dtype=np_type))
+
+
 # Test matrix constructors using explicit type (float16)
 # note that these tests are specifically not using generics / closure
 # args to create kernels dynamically (like the rest of this file)
@@ -1791,6 +1868,13 @@ for dtype in np_float_types:
         TestMat, f"test_determinant_{dtype.__name__}", test_determinant, devices=devices, dtype=dtype
     )
     add_function_test_register_kernel(TestMat, f"test_skew_{dtype.__name__}", test_skew, devices=devices, dtype=dtype)
+    add_function_test_register_kernel(
+        TestMat,
+        f"test_mat_array_type_indexing_{dtype.__name__}",
+        test_mat_array_type_indexing,
+        devices=devices,
+        dtype=dtype,
+    )
 
 
 if __name__ == "__main__":

warp/tests/test_mesh_query_point.py

@@ -11,6 +11,7 @@ import unittest
 import numpy as np
 
 import warp as wp
+import warp.examples
 from warp.tests.unittest_utils import *
 
 
@@ -654,6 +655,192 @@ def test_mesh_query_furthest_point(test, device):
     assert_np_equal(dist_query.numpy(), dist_brute.numpy(), tol=1.0e-3)
 
 
+@wp.func
+def triangle_closest_point_for_test(a: wp.vec3, b: wp.vec3, c: wp.vec3, p: wp.vec3):
+    ab = b - a
+    ac = c - a
+    ap = p - a
+
+    d1 = wp.dot(ab, ap)
+    d2 = wp.dot(ac, ap)
+    if d1 <= 0.0 and d2 <= 0.0:
+        bary = wp.vec3(1.0, 0.0, 0.0)
+        return a, bary
+
+    bp = p - b
+    d3 = wp.dot(ab, bp)
+    d4 = wp.dot(ac, bp)
+    if d3 >= 0.0 and d4 <= d3:
+        bary = wp.vec3(0.0, 1.0, 0.0)
+        return b, bary
+
+    cp = p - c
+    d5 = wp.dot(ab, cp)
+    d6 = wp.dot(ac, cp)
+    if d6 >= 0.0 and d5 <= d6:
+        bary = wp.vec3(0.0, 0.0, 1.0)
+        return c, bary
+
+    vc = d1 * d4 - d3 * d2
+    if vc <= 0.0 and d1 >= 0.0 and d3 <= 0.0:
+        v = d1 / (d1 - d3)
+        bary = wp.vec3(1.0 - v, v, 0.0)
+        return a + v * ab, bary
+
+    vb = d5 * d2 - d1 * d6
+    if vb <= 0.0 and d2 >= 0.0 and d6 <= 0.0:
+        v = d2 / (d2 - d6)
+        bary = wp.vec3(1.0 - v, 0.0, v)
+        return a + v * ac, bary
+
+    va = d3 * d6 - d5 * d4
+    if va <= 0.0 and (d4 - d3) >= 0.0 and (d5 - d6) >= 0.0:
+        v = (d4 - d3) / ((d4 - d3) + (d5 - d6))
+        bary = wp.vec3(0.0, 1.0 - v, v)
+        return b + v * (c - b), bary
+
+    denom = 1.0 / (va + vb + vc)
+    v = vb * denom
+    w = vc * denom
+    bary = wp.vec3(1.0 - v - w, v, w)
+    return a + v * ab + w * ac, bary
+
+
+def load_mesh():
+    from pxr import Usd, UsdGeom
+
+    usd_stage = Usd.Stage.Open(os.path.join(wp.examples.get_asset_directory(), "bunny.usd"))
+    usd_geom = UsdGeom.Mesh(usd_stage.GetPrimAtPath("/root/bunny"))
+
+    vertices = np.array(usd_geom.GetPointsAttr().Get())
+    faces = np.array(usd_geom.GetFaceVertexIndicesAttr().Get())
+
+    return vertices, faces
+
+
+@wp.kernel
+def point_query_aabb_and_closest(
+    query_radius: float,
+    mesh_id: wp.uint64,
+    pts: wp.array(dtype=wp.vec3),
+    pos: wp.array(dtype=wp.vec3),
+    tri_indices: wp.array(dtype=wp.int32, ndim=2),
+    query_results_num_collisions: wp.array(dtype=wp.int32),
+    query_results_min_dist: wp.array(dtype=float),
+    query_results_closest_point_velocity: wp.array(dtype=wp.vec3),
+):
+    p_index = wp.tid()
+    p = pts[p_index]
+
+    lower = wp.vec3(p[0] - query_radius, p[1] - query_radius, p[2] - query_radius)
+    upper = wp.vec3(p[0] + query_radius, p[1] + query_radius, p[2] + query_radius)
+
+    closest_query = wp.mesh_query_point_no_sign(mesh_id, p, query_radius)
+    if closest_query.result:
+        closest_p = wp.mesh_eval_position(mesh_id, closest_query.face, closest_query.u, closest_query.v)
+        closest_p_vel = wp.mesh_eval_velocity(mesh_id, closest_query.face, closest_query.u, closest_query.v)
+
+        query_results_min_dist[p_index] = wp.length(closest_p - p)
+        query_results_closest_point_velocity[p_index] = closest_p_vel
+
+    query = wp.mesh_query_aabb(mesh_id, lower, upper)
+
+    tri_index = wp.int32(0)
+    num_collisions = wp.int32(0)
+    min_dis_to_tris = query_radius
+    while wp.mesh_query_aabb_next(query, tri_index):
+        t1 = tri_indices[tri_index, 0]
+        t2 = tri_indices[tri_index, 1]
+        t3 = tri_indices[tri_index, 2]
+
+        u1 = pos[t1]
+        u2 = pos[t2]
+        u3 = pos[t3]
+
+        closest_p1, barycentric1 = triangle_closest_point_for_test(u1, u2, u3, p)
+
+        dis = wp.length(closest_p1 - p)
+
+        if dis < query_radius:
+            num_collisions = num_collisions + 1
+
+    query_results_num_collisions[p_index] = num_collisions
+
+
+@unittest.skipUnless(USD_AVAILABLE, "Requires usd-core")
+def test_set_mesh_points(test, device):
+    vs, fs = load_mesh()
+
+    vertices1 = wp.array(vs, dtype=wp.vec3, device=device)
+    velocities1_np = np.random.randn(vertices1.shape[0], 3)
+    velocities1 = wp.array(velocities1_np, dtype=wp.vec3, device=device)
+
+    faces = wp.array(fs, dtype=wp.int32, device=device)
+    mesh = wp.Mesh(vertices1, faces, velocities=velocities1)
+    fs_2D = faces.reshape((-1, 3))
+    np.random.seed(12345)
+    n = 1000
+    query_radius = 0.2
+
+    pts1 = wp.array(np.random.randn(n, 3), dtype=wp.vec3, device=device)
+
+    query_results_num_cols1 = wp.zeros(n, dtype=wp.int32, device=device)
+    query_results_min_dist1 = wp.zeros(n, dtype=float, device=device)
+    query_results_closest_point_velocity1 = wp.zeros(n, dtype=wp.vec3, device=device)
+
+    wp.launch(
+        kernel=point_query_aabb_and_closest,
+        inputs=[
+            query_radius,
+            mesh.id,
+            pts1,
+            vertices1,
+            fs_2D,
+            query_results_num_cols1,
+            query_results_min_dist1,
+            query_results_closest_point_velocity1,
+        ],
+        dim=n,
+        device=device,
+    )
+
+    shift = np.random.randn(3)
+
+    vs_higher = vs + shift
+    vertices2 = wp.array(vs_higher, dtype=wp.vec3, device=device)
+
+    velocities2_np = velocities1_np + shift[None, ...]
+    velocities2 = wp.array(velocities2_np, dtype=wp.vec3, device=device)
+
+    pts2 = wp.array(pts1.numpy() + shift, dtype=wp.vec3, device=device)
+
+    mesh.points = vertices2
+    mesh.velocities = velocities2
+
+    query_results_num_cols2 = wp.zeros(n, dtype=wp.int32, device=device)
+    query_results_min_dist2 = wp.zeros(n, dtype=float, device=device)
+    query_results_closest_point_velocity2 = wp.array([shift for i in range(n)], dtype=wp.vec3, device=device)
+
+    wp.launch(
+        kernel=point_query_aabb_and_closest,
+        inputs=[
+            query_radius,
+            mesh.id,
+            pts2,
+            vertices2,
+            fs_2D,
+            query_results_num_cols2,
+            query_results_min_dist2,
+            query_results_closest_point_velocity2,
+        ],
+        dim=n,
+        device=device,
+    )
+
+    test.assertTrue((query_results_num_cols1.numpy() == query_results_num_cols2.numpy()).all())
+    test.assertTrue(((query_results_min_dist1.numpy() - query_results_min_dist2.numpy()) < 1e-5).all())
+
+
 devices = get_test_devices()
 
 
@@ -684,6 +871,7 @@ class TestMeshQueryPoint(unittest.TestCase):
 add_function_test(TestMeshQueryPoint, "test_mesh_query_point", test_mesh_query_point, devices=devices)
 add_function_test(TestMeshQueryPoint, "test_mesh_query_furthest_point", test_mesh_query_furthest_point, devices=devices)
 add_function_test(TestMeshQueryPoint, "test_adj_mesh_query_point", test_adj_mesh_query_point, devices=devices)
+add_function_test(TestMeshQueryPoint, "test_set_mesh_points", test_set_mesh_points, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_model.py

@@ -157,6 +157,19 @@ class TestModel(unittest.TestCase):
         assert builder.body_mass == [1.0, 4.0]
         assert builder.body_inv_mass == [1.0, 0.25]
 
+        # create another builder, test add_builder function
+        builder2 = ModelBuilder()
+        builder2.add_builder(builder)
+        assert builder2.articulation_count == builder.articulation_count
+        assert builder2.joint_count == builder.joint_count
+        assert builder2.body_count == builder.body_count
+        assert builder2.shape_count == builder.shape_count
+        assert builder2.articulation_start == builder.articulation_start
+        # add the same builder again
+        builder2.add_builder(builder)
+        assert builder2.articulation_count == 2 * builder.articulation_count
+        assert builder2.articulation_start == [0, 1, 2, 3]
+
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()

warp/tests/test_module_hashing.py

@@ -25,6 +25,11 @@ def fn():
 @wp.func
 def fn(value: int):
     wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
 """
 
 # should be same hash as FUNC_OVERLOAD_1
@@ -38,6 +43,11 @@ def fn():
 @wp.func
 def fn(value: int):
     wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
 """
 
 # should be different hash than FUNC_OVERLOAD_1 (first overload is different)
@@ -51,6 +61,11 @@ def fn():
 @wp.func
 def fn(value: int):
     wp.print(value)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
 """
 
 # should be different hash than FUNC_OVERLOAD_1 (second overload is different)
@@ -64,6 +79,11 @@ def fn():
 @wp.func
 def fn(value: int):
     wp.print(value + 1)
+
+@wp.kernel
+def k():
+    print(fn())
+    print(fn(99))
 """
 
 FUNC_GENERIC_1 = """# -*- coding: utf-8 -*-
@@ -78,6 +98,11 @@ def generic_fn(x: Any):
 @wp.func
 def generic_fn(x: Any, y: Any):
     return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
 """
 
 # should be same hash as FUNC_GENERIC_1
@@ -93,6 +118,11 @@ def generic_fn(x: Any):
 @wp.func
 def generic_fn(x: Any, y: Any):
     return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
 """
 
 # should be different hash than FUNC_GENERIC_1 (first overload is different)
@@ -108,6 +138,11 @@ def generic_fn(x: Any):
 @wp.func
 def generic_fn(x: Any, y: Any):
     return x * y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
 """
 
 # should be different hash than FUNC_GENERIC_1 (second overload is different)
@@ -123,6 +158,11 @@ def generic_fn(x: Any):
 @wp.func
 def generic_fn(x: Any, y: Any):
     return x + y
+
+@wp.kernel
+def k():
+    print(generic_fn(17))
+    print(generic_fn(17, 42))
 """
 
 

warp/tests/test_multigpu.py

@@ -11,7 +11,7 @@ import numpy as np
 
 import warp as wp
 from warp.tests.unittest_utils import *
-from warp.utils import check_iommu
+from warp.utils import check_p2p
 
 
 @wp.kernel
@@ -105,7 +105,7 @@ class TestMultiGPU(unittest.TestCase):
         assert wp.get_cuda_device() == initial_cuda_device
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
-    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
+    @unittest.skipUnless(check_p2p(), "Peer-to-Peer transfers not supported")
     def test_multigpu_pingpong(self):
         n = 1024 * 1024
 
@@ -129,7 +129,7 @@ class TestMultiGPU(unittest.TestCase):
         assert_np_equal(a1.numpy(), expected)
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
-    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
+    @unittest.skipUnless(check_p2p(), "Peer-to-Peer transfers not supported")
     def test_multigpu_pingpong_streams(self):
         n = 1024 * 1024
 

warp/tests/test_overwrite.py

@@ -7,6 +7,8 @@ import numpy as np
 import warp as wp
 from warp.tests.unittest_utils import *
 
+wp.init()  # For wp.context.runtime.core.is_cutlass_enabled()
+
 # kernels are defined in the global scope, to ensure wp.Kernel objects are not GC'ed in the MGPU case
 # kernel args are assigned array modes during codegen, so wp.Kernel objects generated during codegen
 # must be preserved for overwrite tracking to function
@@ -362,6 +364,9 @@ def test_copy(test, device):
 
 # wp.matmul uses wp.record_func. Ensure array modes are propagated correctly.
 def test_matmul(test, device):
+    if device.is_cuda and not wp.context.runtime.core.is_cutlass_enabled():
+        test.skipTest("Warp was not built with CUTLASS support")
+
     saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
     try:
         wp.config.verify_autograd_array_access = True
@@ -387,6 +392,9 @@ def test_matmul(test, device):
 
 # wp.batched_matmul uses wp.record_func. Ensure array modes are propagated correctly.
 def test_batched_matmul(test, device):
+    if device.is_cuda and not wp.context.runtime.core.is_cutlass_enabled():
+        test.skipTest("Warp was not built with CUTLASS support")
+
     saved_verify_autograd_array_access_setting = wp.config.verify_autograd_array_access
     try:
         wp.config.verify_autograd_array_access = True