[FEATURE] Enhance SAP coupler for coupling between Rigid body and Fem object. (Genesis-Embodied-AI#1458)

Author: Libero0809

genesis/engine/bvh.py

@@ -281,8 +281,8 @@ def _kernel_radix_sort_morton_codes_one_round(self, i: int):
 
         # Reorder morton codes
         for i_b, i_a in ti.ndrange(self.n_batches, self.n_aabbs):
-            code = (self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF
-            idx = ti.i32(self.offset[i_b, i_a] + self.prefix_sum[i_b, ti.i32(code)])
+            code = ti.i32((self.morton_codes[i_b, i_a][1 - (i // 4)] >> ((i % 4) * 8)) & 0xFF)
+            idx = ti.i32(self.offset[i_b, i_a] + self.prefix_sum[i_b, code])
             self.tmp_morton_codes[i_b, idx] = self.morton_codes[i_b, i_a]
 
         # Swap the temporary and original morton codes
@@ -445,14 +445,15 @@ def _kernel_compute_bounds_one_layer(self) -> ti.i32:
 
         return is_done
 
-    @ti.kernel
+    @ti.func
     def query(self, aabbs: ti.template()):
         """
         Query the BVH for intersections with the given AABBs.
 
         The results are stored in the query_result field.
         """
         self.query_result_count[None] = 0
+        overflow = False
 
         n_querys = aabbs.shape[1]
         for i_b, i_q in ti.ndrange(self.n_batches, n_querys):
@@ -474,6 +475,8 @@ def query(self, aabbs: ti.template()):
                         idx = ti.atomic_add(self.query_result_count[None], 1)
                         if idx < self.max_n_query_results:
                             self.query_result[idx] = gs.ti_ivec3(i_b, i_a, i_q)  # Store the AABB index
+                        else:
+                            overflow = True
                     else:
                         # Push children onto the stack
                         if node.right != -1:
@@ -483,6 +486,8 @@ def query(self, aabbs: ti.template()):
                             query_stack[stack_depth] = node.left
                             stack_depth += 1
 
+        return overflow
+
 
 @ti.data_oriented
 class FEMSurfaceTetLBVH(LBVH):
@@ -503,10 +508,13 @@ def filter(self, i_a, i_q):
 
         This is used to avoid self-collisions in FEM surface tets.
 
-        i_a: index of the found AABB
-        i_q: index of the query AABB
+        Parameters
+        ----------
+        i_a:
+            index of the found AABB
+        i_q:
+            index of the query AABB
         """
-
         result = i_a >= i_q
         i_av = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_a]].el2v
         i_qv = self.fem_solver.elements_i[self.fem_solver.surface_elements[i_q]].el2v