IREE Static Library Loader Runtime Demo (dylib-llvm-aot backend) (#6225)
This sample shows how to:
Produce a static library and bytecode module with IREE's compiler
Compile the static library into a program using the static_library_loader
Run the demo with the module using functions exported by the static library
The model compiled into the static library exports a single function simple_mul that returns the multiplication of two tensors. README.md includes full details on building and running the demo.
The sample is enabled through the IREE_BUILD_SAMPLES flag. Its tested/running on linux running the dylib-llvm-aot backend.
Fixes #5666
diff --git a/iree/samples/static_library/CMakeLists.txt b/iree/samples/static_library/CMakeLists.txt
new file mode 100644
index 0000000..4c818b8
--- /dev/null
+++ b/iree/samples/static_library/CMakeLists.txt
@@ -0,0 +1,81 @@
+# 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
+
+if(NOT ${IREE_TARGET_BACKEND_DYLIB-LLVM-AOT}
+ OR NOT ${IREE_HAL_DRIVER_DYLIB}
+ OR NOT ${IREE_BUILD_COMPILER})
+ return()
+endif()
+
+# Set iree-translate binary
+set(_TRANSLATE_TOOL_EXECUTABLE $<TARGET_FILE:iree_tools_iree-translate>)
+
+# Setup args for iree-transalte
+set(_TRANSLATE_ARGS)
+list(APPEND _TRANSLATE_ARGS "-iree-mlir-to-vm-bytecode-module")
+list(APPEND _TRANSLATE_ARGS "-iree-hal-target-backends=dylib-llvm-aot")
+list(APPEND _TRANSLATE_ARGS "-iree-llvm-static-library-output-path=simple_mul.o")
+list(APPEND _TRANSLATE_ARGS "${CMAKE_CURRENT_SOURCE_DIR}/simple_mul.mlir")
+list(APPEND _TRANSLATE_ARGS "-o")
+list(APPEND _TRANSLATE_ARGS "simple_mul.vmfb")
+
+# Custom command for iree-translate to generate static library and bytecode
+add_custom_command(
+ OUTPUT
+ ${CMAKE_CURRENT_BINARY_DIR}/simple_mul.h
+ ${CMAKE_CURRENT_BINARY_DIR}/simple_mul.o
+ ${CMAKE_CURRENT_BINARY_DIR}/simple_mul.vmfb
+ COMMAND ${_TRANSLATE_TOOL_EXECUTABLE} ${_TRANSLATE_ARGS}
+ DEPENDS ${_TRANSLATE_TOOL_EXECUTABLE} "simple_mul.mlir"
+)
+
+# Tell cmake about the simple_mul library so it will link it
+add_library(simple_mul
+ STATIC
+ ${CMAKE_CURRENT_BINARY_DIR}/simple_mul.o)
+
+SET_TARGET_PROPERTIES(
+ simple_mul
+ PROPERTIES
+ LINKER_LANGUAGE C)
+
+# Note: If you're cross compiling the simple_mul for a different backend, you'll
+# need to run iree-translate manually to produce the desired '.vmfb' and '.h/.o'
+# files. Substitute the 'simple_mul' dependency in iree_cc_binary() below with
+# your own static library and the `simple_mul.vmfb` in the iree_c_embed_data()
+# rule. You can use paths to files, i.e. 'path/to/generated/output.vmfb' in
+# place of CMake targets.
+
+# Generate the embed data with the bytecode module
+iree_c_embed_data(
+ NAME
+ simple_mul_c
+ IDENTIFIER
+ iree_samples_static_library_simple_mul
+ GENERATED_SRCS
+ simple_mul.vmfb
+ C_FILE_OUTPUT
+ simple_mul_c.c
+ H_FILE_OUTPUT
+ simple_mul_c.h
+ FLATTEN
+ PUBLIC
+ TESTONLY
+)
+
+iree_cc_binary(
+NAME
+ static_library_demo
+SRCS
+ "static_library_demo.c"
+DEPS
+ ::simple_mul_c
+ iree::runtime
+ iree::hal::local::loaders::static_library_loader
+ iree::hal::local::task_driver
+ iree::task::api
+ simple_mul
+)
diff --git a/iree/samples/static_library/README.md b/iree/samples/static_library/README.md
new file mode 100644
index 0000000..b989f0a
--- /dev/null
+++ b/iree/samples/static_library/README.md
@@ -0,0 +1,79 @@
+# IREE "Static Library" sample
+
+This sample shows how to:
+1. Produce a static library and bytecode module with IREE's compiler
+2. Compile the static library into a program using the `static_library_loader`
+3. Run the demo with the module using functions exported by the static library
+
+The model compiled into the static library exports a single function
+`simple_mul` that returns the multiplication of two tensors:
+
+```mlir
+func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
+{
+ %0 = "mhlo.multiply"(%arg0, %arg1) {name = "mul.1"} : (tensor<4xf32>,
+ tensor<4xf32>) -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
+```
+
+## Background
+
+IREE's `static_library_loader` allows applications to inject a set of static
+libraries that can be resolved at runtime by name. This can be particularly
+useful on "bare metal" or embedded systems running IREE that lack operating
+systems or the ability to load shared libraries in binaries.
+
+When static library output is enabled, `iree-translate` produces a separate
+static library to compile into the target program. At runtime bytecode module
+instructs the VM which static libraries to load exported functions from the
+model.
+
+## Instructions
+_Note: run the following commands from IREE's github repo root._
+
+1. Configure CMake for building the static library then demo. You'll need to set
+the flags building samples, the compiler, and the `dylib-llvm-aot`
+driver/backend. See
+[here](https://google.github.io/iree/building-from-source/getting-started/)
+for general instructions on building using CMake):
+
+ ```shell
+ cmake -B ../iree-build/
+ -DIREE_BUILD_SAMPLES=ON \
+ -DIREE_TARGET_BACKENDS_TO_BUILD=DYLIB-LLVM-AOT \
+ -DIREE_HAL_DRIVERS_TO_BUILD=DYLIB \
+ -DIREE_BUILD_COMPILER=ON \
+ -DCMAKE_BUILD_TYPE=RelWithDebInfo .
+ ```
+
+2. Build the `static_library_demo` CMake target to create the static demo. This
+target has several dependencies that will translate `simple_mul.mlir` into a
+static library (`simple_mul.h` & `simple_mul.c`) as well as a bytecode module
+(`simple_mul.vmfb`) which are finally built into the demo binary:
+
+ ```shell
+ cmake --build ../iree-build/ --target iree_samples_static_library_demo
+ ```
+
+3. Run the sample binary:
+
+ ```shell
+ ../iree-build/iree/samples/static_library/static_library_demo
+
+ # Output: static_library_run passed
+ ```
+
+### Changing compilation options
+
+The steps above build both the compiler for the host (machine doing the
+compiling) and the demo for the target using same options as the host machine.
+If you wish to target a different deployment other than the host, you'll need to
+compile the library and demo with different options.
+
+For example, see
+[documentation](https://google.github.io/iree/building-from-source/android/)
+on cross compiling on Android.
+
+Note: separating the target from the host will require modifying dependencies in
+the demos `CMakeLists.txt`. See included comments for more info.
diff --git a/iree/samples/static_library/simple_mul.mlir b/iree/samples/static_library/simple_mul.mlir
new file mode 100644
index 0000000..2b480bd
--- /dev/null
+++ b/iree/samples/static_library/simple_mul.mlir
@@ -0,0 +1,5 @@
+func @simple_mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32>
+{
+ %0 = "std.mulf"(%arg0, %arg1) {name = "mul.1"} : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+ return %0 : tensor<4xf32>
+}
diff --git a/iree/samples/static_library/static_library_demo.c b/iree/samples/static_library/static_library_demo.c
new file mode 100644
index 0000000..22e8127
--- /dev/null
+++ b/iree/samples/static_library/static_library_demo.c
@@ -0,0 +1,222 @@
+// 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
+
+// A example of static library loading in IREE. See the README.md for more info.
+// Note: this demo requires artifacts from iree-translate before it will run.
+
+#include <stdio.h>
+
+#include "iree/hal/local/loaders/static_library_loader.h"
+#include "iree/hal/local/task_device.h"
+#include "iree/modules/hal/module.h"
+#include "iree/runtime/api.h"
+#include "iree/samples/static_library/simple_mul_c.h"
+#include "iree/task/api.h"
+#include "iree/vm/bytecode_module.h"
+
+// Compiled static library module here to avoid IO:
+#include "iree/samples/static_library/simple_mul.h"
+
+// A function to create the HAL device from the different backend targets.
+// The HAL device is returned based on the implementation, and it must be
+// released by the caller.
+iree_status_t create_device_with_static_loader(iree_hal_device_t** device) {
+ iree_status_t status = iree_ok_status();
+ // Set paramters for the device created in the next step.
+ iree_hal_task_device_params_t params;
+ iree_hal_task_device_params_initialize(¶ms);
+
+ // Load the statically embedded library
+ const iree_hal_executable_library_header_t** static_library =
+ simple_mul_dispatch_0_library_query(
+ IREE_HAL_EXECUTABLE_LIBRARY_LATEST_VERSION, /*reserved=*/NULL);
+ const iree_hal_executable_library_header_t** libraries[1] = {static_library};
+
+ iree_hal_executable_loader_t* library_loader = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_hal_static_library_loader_create(
+ IREE_ARRAYSIZE(libraries), libraries,
+ iree_hal_executable_import_provider_null(), iree_allocator_system(),
+ &library_loader);
+ }
+
+ iree_task_executor_t* executor = NULL;
+ if (iree_status_is_ok(status)) {
+ iree_task_executor_create_from_flags(iree_allocator_system(), &executor);
+ }
+
+ // Create the device and release the executor and loader afterwards.
+ if (iree_status_is_ok(status)) {
+ iree_hal_task_device_create(iree_make_cstring_view("dylib"), ¶ms,
+ executor, 1, &library_loader,
+ iree_allocator_system(), device);
+ }
+ iree_task_executor_release(executor);
+ 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(IREE_API_VERSION_LATEST,
+ &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 dylib device with static loader.
+ iree_hal_device_t* device = NULL;
+ if (iree_status_is_ok(status)) {
+ status = create_device_with_static_loader(&device);
+ }
+ iree_vm_module_t* hal_module = NULL;
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_hal_module_create(device, iree_allocator_system(), &hal_module);
+ }
+
+ // 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 module from the embedded data. Append to the session.
+ const struct iree_file_toc_t* module_file_toc =
+ iree_samples_static_library_simple_mul_create();
+ iree_const_byte_span_t module_data =
+ iree_make_const_byte_span(module_file_toc->data, module_file_toc->size);
+ iree_vm_module_t* bytecode_module = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_vm_bytecode_module_create(module_data, iree_allocator_null(),
+ iree_allocator_system(),
+ &bytecode_module);
+ }
+ if (iree_status_is_ok(status)) {
+ status = iree_runtime_session_append_module(session, bytecode_module);
+ }
+
+ // Lookup the entry point function call.
+ const char kMainFunctionName[] = "module.simple_mul";
+ 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);
+ }
+
+ // Populate initial values for 4 * 2 = 8.
+ const int kElementCount = 4;
+ iree_hal_dim_t shape[1] = {kElementCount};
+ iree_hal_buffer_view_t* arg0_buffer_view = NULL;
+ iree_hal_buffer_view_t* arg1_buffer_view = NULL;
+ float kFloat4[] = {4.0f, 4.0f, 4.0f, 4.0f};
+ float kFloat2[] = {2.0f, 2.0f, 2.0f, 2.0f};
+
+ iree_hal_memory_type_t input_memory_type =
+ IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE;
+ if (iree_status_is_ok(status)) {
+ iree_hal_buffer_view_clone_heap_buffer(
+ iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape),
+ IREE_HAL_ELEMENT_TYPE_FLOAT_32, input_memory_type,
+ IREE_HAL_BUFFER_USAGE_ALL,
+ iree_make_const_byte_span((void*)kFloat4,
+ sizeof(float) * kElementCount),
+ &arg0_buffer_view);
+ }
+ if (iree_status_is_ok(status)) {
+ iree_hal_buffer_view_clone_heap_buffer(
+ iree_hal_device_allocator(device), shape, IREE_ARRAYSIZE(shape),
+ IREE_HAL_ELEMENT_TYPE_FLOAT_32, input_memory_type,
+ IREE_HAL_BUFFER_USAGE_ALL,
+ iree_make_const_byte_span((void*)kFloat2,
+ sizeof(float) * kElementCount),
+ &arg1_buffer_view);
+ }
+
+ // Queue buffer views for input.
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_runtime_call_inputs_push_back_buffer_view(&call, arg0_buffer_view);
+ }
+ iree_hal_buffer_view_release(arg0_buffer_view);
+
+ if (iree_status_is_ok(status)) {
+ status =
+ iree_runtime_call_inputs_push_back_buffer_view(&call, arg1_buffer_view);
+ }
+ iree_hal_buffer_view_release(arg1_buffer_view);
+
+ // Invoke call.
+ if (iree_status_is_ok(status)) {
+ status = iree_runtime_call_invoke(&call, /*flags=*/0);
+ }
+
+ // Retreive output buffer view with results from the invocation.
+ iree_hal_buffer_view_t* ret_buffer_view = NULL;
+ if (iree_status_is_ok(status)) {
+ status = iree_runtime_call_outputs_pop_front_buffer_view(&call,
+ &ret_buffer_view);
+ }
+
+ // Read back the results and ensure we got the right values.
+ iree_hal_buffer_mapping_t mapped_memory = {0};
+ if (iree_status_is_ok(status)) {
+ status = iree_hal_buffer_map_range(
+ iree_hal_buffer_view_buffer(ret_buffer_view),
+ IREE_HAL_MEMORY_ACCESS_READ, 0, IREE_WHOLE_BUFFER, &mapped_memory);
+ }
+ if (iree_status_is_ok(status)) {
+ if (mapped_memory.contents.data_length / sizeof(float) != kElementCount) {
+ status = iree_make_status(IREE_STATUS_UNKNOWN,
+ "result does not match element count ");
+ }
+ }
+ if (iree_status_is_ok(status)) {
+ const float* data = (const float*)mapped_memory.contents.data;
+ for (iree_host_size_t i = 0;
+ i < mapped_memory.contents.data_length / sizeof(float); ++i) {
+ if (data[i] != 8.0f) {
+ status = iree_make_status(IREE_STATUS_UNKNOWN, "result mismatches");
+ }
+ }
+ }
+
+ // Cleanup call and buffers.
+ iree_hal_buffer_unmap_range(&mapped_memory);
+ iree_hal_buffer_view_release(ret_buffer_view);
+ iree_runtime_call_deinitialize(&call);
+
+ // Cleanup session and instance.
+ iree_hal_device_release(device);
+ iree_runtime_session_release(session);
+ iree_runtime_instance_release(instance);
+
+ return status;
+}
+
+int main() {
+ const iree_status_t result = Run();
+ if (!iree_status_is_ok(result)) {
+ iree_status_fprint(stderr, result);
+ iree_status_free(result);
+ return -1;
+ }
+ printf("static_library_run passed\n");
+ return 0;
+}
