compile: a freed closure leaves every thread's compile cache

fork.yaml

@@ -172,6 +172,20 @@ def:
         - "mlx/compile.cpp"
         - "mlx/compile_impl.h"
         - "tests/compile_tests.cpp"
+    - title: "Freed compiled closures leave every thread cache"
+      description: |
+        `compile(std::function)` keys a closure by the address of a heap copy of it. Upstream,
+        the deleter of that copy erases the key only from the cache of the thread that called
+        `compile()`. Each call records its trace in the cache of the calling thread, so a trace
+        from another thread stayed after the copy was freed. The trace kept its arrays alive,
+        and a later closure at the same address could replay it. The deleter now erases the key
+        from every live thread cache through the registry above, before the delete. The test
+        `test freeing a compiled closure erases every thread cache` records a trace on a second
+        thread, frees the closure, and checks that the captured array is uniquely referenced.
+        (PR #43)
+      globs:
+        - "mlx/compile.cpp"
+        - "tests/compile_tests.cpp"
     - title: "Integer floor_divide tests"
       description: |
         The fork keeps its regression tests for integer `floor_divide` from PR #32. Upstream

mlx/compile.cpp

@@ -1302,8 +1302,13 @@ std::function<std::vector<array>(const std::vector<array>&)> compile(
     auto pfun = std::shared_ptr<
         std::function<std::vector<array>(const std::vector<array>&)>>(
         new std::function<std::vector<array>(const std::vector<array>&)>{fun},
-        [cache = detail::compile_cache()](auto* p) {
-          detail::compile_erase(cache, reinterpret_cast<std::uintptr_t>(p));
+        [](auto* p) {
+          // The compiled function can be called on any thread, and each call
+          // records its trace in the cache of the calling thread. Erase the
+          // key from every live thread cache, and do it before the delete so
+          // that no new function can get this address while a trace for it
+          // is still cached.
+          detail::compile_erase(reinterpret_cast<std::uintptr_t>(p));
           delete p;
         });
     fun_id = reinterpret_cast<std::uintptr_t>(pfun.get());

tests/compile_tests.cpp

@@ -921,3 +921,37 @@ TEST_CASE("test compatibility erase clears every live thread cache") {
   first.join();
   second.join();
 }
+
+TEST_CASE("test freeing a compiled closure erases every thread cache") {
+  // compile(std::function) keys a closure by the address of a heap copy of
+  // it. When the last copy of the compiled function is freed, its traces must
+  // go from every thread cache, not only from the cache of the thread that
+  // called compile(). A trace that stays keeps its arrays alive, and a later
+  // closure that gets the same heap address replays it.
+  //
+  // The trace holds the captured array, so the array is uniquely referenced
+  // again only when no cache holds the trace. The array is built from host
+  // data and the worker does not evaluate the output, so no device work holds
+  // a reference to it.
+  auto captured = array({1.0f, 2.0f, 3.0f});
+  auto compiled = compile([captured](const std::vector<array>& inputs) {
+    return std::vector<array>{inputs[0] + captured};
+  });
+
+  std::barrier phase(2);
+  std::thread worker([&]() {
+    {
+      // Record a trace in this thread's cache only.
+      auto outputs = compiled({array({1.0f, 1.0f, 1.0f})});
+    }
+    phase.arrive_and_wait();
+    // Keep this thread, and so its cache, alive while the closure is freed.
+    phase.arrive_and_wait();
+  });
+
+  phase.arrive_and_wait();
+  compiled = nullptr;
+  CHECK(captured.is_uniquely_referenced());
+  phase.arrive_and_wait();
+  worker.join();
+}