Move hibernation to store persistent island by daisy-chaining entity indices

using ContactIsland.entity_idx_to_next_entity_idx_in_hibernated_island

Author: gasnica

genesis/engine/solvers/rigid/constraint_solver_decomp_island.py

@@ -9,6 +9,10 @@
 import genesis.utils.array_class as array_class
 
 from .rigid_solver_decomp_util import func_wakeup_entity_and_its_temp_island
+from .rigid_validate import (
+    validate_entity_hibernation_state_for_all_entities_in_temp_island,
+    validate_next_hibernated_entity_indices_in_entire_scene,
+)
 from .contact_island import ContactIsland
 from genesis.engine.solvers.rigid.rigid_debug import Debug
 
@@ -247,6 +251,12 @@ def add_collision_constraints__and_wakeup_entities(self, i_island: int, i_b: int
                         self.contact_island,
                     )
 
+        # after all collisions added for the island
+        if ti.static(Debug.validate):
+            validate_entity_hibernation_state_for_all_entities_in_temp_island( \
+                i_island, i_b, self._solver.entities_state, self.contact_island, expected_hibernation_state=False)
+
+
     @ti.func
     def add_joint_limit_constraints(self, i_island: int, i_b: int):
         for i_island_entity in range(self.contact_island.island_entity[i_island, i_b].n):

genesis/engine/solvers/rigid/contact_island.py

@@ -5,14 +5,18 @@
 
 import genesis as gs
 import genesis.utils.geom as gu
+from genesis.engine.solvers.rigid.rigid_debug import Debug
+
 
 if TYPE_CHECKING:
     from genesis.engine.solvers.rigid.collider_decomp import Collider
     from genesis.engine.solvers.rigid.rigid_solver_decomp import RigidSolver
 
 INVALID_NEXT_HIBERNATED_ENTITY_IDX = -1
-INVALID_HIBERNATED_ISLAND_ID = -1  # -1 is reserved for "no island", i.e. active or fixed entities
-FIRST_HIBERNATED_ISLAND_ID = 1  # start at 1, leave 0 unused, use -1 for "no island"
+
+
+from genesis.engine.solvers.rigid.rigid_validate import validate_next_hibernated_entity_indices_in_entire_scene
+
 
 @ti.data_oriented
 class ContactIsland:
@@ -87,14 +91,8 @@ def __init__(self, collider: "Collider"):
         #
 
         # Used to make islands persist through hibernation:
-        self.unused__entity_idx_to_next_entity_idx_in_hibernated_island = ti.field(dtype=gs.ti_int, shape=self.solver._batch_shape(self.solver.n_entities))
-        self.unused__entity_idx_to_next_entity_idx_in_hibernated_island.fill(INVALID_NEXT_HIBERNATED_ENTITY_IDX)
-
-        # Warning: do not mistake island_id for island_idx
-        self.next_hibernated_island_id = ti.field(dtype=gs.ti_int, shape=self.solver._B)
-        self.next_hibernated_island_id.fill(FIRST_HIBERNATED_ISLAND_ID)
-        self.hibernated_entity_idx_to_hibernated_island_id = ti.field(dtype=gs.ti_int, shape=self.solver._batch_shape(self.solver.n_entities))
-        self.hibernated_entity_idx_to_hibernated_island_id.fill(INVALID_HIBERNATED_ISLAND_ID)
+        self.entity_idx_to_next_entity_idx_in_hibernated_island = ti.field(dtype=gs.ti_int, shape=self.solver._batch_shape(self.solver.n_entities))
+        self.entity_idx_to_next_entity_idx_in_hibernated_island.fill(INVALID_NEXT_HIBERNATED_ENTITY_IDX)
 
     @ti.kernel
     def clear_island_mapping(self):
@@ -139,23 +137,27 @@ def add_all_contact_edges(self):
                 self.add_edge(link_a, link_b, i_b)
 
     @ti.kernel
-    def temp_hack__add_all_hibernated_island_edges(self):
+    def add_all_hibernated_island_edges(self):
         _B = self.solver._B
         n_entities = self.solver.n_entities
         ti.loop_config(serialize=self.solver._para_level < gs.PARA_LEVEL.ALL)
         for i_b in range(_B):
             for i_e in range(n_entities):
-                for i_e2 in range(i_e + 1, n_entities):
-                    if self.hibernated_entity_idx_to_hibernated_island_id[i_e, i_b] != INVALID_HIBERNATED_ISLAND_ID \
-                        and self.hibernated_entity_idx_to_hibernated_island_id[i_e, i_b] == self.hibernated_entity_idx_to_hibernated_island_id[i_e2, i_b]:
-                        any_link_a = self.solver.entities_info.link_start[i_e]
-                        any_link_b = self.solver.entities_info.link_start[i_e2]
-                        self.add_edge(any_link_a, any_link_b, i_b)
+                next_entity_idx = self.entity_idx_to_next_entity_idx_in_hibernated_island[i_e, i_b]
+                if next_entity_idx != INVALID_NEXT_HIBERNATED_ENTITY_IDX and next_entity_idx != i_e:
+                    if ti.static(Debug.validate):
+                        island_idx_a = self.entity_island[i_e, i_b]
+                        island_idx_b = self.entity_island[next_entity_idx, i_b]
+                        Debug.assertf(0x7ad00012, island_idx_a == island_idx_b)
+                    
+                    any_link_a = self.solver.entities_info.link_start[i_e]
+                    any_link_b = self.solver.entities_info.link_start[next_entity_idx]
+                    self.add_edge(any_link_a, any_link_b, i_b)
 
     def construct(self):
         self.clear_island_mapping()
         self.add_all_contact_edges()
-        self.temp_hack__add_all_hibernated_island_edges()
+        self.add_all_hibernated_island_edges()
         self.preprocess_island__map_entities_to_edges()
         self.construct_islands()
         self.postprocess_island__assign_contact_data_to_temp_islands()
@@ -257,6 +259,8 @@ def construct_islands(self):
                     self.stack[self.n_stack[i_b], i_b] = i_v
                     self.n_stack[i_b] = self.n_stack[i_b] + 1
                     self.entity_island[i_v, i_b] = self.n_islands[i_b]
+                    if ti.static(Debug.validate) and self.n_stack[i_b] > self.stack.shape[0]:
+                        print(f"danger: stack overflow: capacity and size: {self.stack.shape[0]} {self.n_stack[i_b]}")
 
                     while self.n_stack[i_b] > 0:
                         self.n_stack[i_b] = self.n_stack[i_b] - 1
@@ -273,6 +277,8 @@ def construct_islands(self):
                                 self.stack[self.n_stack[i_b], i_b] = next_v
                                 self.n_stack[i_b] = self.n_stack[i_b] + 1
                                 self.entity_island[next_v, i_b] = self.n_islands[i_b]
+                                if ti.static(Debug.validate) and self.n_stack[i_b] > self.stack.shape[0]:
+                                    print(f"danger: stack overflow: capacity and size: {self.stack.shape[0]} {self.n_stack[i_b]}")
 
                     self.n_islands[i_b] = self.n_islands[i_b] + 1
 

genesis/engine/solvers/rigid/rigid_solver_decomp.py

@@ -18,14 +18,18 @@
 import genesis.utils.array_class as array_class
 from genesis.engine.solvers.rigid.contact_island import ContactIsland
 from genesis.engine.solvers.rigid.rigid_debug import Debug
-from genesis.engine.solvers.rigid.rigid_validate import validate_entity_hibernation_state_for_all_entities_in_temp_island
+from genesis.engine.solvers.rigid.rigid_validate import (
+    validate_entity_hibernation_state_for_all_entities_in_temp_island,
+    validate_next_hibernated_entity_indices_in_entire_scene,
+)
 
 from ..base_solver import Solver
 from .collider_decomp import Collider
 from .constraint_solver_decomp import ConstraintSolver
 from .constraint_solver_decomp_island import ConstraintSolverIsland
 from .rigid_solver_decomp_util import func_wakeup_entity_and_its_temp_island
 from ....utils.sdf_decomp import SDF
+from .contact_island import INVALID_NEXT_HIBERNATED_ENTITY_IDX
 
 if TYPE_CHECKING:
     from genesis.engine.scene import Scene
@@ -4039,7 +4043,7 @@ def kernel_step_2(
     )
 
     if ti.static(static_rigid_sim_config.use_hibernation):
-        func_check_all_awake_islands_for_hibernation_or_update_aabb_sort_buffer(
+        func_for_all_awake_islands_either_hiberanate_or_update_aabb_sort_buffer(
             dofs_state=dofs_state,
             entities_state=entities_state,
             entities_info=entities_info,
@@ -4050,7 +4054,7 @@ def kernel_step_2(
             static_rigid_sim_config=static_rigid_sim_config,
             contact_island=contact_island,
         )
-        nondeterministic__func_aggregate_awake_entities(
+        func_aggregate_awake_entities(
             entities_state=entities_state,
             entities_info=entities_info,
             rigid_global_info=rigid_global_info,
@@ -4920,7 +4924,7 @@ def kernel_update_vgeoms(
 
 
 @ti.func
-def func_check_all_awake_islands_for_hibernation_or_update_aabb_sort_buffer(
+def func_for_all_awake_islands_either_hiberanate_or_update_aabb_sort_buffer(
     dofs_state: array_class.DofsState,
     entities_state: array_class.EntitiesState,
     entities_info: array_class.EntitiesInfo,
@@ -4987,14 +4991,12 @@ def func_check_all_awake_islands_for_hibernation_or_update_aabb_sort_buffer(
                     # perform hibernation
                     prev_entity_ref = entity_ref_range.start + entity_ref_range.n - 1
                     prev_entity_idx = ci.entity_id[prev_entity_ref, i_b]
-                    hibernated_island_id = ci.next_hibernated_island_id[i_b]
-                    ci.next_hibernated_island_id[i_b] = hibernated_island_id + 1
-                    # print(f"Hibernated island id: {hibernated_island_id} of {entity_ref_range.n} entities")
+
                     for i in range(entity_ref_range.n):
                         entity_ref = entity_ref_range.start + i
                         entity_idx = ci.entity_id[entity_ref, i_b]
 
-                        func_mark_entity_for_hibernation_and_zero_dof_velocities(
+                        func_hibernate_entity_and_zero_dof_velocities(
                             entity_idx,
                             i_b,
                             entities_state=entities_state,
@@ -5005,14 +5007,13 @@ def func_check_all_awake_islands_for_hibernation_or_update_aabb_sort_buffer(
                         )
 
                         # store entities in the hibernated islands by daisy chaining them
-                        ci.unused__entity_idx_to_next_entity_idx_in_hibernated_island[prev_entity_idx, i_b] = entity_idx
+                        Debug.assertf(0x7ad00014, ci.entity_idx_to_next_entity_idx_in_hibernated_island[prev_entity_idx, i_b] == INVALID_NEXT_HIBERNATED_ENTITY_IDX)
+                        ci.entity_idx_to_next_entity_idx_in_hibernated_island[prev_entity_idx, i_b] = entity_idx
                         prev_entity_idx = entity_idx
 
-                        ci.hibernated_entity_idx_to_hibernated_island_id[entity_idx, i_b] = hibernated_island_id
-
 
 @ti.func
-def nondeterministic__func_aggregate_awake_entities(
+def func_aggregate_awake_entities(
     entities_state: array_class.EntitiesState,
     entities_info: array_class.EntitiesInfo,
     rigid_global_info: array_class.RigidGlobalInfo,
@@ -5032,21 +5033,21 @@ def nondeterministic__func_aggregate_awake_entities(
         next_awake_entity_idx = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1)
         rigid_global_info.awake_entities[next_awake_entity_idx, i_b] = i_e
 
-        dof_start: ti.int32 = entities_info.dof_start[i_e]
         n_dofs = entities_info.n_dofs[i_e]
+        entity_dofs_base_idx: ti.int32 = entities_info.dof_start[i_e]
         awake_dofs_base_idx = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], n_dofs)
         for i in range(n_dofs):
-            rigid_global_info.awake_dofs[awake_dofs_base_idx + i, i_b] = dof_start + i
+            rigid_global_info.awake_dofs[awake_dofs_base_idx + i, i_b] = entity_dofs_base_idx + i
 
-        link_start: ti.int32 = entities_info.link_start[i_e]
         n_links = entities_info.n_links[i_e]
+        entity_links_base_idx: ti.int32 = entities_info.link_start[i_e]
         awake_links_base_idx = ti.atomic_add(rigid_global_info.n_awake_links[i_b], n_links)
         for i in range(n_links):
-            rigid_global_info.awake_links[awake_links_base_idx + i, i_b] = link_start + i
+            rigid_global_info.awake_links[awake_links_base_idx + i, i_b] = entity_links_base_idx + i
 
 
 @ti.func
-def func_mark_entity_for_hibernation_and_zero_dof_velocities(
+def func_hibernate_entity_and_zero_dof_velocities(
     i_e: int,
     i_b: int,
     entities_state: array_class.EntitiesState,

genesis/engine/solvers/rigid/rigid_solver_decomp_util.py

@@ -1,10 +1,14 @@
 import taichi as ti
 
 import genesis.utils.array_class as array_class
-from genesis.engine.solvers.rigid.contact_island import INVALID_HIBERNATED_ISLAND_ID
-from genesis.engine.solvers.rigid.rigid_debug import Debug
-from genesis.engine.solvers.rigid.rigid_validate import validate_temp_island_contains_both_hibernated_and_awake_entities
 
+from genesis.engine.solvers.rigid.contact_island import INVALID_NEXT_HIBERNATED_ENTITY_IDX
+from genesis.engine.solvers.rigid.rigid_debug import Debug
+from genesis.engine.solvers.rigid.rigid_validate import (
+    validate_temp_island_contains_both_hibernated_and_awake_entities,
+    validate_entity_hibernation_state_for_all_entities_in_temp_island,
+    validate_next_hibernated_entity_indices_in_entire_scene,
+)
 
 
 
@@ -23,51 +27,54 @@ def func_wakeup_entity_and_its_temp_island(
     # Note: Original function handled non-hibernated & fixed entities.
     # Now, we require a properly hibernated entity to be passed in.
     island_idx = contact_island.entity_island[i_e, i_b]
-    is_entity_fixed = island_idx == -1
-    Debug.assertf(0x0ad00007, not is_entity_fixed)  # Fixed entities are excluded from hibernation logic
-    Debug.assertf(0x0ad0000b, entities_state.hibernated[i_e, i_b])
+    Debug.assertf(0x7ad00008, not contact_island.island_hibernated[island_idx, i_b])  # Island already expected to be marked as not hibernated
 
-    # Note: temporarily, we have duplicated logic for hibernation: the entity_state.hibernated,
-    # and the new contact_island's arrays that are used to store persistent_hibernated_island_id.
-    hibernated_island_id = contact_island.hibernated_entity_idx_to_hibernated_island_id[i_e, i_b]
-    is_entity_hibernated = hibernated_island_id != INVALID_HIBERNATED_ISLAND_ID
-    Debug.assertf(0x0ad0000c, is_entity_hibernated)  # Entityt must belong to a persistent hibernated island
+    is_entity_fixed = island_idx == -1
+    Debug.assertf(0x7ad00007, not is_entity_fixed)  # Fixed entities are excluded from hibernation logic
+    Debug.assertf(0x7ad0000b, entities_state.hibernated[i_e, i_b])
 
     if ti.static(Debug.validate):
         validate_temp_island_contains_both_hibernated_and_awake_entities(island_idx, i_b, entities_state, contact_island)
 
-    n_entities = entities_info.n_links.shape[0]
-    for entity_idx in range(n_entities): 
-
-
-        if contact_island.hibernated_entity_idx_to_hibernated_island_id[entity_idx, i_b] == hibernated_island_id:
-            Debug.assertf(0x0ad0000b, entities_state.hibernated[entity_idx, i_b])  # Entity expected to be hibernated
-
-            contact_island.hibernated_entity_idx_to_hibernated_island_id[entity_idx, i_b] = INVALID_HIBERNATED_ISLAND_ID
-
-            if entities_state.hibernated[entity_idx, i_b]:
-                entities_state.hibernated[entity_idx, i_b] = False
-                n_awake_entities = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1)
-                rigid_global_info.awake_entities[n_awake_entities, i_b] = entity_idx
-
-                # todo: do single atomic add
-                for i_d in range(entities_info.dof_start[entity_idx], entities_info.dof_end[entity_idx]):
-                    dofs_state.hibernated[i_d, i_b] = False
-                    n_awake_dofs = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], 1)
-                    rigid_global_info.awake_dofs[n_awake_dofs, i_b] = i_d
-
-                # todo: do single atomic add
-                for i_l in range(entities_info.link_start[entity_idx], entities_info.link_end[entity_idx]):
-                    links_state.hibernated[i_l, i_b] = False
-                    n_awake_links = ti.atomic_add(rigid_global_info.n_awake_links[i_b], 1)
-                    rigid_global_info.awake_links[n_awake_links, i_b] = i_l
-
-                for i_g in range(entities_info.geom_start[entity_idx], entities_info.geom_end[entity_idx]):
-                    geoms_state.hibernated[i_g, i_b] = False
-
-                # validation only: un-hibernate the island
-                non_persistent_hibernated_island_idx = contact_island.entity_island[entity_idx, i_b]
-                Debug.assertf(0x0ad00008, non_persistent_hibernated_island_idx != -1)  # Entity being hibernated has invalid temp island index
-                Debug.assertf(0x0ad00009, non_persistent_hibernated_island_idx == island_idx)  # Not matching island indices
-
-                contact_island.island_hibernated[non_persistent_hibernated_island_idx, i_b] = False
+    entity_ref_range = contact_island.island_entity[island_idx, i_b]
+    for ei in range(entity_ref_range.n):
+        entity_ref = entity_ref_range.start + ei
+        entity_idx = contact_island.entity_id[entity_ref, i_b]
+        Debug.assertf(0x7ad0000d, contact_island.entity_island[entity_idx, i_b] == island_idx)  # Unexpected entity from outside of island
+
+        is_entity_hibernated = entities_state.hibernated[entity_idx, i_b]
+        next_hibernated_entity_idx = contact_island.entity_idx_to_next_entity_idx_in_hibernated_island[entity_idx, i_b]
+        Debug.assertf(0x7ad0000e, is_entity_hibernated == (next_hibernated_entity_idx != INVALID_NEXT_HIBERNATED_ENTITY_IDX))  # Inconsistent entity state
+
+        if is_entity_hibernated:
+            contact_island.entity_idx_to_next_entity_idx_in_hibernated_island[entity_idx, i_b] = INVALID_NEXT_HIBERNATED_ENTITY_IDX
+
+            entities_state.hibernated[entity_idx, i_b] = False
+            n_awake_entities = ti.atomic_add(rigid_global_info.n_awake_entities[i_b], 1)
+            rigid_global_info.awake_entities[n_awake_entities, i_b] = entity_idx
+
+            n_dofs = entities_info.n_dofs[entity_idx]
+            base_entity_dof_idx = entities_info.dof_start[entity_idx]
+            base_awake_dof_idx = ti.atomic_add(rigid_global_info.n_awake_dofs[i_b], n_dofs)
+            for i in range(n_dofs):
+                i_d = base_entity_dof_idx + i
+                dofs_state.hibernated[i_d, i_b] = False
+                rigid_global_info.awake_dofs[base_awake_dof_idx + i, i_b] = i_d
+
+            n_links = entities_info.n_links[entity_idx]
+            base_entity_link_idx = entities_info.link_start[entity_idx]
+            base_awake_link_idx = ti.atomic_add(rigid_global_info.n_awake_links[i_b], n_links)
+            for i in range(n_links):
+                i_l = base_entity_link_idx + i
+                links_state.hibernated[i_l, i_b] = False
+                rigid_global_info.awake_links[base_awake_link_idx + i, i_b] = i_l
+
+            for i_g in range(entities_info.geom_start[entity_idx], entities_info.geom_end[entity_idx]):
+                geoms_state.hibernated[i_g, i_b] = False
+
+    # contact_island.island_hibernated[island_idx, i_b] = False
+    if ti.static(Debug.validate):
+        validate_entity_hibernation_state_for_all_entities_in_temp_island( \
+            island_idx, i_b, entities_state, contact_island, expected_hibernation_state=False)
+        validate_next_hibernated_entity_indices_in_entire_scene(i_b, entities_state, contact_island)
+