tflite: mtk_neuron: Add PrepareSyncDriver helper in neuron_delegate_test This CL extract the logic of preparing SyncDriver, making the code looks cleaner. This CL is expected to be no functional changes. BUG=b:374245241 TEST=neuron_delegate_test on navi Change-Id: Ia4b19b5f3a646ccb9fe0bf7f689dcfc64badcd64 Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/tflite/+/6437796 Reviewed-by: Shik Chen <shik@chromium.org> Tested-by: Tommy Chiang <ototot@google.com> Commit-Queue: Tommy Chiang <ototot@google.com>
diff --git a/delegate/mtk_neuron/neuron_delegate_test.cc b/delegate/mtk_neuron/neuron_delegate_test.cc index 859cfbe..506c7a0 100644 --- a/delegate/mtk_neuron/neuron_delegate_test.cc +++ b/delegate/mtk_neuron/neuron_delegate_test.cc
@@ -254,23 +254,39 @@ return buffer; } +std::unique_ptr<SyncDriver> PrepareSyncDriver( + const std::unique_ptr<FlatBufferModel>& model, TfLiteDelegatePtr delegate, + const std::vector<float>& input_buffer) { + if (model == nullptr || delegate == nullptr) { + return nullptr; + } + auto driver = SyncDriver::Create( + std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel())); + if (driver == nullptr) { + return nullptr; + } + driver->SetInputTensorBuffer( + "input", + AllocateTfLiteAlignedTensorBuffer(input_buffer.size() * sizeof(float))); + if (driver->SetInputTensorData("input", input_buffer) != kTfLiteOk) { + return nullptr; + } + if (driver->AllocateBuffers() != kTfLiteOk) { + return nullptr; + } + return driver; +} + TEST(NeuronDelegateTest, InferenceSelfAborted) { TFLiteSettingsT settings = GetDefaultSettings(); auto [model, buffer_size] = GetTimeConsumingModelAndBufferSize(); - ASSERT_NE(model, nullptr); std::vector<float> input = GetRandomizedFloatData(buffer_size); // Get the expected output with the given input. - TfLiteDelegatePtr delegate = GetDelegateFromSettings(&settings); - ASSERT_NE(delegate, nullptr); - auto driver = SyncDriver::Create( - std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel())); + auto driver = + PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input); ASSERT_NE(driver, nullptr); - driver->SetInputTensorBuffer( - "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float))); - ASSERT_EQ(driver->SetInputTensorData("input", input), kTfLiteOk); - ASSERT_EQ(driver->AllocateBuffers(), kTfLiteOk); ASSERT_EQ(driver->Invoke(), kTfLiteOk); std::vector<float> expected_output = driver->GetOutputTensorData<float>("output"); @@ -279,15 +295,8 @@ settings.mtk_neuron_settings->inference_abort_time_ms = 1; // Try inference the model with an abort time. - delegate = GetDelegateFromSettings(&settings); - ASSERT_NE(delegate, nullptr); - driver = SyncDriver::Create( - std::move(delegate), FlatBufferModel::BuildFromModel(model->GetModel())); + driver = PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input); ASSERT_NE(driver, nullptr); - driver->SetInputTensorBuffer( - "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float))); - ASSERT_EQ(driver->SetInputTensorData("input", input), kTfLiteOk); - ASSERT_EQ(driver->AllocateBuffers(), kTfLiteOk); // Check the model inference is properly aborted. absl::Time before_invoke = absl::Now(); @@ -309,30 +318,17 @@ TFLiteSettingsT settings = GetDefaultSettings(); auto [model, buffer_size] = GetTimeConsumingModelAndBufferSize(); - ASSERT_NE(model, nullptr); std::vector<float> input = GetRandomizedFloatData(buffer_size); // Create the first driver without inferencing. - TfLiteDelegatePtr delegate1 = GetDelegateFromSettings(&settings); - ASSERT_NE(delegate1, nullptr); - auto driver1 = SyncDriver::Create( - std::move(delegate1), FlatBufferModel::BuildFromModel(model->GetModel())); + auto driver1 = + PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input); ASSERT_NE(driver1, nullptr); - driver1->SetInputTensorBuffer( - "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float))); - ASSERT_EQ(driver1->SetInputTensorData("input", input), kTfLiteOk); - ASSERT_EQ(driver1->AllocateBuffers(), kTfLiteOk); // Create the second driver without inferencing. - TfLiteDelegatePtr delegate2 = GetDelegateFromSettings(&settings); - ASSERT_NE(delegate2, nullptr); - auto driver2 = SyncDriver::Create( - std::move(delegate2), FlatBufferModel::BuildFromModel(model->GetModel())); + auto driver2 = + PrepareSyncDriver(model, GetDelegateFromSettings(&settings), input); ASSERT_NE(driver2, nullptr); - driver2->SetInputTensorBuffer( - "input", AllocateTfLiteAlignedTensorBuffer(buffer_size * sizeof(float))); - ASSERT_EQ(driver2->SetInputTensorData("input", input), kTfLiteOk); - ASSERT_EQ(driver2->AllocateBuffers(), kTfLiteOk); // Create a child thread to inference the first driver. absl::Duration time_used1;