nnapi/
execution.rs

1use std::{
2    ffi::c_void,
3    mem::size_of,
4    ops::{Deref, DerefMut},
5    ptr::{null_mut, NonNull},
6};
7
8use nnapi_sys::{
9    ANeuralNetworksEvent, ANeuralNetworksExecution, ANeuralNetworksExecution_burstCompute,
10    ANeuralNetworksExecution_create, ANeuralNetworksExecution_free,
11    ANeuralNetworksExecution_setInput, ANeuralNetworksExecution_setOutput,
12    ANeuralNetworksExecution_startCompute, ResultCode,
13};
14
15use crate::{Burst, Compilation, Event, IntoResult};
16
17pub struct Execution {
18    inner: NonNull<ANeuralNetworksExecution>,
19}
20
21impl Execution {
22    pub fn new(compilation: &mut Compilation) -> crate::Result<Self> {
23        let mut execution = null_mut();
24
25        unsafe { ANeuralNetworksExecution_create(&mut **compilation, &mut execution) }
26            .into_result()?;
27
28        Ok(Execution {
29            inner: NonNull::new(execution).ok_or(ResultCode::ANEURALNETWORKS_UNEXPECTED_NULL)?,
30        })
31    }
32
33    #[inline]
34    pub fn set_input<T>(&mut self, input_list_idx: i32, input: &[T]) -> crate::Result<()> {
35        unsafe {
36            self.set_input_raw(
37                input_list_idx,
38                input.as_ptr().cast(),
39                input.len() * size_of::<T>(),
40            )
41        }
42    }
43
44    #[inline]
45    pub unsafe fn set_input_raw(
46        &mut self,
47        input_list_idx: i32,
48        buffer: *const c_void,
49        length: usize,
50    ) -> crate::Result<()> {
51        ANeuralNetworksExecution_setInput(&mut **self, input_list_idx, null_mut(), buffer, length)
52            .into_result()
53    }
54
55    #[inline]
56    pub fn set_output<T>(&mut self, output_list_idx: i32, output: &mut [T]) -> crate::Result<()> {
57        unsafe {
58            self.set_output_raw(
59                output_list_idx,
60                output.as_mut_ptr().cast(),
61                output.len() * size_of::<T>(),
62            )
63        }
64    }
65
66    #[inline]
67    pub unsafe fn set_output_raw(
68        &mut self,
69        output_list_idx: i32,
70        buffer: *mut c_void,
71        length: usize,
72    ) -> crate::Result<()> {
73        unsafe {
74            ANeuralNetworksExecution_setOutput(
75                &mut **self,
76                output_list_idx,
77                null_mut(),
78                buffer,
79                length,
80            )
81        }
82        .into_result()
83    }
84
85    #[inline]
86    pub fn compute(&mut self) -> crate::Result<Event> {
87        let mut end_event: *mut ANeuralNetworksEvent = null_mut();
88        unsafe { ANeuralNetworksExecution_startCompute(&mut **self, &mut end_event) }
89            .into_result()?;
90
91        Ok(Event {
92            inner: NonNull::new(end_event).ok_or(ResultCode::ANEURALNETWORKS_UNEXPECTED_NULL)?,
93        })
94    }
95
96    #[inline]
97    pub fn burst_compute(&mut self, burst: &mut Burst) -> crate::Result<()> {
98        unsafe { ANeuralNetworksExecution_burstCompute(&mut **self, &mut **burst) }.into_result()
99    }
100}
101
102impl Deref for Execution {
103    type Target = ANeuralNetworksExecution;
104
105    #[inline]
106    fn deref(&self) -> &Self::Target {
107        unsafe { self.inner.as_ref() }
108    }
109}
110
111impl DerefMut for Execution {
112    #[inline]
113    fn deref_mut(&mut self) -> &mut Self::Target {
114        unsafe { self.inner.as_mut() }
115    }
116}
117
118impl Drop for Execution {
119    fn drop(&mut self) {
120        unsafe {
121            ANeuralNetworksExecution_free(self.inner.as_mut());
122        }
123    }
124}