Support multiple delegates, especially static ones when no dynamic tensors are present. Also add new tests that focus on multiple delegates.
PiperOrigin-RevId: 373413130
Change-Id: Ifd5bbdaf7c5837558212cf1db5c041df3842274c
diff --git a/tensorflow/lite/c/c_api_experimental_test.cc b/tensorflow/lite/c/c_api_experimental_test.cc
index 2670861..1e170c9 100644
--- a/tensorflow/lite/c/c_api_experimental_test.cc
+++ b/tensorflow/lite/c/c_api_experimental_test.cc
@@ -29,6 +29,7 @@
#include "tensorflow/lite/testing/util.h"
using testing::HasSubstr;
+using tflite::delegates::test_utils::SimpleDelegate;
using tflite::delegates::test_utils::TestDelegate;
namespace {
diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 0f58761..f6e3178 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -1594,44 +1594,8 @@
return kTfLiteDelegateError;
}
- // Restore delegation state if applicable.
- TF_LITE_ENSURE_STATUS(RedoAllDelegates());
-
- if (state_ == kStateInvokableAndImmutable) {
- ReportError(
- "ModifyGraphWithDelegate is disallowed when graph is immutable.");
- return kTfLiteApplicationError;
- }
-
- if (!(delegate->flags & kTfLiteDelegateFlagsAllowDynamicTensors)) {
- int last_execution_plan_index_prepared;
- TF_LITE_ENSURE_OK(
- &context_, PrepareOpsStartingAt(0, execution_plan_,
- &last_execution_plan_index_prepared));
- if (has_dynamic_tensors_) {
- // Make sure that we are in a defined ready state before returning.
- // Plan and allocate tensors before returning.
- TF_LITE_ENSURE_OK(&context_, EnsureMemoryAllocations());
- ReportError(
- "Attempting to use a delegate that only supports static-sized "
- "tensors with a graph that has dynamic-sized tensors.");
- return kTfLiteApplicationError;
- }
- }
-
- const bool was_invokable_before_delegate = state_ == kStateInvokable;
- if (delegates_applied_.empty()) {
- // This is the first delegate being applied, so remember original execution
- // plan.
- // TODO(b/119623453): Restore execution plan to this state if delegate
- // application fails.
- pre_delegation_execution_plan_ = execution_plan_;
- }
-
- // TODO(aselle): Consider if it is worth storing pointers to delegates.
- // Setup additional context interface.
- SwitchToDelegateContext();
-
+ // Resets delegation & leaves graph in consistent state if delegate status is
+ // not okay.
auto reset_delegation_if_not_ok = [this](TfLiteStatus status) {
if (status != kTfLiteOk) {
TF_LITE_ENSURE_STATUS(RemoveAllDelegates());
@@ -1643,14 +1607,59 @@
return kTfLiteOk;
};
- TfLiteStatus status = delegate->Prepare(&context_, delegate);
+ // STEP 1: Verify & prepare graph for delegation.
+ // ==============================================
+ // Restore delegation state if applicable.
+ TF_LITE_ENSURE_STATUS(RedoAllDelegates());
+
+ const bool delegate_supports_dynamic_shapes =
+ delegate->flags & kTfLiteDelegateFlagsAllowDynamicTensors;
+ const auto pre_delegation_state = state_;
+
+ if (state_ == kStateInvokableAndImmutable) {
+ // A delegate that doesn't support dynamic shapes was already applied, so
+ // we can assume tensor shapes have been propagated & there are no dynamic
+ // tensors.
+ // Reset the state to force tensor/op reallocation.
+ state_ = kStateUninvokable;
+ } else if (!delegate_supports_dynamic_shapes) {
+ // Check if graph has dynamic tensors by preparing ops.
+ int last_execution_plan_index_prepared;
+ TF_LITE_ENSURE_STATUS(PrepareOpsStartingAt(
+ 0, execution_plan_, &last_execution_plan_index_prepared));
+ if (has_dynamic_tensors_) {
+ TF_LITE_ENSURE_STATUS(EnsureMemoryAllocations());
+ ReportError(
+ "Attempting to use a delegate that only supports static-sized "
+ "tensors with a graph that has dynamic-sized tensors.");
+ return kTfLiteApplicationError;
+ }
+ }
+
+ if (delegates_applied_.empty()) {
+ // This is the first delegate being applied, so remember original execution
+ // plan.
+ // TODO(b/119623453): Restore execution plan to this state if delegate
+ // application fails.
+ pre_delegation_execution_plan_ = execution_plan_;
+ }
+
+ // STEP 2: Delegate replaces applicable nodes with delegate kernels.
+ // =================================================================
+
+ // Setup additional context interface.
+ SwitchToDelegateContext();
+ TfLiteStatus status = delegate->Prepare(&context_, delegate);
// Remove additional context info.
SwitchToKernelContext();
-
TF_LITE_ENSURE_STATUS(reset_delegation_if_not_ok(status));
- if (!(delegate->flags & kTfLiteDelegateFlagsAllowDynamicTensors)) {
+ // STEP 3: Leave graph in consistent state based on delegate & previous state.
+ // ===========================================================================
+
+ if (!delegate_supports_dynamic_shapes) {
+ // CASE 1: Current delegate does not support dynamic shapes.
// Reset the state to force tensor/op reallocation.
state_ = kStateUninvokable;
TF_LITE_ENSURE_STATUS(
@@ -1658,9 +1667,28 @@
// After using a delegate which doesn't support dynamic tensors, make the
// entire graph immutable.
state_ = kStateInvokableAndImmutable;
- } else if (was_invokable_before_delegate) {
- // If the graph was invokable prior to delegate application, flush
- // allocation now to leave it in a consistent state.
+ } else if (pre_delegation_state == kStateInvokableAndImmutable) {
+ // CASE 2: Current delegate supports dynamic shapes, but a previous one
+ // does not.
+ // Make sure new delegate didn't mark a tensor as dynamic.
+ int last_execution_plan_index_prepared;
+ TF_LITE_ENSURE_STATUS(reset_delegation_if_not_ok(PrepareOpsStartingAt(
+ 0, execution_plan_, &last_execution_plan_index_prepared)));
+ if (has_dynamic_tensors_) {
+ TF_LITE_ENSURE_STATUS(RemoveAllDelegates());
+ ReportError(
+ "Cannot allow dynamic tensors due to previous delegation, resetting "
+ "to original execution plan.");
+ return kTfLiteApplicationError;
+ }
+ // Redo memory allocations & ensure state is set back to original value.
+ TF_LITE_ENSURE_STATUS(
+ reset_delegation_if_not_ok(EnsureMemoryAllocations()));
+ state_ = kStateInvokableAndImmutable;
+ } else if (pre_delegation_state == kStateInvokable) {
+ // CASE 3: Current delegate supports dynamic shapes, and the graph was
+ // previously invokable.
+ // Flush allocation now to leave it in a consistent state.
TF_LITE_ENSURE_STATUS(
reset_delegation_if_not_ok(EnsureMemoryAllocations()));
}
diff --git a/tensorflow/lite/delegates/delegate_test.cc b/tensorflow/lite/delegates/delegate_test.cc
index 66904ea..8992c1e 100644
--- a/tensorflow/lite/delegates/delegate_test.cc
+++ b/tensorflow/lite/delegates/delegate_test.cc
@@ -38,8 +38,10 @@
namespace tflite {
namespace delegates {
+using test_utils::SimpleDelegate;
using test_utils::TestDelegate;
using test_utils::TestFP16Delegation;
+using test_utils::TestTwoDelegates;
namespace {
@@ -130,108 +132,6 @@
}
}
-TEST_F(TestDelegate, SecondDelegationPrepareFailure) {
- // First delegate only supports nodes 1, 2. Gets applied successfully.
- // This delegate should support dynamic tensors, otherwise the second won't be
- // applied.
- delegate_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({1, 2}, kTfLiteDelegateFlagsAllowDynamicTensors));
- // Second delegate supports node 0, but fails during the delegate-node's
- // Prepare.
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0}, kTfLiteDelegateFlagsNone, true /**fail_node_prepare**/));
-
- // Initially, execution plan has 3 nodes.
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- // First delegate should be applied successfully, yielding a plan with 2
- // nodes.
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
- // Second delegate won't get applied.
- // As a result, previous delegate should also get undone, restoring the
- // execution plan to its original state.
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteDelegateError);
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
-
- std::vector<float> input = {1.0f, 2.0f, 3.0f};
- std::vector<float> expected_output = {2.0f, 4.0f, 6.0f};
- constexpr int kOutputTensorIndex = 3;
- TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
-
- // Verify Invoke() behavior.
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
- interpreter_->Invoke();
- for (int i = 0; i < 3; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-}
-
-TEST_F(TestDelegate, SecondDelegationInvokeFailure) {
- delegate_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({1, 2}, kTfLiteDelegateFlagsAllowDynamicTensors));
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0}, kTfLiteDelegateFlagsNone, false /**fail_node_prepare**/,
- 0 /**min_ops_per_subset**/, true /**fail_node_invoke**/));
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- std::vector<float> input = {1.0f, 2.0f, 3.0f};
- // Outputs match the AddOp path, rather than delegate path.
- std::vector<float> expected_output = {2.0f, 4.0f, 6.0f};
- constexpr int kOutputTensorIndex = 3;
-
- // Verify Invoke() behavior to ensure Interpreter isn't broken.
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
- EXPECT_EQ(interpreter_->Invoke(), kTfLiteError);
- EXPECT_EQ(RemoveAllDelegates(), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
- TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
- for (int i = 0; i < 3; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-}
-
-// This test ensures that node indices in multi-delegate application are handled
-// correctly by the TFLite partitioning algorithm.
-TEST_F(TestDelegate, TwoDelegates_ExecutionPlanIndicesDifferent) {
- // First delegate supports nodes 0, 1.
- // After this delegation, the execution plan size is 2.
- delegate_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({0, 1}, kTfLiteDelegateFlagsAllowDynamicTensors));
- // Second delegate supports (original) node index 2.
- // The execution plan has 2 nodes, so this verifies that the partitioning
- // algorithm correctly refers to (original) node indices instead of execution
- // plan indices.
- delegate2_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({2}, kTfLiteDelegateFlagsNone));
-
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- // Verify Invoke works.
- ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
-}
-
TEST_F(TestDelegate, StaticDelegateMakesGraphImmutable) {
delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({0, 1, 2}));
ASSERT_EQ(
@@ -331,39 +231,6 @@
EXPECT_EQ(tensor->buffer_handle, kTfLiteNullBufferHandle);
}
-// We utilize delegation in such a way as to allow node subsets with a minimum
-// number of ops only.
-TEST_F(TestDelegate, TestDelegationWithPartitionPreview) {
- // We set kTfLiteDelegateFlagsAllowDynamicTensors to ensure the second
- // delegate can be applied.
- // Ops 0 and 2 are delegated but end up in the same partition (based on
- // dependency analysis). However, since min_ops_per_subset = 3, no delegation
- // takes place.
- delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0, 2}, kTfLiteDelegateFlagsAllowDynamicTensors,
- false /**fail_node_prepare**/, 3 /**min_ops_per_subset**/));
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate());
-
- // Original execution plan remains.
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- ASSERT_EQ(interpreter_->execution_plan()[0], 0);
- ASSERT_EQ(interpreter_->execution_plan()[1], 1);
- ASSERT_EQ(interpreter_->execution_plan()[2], 2);
-
- // Same ops supported, but min_ops_per_subset = 2.
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0, 2}, kTfLiteDelegateFlagsAllowDynamicTensors,
- false /**fail_node_prepare**/, 2 /**min_ops_per_subset**/));
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate());
-
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
- ASSERT_EQ(interpreter_->execution_plan()[0], 3);
- const auto* node_and_reg = interpreter_->node_and_registration(3);
- ASSERT_EQ(node_and_reg->second.custom_name,
- delegate2_->FakeFusedRegistration().custom_name);
- ASSERT_EQ(interpreter_->execution_plan()[1], 1);
-}
-
TEST_F(TestDelegate, TestResizeInputWithNonDynamicDelegate) {
delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({0, 1, 2}));
ASSERT_EQ(
@@ -377,9 +244,8 @@
ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 3}), kTfLiteOk);
ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 3}), kTfLiteOk);
ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- // This should fail, since the previous application of the delegate will be
- // re-done automatically, making the graph immutable again.
- ASSERT_NE(
+ // ModifyGraphWithDelegate shouldn't fail, but graph won't change.
+ ASSERT_EQ(
interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
kTfLiteOk);
// Ensure graph has been restored to its valid delegated state.
@@ -413,68 +279,6 @@
}
}
-TEST_F(TestDelegate, TestResizeInputWithMultipleDelegates) {
- // First delegate only supports node 0.
- // This delegate should support dynamic tensors, otherwise the second won't be
- // applied.
- delegate_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({0}, kTfLiteDelegateFlagsAllowDynamicTensors));
- // Second delegate supports nodes 1 & 2, and makes the graph immutable.
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({1, 2}));
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteOk);
- // Should be two delegates nodes.
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- // Try resizing input to same shape as before (which should be a No-op).
- ASSERT_EQ(interpreter_->ResizeInputTensor(0, {3}), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- // Resizing input tensors should temporarily restore original execution plan
- // of 3 nodes.
- ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 3}), kTfLiteOk);
- ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 3}), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- // This should fail, since the previous application of the delegate will be
- // re-done automatically, making the graph immutable again.
- ASSERT_NE(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- // Ensure graph has been restored to its valid delegated state.
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
- std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
- constexpr int kOutputTensorIndex = 2;
- TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
-
- // Verify Invoke() behavior.
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
- interpreter_->Invoke();
- for (int i = 0; i < 3; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-
- // Resize again, but call AllocateTensors as usual afterwards.
- ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 4}), kTfLiteOk);
- ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 4}), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 3);
- ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 4 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 4 * sizeof(float));
- interpreter_->Invoke();
- for (int i = 0; i < 4; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-}
-
// If a delegate sets kTfLiteDelegateFlagsRequirePropagatedShapes but not
// kTfLiteDelegateFlagsAllowDynamicTensors, the former is redundant.
TEST_F(TestDelegate, TestRequirePropagatedShapes_NonDynamicDelegate) {
@@ -509,10 +313,8 @@
// propagated by runtime.
int delegate_flags = kTfLiteDelegateFlagsAllowDynamicTensors |
kTfLiteDelegateFlagsRequirePropagatedShapes;
- delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0, 1, 2}, delegate_flags, false /**fail_node_prepare**/,
- 3 /**min_ops_per_subset**/, false /**fail_node_invoke**/,
- true /**automatic_shape_propagation**/));
+ delegate_ = SimpleDelegate::DelegateWithRuntimeShapePropagation(
+ {0, 1, 2}, delegate_flags, 3);
ASSERT_EQ(
interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
kTfLiteOk);
@@ -542,10 +344,8 @@
// Delegate sets both flags and in its Prepare, ensures that shapes have been
// propagated by runtime.
int delegate_flags = kTfLiteDelegateFlagsAllowDynamicTensors;
- delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0, 1, 2}, delegate_flags, false /**fail_node_prepare**/,
- 3 /**min_ops_per_subset**/, false /**fail_node_invoke**/,
- true /**automatic_shape_propagation**/));
+ delegate_ = SimpleDelegate::DelegateWithRuntimeShapePropagation(
+ {0, 1, 2}, delegate_flags, 3);
ASSERT_EQ(
interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
kTfLiteOk);
@@ -556,93 +356,6 @@
ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteError);
}
-TEST_F(TestDelegate, TestRequirePropagatedShapes_MultipleDelegates) {
- // First delegate needs to support dynamic tensors to allow second delegation.
- // This delegate does not require automatic propagation.
- delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {0}, kTfLiteDelegateFlagsAllowDynamicTensors,
- false /**fail_node_prepare**/, 1 /**min_ops_per_subset**/,
- false /**fail_node_invoke**/, false /**automatic_shape_propagation**/));
- // Second delegate supports nodes 1 & 2, and requires automatic shape
- // propagation.
- int delegate_flags = kTfLiteDelegateFlagsAllowDynamicTensors |
- kTfLiteDelegateFlagsRequirePropagatedShapes;
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
- {1, 2}, delegate_flags, false /**fail_node_prepare**/,
- 1 /**min_ops_per_subset**/, false /**fail_node_invoke**/,
- true /**automatic_shape_propagation**/));
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteOk);
- // Should be two delegate nodes.
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 4}), kTfLiteOk);
- ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 4}), kTfLiteOk);
- ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
- std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
- constexpr int kOutputTensorIndex = 2;
- TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
-
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 4 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 4 * sizeof(float));
- interpreter_->Invoke();
- for (int i = 0; i < 4; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-}
-
-TEST_F(TestDelegate, ReleaseNonPersistentMemoryWithDelegates) {
- // First delegate only supports node 0.
- // This delegate should support dynamic tensors, otherwise the second won't be
- // applied.
- delegate_ = std::unique_ptr<SimpleDelegate>(
- new SimpleDelegate({0}, kTfLiteDelegateFlagsAllowDynamicTensors));
- // Second delegate supports nodes 1 & 2, and makes the graph immutable.
- delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({1, 2}));
-
- // No-op.
- ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
-
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
- ASSERT_EQ(
- interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
- kTfLiteOk);
- // Should be two delegates nodes.
- ASSERT_EQ(interpreter_->execution_plan().size(), 2);
-
- ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
- ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
-
- // This should fail, since the graph is immutable.
- ASSERT_NE(
- interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
- kTfLiteOk);
-
- std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
- std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
- constexpr int kOutputTensorIndex = 2;
- TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
-
- // Verify Invoke() behavior.
- memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
- memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
- interpreter_->Invoke();
- for (int i = 0; i < 3; ++i) {
- EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
- }
-
- ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
-}
-
TEST_F(TestDelegate, TestCopyFromBufferInvoke) {
delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({0, 1, 2}));
TfLiteDelegate* delegate = delegate_->get_tf_lite_delegate();
@@ -804,6 +517,338 @@
delegate_->FakeFusedRegistration().custom_name);
}
+TEST_P(TestTwoDelegates, SecondDelegationPrepareFailure) {
+ auto delegate_flag_pair = GetParam();
+ // First delegate only supports nodes 1, 2. Gets applied successfully.
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({1, 2}, delegate_flag_pair.first));
+ // Second delegate supports node 0, but fails during the delegate-node's
+ // Prepare.
+ delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ {0}, delegate_flag_pair.second, true /**fail_node_prepare**/));
+
+ // Initially, execution plan has 3 nodes.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ if (delegate_flag_pair.first == kTfLiteDelegateFlagsAllowDynamicTensors &&
+ delegate_flag_pair.second == kTfLiteDelegateFlagsAllowDynamicTensors) {
+ // If both delegates support dynamic tensors, the execution plan isn't
+ // prepared by ModifyGraphWithDelegate unless the graph was previously
+ // invokable. This is mainly because dynamic tensors anyway need
+ // allocations during Invoke.
+ // But for this test, we call AllocateTensors() to trigger allocations.
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+ }
+ // First delegate should be applied successfully, yielding a plan with 2
+ // nodes.
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ // Second delegate won't get applied.
+ // As a result, previous delegate should also get undone, restoring the
+ // execution plan to its original state.
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteDelegateError);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+
+ std::vector<float> input = {1.0f, 2.0f, 3.0f};
+ std::vector<float> expected_output = {2.0f, 4.0f, 6.0f};
+ constexpr int kOutputTensorIndex = 3;
+ TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
+
+ // Verify Invoke() behavior.
+ memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
+ memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
+ interpreter_->Invoke();
+ for (int i = 0; i < 3; ++i) {
+ EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
+ }
+}
+
+TEST_P(TestTwoDelegates, SecondDelegationInvokeFailure) {
+ auto delegate_flag_pair = GetParam();
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({1, 2}, delegate_flag_pair.first));
+ delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ {0}, delegate_flag_pair.second, false /**fail_node_prepare**/,
+ 0 /**min_ops_per_subset**/, true /**fail_node_invoke**/));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+ if (delegate_flag_pair.first == kTfLiteDelegateFlagsAllowDynamicTensors &&
+ delegate_flag_pair.second == kTfLiteDelegateFlagsAllowDynamicTensors) {
+ // If both delegates support dynamic tensors, the execution plan isn't
+ // prepared by ModifyGraphWithDelegate unless the graph was previously
+ // invokable. This is mainly because dynamic tensors anyway need
+ // allocations during Invoke.
+ // Call AllocateTensors() to trigger allocations.
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+ }
+
+ std::vector<float> input = {1.0f, 2.0f, 3.0f};
+ // Outputs match the AddOp path, rather than delegate path.
+ std::vector<float> expected_output = {2.0f, 4.0f, 6.0f};
+ constexpr int kOutputTensorIndex = 3;
+
+ // Verify Invoke() behavior to ensure Interpreter isn't broken.
+ memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
+ memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
+ EXPECT_EQ(interpreter_->Invoke(), kTfLiteError);
+ EXPECT_EQ(RemoveAllDelegates(), kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ ASSERT_EQ(interpreter_->Invoke(), kTfLiteOk);
+ TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
+ for (int i = 0; i < 3; ++i) {
+ EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
+ }
+}
+
+// This test ensures that node indices in multi-delegate application are handled
+// correctly by the TFLite partitioning algorithm.
+TEST_P(TestTwoDelegates, NodeIndicesCorrectlyHandledAfterDelegation) {
+ auto delegate_flag_pair = GetParam();
+ // First delegate supports nodes 0, 1.
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({0, 1}, delegate_flag_pair.first));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ // Second delegate supports (original) node index 2.
+ // The execution plan has 2 nodes, so this verifies that the partitioning
+ // algorithm correctly refers to (original) node indices instead of execution
+ // plan indices.
+ delegate2_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({2}, delegate_flag_pair.second));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+}
+
+TEST_P(TestTwoDelegates, TestResizeInputTensors) {
+ auto delegate_flag_pair = GetParam();
+ // First delegate only supports node 0.
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({0}, delegate_flag_pair.first));
+ // Second delegate supports nodes 1 & 2.
+ delegate2_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({1, 2}, delegate_flag_pair.second));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ // Should be two delegated nodes.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ // Try resizing input to same shape as before (which should be a No-op).
+ ASSERT_EQ(interpreter_->ResizeInputTensor(0, {3}), kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ // Resize inputs to new shape.
+ ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 4}), kTfLiteOk);
+ ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 4}), kTfLiteOk);
+ if (delegate_flag_pair.first == kTfLiteDelegateFlagsAllowDynamicTensors &&
+ delegate_flag_pair.second == kTfLiteDelegateFlagsAllowDynamicTensors) {
+ // If both delegates support dynamic tensors, execution plan won't be reset.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+ } else {
+ // In the presence of a static delegate, the runtime will reset execution
+ // plan to its original state until AllocateTensors or
+ // ModifyGraphWithDelegate
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ }
+
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+ // Irrespective of whether one or more delegates support dynamic shapes,
+ // execution plan should have 2 (delegated) nodes now.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
+ std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
+ constexpr int kOutputTensorIndex = 2;
+ TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
+
+ // Verify Invoke() behavior.
+ memcpy(interpreter_->typed_tensor<float>(0), input.data(), 4 * sizeof(float));
+ memcpy(interpreter_->typed_tensor<float>(1), input.data(), 4 * sizeof(float));
+ interpreter_->Invoke();
+ for (int i = 0; i < 4; ++i) {
+ EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
+ }
+}
+
+// We utilize delegation in such a way as to allow node subsets with a minimum
+// number of ops only.
+TEST_P(TestTwoDelegates, TestDelegationWithPartitionPreview) {
+ auto delegate_flag_pair = GetParam();
+ // Ops 0 and 2 are delegated but end up in the same partition (based on
+ // dependency analysis). However, since min_ops_per_subset = 3, no delegation
+ // takes place.
+ delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ {0, 2}, delegate_flag_pair.first, false /**fail_node_prepare**/,
+ 3 /**min_ops_per_subset**/));
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate());
+
+ // Original execution plan remains.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ ASSERT_EQ(interpreter_->execution_plan()[0], 0);
+ ASSERT_EQ(interpreter_->execution_plan()[1], 1);
+ ASSERT_EQ(interpreter_->execution_plan()[2], 2);
+
+ // Same ops supported, but min_ops_per_subset = 2.
+ delegate2_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ {0, 2}, delegate_flag_pair.second, false /**fail_node_prepare**/,
+ 2 /**min_ops_per_subset**/));
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate());
+
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+ ASSERT_EQ(interpreter_->execution_plan()[0], 3);
+ const auto* node_and_reg = interpreter_->node_and_registration(3);
+ ASSERT_EQ(node_and_reg->second.custom_name,
+ delegate2_->FakeFusedRegistration().custom_name);
+ ASSERT_EQ(interpreter_->execution_plan()[1], 1);
+}
+
+TEST_P(TestTwoDelegates, TestRequirePropagatedShapes) {
+ // We do not use kTfLiteDelegateFlagsNone in this test, since shape
+ // propagation always requires the delegate to support dynamic tensors. This
+ // delegate does not require automatic propagation.
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({0}, kTfLiteDelegateFlagsAllowDynamicTensors));
+ // Second delegate supports nodes 1 & 2, and requires automatic shape
+ // propagation.
+ int delegate_flags = kTfLiteDelegateFlagsAllowDynamicTensors |
+ kTfLiteDelegateFlagsRequirePropagatedShapes;
+ delegate2_ = SimpleDelegate::DelegateWithRuntimeShapePropagation(
+ {1, 2}, delegate_flags, 1);
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ // Should be two delegate nodes.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ ASSERT_EQ(interpreter_->ResizeInputTensor(0, {1, 4}), kTfLiteOk);
+ ASSERT_EQ(interpreter_->ResizeInputTensor(1, {1, 4}), kTfLiteOk);
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
+ std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
+ constexpr int kOutputTensorIndex = 2;
+ TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
+
+ memcpy(interpreter_->typed_tensor<float>(0), input.data(), 4 * sizeof(float));
+ memcpy(interpreter_->typed_tensor<float>(1), input.data(), 4 * sizeof(float));
+ interpreter_->Invoke();
+ for (int i = 0; i < 4; ++i) {
+ EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
+ }
+}
+
+TEST_P(TestTwoDelegates, ReleaseNonPersistentMemoryWithDelegates) {
+ auto delegate_flag_pair = GetParam();
+ // First delegate only supports node 0.
+ delegate_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({0}, delegate_flag_pair.first));
+ // Second delegate supports nodes 1 & 2, and makes the graph immutable.
+ delegate2_ = std::unique_ptr<SimpleDelegate>(
+ new SimpleDelegate({1, 2}, delegate_flag_pair.second));
+
+ // No-op.
+ ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
+
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ // Should be two delegates nodes.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 2);
+
+ ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+
+ std::vector<float> input = {1.0f, 2.0f, 3.0f, 4.0f};
+ std::vector<float> expected_output = {2.0f, 4.0f, 6.0f, 8.0f};
+ constexpr int kOutputTensorIndex = 2;
+ TfLiteTensor* tensor = interpreter_->tensor(kOutputTensorIndex);
+
+ // Verify Invoke() behavior.
+ memcpy(interpreter_->typed_tensor<float>(0), input.data(), 3 * sizeof(float));
+ memcpy(interpreter_->typed_tensor<float>(1), input.data(), 3 * sizeof(float));
+ interpreter_->Invoke();
+ for (int i = 0; i < 3; ++i) {
+ EXPECT_EQ(tensor->data.f[i], expected_output[i]) << i;
+ }
+
+ ASSERT_EQ(interpreter_->ReleaseNonPersistentMemory(), kTfLiteOk);
+}
+
+// This test ensures that after a static delegate is applied, a future delegate
+// that accepts previous nodes doesn't make them dynamic.
+TEST_F(TestTwoDelegates, DynamicTensorBeforeStaticDelegate) {
+ // First delegate only supports node {1, 2}.
+ // This makes the graph immutable.
+ delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({1, 2}));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ // Second delegate supports node 0, & tries to mark its output as
+ // dynamic. This should result in kTfLiteApplicationError.
+ delegate2_ = SimpleDelegate::DelegateWithDynamicOutput({0});
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteApplicationError);
+ // Execution plan reset to original.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+}
+
+// Same as bove, except a tensor later in the graph is marked static.
+// Even in this case, to be safe, we return an error.
+TEST_F(TestTwoDelegates, DynamicTensorAfterStaticDelegate) {
+ // First delegate only supports node 0.
+ // This makes the graph immutable.
+ delegate_ = std::unique_ptr<SimpleDelegate>(new SimpleDelegate({0}));
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate()),
+ kTfLiteOk);
+ // Second delegate supports node 1, 2 & tries to mark its output as
+ // dynamic. This should result in kTfLiteApplicationError.
+ delegate2_ = SimpleDelegate::DelegateWithDynamicOutput({1, 2});
+ ASSERT_EQ(
+ interpreter_->ModifyGraphWithDelegate(delegate2_->get_tf_lite_delegate()),
+ kTfLiteApplicationError);
+ // Execution plan reset to original.
+ ASSERT_EQ(interpreter_->execution_plan().size(), 3);
+ ASSERT_EQ(interpreter_->AllocateTensors(), kTfLiteOk);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+ TestTwoDelegatesInstance, TestTwoDelegates,
+ ::testing::Values(std::make_pair(kTfLiteDelegateFlagsNone,
+ kTfLiteDelegateFlagsNone),
+ std::make_pair(kTfLiteDelegateFlagsAllowDynamicTensors,
+ kTfLiteDelegateFlagsNone),
+ std::make_pair(kTfLiteDelegateFlagsNone,
+ kTfLiteDelegateFlagsAllowDynamicTensors),
+ std::make_pair(kTfLiteDelegateFlagsAllowDynamicTensors,
+ kTfLiteDelegateFlagsAllowDynamicTensors)));
+
class TestDelegateWithDynamicTensors : public ::testing::Test {
protected:
void SetUp() override {
diff --git a/tensorflow/lite/delegates/delegate_test_util.cc b/tensorflow/lite/delegates/delegate_test_util.cc
index 878b106..0d62b30 100644
--- a/tensorflow/lite/delegates/delegate_test_util.cc
+++ b/tensorflow/lite/delegates/delegate_test_util.cc
@@ -89,12 +89,7 @@
return reg;
}
-void TestDelegate::SetUp() {
- interpreter_.reset(new Interpreter);
- SetUpSubgraph(&interpreter_->primary_subgraph());
-}
-
-void TestDelegate::SetUpSubgraph(Subgraph* subgraph) {
+void TestDelegation::SetUpSubgraph(Subgraph* subgraph) {
subgraph->AddTensors(5);
subgraph->SetInputs({0, 1});
subgraph->SetOutputs({3, 4});
@@ -120,23 +115,44 @@
&node_index_ignored);
}
+void TestDelegate::SetUp() {
+ interpreter_.reset(new Interpreter);
+ SetUpSubgraph(&interpreter_->primary_subgraph());
+}
+
void TestDelegate::TearDown() {
// Interpreter relies on delegate to free the resources properly. Thus
// the life cycle of delegate must be longer than interpreter.
interpreter_.reset();
delegate_.reset();
+ delegate2_.reset();
}
-TestDelegate::SimpleDelegate::SimpleDelegate(
- const std::vector<int>& nodes, int64_t delegate_flags,
- bool fail_node_prepare, int min_ops_per_subset, bool fail_node_invoke,
- bool automatic_shape_propagation, bool custom_op)
+void TestTwoDelegates::SetUp() {
+ interpreter_.reset(new Interpreter);
+ SetUpSubgraph(&interpreter_->primary_subgraph());
+}
+
+void TestTwoDelegates::TearDown() {
+ // Interpreter relies on delegate to free the resources properly. Thus
+ // the life cycle of delegate must be longer than interpreter.
+ interpreter_.reset();
+ delegate_.reset();
+ delegate2_.reset();
+}
+
+SimpleDelegate::SimpleDelegate(const std::vector<int>& nodes,
+ int64_t delegate_flags, bool fail_node_prepare,
+ int min_ops_per_subset, bool fail_node_invoke,
+ bool automatic_shape_propagation, bool custom_op,
+ bool set_output_tensor_dynamic)
: nodes_(nodes),
fail_delegate_node_prepare_(fail_node_prepare),
min_ops_per_subset_(min_ops_per_subset),
fail_delegate_node_invoke_(fail_node_invoke),
automatic_shape_propagation_(automatic_shape_propagation),
- custom_op_(custom_op) {
+ custom_op_(custom_op),
+ set_output_tensor_dynamic_(set_output_tensor_dynamic) {
delegate_.Prepare = [](TfLiteContext* context,
TfLiteDelegate* delegate) -> TfLiteStatus {
auto* simple = static_cast<SimpleDelegate*>(delegate->data_);
@@ -244,7 +260,7 @@
delegate_.flags = delegate_flags;
}
-TfLiteRegistration TestDelegate::SimpleDelegate::FakeFusedRegistration() {
+TfLiteRegistration SimpleDelegate::FakeFusedRegistration() {
TfLiteRegistration reg = {nullptr};
reg.custom_name = "fake_fused_op";
@@ -304,6 +320,13 @@
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
return kTfLiteError;
};
+ } else if (set_output_tensor_dynamic_) {
+ reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
+ TfLiteTensor* output;
+ TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
+ SetTensorToDynamic(output);
+ return kTfLiteOk;
+ };
} else {
reg.prepare = [](TfLiteContext* context, TfLiteNode* node) {
// Set output size to input size
@@ -328,6 +351,26 @@
return reg;
}
+std::unique_ptr<SimpleDelegate>
+SimpleDelegate::DelegateWithRuntimeShapePropagation(
+ const std::vector<int>& nodes, int64_t delegate_flags,
+ int min_ops_per_subset) {
+ return std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ nodes, delegate_flags, false /**fail_node_prepare**/,
+ min_ops_per_subset /**min_ops_per_subset**/, false /**fail_node_invoke**/,
+ true /**automatic_shape_propagation**/));
+}
+
+std::unique_ptr<SimpleDelegate> SimpleDelegate::DelegateWithDynamicOutput(
+ const std::vector<int>& nodes) {
+ // All params default except nodes & set_output_tensor_dynamic.
+ return std::unique_ptr<SimpleDelegate>(new SimpleDelegate(
+ nodes, kTfLiteDelegateFlagsAllowDynamicTensors,
+ false /**fail_node_prepare**/, 0 /**min_ops_per_subset**/,
+ false /**fail_node_invoke**/, false /**automatic_shape_propagation**/,
+ true /**custom_op**/, true /**set_output_tensor_dynamic**/));
+}
+
void TestFP16Delegation::SetUp() {
interpreter_.reset(new Interpreter);
interpreter_->AddTensors(13);
diff --git a/tensorflow/lite/delegates/delegate_test_util.h b/tensorflow/lite/delegates/delegate_test_util.h
index 914fd85..bea97e3 100644
--- a/tensorflow/lite/delegates/delegate_test_util.h
+++ b/tensorflow/lite/delegates/delegate_test_util.h
@@ -18,6 +18,7 @@
#include <stdint.h>
#include <memory>
+#include <utility>
#include <vector>
#include <gtest/gtest.h>
@@ -33,8 +34,73 @@
// tensor inputs to produce a tensor output.
TfLiteRegistration AddOpRegistration();
-// TestDelegate is a friend of Interpreter to access RemoveAllDelegates().
-class TestDelegate : public ::testing::Test {
+class SimpleDelegate {
+ public:
+ // Create a simple implementation of a TfLiteDelegate. We use the C++ class
+ // SimpleDelegate and it can produce a handle TfLiteDelegate that is
+ // value-copyable and compatible with TfLite.
+ //
+ // Parameters:
+ // nodes: Indices of the graph nodes that the delegate will handle.
+ // fail_node_prepare: To simulate failure of Delegate node's Prepare().
+ // min_ops_per_subset: If >0, partitioning preview is used to choose only
+ // those subsets with min_ops_per_subset number of nodes.
+ // fail_node_invoke: To simulate failure of Delegate node's Invoke().
+ // automatic_shape_propagation: This assumes that the runtime will
+ // propagate shapes using the original execution plan.
+ // custom_op: If true, the graph nodes specified in the 'nodes' parameter
+ // should be custom ops with name "my_add"; if false, they should be
+ // the builtin ADD operator.
+ // set_output_tensor_dynamic: If True, this delegate sets output tensor to
+ // as dynamic during kernel Prepare.
+ explicit SimpleDelegate(const std::vector<int>& nodes,
+ int64_t delegate_flags = kTfLiteDelegateFlagsNone,
+ bool fail_node_prepare = false,
+ int min_ops_per_subset = 0,
+ bool fail_node_invoke = false,
+ bool automatic_shape_propagation = false,
+ bool custom_op = true,
+ bool set_output_tensor_dynamic = false);
+
+ static std::unique_ptr<SimpleDelegate> DelegateWithRuntimeShapePropagation(
+ const std::vector<int>& nodes, int64_t delegate_flags,
+ int min_ops_per_subset);
+
+ static std::unique_ptr<SimpleDelegate> DelegateWithDynamicOutput(
+ const std::vector<int>& nodes);
+
+ TfLiteRegistration FakeFusedRegistration();
+
+ TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; }
+
+ int min_ops_per_subset() { return min_ops_per_subset_; }
+
+ private:
+ std::vector<int> nodes_;
+ TfLiteDelegate delegate_;
+ bool fail_delegate_node_prepare_ = false;
+ int min_ops_per_subset_ = 0;
+ bool fail_delegate_node_invoke_ = false;
+ bool automatic_shape_propagation_ = false;
+ bool custom_op_ = true;
+ bool set_output_tensor_dynamic_ = false;
+};
+
+// Base class for single/multiple delegate tests.
+// Friend of Interpreter to access RemoveAllDelegates().
+class TestDelegation {
+ protected:
+ TfLiteStatus RemoveAllDelegates() {
+ return interpreter_->RemoveAllDelegates();
+ }
+
+ void SetUpSubgraph(Subgraph* subgraph);
+
+ std::unique_ptr<Interpreter> interpreter_;
+};
+
+// Tests scenarios involving a single delegate.
+class TestDelegate : public TestDelegation, public ::testing::Test {
protected:
void SetUp() override;
@@ -44,55 +110,20 @@
TfLiteBufferHandle AllocateBufferHandle() { return ++last_allocated_handle_; }
- TfLiteStatus RemoveAllDelegates() {
- return interpreter_->RemoveAllDelegates();
- }
+ std::unique_ptr<SimpleDelegate> delegate_, delegate2_;
+};
- void SetUpSubgraph(Subgraph* subgraph);
-
+// Tests scenarios involving two delegates, parametrized by the first & second
+// delegate's flags.
+class TestTwoDelegates
+ : public TestDelegation,
+ public ::testing::TestWithParam<
+ std::pair<TfLiteDelegateFlags, TfLiteDelegateFlags>> {
protected:
- class SimpleDelegate {
- public:
- // Create a simple implementation of a TfLiteDelegate. We use the C++ class
- // SimpleDelegate and it can produce a handle TfLiteDelegate that is
- // value-copyable and compatible with TfLite.
- //
- // Parameters:
- // nodes: Indices of the graph nodes that the delegate will handle.
- // fail_node_prepare: To simulate failure of Delegate node's Prepare().
- // min_ops_per_subset: If >0, partitioning preview is used to choose only
- // those subsets with min_ops_per_subset number of nodes.
- // fail_node_invoke: To simulate failure of Delegate node's Invoke().
- // automatic_shape_propagation: This assumes that the runtime will
- // propagate shapes using the original execution plan.
- // custom_op: If true, the graph nodes specified in the 'nodes' parameter
- // should be custom ops with name "my_add"; if false, they should be
- // the builtin ADD operator.
- explicit SimpleDelegate(const std::vector<int>& nodes,
- int64_t delegate_flags = kTfLiteDelegateFlagsNone,
- bool fail_node_prepare = false,
- int min_ops_per_subset = 0,
- bool fail_node_invoke = false,
- bool automatic_shape_propagation = false,
- bool custom_op = true);
+ void SetUp() override;
- TfLiteRegistration FakeFusedRegistration();
+ void TearDown() override;
- TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; }
-
- int min_ops_per_subset() { return min_ops_per_subset_; }
-
- private:
- std::vector<int> nodes_;
- TfLiteDelegate delegate_;
- bool fail_delegate_node_prepare_ = false;
- int min_ops_per_subset_ = 0;
- bool fail_delegate_node_invoke_ = false;
- bool automatic_shape_propagation_ = false;
- bool custom_op_ = true;
- };
-
- std::unique_ptr<Interpreter> interpreter_;
std::unique_ptr<SimpleDelegate> delegate_, delegate2_;
};
diff --git a/tensorflow/lite/delegates/interpreter_utils_test.cc b/tensorflow/lite/delegates/interpreter_utils_test.cc
index d29b1f8..f30543a 100644
--- a/tensorflow/lite/delegates/interpreter_utils_test.cc
+++ b/tensorflow/lite/delegates/interpreter_utils_test.cc
@@ -29,6 +29,7 @@
namespace tflite {
namespace delegates {
+using test_utils::SimpleDelegate;
using test_utils::TestDelegate;
using test_utils::TestFP16Delegation;
diff --git a/tensorflow/lite/interpreter.h b/tensorflow/lite/interpreter.h
index 20ed8d2..c6232cd 100644
--- a/tensorflow/lite/interpreter.h
+++ b/tensorflow/lite/interpreter.h
@@ -52,7 +52,7 @@
class InterpreterUtils; // Class for friend declarations.
namespace test_utils {
-class TestDelegate; // Class for friend declarations.
+class TestDelegation; // Class for friend declarations.
} // namespace test_utils
} // namespace delegates
@@ -683,7 +683,7 @@
friend class InterpreterBuilder;
friend class tflite::InterpreterTest;
friend class tflite::delegates::InterpreterUtils;
- friend class tflite::delegates::test_utils::TestDelegate;
+ friend class tflite::delegates::test_utils::TestDelegation;
/// Set the value of an external context.
static void SetExternalContext(struct TfLiteContext* context,
