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}