compile: failed first trace leaves a stale cache entry that returns empty-success on retry
bug
### Describe the bug
The caching lambda built by `mlx::core::detail::compile` marks a cache entry **non-empty before it traces it**. In `mlx/compile.cpp`, the per-call lambda does (paraphrased):
```cpp
auto& entry = compiler_cache().find(fun_id, inputs, shapeless, constants);
if (entry.empty) {
entry.empty = false; // <-- marked filled BEFORE the trace
entry.constants = std::move(constants);
std::tie(entry.inputs, entry.outputs, entry.extra) =
compile_trace(fun, inputs, shapeless); // <-- may throw
...
}
```
If `compile_trace` throws on the **first** trace — e.g. the traced `fun` raises — the entry is left `empty == false` but with no `inputs` / `outputs` / tape filled. A later call with matching inputs then `find()`s that entry, sees `empty == false`, **skips tracing**, and `compile_replace`s an effectively empty tape — returning empty outputs as a spurious success instead of re-raising (or re-tracing).
### To reproduce
Compile a function that raises on its first invocation, then call it twice with the same (matching) inputs:
1. First call: `compile_trace` raises and propagates (expected).
2. Second call: the half-filled entry is hit, the trace is skipped, and **empty outputs are returned as a success**.
A nullary compiled function is the sharpest case (its inputs always match, so the second call is a guaranteed `find` on the same entry).
### Expected behavior
A failed first trace should not leave a usable-looking cache entry. Either:
- set `entry.empty = false` only **after** `compile_trace` (and the dfs / simplify / fuse) complete successfully, or
- roll the entry back (`entry.empty = true`, or erase it) if tracing throws,
so a retry re-traces cleanly and a failure stays a failure.
### Why it matters
Any caller that compiles a callable which can raise inherits a silent wrong-result on retry — the failure is masked as an empty-output success. A binding can work around it by poisoning the compiled wrapper after a first-trace failure, but the clean fix is to not mark the entry filled until it actually is.
1 条评论