warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl

warp/dlpack.py

@@ -65,7 +65,7 @@ class _DLPackTensorHolder:
     """Class responsible for deleting DLManagedTensor memory after ownership is transferred from a capsule."""
 
     def __new__(cls, *args, **kwargs):
-        instance = super(_DLPackTensorHolder, cls).__new__(cls)
+        instance = super().__new__(cls)
         instance.mem_ptr = None
         return instance
 

warp/examples/benchmarks/benchmark_cloth.py

@@ -243,7 +243,7 @@ def run_benchmark(mode, dim, timers, render=False):
             # run one warm-up iteration to accurately measure initialization time (some engines do lazy init)
             positions = integrator.simulate(sim_dt, sim_substeps)
 
-    label = "Dim ({}^2)".format(dim)
+    label = f"Dim ({dim}^2)"
 
     # run simulation
     for _i in range(sim_frames):
@@ -275,7 +275,7 @@ run_benchmark(mode, 128, timers, render=False)
 # write results
 
 for k, v in timers.items():
-    print("{:16} min: {:8.2f} max: {:8.2f} avg: {:8.2f}".format(k, np.min(v), np.max(v), np.mean(v)))
+    print(f"{k:16} min: {np.min(v):8.2f} max: {np.max(v):8.2f} avg: {np.mean(v):8.2f}")
 
 report = open(os.path.join("benchmark.csv"), "a")
 writer = csv.writer(report, delimiter=",")

warp/examples/benchmarks/benchmark_tile_sort.py

@@ -0,0 +1,155 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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 time
+
+import numpy as np
+
+import warp as wp
+
+BLOCK_DIM = 128
+
+
+def create_test_kernel(KEY_TYPE, MAX_SORT_LENGTH):
+    @wp.kernel
+    def tile_sort_kernel(
+        input_keys: wp.array(dtype=KEY_TYPE, ndim=2),
+        input_values: wp.array(dtype=wp.int32, ndim=2),
+        output_keys: wp.array(dtype=KEY_TYPE, ndim=2),
+        output_values: wp.array(dtype=wp.int32, ndim=2),
+    ):
+        batch_id, i = wp.tid()
+
+        # Load input into shared memory
+        keys = wp.tile_load(input_keys[batch_id], shape=MAX_SORT_LENGTH, storage="shared")
+        values = wp.tile_load(input_values[batch_id], shape=MAX_SORT_LENGTH, storage="shared")
+
+        # Perform in-place sorting
+        wp.tile_sort(keys, values)
+
+        # Store sorted shared memory into output arrays
+        wp.tile_store(output_keys[batch_id], keys)
+        wp.tile_store(output_values[batch_id], values)
+
+    return tile_sort_kernel
+
+
+if __name__ == "__main__":
+    wp.config.quiet = True
+    wp.init()
+    wp.clear_kernel_cache()
+    wp.set_module_options({"fast_math": True, "enable_backward": False})
+
+    iterations = 100
+    rng = np.random.default_rng(42)
+
+    shared_benchmark_data = {}
+    cub_segmented_sort_benchmark_data = {}
+
+    array_length = list(range(16, 257, 16))
+
+    print(
+        f"{'Type':<12s} {'Batch Size':<12s} {'Length':<12s} {'Tile Sort (ms)':<16s} {'Cub Segmented Sort (ms)':<24s} {'CubTime/TileTime':<16s}"
+    )
+    print("-" * 100)
+
+    for dtype in [int, float]:
+        for batch_size_exponent in range(5, 11):
+            batch_size = 2**batch_size_exponent
+            for length in array_length:
+                if dtype == int:
+                    np_keys = rng.choice(1000000000, size=(batch_size, length), replace=False)
+                else:  # dtype == float
+                    np_keys = rng.choice(1000000, size=(batch_size, length), replace=False).astype(np.float32)
+
+                np_values = np.tile(np.arange(length), (batch_size, 1))
+
+                # Sort using NumPy for validation
+                np_sorted_keys = np.zeros_like(np_keys)
+                np_sorted_values = np.zeros_like(np_values)
+                for b in range(batch_size):
+                    sorted_indices = np.argsort(np_keys[b])
+                    np_sorted_keys[b] = np_keys[b][sorted_indices]
+                    np_sorted_values[b] = np_values[b][sorted_indices]
+
+                # Generate random keys and iota indexer
+                input_keys = wp.array(np_keys, dtype=dtype, ndim=2, device="cuda")
+                input_values = wp.array(np_values, dtype=int, ndim=2, device="cuda")
+                output_keys = wp.zeros_like(input_keys, device="cuda")
+                output_values = wp.zeros_like(input_values, device="cuda")
+
+                kernel = create_test_kernel(dtype, length)
+
+                cmd = wp.launch_tiled(
+                    kernel,
+                    dim=batch_size,
+                    inputs=[input_keys, input_values, output_keys, output_values],
+                    block_dim=BLOCK_DIM,
+                    record_cmd=True,
+                )
+                # Warmup
+                for _ in range(5):
+                    cmd.launch()
+
+                with wp.ScopedTimer("benchmark", cuda_filter=wp.TIMING_KERNEL, print=False, synchronize=True) as timer:
+                    for _ in range(iterations):
+                        cmd.launch()
+                    wp.synchronize()
+
+                if dtype == int:
+                    keys_match = np.array_equal(output_keys.numpy(), np_sorted_keys)
+                else:  # dtype == float
+                    keys_match = np.allclose(output_keys.numpy(), np_sorted_keys, atol=1e-6)  # Use tolerance for floats
+
+                values_match = np.array_equal(output_values.numpy(), np_sorted_values)
+
+                # Validate results
+                assert keys_match, f"Key sorting mismatch for dtype={dtype}!"
+                assert values_match, f"Value sorting mismatch for dtype={dtype}!"
+
+                timing_results = [result.elapsed for result in timer.timing_results]
+                mean_timing = np.mean(timing_results)
+
+                shared_benchmark_data[length] = mean_timing
+
+                # Allocate memory
+                input_keys = wp.zeros(shape=(batch_size * 2, length), dtype=dtype, device="cuda")
+                input_values = wp.zeros(shape=(batch_size * 2, length), dtype=int, device="cuda")
+
+                # Copy data
+                input_keys.assign(np_keys)
+                input_values.assign(np_values)
+
+                input_keys = input_keys.reshape(-1)
+                input_values = input_values.reshape(-1)
+
+                segments = wp.array(np.arange(0, batch_size + 1) * length, dtype=int, device="cuda")
+
+                # Compare with cub segmented radix sort
+                # Warmup
+                for _ in range(5):
+                    wp.utils.segmented_sort_pairs(input_keys, input_values, batch_size * length, segments)
+
+                t1 = time.time_ns()
+                for _ in range(iterations):
+                    wp.utils.segmented_sort_pairs(input_keys, input_values, batch_size * length, segments)
+                wp.synchronize()
+                t2 = time.time_ns()
+                cub_segmented_sort_benchmark_data[length] = (t2 - t1) / (1_000_000 * iterations)
+
+                # Print results
+                print(
+                    f"{dtype!s:<12s} {batch_size:<12d} {length:<12d} {shared_benchmark_data[length]:<16.4g} {cub_segmented_sort_benchmark_data[length]:<24.4g} {cub_segmented_sort_benchmark_data[length] / shared_benchmark_data[length]:<16.4g}"
+                )

warp/examples/core/example_sample_mesh.py

@@ -176,7 +176,7 @@ def sample_mesh(
     sample = wp.randf(rng)
     tri = wp.lower_bound(cdf, sample)
 
-    # Sample the location in that triangle using random barycentric cordinates.
+    # Sample the location in that triangle using random barycentric coordinates.
     ru = wp.randf(rng)
     rv = wp.randf(rng)
     tri_u = 1.0 - wp.sqrt(ru)

warp/examples/core/example_spin_lock.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+###########################################################################
+# Example Spin Lock
+#
+# Shows how to use a spin lock to synchronize access to a shared resource.
+#
+###########################################################################
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.func
+def spinlock_acquire(lock: wp.array(dtype=wp.int32)):
+    # Try to acquire the lock by setting it to 1 if it's 0
+    while wp.atomic_cas(lock, 0, 0, 1) == 1:
+        pass
+
+
+@wp.func
+def spinlock_release(lock: wp.array(dtype=wp.int32)):
+    # Release the lock by setting it back to 0
+    wp.atomic_exch(lock, 0, 0)
+
+
+@wp.func
+def volatile_read(ptr: wp.array(dtype=wp.int32), index: int):
+    value = wp.atomic_exch(ptr, index, 0)
+    wp.atomic_exch(ptr, index, value)
+    return value
+
+
+@wp.kernel
+def test_spinlock_counter(
+    counter: wp.array(dtype=wp.int32), atomic_counter: wp.array(dtype=wp.int32), lock: wp.array(dtype=wp.int32)
+):
+    # Try to acquire the lock
+    spinlock_acquire(lock)
+
+    # Critical section - increment counter
+    # counter[0] = counter[0] + 1 # This gives wrong results - counter should be marked as volatile
+
+    # Work around since warp arrays cannot be marked as volatile
+    value = volatile_read(counter, 0)
+    counter[0] = value + 1
+
+    # Release the lock
+    spinlock_release(lock)
+
+    # Increment atomic counter for comparison
+    wp.atomic_add(atomic_counter, 0, 1)
+
+
+def test_spinlock(device):
+    # Create a lock array initialized to 0 (unlocked)
+    lock = wp.array([0], dtype=wp.int32, device=device)
+
+    # Create counter arrays initialized to 0
+    counter = wp.array([0], dtype=wp.int32, device=device)
+    atomic_counter = wp.array([0], dtype=wp.int32, device=device)
+
+    # Number of threads to test with
+    n = 1024
+
+    # Launch the test kernel
+    wp.launch(test_spinlock_counter, dim=n, inputs=[counter, atomic_counter, lock], device=device)
+
+    # Verify results
+    assert atomic_counter.numpy()[0] == n, f"Atomic counter should be {n}, got {atomic_counter.numpy()[0]}"
+    assert counter.numpy()[0] == n, f"Counter should be {n}, got {counter.numpy()[0]}"
+    assert lock.numpy()[0] == 0, "Lock was not properly released"
+
+    print(f"Final counter value: {counter.numpy()[0]}")
+    print(f"Final atomic counter value: {atomic_counter.numpy()[0]}")
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    test_spinlock(device="cuda")

warp/examples/core/example_work_queue.py

@@ -0,0 +1,118 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 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.
+
+###########################################################################
+# Example Work Queue
+#
+# Shows how to use a work queue to synchronize access to a shared resource.
+#
+###########################################################################
+
+import warp as wp
+from warp.tests.unittest_utils import *
+
+
+@wp.func
+def volatile_read(ptr: wp.array(dtype=wp.int32), index: int):
+    value = wp.atomic_add(ptr, index, 0)
+    return value
+
+
+@wp.struct
+class WorkQueue:
+    buffer: wp.array(dtype=wp.int32)
+    capacity: int
+    head: wp.array(dtype=wp.int32)
+    tail: wp.array(dtype=wp.int32)
+
+
+@wp.func
+def enqueue(queue: WorkQueue, item: int) -> bool:
+    while True:
+        # Read current head and tail atomically
+        current_tail = volatile_read(queue.tail, 0)
+        current_head = volatile_read(queue.head, 0)
+
+        # Check if queue is full
+        if (current_tail - current_head) >= queue.capacity:
+            return False
+
+        # Try to increment tail atomically
+        index = current_tail % queue.capacity
+        if wp.atomic_cas(queue.tail, 0, current_tail, current_tail + 1) == current_tail:
+            queue.buffer[index] = item
+            return True
+
+        # Retry if another thread changed tail
+
+
+@wp.func
+def dequeue(queue: WorkQueue) -> tuple[bool, int]:
+    while True:
+        # Read current head and tail atomically
+        current_head = volatile_read(queue.head, 0)
+        current_tail = volatile_read(queue.tail, 0)
+
+        # Check if queue is empty
+        if current_head >= current_tail:
+            return False, 0
+
+        # Get item at current head
+        index = current_head % queue.capacity
+        item = queue.buffer[index]
+
+        # Try to increment head atomically
+        if wp.atomic_cas(queue.head, 0, current_head, current_head + 1) == current_head:
+            return True, item
+
+        # Retry if another thread changed head
+
+
+@wp.kernel
+def process_queue(queue: WorkQueue):
+    counter = int(0)
+    while True:
+        success, item = dequeue(queue)
+        if not success:
+            break
+        wp.printf("Processed item: %d\n", item)
+        if item < 1000000:
+            if not enqueue(queue, item + 1000000):
+                wp.printf("Failed to enqueue item: %d\n", item + 1000000)
+        counter = counter + 1
+
+
+def test_work_queue(device):
+    # Create a work queue with capacity 1024
+    capacity = 8192
+    head = wp.array([0], dtype=wp.int32, device=device)
+    tail = wp.array([4096], dtype=wp.int32, device=device)
+    buffer = wp.array(np.arange(4096, dtype=np.int32), dtype=wp.int32, device=device)
+
+    queue = WorkQueue()
+    queue.capacity = capacity
+    queue.head = head
+    queue.tail = tail
+    queue.buffer = buffer
+
+    # Launch processing kernel
+    wp.launch(process_queue, dim=1024, inputs=[queue], device=device)
+
+    wp.synchronize()
+
+
+if __name__ == "__main__":
+    wp.clear_kernel_cache()
+    test_work_queue(device="cuda")

warp/examples/fem/example_adaptive_grid.py

@@ -70,7 +70,7 @@ def mass_form(
 @fem.integrand
 def side_divergence_form(s: fem.Sample, domain: fem.Domain, u: fem.Field, psi: fem.Field):
     # normal velocity jump (non-zero at resolution boundaries)
-    return -wp.dot(fem.jump(u, s), fem.normal(domain, s)) * psi(s)
+    return -wp.dot(fem.jump(u, s), fem.normal(domain, s)) * fem.average(psi, s)
 
 
 @wp.func
@@ -173,7 +173,7 @@ class Example:
         bd_test = fem.make_test(u_space, domain=boundary)
         bd_trial = fem.make_trial(u_space, domain=boundary)
         dirichlet_projector = fem.integrate(
-            noslip_projector_form, fields={"u": bd_test, "v": bd_trial}, nodal=True, output_dtype=float
+            noslip_projector_form, fields={"u": bd_test, "v": bd_trial}, assembly="nodal", output_dtype=float
         )
         fem.normalize_dirichlet_projector(dirichlet_projector)
 
@@ -187,7 +187,7 @@ class Example:
             rho_trial = fem.make_trial(rho_space)
 
         inv_mass_matrix = fem.integrate(
-            mass_form, fields={"u": rho_trial, "v": rho_test}, nodal=True, output_dtype=float
+            mass_form, fields={"u": rho_trial, "v": rho_test}, assembly="nodal", output_dtype=float
         )
         fem_example_utils.invert_diagonal_bsr_matrix(inv_mass_matrix)
 
@@ -269,8 +269,8 @@ if __name__ == "__main__":
 
                 stage = Usd.Stage.Open(os.path.join(warp.examples.get_asset_directory(), "rocks.usd"))
                 mesh = UsdGeom.Mesh(stage.GetPrimAtPath("/root/rocks"))
-                points = np.array((mesh.GetPointsAttr().Get()))
-                counts = np.array((mesh.GetFaceVertexCountsAttr().Get()))
+                points = np.array(mesh.GetPointsAttr().Get())
+                counts = np.array(mesh.GetFaceVertexCountsAttr().Get())
                 indices = np.array(mesh.GetFaceVertexIndicesAttr().Get())
                 ref_geom = (points, counts, indices)
             except Exception:

warp/examples/fem/example_apic_fluid.py

@@ -290,7 +290,7 @@ class Example:
                 vel_projector = fem.integrate(
                     velocity_boundary_projector_form,
                     fields={"u": velocity_trial, "v": velocity_test},
-                    nodal=True,
+                    assembly="nodal",
                     output_dtype=float,
                 )
                 fem.normalize_dirichlet_projector(vel_projector)

warp/examples/fem/example_burgers.py

@@ -146,7 +146,7 @@ class Example:
         # For simplicity, use nodal integration so that inertia matrix is diagonal
         trial = fem.make_trial(space=vector_space, domain=domain)
         matrix_inertia = fem.integrate(
-            vel_mass_form, fields={"u": trial, "v": self._test}, output_dtype=wp.float32, nodal=True
+            vel_mass_form, fields={"u": trial, "v": self._test}, output_dtype=wp.float32, assembly="nodal"
         )
         self._inv_mass_matrix = wp.sparse.bsr_copy(matrix_inertia)
         fem_example_utils.invert_diagonal_bsr_matrix(self._inv_mass_matrix)

warp/examples/fem/example_convection_diffusion.py

@@ -82,7 +82,7 @@ def diffusion_and_inertia_form(s: fem.Sample, phi: fem.Field, psi: fem.Field, dt
 
 
 class Example:
-    def __init__(self, quiet=False, degree=2, resolution=50, tri_mesh=False, viscosity=0.001, ang_vel=1.0):
+    def __init__(self, quiet=False, degree=2, resolution=50, mesh: str = "grid", viscosity=0.001, ang_vel=1.0):
         self._quiet = quiet
 
         self._ang_vel = ang_vel
@@ -91,11 +91,14 @@ class Example:
         self.sim_dt = 1.0 / (ang_vel * res)
         self.current_frame = 0
 
-        if tri_mesh:
-            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(res))
+        if mesh == "tri":
+            positions, tri_vidx = fem_example_utils.gen_trimesh(res=wp.vec2i(resolution))
             geo = fem.Trimesh2D(tri_vertex_indices=tri_vidx, positions=positions, build_bvh=True)
+        elif mesh == "quad":
+            positions, quad_vidx = fem_example_utils.gen_quadmesh(res=wp.vec2i(resolution))
+            geo = fem.Quadmesh2D(quad_vertex_indices=quad_vidx, positions=positions, build_bvh=True)
         else:
-            geo = fem.Grid2D(res=wp.vec2i(res))
+            geo = fem.Grid2D(res=wp.vec2i(resolution))
 
         domain = fem.Cells(geometry=geo)
         scalar_space = fem.make_polynomial_space(geo, degree=degree)
@@ -149,7 +152,7 @@ if __name__ == "__main__":
     parser.add_argument("--num_frames", type=int, default=250, help="Total number of frames.")
     parser.add_argument("--viscosity", type=float, default=0.001, help="Fluid viscosity parameter.")
     parser.add_argument("--ang_vel", type=float, default=1.0, help="Angular velocity.")
-    parser.add_argument("--tri_mesh", action="store_true", help="Use a triangular mesh.")
+    parser.add_argument("--mesh", choices=("grid", "tri", "quad"), default="grid", help="Mesh type.")
     parser.add_argument(
         "--headless",
         action="store_true",
@@ -164,7 +167,7 @@ if __name__ == "__main__":
             quiet=args.quiet,
             degree=args.degree,
             resolution=args.resolution,
-            tri_mesh=args.tri_mesh,
+            mesh=args.mesh,
             viscosity=args.viscosity,
             ang_vel=args.ang_vel,
         )