diff --git a/runtime/samples/CMakeLists.txt b/runtime/samples/CMakeLists.txt index efa1013..9eee09e 100644 --- a/runtime/samples/CMakeLists.txt +++ b/runtime/samples/CMakeLists.txt @@ -2,18 +2,14 @@ include(quidditch_module) quidditch_module(SRC simple_add.mlir) +add_subdirectory(util) + add_executable(IREE_HelloWorld main.c) target_link_libraries( IREE_HelloWorld PRIVATE - snRuntime - iree::base - iree::vm - iree::modules::hal - iree::modules::hal::types - iree::hal::local::local - Quidditch::device::device - Quidditch::loader::loader + samples_util simple_add_module + snRuntime ) diff --git a/runtime/samples/main.c b/runtime/samples/main.c index 49f6e4b..35a4e62 100644 --- a/runtime/samples/main.c +++ b/runtime/samples/main.c @@ -1,202 +1,47 @@ -#include -#include -#include - -#include -#include -#include -#include -#include - #include #include -#include #include -uint32_t snrt_l1_start_addr(); -uint32_t snrt_l1_end_addr(); - -static iree_allocator_inline_storage_t l1_arena; - -static iree_status_t setup_instance_and_device( - iree_allocator_t host_allocator, iree_vm_instance_t** out_instance, - iree_hal_device_t** out_device) { - IREE_ASSERT_ARGUMENT(out_instance); - IREE_ASSERT_ARGUMENT(out_device); - - IREE_RETURN_IF_ERROR(iree_vm_instance_create(IREE_VM_TYPE_CAPACITY_DEFAULT, - host_allocator, out_instance)); - - iree_status_t result = iree_hal_module_register_all_types(*out_instance); - if (!iree_status_is_ok(result)) goto error_release_vm; - - const iree_hal_executable_library_query_fn_t libraries[] = { - add_dispatch_0_library_query}; - - iree_hal_executable_loader_t* loader; - result = quidditch_loader_create(IREE_ARRAYSIZE(libraries), libraries, - iree_hal_executable_import_provider_null(), - host_allocator, &loader); - if (!iree_status_is_ok(result)) goto error_release_vm; - - l1_arena.buffer = (uint8_t*)snrt_l1_start_addr(); - l1_arena.length = 0; - unsigned stack_size_per_core = 1 << SNRT_LOG2_STACK_SIZE; - l1_arena.capacity = - (snrt_l1_end_addr() - snrt_cluster_core_num() * stack_size_per_core) - - snrt_l1_start_addr(); - - iree_hal_allocator_t* device_allocator; - result = - iree_hal_allocator_create_heap(iree_make_cstring_view("quidditch"), - iree_allocator_inline_arena(&l1_arena), - host_allocator, &device_allocator); - if (!iree_status_is_ok(result)) goto error_release_library_loader; - - quidditch_device_params_t params; - quidditch_device_params_initialize(¶ms); - result = - quidditch_device_create(IREE_SV("snitch"), ¶ms, - /*loader_count=*/1, &loader, device_allocator, - host_allocator, out_device); - iree_hal_executable_loader_release(loader); - iree_hal_allocator_release(device_allocator); - return result; - -error_release_library_loader: - iree_hal_executable_loader_release(loader); -error_release_vm: - iree_vm_instance_release(*out_instance); - return result; -} +#include "util/run_model.h" int main() { - if (!snrt_is_dm_core()) return quidditch_dispatch_enter_worker_loop(); - double data[4]; for (int i = 0; i < IREE_ARRAYSIZE(data); i++) { data[i] = (i + 1); } - iree_allocator_t host_allocator = iree_allocator_system(); - - iree_vm_instance_t* vmInstance; - iree_hal_device_t* device; - iree_status_t result = - setup_instance_and_device(host_allocator, &vmInstance, &device); - if (!iree_status_is_ok(result)) { - iree_status_fprint(stderr, result); - iree_status_free(result); - return -1; - } - - iree_vm_module_t* hal_module = NULL; - result = - iree_hal_module_create(vmInstance, /*device_count=*/1, - /*devices=*/&device, IREE_HAL_MODULE_FLAG_NONE, - host_allocator, &hal_module); - if (!iree_status_is_ok(result)) goto error_release_instance_and_device; - - iree_vm_module_t* mlir_module = NULL; - result = test_simple_add_create(vmInstance, host_allocator, &mlir_module); - if (!iree_status_is_ok(result)) goto error_release_hal_module; - - iree_vm_module_t* modules[] = {hal_module, mlir_module}; - - iree_vm_context_t* context; - result = iree_vm_context_create_with_modules( - vmInstance, IREE_VM_CONTEXT_FLAG_NONE, IREE_ARRAYSIZE(modules), modules, - host_allocator, &context); - if (!iree_status_is_ok(result)) goto error_release_mlir_module; - - iree_const_byte_span_t span = iree_make_const_byte_span(data, sizeof(data)); - - iree_hal_buffer_params_t params = { - .usage = IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE, - .access = IREE_HAL_MEMORY_ACCESS_NONE, - .type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL, + model_config_t config = { + .libraries = + (iree_hal_executable_library_query_fn_t[]){ + add_dispatch_0_library_query}, + .num_libraries = 1, + .module_constructor = test_simple_add_create, + .main_function = iree_make_cstring_view("test_simple_add.add"), + + .num_inputs = 2, + .input_data = (const double*[]){data, data}, + .input_sizes = (const iree_host_size_t[]){IREE_ARRAYSIZE(data), + IREE_ARRAYSIZE(data)}, + .input_ranks = (const iree_host_size_t[]){1, 1}, + .input_shapes = + (const iree_hal_dim_t*[]){(iree_hal_dim_t[]){IREE_ARRAYSIZE(data)}, + (iree_hal_dim_t[]){IREE_ARRAYSIZE(data)}}, + + .num_outputs = 1, + .output_data = (double*[]){data}, + .output_sizes = (const iree_host_size_t[]){IREE_ARRAYSIZE(data)}, + .device_allocator = l1_allocator(), }; - iree_hal_buffer_params_canonicalize(¶ms); - - iree_hal_buffer_view_t* buffer = NULL; - result = iree_hal_buffer_view_allocate_buffer_copy( - device, iree_hal_device_allocator(device), 1, - (iree_hal_dim_t[]){IREE_ARRAYSIZE(data)}, IREE_HAL_ELEMENT_TYPE_FLOAT_64, - IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, params, span, &buffer); - if (!iree_status_is_ok(result)) goto error_release_context; - - iree_vm_list_t* inputs = NULL; - result = iree_vm_list_create( - /*element_type=*/iree_vm_make_undefined_type_def(), - /*initial_capacity=*/2, iree_allocator_system(), &inputs); - if (!iree_status_is_ok(result)) goto error_release_context; - - iree_vm_ref_t arg_buffer_view_ref; - arg_buffer_view_ref = iree_hal_buffer_view_move_ref(buffer); - result = iree_vm_list_push_ref_retain(inputs, &arg_buffer_view_ref); - if (!iree_status_is_ok(result)) goto error_release_context; - - result = iree_vm_list_push_ref_move(inputs, &arg_buffer_view_ref); - if (!iree_status_is_ok(result)) goto error_release_context; - - iree_vm_function_t main_function; - IREE_CHECK_OK(iree_vm_context_resolve_function( - context, iree_make_cstring_view("test_simple_add.add"), &main_function)); - iree_vm_list_t* outputs = NULL; - IREE_CHECK_OK(iree_vm_list_create( - /*element_type=*/iree_vm_make_undefined_type_def(), - /*initial_capacity=*/1, iree_allocator_system(), &outputs)); - IREE_CHECK_OK(iree_vm_invoke( - context, main_function, IREE_VM_CONTEXT_FLAG_NONE, - /*policy=*/NULL, inputs, outputs, iree_allocator_system())); + IREE_CHECK_OK(run_model(&config)); - if (!iree_status_is_ok(result)) goto error_release_output; - - iree_hal_buffer_view_t* ret_buffer_view = - iree_vm_list_get_ref_deref(outputs, /*i=*/0, iree_hal_buffer_view_type()); - if (ret_buffer_view == NULL) goto error_release_output; - - iree_hal_buffer_mapping_t mapping; - 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, &mapping); - if (!iree_status_is_ok(result)) goto error_release_output; + if (!snrt_is_dm_core()) return 0; for (int i = 0; i < IREE_ARRAYSIZE(data); i++) { - double value = ((double*)mapping.contents.data)[i]; + double value = data[i]; printf("%f\n", value); if (value == (i + 1) * 2) continue; - - result = iree_make_status(IREE_STATUS_UNKNOWN, "output incorrect"); - break; } - - iree_hal_buffer_unmap_range(&mapping); - -error_release_output: - iree_vm_list_release(outputs); - iree_vm_list_release(inputs); -error_release_context: - iree_vm_context_release(context); -error_release_mlir_module: - iree_vm_module_release(mlir_module); -error_release_hal_module: - iree_vm_module_release(hal_module); - -error_release_instance_and_device: - iree_hal_device_release(device); - iree_vm_instance_release(vmInstance); - -exit: - if (!iree_status_is_ok(result)) { - iree_status_fprint(stderr, result); - iree_status_free(result); - return -1; - } - - quidditch_dispatch_quit(); return 0; } diff --git a/runtime/samples/util/CMakeLists.txt b/runtime/samples/util/CMakeLists.txt new file mode 100644 index 0000000..76fa625 --- /dev/null +++ b/runtime/samples/util/CMakeLists.txt @@ -0,0 +1,15 @@ + +add_library(samples_util run_model.c) +target_link_libraries( + samples_util + PUBLIC + iree::base + iree::vm + PRIVATE + snRuntime + iree::modules::hal + iree::modules::hal::types + iree::hal::local::local + Quidditch::device::device + Quidditch::loader::loader +) diff --git a/runtime/samples/util/run_model.c b/runtime/samples/util/run_model.c new file mode 100644 index 0000000..7c5c547 --- /dev/null +++ b/runtime/samples/util/run_model.c @@ -0,0 +1,181 @@ +#include "run_model.h" + +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +iree_allocator_t l1_allocator() { + uint32_t snrt_l1_start_addr(); + uint32_t snrt_l1_end_addr(); + + static iree_allocator_inline_storage_t l1_arena; + + l1_arena.buffer = (uint8_t*)snrt_l1_start_addr(); + l1_arena.length = 0; + unsigned stack_size_per_core = 1 << SNRT_LOG2_STACK_SIZE; + l1_arena.capacity = + (snrt_l1_end_addr() - snrt_cluster_core_num() * stack_size_per_core) - + snrt_l1_start_addr(); + + return iree_allocator_inline_arena(&l1_arena); +} + +static iree_status_t setup_instance_and_device( + const model_config_t* config, iree_allocator_t host_allocator, + iree_vm_instance_t** out_instance, iree_hal_device_t** out_device) { + IREE_ASSERT_ARGUMENT(out_instance); + IREE_ASSERT_ARGUMENT(out_device); + + IREE_RETURN_IF_ERROR(iree_vm_instance_create(IREE_VM_TYPE_CAPACITY_DEFAULT, + host_allocator, out_instance)); + + iree_status_t result = iree_hal_module_register_all_types(*out_instance); + if (!iree_status_is_ok(result)) goto error_release_vm; + + iree_hal_executable_loader_t* loader; + result = quidditch_loader_create(config->num_libraries, config->libraries, + iree_hal_executable_import_provider_null(), + host_allocator, &loader); + if (!iree_status_is_ok(result)) goto error_release_vm; + + iree_hal_allocator_t* device_allocator; + result = iree_hal_allocator_create_heap(iree_make_cstring_view("quidditch"), + config->device_allocator, + host_allocator, &device_allocator); + if (!iree_status_is_ok(result)) goto error_release_library_loader; + + quidditch_device_params_t params; + quidditch_device_params_initialize(¶ms); + result = + quidditch_device_create(IREE_SV("snitch"), ¶ms, + /*loader_count=*/1, &loader, device_allocator, + host_allocator, out_device); + iree_hal_executable_loader_release(loader); + iree_hal_allocator_release(device_allocator); + return result; + +error_release_library_loader: + iree_hal_executable_loader_release(loader); +error_release_vm: + iree_vm_instance_release(*out_instance); + return result; +} + +iree_status_t run_model(const model_config_t* config) { + if (!snrt_is_dm_core()) { + int ret = quidditch_dispatch_enter_worker_loop(); + if (!ret) return iree_ok_status(); + return iree_make_status(IREE_STATUS_UNKNOWN); + } + + iree_allocator_t host_allocator = iree_allocator_system(); + + iree_vm_instance_t* vmInstance; + iree_hal_device_t* device; + iree_status_t result = + setup_instance_and_device(config, host_allocator, &vmInstance, &device); + IREE_RETURN_IF_ERROR(result); + + iree_vm_module_t* hal_module = NULL; + result = + iree_hal_module_create(vmInstance, /*device_count=*/1, + /*devices=*/&device, IREE_HAL_MODULE_FLAG_NONE, + host_allocator, &hal_module); + if (!iree_status_is_ok(result)) goto error_release_instance_and_device; + + iree_vm_module_t* mlir_module = NULL; + result = config->module_constructor(vmInstance, host_allocator, &mlir_module); + if (!iree_status_is_ok(result)) goto error_release_hal_module; + + iree_vm_module_t* modules[] = {hal_module, mlir_module}; + + iree_vm_context_t* context; + result = iree_vm_context_create_with_modules( + vmInstance, IREE_VM_CONTEXT_FLAG_NONE, IREE_ARRAYSIZE(modules), modules, + host_allocator, &context); + if (!iree_status_is_ok(result)) goto error_release_mlir_module; + + iree_vm_list_t* inputs = NULL; + result = iree_vm_list_create( + /*element_type=*/iree_vm_make_undefined_type_def(), + /*initial_capacity=*/config->num_inputs, iree_allocator_system(), + &inputs); + if (!iree_status_is_ok(result)) goto error_release_context; + + for (iree_host_size_t i = 0; i < config->num_inputs; i++) { + iree_const_byte_span_t span = iree_make_const_byte_span( + config->input_data[i], config->input_sizes[i] * sizeof(double)); + + iree_hal_buffer_params_t params = { + .usage = IREE_HAL_BUFFER_USAGE_DISPATCH_STORAGE, + .access = IREE_HAL_MEMORY_ACCESS_NONE, + .type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL, + }; + iree_hal_buffer_params_canonicalize(¶ms); + + iree_hal_buffer_view_t* buffer = NULL; + result = iree_hal_buffer_view_allocate_buffer_copy( + device, iree_hal_device_allocator(device), config->input_ranks[i], + config->input_shapes[i], IREE_HAL_ELEMENT_TYPE_FLOAT_64, + IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR, params, span, &buffer); + if (!iree_status_is_ok(result)) goto error_release_context; + + iree_vm_ref_t arg_buffer_view_ref; + arg_buffer_view_ref = iree_hal_buffer_view_move_ref(buffer); + result = iree_vm_list_push_ref_retain(inputs, &arg_buffer_view_ref); + if (!iree_status_is_ok(result)) goto error_release_context; + } + + iree_vm_function_t main_function; + IREE_CHECK_OK(iree_vm_context_resolve_function(context, config->main_function, + &main_function)); + + iree_vm_list_t* outputs = NULL; + IREE_CHECK_OK(iree_vm_list_create( + /*element_type=*/iree_vm_make_undefined_type_def(), + /*initial_capacity=*/1, iree_allocator_system(), &outputs)); + IREE_CHECK_OK(iree_vm_invoke( + context, main_function, IREE_VM_CONTEXT_FLAG_NONE, + /*policy=*/NULL, inputs, outputs, iree_allocator_system())); + + if (!iree_status_is_ok(result)) goto error_release_output; + + for (iree_host_size_t i = 0; i < config->num_outputs; i++) { + iree_hal_buffer_view_t* ret_buffer_view = + iree_vm_list_get_ref_deref(outputs, i, iree_hal_buffer_view_type()); + if (ret_buffer_view == NULL) goto error_release_output; + + iree_hal_device_transfer_d2h( + device, iree_hal_buffer_view_buffer(ret_buffer_view), 0, + config->output_data[i], config->output_sizes[i] * sizeof(double), + IREE_HAL_TRANSFER_BUFFER_FLAG_DEFAULT, iree_infinite_timeout()); + } + +error_release_output: + iree_vm_list_release(outputs); + iree_vm_list_release(inputs); +error_release_context: + iree_vm_context_release(context); +error_release_mlir_module: + iree_vm_module_release(mlir_module); +error_release_hal_module: + iree_vm_module_release(hal_module); + +error_release_instance_and_device: + iree_hal_device_release(device); + iree_vm_instance_release(vmInstance); + + if (!iree_status_is_ok(result)) return result; + + quidditch_dispatch_quit(); + return 0; +} diff --git a/runtime/samples/util/run_model.h b/runtime/samples/util/run_model.h new file mode 100644 index 0000000..a2c018c --- /dev/null +++ b/runtime/samples/util/run_model.h @@ -0,0 +1,57 @@ + +#pragma once + +#include +#include +#include +#include +#include + +typedef struct { + /// Array of static library kernels that must be imported. + iree_hal_executable_library_query_fn_t* libraries; + /// Number of elements in 'libraries'. + iree_host_size_t num_libraries; + + /// EmitC host module created by iree-compile that should be executed. + iree_status_t (*module_constructor)(iree_vm_instance_t*, + iree_allocator_t host_allocator, + iree_vm_module_t** module_out); + /// Main function that should be called. + iree_string_view_t main_function; + + /// Number of input tensors. + iree_host_size_t num_inputs; + /// Input tensor data in dense row major encoding. + const double** input_data; + /// Number of elements for each input in 'input_data'. + const iree_host_size_t* input_sizes; + /// Ranks of each input tensor. + const iree_host_size_t* input_ranks; + /// Shapes for each input tensor. + const iree_hal_dim_t** input_shapes; + + /// Number of output tensors. + iree_host_size_t num_outputs; + /// Memory for the output data. + double** output_data; + /// Number of elements for each array in 'output_data'. + const iree_host_size_t* output_sizes; + + /// Allocator to use for "device" allocation. + iree_allocator_t device_allocator; +} model_config_t; + +/// Runs the given IREE module according to the config. Input and output data +/// are copied into and out of the given memory. +iree_status_t run_model(const model_config_t* config); + +/// Allocator which allocates within TCDM (L1) memory. Required for features +/// such as streaming registers. +iree_allocator_t l1_allocator(void); + +/// Allocator which allocates within DRAM (L3) memory. Uses 'malloc' and 'free' +/// behind the scenes. +static inline iree_allocator_t l3_allocator(void) { + return iree_allocator_system(); +}