Skip to content

Commit

Permalink
Adding IREE_EXTERNAL_TOOLING_MODULES cmake flag. (#13367)
Browse files Browse the repository at this point in the history
This allows for custom modules to be linked into IREE tools for the
purposes of profiling, benchmarking, and debugging. As the IREE tools
are not to be deployed this is considered a development feature.
  • Loading branch information
benvanik authored May 2, 2023
1 parent edeb84c commit 5363ea3
Show file tree
Hide file tree
Showing 17 changed files with 753 additions and 29 deletions.
14 changes: 14 additions & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,20 @@ option(IREE_COMPILER_BUILD_SHARED_LIBS "Enables BUILD_SHARED_LIBS CMake mode for
# CI coverage is established.
option(BUILD_SHARED_LIBS "Instructs CMake to build libraries as shared if possible" OFF)

#-------------------------------------------------------------------------------
# IREE command-line tooling configuration
#-------------------------------------------------------------------------------

# External user modules linked into IREE tooling (iree-run-module/etc).
# These are only available in the internal IREE tools and here for testing
# custom modules in standard workflows. This is not a deployment mechanism and
# users integrating IREE into their libraries or applications will need to
# manage the building and registering/resolving the modules themselves.
#
# See runtime/src/iree/tooling/modules/CMakeLists.txt for more information on
# how to declare external modules.
set(IREE_EXTERNAL_TOOLING_MODULES "" CACHE STRING "")

#-------------------------------------------------------------------------------
# Compiler plugins
# IREE compiler plugins can be statically built into the compiler by adding
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/iree/hal/local/executable_plugin_manager.c
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
// as they don't disturb the existing values exposed to the plugin.

#define STATIC_ASSERT_EQ(a, b) \
static_assert((a) == (b), "plugin/runtime API mismatch")
static_assert((int)(a) == (int)(b), "plugin/runtime API mismatch")

STATIC_ASSERT_EQ(IREE_HAL_EXECUTABLE_PLUGIN_STATUS_OK, IREE_STATUS_OK);
STATIC_ASSERT_EQ(IREE_HAL_EXECUTABLE_PLUGIN_STATUS_CANCELLED,
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/iree/tooling/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ iree_runtime_cc_library(
"//runtime/src/iree/modules/hal",
"//runtime/src/iree/modules/hal/inline",
"//runtime/src/iree/modules/hal/loader",
"//runtime/src/iree/modules/vmvx",
"//runtime/src/iree/tooling/modules",
"//runtime/src/iree/vm",
"//runtime/src/iree/vm/bytecode:module",
"//runtime/src/iree/vm/dynamic:module",
Expand Down
2 changes: 1 addition & 1 deletion runtime/src/iree/tooling/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ iree_cc_library(
iree::modules::hal
iree::modules::hal::inline
iree::modules::hal::loader
iree::modules::vmvx
iree::tooling::modules
iree::vm
iree::vm::bytecode::module
iree::vm::dynamic::module
Expand Down
34 changes: 13 additions & 21 deletions runtime/src/iree/tooling/context_util.c
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,10 @@
#include "iree/modules/hal/loader/module.h"
#include "iree/modules/hal/module.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/modules/resolver.h"
#include "iree/vm/bytecode/module.h"
#include "iree/vm/dynamic/module.h"

#if defined(IREE_HAVE_VMVX_MODULE)
#include "iree/modules/vmvx/module.h"
#endif // IREE_HAVE_VMVX_MODULE

//===----------------------------------------------------------------------===//
// Module loading
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -383,7 +380,7 @@ typedef struct {
iree_hal_device_t* device;
iree_hal_allocator_t* device_allocator;
} iree_tooling_resolve_state_t;
static iree_status_t iree_tooling_resolve_module_dependency(
static iree_status_t iree_tooling_resolve_module_dependency_callback(
void* user_data_ptr, const iree_vm_module_dependency_t* dependency) {
iree_tooling_resolve_state_t* state =
(iree_tooling_resolve_state_t*)user_data_ptr;
Expand All @@ -393,9 +390,8 @@ static iree_status_t iree_tooling_resolve_module_dependency(
return iree_ok_status();
}

// Register one of the known modules. If we had a factory mechanism for
// resolving the modules we'd call out to that. Note that today this is not
// recursive but it could be in the future.
// Register one of the known modules. Note that today this is not recursive
// but it could be in the future.
iree_vm_module_t* module = NULL;
if (iree_string_view_equal(dependency->name, IREE_SV("hal"))) {
IREE_RETURN_IF_ERROR(iree_tooling_load_hal_async_module(
Expand All @@ -408,20 +404,12 @@ static iree_status_t iree_tooling_resolve_module_dependency(
} else if (iree_string_view_equal(dependency->name, IREE_SV("hal_loader"))) {
IREE_RETURN_IF_ERROR(iree_tooling_load_hal_loader_module(
state->instance, state->host_allocator, &module));
} else if (iree_string_view_equal(dependency->name, IREE_SV("vmvx"))) {
IREE_RETURN_IF_ERROR(iree_vmvx_module_create(
state->instance, state->host_allocator, &module));
} else if (iree_all_bits_set(dependency->flags,
IREE_VM_MODULE_DEPENDENCY_FLAG_REQUIRED)) {
// Required but not found; fail.
return iree_make_status(
IREE_STATUS_NOT_FOUND,
"required module '%.*s' not registered on the context",
(int)dependency->name.size, dependency->name.data);
} else {
// Optional and not found; skip.
return iree_ok_status();
// Defer to the generic module resolver registry.
IREE_RETURN_IF_ERROR(iree_tooling_resolve_module_dependency(
state->instance, dependency, state->host_allocator, &module));
}
if (!module) return iree_ok_status();

iree_status_t status =
iree_tooling_module_list_push_back(state->resolved_list, module);
Expand Down Expand Up @@ -463,7 +451,8 @@ iree_status_t iree_tooling_resolve_modules(
for (iree_host_size_t i = 0; i < user_module_count; ++i) {
iree_vm_module_t* user_module = user_modules[i];
status = iree_vm_module_enumerate_dependencies(
user_module, iree_tooling_resolve_module_dependency, &resolve_state);
user_module, iree_tooling_resolve_module_dependency_callback,
&resolve_state);
if (!iree_status_is_ok(status)) {
iree_string_view_t module_name = iree_vm_module_name(user_module);
(void)module_name;
Expand Down Expand Up @@ -553,6 +542,9 @@ iree_status_t iree_tooling_create_instance(iree_allocator_t host_allocator,
// HACK: to load modules we need the types registered even though we don't
// know if the types are used.
iree_status_t status = iree_hal_module_register_all_types(instance);
if (iree_status_is_ok(status)) {
status = iree_tooling_register_all_module_types(instance);
}

if (iree_status_is_ok(status)) {
*out_instance = instance;
Expand Down
25 changes: 25 additions & 0 deletions runtime/src/iree/tooling/modules/BUILD.bazel
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
# Copyright 2023 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

load("//build_tools/bazel:build_defs.oss.bzl", "iree_runtime_cc_library")

package(
default_visibility = ["//visibility:public"],
features = ["layering_check"],
licenses = ["notice"], # Apache 2.0
)

iree_runtime_cc_library(
name = "modules",
srcs = ["resolver.c"],
hdrs = ["resolver.h"],
deps = [
"//runtime/src/iree/base",
"//runtime/src/iree/base:tracing",
"//runtime/src/iree/modules/vmvx",
"//runtime/src/iree/vm",
],
)
142 changes: 142 additions & 0 deletions runtime/src/iree/tooling/modules/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
# Copyright 2023 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

# Link in externally defined modules.
# This allows users to conditionally enable modules that live outside of the
# IREE source tree by specifying a few cmake variables.
#
# Drivers are expected to have a CMakeLists.txt that is parsed when enabled.
# If a module is optional it may set an IREE_EXTERNAL_TOOLING_MODULE_{name}_FOUND
# variable to FALSE and be ignored, such as when dependencies are not found or
# other user configuration has disabled them.
#
# Each module provides a static library target name and a function that is
# called at runtime to register the module.
#
# Required variables:
# IREE_EXTERNAL_TOOLING_MODULE_{name}_TARGET: static library target name.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_NAME: module name in MLIR.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_CREATE: module creation name:
# iree_status_t {CREATE}(iree_vm_instance_t* instance,
# iree_allocator_t host_allocator,
# iree_vm_module_t** out_module)
# Optional variables:
# IREE_EXTERNAL_TOOLING_MODULE_{name}_OPTIONAL: true if the module not being
# found is not an error.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_SOURCE_DIR: source directory with a
# CMakeLists.txt included when the module is enabled.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_BINARY_DIR: binary directory for cmake
# outs.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_FOUND: bool to indicate whether the
# module was found and valid for use.
# IREE_EXTERNAL_TOOLING_MODULE_{name}_REGISTER_TYPES: optional type name:
# iree_status_t {REGISTER_TYPES}(iree_vm_instance_t* instance)
set(IREE_EXTERNAL_TOOLING_MODULES_USED)
foreach(_MODULE_NAME ${IREE_EXTERNAL_TOOLING_MODULES})
string(TOUPPER "IREE_EXTERNAL_TOOLING_MODULES_${_MODULE_NAME}" _MODULE_VAR)
string(REGEX REPLACE "-" "_" _MODULE_VAR ${_MODULE_VAR})
message(STATUS "Adding IREE external tooling module: ${_MODULE_NAME}")

set(_MODULE_OPTIONAL ${${_MODULE_VAR}_OPTIONAL})
set(_MODULE_SOURCE_DIR ${${_MODULE_VAR}_SOURCE_DIR})
set(_MODULE_BINARY_DIR ${${_MODULE_VAR}_BINARY_DIR})

# Default to found unless the user overrides it in the module source.
# This allows the module to decide to disable itself even if the user
# requested it.
set(${_MODULE_VAR}_FOUND TRUE CACHE BOOL
"Whether the external module is valid for use.")

# Include the module source CMakeLists.txt if required.
# Users may have already defined the targets and not need this.
if(_MODULE_SOURCE_DIR)
if(NOT EXISTS "${_MODULE_SOURCE_DIR}/CMakeLists.txt")
message(FATAL_ERROR "External module CMakeLists.txt not found at "
"${_MODULE_SOURCE_DIR}")
endif()
add_subdirectory(${_MODULE_SOURCE_DIR} ${_MODULE_BINARY_DIR})
endif()

# If found then add to the list of valid modules.
if(${${_MODULE_VAR}_FOUND})
list(APPEND IREE_EXTERNAL_TOOLING_MODULES_USED ${_MODULE_NAME})
else()
if(${_MODULE_OPTIONAL})
message(STATUS "Optional external module '${_MODULE_NAME}' requested "
"but not found; disabling and continuing")
else()
message(FATAL_ERROR "External module '${_MODULE_NAME}' not found; may "
"have unavailable dependencies")
endif()
endif()
endforeach()

# Produce resolve_external.c that contains a resolver switch for all modules.
# This will be called by resolver.c after internal modules are tried.
set(_RESOLVER_EXTERNAL_C_SRC)
set(_RESOLVER_EXTERNAL_COPTS)
set(_RESOLVER_EXTERNAL_DEPS)
if(IREE_EXTERNAL_TOOLING_MODULES_USED)
message(STATUS "Registering external tooling modules: ${IREE_EXTERNAL_TOOLING_MODULES_USED}")

set(_RESOLVER_EXTERNAL_COPTS "-DIREE_HAVE_EXTERNAL_TOOLING_MODULES=1")

# Build the list of deps and our source code lines.
set(_RESOLVER_EXTERNAL_DEPS)
set(_RESOLVER_EXTERNAL_DECLS)
set(_RESOLVER_EXTERNAL_TYPES)
set(_RESOLVER_EXTERNAL_CHECKS)
foreach(_MODULE_NAME ${IREE_EXTERNAL_TOOLING_MODULES_USED})
string(TOUPPER "IREE_EXTERNAL_TOOLING_MODULE_${_MODULE_NAME}" _MODULE_VAR)
string(REGEX REPLACE "-" "_" _MODULE_VAR ${_MODULE_VAR})
set(_MODULE_TARGET ${${_MODULE_VAR}_TARGET})
set(_MODULE_NAME ${${_MODULE_VAR}_NAME})
set(_MODULE_REGISTER_TYPES ${${_MODULE_VAR}_REGISTER_TYPES})
set(_MODULE_CREATE ${${_MODULE_VAR}_CREATE})
list(APPEND _RESOLVER_EXTERNAL_DEPS ${_MODULE_TARGET})
if(_MODULE_REGISTER_TYPES)
list(APPEND _RESOLVER_EXTERNAL_DECLS
"extern iree_status_t ${_MODULE_REGISTER_TYPES}(iree_vm_instance_t* instance);\n")
list(APPEND _RESOLVER_EXTERNAL_REGISTER_TYPES
"IREE_RETURN_IF_ERROR(${_MODULE_REGISTER_TYPES}(instance));\n")
endif()
list(APPEND _RESOLVER_EXTERNAL_DECLS
"extern iree_status_t ${_MODULE_CREATE}(iree_vm_instance_t* instance, iree_allocator_t host_allocator, iree_vm_module_t** out_module);\n")
list(APPEND _RESOLVER_EXTERNAL_TRY_CREATES
"if (iree_string_view_equal(dependency->name, IREE_SV(\"${_MODULE_NAME}\"))) {\n IREE_RETURN_IF_ERROR(${_MODULE_CREATE}(instance, host_allocator, out_module));\n }\n")
endforeach()

# Read template file and substitute variables.
set(_RESOLVER_EXTERNAL_C_TPL "${CMAKE_CURRENT_SOURCE_DIR}/resolve_external.c.in")
set(_RESOLVER_EXTERNAL_C_SRC "${CMAKE_CURRENT_BINARY_DIR}/resolve_external.c")
file(READ ${_RESOLVER_EXTERNAL_C_TPL} _RESOLVER_EXTERNAL_TEMPLATE)
file(
CONFIGURE OUTPUT ${_RESOLVER_EXTERNAL_C_SRC}
CONTENT "${_RESOLVER_EXTERNAL_TEMPLATE}"
)
endif()

set(_RESOLVER_INTERNAL_DEPS)
list(APPEND _RESOLVER_INTERNAL_DEPS iree::modules::vmvx)

iree_cc_library(
NAME
modules
HDRS
"resolver.h"
SRCS
"resolver.c"
${_RESOLVER_EXTERNAL_C_SRC}
COPTS
${_RESOLVER_EXTERNAL_COPTS}
DEPS
iree::base
iree::base::tracing
iree::vm
${_RESOLVER_INTERNAL_DEPS}
${_RESOLVER_EXTERNAL_DEPS}
PUBLIC
)
22 changes: 22 additions & 0 deletions runtime/src/iree/tooling/modules/resolve_external.c.in
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
// Copyright 2023 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

#include "iree/tooling/modules/resolver.h"

${_RESOLVER_EXTERNAL_DECLS}

iree_status_t iree_tooling_register_external_module_types(
iree_vm_instance_t* instance) {
${_RESOLVER_EXTERNAL_REGISTER_TYPES}
return iree_ok_status();
}

iree_status_t iree_tooling_try_resolve_external_module_dependency(
iree_vm_instance_t* instance, const iree_vm_module_dependency_t* dependency,
iree_allocator_t host_allocator, iree_vm_module_t** out_module) {
${_RESOLVER_EXTERNAL_TRY_CREATES}
return iree_ok_status();
}
74 changes: 74 additions & 0 deletions runtime/src/iree/tooling/modules/resolver.c
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
// Copyright 2023 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

#include "iree/tooling/modules/resolver.h"

#include "iree/base/tracing.h"

#if defined(IREE_HAVE_VMVX_MODULE)
#include "iree/modules/vmvx/module.h"
#endif // IREE_HAVE_VMVX_MODULE

#if defined(IREE_HAVE_EXTERNAL_TOOLING_MODULES)
// Defined in the generated registry_external.c file:
extern iree_status_t iree_tooling_register_external_module_types(
iree_vm_instance_t* instance);
extern iree_status_t iree_tooling_try_resolve_external_module_dependency(
iree_vm_instance_t* instance, const iree_vm_module_dependency_t* dependency,
iree_allocator_t host_allocator, iree_vm_module_t** out_module);
#else
static iree_status_t iree_tooling_register_external_module_types(
iree_vm_instance_t* instance) {
return iree_ok_status();
}
static iree_status_t iree_tooling_try_resolve_external_module_dependency(
iree_vm_instance_t* instance, const iree_vm_module_dependency_t* dependency,
iree_allocator_t host_allocator, iree_vm_module_t** out_module) {
*out_module = NULL;
return iree_ok_status();
}
#endif // IREE_HAVE_EXTERNAL_TOOLING_MODULES

iree_status_t iree_tooling_register_all_module_types(
iree_vm_instance_t* instance) {
IREE_RETURN_IF_ERROR(iree_tooling_register_external_module_types(instance));
return iree_ok_status();
}

iree_status_t iree_tooling_resolve_module_dependency(
iree_vm_instance_t* instance, const iree_vm_module_dependency_t* dependency,
iree_allocator_t host_allocator, iree_vm_module_t** out_module) {
IREE_ASSERT_ARGUMENT(instance);
IREE_ASSERT_ARGUMENT(dependency);
IREE_ASSERT_ARGUMENT(out_module);
*out_module = NULL;
IREE_TRACE_ZONE_BEGIN(z0);
IREE_TRACE_ZONE_APPEND_TEXT(z0, dependency->name.data, dependency->name.size);

iree_vm_module_t* module = NULL;
if (iree_string_view_equal(dependency->name, IREE_SV("vmvx"))) {
// VMVX module used on the host side for the inline HAL.
IREE_RETURN_AND_END_ZONE_IF_ERROR(
z0, iree_vmvx_module_create(instance, host_allocator, &module));
} else {
// Try to resolve the module from externally-defined modules.
// If the module is not found this will succeed but module will be NULL.
IREE_RETURN_AND_END_ZONE_IF_ERROR(
z0, iree_tooling_try_resolve_external_module_dependency(
instance, dependency, host_allocator, &module));
}

IREE_TRACE_ZONE_END(z0);
if (!module && iree_all_bits_set(dependency->flags,
IREE_VM_MODULE_DEPENDENCY_FLAG_REQUIRED)) {
// Required but not found; fail.
return iree_make_status(IREE_STATUS_NOT_FOUND,
"required module '%.*s' not available in the build",
(int)dependency->name.size, dependency->name.data);
}
*out_module = module;
return iree_ok_status();
}
Loading

0 comments on commit 5363ea3

Please sign in to comment.