diff --git a/tensorflow/lite/interpreter.cc b/tensorflow/lite/interpreter.cc index cf8d56faf65..f97d6d805d0 100644 --- a/tensorflow/lite/interpreter.cc +++ b/tensorflow/lite/interpreter.cc @@ -180,17 +180,20 @@ TfLiteStatus Interpreter::AllocateTensors() { // Apply the default delegate that TFLite will enable at this point to allow // other user-level delegates to be applied first. if (!lazy_delegate_providers_.empty()) { + // We only apply lazy delegate providers once. + std::vector<TfLiteDelegatePtr> delegate_providers; + delegate_providers.swap(lazy_delegate_providers_); + TFLITE_LOG(TFLITE_LOG_INFO, "Applying %zu TensorFlow Lite delegate(s) lazily.", - lazy_delegate_providers_.size()); + delegate_providers.size()); // At the momement, XNNPACK delegate is the only one that might be applied // by default, in which case, the execution will fall back to default // implementation if the XNNPACK delegate fails to be applied. Therefore, we // ignore the return status here and let it fall through the rest of the // code. - for (size_t i = 0; i < lazy_delegate_providers_.size(); ++i) { - auto status = - ModifyGraphWithDelegate(std::move(lazy_delegate_providers_[i])); + for (size_t i = 0; i < delegate_providers.size(); ++i) { + auto status = ModifyGraphWithDelegate(std::move(delegate_providers[i])); switch (status) { case kTfLiteOk: TFLITE_LOG(TFLITE_LOG_INFO, @@ -225,7 +228,6 @@ TfLiteStatus Interpreter::AllocateTensors() { return kTfLiteError; } } - lazy_delegate_providers_.clear(); } return primary_subgraph().AllocateTensors(); diff --git a/tensorflow/lite/interpreter_test.cc b/tensorflow/lite/interpreter_test.cc index 6f5c6b3930d..dd26c298243 100644 --- a/tensorflow/lite/interpreter_test.cc +++ b/tensorflow/lite/interpreter_test.cc @@ -1804,8 +1804,14 @@ class TestLazyDelegateProvider : public InterpreterTest { }; void InitWithLazyDelegate(int64_t delegate_flags, - bool create_dyanmic_tensor = false) { - TfLiteRegistration reg = {nullptr, nullptr, nullptr, nullptr}; + bool create_dyanmic_tensor = false, + bool return_error = false) { + TfLiteRegistration reg = {nullptr}; + if (return_error) { + reg.prepare = [](TfLiteContext* context, TfLiteNode* node) { + return kTfLiteError; + }; + } ASSERT_EQ(interpreter_.AddTensors(2), kTfLiteOk); interpreter_.SetInputs({0}); interpreter_.SetOutputs({1}); @@ -1836,6 +1842,16 @@ TEST_F(TestLazyDelegateProvider, ApplicationSuccess) { EXPECT_TRUE(HasDelegates()); } +TEST_F(TestLazyDelegateProvider, ApplicationFailure) { + InitWithLazyDelegate(kTfLiteDelegateFlagsNone, + false /* create_dyanmic_tensor */, + true /* return_error */); + EXPECT_EQ(kTfLiteError, interpreter_.AllocateTensors()); + // We clear Interpreter::lazy_delegate_providers_ after they are tried out. + EXPECT_TRUE(mutable_lazy_delegate_providers()->empty()); + EXPECT_FALSE(HasDelegates()); +} + TEST_F(TestLazyDelegateProvider, ApplicationSkipped) { InitWithLazyDelegate(kTfLiteDelegateFlagsNone, true /* create_dyanmic_tensor */);