Repository: Genesis-Embodied-AI/Genesis
File: genesis/engine/solvers/rigid/rigid_solver_decomp.py (Python)

Code before:
                    and self.equalities_info[i_e, i_b].eq_obj1id == link1_idx
                    and self.equalities_info[i_e, i_b].eq_obj2id == link2_idx
                ):
                    if i_e < self.constraint_solver.ti_n_equalities[i_b] - 1:
                        self.equalities_info[i_e, i_b] = self.equalities_info[
                            self.constraint_solver.ti_n_equalities[i_b] - 1, i_b
                        ]
                    self.constraint_solver.ti_n_equalities[i_b] = self.constraint_solver.ti_n_equalities[i_b] - 1

    @gs.assert_built
    def get_weld_constraints(self, as_tensor: bool = True, to_torch: bool = True):
        n_welds = tuple(self.constraint_solver.ti_n_equalities.to_numpy())
        n_envs = len(n_welds)
        max_welds = max(n_welds) if n_welds else 0
        if max_welds == 0:
            empty = torch.empty if to_torch else np.empty
            obj_a = empty((0,), dtype=np.int32)
            obj_b = empty((0,), dtype=np.int32)
            if n_envs == 1:
                return {"obj_a": obj_a, "obj_b": obj_b}
            return {
                "obj_a": [obj_a] * n_envs,
                "obj_b": [obj_b] * n_envs,
            }

        total = n_envs * max_welds
        if to_torch:
            buf = torch.full((total, 3), -1, dtype=gs.tc_int, device=gs.device)
        else:
            buf = np.full((total, 3), -1, dtype=np.int32)

        self._kernel_collect_welds(buf)

        if to_torch:
            buf = buf.view(n_envs, max_welds, 3)
        else:
            buf = buf.reshape(n_envs, max_welds, 3)
        env_idx = buf[..., 0]
        obj_a = buf[..., 1]
        obj_b = buf[..., 2]

        if as_tensor:
            return {"env": env_idx, "obj_a": obj_a, "obj_b": obj_b}
        else:
            result_a, result_b = [], []
            for e, count in enumerate(n_welds):
                result_a.append(obj_a[e, :count].copy())
                result_b.append(obj_b[e, :count].copy())
            return {"obj_a": result_a, "obj_b": result_b}

    @ti.kernel

Review comment (bug):
This is wrong. See my previous comment:

> Could you follow more strictly the pattern in RigidSolver.get_contacts? You should not pass a 3D tensor because it would allocate more memory than necessary but most importantly, it will break "C"-contiguity, which would be problematic if forwarded as input argument of another taichi kernel.

Code after:
        for i in ti.static(range(4)):
            links_state.quat[i_l, envs_idx[i_b_]][i] = links_quat[envs_idx[i_b_], i_l, i]
        links_state.mass_shift[i_l, envs_idx[i_b_]] = mass_shift[envs_idx[i_b_], i_l]

    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
    for i_l, i_b_ in ti.ndrange(n_geoms, envs_idx.shape[0]):
        geoms_state.friction_ratio[i_l, envs_idx[i_b_]] = friction_ratio[envs_idx[i_b_], i_l]


@ti.kernel
def kernel_set_links_pos(
    relative: ti.i32,
    pos: ti.types.ndarray(),
    links_idx: ti.types.ndarray(),
    envs_idx: ti.types.ndarray(),
    links_info: array_class.LinksInfo,
    links_state: array_class.LinksState,
    rigid_global_info: array_class.RigidGlobalInfo,
    static_rigid_sim_config: ti.template(),
):

    ti.loop_config(serialize=static_rigid_sim_config.para_level < gs.PARA_LEVEL.ALL)
    for i_l_, i_b_ in ti.ndrange(links_idx.shape[0], envs_idx.shape[0]):
        i_b = envs_idx[i_b_]
        i_l = links_idx[i_l_]
        I_l = [i_l, i_b] if ti.static(static_rigid_sim_config.batch_links_info) else i_l

        if links_info.parent_idx[I_l] == -1 and links_info.is_fixed[I_l]:
            for i in ti.static(range(3)):
                links_state.pos[i_l, i_b][i] = pos[i_b_, i_l_, i]
            if relative:
                for i in ti.static(range(3)):
                    links_state.pos[i_l, i_b][i] = links_state.pos[i_l, i_b][i] + links_info.pos[I_l][i]
        else:
            q_start = links_info.q_start[I_l]
            for i in ti.static(range(3)):
                rigid_global_info.qpos[q_start + i, i_b] = pos[i_b_, i_l_, i]
            if relative:
                for i in ti.static(range(3)):
                    rigid_global_info.qpos[q_start + i, i_b] = (
                        rigid_global_info.qpos[q_start + i, i_b] + rigid_global_info.qpos0[q_start + i, i_b]
                    )


@ti.kernel
def kernel_set_links_quat(
    relative: ti.i32,
    quat: ti.types.ndarray(),
    links_idx: ti.types.ndarray(),
    envs_idx: ti.types.ndarray(),
    links_info: array_class.LinksInfo,