Repository: Genesis-Embodied-AI/Genesis
File: tests/test_fastcache_unified.py (Python)

Code before:
    )
    scene.add_entity(
        gs.morphs.Box(
            size=(0.4, 0.4, 0.4),
            pos=(0.0, 0.0, 0.18),
        )
    )

    scene.build(n_envs=args.num_env, env_spacing=(0.5, 0.5))
    scene.rigid_solver.collider.detection()
    gs.ti.sync()

    from genesis.engine.solvers.rigid.rigid_solver_decomp import kernel_step_1

    assert kernel_step_1._primal.fe_ll_cache_observations.cache_hit == args.expected_fe_ll_cache_hit
    assert kernel_step_1._primal.src_ll_cache_observations.cache_key_generated == args.expected_use_src_ll_cache
    assert kernel_step_1._primal.src_ll_cache_observations.cache_loaded == args.expected_src_ll_cache_hit


@pytest.mark.required
@pytest.mark.parametrize("enable_pure", [False, True])  # should not affect result
# note that using `backend` instead of `test_backend`, breaks genesis pytest...
@pytest.mark.parametrize("test_backend", ["cpu", "gpu"])  # should not affect result
@pytest.mark.parametrize("use_ndarray", [False, True])
def test_gs_num_envs(use_ndarray: bool, enable_pure: bool, test_backend: str, tmp_path: pathlib.Path) -> None:
    if platform.system() == "Darwin" and test_backend == "gpu" and use_ndarray and enable_pure:
        pytest.skip(
            "fast cache not supported on mac gpus when using ndarray,"
            " because mac gpu only supports up to 31 kernel parameters, and we need more than that."
        )

    # change num_env each time, and check effect on reading from cache
    for it, num_env in enumerate([3, 5, 7]):
        cmd_line = [
            sys.executable,
            __file__,
            gs_num_envs_child.__name__,
            "--test-backend",
            test_backend,
            "--num-env",
            str(num_env),
        ]
        env = dict(os.environ)
        env["GS_BETA_PURE"] = "1" if enable_pure else "0"
        env["GS_USE_NDARRAY"] = "1" if use_ndarray else "0"
        env["TI_OFFLINE_CACHE_FILE_PATH"] = str(tmp_path)
        # notes:
        # - if we use pure, we won't get as far as fe-ll-cache
        # - ndarray and pure therefore wont ever use fe-ll-cache (first time, nothing in cache; after that hit src-ll cache)
        # - not use ndarray will always try using the fe-ll-cache, but cache will be empty on first it
        #   but since we are changing num envs each time, using fields will never get a cache hit either

Review comment (performance):
```python
         pytest.skip(
            "fast cache not supported on mac gpus when using ndarray, because mac gpu only supports up to 31 kernel"
            " parameters, and we need more than that."
        )
```

Code after:
    scene.add_entity(
        gs.morphs.Box(
            size=(0.4, 0.4, 0.4),
            pos=(0.0, 0.0, 0.18),
        )
    )

    scene.build(n_envs=args.num_env, env_spacing=(0.5, 0.5))
    scene.rigid_solver.collider.detection()
    gs.ti.sync()

    from genesis.engine.solvers.rigid.rigid_solver_decomp import kernel_step_1

    assert kernel_step_1._primal.fe_ll_cache_observations.cache_hit == args.expected_fe_ll_cache_hit
    assert kernel_step_1._primal.src_ll_cache_observations.cache_key_generated == args.expected_use_src_ll_cache
    assert kernel_step_1._primal.src_ll_cache_observations.cache_loaded == args.expected_src_ll_cache_hit


@pytest.mark.required
@pytest.mark.parametrize("enable_pure", [False, True])  # should not affect result
# note that using `backend` instead of `test_backend`, breaks genesis pytest...
@pytest.mark.parametrize("test_backend", ["cpu", "gpu"])  # should not affect result
@pytest.mark.parametrize("use_ndarray", [False, True])
def test_gs_num_envs(use_ndarray: bool, enable_pure: bool, test_backend: str, tmp_path: pathlib.Path) -> None:
    if sys.platform == "Darwin" and test_backend == "gpu" and use_ndarray and enable_pure:
        pytest.skip(
            "fast cache not supported on mac gpus when using ndarray, because mac gpu only supports up to 31 kernel"
            " parameters, and we need more than that."
        )

    # change num_env each time, and check effect on reading from cache
    for it, num_env in enumerate([3, 5, 7]):
        cmd_line = [
            sys.executable,
            __file__,
            gs_num_envs_child.__name__,
            "--test-backend",
            test_backend,
            "--num-env",
            str(num_env),
        ]
        env = dict(os.environ)
        env["GS_BETA_PURE"] = "1" if enable_pure else "0"
        env["GS_USE_NDARRAY"] = "1" if use_ndarray else "0"
        env["TI_OFFLINE_CACHE_FILE_PATH"] = str(tmp_path)
        # notes:
        # - if we use pure, we won't get as far as fe-ll-cache
        # - ndarray and pure therefore wont ever use fe-ll-cache (first time, nothing in cache; after that hit src-ll cache)
        # - not use ndarray will always try using the fe-ll-cache, but cache will be empty on first it
        #   but since we are changing num envs each time, using fields will never get a cache hit either
        # soooo we are left only with (not pure) and (ndarray) and (it > 0)