blob: 731f34e31ae5996652d42fd69e9ce22ee7dbc009 [file] [log] [blame]
// An example based on iree/samples/simple_embedding.
#include "samples/util/util.h"
#include <springbok.h>
#include <stdio.h>
#include "iree/base/api.h"
#include "iree/hal/api.h"
#include "iree/modules/hal/module.h"
#include "iree/vm/api.h"
#include "iree/vm/bytecode_module.h"
#include "samples/device/device.h"
extern const MlModel kModel;
// Prepare the input buffers and buffer_views based on the data type. They must
// be released by the caller.
static iree_status_t prepare_input_hal_buffer_views(
const MlModel *model, iree_hal_device_t *device, void **arg_buffers,
iree_hal_buffer_view_t **arg_buffer_views) {
iree_status_t result = iree_ok_status();
// Prepare the input buffer, and populate the initial value.
// The input buffer must be released by the caller.
iree_byte_span_t *byte_span[MAX_MODEL_INPUT_NUM] = {NULL};
result = load_input_data(model, arg_buffers, byte_span);
// Wrap buffers in shaped buffer views.
// The buffers can be mapped on the CPU and that can also be used
// on the device. Not all devices support this, but the ones we have now do.
iree_hal_memory_type_t input_memory_type =
IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
for (int i = 0; i < model->num_input; ++i) {
if (iree_status_is_ok(result)) {
result = iree_hal_buffer_view_wrap_or_clone_heap_buffer(
iree_hal_device_allocator(device), model->input_shape[i],
model->num_input_dim[i], model->hal_element_type,
IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, input_memory_type,
IREE_HAL_MEMORY_ACCESS_READ, IREE_HAL_BUFFER_USAGE_ALL, *byte_span[i],
iree_allocator_null(), &(arg_buffer_views[i]));
}
if (byte_span[i] != NULL) {
free(byte_span[i]);
}
}
return result;
}
iree_status_t run(const MlModel *model) {
IREE_RETURN_IF_ERROR(iree_hal_module_register_types());
iree_vm_instance_t *instance = NULL;
iree_status_t result =
iree_vm_instance_create(iree_allocator_system(), &instance);
iree_hal_device_t *device = NULL;
if (iree_status_is_ok(result)) {
result = create_sample_device(iree_allocator_system(), &device);
}
iree_vm_module_t *hal_module = NULL;
if (iree_status_is_ok(result)) {
result =
iree_hal_module_create(device, iree_allocator_system(), &hal_module);
}
// Load bytecode or C module.
iree_vm_module_t *module = NULL;
if (iree_status_is_ok(result)) {
result = create_module(&module);
}
// Allocate a context that will hold the module state across invocations.
iree_vm_context_t *context = NULL;
iree_vm_module_t *modules[] = {hal_module, module};
if (iree_status_is_ok(result)) {
result = iree_vm_context_create_with_modules(
instance, IREE_VM_CONTEXT_FLAG_NONE, &modules[0],
IREE_ARRAYSIZE(modules), iree_allocator_system(), &context);
}
iree_vm_module_release(hal_module);
iree_vm_module_release(module);
// Lookup the entry point function.
// Note that we use the synchronous variant which operates on pure type/shape
// erased buffers.
iree_vm_function_t main_function;
if (iree_status_is_ok(result)) {
result = (iree_vm_context_resolve_function(
context, iree_make_cstring_view(model->entry_func), &main_function));
}
// Prepare the input buffers.
void *arg_buffers[MAX_MODEL_INPUT_NUM] = {NULL};
iree_hal_buffer_view_t *arg_buffer_views[MAX_MODEL_INPUT_NUM] = {NULL};
if (iree_status_is_ok(result)) {
result = prepare_input_hal_buffer_views(model, device, arg_buffers,
arg_buffer_views);
}
// Setup call inputs with our buffers.
iree_vm_list_t *inputs = NULL;
if (iree_status_is_ok(result)) {
result = iree_vm_list_create(
/*element_type=*/NULL, /*capacity=*/model->num_input,
iree_allocator_system(), &inputs);
}
iree_vm_ref_t arg_buffer_view_ref;
for (int i = 0; i < model->num_input; ++i) {
arg_buffer_view_ref = iree_hal_buffer_view_move_ref(arg_buffer_views[i]);
if (iree_status_is_ok(result)) {
result = iree_vm_list_push_ref_move(inputs, &arg_buffer_view_ref);
}
}
// Prepare outputs list to accept the results from the invocation.
// The output vm list is allocated statically.
iree_vm_list_t *outputs = NULL;
if (iree_status_is_ok(result)) {
result = iree_vm_list_create(
/*element_type=*/NULL,
/*capacity=*/1, iree_allocator_system(), &outputs);
}
// Invoke the function.
if (iree_status_is_ok(result)) {
result = iree_vm_invoke(context, main_function, IREE_VM_CONTEXT_FLAG_NONE,
/*policy=*/NULL, inputs, outputs,
iree_allocator_system());
}
// Validate output and gather buffers.
iree_hal_buffer_mapping_t mapped_memories[MAX_MODEL_OUTPUTS] = {{0}};
for (int index_output = 0; index_output < model->num_output; index_output++) {
iree_hal_buffer_view_t *ret_buffer_view = NULL;
if (iree_status_is_ok(result)) {
// Get the result buffers from the invocation.
ret_buffer_view = (iree_hal_buffer_view_t *)iree_vm_list_get_ref_deref(
outputs, index_output, iree_hal_buffer_view_get_descriptor());
if (ret_buffer_view == NULL) {
result = iree_make_status(IREE_STATUS_NOT_FOUND,
"can't find return buffer view");
}
}
if (iree_status_is_ok(result)) {
result = iree_hal_buffer_map_range(
iree_hal_buffer_view_buffer(ret_buffer_view),
IREE_HAL_MAPPING_MODE_SCOPED, IREE_HAL_MEMORY_ACCESS_READ, 0,
IREE_WHOLE_BUFFER, &mapped_memories[index_output]);
}
if (iree_status_is_ok(result)) {
if (index_output > model->num_output ||
mapped_memories[index_output].contents.data_length / model->output_size_bytes !=
model->output_length[index_output]) {
result = iree_make_status(IREE_STATUS_UNKNOWN, "output length mismatches");
}
}
}
// Post-process memory into model output.
if (iree_status_is_ok(result)) {
MlOutput output = {.result = NULL, .len = 0};
result = process_output(model, mapped_memories, &output);
// TODO(jesionowski): Populate CSRs with `output` after validating result.
}
for (int index_output = 0; index_output < model->num_output; index_output++) {
if (mapped_memories[index_output].contents.data != NULL) {
iree_hal_buffer_unmap_range(&mapped_memories[index_output]);
}
}
iree_vm_list_release(inputs);
iree_vm_list_release(outputs);
for (int i = 0; i < model->num_input; ++i) {
if (arg_buffers[i] != NULL) {
iree_aligned_free(arg_buffers[i]);
}
}
iree_vm_context_release(context);
IREE_IGNORE_ERROR(iree_hal_allocator_statistics_fprint(
stdout, iree_hal_device_allocator(device)));
iree_hal_device_release(device);
iree_vm_instance_release(instance);
return result;
}
int main() {
const MlModel *model_ptr = &kModel;
const iree_status_t result = run(model_ptr);
int ret = (int)iree_status_code(result);
if (!iree_status_is_ok(result)) {
iree_status_fprint(stderr, result);
iree_status_free(result);
} else {
LOG_INFO("%s finished successfully", model_ptr->model_name);
}
return ret;
}