Phase 2: convert all op NIFs to async enif_send

.dialyzer_ignore.exs

@@ -26,7 +26,10 @@
   {"lib/emily/backend.ex", :callback_arg_type_mismatch},
   # Same root cause — `wrap/3`'s success typing ends up wider than the
   # declared `@spec wrap(ref(), tensor(), reference()) :: tensor()`.
-  {"lib/emily/backend.ex", :invalid_contract, 71},
+  # Line omitted because dialyzer's reported position varies with the
+  # inferred typing that flows in from Emily.Native (which became
+  # wider when the op NIFs were converted to async).
+  {"lib/emily/backend.ex", :invalid_contract},
 
   # ---------------------------------------------------------------
   # Emily.Fast — `Nx.Defn.Expr.optional/3` untyped return
@@ -53,7 +56,8 @@
   # `Nx.Type.t()` specific union. The NIF does return a valid
   # `Nx.Type.t()` tuple; the width is a declaration limitation in
   # the stub module.
-  {"lib/emily/quantization.ex", :invalid_contract, 67},
-  {"lib/emily/quantized_weight.ex", :invalid_contract, 78},
-  {"lib/emily/quantized_weight.ex", :invalid_contract, 112}
+  # Line numbers omitted: positions shift with inferred typing from
+  # Emily.Native, which widened when op NIFs moved to the async model.
+  {"lib/emily/quantization.ex", :invalid_contract},
+  {"lib/emily/quantized_weight.ex", :invalid_contract}
 ]

c_src/emily/async.hpp

@@ -140,4 +140,22 @@ fine::Term async_reply(ErlNifEnv *env,
   return fine::Term(ref_to_return);
 }
 
+// Convenience wrapper over async_reply for lambdas that return a
+// fine-encodable value (typically a `fine::ResourcePtr<Tensor>` or
+// a tuple of them). The helper wraps the computation, encodes the
+// result in msg_env, and posts `{ref, {:ok, encoded}}` back.
+//
+// `build` signature:
+//     (mx::Stream &stream) -> T   (T is any fine-encodable type)
+template <typename F>
+fine::Term async_encoded(ErlNifEnv *env,
+                         fine::ResourcePtr<WorkerThread> w,
+                         F &&build) {
+  return async_reply(
+      env, w,
+      [build = std::forward<F>(build)](mx::Stream &s, ErlNifEnv *msg_env) mutable {
+        return fine::encode(msg_env, build(s));
+      });
+}
+
 }  // namespace emily

c_src/ops/binary.cpp

@@ -1,29 +1,32 @@
 // Binary elementwise: arithmetic, compare, logical, bitwise.
 
+#include "../emily/async.hpp"
 #include "../emily/tensor.hpp"
 #include "../emily/worker.hpp"
 
 #include <fine.hpp>
 #include <mlx/mlx.h>
 
 namespace mx = mlx::core;
+using emily::async_encoded;
 using emily::Tensor;
-using emily::WorkerThread;
 using emily::wrap;
+using emily::WorkerThread;
 
 namespace {
 
-#define EMILY_BINARY(nif_name, mlx_fn)                                         \
-  fine::ResourcePtr<Tensor> nif_name(                                          \
-      ErlNifEnv *,                                                             \
+#define EMILY_BINARY(op_name, mlx_fn)                                          \
+  fine::Term op_name##_nif(                                                    \
+      ErlNifEnv *env,                                                          \
       fine::ResourcePtr<WorkerThread> w,                                       \
       fine::ResourcePtr<Tensor> a,                                             \
       fine::ResourcePtr<Tensor> b) {                                           \
-    return w->run_sync([&](mx::Stream &s) {                                    \
-      return wrap(mlx_fn(a->array, b->array, s));                              \
-    });                                                                        \
+    return async_encoded(env, w,                                               \
+        [a = std::move(a), b = std::move(b)](mx::Stream &s) {                  \
+          return wrap(mlx_fn(a->array, b->array, s));                          \
+        });                                                                    \
   }                                                                            \
-  FINE_NIF(nif_name, 0);
+  FINE_NIF(op_name##_nif, 0);
 
 // Arithmetic
 EMILY_BINARY(add,           mx::add)

c_src/ops/cast.cpp

@@ -1,5 +1,6 @@
 // Dtype cast.
 
+#include "../emily/async.hpp"
 #include "../emily/tensor.hpp"
 #include "../emily/worker.hpp"
 
@@ -10,33 +11,34 @@
 #include <tuple>
 
 namespace mx = mlx::core;
+using emily::async_encoded;
 using emily::Tensor;
-using emily::WorkerThread;
 using emily::to_mlx_dtype;
 using emily::wrap;
+using emily::WorkerThread;
 
 namespace {
 
-fine::ResourcePtr<Tensor> astype(
-    ErlNifEnv *,
+fine::Term astype_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     fine::ResourcePtr<Tensor> a,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [a = std::move(a), dtype](mx::Stream &s) {
     return wrap(mx::astype(a->array, to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(astype, 0);
+FINE_NIF(astype_nif, 0);
 
-fine::ResourcePtr<Tensor> bitcast(
-    ErlNifEnv *,
+fine::Term bitcast_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     fine::ResourcePtr<Tensor> a,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [a = std::move(a), dtype](mx::Stream &s) {
     return wrap(mx::view(a->array, to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(bitcast, 0);
+FINE_NIF(bitcast_nif, 0);
 
 } // namespace

c_src/ops/conv.cpp

@@ -1,5 +1,6 @@
 // Convolutions.
 
+#include "../emily/async.hpp"
 #include "../emily/tensor.hpp"
 #include "../emily/worker.hpp"
 
@@ -10,15 +11,16 @@
 #include <vector>
 
 namespace mx = mlx::core;
+using emily::async_encoded;
 using emily::Tensor;
-using emily::WorkerThread;
 using emily::to_int_vec;
 using emily::wrap;
+using emily::WorkerThread;
 
 namespace {
 
-fine::ResourcePtr<Tensor> conv_general(
-    ErlNifEnv *,
+fine::Term conv_general_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     fine::ResourcePtr<Tensor> input,
     fine::ResourcePtr<Tensor> weight,
@@ -27,14 +29,17 @@ fine::ResourcePtr<Tensor> conv_general(
     std::tuple<std::vector<int64_t>, std::vector<int64_t>> dilation,
     int64_t groups,
     bool flip) {
-  return w->run_sync([&](mx::Stream &s) {
-    return wrap(mx::conv_general(
-        input->array, weight->array, to_int_vec(stride),
-        to_int_vec(std::get<0>(padding)), to_int_vec(std::get<1>(padding)),
-        to_int_vec(std::get<0>(dilation)), to_int_vec(std::get<1>(dilation)),
-        static_cast<int>(groups), flip, s));
-  });
+  return async_encoded(env, w,
+      [input = std::move(input), weight = std::move(weight),
+       stride = std::move(stride), padding = std::move(padding),
+       dilation = std::move(dilation), groups, flip](mx::Stream &s) {
+        return wrap(mx::conv_general(
+            input->array, weight->array, to_int_vec(stride),
+            to_int_vec(std::get<0>(padding)), to_int_vec(std::get<1>(padding)),
+            to_int_vec(std::get<0>(dilation)), to_int_vec(std::get<1>(dilation)),
+            static_cast<int>(groups), flip, s));
+      });
 }
-FINE_NIF(conv_general, 0);
+FINE_NIF(conv_general_nif, 0);
 
 } // namespace

c_src/ops/creation.cpp

@@ -1,5 +1,6 @@
 // Creation ops: zeros, ones, full, arange, eye.
 
+#include "../emily/async.hpp"
 #include "../emily/tensor.hpp"
 #include "../emily/worker.hpp"
 
@@ -11,74 +12,76 @@
 #include <vector>
 
 namespace mx = mlx::core;
+using emily::async_encoded;
 using emily::Tensor;
-using emily::WorkerThread;
 using emily::to_mlx_dtype;
 using emily::to_mlx_shape;
 using emily::wrap;
+using emily::WorkerThread;
 
 namespace {
 
-fine::ResourcePtr<Tensor> zeros(
-    ErlNifEnv *,
+fine::Term zeros_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     std::vector<int64_t> shape,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [shape = std::move(shape), dtype](mx::Stream &s) {
     return wrap(mx::zeros(to_mlx_shape(shape), to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(zeros, 0);
+FINE_NIF(zeros_nif, 0);
 
-fine::ResourcePtr<Tensor> ones(
-    ErlNifEnv *,
+fine::Term ones_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     std::vector<int64_t> shape,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [shape = std::move(shape), dtype](mx::Stream &s) {
     return wrap(mx::ones(to_mlx_shape(shape), to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(ones, 0);
+FINE_NIF(ones_nif, 0);
 
-fine::ResourcePtr<Tensor> full(
-    ErlNifEnv *,
+fine::Term full_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     std::vector<int64_t> shape,
     fine::ResourcePtr<Tensor> value,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
-    return wrap(mx::full(to_mlx_shape(shape), value->array,
-                         to_mlx_dtype(dtype), s));
-  });
+  return async_encoded(env, w,
+      [shape = std::move(shape), value = std::move(value), dtype](mx::Stream &s) {
+        return wrap(mx::full(to_mlx_shape(shape), value->array,
+                             to_mlx_dtype(dtype), s));
+      });
 }
-FINE_NIF(full, 0);
+FINE_NIF(full_nif, 0);
 
-fine::ResourcePtr<Tensor> arange(
-    ErlNifEnv *,
+fine::Term arange_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     double start,
     double stop,
     double step,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [start, stop, step, dtype](mx::Stream &s) {
     return wrap(mx::arange(start, stop, step, to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(arange, 0);
+FINE_NIF(arange_nif, 0);
 
-fine::ResourcePtr<Tensor> eye(
-    ErlNifEnv *,
+fine::Term eye_nif(
+    ErlNifEnv *env,
     fine::ResourcePtr<WorkerThread> w,
     int64_t n,
     int64_t m,
     int64_t k,
     std::tuple<fine::Atom, int64_t> dtype) {
-  return w->run_sync([&](mx::Stream &s) {
+  return async_encoded(env, w, [n, m, k, dtype](mx::Stream &s) {
     return wrap(mx::eye(static_cast<int>(n), static_cast<int>(m),
                         static_cast<int>(k), to_mlx_dtype(dtype), s));
   });
 }
-FINE_NIF(eye, 0);
+FINE_NIF(eye_nif, 0);
 
 } // namespace