warp-lang 0.11.0__py3-none-manylinux2014_x86_64.whl → 1.0.0__py3-none-manylinux2014_x86_64.whl

warp/tests/test_sim_kinematics.py

@@ -0,0 +1,97 @@
+# 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 warp as wp
+from warp.tests.unittest_utils import *
+
+import math
+import os
+
+import numpy as np
+
+import warp as wp
+import warp.sim
+
+wp.init()
+
+
+def test_fk_ik(test, device):
+
+    builder = wp.sim.ModelBuilder()
+
+    num_envs = 1
+
+    for i in range(num_envs):
+        wp.sim.parse_mjcf(
+            os.path.join(os.path.dirname(__file__), "../examples/assets/nv_ant.xml"),
+            builder,
+            stiffness=0.0,
+            damping=1.0,
+            armature=0.1,
+            contact_ke=1.0e4,
+            contact_kd=1.0e2,
+            contact_kf=1.0e2,
+            contact_mu=0.75,
+            limit_ke=1.0e3,
+            limit_kd=1.0e1,
+            up_axis="y",
+        )
+
+        coord_count = 15
+        dof_count = 14
+
+        coord_start = i * coord_count
+        dof_start = i * dof_count
+
+        # base
+        builder.joint_q[coord_start : coord_start + 3] = [i * 2.0, 0.70, 0.0]
+        builder.joint_q[coord_start + 3 : coord_start + 7] = wp.quat_from_axis_angle(
+            wp.vec3(1.0, 0.0, 0.0), -math.pi * 0.5
+        )
+
+        # joints
+        builder.joint_q[coord_start + 7 : coord_start + coord_count] = [0.0, 1.0, 0.0, -1.0, 0.0, -1.0, 0.0, 1.0]
+        builder.joint_qd[dof_start + 6 : dof_start + dof_count] = [1.0, 1.0, 1.0, -1.0, 1.0, -1.0, 1.0, 1.0]
+
+    # finalize model
+    model = builder.finalize()
+    model.ground = True
+    model.joint_attach_ke *= 16.0
+    model.joint_attach_kd *= 4.0
+
+    state = model.state()
+
+    # save a copy of joint values
+    q_fk = model.joint_q.numpy()
+    qd_fk = model.joint_qd.numpy()
+
+    wp.sim.eval_fk(model, model.joint_q, model.joint_qd, None, state)
+
+    q_ik = wp.zeros_like(model.joint_q)
+    qd_ik = wp.zeros_like(model.joint_qd)
+
+    wp.sim.eval_ik(model, state, q_ik, qd_ik)
+
+    assert_np_equal(q_fk, q_ik.numpy(), tol=1e-6)
+    assert_np_equal(qd_fk, qd_ik.numpy(), tol=1e-6)
+
+
+devices = get_test_devices()
+
+
+class TestSimKinematics(unittest.TestCase):
+    pass
+
+
+add_function_test(TestSimKinematics, "test_fk_ik", test_fk_ik, devices=devices)
+
+
+if __name__ == "__main__":
+    wp.build.clear_kernel_cache()
+    unittest.main(verbosity=2, failfast=True)

warp/tests/test_snippet.py

@@ -13,9 +13,9 @@ def test_basic(test, device):
     out[tid] = a * x[tid] + y[tid];
     """
     adj_snippet = """
-    adj_a = x[tid] * adj_out[tid];
-    adj_x[tid] = a * adj_out[tid];
-    adj_y[tid] = adj_out[tid];
+    adj_a += x[tid] * adj_out[tid];
+    adj_x[tid] += a * adj_out[tid];
+    adj_y[tid] += adj_out[tid];
     """
 
     @wp.func_native(snippet, adj_snippet)
@@ -86,6 +86,7 @@ def test_shared_memory(test, device):
 
     @wp.func_native(snippet)
     def reverse(d: wp.array(dtype=int), N: int, tid: int):
+        """Reverse the array d in place using shared memory."""
         return
 
     @wp.kernel
@@ -100,6 +101,7 @@ def test_shared_memory(test, device):
     wp.launch(kernel=reverse_kernel, dim=N, inputs=[x, N], device=device)
 
     assert_np_equal(x.numpy(), y)
+    assert reverse.__doc__ == "Reverse the array d in place using shared memory."
 
 
 def test_cpu_snippet(test, device):
@@ -130,6 +132,124 @@ def test_cpu_snippet(test, device):
     assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
 
 
+def test_custom_replay_grad(test, device):
+    num_threads = 16
+    counter = wp.zeros(1, dtype=wp.int32, device=device)
+    thread_ids = wp.zeros(num_threads, dtype=wp.int32, device=device)
+    inputs = wp.array(np.arange(num_threads, dtype=np.float32), device=device, requires_grad=True)
+    outputs = wp.zeros_like(inputs)
+
+    snippet = """
+        int next_index = atomicAdd(counter, 1);
+        thread_values[tid] = next_index;
+        """
+    replay_snippet = ""
+
+    @wp.func_native(snippet, replay_snippet=replay_snippet)
+    def reversible_increment(
+        counter: wp.array(dtype=int), thread_values: wp.array(dtype=int), tid: int
+    ):
+        ...
+
+    @wp.kernel
+    def run_atomic_add(
+        input: wp.array(dtype=float),
+        counter: wp.array(dtype=int),
+        thread_values: wp.array(dtype=int),
+        output: wp.array(dtype=float),
+    ):
+        tid = wp.tid()
+        reversible_increment(counter, thread_values, tid)
+        idx = thread_values[tid]
+        output[idx] = input[idx] ** 2.0
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(
+            run_atomic_add, dim=num_threads, inputs=[inputs, counter, thread_ids], outputs=[outputs], device=device
+        )
+
+    tape.backward(grads={outputs: wp.array(np.ones(num_threads, dtype=np.float32), device=device)})
+    assert_np_equal(inputs.grad.numpy(), 2.0 * inputs.numpy(), tol=1e-4)
+
+
+def test_replay_simplification(test, device):
+    num_threads = 8
+    x = wp.array(1.0 + np.arange(num_threads, dtype=np.float32), device=device, requires_grad=True)
+    y = wp.zeros_like(x)
+    z = wp.zeros_like(x)
+
+    snippet = "y[tid] = powf(x[tid], 2.0);"
+    replay_snippet = "y[tid] = x[tid];"
+    adj_snippet = "adj_x[tid] += 2.0 * adj_y[tid];"
+
+    @wp.func_native(snippet, adj_snippet=adj_snippet, replay_snippet=replay_snippet)
+    def square(x: wp.array(dtype=float), y: wp.array(dtype=float), tid: int):
+        ...
+
+    @wp.kernel
+    def log_square_kernel(
+        x: wp.array(dtype=float),
+        y: wp.array(dtype=float),
+        z: wp.array(dtype=float)
+    ):
+        tid = wp.tid()
+        square(x, y, tid)
+        z[tid] = wp.log(y[tid])
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(log_square_kernel, dim=num_threads, inputs=[x, y], outputs=[z], device=device)
+
+    tape.backward(grads={z: wp.array(np.ones(num_threads, dtype=np.float32), device=device)})
+    assert_np_equal(x.grad.numpy(), 2.0 / (1.0 + np.arange(num_threads)), tol=1e-6)
+
+
+def test_recompile_snippet(test, device):
+    snippet = """
+    int inc = 1;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):
+        ...
+
+    @wp.kernel
+    def increment(x: wp.array(dtype=wp.int32), out: wp.array(dtype=wp.int32)):
+        tid = wp.tid()
+        increment_snippet(x, out, tid)
+
+    N = 128
+    x = wp.array(np.arange(N, dtype=np.int32), dtype=wp.int32, device=device)
+    out = wp.zeros(N, dtype=wp.int32, device=device)
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), np.arange(1, N + 1, 1, dtype=np.int32))
+
+    snippet = """
+    int inc = 2;
+    out[tid] = x[tid] + inc;
+    """
+
+    @wp.func_native(snippet)
+    def increment_snippet(
+        x: wp.array(dtype=wp.int32),
+        out: wp.array(dtype=wp.int32),
+        tid: int,
+    ):
+        ...
+
+    wp.launch(kernel=increment, dim=N, inputs=[x], outputs=[out], device=device)
+
+    assert_np_equal(out.numpy(), 1 + np.arange(1, N + 1, 1, dtype=np.int32))
+
+
 class TestSnippets(unittest.TestCase):
     pass
 
@@ -137,6 +257,9 @@ class TestSnippets(unittest.TestCase):
 add_function_test(TestSnippets, "test_basic", test_basic, devices=get_unique_cuda_test_devices())
 add_function_test(TestSnippets, "test_shared_memory", test_shared_memory, devices=get_unique_cuda_test_devices())
 add_function_test(TestSnippets, "test_cpu_snippet", test_cpu_snippet, devices=["cpu"])
+add_function_test(TestSnippets, "test_custom_replay_grad", test_custom_replay_grad, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_replay_simplification", test_replay_simplification, devices=get_unique_cuda_test_devices())
+add_function_test(TestSnippets, "test_recompile_snippet", test_recompile_snippet, devices=get_unique_cuda_test_devices())
 
 
 if __name__ == "__main__":

warp/tests/test_streams.py

@@ -10,6 +10,7 @@ import unittest
 import numpy as np
 
 import warp as wp
+from warp.utils import check_iommu
 from warp.tests.unittest_utils import *
 
 wp.init()
@@ -37,19 +38,34 @@ def sum(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=fl
 N = 10 * 1024 * 1024
 
 
-def test_stream_arg_implicit_sync(test, device):
-    # wp.zeros() and array.numpy() should not require explicit sync
+def test_stream_set(test, device):
 
+    device = wp.get_device(device)
+
+    old_stream = device.stream
+    new_stream = wp.Stream(device)
+
+    try:
+        wp.set_stream(new_stream, device)
+
+        test.assertTrue(device.has_stream)
+        test.assertEqual(device.stream, new_stream)
+
+    finally:
+        # restore original stream
+        wp.set_stream(old_stream, device)
+
+
+def test_stream_arg_explicit_sync(test, device):
     a = wp.zeros(N, dtype=float, device=device)
-    b = wp.empty(N, dtype=float, device=device)
+    b = wp.full(N, 42, dtype=float, device=device)
     c = wp.empty(N, dtype=float, device=device)
 
+    old_stream = wp.get_stream(device)
     new_stream = wp.Stream(device)
 
-    # Exercise code path
-    wp.set_stream(new_stream, device)
-
-    test.assertTrue(wp.get_device(device).has_stream)
+    # allocations need to be explicitly synced before launching work using stream arguments
+    new_stream.wait_stream(old_stream)
 
     # launch work on new stream
     wp.launch(inc, dim=a.size, inputs=[a], stream=new_stream)
@@ -64,17 +80,17 @@ def test_stream_arg_implicit_sync(test, device):
 
 
 def test_stream_scope_implicit_sync(test, device):
-    # wp.zeros() and array.numpy() should not require explicit sync
 
     with wp.ScopedDevice(device):
         a = wp.zeros(N, dtype=float)
-        b = wp.empty(N, dtype=float)
+        b = wp.full(N, 42, dtype=float)
         c = wp.empty(N, dtype=float)
 
         old_stream = wp.get_stream()
         new_stream = wp.Stream()
 
         # launch work on new stream
+        # allocations are implicitly synced when entering wp.ScopedStream
         with wp.ScopedStream(new_stream):
             assert wp.get_stream() == new_stream
 
@@ -309,103 +325,116 @@ class TestStreams(unittest.TestCase):
             cpu_stream = cpu_device.stream  # noqa: F841
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
     def test_stream_arg_graph_mgpu(self):
         wp.load_module(device="cuda:0")
         wp.load_module(device="cuda:1")
 
-        # resources on GPU 0
-        stream0 = wp.get_stream("cuda:0")
-        a0 = wp.zeros(N, dtype=float, device="cuda:0")
-        b0 = wp.empty(N, dtype=float, device="cuda:0")
-        c0 = wp.empty(N, dtype=float, device="cuda:0")
+        # Peer-to-peer copies are not possible during graph capture if the arrays were
+        # allocated using pooled allocators and mempool access is not enabled.
+        # Here, we force default CUDA allocators and pre-allocate the memory.
+        with wp.ScopedMempool("cuda:0", False), wp.ScopedMempool("cuda:1", False):
 
-        # resources on GPU 1
-        stream1 = wp.get_stream("cuda:1")
-        a1 = wp.zeros(N, dtype=float, device="cuda:1")
+            # resources on GPU 0
+            stream0 = wp.get_stream("cuda:0")
+            a0 = wp.zeros(N, dtype=float, device="cuda:0")
+            b0 = wp.empty(N, dtype=float, device="cuda:0")
+            c0 = wp.empty(N, dtype=float, device="cuda:0")
 
-        # start recording on stream0
-        wp.capture_begin(stream=stream0, force_module_load=False)
-        try:
-            # branch into stream1
-            stream1.wait_stream(stream0)
+            # resources on GPU 1
+            stream1 = wp.get_stream("cuda:1")
+            a1 = wp.zeros(N, dtype=float, device="cuda:1")
 
-            # launch concurrent kernels on each stream
-            wp.launch(inc, dim=N, inputs=[a0], stream=stream0)
-            wp.launch(inc, dim=N, inputs=[a1], stream=stream1)
+            # start recording on stream0
+            wp.capture_begin(stream=stream0, force_module_load=False)
+            try:
+                # branch into stream1
+                stream1.wait_stream(stream0)
 
-            # wait for stream1 to finish
-            stream0.wait_stream(stream1)
+                # launch concurrent kernels on each stream
+                wp.launch(inc, dim=N, inputs=[a0], stream=stream0)
+                wp.launch(inc, dim=N, inputs=[a1], stream=stream1)
 
-            # copy values from stream1
-            wp.copy(b0, a1, stream=stream0)
+                # wait for stream1 to finish
+                stream0.wait_stream(stream1)
+
+                # copy values from stream1
+                wp.copy(b0, a1, stream=stream0)
 
-            # compute sum
-            wp.launch(sum, dim=N, inputs=[a0, b0, c0], stream=stream0)
-        finally:
-            # finish recording on stream0
-            g = wp.capture_end(stream=stream0)
+                # compute sum
+                wp.launch(sum, dim=N, inputs=[a0, b0, c0], stream=stream0)
+            finally:
+                # finish recording on stream0
+                g = wp.capture_end(stream=stream0)
 
-        # replay
-        num_iters = 10
-        for _ in range(num_iters):
-            wp.capture_launch(g, stream=stream0)
+            # replay
+            num_iters = 10
+            for _ in range(num_iters):
+                wp.capture_launch(g, stream=stream0)
 
-        # check results
-        assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
+            # check results
+            assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
     @unittest.skipUnless(len(wp.get_cuda_devices()) > 1, "Requires at least two CUDA devices")
+    @unittest.skipUnless(check_iommu(), "IOMMU seems enabled")
     def test_stream_scope_graph_mgpu(self):
         wp.load_module(device="cuda:0")
         wp.load_module(device="cuda:1")
 
-        # resources on GPU 0
-        with wp.ScopedDevice("cuda:0"):
-            stream0 = wp.get_stream()
-            a0 = wp.zeros(N, dtype=float)
-            b0 = wp.empty(N, dtype=float)
-            c0 = wp.empty(N, dtype=float)
-
-        # resources on GPU 1
-        with wp.ScopedDevice("cuda:1"):
-            stream1 = wp.get_stream()
-            a1 = wp.zeros(N, dtype=float)
-
-        # capture graph
-        with wp.ScopedDevice("cuda:0"):
-            # start recording
-            wp.capture_begin(force_module_load=False)
-            try:
-                with wp.ScopedDevice("cuda:1"):
-                    # branch into stream1
-                    wp.wait_stream(stream0)
+        # Peer-to-peer copies are not possible during graph capture if the arrays were
+        # allocated using pooled allocators and mempool access is not enabled.
+        # Here, we force default CUDA allocators and pre-allocate the memory.
+        with wp.ScopedMempool("cuda:0", False), wp.ScopedMempool("cuda:1", False):
 
-                    wp.launch(inc, dim=N, inputs=[a1])
+            # resources on GPU 0
+            with wp.ScopedDevice("cuda:0"):
+                stream0 = wp.get_stream()
+                a0 = wp.zeros(N, dtype=float)
+                b0 = wp.empty(N, dtype=float)
+                c0 = wp.empty(N, dtype=float)
 
-                wp.launch(inc, dim=N, inputs=[a0])
+            # resources on GPU 1
+            with wp.ScopedDevice("cuda:1"):
+                stream1 = wp.get_stream()
+                a1 = wp.zeros(N, dtype=float)
 
-                # wait for stream1 to finish
-                wp.wait_stream(stream1)
+            # capture graph
+            with wp.ScopedDevice("cuda:0"):
+                # start recording
+                wp.capture_begin(force_module_load=False)
+                try:
+                    with wp.ScopedDevice("cuda:1"):
+                        # branch into stream1
+                        wp.wait_stream(stream0)
 
-                # copy values from stream1
-                wp.copy(b0, a1)
+                        wp.launch(inc, dim=N, inputs=[a1])
 
-                # compute sum
-                wp.launch(sum, dim=N, inputs=[a0, b0, c0])
-            finally:
-                # finish recording
-                g = wp.capture_end()
+                    wp.launch(inc, dim=N, inputs=[a0])
 
-        # replay
-        with wp.ScopedDevice("cuda:0"):
-            num_iters = 10
-            for _ in range(num_iters):
-                wp.capture_launch(g)
+                    # wait for stream1 to finish
+                    wp.wait_stream(stream1)
+
+                    # copy values from stream1
+                    wp.copy(b0, a1)
+
+                    # compute sum
+                    wp.launch(sum, dim=N, inputs=[a0, b0, c0])
+                finally:
+                    # finish recording
+                    g = wp.capture_end()
+
+            # replay
+            with wp.ScopedDevice("cuda:0"):
+                num_iters = 10
+                for _ in range(num_iters):
+                    wp.capture_launch(g)
 
-        # check results
-        assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
+            # check results
+            assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
 
-add_function_test(TestStreams, "test_stream_arg_implicit_sync", test_stream_arg_implicit_sync, devices=devices)
+add_function_test(TestStreams, "test_stream_set", test_stream_set, devices=devices)
+add_function_test(TestStreams, "test_stream_arg_explicit_sync", test_stream_arg_explicit_sync, devices=devices)
 add_function_test(TestStreams, "test_stream_scope_implicit_sync", test_stream_scope_implicit_sync, devices=devices)
 
 add_function_test(TestStreams, "test_stream_arg_synchronize", test_stream_arg_synchronize, devices=devices)

warp/tests/test_torch.py

@@ -490,10 +490,14 @@ def test_torch_graph_torch_stream(test, device):
 
     # capture graph
     with wp.ScopedStream(warp_stream), torch.cuda.graph(g, stream=torch_stream):
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
+        wp.capture_begin(force_module_load=False, external=True)
+        try:
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+        finally:
+            wp.capture_end()
 
     # replay graph
     num_iters = 10
@@ -522,10 +526,14 @@ def test_torch_graph_warp_stream(test, device):
 
     # capture graph
     with wp.ScopedDevice(device), torch.cuda.graph(g, stream=torch_stream):
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
-        t += 1.0
-        wp.launch(inc, dim=n, inputs=[a])
+        wp.capture_begin(force_module_load=False, external=True)
+        try:
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+            t += 1.0
+            wp.launch(inc, dim=n, inputs=[a])
+        finally:
+            wp.capture_end()
 
     # replay graph
     num_iters = 10

warp/tests/test_utils.py

@@ -267,55 +267,60 @@ class TestUtils(unittest.TestCase):
     def test_warn(self):
         # Multiple warnings get printed out each time.
         with contextlib.redirect_stdout(io.StringIO()) as f:
-            frame_info = inspect.getframeinfo(inspect.currentframe())
             wp.utils.warn("hello, world!")
             wp.utils.warn("hello, world!")
 
         expected = (
-            "{}:{}: {}\n"
-            "{}:{}: {}\n"
-        ).format(
-            frame_info.filename,
-            frame_info.lineno + 1,
-            "UserWarning: hello, world!\n  wp.utils.warn(\"hello, world!\")",
-            frame_info.filename,
-            frame_info.lineno + 2,
-            "UserWarning: hello, world!\n  wp.utils.warn(\"hello, world!\")",
+            "Warp UserWarning: hello, world!\n"
+            "Warp UserWarning: hello, world!\n"
         )
+
         self.assertEqual(f.getvalue(), expected)
 
+        # Test verbose warnings
+        saved_verbosity = wp.config.verbose_warnings
+        try:
+            wp.config.verbose_warnings = True
+            with contextlib.redirect_stdout(io.StringIO()) as f:
+                frame_info = inspect.getframeinfo(inspect.currentframe())
+                wp.utils.warn("hello, world!")
+                wp.utils.warn("hello, world!")
+
+            expected = (
+                f"Warp UserWarning: hello, world! ({frame_info.filename}:{frame_info.lineno + 1})\n"
+                "  wp.utils.warn(\"hello, world!\")\n"
+                f"Warp UserWarning: hello, world! ({frame_info.filename}:{frame_info.lineno + 2})\n"
+                "  wp.utils.warn(\"hello, world!\")\n"
+            )
+
+            self.assertEqual(f.getvalue(), expected)
+
+        finally:
+            # make sure to restore warning verbosity
+            wp.config.verbose_warnings = saved_verbosity
+
+
         # Multiple similar deprecation warnings get printed out only once.
         with contextlib.redirect_stdout(io.StringIO()) as f:
-            frame_info = inspect.getframeinfo(inspect.currentframe())
             wp.utils.warn("hello, world!", category=DeprecationWarning)
             wp.utils.warn("hello, world!", category=DeprecationWarning)
 
         expected = (
-            "{}:{}: {}\n"
-        ).format(
-            frame_info.filename,
-            frame_info.lineno + 1,
-            "DeprecationWarning: hello, world!\n  wp.utils.warn(\"hello, world!\", category=DeprecationWarning)",
+            "Warp DeprecationWarning: hello, world!\n"
         )
+
         self.assertEqual(f.getvalue(), expected)
 
         # Multiple different deprecation warnings get printed out each time.
         with contextlib.redirect_stdout(io.StringIO()) as f:
-            frame_info = inspect.getframeinfo(inspect.currentframe())
             wp.utils.warn("foo", category=DeprecationWarning)
             wp.utils.warn("bar", category=DeprecationWarning)
 
         expected = (
-            "{}:{}: {}\n"
-            "{}:{}: {}\n"
-        ).format(
-            frame_info.filename,
-            frame_info.lineno + 1,
-            "DeprecationWarning: foo\n  wp.utils.warn(\"foo\", category=DeprecationWarning)",
-            frame_info.filename,
-            frame_info.lineno + 2,
-            "DeprecationWarning: bar\n  wp.utils.warn(\"bar\", category=DeprecationWarning)",
+            "Warp DeprecationWarning: foo\n"
+            "Warp DeprecationWarning: bar\n"
         )
+
         self.assertEqual(f.getvalue(), expected)
 
     def test_transform_expand(self):
@@ -425,7 +430,7 @@ class TestUtils(unittest.TestCase):
             with wp.ScopedTimer("hello", detailed=True):
                 pass
 
-        self.assertRegex(f.getvalue(), r"^         2 function calls in \d+\.\d+ seconds")
+        self.assertRegex(f.getvalue(), r"^         4 function calls in \d+\.\d+ seconds")
         self.assertRegex(f.getvalue(), r"hello took \d+\.\d+ ms$")