1 // Copyright 2020 Google LLC
2 //
3 // This source code is licensed under the BSD-style license found in the
4 // LICENSE file in the root directory of this source tree.
5
6 #include <assert.h>
7 #include <math.h>
8 #include <stddef.h>
9 #include <stdint.h>
10
11 #include <xnnpack.h>
12 #include <xnnpack/log.h>
13 #include <xnnpack/operator.h>
14 #include <xnnpack/params.h>
15 #include <xnnpack/subgraph.h>
16 #include <xnnpack/subgraph-validation.h>
17
18
create_square_root_operator(const struct xnn_node * node,const struct xnn_value * values,size_t num_values,struct xnn_operator_data * opdata,const struct xnn_caches * caches)19 static enum xnn_status create_square_root_operator(
20 const struct xnn_node* node,
21 const struct xnn_value* values,
22 size_t num_values,
23 struct xnn_operator_data* opdata,
24 const struct xnn_caches* caches)
25 {
26 assert(node->compute_type == xnn_compute_type_fp32);
27
28 assert(node->num_inputs == 1);
29 const uint32_t input_id = node->inputs[0];
30 assert(input_id != XNN_INVALID_VALUE_ID);
31 assert(input_id < num_values);
32
33 assert(node->num_outputs == 1);
34 const uint32_t output_id = node->outputs[0];
35 assert(output_id != XNN_INVALID_VALUE_ID);
36 assert(output_id < num_values);
37
38 const size_t num_input_dims = values[input_id].shape.num_dims;
39 const size_t channel_dim = num_input_dims == 0 ? 1 : values[input_id].shape.dim[num_input_dims - 1];
40
41 enum xnn_status status;
42 switch (node->compute_type) {
43 case xnn_compute_type_fp32:
44 status = xnn_create_square_root_nc_f32(
45 channel_dim /* channels */, channel_dim /* input stride */, channel_dim /* output stride */,
46 node->flags,
47 &opdata->operator_objects[0]);
48 break;
49 #ifndef XNN_NO_F16_OPERATORS
50 case xnn_compute_type_fp16:
51 status = xnn_create_square_root_nc_f16(
52 channel_dim /* channels */, channel_dim /* input stride */, channel_dim /* output stride */,
53 node->flags,
54 &opdata->operator_objects[0]);
55 break;
56 #endif // !defined(XNN_NO_F16_OPERATORS)
57 default:
58 XNN_UNREACHABLE;
59 }
60 if (status == xnn_status_success) {
61 opdata->batch_size = xnn_shape_multiply_non_channel_dims(&values[input_id].shape);
62 opdata->inputs[0] = input_id;
63 opdata->outputs[0] = output_id;
64 }
65 return status;
66 }
67
setup_square_root_operator(const struct xnn_operator_data * opdata,const struct xnn_blob * blobs,size_t num_blobs,pthreadpool_t threadpool)68 static enum xnn_status setup_square_root_operator(
69 const struct xnn_operator_data* opdata,
70 const struct xnn_blob* blobs,
71 size_t num_blobs,
72 pthreadpool_t threadpool)
73 {
74 const uint32_t input_id = opdata->inputs[0];
75 assert(input_id != XNN_INVALID_VALUE_ID);
76 assert(input_id < num_blobs);
77
78 const uint32_t output_id = opdata->outputs[0];
79 assert(output_id != XNN_INVALID_VALUE_ID);
80 assert(output_id < num_blobs);
81
82 const struct xnn_blob* input_blob = blobs + input_id;
83 const void* input_data = input_blob->data;
84 assert(input_data != NULL);
85
86 const struct xnn_blob* output_blob = blobs + output_id;
87 void* output_data = output_blob->data;
88 assert(output_data != NULL);
89
90 switch (opdata->operator_objects[0]->type) {
91 case xnn_operator_type_square_root_nc_f32:
92 return xnn_setup_square_root_nc_f32(
93 opdata->operator_objects[0],
94 opdata->batch_size,
95 input_data,
96 output_data,
97 threadpool);
98 #ifndef XNN_NO_F16_OPERATORS
99 case xnn_operator_type_square_root_nc_f16:
100 return xnn_setup_square_root_nc_f16(
101 opdata->operator_objects[0],
102 opdata->batch_size,
103 input_data,
104 output_data,
105 threadpool);
106 #endif // !defined(XNN_NO_F16_OPERATORS)
107 default:
108 XNN_UNREACHABLE;
109 }
110 }
111
xnn_define_square_root(xnn_subgraph_t subgraph,uint32_t input_id,uint32_t output_id,uint32_t flags)112 enum xnn_status xnn_define_square_root(
113 xnn_subgraph_t subgraph,
114 uint32_t input_id,
115 uint32_t output_id,
116 uint32_t flags)
117 {
118 enum xnn_status status;
119 if ((status = xnn_subgraph_check_xnnpack_initialized(xnn_node_type_square_root)) != xnn_status_success) {
120 return status;
121 }
122
123 if (input_id >= subgraph->num_values) {
124 xnn_log_error(
125 "failed to define %s operator with input ID #%" PRIu32 ": invalid Value ID",
126 xnn_node_type_to_string(xnn_node_type_square_root), input_id);
127 return xnn_status_invalid_parameter;
128 }
129
130 const struct xnn_value* input_value = &subgraph->values[input_id];
131 status = xnn_subgraph_check_input_type_dense(xnn_node_type_square_root, input_id, input_value);
132 if (status != xnn_status_success) {
133 return status;
134 }
135
136 switch (input_value->datatype) {
137 case xnn_datatype_fp32:
138 break;
139 default:
140 xnn_log_error(
141 "failed to define %s operator with input ID #%" PRIu32 ": unsupported Value datatype %s (%d)",
142 xnn_node_type_to_string(xnn_node_type_square_root), input_id,
143 xnn_datatype_to_string(input_value->datatype), input_value->datatype);
144 return xnn_status_invalid_parameter;
145 }
146
147 status = xnn_subgraph_check_output_node_id(xnn_node_type_square_root, output_id, subgraph->num_values);
148 if (status != xnn_status_success) {
149 return status;
150 }
151
152 const struct xnn_value* output_value = &subgraph->values[output_id];
153 status = xnn_subgraph_check_output_type_dense(xnn_node_type_square_root, output_id, output_value);
154 if (status != xnn_status_success) {
155 return status;
156 }
157
158 switch (output_value->datatype) {
159 case xnn_datatype_fp32:
160 break;
161 default:
162 xnn_log_error(
163 "failed to define %s operator with output ID #%" PRIu32 ": unsupported Value datatype %s (%d)",
164 xnn_node_type_to_string(xnn_node_type_square_root), output_id,
165 xnn_datatype_to_string(output_value->datatype), output_value->datatype);
166 return xnn_status_invalid_parameter;
167 }
168
169 struct xnn_node* node = xnn_subgraph_new_node(subgraph);
170 if (node == NULL) {
171 return xnn_status_out_of_memory;
172 }
173
174 node->type = xnn_node_type_square_root;
175 node->compute_type = xnn_compute_type_fp32;
176 node->num_inputs = 1;
177 node->inputs[0] = input_id;
178 node->num_outputs = 1;
179 node->outputs[0] = output_id;
180 node->flags = flags;
181
182 node->create = create_square_root_operator;
183 node->setup = setup_square_root_operator;
184
185 return xnn_status_success;
186 }
187