apple_mpsgraph/

rnn.rs

use crate::error::{Error, Result};
use crate::ffi;
use crate::graph::Tensor;
use crate::types::collect_owned_tensors;
use core::ffi::{c_char, c_void};
use core::ptr;
use std::ffi::CString;

fn optional_cstring(name: Option<&str>) -> Option<CString> {
    name.and_then(|value| CString::new(value).ok())
}

#[allow(clippy::ref_option)]
fn cstring_ptr(value: &Option<CString>) -> *const c_char {
    value.as_ref().map_or(ptr::null(), |value| value.as_ptr())
}

fn optional_tensor_ptr(tensor: Option<&Tensor>) -> *mut c_void {
    tensor.map_or(ptr::null_mut(), Tensor::as_ptr)
}

fn wrap_tensor_array(box_handle: *mut c_void) -> Option<Vec<Tensor>> {
    if box_handle.is_null() {
        None
    } else {
        Some(collect_owned_tensors(box_handle))
    }
}

/// `MPSGraphRNNActivation` constants.
pub mod rnn_activation {
    pub const NONE: usize = 0;
    pub const RELU: usize = 1;
    pub const TANH: usize = 2;
    pub const SIGMOID: usize = 3;
    pub const HARD_SIGMOID: usize = 4;
}

macro_rules! descriptor_handle {
    ($name:ident) => {
        pub struct $name {
            ptr: *mut c_void,
        }

        unsafe impl Send for $name {}
        unsafe impl Sync for $name {}

        impl Drop for $name {
            fn drop(&mut self) {
                if !self.ptr.is_null() {
                    // SAFETY: `ptr` is a +1 retained Swift/ObjC object pointer owned by this wrapper.
                    unsafe { ffi::mpsgraph_object_release(self.ptr) };
                    self.ptr = ptr::null_mut();
                }
            }
        }

        impl $name {
            #[must_use]
            pub(crate) const fn as_ptr(&self) -> *mut c_void {
                self.ptr
            }
        }
    };
}

macro_rules! bool_getter_setter {
    ($getter:ident, $setter:ident, $ffi_get:ident, $ffi_set:ident, $msg:literal) => {
        #[must_use]
        pub fn $getter(&self) -> bool {
            // SAFETY: `self.ptr` is a live descriptor handle.
            unsafe { ffi::$ffi_get(self.ptr) }
        }

        pub fn $setter(&self, value: bool) -> Result<()> {
            // SAFETY: `self.ptr` is a live descriptor handle.
            let ok = unsafe { ffi::$ffi_set(self.ptr, value) };
            if ok {
                Ok(())
            } else {
                Err(Error::OperationFailed($msg))
            }
        }
    };
}

macro_rules! activation_getter_setter {
    ($getter:ident, $setter:ident, $ffi_get:ident, $ffi_set:ident, $msg:literal) => {
        #[must_use]
        pub fn $getter(&self) -> usize {
            // SAFETY: `self.ptr` is a live descriptor handle.
            unsafe { ffi::$ffi_get(self.ptr) }
        }

        pub fn $setter(&self, value: usize) -> Result<()> {
            // SAFETY: `self.ptr` is a live descriptor handle.
            let ok = unsafe { ffi::$ffi_set(self.ptr, value) };
            if ok {
                Ok(())
            } else {
                Err(Error::OperationFailed($msg))
            }
        }
    };
}

descriptor_handle!(SingleGateRNNDescriptor);
impl SingleGateRNNDescriptor {
    #[must_use]
    pub fn new() -> Option<Self> {
        // SAFETY: pure constructor.
        let ptr = unsafe { ffi::mpsgraph_single_gate_rnn_descriptor_new() };
        if ptr.is_null() {
            None
        } else {
            Some(Self { ptr })
        }
    }

    bool_getter_setter!(
        reverse,
        set_reverse,
        mpsgraph_single_gate_rnn_descriptor_reverse,
        mpsgraph_single_gate_rnn_descriptor_set_reverse,
        "failed to set single-gate RNN reverse"
    );
    bool_getter_setter!(
        bidirectional,
        set_bidirectional,
        mpsgraph_single_gate_rnn_descriptor_bidirectional,
        mpsgraph_single_gate_rnn_descriptor_set_bidirectional,
        "failed to set single-gate RNN bidirectional"
    );
    bool_getter_setter!(
        training,
        set_training,
        mpsgraph_single_gate_rnn_descriptor_training,
        mpsgraph_single_gate_rnn_descriptor_set_training,
        "failed to set single-gate RNN training"
    );
    activation_getter_setter!(
        activation,
        set_activation,
        mpsgraph_single_gate_rnn_descriptor_activation,
        mpsgraph_single_gate_rnn_descriptor_set_activation,
        "failed to set single-gate RNN activation"
    );
}

descriptor_handle!(LSTMDescriptor);
impl LSTMDescriptor {
    #[must_use]
    pub fn new() -> Option<Self> {
        // SAFETY: pure constructor.
        let ptr = unsafe { ffi::mpsgraph_lstm_descriptor_new() };
        if ptr.is_null() {
            None
        } else {
            Some(Self { ptr })
        }
    }

    bool_getter_setter!(
        reverse,
        set_reverse,
        mpsgraph_lstm_descriptor_reverse,
        mpsgraph_lstm_descriptor_set_reverse,
        "failed to set LSTM reverse"
    );
    bool_getter_setter!(
        bidirectional,
        set_bidirectional,
        mpsgraph_lstm_descriptor_bidirectional,
        mpsgraph_lstm_descriptor_set_bidirectional,
        "failed to set LSTM bidirectional"
    );
    bool_getter_setter!(
        produce_cell,
        set_produce_cell,
        mpsgraph_lstm_descriptor_produce_cell,
        mpsgraph_lstm_descriptor_set_produce_cell,
        "failed to set LSTM produceCell"
    );
    bool_getter_setter!(
        training,
        set_training,
        mpsgraph_lstm_descriptor_training,
        mpsgraph_lstm_descriptor_set_training,
        "failed to set LSTM training"
    );
    bool_getter_setter!(
        forget_gate_last,
        set_forget_gate_last,
        mpsgraph_lstm_descriptor_forget_gate_last,
        mpsgraph_lstm_descriptor_set_forget_gate_last,
        "failed to set LSTM forgetGateLast"
    );
    activation_getter_setter!(
        input_gate_activation,
        set_input_gate_activation,
        mpsgraph_lstm_descriptor_input_gate_activation,
        mpsgraph_lstm_descriptor_set_input_gate_activation,
        "failed to set LSTM inputGateActivation"
    );
    activation_getter_setter!(
        forget_gate_activation,
        set_forget_gate_activation,
        mpsgraph_lstm_descriptor_forget_gate_activation,
        mpsgraph_lstm_descriptor_set_forget_gate_activation,
        "failed to set LSTM forgetGateActivation"
    );
    activation_getter_setter!(
        cell_gate_activation,
        set_cell_gate_activation,
        mpsgraph_lstm_descriptor_cell_gate_activation,
        mpsgraph_lstm_descriptor_set_cell_gate_activation,
        "failed to set LSTM cellGateActivation"
    );
    activation_getter_setter!(
        output_gate_activation,
        set_output_gate_activation,
        mpsgraph_lstm_descriptor_output_gate_activation,
        mpsgraph_lstm_descriptor_set_output_gate_activation,
        "failed to set LSTM outputGateActivation"
    );
    activation_getter_setter!(
        activation,
        set_activation,
        mpsgraph_lstm_descriptor_activation,
        mpsgraph_lstm_descriptor_set_activation,
        "failed to set LSTM activation"
    );
}

descriptor_handle!(GRUDescriptor);
impl GRUDescriptor {
    #[must_use]
    pub fn new() -> Option<Self> {
        // SAFETY: pure constructor.
        let ptr = unsafe { ffi::mpsgraph_gru_descriptor_new() };
        if ptr.is_null() {
            None
        } else {
            Some(Self { ptr })
        }
    }

    bool_getter_setter!(
        reverse,
        set_reverse,
        mpsgraph_gru_descriptor_reverse,
        mpsgraph_gru_descriptor_set_reverse,
        "failed to set GRU reverse"
    );
    bool_getter_setter!(
        bidirectional,
        set_bidirectional,
        mpsgraph_gru_descriptor_bidirectional,
        mpsgraph_gru_descriptor_set_bidirectional,
        "failed to set GRU bidirectional"
    );
    bool_getter_setter!(
        training,
        set_training,
        mpsgraph_gru_descriptor_training,
        mpsgraph_gru_descriptor_set_training,
        "failed to set GRU training"
    );
    bool_getter_setter!(
        reset_gate_first,
        set_reset_gate_first,
        mpsgraph_gru_descriptor_reset_gate_first,
        mpsgraph_gru_descriptor_set_reset_gate_first,
        "failed to set GRU resetGateFirst"
    );
    bool_getter_setter!(
        reset_after,
        set_reset_after,
        mpsgraph_gru_descriptor_reset_after,
        mpsgraph_gru_descriptor_set_reset_after,
        "failed to set GRU resetAfter"
    );
    bool_getter_setter!(
        flip_z,
        set_flip_z,
        mpsgraph_gru_descriptor_flip_z,
        mpsgraph_gru_descriptor_set_flip_z,
        "failed to set GRU flipZ"
    );
    activation_getter_setter!(
        update_gate_activation,
        set_update_gate_activation,
        mpsgraph_gru_descriptor_update_gate_activation,
        mpsgraph_gru_descriptor_set_update_gate_activation,
        "failed to set GRU updateGateActivation"
    );
    activation_getter_setter!(
        reset_gate_activation,
        set_reset_gate_activation,
        mpsgraph_gru_descriptor_reset_gate_activation,
        mpsgraph_gru_descriptor_set_reset_gate_activation,
        "failed to set GRU resetGateActivation"
    );
    activation_getter_setter!(
        output_gate_activation,
        set_output_gate_activation,
        mpsgraph_gru_descriptor_output_gate_activation,
        mpsgraph_gru_descriptor_set_output_gate_activation,
        "failed to set GRU outputGateActivation"
    );
}

impl crate::graph::Graph {
    #[allow(clippy::too_many_arguments)]
    pub fn single_gate_rnn(
        &self,
        source: &Tensor,
        recurrent_weight: &Tensor,
        input_weight: Option<&Tensor>,
        bias: Option<&Tensor>,
        init_state: Option<&Tensor>,
        mask: Option<&Tensor>,
        descriptor: &SingleGateRNNDescriptor,
        name: Option<&str>,
    ) -> Option<Vec<Tensor>> {
        let name = optional_cstring(name);
        // SAFETY: all handles remain valid for the duration of the call.
        let box_handle = unsafe {
            ffi::mpsgraph_graph_single_gate_rnn(
                self.as_ptr(),
                source.as_ptr(),
                recurrent_weight.as_ptr(),
                optional_tensor_ptr(input_weight),
                optional_tensor_ptr(bias),
                optional_tensor_ptr(init_state),
                optional_tensor_ptr(mask),
                descriptor.as_ptr(),
                cstring_ptr(&name),
            )
        };
        wrap_tensor_array(box_handle)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn lstm(
        &self,
        source: &Tensor,
        recurrent_weight: &Tensor,
        input_weight: Option<&Tensor>,
        bias: Option<&Tensor>,
        init_state: Option<&Tensor>,
        init_cell: Option<&Tensor>,
        mask: Option<&Tensor>,
        peephole: Option<&Tensor>,
        descriptor: &LSTMDescriptor,
        name: Option<&str>,
    ) -> Option<Vec<Tensor>> {
        let name = optional_cstring(name);
        // SAFETY: all handles remain valid for the duration of the call.
        let box_handle = unsafe {
            ffi::mpsgraph_graph_lstm(
                self.as_ptr(),
                source.as_ptr(),
                recurrent_weight.as_ptr(),
                optional_tensor_ptr(input_weight),
                optional_tensor_ptr(bias),
                optional_tensor_ptr(init_state),
                optional_tensor_ptr(init_cell),
                optional_tensor_ptr(mask),
                optional_tensor_ptr(peephole),
                descriptor.as_ptr(),
                cstring_ptr(&name),
            )
        };
        wrap_tensor_array(box_handle)
    }

    #[allow(clippy::too_many_arguments)]
    pub fn gru(
        &self,
        source: &Tensor,
        recurrent_weight: &Tensor,
        input_weight: Option<&Tensor>,
        bias: Option<&Tensor>,
        init_state: Option<&Tensor>,
        mask: Option<&Tensor>,
        secondary_bias: Option<&Tensor>,
        descriptor: &GRUDescriptor,
        name: Option<&str>,
    ) -> Option<Vec<Tensor>> {
        let name = optional_cstring(name);
        // SAFETY: all handles remain valid for the duration of the call.
        let box_handle = unsafe {
            ffi::mpsgraph_graph_gru(
                self.as_ptr(),
                source.as_ptr(),
                recurrent_weight.as_ptr(),
                optional_tensor_ptr(input_weight),
                optional_tensor_ptr(bias),
                optional_tensor_ptr(init_state),
                optional_tensor_ptr(mask),
                optional_tensor_ptr(secondary_bias),
                descriptor.as_ptr(),
                cstring_ptr(&name),
            )
        };
        wrap_tensor_array(box_handle)
    }
}