diff --git a/libshortfin/bindings/python/CMakeLists.txt b/libshortfin/bindings/python/CMakeLists.txt index 653648961..66e6476a6 100644 --- a/libshortfin/bindings/python/CMakeLists.txt +++ b/libshortfin/bindings/python/CMakeLists.txt @@ -28,7 +28,7 @@ set_target_properties(shortfin_python_extension target_link_libraries(shortfin_python_extension # TODO: This should be configurable as to whether we link to the static # or dynamic version. - PRIVATE shortfin-static + PRIVATE shortfin ) nanobind_add_stub( diff --git a/libshortfin/src/shortfin/local/program.cc b/libshortfin/src/shortfin/local/program.cc index cb278ac54..c7a045eeb 100644 --- a/libshortfin/src/shortfin/local/program.cc +++ b/libshortfin/src/shortfin/local/program.cc @@ -241,6 +241,13 @@ std::vector Program::exports() const { // ProgramInvocation // -------------------------------------------------------------------------- // +iree_vm_list_t *ProgramInvocation::arg_list() { + // The arg list is located immediately after this, allocated as a trailing + // data structure. + return reinterpret_cast(reinterpret_cast(this) + + sizeof(*this)); +} + void ProgramInvocation::Deleter::operator()(ProgramInvocation *inst) { inst->~ProgramInvocation(); uint8_t *memory = static_cast(static_cast(inst)); @@ -311,7 +318,6 @@ ProgramInvocation::Ptr ProgramInvocation::New( vm_context.release(); // Ref transfer to ProgramInvocation. inst->state.params.function = vm_function; inst->state.params.invocation_model = invocation_model; - inst->state.params.arg_list = arg_list; inst->result_list_ = result_list; return inst; } @@ -324,14 +330,12 @@ void ProgramInvocation::CheckNotScheduled() { void ProgramInvocation::AddArg(iree::vm_opaque_ref ref) { CheckNotScheduled(); - SHORTFIN_THROW_IF_ERROR( - iree_vm_list_push_ref_move(state.params.arg_list, &ref)); + SHORTFIN_THROW_IF_ERROR(iree_vm_list_push_ref_move(arg_list(), &ref)); } void ProgramInvocation::AddArg(iree_vm_ref_t *ref) { CheckNotScheduled(); - SHORTFIN_THROW_IF_ERROR( - iree_vm_list_push_ref_retain(state.params.arg_list, ref)); + SHORTFIN_THROW_IF_ERROR(iree_vm_list_push_ref_retain(arg_list(), ref)); } iree_status_t ProgramInvocation::FinalizeCallingConvention( @@ -395,7 +399,7 @@ ProgramInvocation::Future ProgramInvocation::Invoke( auto schedule = [](ProgramInvocation *raw_invocation, Worker *worker, iree_vm_context_t *owned_context, - iree_vm_function_t function, iree_vm_list_t *arg_list, + iree_vm_function_t function, ProgramInvocationModel invocation_model, std::optional failure_future) { auto complete_callback = @@ -436,8 +440,8 @@ ProgramInvocation::Future ProgramInvocation::Invoke( status = invocation->scope()->scheduler().FlushWithStatus(); } if (iree_status_is_ok(status)) { - status = invocation->FinalizeCallingConvention(arg_list, function, - invocation_model); + status = invocation->FinalizeCallingConvention( + invocation->arg_list(), function, invocation_model); } if (iree_status_is_ok(status)) { status = iree_vm_async_invoke(worker->loop(), @@ -445,7 +449,7 @@ ProgramInvocation::Future ProgramInvocation::Invoke( owned_context, function, /*flags=*/IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/nullptr, - /*inputs=*/arg_list, + /*inputs=*/invocation->arg_list(), /*outputs=*/invocation->result_list_, iree_allocator_system(), +complete_callback, /*user_data=*/invocation.get()); @@ -478,12 +482,12 @@ ProgramInvocation::Future ProgramInvocation::Invoke( if (&worker == Worker::GetCurrent()) { // On the same worker: fast-path directly to the loop. schedule(invocation.release(), &worker, params.context, params.function, - params.arg_list, params.invocation_model, /*failure_future=*/{}); + params.invocation_model, /*failure_future=*/{}); } else { // Cross worker coordination: submit an external task to bootstrap. auto bound_schedule = std::bind(schedule, invocation.release(), &worker, params.context, - params.function, params.arg_list, params.invocation_model, + params.function, params.invocation_model, /*failure_future=*/fork_future); worker.CallThreadsafe(bound_schedule); } diff --git a/libshortfin/src/shortfin/local/program.h b/libshortfin/src/shortfin/local/program.h index f0c6b5090..e701b0489 100644 --- a/libshortfin/src/shortfin/local/program.h +++ b/libshortfin/src/shortfin/local/program.h @@ -131,6 +131,9 @@ class SHORTFIN_API ProgramInvocation { ProgramInvocation(); void CheckNotScheduled(); + // Returns a pointer to the trailing arg list. + iree_vm_list_t *arg_list(); + // Accesses the invocation owned wait fence, creating it if needed. iree_hal_fence_t *wait_fence(); @@ -154,7 +157,6 @@ class SHORTFIN_API ProgramInvocation { iree_vm_context_t *context; iree_vm_function_t function; ProgramInvocationModel invocation_model; - iree_vm_list_t *arg_list = nullptr; }; union State { State() { new (¶ms) Params(); }