| // Copyright 2019 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/vm/invocation.h" |
| |
| #include <stddef.h> |
| #include <stdint.h> |
| #include <string.h> |
| |
| #include "iree/base/api.h" |
| #include "iree/base/internal/debugging.h" |
| #include "iree/vm/ref.h" |
| #include "iree/vm/stack.h" |
| #include "iree/vm/value.h" |
| |
| //===----------------------------------------------------------------------===// |
| // Invocation utilities for I/O |
| //===----------------------------------------------------------------------===// |
| |
| // Releases reference counted values in |storage|. |
| static void iree_vm_invoke_release_io_refs(iree_string_view_t cconv_fragment, |
| iree_byte_span_t storage) { |
| if (!storage.data_length) return; |
| if (cconv_fragment.size == 0 || cconv_fragment.data[0] != '0') return; |
| uint8_t* p = storage.data; |
| for (iree_host_size_t i = 1; i < cconv_fragment.size; ++i) { |
| char c = cconv_fragment.data[i]; |
| switch (c) { |
| default: |
| IREE_ASSERT_UNREACHABLE("calling convention/FFI mismatch"); |
| break; |
| case IREE_VM_CCONV_TYPE_VOID: |
| break; |
| case IREE_VM_CCONV_TYPE_I32: |
| case IREE_VM_CCONV_TYPE_F32: |
| p += sizeof(int32_t); |
| break; |
| case IREE_VM_CCONV_TYPE_I64: |
| case IREE_VM_CCONV_TYPE_F64: |
| p += sizeof(int64_t); |
| break; |
| case IREE_VM_CCONV_TYPE_REF: |
| iree_vm_ref_release((iree_vm_ref_t*)p); |
| p += sizeof(iree_vm_ref_t); |
| break; |
| } |
| } |
| } |
| |
| // Releases storage for arguments. |
| static void iree_vm_invoke_release_argument_storage( |
| iree_string_view_t cconv_fragment, iree_byte_span_t storage, |
| bool is_heap_alloc, iree_allocator_t host_allocator) { |
| iree_vm_invoke_release_io_refs(cconv_fragment, storage); |
| if (is_heap_alloc) { |
| iree_allocator_free(host_allocator, storage.data); |
| } |
| } |
| |
| // Releases storage for results. |
| static void iree_vm_invoke_release_result_storage( |
| iree_string_view_t cconv_fragment, iree_byte_span_t storage, |
| void* stack_storage, iree_allocator_t host_allocator) { |
| iree_vm_invoke_release_io_refs(cconv_fragment, storage); |
| if (storage.data != stack_storage) { |
| iree_allocator_free(host_allocator, storage.data); |
| } |
| } |
| |
| // Marshals caller arguments from the variant list to the ABI convention. |
| static iree_status_t iree_vm_invoke_marshal_inputs( |
| iree_string_view_t cconv_arguments, const iree_vm_list_t* inputs, |
| iree_byte_span_t arguments) { |
| // We are 1:1 right now with no variadic args, so do a quick verification on |
| // the input list. |
| iree_host_size_t expected_input_count = |
| cconv_arguments.size > 0 |
| ? (cconv_arguments.data[0] == 'v' ? 0 : cconv_arguments.size) |
| : 0; |
| if (IREE_UNLIKELY(!inputs)) { |
| if (IREE_UNLIKELY(expected_input_count > 0)) { |
| return iree_make_status( |
| IREE_STATUS_INVALID_ARGUMENT, |
| "no input provided to a function that has inputs"); |
| } |
| return iree_ok_status(); |
| } else if (IREE_UNLIKELY(expected_input_count != iree_vm_list_size(inputs))) { |
| return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, |
| "input list and function mismatch; expected %zu " |
| "arguments but passed %zu", |
| expected_input_count, iree_vm_list_size(inputs)); |
| } |
| |
| uint8_t* p = arguments.data; |
| for (iree_host_size_t cconv_i = 0, arg_i = 0; cconv_i < cconv_arguments.size; |
| ++cconv_i, ++arg_i) { |
| switch (cconv_arguments.data[cconv_i]) { |
| case IREE_VM_CCONV_TYPE_VOID: |
| break; |
| case IREE_VM_CCONV_TYPE_I32: { |
| iree_vm_value_t value; |
| IREE_RETURN_IF_ERROR(iree_vm_list_get_value_as( |
| inputs, arg_i, IREE_VM_VALUE_TYPE_I32, &value)); |
| memcpy(p, &value.i32, sizeof(int32_t)); |
| p += sizeof(int32_t); |
| } break; |
| case IREE_VM_CCONV_TYPE_I64: { |
| iree_vm_value_t value; |
| IREE_RETURN_IF_ERROR(iree_vm_list_get_value_as( |
| inputs, arg_i, IREE_VM_VALUE_TYPE_I64, &value)); |
| memcpy(p, &value.i64, sizeof(int64_t)); |
| p += sizeof(int64_t); |
| } break; |
| case IREE_VM_CCONV_TYPE_F32: { |
| iree_vm_value_t value; |
| IREE_RETURN_IF_ERROR(iree_vm_list_get_value_as( |
| inputs, arg_i, IREE_VM_VALUE_TYPE_F32, &value)); |
| memcpy(p, &value.f32, sizeof(float)); |
| p += sizeof(float); |
| } break; |
| case IREE_VM_CCONV_TYPE_F64: { |
| iree_vm_value_t value; |
| IREE_RETURN_IF_ERROR(iree_vm_list_get_value_as( |
| inputs, arg_i, IREE_VM_VALUE_TYPE_F64, &value)); |
| memcpy(p, &value.f64, sizeof(double)); |
| p += sizeof(double); |
| } break; |
| case IREE_VM_CCONV_TYPE_REF: { |
| // TODO(benvanik): see if we can't remove this retain by instead relying |
| // on the caller still owning the list. |
| IREE_RETURN_IF_ERROR( |
| iree_vm_list_get_ref_retain(inputs, arg_i, (iree_vm_ref_t*)p)); |
| p += sizeof(iree_vm_ref_t); |
| } break; |
| } |
| } |
| return iree_ok_status(); |
| } |
| |
| // Marshals callee results from the ABI convention to the variant list. |
| static iree_status_t iree_vm_invoke_marshal_outputs( |
| iree_string_view_t cconv_results, iree_byte_span_t results, |
| iree_vm_list_t* outputs) { |
| iree_host_size_t expected_output_count = |
| cconv_results.size > 0 |
| ? (cconv_results.data[0] == 'v' ? 0 : cconv_results.size) |
| : 0; |
| if (IREE_UNLIKELY(!outputs)) { |
| if (IREE_UNLIKELY(expected_output_count > 0)) { |
| return iree_make_status( |
| IREE_STATUS_INVALID_ARGUMENT, |
| "no output provided to a function that has outputs"); |
| } |
| return iree_ok_status(); |
| } |
| |
| // Resize the output list to hold all results (and kill anything that may |
| // have been in there). |
| // TODO(benvanik): list method for resetting to a new size. |
| IREE_RETURN_IF_ERROR(iree_vm_list_resize(outputs, 0)); |
| IREE_RETURN_IF_ERROR(iree_vm_list_resize(outputs, expected_output_count)); |
| |
| uint8_t* p = results.data; |
| for (iree_host_size_t cconv_i = 0, arg_i = 0; cconv_i < cconv_results.size; |
| ++cconv_i, ++arg_i) { |
| switch (cconv_results.data[cconv_i]) { |
| case IREE_VM_CCONV_TYPE_VOID: |
| break; |
| case IREE_VM_CCONV_TYPE_I32: { |
| iree_vm_value_t value = iree_vm_value_make_i32(*(int32_t*)p); |
| IREE_RETURN_IF_ERROR(iree_vm_list_set_value(outputs, arg_i, &value)); |
| p += sizeof(int32_t); |
| } break; |
| case IREE_VM_CCONV_TYPE_I64: { |
| iree_vm_value_t value = iree_vm_value_make_i64(*(int64_t*)p); |
| IREE_RETURN_IF_ERROR(iree_vm_list_set_value(outputs, arg_i, &value)); |
| p += sizeof(int64_t); |
| } break; |
| case IREE_VM_CCONV_TYPE_F32: { |
| iree_vm_value_t value = iree_vm_value_make_f32(*(float*)p); |
| IREE_RETURN_IF_ERROR(iree_vm_list_set_value(outputs, arg_i, &value)); |
| p += sizeof(float); |
| } break; |
| case IREE_VM_CCONV_TYPE_F64: { |
| iree_vm_value_t value = iree_vm_value_make_f64(*(double*)p); |
| IREE_RETURN_IF_ERROR(iree_vm_list_set_value(outputs, arg_i, &value)); |
| p += sizeof(double); |
| } break; |
| case IREE_VM_CCONV_TYPE_REF: { |
| IREE_RETURN_IF_ERROR( |
| iree_vm_list_set_ref_move(outputs, arg_i, (iree_vm_ref_t*)p)); |
| p += sizeof(iree_vm_ref_t); |
| } break; |
| } |
| } |
| return iree_ok_status(); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Fiber tracing support |
| //===----------------------------------------------------------------------===// |
| |
| // Fibers are tricky things to instrument as tooling support is often lacking. |
| // We support two major modes (beyond when tracing is entirely disabled): |
| // IREE_TRACING_FEATURE_FIBERS: use Tracy's native fiber support. |
| // Does not support concurrent/interleaved coroutines. |
| // !IREE_TRACING_FEATURE_FIBERS: emulated support by trace stack fiddling. |
| // Supports concurrent/interleaved coroutines but messes with statistics |
| // as the trace stack is suspended/resumed and zones get extra counts. |
| // |
| // To make concurrent coroutines work when Tracy's fiber support is enabled we |
| // go from treating each context as a fiber to treating each invocation as one. |
| // This has the side-effect of creating one fiber per invocation and in |
| // benchmarks that can be really noisy; best that can be done there is disabling |
| // native fiber support. |
| static iree_vm_invocation_id_t iree_vm_invoke_allocate_id( |
| iree_vm_context_t* context, const iree_vm_function_t* function) { |
| #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_FIBERS |
| if (iree_vm_context_flags(context) & IREE_VM_CONTEXT_FLAG_CONCURRENT) { |
| // Native Tracy fiber support does not handle interleaved coroutines. |
| // Instead we'll allocate a unique ID per invocation. |
| // The string must remain live for the lifetime of the process. |
| // TODO(benvanik): name it based on the function? |
| static iree_atomic_int32_t next_invocation_id = IREE_ATOMIC_VAR_INIT(1); |
| uint32_t invocation_id = iree_atomic_fetch_add_int32( |
| &next_invocation_id, 1, iree_memory_order_relaxed); |
| IREE_LEAK_CHECK_DISABLE_PUSH(); |
| char* name = (char*)malloc(32); |
| snprintf(name, 32, "invoke-%04d", invocation_id - 1); |
| IREE_LEAK_CHECK_DISABLE_POP(); |
| return (iree_vm_invocation_id_t)name; |
| } else { |
| // Non-concurrent (sequential) execution can just reuse the context ID. |
| return (iree_vm_invocation_id_t)iree_vm_context_id(context); |
| } |
| #else |
| return (iree_vm_invocation_id_t)iree_vm_context_id(context); |
| #endif // IREE_TRACING_FEATURE_FIBERS |
| } |
| |
| #if IREE_TRACING_FEATURES & IREE_TRACING_FEATURE_INSTRUMENTATION |
| |
| static void iree_vm_invoke_fiber_enter(iree_vm_invocation_id_t invocation_id) { |
| if (!invocation_id) return; |
| IREE_TRACE_FIBER_ENTER(invocation_id); |
| } |
| |
| static void iree_vm_invoke_fiber_reenter(iree_vm_invocation_id_t invocation_id, |
| iree_vm_stack_t* stack) { |
| if (!invocation_id) return; |
| IREE_TRACE_FIBER_ENTER(invocation_id); |
| iree_vm_stack_resume_trace_zones(stack); |
| } |
| |
| static void iree_vm_invoke_fiber_leave(iree_vm_invocation_id_t invocation_id, |
| iree_vm_stack_t* stack) { |
| if (!invocation_id) return; |
| if (stack) iree_vm_stack_suspend_trace_zones(stack); |
| IREE_TRACE_FIBER_LEAVE(); |
| } |
| |
| #endif // IREE_TRACING_FEATURE_INSTRUMENTATION |
| |
| //===----------------------------------------------------------------------===// |
| // Synchronous invocation |
| //===----------------------------------------------------------------------===// |
| |
| IREE_API_EXPORT iree_status_t iree_vm_invoke( |
| iree_vm_context_t* context, iree_vm_function_t function, |
| iree_vm_invocation_flags_t flags, const iree_vm_invocation_policy_t* policy, |
| const iree_vm_list_t* inputs, iree_vm_list_t* outputs, |
| iree_allocator_t host_allocator) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Bound the synchronous invocation to the timeout specified by the user |
| // regardless of what the target of the invocation wants when it waits. |
| // TODO(benvanik): add a timeout arg to iree_vm_invoke. |
| // For now we only use the timeouts specified on the wait operations. |
| iree_timeout_t timeout = iree_infinite_timeout(); |
| iree_time_t deadline_ns = iree_timeout_as_deadline_ns(timeout); |
| |
| // Allocate an invocation ID for tracing. |
| iree_vm_invocation_id_t invocation_id = |
| iree_any_bit_set(flags, IREE_VM_INVOCATION_FLAG_TRACE_INLINE) |
| ? 0 |
| : iree_vm_invoke_allocate_id(context, &function); |
| (void)invocation_id; // unused when tracing is disabled |
| |
| // Begin a zone outside the fiber to represent one tick of the loop. |
| IREE_TRACE_ZONE_BEGIN_NAMED(zi, "iree_vm_invoke_tick"); |
| // Enter the fiber to start attributing zones to the context. |
| IREE_TRACE(iree_vm_invoke_fiber_enter(invocation_id)); |
| |
| // Perform the initial invocation step, which if synchronous may fully |
| // complete the invocation before returning. If it yields we'll need to resume |
| // it, possibly after taking care of pending waits. |
| iree_vm_invoke_state_t state = {0}; |
| iree_status_t status = iree_vm_begin_invoke(&state, context, function, flags, |
| policy, inputs, host_allocator); |
| while (iree_status_is_deferred(status)) { |
| // Grab the wait frame from the stack holding the wait parameters. |
| // This is optional: if an invocation yields for cooperative scheduling |
| // purposes there will not be a wait frame on the stack and we'll just |
| // resume it below. |
| iree_vm_stack_frame_t* current_frame = |
| iree_vm_stack_current_frame(state.stack); |
| if (IREE_UNLIKELY(!current_frame)) { |
| // Unbalanced stack. |
| status = iree_make_status(IREE_STATUS_FAILED_PRECONDITION, |
| "unbalanced stack after yield"); |
| break; // bail and don't attempt a resume |
| } else if (current_frame->type == IREE_VM_STACK_FRAME_WAIT) { |
| // Perform the wait operation synchronously. |
| // We do this outside of the fiber to match accounting with async |
| // executors. |
| IREE_TRACE(iree_vm_invoke_fiber_leave(invocation_id, state.stack)); |
| IREE_TRACE_ZONE_END(zi); |
| |
| iree_vm_wait_frame_t* wait_frame = |
| (iree_vm_wait_frame_t*)iree_vm_stack_frame_storage(current_frame); |
| status = iree_vm_wait_invoke(&state, wait_frame, deadline_ns); |
| |
| // Restore tick zone and re-enter the fiber for the resume. |
| IREE_TRACE_ZONE_BEGIN_NAMED(zi_next, "iree_vm_invoke_tick"); |
| zi = zi_next; |
| IREE_TRACE(iree_vm_invoke_fiber_reenter(invocation_id, state.stack)); |
| if (!iree_status_is_ok(status)) break; |
| } |
| |
| // Resume the invocation after its wait completes (if it wasn't just a |
| // simple yield for cooperation). This may yield again and require another |
| // tick or complete with OK (or an error). |
| status = iree_vm_resume_invoke(&state); |
| } |
| |
| // If the invoke process itself was successful we can end the invocation |
| // cleanly and get the invocation status as returned by the target function. |
| iree_status_t invoke_status = iree_ok_status(); |
| if (iree_status_is_ok(status)) { |
| status = iree_vm_end_invoke(&state, outputs, &invoke_status); |
| } |
| |
| // Otherwise if we failed to invoke we need to tear down the state to release |
| // all resources retained by the stack. |
| if (!iree_status_is_ok(status)) { |
| // Cleanup the invocation state if the end wasn't able to. |
| // This may leave the context in an unexpected state but the caller is |
| // expected to tear down everything if this happens. |
| iree_vm_abort_invoke(&state); |
| } |
| |
| // Leave the fiber context now that execution has completed. |
| IREE_TRACE(iree_vm_invoke_fiber_leave(invocation_id, state.stack)); |
| IREE_TRACE_ZONE_END(zi); |
| |
| // If we succeeded at invoking the status will be OK and the invoke_status |
| // will hold the status returned by the invokee. If we failed at invoking |
| // the invoke_status won't be set. |
| IREE_ASSERT(iree_status_is_ok(status) || |
| (!iree_status_is_ok(status) && iree_status_is_ok(invoke_status))); |
| status = !iree_status_is_ok(invoke_status) ? invoke_status : status; |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Asynchronous invocation |
| //===----------------------------------------------------------------------===// |
| |
| // Argument storage larger than this will require a heap allocation. |
| #define IREE_VM_STACK_MAX_ARGUMENT_ALLOCA_SIZE (iree_host_size_t)(16 * 1024) |
| |
| // WARNING: this function cannot have any trace markers that span the begin |
| // call; the begin may yield with zones still open. |
| IREE_API_EXPORT iree_status_t iree_vm_begin_invoke( |
| iree_vm_invoke_state_t* state, iree_vm_context_t* context, |
| iree_vm_function_t function, iree_vm_invocation_flags_t flags, |
| const iree_vm_invocation_policy_t* policy, const iree_vm_list_t* inputs, |
| iree_allocator_t host_allocator) { |
| IREE_ASSERT_ARGUMENT(context); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Force tracing if specified on the context. |
| if (iree_vm_context_flags(context) & IREE_VM_CONTEXT_FLAG_TRACE_EXECUTION) { |
| flags |= IREE_VM_INVOCATION_FLAG_TRACE_EXECUTION; |
| } |
| |
| // Grab function metadata used for marshaling inputs/outputs. |
| iree_vm_function_signature_t signature = |
| iree_vm_function_signature(&function); |
| iree_string_view_t cconv_arguments = iree_string_view_empty(); |
| iree_string_view_t cconv_results = iree_string_view_empty(); |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_vm_function_call_get_cconv_fragments( |
| &signature, &cconv_arguments, &cconv_results)); |
| |
| // Allocate argument storage on the native stack. It only needs to survive the |
| // begin call as it's consumed by the invokee. |
| iree_byte_span_t arguments = iree_make_byte_span(NULL, 0); |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, |
| iree_vm_function_call_compute_cconv_fragment_size( |
| cconv_arguments, /*segment_size_list=*/NULL, &arguments.data_length)); |
| const bool arguments_on_heap = |
| arguments.data_length > IREE_VM_STACK_MAX_ARGUMENT_ALLOCA_SIZE; |
| if (!arguments_on_heap) { |
| // Arguments fit on the native stack without too much worry about |
| // overflowing. This is the fast path (effectively just an $sp bump). |
| arguments.data = iree_alloca(arguments.data_length); |
| } else { |
| // Couldn't inline, do a heap allocation that we'll keep until this function |
| // returns. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_allocator_malloc(host_allocator, arguments.data_length, |
| (void**)&arguments.data)); |
| } |
| memset(arguments.data, 0, arguments.data_length); |
| |
| // Allocate the result storage that will be populated by the invokee. This |
| // must survive until the end() so we slice it off the bottom of the stack |
| // storage. This reduces the overall available stack space but not by much, |
| // and if the stack needs to dynamically grow the inlined storage will still |
| // be available. |
| iree_byte_span_t results = iree_make_byte_span(NULL, 0); |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_vm_function_call_compute_cconv_fragment_size( |
| cconv_results, /*segment_size_list=*/NULL, &results.data_length)); |
| iree_host_size_t reserved_storage_size = 0; |
| if (results.data_length <= sizeof(state->stack_storage) / 4) { |
| // Results fit in the inlined storage and we can avoid a heap allocation. |
| // If we exceed the maximum we'll heap allocate below inside the stack. |
| results.data = state->stack_storage; |
| reserved_storage_size = |
| iree_host_align(results.data_length, iree_max_align_t); |
| } else { |
| // Couldn't inline, do a heap allocation we'll have to hang on to and |
| // clean up when the invocation state is released. |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_allocator_malloc(host_allocator, results.data_length, |
| (void**)&results.data)); |
| } |
| memset(results.data, 0, results.data_length); |
| |
| // Marshal the input arguments into the VM ABI and preallocate the result |
| // buffer. If marshaling fails we need to cleanup the arguments. |
| // NOTE: today we don't support variadic arguments through this interface. |
| iree_status_t status = |
| iree_vm_invoke_marshal_inputs(cconv_arguments, inputs, arguments); |
| if (!iree_status_is_ok(status)) { |
| iree_vm_invoke_release_argument_storage(cconv_arguments, arguments, |
| arguments_on_heap, host_allocator); |
| iree_vm_invoke_release_result_storage(cconv_results, results, |
| state->stack_storage, host_allocator); |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| // Initialize the stack with the inline storage. |
| // We (probably) sliced off the head of the storage above to use for results |
| // and perform an offset here to account for that. |
| iree_vm_stack_t* stack = NULL; |
| status = iree_vm_stack_initialize( |
| iree_make_byte_span(state->stack_storage + reserved_storage_size, |
| sizeof(state->stack_storage) - reserved_storage_size), |
| flags, iree_vm_context_state_resolver(context), host_allocator, &stack); |
| if (!iree_status_is_ok(status)) { |
| iree_vm_invoke_release_argument_storage(cconv_arguments, arguments, |
| arguments_on_heap, host_allocator); |
| iree_vm_invoke_release_result_storage(cconv_results, results, |
| state->stack_storage, host_allocator); |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| // NOTE: at this point the stack must be properly deinitialized if we bail. |
| |
| // Initialize state now that we are confident we're returning OK. |
| // If we return a failure the user won't know they have to end() and clean |
| // these up. |
| state->context = context; |
| state->cconv_results = cconv_results; |
| state->results = results; |
| iree_vm_context_retain(context); |
| state->stack = stack; |
| |
| // NOTE: we must end the zone here as the begin_call will return with |
| // unbalanced zones if we yield. |
| IREE_TRACE_ZONE_END(z0); |
| |
| // Execute the target function until the first yield point is reached or it |
| // completes. A result of OK indicates successful completion while DEFERRED |
| // indicates that the invocation needs to be resumed/waited again. |
| iree_vm_function_call_t call = { |
| .function = function, |
| .arguments = arguments, |
| .results = results, |
| }; |
| state->status = |
| function.module->begin_call(function.module->self, stack, call); |
| |
| // Arguments should no longer be required - they were either consumed by the |
| // begin_call or need to be cleaned up before we return. |
| iree_vm_invoke_release_argument_storage(cconv_arguments, call.arguments, |
| arguments_on_heap, host_allocator); |
| |
| // The call may have yielded, either for cooperative scheduling purposes or |
| // for a wait operation (in which case the top of the stack will have a wait |
| // frame). |
| if (iree_status_is_deferred(state->status)) { |
| return iree_status_from_code(IREE_STATUS_DEFERRED); |
| } |
| |
| // NOTE: the begin-invoke was ok, but the operation itself may have failed. |
| return iree_ok_status(); |
| } |
| |
| // WARNING: this function cannot have any trace markers that span the resume |
| // call; the resume may yield with zones still open. |
| IREE_API_EXPORT iree_status_t |
| iree_vm_resume_invoke(iree_vm_invoke_state_t* state) { |
| IREE_ASSERT_ARGUMENT(state); |
| |
| // In a stackless world resuming may pop a stack frame that needs to be |
| // executed inline. We run here until either all stack frames have been popped |
| // (indicating the invocation has completed) or we yield/error and want to |
| // return to the scheduler. |
| do { |
| if (iree_status_is_deferred(state->status)) { |
| // Wait required; top of the stack should be a wait frame. |
| IREE_ASSERT_EQ(iree_vm_stack_current_frame(state->stack)->type, |
| IREE_VM_STACK_FRAME_WAIT); |
| return iree_status_from_code(IREE_STATUS_DEFERRED); |
| } else if (!iree_status_is_ok(state->status)) { |
| // Invocation previously failed so return immediately. The user should |
| // then call end() to get the result. By returning OK here we are telling |
| // the user the resume operation succeeded. |
| return iree_ok_status(); |
| } |
| |
| // Get the top execution frame of the stack where we will resume execution. |
| iree_vm_stack_frame_t* resume_frame = iree_vm_stack_top(state->stack); |
| if (IREE_UNLIKELY(!resume_frame)) { |
| return iree_make_status(IREE_STATUS_FAILED_PRECONDITION, |
| "resume called with no parent frame"); |
| } |
| |
| // Call into the VM to resume the function. It may complete (returning OK), |
| // defer to be waited/resumed later, or fail. |
| iree_vm_function_t resume_function = resume_frame->function; |
| state->status = resume_function.module->resume_call( |
| resume_function.module->self, state->stack, state->results); |
| |
| // If the call yielded then return that so the user knows to resume again. |
| if (iree_status_is_deferred(state->status)) { |
| return iree_status_from_code(IREE_STATUS_DEFERRED); |
| } |
| |
| // Stack resume: if the resume succeeded but the stack is not empty it means |
| // we've got to resume the parent frame. When we do a full yield up to the |
| // scheduler and then resume we're calling into the VM stack top from the |
| // host stack bottom - to have the same behavior as a normal stack pop we've |
| // got to continue running. To keep the trace cleaner and reduce overhead we |
| // jump back up and pop the next frame, which also helps us avoid |
| // introducing latency between pops where otherwise there should be none. |
| } while (iree_status_is_ok(state->status) && |
| iree_vm_stack_current_frame(state->stack) != NULL); |
| |
| // We're indicating the resume operation was successful, not the result of the |
| // VM call; the user will call end() to get that. |
| return iree_ok_status(); |
| } |
| |
| IREE_API_EXPORT iree_status_t |
| iree_vm_wait_invoke(iree_vm_invoke_state_t* state, |
| iree_vm_wait_frame_t* wait_frame, iree_time_t deadline_ns) { |
| IREE_ASSERT_ARGUMENT(state); |
| if (IREE_UNLIKELY(!iree_status_is_deferred(state->status))) { |
| // Can only wait if the invocation is actually waiting. |
| // We could make this OK and act as a no-op but it can be useful for |
| // ensuring scheduler implementations don't do extraneous work. |
| return iree_make_status( |
| IREE_STATUS_FAILED_PRECONDITION, |
| "wait-invoke attempted on a non-waiting invocation"); |
| } |
| |
| // Combine the wait-invoke deadline with the one specified by the wait |
| // operation itself. This allows schedulers to timeslice waits without |
| // worrying whether user programs request to wait forever. |
| iree_time_t min_deadline_ns = iree_min(deadline_ns, wait_frame->deadline_ns); |
| |
| // Perform the wait operation, blocking the calling thread until it completes, |
| // fails, or hits the min_deadline_ns. |
| if (wait_frame->wait_type == IREE_VM_WAIT_UNTIL) { |
| wait_frame->wait_status = iree_wait_until(min_deadline_ns) |
| ? iree_ok_status() |
| : iree_status_from_code(IREE_STATUS_ABORTED); |
| } else if (wait_frame->count == 1) { |
| wait_frame->wait_status = iree_wait_source_wait_one( |
| wait_frame->wait_sources[0], iree_make_deadline(min_deadline_ns)); |
| } else { |
| // TODO(benvanik): multi-wait when running synchronously. This is already |
| // supported by iree_loop_inline_t and maybe we can just reuse that. These |
| // are not currently emitted by the compiler. |
| return iree_make_status( |
| IREE_STATUS_UNIMPLEMENTED, |
| "multi-wait in synchronous invocations not yet implemented"); |
| } |
| |
| // Reset status to OK - the next resume will pick back up in the waiter. |
| iree_status_free(state->status); |
| state->status = iree_ok_status(); |
| |
| // OK here indicates we performed the wait and not the result of the wait. |
| return iree_ok_status(); |
| } |
| |
| IREE_API_EXPORT iree_status_t iree_vm_end_invoke(iree_vm_invoke_state_t* state, |
| iree_vm_list_t* outputs, |
| iree_status_t* out_status) { |
| IREE_ASSERT_ARGUMENT(state); |
| IREE_ASSERT_ARGUMENT(out_status); |
| *out_status = iree_ok_status(); |
| |
| // Suspend stack frame tracing zones; if the invocation failed the failing |
| // frames will still be on the stack and thus also still have their trace |
| // zones active. If we were to begin a zone here and then deinit the stack |
| // we'd end up with unbalanced zones. |
| iree_vm_stack_suspend_trace_zones(state->stack); |
| |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Grab operation status. If this is not OK it's because the operation failed |
| // or the user is calling this with a wait frame on the stack. |
| iree_status_t invoke_status = state->status; |
| if (IREE_UNLIKELY(iree_status_is_deferred(invoke_status))) { |
| IREE_TRACE_ZONE_END(z0); |
| return iree_make_status(IREE_STATUS_FAILED_PRECONDITION, |
| "end-invoke attempted on a waiting invocation"); |
| } else if (IREE_UNLIKELY(!iree_status_is_ok(invoke_status))) { |
| // Annotate failures with the stack trace (if compiled in). |
| invoke_status = IREE_VM_STACK_ANNOTATE_BACKTRACE_IF_ENABLED(state->stack, |
| invoke_status); |
| } |
| |
| // If the operation succeeded marshal the outputs from the stack buffers into |
| // the user-provided storage. The outputs list will retain all results. |
| if (iree_status_is_ok(invoke_status)) { |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( |
| z0, iree_vm_invoke_marshal_outputs(state->cconv_results, state->results, |
| outputs)); |
| } |
| |
| // Cleanup the invocation resources. |
| *out_status = invoke_status; // takes ownership |
| state->status = iree_ok_status(); |
| iree_vm_abort_invoke(state); |
| |
| IREE_TRACE_ZONE_END(z0); |
| return iree_ok_status(); |
| } |
| |
| IREE_API_EXPORT void iree_vm_abort_invoke(iree_vm_invoke_state_t* state) { |
| // We expect that the caller has already suspended the stack tracing zones, |
| // but in failure cases we can end up here and want to ensure that things are |
| // cleaned up. If we were to begin a zone now the stack deinit would lead to |
| // unbalanced zones. |
| if (state->stack) iree_vm_stack_suspend_trace_zones(state->stack); |
| |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| iree_allocator_t host_allocator = state->stack |
| ? iree_vm_stack_allocator(state->stack) |
| : iree_allocator_null(); |
| |
| if (state->stack) { |
| iree_vm_stack_deinitialize(state->stack); |
| state->stack = NULL; |
| } |
| |
| if (!iree_byte_span_is_empty(state->results)) { |
| iree_vm_invoke_release_result_storage(state->cconv_results, state->results, |
| state->stack_storage, host_allocator); |
| state->results = iree_byte_span_empty(); |
| } |
| |
| if (state->context) { |
| iree_vm_context_release(state->context); |
| state->context = NULL; |
| } |
| |
| iree_status_free(state->status); |
| state->status = iree_status_from_code(IREE_STATUS_INTERNAL); |
| |
| IREE_TRACE_ZONE_END(z0); |
| } |
| |
| //===----------------------------------------------------------------------===// |
| // Loop-based asynchronous invocation |
| //===----------------------------------------------------------------------===// |
| |
| static iree_status_t iree_vm_async_begin_invoke(void* user_data, |
| iree_loop_t loop, |
| iree_status_t loop_status); |
| static iree_status_t iree_vm_async_resume_invoke(void* user_data, |
| iree_loop_t loop, |
| iree_status_t loop_status); |
| static iree_status_t iree_vm_async_tick_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop); |
| static iree_status_t iree_vm_async_end_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop); |
| static iree_status_t iree_vm_async_complete_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop, |
| iree_status_t status); |
| |
| IREE_API_EXPORT iree_status_t iree_vm_async_invoke( |
| iree_loop_t loop, iree_vm_async_invoke_state_t* state, |
| iree_vm_context_t* context, iree_vm_function_t function, |
| iree_vm_invocation_flags_t flags, const iree_vm_invocation_policy_t* policy, |
| iree_vm_list_t* inputs, iree_vm_list_t* outputs, |
| iree_allocator_t host_allocator, |
| iree_vm_async_invoke_callback_fn_t callback, void* user_data) { |
| IREE_ASSERT_ARGUMENT(state); |
| IREE_ASSERT_ARGUMENT(context); |
| IREE_TRACE_ZONE_BEGIN(z0); |
| |
| // Initialize to the pre-begin state. |
| state->begin_params.context = context; |
| iree_vm_context_retain(context); |
| state->begin_params.function = function; |
| state->begin_params.flags = flags; |
| state->begin_params.policy = policy; |
| state->begin_params.inputs = inputs; |
| iree_vm_list_retain(inputs); |
| state->deadline_ns = IREE_TIME_INFINITE_FUTURE; |
| state->host_allocator = host_allocator; |
| state->outputs = outputs; |
| iree_vm_list_retain(outputs); |
| state->callback = callback; |
| state->user_data = user_data; |
| |
| // Launch the invocation; if this fails we'll need to cleanup the state we've |
| // already initialized. |
| // NOTE: based on the loop type THIS MAY COMPLETE THE INVOCATION IMMEDIATELY. |
| iree_status_t status = iree_loop_call(loop, IREE_LOOP_PRIORITY_DEFAULT, |
| iree_vm_async_begin_invoke, state); |
| if (!iree_status_is_ok(status)) { |
| iree_vm_list_release(state->outputs); |
| iree_vm_list_release(state->begin_params.inputs); |
| iree_vm_context_release(state->begin_params.context); |
| } |
| |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| // Begins the invocation from the first loop callback. |
| // The begin_params on the state will have everything we need to initialize the |
| // call but since we alias with the base invocation state we must be sure to |
| // copy out the args first. |
| // |
| // Note that |status| may indicate a failure already, such as if the loop |
| // aborted. In that case we need to clean up the state before issuing the user |
| // callback so they can do the same. |
| static iree_status_t iree_vm_async_begin_invoke(void* user_data, |
| iree_loop_t loop, |
| iree_status_t loop_status) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| iree_vm_async_invoke_state_t* state = |
| (iree_vm_async_invoke_state_t*)user_data; |
| |
| // Check to see if the loop has failed before we even begin. |
| if (IREE_UNLIKELY(!iree_status_is_ok(loop_status))) { |
| // We release our retained resources because we don't guarantee they live to |
| // the callback. This allows callbacks to reuse memory. |
| iree_vm_list_release(state->outputs); |
| iree_vm_list_release(state->begin_params.inputs); |
| iree_vm_context_release(state->begin_params.context); |
| |
| // Issue user callback notifying them of the failure and pass along the loop |
| // status; this is likely something like IREE_STATUS_ABORTED. |
| iree_status_t callback_status = |
| state->callback(state->user_data, loop, loop_status, NULL); |
| IREE_TRACE_ZONE_END(z0); |
| return callback_status; |
| } |
| |
| // Pull fields locally so that we can reuse the aliased storage. |
| // Note that we have ownership of all these and must release them if we fail |
| // to begin the invocation. |
| iree_vm_context_t* context = state->begin_params.context; |
| iree_vm_function_t function = state->begin_params.function; |
| iree_vm_invocation_flags_t flags = state->begin_params.flags; |
| const iree_vm_invocation_policy_t* policy = state->begin_params.policy; |
| iree_vm_list_t* inputs = state->begin_params.inputs; |
| |
| // Allocate an invocation ID for tracing. |
| IREE_TRACE({ |
| state->invocation_id = |
| iree_any_bit_set(flags, IREE_VM_INVOCATION_FLAG_TRACE_INLINE) |
| ? 0 |
| : iree_vm_invoke_allocate_id(context, &function); |
| }); |
| |
| // Try to begin the invocation. This may fail if the parameters are invalid. |
| // It may also complete inline if the entire invocation can be handled without |
| // blocking (in which case begin_status is OK). |
| IREE_TRACE(iree_vm_invoke_fiber_enter(state->invocation_id)); |
| iree_status_t status = |
| iree_vm_begin_invoke(&state->base, context, function, flags, policy, |
| inputs, state->host_allocator); |
| if (iree_status_is_ok(status) || iree_status_is_deferred(status)) { |
| // Ownership transferred. |
| iree_vm_list_release(inputs); |
| inputs = NULL; |
| iree_vm_context_release(context); |
| context = NULL; |
| } |
| if (iree_status_is_deferred(status)) { |
| IREE_TRACE({ |
| iree_vm_invoke_fiber_leave(state->invocation_id, state->base.stack); |
| }); |
| // Deferred until a wait completes or the next tick. |
| status = iree_vm_async_tick_invoke(state, loop); |
| } else if (iree_status_is_ok(status)) { |
| // Completed synchronously. This is the happy path and lets us complete the |
| // entire invocation in a single loop operation. |
| status = iree_vm_async_end_invoke(state, loop); |
| } else { |
| IREE_TRACE(iree_vm_invoke_fiber_leave(state->invocation_id, NULL)); |
| // Failed to begin the invocation; release resources and call back. |
| // We know the state wasn't fully initialized and don't need to clean it up. |
| iree_vm_list_release(state->outputs); |
| iree_vm_list_release(inputs); |
| iree_vm_context_release(context); |
| status = state->callback(state->user_data, loop, status, NULL); |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| // If we began but failed to tick/end we need to propagate that to the user |
| // and clean up our state. |
| if (!iree_status_is_ok(status)) { |
| status = iree_vm_async_complete_invoke(state, loop, status); |
| } |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static iree_status_t iree_vm_async_resume_invoke(void* user_data, |
| iree_loop_t loop, |
| iree_status_t loop_status) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| iree_vm_async_invoke_state_t* state = |
| (iree_vm_async_invoke_state_t*)user_data; |
| |
| // Resume the invocation and execute the next step. |
| IREE_TRACE({ |
| iree_vm_invoke_fiber_reenter(state->invocation_id, state->base.stack); |
| }); |
| iree_status_t status = iree_vm_resume_invoke(&state->base); |
| if (iree_status_is_deferred(status)) { |
| IREE_TRACE({ |
| iree_vm_invoke_fiber_leave(state->invocation_id, state->base.stack); |
| }); |
| // Deferred on a wait or yield. Enqueue waits/a resume. |
| status = iree_vm_async_tick_invoke(state, loop); |
| } else if (iree_status_is_ok(status)) { |
| // Completed synchronously. |
| status = iree_vm_async_end_invoke(state, loop); |
| } else { |
| IREE_TRACE({ |
| iree_vm_invoke_fiber_leave(state->invocation_id, state->base.stack); |
| }); |
| } |
| |
| // If we failed to tick/end we need to propagate that to the user and clean up |
| // our state. |
| if (!iree_status_is_ok(status)) { |
| status = iree_vm_async_complete_invoke(state, loop, status); |
| } |
| IREE_TRACE_ZONE_END(z0); |
| return status; |
| } |
| |
| static iree_status_t iree_vm_async_wake_invoke(void* user_data, |
| iree_loop_t loop, |
| iree_status_t loop_status) { |
| IREE_TRACE_ZONE_BEGIN(z0); |
| iree_vm_async_invoke_state_t* state = |
| (iree_vm_async_invoke_state_t*)user_data; |
| |
| // If we were aborted then we need to tear everything down. |
| // TODO(benvanik): maybe allow the failures through to the target? It'd be |
| // impossible to tell when the loop was in an invalid state if we did. May |
| // need to rework the loop callback on waits so that we can differentiate. |
| if (iree_status_is_aborted(loop_status)) { |
| IREE_TRACE_ZONE_END(z0); |
| return iree_vm_async_complete_invoke(state, loop, loop_status); |
| } |
| |
| // The loop_status we receive here is the result of the wait operation and |
| // something we need to propagate to the waiter. |
| iree_vm_stack_frame_t* current_frame = |
| iree_vm_stack_current_frame(state->base.stack); |
| iree_vm_wait_frame_t* wait_frame = |
| (iree_vm_wait_frame_t*)iree_vm_stack_frame_storage(current_frame); |
| wait_frame->wait_status = loop_status; |
| |
| IREE_ASSERT(iree_status_is_deferred(state->base.status)); |
| iree_status_free(state->base.status); |
| state->base.status = iree_ok_status(); |
| |
| IREE_TRACE_ZONE_END(z0); |
| |
| // Resume the invocation and execute the next step. |
| // We do this inline instead of enqueuing a resume so that we avoid a needless |
| // operation in the loop. The invocation may immediately wait again and we |
| // want to keep the total wait-to-wait latency low. |
| return iree_vm_async_resume_invoke(user_data, loop, iree_ok_status()); |
| } |
| |
| static iree_status_t iree_vm_async_tick_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop) { |
| // Grab the wait frame from the stack holding the wait parameters. |
| // This is optional: if an invocation yields for cooperative scheduling |
| // purposes there will not be a wait frame on the stack and we'll just |
| // resume it below. |
| iree_vm_stack_frame_t* current_frame = |
| iree_vm_stack_current_frame(state->base.stack); |
| if (IREE_UNLIKELY(!current_frame)) { |
| // Unbalanced stack. |
| return iree_make_status(IREE_STATUS_FAILED_PRECONDITION, |
| "unbalanced stack after yield"); |
| } else if (current_frame->type == IREE_VM_STACK_FRAME_WAIT) { |
| // Wait on a wait source. |
| iree_vm_wait_frame_t* wait_frame = |
| (iree_vm_wait_frame_t*)iree_vm_stack_frame_storage(current_frame); |
| |
| // Combine the wait-invoke deadline with the one specified by the wait |
| // operation itself. This allows schedulers to timeslice waits without |
| // worrying whether user programs request to wait forever. |
| iree_timeout_t timeout = iree_make_deadline( |
| iree_min(state->deadline_ns, wait_frame->deadline_ns)); |
| switch (wait_frame->wait_type) { |
| default: |
| case IREE_VM_WAIT_UNTIL: |
| return iree_loop_wait_until(loop, timeout, iree_vm_async_wake_invoke, |
| state); |
| case IREE_VM_WAIT_ANY: |
| return iree_loop_wait_any(loop, wait_frame->count, |
| wait_frame->wait_sources, timeout, |
| iree_vm_async_wake_invoke, state); |
| case IREE_VM_WAIT_ALL: |
| return iree_loop_wait_all(loop, wait_frame->count, |
| wait_frame->wait_sources, timeout, |
| iree_vm_async_wake_invoke, state); |
| } |
| } else { |
| // Resume from a yield point (cooperative scheduling). |
| return iree_loop_call(loop, IREE_LOOP_PRIORITY_DEFAULT, |
| iree_vm_async_resume_invoke, state); |
| } |
| } |
| |
| static iree_status_t iree_vm_async_end_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop) { |
| // End the invocation and retrieve the results. |
| iree_status_t invoke_status = iree_ok_status(); |
| IREE_RETURN_IF_ERROR( |
| iree_vm_end_invoke(&state->base, state->outputs, &invoke_status)); |
| IREE_TRACE({ |
| // We leave the fiber before completing so that the callback is attributed |
| // to the thread running it instead. |
| iree_vm_invoke_fiber_leave(state->invocation_id, state->base.stack); |
| }); |
| return iree_vm_async_complete_invoke(state, loop, invoke_status); |
| } |
| |
| static iree_status_t iree_vm_async_complete_invoke( |
| iree_vm_async_invoke_state_t* state, iree_loop_t loop, |
| iree_status_t status) { |
| // Release all resources if we didn't already clean them up. |
| if (!iree_status_is_ok(status)) { |
| iree_vm_abort_invoke(&state->base); |
| iree_vm_list_release(state->outputs); |
| state->outputs = NULL; |
| } |
| |
| // Issue callback. |
| iree_vm_list_t* outputs = state->outputs; |
| return state->callback(state->user_data, loop, status, outputs); |
| } |
