Adding iree_vm_context_notify and signals to the VM. This allows for hosting programs to signal the modules in a VM of a system event. Both the system module interface (bytecode, future WASM, etc) and the user module are notified, with the latter being via an optional exported `__notify` function.
diff --git a/iree/modules/hal/module.c b/iree/modules/hal/module.c index ccd5cbc..4a4d6f3 100644 --- a/iree/modules/hal/module.c +++ b/iree/modules/hal/module.c
@@ -178,6 +178,20 @@ iree_allocator_free(state->host_allocator, state); } +static iree_status_t IREE_API_PTR iree_hal_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + iree_hal_module_state_t* state = (iree_hal_module_state_t*)module_state; + switch (signal) { + case IREE_VM_SIGNAL_SUSPEND: + case IREE_VM_SIGNAL_LOW_MEMORY: + // TODO(benvanik): trims for the deferred_releases list and our other + // tables. + return iree_hal_device_trim(state->shared_device); + default: + return iree_ok_status(); + } +} + //===----------------------------------------------------------------------===// // Experimental APIs //===----------------------------------------------------------------------===// @@ -1419,6 +1433,7 @@ .destroy = iree_hal_module_destroy, .alloc_state = iree_hal_module_alloc_state, .free_state = iree_hal_module_free_state, + .notify = iree_hal_module_notify, }; // Allocate shared module state.
diff --git a/iree/vm/bytecode_module.c b/iree/vm/bytecode_module.c index af15527..f64ac0f 100644 --- a/iree/vm/bytecode_module.c +++ b/iree/vm/bytecode_module.c
@@ -788,6 +788,11 @@ return iree_ok_status(); } +static iree_status_t IREE_API_PTR iree_vm_bytecode_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + return iree_ok_status(); +} + static iree_status_t iree_vm_bytecode_module_begin_call( void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) { @@ -914,6 +919,7 @@ module->interface.alloc_state = iree_vm_bytecode_module_alloc_state; module->interface.free_state = iree_vm_bytecode_module_free_state; module->interface.resolve_import = iree_vm_bytecode_module_resolve_import; + module->interface.notify = iree_vm_bytecode_module_notify; module->interface.begin_call = iree_vm_bytecode_module_begin_call; module->interface.get_function_reflection_attr = iree_vm_bytecode_module_get_function_reflection_attr;
diff --git a/iree/vm/context.c b/iree/vm/context.c index 4fc7938..e060dd6 100644 --- a/iree/vm/context.c +++ b/iree/vm/context.c
@@ -487,3 +487,127 @@ (int)module_name.size, module_name.data, (int)full_name.size, full_name.data); } + +// Calls the '__notify(i32)' function in |module|, if present. +static iree_status_t iree_vm_context_call_module_notify( + iree_vm_stack_t* stack, iree_vm_module_t* module, + iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + // Single i32 argument with the signal number. + uint32_t signal_arg = (uint32_t)signal; + iree_vm_function_call_t call; + memset(&call, 0, sizeof(call)); + call.arguments = iree_make_byte_span(&signal_arg, sizeof(signal_arg)); + + // Try to find the function. Modules are not required to export it. + iree_status_t status = iree_vm_module_lookup_function_by_name( + module, IREE_VM_FUNCTION_LINKAGE_EXPORT, + iree_make_cstring_view("__notify"), &call.function); + if (iree_status_is_not_found(status)) { + // Function doesn't exist; that's ok as this was an optional call. + return iree_status_ignore(status); + } else if (!iree_status_is_ok(status)) { + // Failed during trim. + return status; + } + + // Call the resolved function. + iree_vm_execution_result_t result; + status = module->begin_call(module->self, stack, &call, &result); + if (!iree_status_is_ok(status)) { + status = IREE_VM_STACK_ANNOTATE_BACKTRACE_IF_ENABLED(stack, status); + } + + // TODO(benvanik): ensure completed synchronously. + + return status; +} + +// Calls the module notify methods in registration order. +static iree_status_t iree_vm_context_notify_forward(iree_vm_stack_t* stack, + iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_ok_status(); + for (iree_host_size_t i = 0; i < context->list.count; ++i) { + iree_vm_module_t* module = context->list.modules[i]; + iree_vm_module_state_t* module_state = context->list.module_states[i]; + + // Call the module internal interface notify method. + // This handles the resources owned by the module implementation itself + // such as JITed binaries or other module infrastructure. + status = module->notify(module->self, module_state, signal); + if (!iree_status_is_ok(status)) break; + + // Call the user-level notify method. + // This may new use the reallocated resources from the module internal + // implementation above. + status = + iree_vm_context_call_module_notify(stack, module, module_state, signal); + if (!iree_status_is_ok(status)) break; + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +// Calls the module notify methods in reverse registration order. +static iree_status_t iree_vm_context_notify_reverse(iree_vm_stack_t* stack, + iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + iree_status_t status = iree_ok_status(); + for (int i = (int)context->list.count - 1; i >= 0; --i) { + iree_vm_module_t* module = context->list.modules[i]; + iree_vm_module_state_t* module_state = context->list.module_states[i]; + + // Call the user-level notify method first. + // This allows users to drop any state that they can rematerialize and + // return the resources to pools/caches to be trimmed below. + status = + iree_vm_context_call_module_notify(stack, module, module_state, signal); + if (!iree_status_is_ok(status)) break; + + // Call the module internal interface notify method. + // This handles the resources owned by the module implementation itself + // such as JITed binaries or other module infrastructure. Since we've + // already called the user-level function we likely have all of the + // resources that could be returned to pools there for this to reclaim. + status = module->notify(module->self, module_state, signal); + if (!iree_status_is_ok(status)) break; + } + IREE_TRACE_ZONE_END(z0); + return status; +} + +IREE_API_EXPORT iree_status_t iree_vm_context_notify(iree_vm_context_t* context, + iree_vm_signal_t signal) { + IREE_TRACE_ZONE_BEGIN(z0); + IREE_TRACE_ZONE_APPEND_VALUE(z0, (uint64_t)signal); + + // VM stack used to call into module __init methods. + IREE_VM_INLINE_STACK_INITIALIZE( + stack, + context->flags & IREE_VM_CONTEXT_FLAG_TRACE_EXECUTION + ? IREE_VM_INVOCATION_FLAG_TRACE_EXECUTION + : IREE_VM_INVOCATION_FLAG_NONE, + iree_vm_context_state_resolver(context), context->allocator); + + // Resumes are walked forward while suspends are walked backward. + // This follows the expected construction/destruction pattern where for + // example on suspend one would walk user modules to release resources back + // to system module pools before the system modules then clean up the pools. + iree_status_t status = iree_ok_status(); + switch (signal) { + default: + case IREE_VM_SIGNAL_RESUME: + status = iree_vm_context_notify_forward(stack, context, signal); + break; + case IREE_VM_SIGNAL_SUSPEND: + case IREE_VM_SIGNAL_LOW_MEMORY: + status = iree_vm_context_notify_reverse(stack, context, signal); + break; + } + + iree_vm_stack_deinitialize(stack); + IREE_TRACE_ZONE_END(z0); + return status; +}
diff --git a/iree/vm/context.h b/iree/vm/context.h index 71c36d3..b58bca6 100644 --- a/iree/vm/context.h +++ b/iree/vm/context.h
@@ -106,6 +106,10 @@ const iree_vm_context_t* context, iree_string_view_t full_name, iree_vm_function_t* out_function); +// Notifies all modules in the context of a system signal. +IREE_API_EXPORT iree_status_t iree_vm_context_notify(iree_vm_context_t* context, + iree_vm_signal_t signal); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus
diff --git a/iree/vm/module.h b/iree/vm/module.h index f164527..832d5aa 100644 --- a/iree/vm/module.h +++ b/iree/vm/module.h
@@ -24,11 +24,9 @@ typedef struct iree_vm_stack_t iree_vm_stack_t; typedef struct iree_vm_stack_frame_t iree_vm_stack_frame_t; -// An opaque offset into a source map that a source resolver can calculate. -// Do not assume that iree_vm_source_offset_t+1 means the next byte offset as -// backends are free to treat these as everything from pointers to machine code -// to hash codes. -typedef int64_t iree_vm_source_offset_t; +//===----------------------------------------------------------------------===// +// Module / function reflection +//===----------------------------------------------------------------------===// // A key-value pair of module/function reflection information. typedef struct iree_vm_reflection_attr_t { @@ -36,24 +34,6 @@ iree_string_view_t value; } iree_vm_reflection_attr_t; -// A variable-length list of registers. -// -// This structure is an overlay for the bytecode that is serialized in a -// matching format, though it can be stack allocated as needed. -// -// TODO(benvanik): this should be made private to the bytecode module, but is -// used for toll-free variadic argument lists here. We could just define an -// identical structure (and static_assert) to at least rename it to something -// sensible (iree_vm_segment_size_list_t). -typedef struct iree_vm_register_list_t { - uint16_t size; - uint16_t registers[]; -} iree_vm_register_list_t; -static_assert(iree_alignof(iree_vm_register_list_t) == 2, - "expecting byte alignment (to avoid padding)"); -static_assert(offsetof(iree_vm_register_list_t, registers) == 2, - "expect no padding in the struct"); - // Describes the type of a function reference. typedef enum iree_vm_function_linkage_e { // Function is internal to the module and may not be reflectable. @@ -136,6 +116,28 @@ // VM functions and accessing this state. typedef struct iree_vm_module_state_t iree_vm_module_state_t; +//===----------------------------------------------------------------------===// +// Function calls and coroutines +//===----------------------------------------------------------------------===// + +// A variable-length list of registers. +// +// This structure is an overlay for the bytecode that is serialized in a +// matching format, though it can be stack allocated as needed. +// +// TODO(benvanik): this should be made private to the bytecode module, but is +// used for toll-free variadic argument lists here. We could just define an +// identical structure (and static_assert) to at least rename it to something +// sensible (iree_vm_segment_size_list_t). +typedef struct iree_vm_register_list_t { + uint16_t size; + uint16_t registers[]; +} iree_vm_register_list_t; +static_assert(iree_alignof(iree_vm_register_list_t) == 2, + "expecting byte alignment (to avoid padding)"); +static_assert(offsetof(iree_vm_register_list_t, registers) == 2, + "expect no padding in the struct"); + // Function call data. // // Arguments and results are encoded following a standard format shared across @@ -261,6 +263,16 @@ int reserved; } iree_vm_execution_result_t; +//===----------------------------------------------------------------------===// +// Source locations +//===----------------------------------------------------------------------===// + +// An opaque offset into a source map that a source resolver can calculate. +// Do not assume that iree_vm_source_offset_t+1 means the next byte offset as +// backends are free to treat these as everything from pointers to machine code +// to hash codes. +typedef int64_t iree_vm_source_offset_t; + // Controls how source locations are formatted into strings. enum iree_vm_source_location_format_flag_bits_e { IREE_VM_SOURCE_LOCATION_FORMAT_FLAG_NONE = 0u, @@ -290,6 +302,34 @@ iree_vm_source_location_format_flags_t flags, iree_string_builder_t* builder); +//===----------------------------------------------------------------------===// +// iree_vm_module_t +//===----------------------------------------------------------------------===// + +// Indicates an event that can be signaled in modules from the hosting program. +typedef enum iree_vm_signal_e { + // Program is resuming from a suspended state. + // Modules may reallocate memory for pools and caches. + // + // Modules are walked in registration order (A->B->C). + IREE_VM_SIGNAL_RESUME = 0, + + // Program is entering a suspended state. + // Modules should drop any transient memory that is possible to reallocate + // upon resume. + // + // Modules are walked in reverse registration order (C->B->A). + IREE_VM_SIGNAL_SUSPEND = 1, + + // Program has received a low memory alert. + // Modules must aggressively drop all possible memory even if expensive to + // rematerialize it. On some platforms this is sent as a threat that if + // sufficient memory is not unwired/freed ASAP the process will be killed. + // + // Modules are walked in reverse registration order (C->B->A). + IREE_VM_SIGNAL_LOW_MEMORY = 2, +} iree_vm_signal_t; + // Defines an interface that can be used to reflect and execute functions on a // module. // @@ -349,6 +389,11 @@ iree_host_size_t ordinal, const iree_vm_function_t* function, const iree_vm_function_signature_t* signature); + // Notifies the module of a system signal. + iree_status_t(IREE_API_PTR* notify)(void* self, + iree_vm_module_state_t* module_state, + iree_vm_signal_t signal); + // Begins a function call with the given |call| arguments. // Execution may yield in the case of asynchronous code and require one or // more calls to the resume method to complete.
diff --git a/iree/vm/native_module.c b/iree/vm/native_module.c index e941089..c45a75c 100644 --- a/iree/vm/native_module.c +++ b/iree/vm/native_module.c
@@ -270,6 +270,15 @@ "native module does not support imports"); } +static iree_status_t IREE_API_PTR iree_vm_native_module_notify( + void* self, iree_vm_module_state_t* module_state, iree_vm_signal_t signal) { + iree_vm_native_module_t* module = (iree_vm_native_module_t*)self; + if (module->user_interface.notify) { + return module->user_interface.notify(module->self, module_state, signal); + } + return iree_ok_status(); +} + static iree_status_t IREE_API_PTR iree_vm_native_module_begin_call( void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) { @@ -428,6 +437,7 @@ module->base_interface.alloc_state = iree_vm_native_module_alloc_state; module->base_interface.free_state = iree_vm_native_module_free_state; module->base_interface.resolve_import = iree_vm_native_module_resolve_import; + module->base_interface.notify = iree_vm_native_module_notify; module->base_interface.begin_call = iree_vm_native_module_begin_call; module->base_interface.resume_call = iree_vm_native_module_resume_call;
diff --git a/iree/vm/native_module_cc.h b/iree/vm/native_module_cc.h index df1e3b9..015fdb8 100644 --- a/iree/vm/native_module_cc.h +++ b/iree/vm/native_module_cc.h
@@ -78,6 +78,7 @@ interface_.alloc_state = NativeModule::ModuleAllocState; interface_.free_state = NativeModule::ModuleFreeState; interface_.resolve_import = NativeModule::ModuleResolveImport; + interface_.notify = NativeModule::ModuleNotify; interface_.begin_call = NativeModule::ModuleBeginCall; } @@ -91,6 +92,11 @@ virtual StatusOr<std::unique_ptr<State>> CreateState( iree_allocator_t allocator) = 0; + // Notifies the module a signal has been raised. + virtual Status Notify(State* state, iree_vm_signal_t signal) { + return OkStatus(); + } + private: static NativeModule* FromModulePointer(void* self) { return reinterpret_cast<NativeModule*>(self); @@ -201,6 +207,13 @@ "C++ API does not support imports"); } + static iree_status_t ModuleNotify(void* self, + iree_vm_module_state_t* module_state, + iree_vm_signal_t signal) { + auto* module = FromModulePointer(self); + return module->Notify(FromStatePointer(module_state), signal); + } + static iree_status_t ModuleBeginCall(void* self, iree_vm_stack_t* stack, const iree_vm_function_call_t* call, iree_vm_execution_result_t* out_result) {