warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.0__py3-none-win_amd64.whl

warp/tests/test_array_reduce.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+from warp.utils import array_inner, array_sum
+
+wp.init()
+
+
+def make_test_array_sum(dtype):
+    N = 1000
+
+    def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
+        cols = wp.types.type_length(dtype)
+
+        values_np = rng.random(size=(N, cols))
+        values = wp.array(values_np, device=device, dtype=dtype)
+
+        vsum = array_sum(values)
+        ref_vsum = values_np.sum(axis=0)
+
+        assert_np_equal(vsum / N, ref_vsum / N, 0.0001)
+
+    return test_array_sum
+
+
+def make_test_array_sum_axis(dtype):
+    I = 5
+    J = 10
+    K = 2
+
+    N = I * J * K
+
+    def test_array_sum(test, device):
+        rng = np.random.default_rng(123)
+
+        values_np = rng.random(size=(I, J, K))
+        values = wp.array(values_np, shape=(I, J, K), device=device, dtype=dtype)
+
+        for axis in range(3):
+            vsum = array_sum(values, axis=axis)
+            ref_vsum = values_np.sum(axis=axis)
+
+            assert_np_equal(vsum.numpy() / N, ref_vsum / N, 0.0001)
+
+    return test_array_sum
+
+
+def test_array_sum_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    assert_np_equal(array_sum(values), np.zeros(2))
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_sum(values, axis=0).numpy(), np.zeros(3))
+
+
+def make_test_array_inner(dtype):
+    N = 1000
+
+    def test_array_inner(test, device):
+        rng = np.random.default_rng(123)
+
+        cols = wp.types.type_length(dtype)
+
+        a_np = rng.random(size=(N, cols))
+        b_np = rng.random(size=(N, cols))
+
+        a = wp.array(a_np, device=device, dtype=dtype)
+        b = wp.array(b_np, device=device, dtype=dtype)
+
+        ab = array_inner(a, b)
+        ref_ab = np.dot(a_np.flatten(), b_np.flatten())
+
+        test.assertAlmostEqual(ab / N, ref_ab / N, places=5)
+
+    return test_array_inner
+
+
+def make_test_array_inner_axis(dtype):
+    I = 5
+    J = 10
+    K = 2
+
+    N = I * J * K
+
+    def test_array_inner(test, device):
+        rng = np.random.default_rng(123)
+
+        a_np = rng.random(size=(I, J, K))
+        b_np = rng.random(size=(I, J, K))
+
+        a = wp.array(a_np, shape=(I, J, K), device=device, dtype=dtype)
+        b = wp.array(b_np, shape=(I, J, K), device=device, dtype=dtype)
+
+        ab = array_inner(a, b, axis=0)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [1, 2])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+        ab = array_inner(a, b, axis=1)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [0, 2])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+        ab = array_inner(a, b, axis=2)
+        ref_ab = np.einsum(a_np, [0, 1, 2], b_np, [0, 1, 2], [0, 1])
+        assert_np_equal(ab.numpy() / N, ref_ab / N, 0.0001)
+
+    return test_array_inner
+
+
+def test_array_inner_empty(test, device):
+    values = wp.array([], device=device, dtype=wp.vec2)
+    test.assertEqual(array_inner(values, values), 0.0)
+
+    values = wp.array([], shape=(0, 3), device=device, dtype=float)
+    assert_np_equal(array_inner(values, values, axis=0).numpy(), np.zeros(3))
+
+
+devices = get_test_devices()
+
+
+class TestArrayReduce(unittest.TestCase):
+    pass
+
+
+add_function_test(TestArrayReduce, "test_array_sum_double", make_test_array_sum(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_vec3", make_test_array_sum(wp.vec3), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_axis_float", make_test_array_sum_axis(wp.float32), devices=devices)
+add_function_test(TestArrayReduce, "test_array_sum_empty", test_array_sum_empty, devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_double", make_test_array_inner(wp.float64), devices=devices)
+add_function_test(TestArrayReduce, "test_array_inner_vec3", make_test_array_inner(wp.vec3), devices=devices)
+add_function_test(
+    TestArrayReduce, "test_array_inner_axis_float", make_test_array_inner_axis(wp.float32), devices=devices
+)
+add_function_test(TestArrayReduce, "test_array_inner_empty", test_array_inner_empty, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)

warp/tests/test_atomic.py

@@ -5,14 +5,12 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-# include parent path
+import unittest
+
 import numpy as np
-import math
 
 import warp as wp
-from warp.tests.test_base import *
-
-import unittest
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -34,8 +32,7 @@ def make_atomic_test(type):
     # register a custom kernel (no decorator) function
     # this lets us register the same function definition
     # against multiple symbols, with different arg types
-    module = wp.get_module(test_atomic_kernel.__module__)
-    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel", module=module)
+    kernel = wp.Kernel(func=test_atomic_kernel, key=f"test_atomic_{type.__name__}_kernel")
 
     def test_atomic(test, device):
         n = 1024
@@ -54,20 +51,60 @@ def make_atomic_test(type):
             base = rng.random(size=(1, *type._shape_), dtype=float)
             val = rng.random(size=(n, *type._shape_), dtype=float)
 
-        add_array = wp.array(base, dtype=type, device=device)
-        min_array = wp.array(base, dtype=type, device=device)
-        max_array = wp.array(base, dtype=type, device=device)
+        add_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        min_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        max_array = wp.array(base, dtype=type, device=device, requires_grad=True)
+        add_array.zero_()
+        min_array.fill_(10000)
+        max_array.fill_(-10000)
 
-        val_array = wp.array(val, dtype=type, device=device)
+        val_array = wp.array(val, dtype=type, device=device, requires_grad=True)
 
-        wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
-
-        val = np.append(val, [base[0]], axis=0)
+        tape = wp.Tape()
+        with tape:
+            wp.launch(kernel, n, inputs=[add_array, min_array, max_array, val_array], device=device)
 
         assert_np_equal(add_array.numpy(), np.sum(val, axis=0), tol=1.0e-2)
         assert_np_equal(min_array.numpy(), np.min(val, axis=0), tol=1.0e-2)
         assert_np_equal(max_array.numpy(), np.max(val, axis=0), tol=1.0e-2)
 
+        if type != wp.int32:
+            add_array.grad.fill_(1)
+            tape.backward()
+            assert_np_equal(val_array.grad.numpy(), np.ones_like(val))
+            tape.zero()
+
+            min_array.grad.fill_(1)
+            tape.backward()
+            min_grad_array = np.zeros_like(val)
+            argmin = val.argmin(axis=0)
+            if val.ndim == 1:
+                min_grad_array[argmin] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    min_grad_array[argmin[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        min_grad_array[argmin[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), min_grad_array)
+            tape.zero()
+
+            max_array.grad.fill_(1)
+            tape.backward()
+            max_grad_array = np.zeros_like(val)
+            argmax = val.argmax(axis=0)
+            if val.ndim == 1:
+                max_grad_array[argmax] = 1
+            elif val.ndim == 2:
+                for i in range(val.shape[1]):
+                    max_grad_array[argmax[i], i] = 1
+            elif val.ndim == 3:
+                for i in range(val.shape[1]):
+                    for j in range(val.shape[2]):
+                        max_grad_array[argmax[i, j], i, j] = 1
+            assert_np_equal(val_array.grad.numpy(), max_grad_array)
+
     return test_atomic
 
 
@@ -82,24 +119,23 @@ test_atomic_mat33 = make_atomic_test(wp.mat33)
 test_atomic_mat44 = make_atomic_test(wp.mat44)
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
 
-    class TestAtomic(parent):
-        pass
+class TestAtomic(unittest.TestCase):
+    pass
 
-    add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
-    add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
 
-    return TestAtomic
+add_function_test(TestAtomic, "test_atomic_int", test_atomic_int, devices=devices)
+add_function_test(TestAtomic, "test_atomic_float", test_atomic_float, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec2", test_atomic_vec2, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec3", test_atomic_vec3, devices=devices)
+add_function_test(TestAtomic, "test_atomic_vec4", test_atomic_vec4, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat22", test_atomic_mat22, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat33", test_atomic_mat33, devices=devices)
+add_function_test(TestAtomic, "test_atomic_mat44", test_atomic_mat44, devices=devices)
 
 
 if __name__ == "__main__":
-    c = register(unittest.TestCase)
+    wp.build.clear_kernel_cache()
     unittest.main(verbosity=2)

warp/tests/test_bool.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2023 NVIDIA CORPORATION.  All rights reserved.
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import unittest
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+wp.init()
+
+TRUE_CONSTANT = wp.constant(True)
+
+
+@wp.func
+def identity_function(input_bool: wp.bool, plain_bool: bool):
+    return input_bool and plain_bool
+
+
+@wp.kernel
+def identity_test(data: wp.array(dtype=wp.bool)):
+    i = wp.tid()
+
+    data[i] = data[i] and True
+    data[i] = data[i] and wp.bool(True)
+    data[i] = data[i] and not False
+    data[i] = data[i] and not wp.bool(False)
+    data[i] = identity_function(data[i], True)
+
+    if data[i]:
+        data[i] = True
+    else:
+        data[i] = False
+
+    if not data[i]:
+        data[i] = False
+    else:
+        data[i] = True
+
+    if data[i] and True:
+        data[i] = True
+    else:
+        data[i] = False
+
+    if data[i] or False:
+        data[i] = True
+    else:
+        data[i] = False
+
+
+def test_bool_identity_ops(test, device):
+    rng = np.random.default_rng(123)
+
+    dim_x = 10
+
+    rand_np = rng.random(dim_x) > 0.5
+
+    data_array = wp.array(data=rand_np, device=device)
+
+    test.assertEqual(data_array.dtype, wp.bool)
+
+    wp.launch(identity_test, dim=data_array.shape, inputs=[data_array], device=device)
+
+    assert_np_equal(data_array.numpy(), rand_np)
+
+
+@wp.kernel
+def check_compile_constant(result: wp.array(dtype=wp.bool)):
+    if TRUE_CONSTANT:
+        result[0] = TRUE_CONSTANT
+    else:
+        result[0] = False
+
+
+def test_bool_constant(test, device):
+    compile_constant_value = wp.zeros(1, dtype=wp.bool, device=device)
+    wp.launch(check_compile_constant, 1, inputs=[compile_constant_value], device=device)
+    test.assertTrue(compile_constant_value.numpy()[0])
+
+
+devices = get_test_devices()
+
+
+class TestBool(unittest.TestCase):
+    pass
+
+
+add_function_test(TestBool, "test_bool_identity_ops", test_bool_identity_ops, devices=devices)
+add_function_test(TestBool, "test_bool_constant", test_bool_constant, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2)