Improve hydroelastic contact performance (#4142)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: nvtw

asv/benchmarks/simulation/bench_contacts.py

@@ -317,6 +317,67 @@ def time_simulate(self):
         wp.synchronize_device()
 
 
+class _ExampleCollideBenchmark:
+    """Collision-only timing of an example scene sized so contact generation dominates.
+
+    The ``*Defaults`` benchmarks time whole simulation frames, where the solver
+    hides most of the collision cost. This base class builds the same example
+    with more worlds, captures ``CollisionPipeline.collide`` alone into a CUDA
+    graph, and replays it, so contact-generation changes show up directly while
+    the setup cost stays close to the defaults benchmarks.
+    """
+
+    module_names: ClassVar[list[str]] = []
+    world_count = 200
+    launch_count = 20
+    repeat = pr_gate_repeat(5)
+    number = 1
+
+    def setup_cache(self):
+        _download_external_git_folder(ISAACGYM_ENVS_REPO_URL, ISAACGYM_NUT_BOLT_FOLDER)
+
+    def setup(self):
+        device = wp.get_device()
+        if not device.is_cuda:
+            raise SkipNotImplemented
+        example_cls = _import_example_class(self.module_names)
+        args = newton.examples.default_args(example_cls.create_parser())
+        args.world_count = self.world_count
+        args.num_per_world = 1
+        self.example = example_cls(ViewerNull(num_frames=1), args)
+        self.pipeline = self.example.collision_pipeline
+        self.state = self.example.state_0
+        self.contacts = self.example.contacts
+
+        for _ in range(3):
+            self.pipeline.collide(self.state, self.contacts)
+        wp.synchronize_device()
+        if int(self.contacts.rigid_contact_count.numpy()[0]) == 0:
+            raise RuntimeError("collide benchmark scene produced no contacts")
+
+        with wp.ScopedCapture(device=device) as capture:
+            self.pipeline.collide(self.state, self.contacts)
+        self.graph = capture.graph
+
+    @skip_benchmark_if(wp.get_cuda_device_count() == 0)
+    def time_collide(self):
+        for _ in range(self.launch_count):
+            wp.capture_launch(self.graph)
+        wp.synchronize_device()
+
+
+class FastExampleContactSdfCollide(_ExampleCollideBenchmark):
+    """Collision-only benchmark of the mesh-SDF nut-bolt scene at 200 worlds."""
+
+    module_names: ClassVar[list[str]] = ["newton.examples.contacts.example_nut_bolt_sdf"]
+
+
+class FastExampleContactHydroCollide(_ExampleCollideBenchmark):
+    """Collision-only benchmark of the hydroelastic nut-bolt scene at 200 worlds."""
+
+    module_names: ClassVar[list[str]] = ["newton.examples.contacts.example_nut_bolt_hydro"]
+
+
 class FastExampleContactPyramidDefaults:
     """Benchmark the box pyramid example with default configuration."""
 
@@ -528,6 +589,8 @@ def time_collide(self, case):
     benchmark_list = {
         "FastExampleContactSdfDefaults": FastExampleContactSdfDefaults,
         "FastExampleContactHydroWorkingDefaults": FastExampleContactHydroWorkingDefaults,
+        "FastExampleContactSdfCollide": FastExampleContactSdfCollide,
+        "FastExampleContactHydroCollide": FastExampleContactHydroCollide,
         "FastExampleContactPyramidDefaults": FastExampleContactPyramidDefaults,
         "FastConvexCollision": FastConvexCollision,
         "BroadPhaseCollision": BroadPhaseCollision,

changelog/+sdf-contact-performance-4d1c8e27.changed.md

@@ -0,0 +1,6 @@
+Speed up mesh-mesh SDF and hydroelastic SDF contact generation without changing the produced contacts:
+
+  - Hydroelastic octree refinement scans only the active records instead of the full worst-case buffers, and the marching-cubes kernel reads voxel corners with fewer texture lookups.
+  - The mesh-SDF edge kernels use 128-thread blocks with overlapped pre-prune hashtable probes, and the contact reducer clears its active entries with more threads.
+
+No migration is needed.

newton/_src/geometry/contact_reduction_global.py

@@ -59,6 +59,7 @@
 from newton._src.geometry.hashtable import (
     HASHTABLE_EMPTY_KEY,
     HashTable,
+    hashtable_find,
     hashtable_find_or_insert,
 )
 
@@ -1096,8 +1097,11 @@ def clear_active(self):
         later-scheduled blocks (or even later-issued warps/lanes under
         independent thread scheduling), causing some entries to be skipped.
         """
-        # Use fixed thread count for efficient GPU utilization
-        num_threads = min(1024, self.hashtable.capacity)
+        # The clear is a grid-stride loop over the active entries with several
+        # scattered stores each, so it needs many resident warps to hide the
+        # store latency; the active count is only known on the device, and
+        # surplus threads exit immediately.
+        num_threads = min(65536, self.hashtable.capacity)
 
         wp.launch(
             _clear_active_kernel,
@@ -1567,41 +1571,55 @@ def _export_and_reduce_contact_centered_two_spatial_depths(
     pos_2d = project_point_to_plane(bin_id, centered_position)
     key = make_contact_key(shape_a, shape_b, bin_id)
 
+    # === Voxel bin: inner depth coverage ===
+    voxel_idx = compute_voxel_index(position_local, aabb_lower_voxel, aabb_upper_voxel, voxel_res)
+    voxel_idx = wp.clamp(voxel_idx, 0, wp.static(NUM_VOXEL_DEPTH_SLOTS - 1))
+
+    voxels_per_group = wp.static(NUM_SPATIAL_DIRECTIONS + 1)
+    voxel_group = voxel_idx // voxels_per_group
+    voxel_local_slot = voxel_idx % voxels_per_group
+    voxel_bin_id = wp.static(NUM_NORMAL_BINS) + voxel_group
+    voxel_key = make_contact_key(shape_a, shape_b, voxel_bin_id)
+
+    # Resolve both keys up front so their probes and the slot reads below can
+    # overlap. Missing voxel keys are published only after a contact ID is
+    # available; deleting a speculative key after publication would race with
+    # concurrent threads that have already found it.
     entry_idx = hashtable_find_or_insert(key, reducer_data.ht_keys, reducer_data.ht_active_slots)
-    might_win = False
+    voxel_entry_idx = -1
+    if use_inner:
+        voxel_entry_idx = hashtable_find(voxel_key, reducer_data.ht_keys)
 
+    might_win = False
     if entry_idx >= 0:
+        # Read every slot before comparing so the loads issue back to back.
+        slot_values = replaced_values_vec_type()
+        for dir_i in range(wp.static(NUM_SPATIAL_DIRECTIONS + 1)):
+            slot_values[dir_i] = reducer_data.ht_values[dir_i * ht_capacity + entry_idx]
+        voxel_slot_value = wp.uint64(0)
+        if voxel_entry_idx >= 0:
+            voxel_slot_value = reducer_data.ht_values[voxel_local_slot * ht_capacity + voxel_entry_idx]
+
         if use_inner:
             if deterministic != 0:
                 max_depth_probe = _make_preprune_probe_det(-depth, fingerprint)
             else:
                 max_depth_probe = _make_contact_value_fast(-depth, 0, 0)
-            if reducer_data.ht_values[wp.static(NUM_SPATIAL_DIRECTIONS) * ht_capacity + entry_idx] < max_depth_probe:
+            if slot_values[wp.static(NUM_SPATIAL_DIRECTIONS)] < max_depth_probe:
+                might_win = True
+            if voxel_entry_idx >= 0 and voxel_slot_value < max_depth_probe:
+                might_win = True
+            if voxel_entry_idx < 0:
                 might_win = True
 
         for dir_i in range(wp.static(NUM_SPATIAL_DIRECTIONS)):
-            if not might_win:
-                dir_2d = get_spatial_direction_2d(dir_i)
-                score = wp.dot(pos_2d, dir_2d)
-                probe = make_spatial_preprune_probe(score, use_inner, fingerprint, deterministic)
-                if reducer_data.ht_values[dir_i * ht_capacity + entry_idx] < probe:
-                    might_win = True
+            dir_2d = get_spatial_direction_2d(dir_i)
+            score = wp.dot(pos_2d, dir_2d)
+            probe = make_spatial_preprune_probe(score, use_inner, fingerprint, deterministic)
+            if slot_values[dir_i] < probe:
+                might_win = True
     else:
         wp.atomic_add(reducer_data.ht_insert_failures, 0, 1)
-
-    # === Voxel bin: inner depth coverage ===
-    voxel_idx = compute_voxel_index(position_local, aabb_lower_voxel, aabb_upper_voxel, voxel_res)
-    voxel_idx = wp.clamp(voxel_idx, 0, wp.static(NUM_VOXEL_DEPTH_SLOTS - 1))
-
-    voxels_per_group = wp.static(NUM_SPATIAL_DIRECTIONS + 1)
-    voxel_group = voxel_idx // voxels_per_group
-    voxel_local_slot = voxel_idx % voxels_per_group
-    voxel_bin_id = wp.static(NUM_NORMAL_BINS) + voxel_group
-    voxel_key = make_contact_key(shape_a, shape_b, voxel_bin_id)
-
-    voxel_entry_idx = -1
-    if use_inner and not might_win:
-        voxel_entry_idx = hashtable_find_or_insert(voxel_key, reducer_data.ht_keys, reducer_data.ht_active_slots)
         if voxel_entry_idx >= 0:
             if deterministic != 0:
                 voxel_probe = _make_preprune_probe_det(-depth, fingerprint)
@@ -1613,11 +1631,6 @@ def _export_and_reduce_contact_centered_two_spatial_depths(
     if not might_win:
         return -1
 
-    # Compete with reserved ID zero before materializing contact geometry, so
-    # stale pre-prune survivors consume no buffer space.
-    if use_inner and voxel_entry_idx < 0:
-        voxel_entry_idx = hashtable_find_or_insert(voxel_key, reducer_data.ht_keys, reducer_data.ht_active_slots)
-
     won_mask = int(0)
     replaced_values = replaced_values_vec_type()
     if use_inner and entry_idx >= 0:
@@ -1664,7 +1677,8 @@ def _export_and_reduce_contact_centered_two_spatial_depths(
             won_mask |= 1 << wp.static(NUM_SPATIAL_DIRECTIONS + 1)
             replaced_values[wp.static(NUM_SPATIAL_DIRECTIONS + 1)] = previous_value
 
-    if won_mask == 0:
+    voxel_entry_missing = use_inner and voxel_entry_idx < 0
+    if won_mask == 0 and not voxel_entry_missing:
         return -1
 
     # Avoid allocating candidates superseded during their own slot updates.
@@ -1693,7 +1707,11 @@ def _export_and_reduce_contact_centered_two_spatial_depths(
         if reducer_data.ht_values[voxel_local_slot * ht_capacity + voxel_entry_idx] == provisional_value:
             still_wins = True
 
-    if not still_wins:
+    # Without a surviving slot win, a contact may still claim the voxel slot of
+    # an entry that is not published yet. That claim happens after the contact
+    # ID exists, so a losing claimant returns its ID below.
+    voxel_only = voxel_entry_missing and not still_wins
+    if not still_wins and not voxel_only:
         return -1
     contact_id = export_contact_to_buffer(shape_a, shape_b, position, normal, depth, fingerprint, reducer_data)
     if contact_id < 0:
@@ -1756,9 +1774,30 @@ def _export_and_reduce_contact_centered_two_spatial_depths(
             voxel_entry_idx = hashtable_find_or_insert(voxel_key, reducer_data.ht_keys, reducer_data.ht_active_slots)
         if voxel_entry_idx >= 0:
             voxel_value = make_contact_value(-depth, fingerprint, contact_id, deterministic)
-            reduction_update_slot(voxel_entry_idx, voxel_local_slot, voxel_value, reducer_data.ht_values, ht_capacity)
+            if voxel_only:
+                previous_value = reduction_try_update_slot(
+                    voxel_entry_idx,
+                    voxel_local_slot,
+                    voxel_value,
+                    reducer_data.ht_values,
+                    ht_capacity,
+                )
+                if previous_value >= voxel_value:
+                    reclaim_contact_id(contact_id, reducer_data)
+                    return -1
+            else:
+                reduction_update_slot(
+                    voxel_entry_idx,
+                    voxel_local_slot,
+                    voxel_value,
+                    reducer_data.ht_values,
+                    ht_capacity,
+                )
         else:
             wp.atomic_add(reducer_data.ht_insert_failures, 0, 1)
+            if voxel_only:
+                reclaim_contact_id(contact_id, reducer_data)
+                return -1
 
     return contact_id
 

newton/_src/geometry/narrow_phase.py

@@ -2691,15 +2691,18 @@ class (``gjk_candidate_pairs_count``, ``split_gjk_work_count``,
         self.mesh_triangle_block_dim = 32 if device_obj.is_cuda else self.block_dim
 
         # Dynamic block allocation for mesh-mesh and mesh-plane contacts.
-        # On CUDA we partition toward ~4 blocks per SM and launch twice as many
-        # mesh-mesh blocks to reduce serial work in long per-pair queues. On CPU
-        # there is no SM notion so we pick 64 as a modest parallelism
-        # target that splits pair work across OpenMP threads without
-        # over-subscribing on small scenes.
+        # On CUDA we partition toward ~4 blocks per SM. The 128-thread mesh-mesh
+        # kernel keeps four blocks resident per SM, and launching four times
+        # that many blocks (several waves of small chunks) balances the uneven
+        # per-chunk edge work: with only two waves, scenes around 100 pairs
+        # measured 25% slower because busy chunks dominated the tail. On CPU
+        # there is no SM notion so we pick 64 as a modest parallelism target
+        # that splits pair work across OpenMP threads without over-subscribing
+        # on small scenes.
         if self.reduce_contacts:
             target_blocks = device_obj.sm_count * 4 if device_obj.is_cuda else 64
             # Mesh-mesh
-            self.num_mesh_mesh_blocks = target_blocks * 2 if device_obj.is_cuda else target_blocks
+            self.num_mesh_mesh_blocks = target_blocks * 4 if device_obj.is_cuda else target_blocks
             self.mesh_mesh_target_blocks = target_blocks
             mesh_mesh_scan_size = self.max_mesh_mesh_pairs + 1
             self.mesh_mesh_block_offsets = wp.zeros(mesh_mesh_scan_size, dtype=wp.int32, device=device)

newton/_src/geometry/sdf_contact.py

@@ -49,8 +49,10 @@
 # ``mesh_sdf_collision_kernel`` and ``mesh_sdf_collision_global_reduce_kernel``.
 # Both kernels assume ``wp.block_dim() == MESH_SDF_BLOCK_DIM`` so that the
 # tile-stack capacity below correctly sizes the cooperative push overflow
-# margin.
-MESH_SDF_BLOCK_DIM = 256
+# margin. 128 threads (four resident blocks per SM at the 128-register cap)
+# measured about 8% faster than 256 threads with two blocks: the culling loop
+# barriers wait on four warps instead of eight at the same occupancy.
+MESH_SDF_BLOCK_DIM = 128
 
 # Capacity of the cooperative edge-selection tile stack. Sized to
 # ``2 * MESH_SDF_BLOCK_DIM`` so that the inner push loop can never
@@ -1541,7 +1543,7 @@ def mesh_sdf_collision_kernel(
 
     # The tiled launch provides exactly ``total_num_blocks`` blocks and the
     # kernel strides those blocks over all active combinations itself.
-    @wp.kernel(enable_backward=False, launch_bounds=(256, 2), grid_stride=False, module=_module)
+    @wp.kernel(enable_backward=False, launch_bounds=(MESH_SDF_BLOCK_DIM, 4), grid_stride=False, module=_module)
     def mesh_sdf_collision_global_reduce_kernel(
         shape_data: wp.array[wp.vec4],
         shape_transform: wp.array[wp.transform],