1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_ 16 #define TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_ 17 18 #include <stdint.h> 19 20 #include <memory> 21 #include <vector> 22 23 #include <gtest/gtest.h> 24 #include "third_party/eigen3/Eigen/Core" 25 #include "tensorflow/lite/interpreter.h" 26 #include "tensorflow/lite/kernels/internal/compatibility.h" 27 28 namespace tflite { 29 namespace delegates { 30 namespace test_utils { 31 32 // Build a kernel registration for a custom addition op that adds its two 33 // tensor inputs to produce a tensor output. 34 TfLiteRegistration AddOpRegistration(); 35 36 // TestDelegate is a friend of Interpreter to access RemoveAllDelegates(). 37 class TestDelegate : public ::testing::Test { 38 protected: 39 void SetUp() override; 40 41 void TearDown() override; 42 43 TfLiteBufferHandle last_allocated_handle_ = kTfLiteNullBufferHandle; 44 AllocateBufferHandle()45 TfLiteBufferHandle AllocateBufferHandle() { return ++last_allocated_handle_; } 46 RemoveAllDelegates()47 TfLiteStatus RemoveAllDelegates() { 48 return interpreter_->RemoveAllDelegates(); 49 } 50 51 void SetUpSubgraph(Subgraph* subgraph); 52 53 protected: 54 class SimpleDelegate { 55 public: 56 // Create a simple implementation of a TfLiteDelegate. We use the C++ class 57 // SimpleDelegate and it can produce a handle TfLiteDelegate that is 58 // value-copyable and compatible with TfLite. 59 // 60 // Parameters: 61 // nodes: Indices of the graph nodes that the delegate will handle. 62 // fail_node_prepare: To simulate failure of Delegate node's Prepare(). 63 // min_ops_per_subset: If >0, partitioning preview is used to choose only 64 // those subsets with min_ops_per_subset number of nodes. 65 // fail_node_invoke: To simulate failure of Delegate node's Invoke(). 66 // automatic_shape_propagation: This assumes that the runtime will 67 // propagate shapes using the original execution plan. 68 // custom_op: If true, the graph nodes specified in the 'nodes' parameter 69 // should be custom ops with name "my_add"; if false, they should be 70 // the builtin ADD operator. 71 explicit SimpleDelegate(const std::vector<int>& nodes, 72 int64_t delegate_flags = kTfLiteDelegateFlagsNone, 73 bool fail_node_prepare = false, 74 int min_ops_per_subset = 0, 75 bool fail_node_invoke = false, 76 bool automatic_shape_propagation = false, 77 bool custom_op = true); 78 79 TfLiteRegistration FakeFusedRegistration(); 80 get_tf_lite_delegate()81 TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; } 82 min_ops_per_subset()83 int min_ops_per_subset() { return min_ops_per_subset_; } 84 85 private: 86 std::vector<int> nodes_; 87 TfLiteDelegate delegate_; 88 bool fail_delegate_node_prepare_ = false; 89 int min_ops_per_subset_ = 0; 90 bool fail_delegate_node_invoke_ = false; 91 bool automatic_shape_propagation_ = false; 92 bool custom_op_ = true; 93 }; 94 95 std::unique_ptr<Interpreter> interpreter_; 96 std::unique_ptr<SimpleDelegate> delegate_, delegate2_; 97 }; 98 99 // Tests delegate functionality related to FP16 graphs. 100 // Model architecture: 101 // 1->DEQ->2 4->DEQ->5 7->DEQ->8 10->DEQ->11 102 // | | | | 103 // 0----->ADD->3----->ADD->6----->MUL->9------>ADD-->12 104 // Input: 0, Output:12. 105 // All constants are 2, so the function is: (x + 2 + 2) * 2 + 2 = 2x + 10 106 // 107 // Delegate only supports ADD, so can have up to two delegated partitions. 108 // TODO(b/156707497): Add more cases here once we have landed CPU kernels 109 // supporting FP16. 110 class TestFP16Delegation : public ::testing::TestWithParam<int> { 111 protected: 112 void SetUp() override; 113 114 void VerifyInvoke(); 115 TearDown()116 void TearDown() override { interpreter_.reset(); } 117 118 protected: 119 class FP16Delegate { 120 public: 121 // Uses FP16GraphPartitionHelper to accept ADD nodes with fp16 input. 122 explicit FP16Delegate(int num_delegated_subsets, 123 bool fail_node_prepare = false, 124 bool fail_node_invoke = false); 125 126 TfLiteRegistration FakeFusedRegistration(); 127 get_tf_lite_delegate()128 TfLiteDelegate* get_tf_lite_delegate() { return &delegate_; } 129 num_delegated_subsets()130 int num_delegated_subsets() { return num_delegated_subsets_; } 131 132 private: 133 TfLiteDelegate delegate_; 134 int num_delegated_subsets_; 135 bool fail_delegate_node_prepare_ = false; 136 bool fail_delegate_node_invoke_ = false; 137 }; 138 139 std::unique_ptr<Interpreter> interpreter_; 140 std::unique_ptr<FP16Delegate> delegate_; 141 Eigen::half float16_const_; 142 }; 143 144 } // namespace test_utils 145 } // namespace delegates 146 } // namespace tflite 147 148 #endif // TENSORFLOW_LITE_DELEGATES_DELEGATE_TEST_UTIL_ 149