clear Mooncake OpaqueClosure Caches

[AstitvaAggarwal]

This PR closes #1116 by adding clear_mooncake_caches!(; gc::Bool=false) as a public API function that clears all Julia-level caches Mooncake holds globally.

Problem :

• When Mooncake differentiates a function for the first time, it compiles a reverse-mode rule and stores it in GLOBAL_INTERPRETERS across three caches : oc_cache (compiled OpaqueClosures), code_cache (CodeInstance objects), and inf_cache (InferenceResult objects). These caches never shrink on their own.

• In long-running Mooncake sessions we end up differentiating many distinct function signatures - these accumulate indefinitely. (The Discourse thread linked in Clear Mooncake internal caches #1116 reported the user hitting ~4.5 GB from Mooncake's caches with no way to reclaim it)

Solution :

• Our function empties oc_cache, code_cache.dict, and inf_cache in-place for both ForwardMode and ReverseMode interpreters. After calling it, Mooncake re-derives rules from scratch on the next use, and the GC can reclaim the now-unreferenced DerivedRule, CodeInstance, and InferenceResult objects.

• The gc keyword (default false) optionally runs GC.gc(true) immediately after clearing for users who want memory reclaimed right away rather than at the next GC cycle.

What is exactly freed :

Clearing the caches makes the Julia-level objects (DerivedRule, CodeInstance, InferenceResult) GC eligible. The JIT-compiled native machine code in main memory is allocated by LLVM outside the GC heap and cannot be reclaimed from user code and remains for the lifetime of the process (at least as of Julia v1.12), to clear this we must restart the REPL instance itself.

CI Summary — GitHub Actions

Documentation Preview

Mooncake.jl documentation for PR #1131 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1131/

Performance

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 170.0 ns │     1.65 │        1.76 │   0.706 │        3.42 │    6.6 │
│                  _sum_1000 │ 952.0 ns │     6.91 │        1.04 │  4830.0 │        34.1 │   1.08 │
│               sum_sin_1000 │   6.9 μs │     3.52 │        1.41 │    1.66 │        10.2 │    1.8 │
│              _sum_sin_1000 │   5.9 μs │     3.62 │        2.01 │   257.0 │        11.8 │   2.21 │
│                   kron_sum │ 212.0 μs │     11.8 │        3.27 │    22.8 │       348.0 │   19.8 │
│              kron_view_sum │ 292.0 μs │     10.7 │        4.71 │    23.2 │       260.0 │   12.6 │
│      naive_map_sin_cos_exp │  2.63 μs │     2.82 │        1.43 │ missing │        5.94 │   1.94 │
│            map_sin_cos_exp │   2.6 μs │     3.05 │        1.46 │    1.95 │        5.03 │   2.42 │
│      broadcast_sin_cos_exp │  2.68 μs │     2.77 │        1.43 │    3.95 │        1.27 │   1.92 │
│                 simple_mlp │ 396.0 μs │     4.38 │        2.45 │    1.35 │        6.63 │   2.57 │
│                     gp_lml │ 448.0 μs │     5.71 │        1.69 │     5.8 │     missing │   4.22 │
│ turing_broadcast_benchmark │  2.18 ms │     4.47 │        2.79 │ missing │        22.4 │   1.84 │
│         large_single_block │ 401.0 ns │     5.75 │        2.05 │  4310.0 │        34.1 │   2.15 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!