blob: a105a6feddeb514e2d0d3b7e784ff4ab232a4489 [file] [log] [blame]
// Copyright 2021 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// clang-format off
#cmakedefine IREE_STATIC_LIB_QUERY_FN @IREE_STATIC_LIB_QUERY_FN@
#cmakedefine IREE_STATIC_LIB_HDR @IREE_STATIC_LIB_HDR@
#cmakedefine IREE_MODULE_CREATE_FN @IREE_MODULE_CREATE_FN@
#cmakedefine IREE_MODULE_HDR "@IREE_MODULE_HDR@"
#cmakedefine IREE_EMITC_HDR "@IREE_EMITC_HDR@"
#cmakedefine IREE_MODULE_MAIN_FN @IREE_MODULE_MAIN_FN@
#cmakedefine IREE_INPUT_NUM @IREE_INPUT_NUM@
#cmakedefine IREE_INPUT_DIM_MAX @IREE_INPUT_DIM_MAX@
#cmakedefine IREE_INPUT_DIM_ARR @IREE_INPUT_DIM_ARR@
#cmakedefine IREE_INPUT_TYPE @IREE_INPUT_TYPE@
#cmakedefine IREE_INPUT_SHAPE_ARR @IREE_INPUT_SHAPE_ARR@
#cmakedefine IREE_INPUT_SIZE_ARR @IREE_INPUT_SIZE_ARR@
#cmakedefine IREE_HAL_TYPE @IREE_HAL_TYPE@
#cmakedefine IREE_EXE_NAME @IREE_EXE_NAME@
// clang-format on
#include <stdio.h>
#include IREE_STATIC_LIB_HDR
#include "iree/hal/drivers/local_sync/sync_device.h"
#include "iree/hal/local/loaders/static_library_loader.h"
#include "iree/modules/hal/module.h"
#include "iree/runtime/api.h"
#ifdef EMITC_IMPLEMENTATION
#include IREE_EMITC_HDR
#else
#include IREE_MODULE_HDR
#include "iree/vm/bytecode/module.h"
#endif
// A function to create the bytecode or c module.
iree_status_t create_module(iree_vm_instance_t* instance,
iree_vm_module_t** out_module) {
#ifdef EMITC_IMPLEMENTATION
return module_create(instance, iree_vm_instance_allocator(instance),
out_module);
#else
const struct iree_file_toc_t* module_file_toc =
IREE_MODULE_CREATE_FN;
iree_const_byte_span_t module_data =
iree_make_const_byte_span(module_file_toc->data, module_file_toc->size);
return iree_vm_bytecode_module_create(
instance, module_data, iree_allocator_null(),
iree_vm_instance_allocator(instance), out_module);
#endif
}
iree_status_t create_device_with_static_loader(iree_allocator_t host_allocator,
iree_hal_device_t** out_device) {
// Set paramters for the device created in the next step.
iree_hal_sync_device_params_t params;
iree_hal_sync_device_params_initialize(&params);
// Register the statically linked executable library.
const iree_hal_executable_library_query_fn_t libraries[] = {
IREE_STATIC_LIB_QUERY_FN,
};
iree_hal_executable_loader_t* library_loader = NULL;
iree_status_t status = iree_hal_static_library_loader_create(
IREE_ARRAYSIZE(libraries), libraries,
iree_hal_executable_import_provider_null(), host_allocator,
&library_loader);
// Use the default host allocator for buffer allocations.
iree_string_view_t identifier = iree_make_cstring_view("local-sync");
iree_hal_allocator_t* device_allocator = NULL;
if (iree_status_is_ok(status)) {
status = iree_hal_allocator_create_heap(identifier, host_allocator,
host_allocator, &device_allocator);
}
// Create the device and release the executor and loader afterwards.
if (iree_status_is_ok(status)) {
status = iree_hal_sync_device_create(
identifier, &params, /*loader_count=*/1, &library_loader,
device_allocator, host_allocator, out_device);
}
iree_hal_allocator_release(device_allocator);
iree_hal_executable_loader_release(library_loader);
return status;
}
iree_status_t Run() {
iree_status_t status = iree_ok_status();
// Instance configuration (this should be shared across sessions).
iree_runtime_instance_options_t instance_options;
iree_runtime_instance_options_initialize(&instance_options);
iree_runtime_instance_options_use_all_available_drivers(&instance_options);
iree_runtime_instance_t* instance = NULL;
if (iree_status_is_ok(status)) {
status = iree_runtime_instance_create(&instance_options,
iree_allocator_system(), &instance);
}
// Create local device with static loader.
iree_hal_device_t* device = NULL;
if (iree_status_is_ok(status)) {
status = create_device_with_static_loader(iree_allocator_system(), &device);
}
// Session configuration (one per loaded module to hold module state).
iree_runtime_session_options_t session_options;
iree_runtime_session_options_initialize(&session_options);
iree_runtime_session_t* session = NULL;
if (iree_status_is_ok(status)) {
status = iree_runtime_session_create_with_device(
instance, &session_options, device,
iree_runtime_instance_host_allocator(instance), &session);
}
// Load bytecode/c module from the embedded data. Append to the session.
iree_vm_module_t* module = NULL;
if (iree_status_is_ok(status)) {
status =
create_module(iree_runtime_instance_vm_instance(instance), &module);
}
if (iree_status_is_ok(status)) {
status = iree_runtime_session_append_module(session, module);
}
// Lookup the entry point function call.
const char kMainFunctionName[] = IREE_MODULE_MAIN_FN;
iree_runtime_call_t call;
memset(&call, 0, sizeof(call));
if (iree_status_is_ok(status)) {
status = iree_runtime_call_initialize_by_name(
session, iree_make_cstring_view(kMainFunctionName), &call);
}
// Prepare input.
iree_hal_dim_t shape[][IREE_INPUT_DIM_MAX] = IREE_INPUT_SHAPE_ARR;
iree_hal_dim_t size[] = IREE_INPUT_SIZE_ARR;
iree_hal_dim_t rank[] = IREE_INPUT_DIM_ARR;
iree_hal_buffer_view_t* arg_buffer_view = NULL;
IREE_INPUT_TYPE* arg_buffers[IREE_INPUT_NUM] = {NULL};
for (int i = 0; i < IREE_INPUT_NUM; ++i) {
if (iree_status_is_ok(status)) {
arg_buffers[i] = (IREE_INPUT_TYPE*)malloc(
sizeof(IREE_INPUT_TYPE) * size[i]);
if (arg_buffers[i] == NULL) {
status = iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED);
}
}
if (iree_status_is_ok(status)) {
status = iree_hal_buffer_view_allocate_buffer_copy(
device, iree_hal_device_allocator(device), rank[i], shape[i],
IREE_HAL_TYPE, IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR,
(iree_hal_buffer_params_t){
.type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL,
.usage = IREE_HAL_BUFFER_USAGE_DEFAULT,
},
iree_make_const_byte_span((void*)arg_buffers[i],
sizeof(IREE_INPUT_TYPE) * size[i]),
&arg_buffer_view);
}
// Queue buffer views for input.
if (iree_status_is_ok(status)) {
status =
iree_runtime_call_inputs_push_back_buffer_view(&call,
arg_buffer_view);
}
iree_hal_buffer_view_release(arg_buffer_view);
}
// Invoke call.
if (iree_status_is_ok(status)) {
status = iree_runtime_call_invoke(&call, /*flags=*/0);
}
// Cleanup call and buffers.
iree_runtime_call_deinitialize(&call);
for (int i = 0; i < IREE_INPUT_NUM; ++i) {
if (arg_buffers[i] != NULL) {
free(arg_buffers[i]);
}
}
// Cleanup session and instance.
iree_hal_device_release(device);
iree_runtime_session_release(session);
iree_runtime_instance_release(instance);
iree_vm_module_release(module);
return status;
}
int main() {
const char* kName = IREE_EXE_NAME;
const iree_status_t result = Run();
if (!iree_status_is_ok(result)) {
fprintf(stderr, "%s encountered error:\n", kName);
iree_status_fprint(stderr, result);
iree_status_free(result);
return -1;
}
fprintf(stdout, "%s finished successfully\n", kName);
return 0;
}