1 /* Copyright 2019 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 
16 // This module provides helper functions for testing the interaction between
17 // control flow ops and subgraphs.
18 // For convenience, we mostly only use `kTfLiteInt32` in this module.
19 
20 #ifndef TENSORFLOW_LITE_KERNELS_SUBGRAPH_TEST_UTIL_H_
21 #define TENSORFLOW_LITE_KERNELS_SUBGRAPH_TEST_UTIL_H_
22 
23 #include <stdint.h>
24 
25 #include <memory>
26 #include <vector>
27 
28 #include <gtest/gtest.h>
29 #include "tensorflow/lite/core/subgraph.h"
30 #include "tensorflow/lite/interpreter.h"
31 
32 namespace tflite {
33 namespace subgraph_test_util {
34 
35 // TODO(ycling): This file should be renamed as
36 // `control_flow_test_util` to avoid confusion. I'll do it immediately
37 // in a separated change.
38 class SubgraphBuilder {
39  public:
40   ~SubgraphBuilder();
41 
42   // Build a subgraph with a single Add op.
43   // 2 inputs. 1 output.
44   void BuildAddSubgraph(Subgraph* subgraph);
45 
46   // Build a subgraph with a single Mul op.
47   // 2 inputs. 1 output.
48   void BuildMulSubgraph(Subgraph* subgraph);
49 
50   // Build a subgraph with a single Pad op.
51   // 2 inputs. 1 output.
52   void BuildPadSubgraph(Subgraph* subgraph);
53 
54   // Build a subgraph with a single If op.
55   // 3 inputs:
56   //   The 1st input is condition with boolean type.
57   //   The 2nd and 3rd inputs are feed input the branch subgraphs.
58   // 1 output.
59   void BuildIfSubgraph(Subgraph* subgraph);
60 
61   // Build a subgraph with a single Less op.
62   // The subgraph is used as the condition subgraph for testing `While` op.
63   // 2 inputs:
64   //   The 1st input is a counter with `kTfLiteInt32` type.
65   //   The 2nd input is ignored in this subgraph.
66   // 1 output with `kTfLiteBool` type.
67   //   Equivalent to (input < rhs).
68   void BuildLessEqualCondSubgraph(Subgraph* subgraph, int rhs);
69 
70   // An accumulate loop body subgraph. Used to produce triangle number
71   // sequence. 2 inputs and 2 outputs
72   //   Equivalent to (counter, value) -> (counter + 1, counter + 1 + value)
73   void BuildAccumulateLoopBodySubgraph(Subgraph* subgraph);
74 
75   // A pad loop body subgraph. When used in a loop it will repeatively enlarge
76   // the
77   //   tensor.
78   // 2 inputs and 2 outputs.
79   //   Equivalent to (counter, value) -> (counter + 1, tf.pad(value, padding))
80   // Note the padding is created as a constant tensor.
81   void BuildPadLoopBodySubgraph(Subgraph* subgraph,
82                                 const std::vector<int> padding);
83 
84   // Build a subgraph with a single While op.
85   // 2 inputs, 2 outputs.
86   void BuildWhileSubgraph(Subgraph* subgraph);
87 
88   // Build a subgraph that assigns a random value to a variable.
89   // No input/output.
90   void BuildAssignRandomValueToVariableSubgraph(Subgraph* graph);
91 
92   // Build a subgraph with CallOnce op and ReadVariable op.
93   // No input and 1 output.
94   void BuildCallOnceAndReadVariableSubgraph(Subgraph* graph);
95 
96   // Build a subgraph with a single Less op.
97   // The subgraph is used as the condition subgraph for testing `While` op.
98   // 3 inputs:
99   //   The 1st and 2nd inputs are string tensors, which will be ignored.
100   //   The 3rd input is an integner value as a counter in this subgraph.
101   // 1 output with `kTfLiteBool` type.
102   //   Equivalent to (int_val < rhs).
103   void BuildLessEqualCondSubgraphWithDynamicTensor(Subgraph* subgraph, int rhs);
104 
105   // Build a subgraph with a single While op, which has 3 inputs and 3 outputs.
106   // This subgraph is used for creating/invoking dynamic allocated tensors based
107   // on string tensors.
108   //   Equivalent to (str1, str2, int_val) ->
109   //                 (str1, Fill(str1, int_val + 1), int_val + 1).
110   void BuildBodySubgraphWithDynamicTensor(Subgraph* subgraph);
111 
112   // Build a subgraph with a single While op, that contains 3 inputs and 3
113   // outputs (str1, str2, int_val).
114   void BuildWhileSubgraphWithDynamicTensor(Subgraph* subgraph);
115 
116  private:
117   void CreateConstantInt32Tensor(Subgraph* subgraph, int tensor_index,
118                                  const std::vector<int>& shape,
119                                  const std::vector<int>& data);
120   std::vector<void*> buffers_;
121 };
122 
123 class ControlFlowOpTest : public ::testing::Test {
124  public:
ControlFlowOpTest()125   ControlFlowOpTest()
126       : interpreter_(new Interpreter), builder_(new SubgraphBuilder) {}
127 
~ControlFlowOpTest()128   ~ControlFlowOpTest() override {
129     interpreter_.reset();
130     builder_.reset();
131   }
132 
133  protected:
134   std::unique_ptr<Interpreter> interpreter_;
135   std::unique_ptr<SubgraphBuilder> builder_;
136 };
137 
138 // Fill a `TfLiteTensor` with a 32-bits integer vector.
139 // Preconditions:
140 // * The tensor must have `kTfLiteInt32` type.
141 // * The tensor must be allocated.
142 // * The element count of the tensor must be equal to the length or
143 //   the vector.
144 void FillIntTensor(TfLiteTensor* tensor, const std::vector<int32_t>& data);
145 
146 // Fill a `TfLiteTensor` with a string value.
147 // Preconditions:
148 // * The tensor must have `kTfLitString` type.
149 void FillScalarStringTensor(TfLiteTensor* tensor, const std::string& data);
150 
151 // Check if the scalar string data of a tensor is as expected.
152 void CheckScalarStringTensor(const TfLiteTensor* tensor,
153                              const std::string& data);
154 
155 // Check if the shape and string data of a tensor is as expected.
156 void CheckStringTensor(const TfLiteTensor* tensor,
157                        const std::vector<int>& shape,
158                        const std::vector<std::string>& data);
159 
160 // Check if the shape and int32 data of a tensor is as expected.
161 void CheckIntTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
162                     const std::vector<int32_t>& data);
163 // Check if the shape and bool data of a tensor is as expected.
164 void CheckBoolTensor(const TfLiteTensor* tensor, const std::vector<int>& shape,
165                      const std::vector<bool>& data);
166 
167 }  // namespace subgraph_test_util
168 }  // namespace tflite
169 
170 #endif  // TENSORFLOW_LITE_KERNELS_SUBGRAPH_TEST_UTIL_H_
171