mlx-sys 0.1.0-alpha

//! TODO: rename mod to something closer to transforms

use cxx::{CxxVector, UniquePtr};

use crate::array::ffi::array;

pub trait Function<Args> {
    type Output;

    fn execute(&self, args: Args) -> Self::Output;
}

impl<Args, F, R> Function<Args> for F
where
    F: Fn(Args) -> R,
{
    type Output = R;

    fn execute(&self, args: Args) -> Self::Output {
        self(args)
    }
}

#[repr(transparent)]
pub struct UnaryFn(pub Box<dyn for<'a> Function<&'a array, Output = UniquePtr<array>> + 'static>);

impl<F> From<F> for UnaryFn
where
    F: for<'a> Function<&'a array, Output = UniquePtr<array>> + 'static,
{
    fn from(f: F) -> Self {
        Self(Box::new(f))
    }
}

pub fn execute_unary_fn(f: &UnaryFn, args: &array) -> UniquePtr<array> {
    f.0.execute(args)
}

#[repr(transparent)]
pub struct MultiaryFn(
    pub  Box<
        dyn for<'a> Function<&'a CxxVector<array>, Output = UniquePtr<CxxVector<array>>> + 'static,
    >,
);

impl<F> From<F> for MultiaryFn
where
    F: for<'a> Function<&'a CxxVector<array>, Output = UniquePtr<CxxVector<array>>> + 'static,
{
    fn from(f: F) -> Self {
        Self(Box::new(f))
    }
}

pub fn execute_multiary_fn(f: &MultiaryFn, args: &CxxVector<array>) -> UniquePtr<CxxVector<array>> {
    f.0.execute(args)
}

#[repr(transparent)]
pub struct MultiInputSingleOutputFn(
    pub Box<dyn for<'a> Function<&'a CxxVector<array>, Output = UniquePtr<array>> + 'static>,
);

impl<F> From<F> for MultiInputSingleOutputFn
where
    F: for<'a> Function<&'a CxxVector<array>, Output = UniquePtr<array>> + 'static,
{
    fn from(f: F) -> Self {
        Self(Box::new(f))
    }
}

fn execute_multi_input_single_output_fn(
    f: &MultiInputSingleOutputFn,
    args: &CxxVector<array>,
) -> UniquePtr<array> {
    f.0.execute(args)
}

#[repr(transparent)]
pub struct PairInputSingleOutputFn(
    pub Box<dyn for<'a> Function<[&'a array; 2], Output = UniquePtr<array>> + 'static>,
);

impl<F> From<F> for PairInputSingleOutputFn
where
    F: for<'a> Function<[&'a array; 2], Output = UniquePtr<array>> + 'static,
{
    fn from(f: F) -> Self {
        Self(Box::new(f))
    }
}

// This runs into problem with the c++ binding if we use [&array; 2] as the input type
fn execute_pair_input_single_output_fn(
    f: &PairInputSingleOutputFn,
    first: &array,
    second: &array,
) -> UniquePtr<array> {
    f.0.execute([first, second])
}

// TODO: change visibility and then re-export
#[cxx::bridge]
pub mod ffi {
    extern "C++" {
        include!("mlx/array.h");

        #[namespace = "mlx::core"]
        type array = crate::array::ffi::array;
    }

    extern "Rust" {
        #[namespace = "mlx_cxx"]
        type UnaryFn;

        #[namespace = "mlx_cxx"]
        type MultiaryFn;

        #[namespace = "mlx_cxx"]
        type MultiInputSingleOutputFn;

        #[namespace = "mlx_cxx"]
        type PairInputSingleOutputFn;

        #[namespace = "mlx_cxx"]
        fn execute_unary_fn(f: &UnaryFn, x: &array) -> UniquePtr<array>;

        #[namespace = "mlx_cxx"]
        fn execute_multiary_fn(
            f: &MultiaryFn,
            xs: &CxxVector<array>,
        ) -> UniquePtr<CxxVector<array>>;

        #[namespace = "mlx_cxx"]
        fn execute_multi_input_single_output_fn(
            f: &MultiInputSingleOutputFn,
            xs: &CxxVector<array>,
        ) -> UniquePtr<array>;

        #[namespace = "mlx_cxx"]
        fn execute_pair_input_single_output_fn(
            f: &PairInputSingleOutputFn,
            first: &array,
            second: &array,
        ) -> UniquePtr<array>;
    }

    unsafe extern "C++" {
        include!("mlx-cxx/transforms.hpp");

        #[namespace = "mlx_cxx"]
        type CxxUnaryFn = crate::transforms::ffi::CxxUnaryFn;

        #[namespace = "mlx_cxx"]
        type CxxMultiaryFn = crate::transforms::ffi::CxxMultiaryFn;

        #[namespace = "mlx_cxx"]
        type CxxMultiInputSingleOutputFn = crate::transforms::ffi::CxxMultiInputSingleOutputFn;

        #[namespace = "mlx_cxx"]
        type CxxPairInputSingleOutputFn = crate::transforms::ffi::CxxPairInputSingleOutputFn;

        #[namespace = "mlx_cxx"]
        type CxxSingleInputPairOutputFn = crate::transforms::ffi::CxxSingleInputPairOutputFn;

        #[namespace = "mlx::core"]
        #[cxx_name = "ValueAndGradFn"]
        type CxxValueAndGradFn = crate::transforms::ffi::CxxValueAndGradFn;

        #[namespace = "mlx::core"]
        #[cxx_name = "SimpleValueAndGradFn"]
        type CxxSimpleValueAndGradFn = crate::transforms::ffi::CxxSimpleValueAndGradFn;

        // TODO: This is for test only. Remove later
        #[namespace = "mlx_cxx"]
        fn accept_rust_unary_fn(f: &UnaryFn) -> i32;

        #[namespace = "mlx_cxx"]
        #[rust_name = "vjp_unary_fn"]
        unsafe fn vjp(
            f: *const UnaryFn,
            primal: &array,
            cotangent: &array,
        ) -> Result<[UniquePtr<array>; 2]>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "vjp_multiary_fn"]
        unsafe fn vjp(
            f: *const MultiaryFn,
            primal: &[UniquePtr<array>],
            cotangent: &[UniquePtr<array>],
        ) -> Result<[UniquePtr<CxxVector<array>>; 2]>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "jvp_unary_fn"]
        unsafe fn jvp(
            f: *const UnaryFn,
            primal: &array,
            tangent: &array,
        ) -> Result<[UniquePtr<array>; 2]>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "jvp_multiary_fn"]
        unsafe fn jvp(
            f: *const MultiaryFn,
            primal: &[UniquePtr<array>],
            tangent: &[UniquePtr<array>],
        ) -> Result<[UniquePtr<CxxVector<array>>; 2]>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "value_and_grad_multiary_fn_argnums"]
        unsafe fn value_and_grad(
            f: *const MultiaryFn,
            argnums: &CxxVector<i32>,
        ) -> Result<UniquePtr<CxxValueAndGradFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "value_and_grad_multiary_fn_argnum"]
        unsafe fn value_and_grad(
            f: *const MultiaryFn,
            argnum: i32,
        ) -> Result<UniquePtr<CxxValueAndGradFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "value_and_grad_unary_fn"]
        unsafe fn value_and_grad(
            f: *const UnaryFn,
        ) -> Result<UniquePtr<CxxSingleInputPairOutputFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "value_and_grad_multi_input_single_output_fn"]
        unsafe fn value_and_grad(
            f: *const MultiInputSingleOutputFn,
            argnums: &CxxVector<i32>,
        ) -> Result<UniquePtr<CxxSimpleValueAndGradFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "grad_multi_input_single_output_fn_argnums"]
        unsafe fn grad(
            f: *const MultiInputSingleOutputFn,
            argnums: &CxxVector<i32>,
        ) -> Result<UniquePtr<CxxMultiaryFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "grad_multi_input_single_output_fn_argnum"]
        unsafe fn grad(
            f: *const MultiInputSingleOutputFn,
            argnum: i32,
        ) -> Result<UniquePtr<CxxMultiaryFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "grad_unary_fn"]
        unsafe fn grad(f: *const UnaryFn) -> Result<UniquePtr<CxxUnaryFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "vmap_unary_fn"]
        unsafe fn vmap(
            f: *const UnaryFn,
            in_axis: i32,
            out_axis: i32,
        ) -> Result<UniquePtr<CxxUnaryFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "vmap_pair_input_single_output_fn"]
        unsafe fn vmap(
            f: *const PairInputSingleOutputFn,
            in_axis_a: i32,
            in_axis_b: i32,
            out_axis: i32,
        ) -> Result<UniquePtr<CxxPairInputSingleOutputFn>>;

        #[namespace = "mlx_cxx"]
        #[rust_name = "vmap_multiary_fn"]
        unsafe fn vmap(
            f: *const MultiaryFn,
            in_axes: &CxxVector<i32>,
            out_axes: &CxxVector<i32>,
        ) -> Result<UniquePtr<CxxMultiaryFn>>;
    }
}

// #[cfg(test)]
// mod tests {
//     use crate::cxx_vec;

//     #[test]
//     fn test_accept_rust_unary_fn() {
//         let f = |x: &crate::array::ffi::array| -> cxx::UniquePtr<crate::array::ffi::array> {
//             crate::array::ffi::array_new_bool(true)
//         };
//         let f = super::UnaryFn::from(f);
//         let o = super::ffi::accept_rust_unary_fn(&f);
//         println!("{}", o);
//     }

//     #[test]
//     fn test_vjp_unary_fn() {
//         use std::sync::Arc;

//         let shape = cxx_vec!(3i32);
//         let b = Arc::new(crate::array::ffi::array_from_slice_float32(
//             &[1.0, 1.0, 1.0],
//             &shape,
//         ));
//         let f = move |arr: &crate::array::ffi::array| -> cxx::UniquePtr<crate::array::ffi::array> {
//             crate::ops::ffi::multiply(arr, &**b, Default::default())
//         };
//         let f = super::UnaryFn::from(f);

//         let primal = crate::array::ffi::array_from_slice_float32(&[1.0, 1.0, 1.0], &shape);
//         let cotangent = crate::array::ffi::array_from_slice_float32(&[1.0, 1.0, 1.0], &shape);
//         unsafe {
//             let f_ptr: *const crate::function::UnaryFn = &f;
//             let [p, c] = super::ffi::vjp_unary_fn(f_ptr, &primal, &cotangent);
//         }
//     }
// }