Packages

emily

0.1.2
1.0.0 0.7.2 0.7.1 0.6.1 0.5.1 0.4.0 0.3.5 0.3.3 0.3.2 0.1.2

Elixir bindings and Nx backend for Apple MLX

HexDocs

Current section

Files

Jump to
Readme
10 Versions
6 Dependencies
Dependants
Files
Activity
emily c_src ops binary.cpp
Raw

c_src/ops/binary.cpp 

// 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::wrap;

using emily::WorkerThread;



namespace {



#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 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(op_name##_nif, 0);



// Arithmetic

EMILY_BINARY(add,           mx::add)

EMILY_BINARY(subtract,      mx::subtract)

EMILY_BINARY(multiply,      mx::multiply)

EMILY_BINARY(divide,        mx::divide)

EMILY_BINARY(floor_divide,  mx::floor_divide)

EMILY_BINARY(remainder,     mx::remainder)

EMILY_BINARY(power,         mx::power)

EMILY_BINARY(maximum,       mx::maximum)

EMILY_BINARY(minimum,       mx::minimum)

EMILY_BINARY(logaddexp,     mx::logaddexp)

EMILY_BINARY(arctan2,       mx::arctan2)



// Compare

EMILY_BINARY(equal,         mx::equal)

EMILY_BINARY(not_equal,     mx::not_equal)

EMILY_BINARY(less,          mx::less)

EMILY_BINARY(less_equal,    mx::less_equal)

EMILY_BINARY(greater,       mx::greater)

EMILY_BINARY(greater_equal, mx::greater_equal)



// Logical

EMILY_BINARY(logical_and,   mx::logical_and)

EMILY_BINARY(logical_or,    mx::logical_or)



// Bitwise

EMILY_BINARY(bitwise_and,   mx::bitwise_and)

EMILY_BINARY(bitwise_or,    mx::bitwise_or)

EMILY_BINARY(bitwise_xor,   mx::bitwise_xor)

EMILY_BINARY(left_shift,    mx::left_shift)

EMILY_BINARY(right_shift,   mx::right_shift)



#undef EMILY_BINARY



} // namespace