diff --git a/cpp/HybridTfliteModel.cpp b/cpp/HybridTfliteModel.cpp index 020b60fa..439de90e 100644 --- a/cpp/HybridTfliteModel.cpp +++ b/cpp/HybridTfliteModel.cpp @@ -28,11 +28,28 @@ HybridTfliteModel::HybridTfliteModel(std::shared_ptr interpre } } +void HybridTfliteModel::dispose() { + // The lock waits out an in-flight inference on another thread - freeing the + // interpreter under a running TfLiteInterpreterInvoke would be a native crash. + std::lock_guard lock(_lifecycleMutex); + if (_interpreter == nullptr) { + return; // already disposed + } + // We hold the only reference to the interpreter, so this runs its deleter + // now: TfLiteInterpreterDelete, then the delegates. The model bytes are + // released once nobody else (e.g. JS) references them either. + _interpreter.reset(); + _modelData.reset(); + _outputBuffers.clear(); +} + std::vector HybridTfliteModel::getDelegates() { return _delegates; } std::vector HybridTfliteModel::getInputs() { + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); int count = TfLiteInterpreterGetInputTensorCount(_interpreter.get()); std::vector tensors; tensors.reserve(count); @@ -55,6 +72,8 @@ std::vector HybridTfliteModel::getInputs() { } std::vector HybridTfliteModel::getOutputs() { + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); int count = TfLiteInterpreterGetOutputTensorCount(_interpreter.get()); std::vector tensors; tensors.reserve(count); @@ -148,6 +167,10 @@ void HybridTfliteModel::invoke() { std::vector> HybridTfliteModel::runSync(const std::vector>& input) { + // Held for the whole inference so dispose() can never free the interpreter + // mid-invoke. The disposed-throw is a catchable JS error on any runtime. + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); copyInputBuffers(input); invoke(); return copyOutputBuffers(); @@ -155,12 +178,20 @@ HybridTfliteModel::runSync(const std::vector>& inpu std::shared_ptr>>> HybridTfliteModel::run(const std::vector>& input) { - // Copy input buffers on caller (JS) thread first — input ArrayBuffers are - // non-owning JS buffers that may be GC'd if we access them async. - copyInputBuffers(input); + { + // Copy input buffers on caller (JS) thread first — input ArrayBuffers are + // non-owning JS buffers that may be GC'd if we access them async. + std::lock_guard lock(_lifecycleMutex); + throwIfDisposed(); + copyInputBuffers(input); + } std::shared_ptr sharedThis = shared_cast(); return Promise>>::async( [sharedThis]() -> std::vector> { + // Re-acquire on the async thread: dispose() may have landed between + // the input copy above and this lambda running. + std::lock_guard lock(sharedThis->_lifecycleMutex); + sharedThis->throwIfDisposed(); sharedThis->invoke(); return sharedThis->copyOutputBuffers(); }); diff --git a/cpp/HybridTfliteModel.hpp b/cpp/HybridTfliteModel.hpp index d0a137ca..b9b04969 100644 --- a/cpp/HybridTfliteModel.hpp +++ b/cpp/HybridTfliteModel.hpp @@ -2,6 +2,8 @@ #include "HybridTfliteModelSpec.hpp" #include +#include +#include #include #include @@ -22,6 +24,17 @@ class HybridTfliteModel : public HybridTfliteModelSpec { std::vector delegates); ~HybridTfliteModel() override = default; + /** + * Free the interpreter, its delegates and our reference to the model bytes + * NOW, without waiting for every runtime's GC to drop its reference (worklet + * runtimes may not GC for a long time, especially while backgrounded). + * Thread-safe: blocks until an in-flight inference on another thread has + * completed. Every later runSync/run/getInputs/getOutputs call throws a + * catchable JS error. Idempotent. Called by Nitro when JS invokes + * `model.dispose()`. + */ + void dispose() override; + // Properties (from HybridTfliteModelSpec) std::vector getDelegates() override; std::vector getInputs() override; @@ -39,11 +52,23 @@ class HybridTfliteModel : public HybridTfliteModelSpec { std::vector> copyOutputBuffers(); std::shared_ptr getOutputBufferForTensor(const TfLiteTensor* tensor); + // Caller must hold _lifecycleMutex. A disposed model has released its + // interpreter; std::runtime_error surfaces as a catchable JS error. + void throwIfDisposed() const { + if (_interpreter == nullptr) { + throw std::runtime_error("TFLite: Model was disposed!"); + } + } + private: std::shared_ptr _interpreter; std::vector _delegates; std::shared_ptr _modelData; std::unordered_map> _outputBuffers; + + // Serializes inference against dispose(). Uncontended in normal operation + // (one lock per inference, ~ns vs ~ms inference cost). + std::mutex _lifecycleMutex; }; } // namespace margelo::nitro::tflite diff --git a/cpp/HybridTfliteModule.cpp b/cpp/HybridTfliteModule.cpp index 91102012..631ee879 100644 --- a/cpp/HybridTfliteModule.cpp +++ b/cpp/HybridTfliteModule.cpp @@ -2,11 +2,18 @@ #include "TfliteHelpers.hpp" #include +#include +#include #if defined(ANDROID) #include +#include +#include #elif defined(__APPLE__) #include +#if FAST_TFLITE_ENABLE_CORE_ML +#include +#endif #else #error "Invalid Platform!" #endif @@ -32,6 +39,37 @@ TfLiteDelegate* getDelegate(TensorflowModelDelegate delegateType) { "\"!"); } +/** + * TFLite's C API does not transfer delegate ownership to the interpreter: the + * caller must keep a delegate alive for the interpreter's lifetime and free it + * afterwards with the delegate's own delete function. + */ +struct DelegateDeleter { + TensorflowModelDelegate delegateType; + + void operator()(TfLiteDelegate* delegate) const { + switch (delegateType) { +#if defined(__APPLE__) && FAST_TFLITE_ENABLE_CORE_ML + case TensorflowModelDelegate::CORE_ML: + TfLiteCoreMlDelegateDelete(delegate); + return; +#endif +#if defined(ANDROID) + case TensorflowModelDelegate::ANDROID_GPU: + TfLiteGpuDelegateV2Delete(delegate); + return; + case TensorflowModelDelegate::NNAPI: + TfLiteNnapiDelegateDelete(delegate); + return; +#endif + default: + // getDelegate() throws for every other type on this platform. + return; + } + } +}; +using OwnedDelegate = std::unique_ptr; + std::shared_ptr HybridTfliteModule::createModel(const std::shared_ptr& modelData, const std::vector& delegates) { @@ -50,20 +88,39 @@ HybridTfliteModule::createModel(const std::shared_ptr& modelData, // Add all hardware accelerated delegates (e.g. GPU, NPU, ...) // if any. The default CPU delegate will always be available. + std::vector effectiveDelegates; + std::vector ownedDelegates; + effectiveDelegates.reserve(delegates.size()); + ownedDelegates.reserve(delegates.size()); for (const TensorflowModelDelegate& delegateType : delegates) { - TfLiteDelegate* delegate = getDelegate(delegateType); - TfLiteInterpreterOptionsAddDelegate(options.get(), delegate); + OwnedDelegate delegate(getDelegate(delegateType), DelegateDeleter{delegateType}); + if (delegate == nullptr) { + // e.g. CoreML on devices without a Neural Engine — fall back to CPU + // instead of registering a null delegate with the interpreter. + continue; + } + TfLiteInterpreterOptionsAddDelegate(options.get(), delegate.get()); + effectiveDelegates.push_back(delegateType); + ownedDelegates.push_back(std::move(delegate)); } TfLiteInterpreter* rawInterpreter = TfLiteInterpreterCreate(model.get(), options.get()); if (rawInterpreter == nullptr) { + // `ownedDelegates` frees the delegates on unwind. throw std::runtime_error("Failed to create TFLite interpreter!"); } + // The delegates travel with the interpreter and are freed right after it, + // so they can never be deleted while the interpreter still uses them. const std::shared_ptr interpreter( - rawInterpreter, [modelData](TfLiteInterpreter* value) { TfLiteInterpreterDelete(value); }); + rawInterpreter, + [modelData, ownedDelegates = std::move(ownedDelegates)](TfLiteInterpreter* value) mutable { + TfLiteInterpreterDelete(value); + ownedDelegates.clear(); + }); - // Wrap in HybridTfliteModel — stores shared_ptr to keep model data bytes alive - return std::make_shared(interpreter, modelData, delegates); + // Wrap in HybridTfliteModel — stores shared_ptr to keep model data bytes alive. + // Only the delegates that were actually registered are reported via `getDelegates()`. + return std::make_shared(interpreter, modelData, effectiveDelegates); } } // namespace margelo::nitro::tflite