perf(nif): cache reply and dtype atoms at load

Atoms on the async-reply path (`:ok` on every op, plus
`:error`/`:stopped`/`:argument`/`:runtime`/`:unknown`) and the Nx dtype
"kind" atoms (`:f`/`:bf`/`:s`/`:u`/`:c`/`:pred`) were built at point of
use with enif_make_atom / fine::Atom("..."), redoing the atom-table
lookup on each call.

Define them once as static fine::Atom objects (emily/atoms.hpp). fine
registers namespace-scope atoms and creates their terms in the NIF load
callback, and Encoder<Atom> then returns the cached term directly. The
reply builders and from_mlx_dtype now reference these cached atoms.

Correctness-neutral (verified by the dtype and async tests); removes the
per-reply atom-table lookup from the hot path.

Closes #135

Author: ausimian

c_src/emily/async.hpp

@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "atoms.hpp"
 #include "worker.hpp"
 
 #include <fine.hpp>
@@ -66,16 +67,16 @@ error_reason_from_current_exception(ErlNifEnv *msg_env) {
   try {
     throw;  // re-raise the current exception to classify it
   } catch (const std::invalid_argument &e) {
-    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "argument"),
+    return enif_make_tuple2(msg_env, fine::encode(msg_env, emily::atoms::argument),
                             make_binary_from_cstr(msg_env, e.what()));
   } catch (const std::runtime_error &e) {
-    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
+    return enif_make_tuple2(msg_env, fine::encode(msg_env, emily::atoms::runtime),
                             make_binary_from_cstr(msg_env, e.what()));
   } catch (const std::exception &e) {
-    return enif_make_tuple2(msg_env, enif_make_atom(msg_env, "runtime"),
+    return enif_make_tuple2(msg_env, fine::encode(msg_env, emily::atoms::runtime),
                             make_binary_from_cstr(msg_env, e.what()));
   } catch (...) {
-    return enif_make_atom(msg_env, "unknown");
+    return fine::encode(msg_env, emily::atoms::unknown);
   }
 }
 
@@ -118,19 +119,19 @@ fine::Term async_reply(ErlNifEnv *env,
         // of hanging on a reply that will never come.
         reply = enif_make_tuple2(
             msg_env, ref_in_msg,
-            enif_make_tuple2(msg_env, enif_make_atom(msg_env, "error"),
-                             enif_make_atom(msg_env, "stopped")));
+            enif_make_tuple2(msg_env, fine::encode(msg_env, emily::atoms::error),
+                             fine::encode(msg_env, emily::atoms::stopped)));
       } else {
         try {
           ERL_NIF_TERM payload = build_payload(s, msg_env);
           ERL_NIF_TERM ok_tuple = enif_make_tuple2(
-              msg_env, enif_make_atom(msg_env, "ok"), payload);
+              msg_env, fine::encode(msg_env, emily::atoms::ok), payload);
           reply = enif_make_tuple2(msg_env, ref_in_msg, ok_tuple);
         } catch (...) {
           reply = enif_make_tuple2(
               msg_env, ref_in_msg,
               enif_make_tuple2(
-                  msg_env, enif_make_atom(msg_env, "error"),
+                  msg_env, fine::encode(msg_env, emily::atoms::error),
                   __async::error_reason_from_current_exception(msg_env)));
         }
       }

c_src/emily/atoms.hpp

@@ -0,0 +1,37 @@
+// Cached atom terms.
+//
+// fine::Atom objects defined at namespace scope register themselves on
+// construction and have their terms created once in the NIF load callback
+// (fine's load runs fine::__private__::init_atoms — see deps/fine).
+// Encoding such an atom (via fine::encode) then returns the prebuilt term
+// directly (Encoder<Atom> uses the cached `term`), with no per-call
+// atom-table lookup.
+//
+// Defining the reply and dtype atoms here — rather than building them at
+// point of use with enif_make_atom / fine::Atom("...") — keeps every atom
+// on the hot async-reply path (notably `:ok`, posted for every successful
+// op) off the lookup.
+
+#pragma once
+
+#include <fine.hpp>
+
+namespace emily::atoms {
+
+// Async reply atoms (see emily/async.hpp, emily_nif.cpp).
+inline auto ok = fine::Atom("ok");
+inline auto error = fine::Atom("error");
+inline auto stopped = fine::Atom("stopped");
+inline auto argument = fine::Atom("argument");
+inline auto runtime = fine::Atom("runtime");
+inline auto unknown = fine::Atom("unknown");
+
+// Nx dtype "kind" atoms (see emily/dtype.hpp).
+inline auto f = fine::Atom("f");
+inline auto bf = fine::Atom("bf");
+inline auto s = fine::Atom("s");
+inline auto u = fine::Atom("u");
+inline auto c = fine::Atom("c");
+inline auto pred = fine::Atom("pred");
+
+} // namespace emily::atoms

c_src/emily/dtype.hpp

@@ -6,6 +6,8 @@
 
 #pragma once
 
+#include "atoms.hpp"
+
 #include <fine.hpp>
 #include <mlx/mlx.h>
 
@@ -42,19 +44,19 @@ inline mx::Dtype to_mlx_dtype(const std::tuple<fine::Atom, int64_t> &t) {
 }
 
 inline std::tuple<fine::Atom, int64_t> from_mlx_dtype(mx::Dtype dtype) {
-  if (dtype == mx::float32)   return {fine::Atom("f"),    32};
-  if (dtype == mx::float16)   return {fine::Atom("f"),    16};
-  if (dtype == mx::bfloat16)  return {fine::Atom("bf"),   16};
-  if (dtype == mx::int8)      return {fine::Atom("s"),     8};
-  if (dtype == mx::int16)     return {fine::Atom("s"),    16};
-  if (dtype == mx::int32)     return {fine::Atom("s"),    32};
-  if (dtype == mx::int64)     return {fine::Atom("s"),    64};
-  if (dtype == mx::uint8)     return {fine::Atom("u"),     8};
-  if (dtype == mx::uint16)    return {fine::Atom("u"),    16};
-  if (dtype == mx::uint32)    return {fine::Atom("u"),    32};
-  if (dtype == mx::uint64)    return {fine::Atom("u"),    64};
-  if (dtype == mx::complex64) return {fine::Atom("c"),    64};
-  if (dtype == mx::bool_)     return {fine::Atom("pred"),  1};
+  if (dtype == mx::float32)   return {atoms::f,    32};
+  if (dtype == mx::float16)   return {atoms::f,    16};
+  if (dtype == mx::bfloat16)  return {atoms::bf,   16};
+  if (dtype == mx::int8)      return {atoms::s,     8};
+  if (dtype == mx::int16)     return {atoms::s,    16};
+  if (dtype == mx::int32)     return {atoms::s,    32};
+  if (dtype == mx::int64)     return {atoms::s,    64};
+  if (dtype == mx::uint8)     return {atoms::u,     8};
+  if (dtype == mx::uint16)    return {atoms::u,    16};
+  if (dtype == mx::uint32)    return {atoms::u,    32};
+  if (dtype == mx::uint64)    return {atoms::u,    64};
+  if (dtype == mx::complex64) return {atoms::c,    64};
+  if (dtype == mx::bool_)     return {atoms::pred,  1};
   throw std::runtime_error("unmapped mlx dtype");
 }
 

c_src/emily_nif.cpp

@@ -4,6 +4,7 @@
 // defined here via emily/tensor.hpp.
 
 #include "emily/async.hpp"
+#include "emily/atoms.hpp"
 #include "emily/tensor.hpp"
 #include "emily/worker.hpp"
 
@@ -138,7 +139,7 @@ fine::Term eval_nif(ErlNifEnv *env,
       env, w,
       [tensor](mx::Stream &, ErlNifEnv *msg_env) {
         mx::eval(tensor->array);
-        return enif_make_atom(msg_env, "ok");
+        return fine::encode(msg_env, emily::atoms::ok);
       });
 }
 FINE_NIF(eval_nif, 0);