warp-lang 1.7.2rc1__py3-none-win_amd64.whl → 1.8.1__py3-none-win_amd64.whl

warp/tests/tile/test_tile_matmul.py

@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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
+from typing import Any
+
+import numpy as np
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+TILE_M = wp.constant(8)
+TILE_N = wp.constant(4)
+TILE_K = wp.constant(8)
+
+# num threads per-tile
+TILE_DIM = 64
+
+
+@wp.kernel
+def tile_grouped_gemm(A: wp.array3d(dtype=float), B: wp.array3d(dtype=float), C: wp.array3d(dtype=float)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(A[i], shape=(TILE_M, TILE_K))
+    b = wp.tile_load(B[i], shape=(TILE_K, TILE_N))
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=wp.float32)
+
+    wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C[i], sum)
+
+
+def test_tile_grouped_gemm(test, device):
+    batch_count = 56
+
+    M = TILE_M
+    N = TILE_N
+    K = TILE_K
+
+    rng = np.random.default_rng(42)
+    A = rng.random((batch_count, M, K), dtype=np.float32)
+    B = rng.random((batch_count, K, N), dtype=np.float32)
+    C = A @ B
+
+    A_wp = wp.array(A, requires_grad=True, device=device)
+    B_wp = wp.array(B, requires_grad=True, device=device)
+    C_wp = wp.zeros((batch_count, TILE_M, TILE_N), requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_grouped_gemm, dim=[batch_count], inputs=[A_wp, B_wp, C_wp], block_dim=TILE_DIM, device=device
+        )
+
+    # TODO: 32 mismatched elements
+    assert_np_equal(C_wp.numpy(), C, 1e-6)
+
+
+@wp.kernel
+def tile_gemm(A: wp.array2d(dtype=Any), B: wp.array2d(dtype=Any), C: wp.array2d(dtype=Any)):
+    # output tile index
+    i, j = wp.tid()
+
+    sum = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=A.dtype)
+
+    M = A.shape[0]
+    N = B.shape[1]
+    K = A.shape[1]
+
+    count = int(K / TILE_K)
+
+    for k in range(0, count):
+        a = wp.tile_load(A, shape=(TILE_M, TILE_K), offset=(i * TILE_M, k * TILE_K))
+        b = wp.tile_load(B, shape=(TILE_K, TILE_N), offset=(k * TILE_K, j * TILE_N))
+
+        # sum += a*b
+        wp.tile_matmul(a, b, sum)
+
+    wp.tile_store(C, sum, offset=(i * TILE_M, j * TILE_N))
+
+
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float16), "B": wp.array2d(dtype=wp.float16), "C": wp.array2d(dtype=wp.float16)}
+)
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float32), "B": wp.array2d(dtype=wp.float32), "C": wp.array2d(dtype=wp.float32)}
+)
+wp.overload(
+    tile_gemm, {"A": wp.array2d(dtype=wp.float64), "B": wp.array2d(dtype=wp.float64), "C": wp.array2d(dtype=wp.float64)}
+)
+
+
+def test_tile_gemm(dtype):
+    def test(test, device):
+        M = TILE_M * 7
+        K = TILE_K * 6
+        N = TILE_N * 5
+
+        rng = np.random.default_rng(42)
+        A = rng.random((M, K), dtype=float).astype(wp.dtype_to_numpy(dtype))
+        B = rng.random((K, N), dtype=float).astype(wp.dtype_to_numpy(dtype))
+        C = np.zeros((M, N), dtype=float).astype(wp.dtype_to_numpy(dtype))
+
+        A_wp = wp.array(A, requires_grad=True, device=device)
+        B_wp = wp.array(B, requires_grad=True, device=device)
+        C_wp = wp.array(C, requires_grad=True, device=device)
+
+        with wp.Tape() as tape:
+            wp.launch_tiled(
+                tile_gemm,
+                dim=(int(M / TILE_M), int(N / TILE_N)),
+                inputs=[A_wp, B_wp, C_wp],
+                block_dim=TILE_DIM,
+                device=device,
+            )
+
+        assert_np_equal(C_wp.numpy(), A @ B, tol=1.0e-1)
+
+        adj_C = np.ones_like(C)
+
+        tape.backward(grads={C_wp: wp.array(adj_C, device=device)})
+
+        assert_np_equal(A_wp.grad.numpy(), adj_C @ B.T, tol=1.0e-1)
+        assert_np_equal(B_wp.grad.numpy(), A.T @ adj_C, 1.0e-1)
+
+    return test
+
+
+@wp.kernel
+def test_tile_transpose_matmul_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+    x = wp.tile_load(input, shape=(TILE_M, TILE_N))
+    y = wp.tile_transpose(x)
+
+    z = wp.tile_zeros(dtype=float, shape=(TILE_N, TILE_N))
+    wp.tile_matmul(y, x, z)
+
+    wp.tile_store(output, z)
+
+
+def test_tile_transpose_matmul(test, device):
+    rng = np.random.default_rng(42)
+    input = wp.array(rng.random((TILE_M, TILE_N), dtype=np.float32), device=device)
+    output = wp.zeros((TILE_N, TILE_N), dtype=float, device=device)
+
+    wp.launch_tiled(
+        test_tile_transpose_matmul_kernel, dim=[1], inputs=[input, output], block_dim=TILE_DIM, device=device
+    )
+
+    assert_np_equal(output.numpy(), input.numpy().T @ input.numpy())
+
+
+class TestTileMatmul(unittest.TestCase):
+    pass
+
+
+devices = get_test_devices()
+
+add_function_test(TestTileMatmul, "test_tile_gemm_fp16", test_tile_gemm(wp.float16), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_gemm_fp32", test_tile_gemm(wp.float32), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_gemm_fp64", test_tile_gemm(wp.float64), devices=devices)
+add_function_test(TestTileMatmul, "test_tile_grouped_gemm", test_tile_grouped_gemm, devices=devices)
+add_function_test(TestTileMatmul, "test_tile_transpose_matmul", test_tile_transpose_matmul, devices=devices)
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/tile/test_tile_reduce.py

@@ -73,6 +73,46 @@ def test_tile_reduce_sum(test, device):
     assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
 
 
+@wp.kernel
+def tile_sum_to_shared_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
+    i, lane = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    s = wp.tile_sum(a)
+    v = s[0]  # force shared storage for s
+    wp.tile_store(output, s * 0.5, offset=i)
+
+
+def test_tile_sum_to_shared(test, device):
+    batch_count = 1
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, TILE_DIM), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device, dtype=float)
+    output_wp = wp.zeros(batch_count, requires_grad=True, device=device, dtype=float)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_sum_to_shared_kernel,
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    sum_wp = output_wp.numpy()
+    for i in range(batch_count):
+        sum_np = np.sum(input[i], axis=0) * 0.5
+        assert_np_equal(sum_wp[i], sum_np, tol=0.0001)
+
+    output_wp.grad.fill_(1.0)
+
+    tape.backward()
+
+    assert_np_equal(input_wp.grad.numpy(), np.ones_like(input) * 0.5, tol=1.0e-4)
+
+
 @wp.kernel
 def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -84,6 +124,13 @@ def tile_min_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float
     wp.tile_store(output, m, offset=i)
 
 
+@wp.kernel
+def tile_min_kernel_edge_case(x: wp.array2d(dtype=float), y: wp.array(dtype=float)):
+    t = wp.tile_load(x, shape=(3, 3))
+    min = wp.tile_min(t)
+    wp.tile_store(y, min)
+
+
 def test_tile_reduce_min(test, device):
     batch_count = 56
 
@@ -105,6 +152,62 @@ def test_tile_reduce_min(test, device):
         min_np = np.min(input[i])
         test.assertAlmostEqual(min_wp[i], min_np, places=4)
 
+    # test edge case: tile is multiple warps in size but at least one is empty
+    x = wp.array(np.array([[2.0, 2.0, 3.0], [4.0, 1.0, 6.0], [7.0, 3.0, 9.0]]), dtype=float, device=device)
+    y = wp.zeros(1, dtype=float, device=device)
+
+    wp.launch_tiled(tile_min_kernel_edge_case, dim=1, inputs=[x, y], block_dim=64, device=device)
+
+    assert_np_equal(y.numpy(), np.array([1.0]))
+
+
+@wp.kernel
+def tile_argmin_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=int)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    m = wp.tile_argmin(a)
+
+    wp.tile_store(output, m, offset=i)
+
+
+@wp.kernel
+def tile_argmin_kernel_edge_case(x: wp.array2d(dtype=float), y: wp.array(dtype=int)):
+    t = wp.tile_load(x, shape=(3, 3))
+    min = wp.tile_argmin(t)
+    wp.tile_store(y, min)
+
+
+def test_tile_reduce_argmin(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, dtype=wp.int32, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_argmin_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    argmin_wp = output_wp.numpy()
+    for i in range(batch_count):
+        argmin_np = np.argmin(input[i])
+        test.assertAlmostEqual(argmin_wp[i], argmin_np, places=4)
+
+    # test edge case: tile is multiple warps in size but at least one is empty
+    x = wp.array(np.array([[2.0, 2.0, 3.0], [4.0, 1.0, 6.0], [7.0, 3.0, 9.0]]), dtype=float, device=device)
+    y = wp.zeros(1, dtype=int, device=device)
+
+    wp.launch_tiled(tile_argmin_kernel_edge_case, dim=1, inputs=[x, y], block_dim=64, device=device)
+
+    assert_np_equal(y.numpy(), np.array([4]))
+
 
 @wp.kernel
 def tile_max_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
@@ -139,6 +242,39 @@ def test_tile_reduce_max(test, device):
         test.assertAlmostEqual(max_wp[i], max_np, places=4)
 
 
+@wp.kernel
+def tile_argmax_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=int)):
+    # output tile index
+    i = wp.tid()
+
+    a = wp.tile_load(input[i], shape=TILE_DIM)
+    m = wp.tile_argmax(a)
+
+    wp.tile_store(output, m, offset=i)
+
+
+def test_tile_reduce_argmax(test, device):
+    batch_count = 56
+
+    N = TILE_DIM
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array(input, requires_grad=True, device=device)
+    output_wp = wp.zeros(batch_count, dtype=wp.int32, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            tile_argmax_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+        )
+
+    argmax_wp = output_wp.numpy()
+    for i in range(batch_count):
+        argmax_np = np.argmax(input[i])
+        test.assertAlmostEqual(argmax_wp[i], argmax_np, places=4)
+
+
 @wp.kernel
 def tile_reduce_custom_kernel(input: wp.array2d(dtype=float), output: wp.array(dtype=float)):
     # output tile index
@@ -176,6 +312,79 @@ def test_tile_reduce_custom(test, device):
         test.assertAlmostEqual(prod_wp[i], prod_np, places=4)
 
 
+def create_tile_scan_inclusive_kernel(tile_dim: int):
+    @wp.kernel(module="unique")
+    def tile_scan_inclusive_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+        i = wp.tid()
+        t = wp.tile_load(input[i], shape=tile_dim)
+        t = wp.tile_scan_inclusive(t)
+        wp.tile_store(output[i], t)
+
+    return tile_scan_inclusive_kernel
+
+
+def test_tile_scan_inclusive(test, device):
+    batch_count = 56
+    N = 1234
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array2d(input, requires_grad=True, device=device)
+    output_wp = wp.zeros_like(input_wp, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            create_tile_scan_inclusive_kernel(N),
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    scan_wp = output_wp.numpy()
+    for i in range(batch_count):
+        scan_np = np.cumsum(input[i])
+        np.testing.assert_allclose(scan_wp[i], scan_np, rtol=1e-5, atol=1e-6)
+
+
+def create_tile_scan_exclusive_kernel(tile_dim: int):
+    @wp.kernel(module="unique")
+    def tile_scan_exclusive_kernel(input: wp.array2d(dtype=float), output: wp.array2d(dtype=float)):
+        i = wp.tid()
+        t = wp.tile_load(input[i], shape=tile_dim)
+        t = wp.tile_scan_exclusive(t)
+        wp.tile_store(output[i], t)
+
+    return tile_scan_exclusive_kernel
+
+
+def test_tile_scan_exclusive(test, device):
+    batch_count = 56
+    N = 1234
+
+    rng = np.random.default_rng(42)
+    input = rng.random((batch_count, N), dtype=np.float32)
+
+    input_wp = wp.array2d(input, requires_grad=True, device=device)
+    output_wp = wp.zeros_like(input_wp, requires_grad=True, device=device)
+
+    with wp.Tape() as tape:
+        wp.launch_tiled(
+            create_tile_scan_exclusive_kernel(N),
+            dim=[batch_count],
+            inputs=[input_wp, output_wp],
+            block_dim=TILE_DIM,
+            device=device,
+        )
+
+    scan_wp = output_wp.numpy()
+    for i in range(batch_count):
+        scan_np = np.zeros(N, dtype=np.float32)
+        scan_np[1:] = np.cumsum(input[i][:-1])
+        np.testing.assert_allclose(scan_wp[i], scan_np, rtol=1e-5, atol=1e-6)
+
+
 @wp.struct
 class KeyValue:
     key: wp.int32
@@ -259,7 +468,7 @@ def test_tile_reduce_grouped_sum(test, device):
 
     with wp.Tape() as tape:
         wp.launch_tiled(
-            tile_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
+            tile_grouped_sum_kernel, dim=[batch_count], inputs=[input_wp, output_wp], block_dim=TILE_DIM, device=device
         )
 
     sum_wp = output_wp.numpy()
@@ -359,17 +568,17 @@ def test_untile_vector_kernel(input: wp.array(dtype=wp.vec3), output: wp.array(d
 
 
 def test_tile_untile_vector(test, device):
-    input = wp.full(16, wp.vec3(1.0, 2.0, 3.0), requires_grad=True, device=device)
+    input = wp.full(TILE_DIM, wp.vec3(1.0, 2.0, 3.0), requires_grad=True, device=device)
     output = wp.zeros_like(input, device=device)
 
     with wp.Tape() as tape:
-        wp.launch(test_untile_vector_kernel, dim=16, inputs=[input, output], block_dim=16, device=device)
+        wp.launch(test_untile_vector_kernel, dim=TILE_DIM, inputs=[input, output], block_dim=TILE_DIM, device=device)
 
     output.grad = wp.ones_like(output, device=device)
     tape.backward()
 
     assert_np_equal(output.numpy(), input.numpy())
-    assert_np_equal(input.grad.numpy(), np.ones((16, 3)))
+    assert_np_equal(input.grad.numpy(), np.ones((TILE_DIM, 3)))
 
 
 @wp.kernel
@@ -423,6 +632,91 @@ def test_tile_arange(test, device):
     assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
 
 
+@wp.kernel(module="unique")
+def tile_strided_loop_kernel(arr: wp.array(dtype=float), max_val: wp.array(dtype=float)):
+    tid, lane = wp.tid()
+
+    num_threads = wp.block_dim()
+
+    thread_max = wp.float32(-wp.inf)
+
+    length = arr.shape[0]
+    upper = ((length + num_threads - 1) // num_threads) * num_threads
+    for el_id in range(lane, upper, num_threads):
+        if el_id < length:
+            val = arr[el_id]
+        else:
+            val = wp.float32(-wp.inf)
+
+        t = wp.tile(val)
+        local_max = wp.tile_max(t)
+
+        thread_max = wp.max(thread_max, local_max[0])
+
+    if lane == 0:
+        max_val[0] = thread_max
+
+
+def test_tile_strided_loop(test, device):
+    N = 5  # Length of array
+
+    rng = np.random.default_rng(42)
+    input = rng.random(N, dtype=np.float32)
+
+    input_wp = wp.array(input, device=device)
+    output_wp = wp.zeros(1, dtype=wp.float32, device=device)
+
+    wp.launch_tiled(
+        tile_strided_loop_kernel,
+        dim=[1],
+        inputs=[input_wp, output_wp],
+        device=device,
+        block_dim=128,
+    )
+
+    max_wp = output_wp.numpy()
+    max_np = np.max(input)
+    test.assertAlmostEqual(max_wp[0], max_np, places=4)
+
+
+@wp.kernel
+def test_tile_reduce_matrix_kernel(y: wp.array(dtype=wp.mat33)):
+    i = wp.tid()
+    I = wp.identity(3, dtype=wp.float32)
+    m = wp.float32(i) * I
+
+    t = wp.tile(m, preserve_type=True)
+    sum = wp.tile_reduce(wp.add, t)
+
+    wp.tile_atomic_add(y, sum)
+
+
+def test_tile_reduce_matrix(test, device):
+    y = wp.zeros(shape=1, dtype=wp.mat33, device=device)
+
+    wp.launch(test_tile_reduce_matrix_kernel, dim=TILE_DIM, inputs=[], outputs=[y], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(y.numpy().squeeze(), np.eye(3, dtype=np.float32) * 2016.0)
+
+
+@wp.kernel
+def test_tile_reduce_vector_kernel(out: wp.array(dtype=wp.vec3)):
+    v = wp.vec3f(1.0)
+    v_tile = wp.tile(v, preserve_type=True)
+
+    sum = wp.tile_reduce(wp.add, v_tile)
+
+    wp.tile_atomic_add(out, sum)
+
+
+def test_tile_reduce_vector(test, device):
+    out = wp.zeros(1, dtype=wp.vec3, device=device)
+
+    wp.launch(kernel=test_tile_reduce_vector_kernel, dim=8, inputs=[], outputs=[out], block_dim=TILE_DIM, device=device)
+
+    assert_np_equal(out.numpy(), np.array([[8.0, 8.0, 8.0]]))
+
+
 devices = get_test_devices()
 
 
@@ -431,16 +725,24 @@ class TestTileReduce(unittest.TestCase):
 
 
 add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum, devices=devices)
+add_function_test(TestTileReduce, "test_tile_sum_to_shared", test_tile_sum_to_shared, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_min", test_tile_reduce_min, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_max", test_tile_reduce_max, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_argmin", test_tile_reduce_argmin, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_argmax", test_tile_reduce_argmax, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_custom", test_tile_reduce_custom, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_custom_struct", test_tile_reduce_custom_struct, devices=devices)
-add_function_test(TestTileReduce, "test_tile_reduce_grouped_sum", test_tile_reduce_sum, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_grouped_sum", test_tile_reduce_grouped_sum, devices=devices)
 add_function_test(TestTileReduce, "test_tile_reduce_simt", test_tile_reduce_simt, devices=devices)
 add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devices)
 add_function_test(TestTileReduce, "test_tile_arange", test_tile_arange, devices=devices)
 add_function_test(TestTileReduce, "test_tile_untile_scalar", test_tile_untile_scalar, devices=devices)
 add_function_test(TestTileReduce, "test_tile_untile_vector", test_tile_untile_vector, devices=devices)
+add_function_test(TestTileReduce, "test_tile_strided_loop", test_tile_strided_loop, devices=devices)
+add_function_test(TestTileReduce, "test_tile_scan_inclusive", test_tile_scan_inclusive, devices=devices)
+add_function_test(TestTileReduce, "test_tile_scan_exclusive", test_tile_scan_exclusive, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_matrix", test_tile_reduce_matrix, devices=devices)
+add_function_test(TestTileReduce, "test_tile_reduce_vector", test_tile_reduce_vector, devices=devices)
 
 if __name__ == "__main__":
     wp.clear_kernel_cache()