Honor mooncake_overlay in primitive inference

[yebai]

MooncakeInterpreter previously routed every primitive call site through NativeInterpreter for CallMeta, but NativeInterpreter is overlay-blind: when a primitive has a @mooncake_overlay that changes its return type, the inferred type at the call site was wrong, and downstream dispatch compiled against the original type. Detect overlay matches via Method.external_mt and route them through @invoke instead, keeping the NativeInterpreter fast path for non-overlay primitives.

Fixes #1169. Alternative to #1168.

CI Summary — GitHub Actions

Documentation Preview

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

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 │ 190.0 ns │     1.53 │        1.53 │   0.737 │        3.43 │   6.64 │
│             _sum_1000 │  1.08 μs │     6.72 │        1.03 │  4700.0 │        35.5 │   1.07 │
│          sum_sin_1000 │  7.42 μs │     2.57 │        1.11 │    1.65 │        11.0 │   1.77 │
│         _sum_sin_1000 │  4.64 μs │     3.81 │        2.63 │   338.0 │        17.6 │   3.05 │
│              kron_sum │ 193.0 μs │     13.1 │        3.31 │    10.7 │       520.0 │   13.7 │
│         kron_view_sum │ 263.0 μs │     12.9 │        5.27 │    20.9 │       450.0 │    7.2 │
│ naive_map_sin_cos_exp │   2.2 μs │     2.81 │        1.55 │ missing │        8.48 │   2.17 │
│       map_sin_cos_exp │  2.15 μs │     3.35 │        1.66 │     1.6 │        7.34 │   2.73 │
│ broadcast_sin_cos_exp │  2.31 μs │     2.93 │        1.52 │    4.25 │        1.39 │   2.06 │
│            simple_mlp │ 344.0 μs │     5.04 │        2.91 │    1.61 │        9.67 │   3.21 │
│                gp_lml │ 360.0 μs │     6.17 │        1.79 │    2.76 │     missing │   3.18 │
│    large_single_block │ 471.0 ns │     5.94 │        1.96 │  4190.0 │        31.8 │   2.08 │
└───────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[sunxd3]

On this issue, I slightly lean towards #1168 because indeed it's more general and the complexity diff w.r.t this PR is not too bad. A minor issue with the current code is the use of external_mt — not blocking, but less ideal in terms of interfacing with julia internals.

[yebai]

Note that there is already any_matches_primitive defined, which accesses Julia internals in similar ways.

[sunxd3]

Thanks @yebai for the follow-up.

After some thought and investigation, I now lean toward this PR over #1168 for the scope of issue #1169. Use of external_mt, IMO, is fine and a small price to pay; the alternatives seem worse as far as I can tell.

Before merging, there is one particular case I want to think through. Maybe it is not that important, but I would like to be clear.

---

Comparing on four versions of Mooncake:

pre-#1115                 4e2df06c0  c7c30f259^
main                      b79d67e1e
#1168                     ac5a35f72
#1170                     a0f11465e

on

using Mooncake

struct A end
struct B end

helper(::A) = A()
Mooncake.@mooncake_overlay helper(::A) = B()

primitive_wrapper(x::A) = helper(x)
Mooncake.@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{
    typeof(primitive_wrapper), A
}

caller() = primitive_wrapper(A())
caller_arg(x::A) = primitive_wrapper(x)

interp = Mooncake.MooncakeInterpreter(Mooncake.DefaultCtx, Mooncake.ReverseMode)

for sig in (
    Tuple{typeof(primitive_wrapper), A},
    Tuple{typeof(caller)},
    Tuple{typeof(caller_arg), A},
)
    ir, rt = Base.code_ircode_by_type(sig; interp)[1]
    println(sig, " => ", rt)
end

Results:

Version	primitive_wrapper(::A)	caller()	caller_arg(::A)
pre-#1115 4e2df06c0	B	B	B
#1168 ac5a35f72	B	B	B
main b79d67e1e	B	A	A
#1170 a0f11465e	B	A	A

---

This result indicates that #1115 introduced a behavior change, and #1168 reverted it. @AstitvaAggarwal and I had an offline chat and we'll dig deeper on this.

[sunxd3]

@AstitvaAggarwal I added a developer doc on primitives and overlay, can you give it a read and let me know your thoughts?

[AstitvaAggarwal]

looks good, dont think anything is incorrect.