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 #include "tensorflow/lite/delegates/hexagon/builders/reshape_builder.h"
16
17 #include <stdint.h>
18
19 #include <limits>
20
21 #include "tensorflow/lite/c/builtin_op_data.h"
22 #include "tensorflow/lite/delegates/hexagon/hexagon_nn/hexagon_nn.h"
23 #include "tensorflow/lite/kernels/kernel_util.h"
24
25 namespace tflite {
26 namespace delegates {
27 namespace hexagon {
28 namespace {
29
PopulateOutputShapeFromTensor(const TfLiteTensor * shape_tensor,std::vector<int> * output_shape)30 void PopulateOutputShapeFromTensor(const TfLiteTensor* shape_tensor,
31 std::vector<int>* output_shape) {
32 for (int i = 0; i < shape_tensor->dims->data[0]; ++i) {
33 output_shape->push_back(shape_tensor->data.i32[i]);
34 }
35 }
36
PopulateShapeFromParam(const TfLiteReshapeParams * params,std::vector<int> * output_shape)37 void PopulateShapeFromParam(const TfLiteReshapeParams* params,
38 std::vector<int>* output_shape) {
39 // The function is returned above this line if the shape tensor is usable.
40 // Now fallback to the shape parameter in `TfLiteReshapeParams`.
41 int num_dimensions = params->num_dimensions;
42 if (num_dimensions == 1 && params->shape[0] == 0) {
43 // Legacy tflite models use a shape parameter of [0] to indicate scalars,
44 // so adjust accordingly. TODO(b/111614235): Allow zero-sized buffers during
45 // toco conversion.
46 num_dimensions = 0;
47 }
48 for (int i = 0; i < num_dimensions; ++i) {
49 output_shape->push_back(params->shape[i]);
50 }
51 }
52 } // namespace
53
PopulateSubGraph(const TfLiteIntArray * inputs,const TfLiteIntArray * outputs,TfLiteContext * context)54 TfLiteStatus ReshapeOpBuilder::PopulateSubGraph(const TfLiteIntArray* inputs,
55 const TfLiteIntArray* outputs,
56 TfLiteContext* context) {
57 // Input data tensor.
58 AddInput(graph_builder_->GetHexagonTensorId(inputs->data[0]));
59
60 // Output shape.
61 TfLiteTensor* shape_tensor = nullptr;
62 bool output_shape_is_dynamic = false;
63 if (inputs->size == 2) {
64 shape_tensor = &context->tensors[inputs->data[1]];
65 bool is_shape_tensor =
66 (shape_tensor->dims->size == 1 && shape_tensor->type == kTfLiteInt32);
67 // If tensor shape is dynamic, pass it along directly.
68 if (shape_tensor->allocation_type != kTfLiteMmapRo && is_shape_tensor) {
69 output_shape_is_dynamic = true;
70 AddInput(graph_builder_->GetHexagonTensorId(inputs->data[1]));
71 }
72 if (!is_shape_tensor) {
73 shape_tensor = nullptr;
74 }
75 }
76 if (!output_shape_is_dynamic) {
77 if (shape_tensor) {
78 PopulateOutputShapeFromTensor(shape_tensor, &output_shape_);
79 } else {
80 const TfLiteReshapeParams* reshape_params =
81 reinterpret_cast<const TfLiteReshapeParams*>(builtin_data_);
82 PopulateShapeFromParam(reshape_params, &output_shape_);
83 }
84 int num_elements_in_shape = static_cast<int>(output_shape_.size());
85 output_shape_shape_ = {1, 1, 1, num_elements_in_shape};
86 auto* shape_node = graph_builder_->AddConstNodeWithData(
87 output_shape_shape_.data(),
88 reinterpret_cast<char*>(output_shape_.data()),
89 sizeof(int) * num_elements_in_shape);
90 AddInput(TensorID(shape_node->GetID(), 0));
91 }
92
93 // Hexagon output for this node.
94 int output_batch_size, output_height_size, output_width_size,
95 output_depth_size;
96 GetDims(&output_batch_size, &output_height_size, &output_width_size,
97 &output_depth_size, context->tensors[outputs->data[0]].dims);
98 node_output_ = AddOutput(sizeof(uint8_t), 4,
99 {output_batch_size, output_height_size,
100 output_width_size, output_depth_size});
101
102 return kTfLiteOk;
103 }
104
RegisterOutputs(const TfLiteIntArray * outputs,TfLiteContext * context)105 TfLiteStatus ReshapeOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
106 TfLiteContext* context) {
107 // Should be only 1 output.
108 graph_builder_->AddTensorWithID(outputs->data[0], node_output_.first,
109 node_output_.second);
110 return kTfLiteOk;
111 }
112
~ReshapeOpBuilder()113 ReshapeOpBuilder::~ReshapeOpBuilder() {}
114
CreateReshapeBuilder(GraphBuilder * graph_builder,int op_type)115 OpBuilder* CreateReshapeBuilder(GraphBuilder* graph_builder, int op_type) {
116 return new ReshapeOpBuilder(graph_builder, op_type);
117 }
118
119 } // namespace hexagon
120 } // namespace delegates
121 } // namespace tflite
122