diff options
author | 2018-08-15 17:10:25 -0700 | |
---|---|---|
committer | 2018-08-15 17:17:23 -0700 | |
commit | ec9d2891c3b1a5b9fd851f67a3cccf8ebaff38ad (patch) | |
tree | 010be4c118c80f17478e4536b36e91d57d07c742 /tensorflow/contrib/lite/interpreter_test.cc | |
parent | c3ef5d7034a50ca1b500c6fabea9250d38628884 (diff) |
Internal change
PiperOrigin-RevId: 208910402
Diffstat (limited to 'tensorflow/contrib/lite/interpreter_test.cc')
-rw-r--r-- | tensorflow/contrib/lite/interpreter_test.cc | 58 |
1 files changed, 58 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/interpreter_test.cc b/tensorflow/contrib/lite/interpreter_test.cc index f00697826c..5bcf0927d8 100644 --- a/tensorflow/contrib/lite/interpreter_test.cc +++ b/tensorflow/contrib/lite/interpreter_test.cc @@ -26,6 +26,13 @@ namespace tflite { // InterpreterTest is a friend of Interpreter, so it can access context_. class InterpreterTest : public ::testing::Test { + public: + template <typename Delegate> + static TfLiteStatus ModifyGraphWithDelegate( + Interpreter* interpreter, std::unique_ptr<Delegate> delegate) { + return interpreter->ModifyGraphWithDelegate(std::move(delegate)); + } + protected: TfLiteContext* GetInterpreterContext() { return &interpreter_.context_; } @@ -1302,6 +1309,57 @@ TEST_F(TestDelegateWithDynamicTensors, AllowDynamicTensors) { ASSERT_EQ(interpreter_->execution_plan()[0], 1); } +TEST(TestDelegateOwnership, ProperlyDisposed) { + struct TfLiteInterpreterOwnedDelegate : public TfLiteDelegate { + TfLiteInterpreterOwnedDelegate(bool* destroyed, bool* prepared) + : destroyed(destroyed), prepared(prepared) { + Prepare = [](TfLiteContext*, TfLiteDelegate* delegate) -> TfLiteStatus { + *static_cast<TfLiteInterpreterOwnedDelegate*>(delegate)->prepared = + true; + return kTfLiteOk; + }; + } + ~TfLiteInterpreterOwnedDelegate() { *destroyed = true; } + + bool* destroyed; + bool* prepared; + }; + + // Construct a delegate with flags for indicating preparation/destruction. + bool destroyed = false; + bool prepared = false; + std::unique_ptr<TfLiteInterpreterOwnedDelegate> delegate( + new TfLiteInterpreterOwnedDelegate(&destroyed, &prepared)); + { + // Create an interpreter and assemble a simple graph. + Interpreter interpreter; + TfLiteRegistration registration = {nullptr, nullptr, nullptr, nullptr}; + ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk); + ASSERT_EQ(interpreter.SetInputs({0}), kTfLiteOk); + ASSERT_EQ(interpreter.SetOutputs({1}), kTfLiteOk); + ASSERT_EQ(interpreter.AddNodeWithParameters({0}, {1}, nullptr, 0, nullptr, + ®istration), + kTfLiteOk); + + // Pass delegate ownership to that interpreter. + ASSERT_EQ(InterpreterTest::ModifyGraphWithDelegate(&interpreter, + std::move(delegate)), + kTfLiteOk); + + // The delegate should be prepared as normal, and should be preserved. + EXPECT_TRUE(prepared); + EXPECT_FALSE(destroyed); + + // Interpreter interaction should not impact the delegate's validity. + interpreter.AllocateTensors(); + interpreter.Invoke(); + EXPECT_FALSE(destroyed); + } + + // Only after the interpreter is destroyed should the delegate be destroyed. + EXPECT_TRUE(destroyed); +} + } // namespace } // namespace tflite |