Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 1 | #ifndef SW_VEC_IREE_SAMPLES_UTIL_MODEL_API_H_ |
| 2 | #define SW_VEC_IREE_SAMPLES_UTIL_MODEL_API_H_ |
| 3 | |
| 4 | // Define ML model configuration and model-specific utility APIs. |
| 5 | |
| 6 | #include "iree/base/api.h" |
| 7 | #include "iree/hal/api.h" |
Lun Dong | 72890d0 | 2021-12-03 18:51:03 -0800 | [diff] [blame] | 8 | #include "iree/hal/local/executable_library.h" |
Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 9 | #include "iree/modules/hal/module.h" |
| 10 | #include "iree/vm/api.h" |
| 11 | #include "iree/vm/bytecode_module.h" |
| 12 | |
| 13 | #define MAX_MODEL_INPUT_NUM 2 |
| 14 | #define MAX_MODEL_INPUT_DIM 4 |
| 15 | #define MAX_MODEL_OUTPUTS 12 |
| 16 | #define MAX_ENTRY_FUNC_NAME 20 |
| 17 | |
| 18 | typedef struct { |
| 19 | int num_input; |
| 20 | int num_input_dim[MAX_MODEL_INPUT_NUM]; |
| 21 | iree_hal_dim_t input_shape[MAX_MODEL_INPUT_NUM][MAX_MODEL_INPUT_DIM]; |
| 22 | int input_length[MAX_MODEL_INPUT_NUM]; |
| 23 | int input_size_bytes[MAX_MODEL_INPUT_NUM]; |
| 24 | int num_output; |
| 25 | int output_length[MAX_MODEL_OUTPUTS]; |
| 26 | int output_size_bytes; |
| 27 | enum iree_hal_element_types_t hal_element_type; |
| 28 | char entry_func[MAX_ENTRY_FUNC_NAME]; |
| 29 | char model_name[]; |
| 30 | } MlModel; |
| 31 | |
Adam Jesionowski | da88093 | 2021-12-20 11:24:07 -0800 | [diff] [blame] | 32 | typedef struct { |
| 33 | void* result; |
| 34 | uint32_t len; |
| 35 | } MlOutput; |
| 36 | |
Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 37 | // Load the VM bytecode module from the embedded c library into memory. |
| 38 | const iree_const_byte_span_t load_bytecode_module_data(); |
| 39 | |
Lun Dong | 72890d0 | 2021-12-03 18:51:03 -0800 | [diff] [blame] | 40 | // Load the statically embedded library |
| 41 | const iree_hal_executable_library_header_t **library_query( |
| 42 | iree_hal_executable_library_version_t max_version, void *reserved); |
| 43 | |
| 44 | // Function to create the C module. |
| 45 | iree_status_t create_c_module(iree_vm_module_t **module); |
| 46 | |
Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 47 | // For each ML workload, based on the model configuration, allocate the buffer |
| 48 | // and prepare the data. It can be loaded from a embedded image binary, a |
| 49 | // randomly generated stream, or a pointer from the sensor/ISP output. |
Lun Dong | dbc0ab8 | 2022-01-07 18:18:10 +0000 | [diff] [blame^] | 50 | iree_status_t load_input_data(const MlModel *model, void **buffer, |
| 51 | iree_byte_span_t **byte_span); |
Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 52 | |
Adam Jesionowski | da88093 | 2021-12-20 11:24:07 -0800 | [diff] [blame] | 53 | // Process the ML execution output into the final data to be sent to the |
| 54 | // host. The final format is model dependent, so the address and size |
| 55 | // are returned via `output.` |
| 56 | iree_status_t process_output(const MlModel *model, |
| 57 | iree_hal_buffer_mapping_t *buffers, |
| 58 | MlOutput *output); |
Cindy Liu | 2449a16 | 2021-11-12 22:48:02 -0800 | [diff] [blame] | 59 | |
| 60 | #endif // SW_VEC_IREE_SAMPLES_UTIL_MODEL_API_H_ |