emily

// emily_nif.cpp — core NIFs: tensor resource, round-trip, eval.
//
// Op NIFs live in c_src/ops/*.cpp; they share the Tensor resource
// defined here via emily/tensor.hpp.

#include "emily/async.hpp"
#include "emily/tensor.hpp"
#include "emily/worker.hpp"

#include <fine.hpp>
#include <mlx/mlx.h>

#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

namespace mx = mlx::core;
using emily::Tensor;
using emily::WorkerThread;
using emily::from_mlx_dtype;
using emily::to_mlx_dtype;
using emily::to_mlx_shape;
using emily::wrap;

FINE_RESOURCE(Tensor);

// ---------- Core NIFs ----------

// from_binary/3 — build a lazy MLX array from a BEAM binary.
// Regular scheduler: MLX copies the buffer into its own storage during
// construction, so this is cheap and bounded.
//
// BEAM→MLX zero-copy is not possible with the current allocator API:
// on Metal, allocator::Buffer stores an MTL::Buffer*, so wrapping a
// BEAM heap pointer would crash on GPU dispatch.
fine::ResourcePtr<Tensor> from_binary(
    ErlNifEnv *,
    ErlNifBinary data,
    std::vector<int64_t> shape,
    std::tuple<fine::Atom, int64_t> dtype_tuple) {

  auto dtype = to_mlx_dtype(dtype_tuple);
  auto shape_ints = to_mlx_shape(shape);

  int64_t nelem = 1;
  for (auto d : shape_ints) {
    nelem *= d;
  }

  size_t expected = static_cast<size_t>(nelem) * dtype.size();
  if (data.size != expected) {
    throw std::invalid_argument(
        "binary size mismatch: expected " + std::to_string(expected) +
        " got " + std::to_string(data.size));
  }

  // Allocate an MLX-owned buffer, memcpy into it, hand ownership to
  // the array with a matching deleter. See comment above for why we
  // don't alias the BEAM binary directly.
  auto buf = mx::allocator::malloc(expected);
  std::memcpy(buf.raw_ptr(), data.data, expected);
  auto deleter = [](mx::allocator::Buffer b) { mx::allocator::free(b); };

  mx::array arr(buf, std::move(shape_ints), dtype, deleter);
  return wrap(std::move(arr));
}
FINE_NIF(from_binary, 0);

// to_binary_nif/2 — materialize the array on the worker thread and
// return its bytes as a BEAM resource binary aliasing MLX storage
// (no memcpy). Async: the NIF returns a ref immediately; the worker
// posts `{ref, {:ok, bin}}` back to the caller once `mx::contiguous +
// mx::eval` completes. The Elixir wrapper `Emily.Native.to_binary/2`
// awaits via `Emily.Native.Async.call/1`.
//
// Pinning: `fine::make_resource_binary` called on the worker-allocated
// msg_env bumps the Tensor resource's refcount; the binary carries
// that ref into the receiving process, where it stays alive until
// the binary is GC'd. Validated by Spike B.
fine::Term to_binary_nif(ErlNifEnv *env,
                         fine::ResourcePtr<WorkerThread> w,
                         fine::ResourcePtr<Tensor> tensor) {
  return emily::async_reply(
      env, w,
      [tensor = std::move(tensor)](mx::Stream &s, ErlNifEnv *msg_env) {
        auto materialized = mx::contiguous(tensor->array, false, s);
        mx::eval(materialized);

        if (!materialized.flags().row_contiguous) {
          throw std::runtime_error(
              "to_binary: array is not row-contiguous after mx::contiguous");
        }

        auto nbytes = materialized.nbytes();
        auto data = reinterpret_cast<const char *>(materialized.data<void>());

        auto pin = wrap(std::move(materialized));
        return fine::Term(
            fine::make_resource_binary(msg_env, std::move(pin), data, nbytes));
      });
}
FINE_NIF(to_binary_nif, 0);

// shape/1 — return the array's shape as a list of ints.
std::vector<int64_t> shape(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
  const auto &s = tensor->array.shape();
  return std::vector<int64_t>(s.begin(), s.end());
}
FINE_NIF(shape, 0);

// dtype/1 — return the array's dtype as an {atom, bits} tuple.
std::tuple<fine::Atom, int64_t> dtype(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
  return from_mlx_dtype(tensor->array.dtype());
}
FINE_NIF(dtype, 0);

// eval_nif/2 — force evaluation of the lazy graph rooted at this
// tensor. Async: the NIF enqueues the eval onto the worker and
// returns a ref synchronously; the worker posts {ref, {:ok, :ok}}
// back once eval completes. The Elixir wrapper `Emily.Native.eval/2`
// awaits via `Emily.Native.Async.call/1`.
//
// Runs on a regular scheduler — enqueueing is sub-microsecond and
// the scheduler is never blocked on MLX work.
fine::Term eval_nif(ErlNifEnv *env,
                    fine::ResourcePtr<WorkerThread> w,
                    fine::ResourcePtr<Tensor> tensor) {
  return emily::async_reply(
      env, w,
      [tensor](mx::Stream &, ErlNifEnv *msg_env) {
        mx::eval(tensor->array);
        return enif_make_atom(msg_env, "ok");
      });
}
FINE_NIF(eval_nif, 0);

FINE_INIT("Elixir.Emily.Native");