1 /**
2 * Copyright 2021 Huawei Technologies Co., Ltd
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 #include "nnacl/infer/control/tensorlist_reserve_infer.h"
18 #include "nnacl/infer/infer_register.h"
19 #include "nnacl/tensorlist_parameter.h"
20 #include "nnacl/tensorlist_c_utils.h"
21 #include "nnacl/tensor_c_utils.h"
22
TensorListReserveInferShape(const TensorC * const * inputs,size_t inputs_size,TensorC ** outputs,size_t outputs_size,OpParameter * parameter)23 int TensorListReserveInferShape(const TensorC *const *inputs, size_t inputs_size, TensorC **outputs,
24 size_t outputs_size, OpParameter *parameter) {
25 int check_ret = CheckAugmentWithMinSize(inputs, inputs_size, outputs, outputs_size, parameter, 2, 1);
26 if (check_ret != NNACL_OK) {
27 return check_ret;
28 }
29
30 TensorListParameter *reserve_param = (TensorListParameter *)parameter;
31 const TensorC *input0 = inputs[0];
32 int ele_shape_type = input0->data_type_;
33 if (ele_shape_type != kNumberTypeInt && ele_shape_type != kNumberTypeInt32) {
34 return NNACL_ERR;
35 }
36
37 TensorListC *output = (TensorListC *)(outputs[0]);
38 output->data_type_ = kObjectTypeTensorType;
39 output->format_ = Format_NHWC;
40 output->tensors_data_type_ = reserve_param->element_dtype_;
41
42 if (input0->data_ == NULL) {
43 return NNACL_INFER_INVALID;
44 }
45 int *ele_shape_ptr = (int *)(input0->data_);
46
47 const TensorC *input1 = inputs[1];
48 int num_ele_type = input1->data_type_;
49 if (num_ele_type != kNumberTypeInt && ele_shape_type != kNumberTypeInt32) {
50 return NNACL_ERR;
51 }
52 if (input1->data_ == NULL) {
53 return NNACL_INFER_INVALID;
54 }
55 if (GetElementNum(input1) != 1) {
56 return NNACL_ERR;
57 }
58 int num_elements = ((int *)(input1->data_))[0];
59 ShapeSet(output->element_shape_, &(output->element_shape_size_), ele_shape_ptr, (size_t)GetElementNum(input0));
60 output->element_num_ = (size_t)(num_elements);
61
62 vvector tmp_shape;
63 tmp_shape.size_ = (size_t)(num_elements);
64 tmp_shape.shape_ = (int **)malloc(tmp_shape.size_ * sizeof(int *));
65 if (tmp_shape.shape_ == NULL) {
66 return NNACL_NULL_PTR;
67 }
68 tmp_shape.shape_size_ = (int *)malloc(tmp_shape.size_ * sizeof(int));
69 if (tmp_shape.shape_size_ == NULL) {
70 free(tmp_shape.shape_);
71 return NNACL_NULL_PTR;
72 }
73
74 for (size_t i = 0; i < num_elements; ++i) {
75 tmp_shape.shape_size_[i] = output->element_shape_size_;
76 tmp_shape.shape_[i] = output->element_shape_;
77 }
78 int ret = MallocTensorListData(output, reserve_param->element_dtype_, &tmp_shape);
79 free(tmp_shape.shape_size_);
80 free(tmp_shape.shape_);
81 return ret;
82 }
83
84 REG_INFER(TensorListReserve, PrimType_TensorListReserve, TensorListReserveInferShape)
85