warp-lang 1.0.0b5__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/tests/test_math.py

@@ -5,13 +5,13 @@
 # distribution of this software and related documentation without an express
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
-from typing import NamedTuple
 import unittest
+from typing import NamedTuple
 
 import numpy as np
 
 import warp as wp
-from warp.tests.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
@@ -176,19 +176,18 @@ def test_mat_type(test, device):
         raise ValueError("mat to string error")
 
 
-def register(parent):
-    devices = get_test_devices()
+devices = get_test_devices()
+
+
+class TestMath(unittest.TestCase):
+    pass
 
-    class TestMath(parent):
-        pass
 
-    add_function_test(TestMath, "test_scalar_math", test_scalar_math, devices=devices)
-    add_function_test(TestMath, "test_vec_type", test_vec_type, devices=devices)
-    add_function_test(TestMath, "test_mat_type", test_mat_type, devices=devices)
-    return TestMath
+add_function_test(TestMath, "test_scalar_math", test_scalar_math, devices=devices)
+add_function_test(TestMath, "test_vec_type", test_vec_type, devices=devices)
+add_function_test(TestMath, "test_mat_type", test_mat_type, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2)

warp/tests/test_matmul.py

@@ -1,11 +1,21 @@
-import numpy as np
+# Copyright (c) 2022 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.test_base import *
+from warp.tests.unittest_utils import *
 
 wp.init()
 
+from warp.context import runtime  # noqa: E402
+
 
 class gemm_test_bed_runner:
     def __init__(self, dtype, device):
@@ -21,63 +31,54 @@ class gemm_test_bed_runner:
                 np.ceil(rng.uniform(low=low, high=high, size=(m, k))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             B = wp.array2d(
                 np.ceil(rng.uniform(low=low, high=high, size=(k, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             C = wp.array2d(
                 np.ceil(rng.uniform(low=low, high=high, size=(m, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
-            D = wp.array2d(
-                np.zeros((m, n)),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True)
+            D = wp.array2d(np.zeros((m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
         else:
             A = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             B = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             C = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
-            )
-            D = wp.array3d(
-                np.zeros((batch_count, m, n)),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
+            D = wp.array3d(np.zeros((batch_count, m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
         return A, B, C, D
 
     def run_and_verify(self, m, n, k, batch_count, alpha, beta):
         A, B, C, D = self.alloc(m, n, k, batch_count)
         ones = wp.zeros_like(D)
         ones.fill_(1.0)
-        
+
         if batch_count == 1:
             tape = wp.Tape()
             with tape:
                 wp.matmul(A, B, C, D, alpha, beta, False, self.device)
-            tape.backward(grads={D : ones})
-            
+            tape.backward(grads={D: ones})
+
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C.numpy()
             assert np.array_equal(D_np, D.numpy())
 
@@ -89,8 +90,8 @@ class gemm_test_bed_runner:
             tape = wp.Tape()
             with tape:
                 wp.batched_matmul(A, B, C, D, alpha, beta, False, self.device)
-            tape.backward(grads={D : ones})
-            
+            tape.backward(grads={D: ones})
+
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
             assert np.array_equal(D_np, D.numpy())
 
@@ -132,75 +133,45 @@ class gemm_test_bed_runner_transpose:
                 np.ceil(rng.uniform(low=low, high=high, size=(m, k))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             B = wp.array2d(
                 np.ceil(rng.uniform(low=low, high=high, size=(k, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             C = wp.array2d(
                 np.ceil(rng.uniform(low=low, high=high, size=(m, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
-            )
-            D = wp.array2d(
-                np.zeros((m, n)),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
-            )
-            AT = wp.array2d(
-                A.numpy().transpose([1, 0]),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
-            )
-            BT = wp.array2d(
-                B.numpy().transpose([1, 0]),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
+            D = wp.array2d(np.zeros((m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+            AT = wp.array2d(A.numpy().transpose([1, 0]), dtype=self.dtype, device=self.device, requires_grad=True)
+            BT = wp.array2d(B.numpy().transpose([1, 0]), dtype=self.dtype, device=self.device, requires_grad=True)
         else:
             A = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             B = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
             C = wp.array3d(
                 np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
                 dtype=self.dtype,
                 device=self.device,
-                requires_grad=True
-            )
-            D = wp.array3d(
-                np.zeros((batch_count, m, n)),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
-            )
-            AT = wp.array3d(
-                A.numpy().transpose([0, 2, 1]),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
-            )
-            BT = wp.array3d(
-                B.numpy().transpose([0, 2, 1]),
-                dtype=self.dtype,
-                device=self.device,
-                requires_grad=True
+                requires_grad=True,
             )
+            D = wp.array3d(np.zeros((batch_count, m, n)), dtype=self.dtype, device=self.device, requires_grad=True)
+            AT = wp.array3d(A.numpy().transpose([0, 2, 1]), dtype=self.dtype, device=self.device, requires_grad=True)
+            BT = wp.array3d(B.numpy().transpose([0, 2, 1]), dtype=self.dtype, device=self.device, requires_grad=True)
         return A, B, C, D, AT, BT
 
     def run_and_verify(self, m, n, k, batch_count, alpha, beta):
@@ -219,17 +190,17 @@ class gemm_test_bed_runner_transpose:
         ones3.fill_(1.0)
 
         if batch_count == 1:
-            ATT1 = AT1.transpose([1, 0])        
+            ATT1 = AT1.transpose([1, 0])
             BTT1 = BT1.transpose([1, 0])
-            ATT2 = AT2.transpose([1, 0])        
+            ATT2 = AT2.transpose([1, 0])
             BTT2 = BT2.transpose([1, 0])
             tape = wp.Tape()
             with tape:
                 wp.matmul(A, BTT1, C1, D1, alpha, beta, False, self.device)
                 wp.matmul(ATT1, B, C2, D2, alpha, beta, False, self.device)
                 wp.matmul(ATT2, BTT2, C3, D3, alpha, beta, False, self.device)
-            tape.backward(grads={D1 : ones1, D2 : ones2, D3 : ones3})
-            
+            tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
+
             D_np = alpha * (A.numpy() @ B.numpy()) + beta * C1.numpy()
             assert np.array_equal(D_np, D1.numpy())
             assert np.array_equal(D_np, D2.numpy())
@@ -240,7 +211,7 @@ class gemm_test_bed_runner_transpose:
             adj_C_np = beta * ones1.numpy()
 
         else:
-            ATT1 = AT1.transpose([0, 2, 1])        
+            ATT1 = AT1.transpose([0, 2, 1])
             BTT1 = BT1.transpose([0, 2, 1])
             ATT2 = AT2.transpose([0, 2, 1])
             BTT2 = BT2.transpose([0, 2, 1])
@@ -249,8 +220,8 @@ class gemm_test_bed_runner_transpose:
                 wp.batched_matmul(A, BTT1, C1, D1, alpha, beta, False, self.device)
                 wp.batched_matmul(ATT1, B, C2, D2, alpha, beta, False, self.device)
                 wp.batched_matmul(ATT2, BTT2, C3, D3, alpha, beta, False, self.device)
-            tape.backward(grads={D1 : ones1, D2 : ones2, D3 : ones3})
-            
+            tape.backward(grads={D1: ones1, D2: ones2, D3: ones3})
+
             D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C1.numpy()
             assert np.array_equal(D_np, D1.numpy())
             assert np.array_equal(D_np, D2.numpy())
@@ -288,11 +259,13 @@ def test_f16(test, device):
     gemm_test_bed_runner_transpose(wp.float16, device).run()
 
 
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_f32(test, device):
     gemm_test_bed_runner(wp.float32, device).run()
     gemm_test_bed_runner_transpose(wp.float32, device).run()
 
 
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_f64(test, device):
     gemm_test_bed_runner(wp.float64, device).run()
     gemm_test_bed_runner_transpose(wp.float64, device).run()
@@ -304,6 +277,7 @@ def matrix_sum_kernel(arr: wp.array2d(dtype=float), loss: wp.array(dtype=float))
     wp.atomic_add(loss, 0, arr[i, j])
 
 
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_tape(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -331,6 +305,7 @@ def test_tape(test, device):
 
     tape.backward(loss=loss)
     A_grad = A.grad.numpy()
+    tape.reset()
 
     # test adjoint
     D.grad = wp.array2d(np.ones((m, n)), dtype=float, device=device)
@@ -342,6 +317,7 @@ def test_tape(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_operator(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -377,6 +353,7 @@ def test_operator(test, device):
     assert_array_equal(A.grad, wp.zeros_like(A))
 
 
+@unittest.skipUnless(runtime.core.is_cutlass_enabled(), "Warp was not built with CUTLASS support")
 def test_large_batch_count(test, device):
     rng = np.random.default_rng(42)
     low = -4.5
@@ -386,31 +363,38 @@ def test_large_batch_count(test, device):
     k = 4
     batch_count = 65535 * 2 + int(65535 / 2)
     A = wp.array3d(
-        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))), dtype=float, device=device, requires_grad=True
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, k))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
     )
     B = wp.array3d(
-        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))), dtype=float, device=device, requires_grad=True
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, k, n))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
     )
     C = wp.array3d(
-        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))), dtype=float, device=device, requires_grad=True
-    )
-    D = wp.array3d(
-        np.zeros((batch_count, m, n)), dtype=float, device=device, requires_grad=True
+        np.ceil(rng.uniform(low=low, high=high, size=(batch_count, m, n))),
+        dtype=float,
+        device=device,
+        requires_grad=True,
     )
+    D = wp.array3d(np.zeros((batch_count, m, n)), dtype=float, device=device, requires_grad=True)
     ones = wp.zeros_like(D)
     ones.fill_(1.0)
 
     alpha = 1.0
     beta = 1.0
-    
+
     tape = wp.Tape()
     with tape:
         wp.batched_matmul(A, B, C, D, alpha=alpha, beta=beta, allow_tf32x3_arith=False, device=device)
-    tape.backward(grads={D : ones})
+    tape.backward(grads={D: ones})
 
     D_np = alpha * np.matmul(A.numpy(), B.numpy()) + beta * C.numpy()
     assert np.array_equal(D_np, D.numpy())
-    
+
     adj_A_np = alpha * np.matmul(ones.numpy(), B.numpy().transpose((0, 2, 1)))
     adj_B_np = alpha * np.matmul(A.numpy().transpose((0, 2, 1)), ones.numpy())
     adj_C_np = beta * ones.numpy()
@@ -420,30 +404,50 @@ def test_large_batch_count(test, device):
     assert np.array_equal(adj_C_np, C.grad.numpy())
 
 
-def register(parent):
-    devices = [d for d in get_test_devices()]
+def test_adjoint_accumulation(test, device):
+    a_np = np.ones(shape=(2,3))
+    b_np = np.ones(shape=(3,2))
+    c_np = np.zeros(shape=(2,2))
+    d_np = np.zeros(shape=(2,2))
 
-    class TestMatmul(parent):
-        pass
+    a_wp = wp.from_numpy(a_np, dtype=float, requires_grad=True)
+    b_wp = wp.from_numpy(b_np, dtype=float, requires_grad=True)
+    c_wp = wp.from_numpy(c_np, dtype=float, requires_grad=True)
+    d1_wp = wp.from_numpy(d_np, dtype=float, requires_grad=True)
+    d2_wp = wp.from_numpy(d_np, dtype=float, requires_grad=True)
 
-    if devices:
-        # check if CUTLASS is available
-        from warp.context import runtime
+    tape = wp.Tape()
+
+    with tape:
+        wp.matmul(a_wp, b_wp, c_wp, d1_wp, alpha=1.0, beta=1.0)
+        wp.matmul(a_wp, b_wp, d1_wp, d2_wp, alpha=1.0, beta=1.0)
+
+    d_grad = wp.zeros_like(d2_wp)
+    d_grad.fill_(1.)
+    grads = {d2_wp : d_grad}
+    tape.backward(grads=grads) 
+
+    assert np.array_equal(a_wp.grad.numpy(), 4.0 * np.ones(shape=(2,3)))
+    assert np.array_equal(b_wp.grad.numpy(), 4.0 * np.ones(shape=(3,2)))
+    assert np.array_equal(c_wp.grad.numpy(), np.ones(shape=(2,2)))
+
+
+devices = get_test_devices()
+
+
+class TestMatmul(unittest.TestCase):
+    pass
 
-        if runtime.core.is_cutlass_enabled():
-            # add_function_test(TestMatmul, "test_f16", test_f16, devices=devices)
-            add_function_test(TestMatmul, "test_f32", test_f32, devices=devices)
-            add_function_test(TestMatmul, "test_f64", test_f64, devices=devices)
-            add_function_test(TestMatmul, "test_tape", test_tape, devices=devices)
-            add_function_test(TestMatmul, "test_operator", test_operator, devices=devices)
-            add_function_test(TestMatmul, "test_large_batch_count", test_large_batch_count, devices=devices)
-        else:
-            print("Skipping matmul tests because CUTLASS is not supported in this build")
 
-    return TestMatmul
+# add_function_test(TestMatmul, "test_f16", test_f16, devices=devices)
+add_function_test(TestMatmul, "test_f32", test_f32, devices=devices)
+add_function_test(TestMatmul, "test_f64", test_f64, devices=devices)
+add_function_test(TestMatmul, "test_tape", test_tape, devices=devices)
+add_function_test(TestMatmul, "test_operator", test_operator, devices=devices)
+add_function_test(TestMatmul, "test_large_batch_count", test_large_batch_count, devices=devices)
+add_function_test(TestMatmul, "test_adjoint_accumulation", test_adjoint_accumulation, devices=devices)
 
 
 if __name__ == "__main__":
     wp.build.clear_kernel_cache()
-    _ = register(unittest.TestCase)
     unittest.main(verbosity=2, failfast=False)