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_argmax_pooling_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_argmax_pooling_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 == 2);
34 const uint32_t output_value_id = node->outputs[0];
35 assert(output_value_id != XNN_INVALID_VALUE_ID);
36 assert(output_value_id < num_values);
37 const uint32_t output_index_id = node->outputs[1];
38 assert(output_index_id != XNN_INVALID_VALUE_ID);
39 assert(output_index_id < num_values);
40
41 const size_t channel_dim = values[input_id].shape.dim[3];
42 assert(channel_dim == values[output_value_id].shape.dim[3]);
43 assert(channel_dim == values[output_index_id].shape.dim[3]);
44
45 const enum xnn_status status = xnn_create_argmax_pooling2d_nhwc_f32(
46 node->params.pooling_2d.padding_top,
47 node->params.pooling_2d.padding_right,
48 node->params.pooling_2d.padding_bottom,
49 node->params.pooling_2d.padding_left,
50 node->params.pooling_2d.pooling_height,
51 node->params.pooling_2d.pooling_width,
52 channel_dim /* channels */, channel_dim /* input stride */, channel_dim /* output stride */,
53 node->flags,
54 &opdata->operator_objects[0]);
55 if (status == xnn_status_success) {
56 opdata->batch_size = values[input_id].shape.dim[0];
57 opdata->input_height = values[input_id].shape.dim[1];
58 opdata->input_width = values[input_id].shape.dim[2];
59 opdata->inputs[0] = input_id;
60 opdata->outputs[0] = output_value_id;
61 opdata->outputs[1] = output_index_id;
62 }
63 return status;
64 }
65
setup_argmax_pooling_operator(const struct xnn_operator_data * opdata,const struct xnn_blob * blobs,size_t num_blobs,pthreadpool_t threadpool)66 static enum xnn_status setup_argmax_pooling_operator(
67 const struct xnn_operator_data* opdata,
68 const struct xnn_blob* blobs,
69 size_t num_blobs,
70 pthreadpool_t threadpool)
71 {
72 const uint32_t input_id = opdata->inputs[0];
73 assert(input_id != XNN_INVALID_VALUE_ID);
74 assert(input_id < num_blobs);
75
76 const uint32_t output_value_id = opdata->outputs[0];
77 assert(output_value_id != XNN_INVALID_VALUE_ID);
78 assert(output_value_id < num_blobs);
79
80 const uint32_t output_index_id = opdata->outputs[1];
81 assert(output_index_id != XNN_INVALID_VALUE_ID);
82 assert(output_index_id < num_blobs);
83
84 const struct xnn_blob* input_blob = blobs + input_id;
85 const void* input_data = input_blob->data;
86 assert(input_data != NULL);
87
88 const struct xnn_blob* output_value_blob = blobs + output_value_id;
89 void* output_value_data = output_value_blob->data;
90 assert(output_value_data != NULL);
91
92 const struct xnn_blob* output_index_blob = blobs + output_index_id;
93 void* output_index_data = output_index_blob->data;
94 assert(output_index_data != NULL);
95
96 return xnn_setup_argmax_pooling2d_nhwc_f32(
97 opdata->operator_objects[0],
98 opdata->batch_size,
99 opdata->input_height,
100 opdata->input_width,
101 input_data,
102 output_value_data,
103 output_index_data,
104 threadpool);
105 }
106
xnn_define_argmax_pooling_2d(xnn_subgraph_t subgraph,uint32_t input_padding_top,uint32_t input_padding_right,uint32_t input_padding_bottom,uint32_t input_padding_left,uint32_t pooling_height,uint32_t pooling_width,uint32_t input_id,uint32_t output_value_id,uint32_t output_index_id,uint32_t flags)107 enum xnn_status xnn_define_argmax_pooling_2d(
108 xnn_subgraph_t subgraph,
109 uint32_t input_padding_top,
110 uint32_t input_padding_right,
111 uint32_t input_padding_bottom,
112 uint32_t input_padding_left,
113 uint32_t pooling_height,
114 uint32_t pooling_width,
115 uint32_t input_id,
116 uint32_t output_value_id,
117 uint32_t output_index_id,
118 uint32_t flags)
119 {
120 enum xnn_status status;
121 if ((status = xnn_subgraph_check_xnnpack_initialized(xnn_node_type_argmax_pooling_2d)) != xnn_status_success) {
122 return status;
123 }
124
125 const uint32_t pooling_size = pooling_height * pooling_width;
126 if (pooling_size == 0) {
127 xnn_log_error(
128 "failed to define %s operator with %" PRIu32 "x%" PRIu32 " pooling size: "
129 "pooling size dimensions must be non-zero",
130 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), pooling_width, pooling_height);
131 return xnn_status_invalid_parameter;
132 }
133
134 if (pooling_size == 1) {
135 xnn_log_error(
136 "failed to define %s operator with 1 pooling element: 1x1 pooling is meaningless",
137 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d));
138 return xnn_status_invalid_parameter;
139 }
140
141 if ((status = xnn_subgraph_check_input_node_id(xnn_node_type_argmax_pooling_2d, input_id, subgraph->num_values))
142 != xnn_status_success) {
143 return status;
144 }
145
146 const struct xnn_value* input_value = &subgraph->values[input_id];
147 status = xnn_subgraph_check_input_type_dense(xnn_node_type_argmax_pooling_2d, input_id, input_value);
148 if (status != xnn_status_success) {
149 return status;
150 }
151
152 switch (input_value->datatype) {
153 case xnn_datatype_fp32:
154 break;
155 default:
156 xnn_log_error(
157 "failed to define %s operator with input ID #%" PRIu32 ": unsupported Value datatype %s (%d)",
158 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), input_id,
159 xnn_datatype_to_string(input_value->datatype), input_value->datatype);
160 return xnn_status_invalid_parameter;
161 }
162
163 if (output_value_id >= subgraph->num_values) {
164 xnn_log_error(
165 "failed to define %s operator with output value ID #%" PRIu32 ": invalid Value ID",
166 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), output_value_id);
167 return xnn_status_invalid_parameter;
168 }
169
170 const struct xnn_value* output_value_value = &subgraph->values[output_value_id];
171 if (output_value_value->type != xnn_value_type_dense_tensor) {
172 xnn_log_error(
173 "failed to define %s operator with output value ID #%" PRIu32 ": unsupported Value type %d (expected dense tensor)",
174 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), output_value_id, output_value_value->type);
175 return xnn_status_invalid_parameter;
176 }
177
178 switch (output_value_value->datatype) {
179 case xnn_datatype_fp32:
180 break;
181 default:
182 xnn_log_error(
183 "failed to define %s operator with output value ID #%" PRIu32 ": unsupported Value datatype %s (%d)",
184 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), output_value_id,
185 xnn_datatype_to_string(output_value_value->datatype), output_value_value->datatype);
186 return xnn_status_invalid_parameter;
187 }
188
189 if (output_index_id >= subgraph->num_values) {
190 xnn_log_error(
191 "failed to define %s operator with output index ID #%" PRIu32 ": invalid Value ID",
192 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), output_index_id);
193 return xnn_status_invalid_parameter;
194 }
195
196 const struct xnn_value* output_index_value = &subgraph->values[output_index_id];
197 if (output_index_value->type != xnn_value_type_dense_tensor) {
198 xnn_log_error(
199 "failed to define %s operator with output index ID #%" PRIu32 ": unsupported Value type %d (expected dense tensor)",
200 xnn_node_type_to_string(xnn_node_type_argmax_pooling_2d), output_index_id, output_index_value->type);
201 return xnn_status_invalid_parameter;
202 }
203
204 struct xnn_node* node = xnn_subgraph_new_node(subgraph);
205 if (node == NULL) {
206 return xnn_status_out_of_memory;
207 }
208
209 node->type = xnn_node_type_argmax_pooling_2d;
210 node->compute_type = xnn_compute_type_fp32;
211 node->params.pooling_2d.padding_top = input_padding_top;
212 node->params.pooling_2d.padding_right = input_padding_right;
213 node->params.pooling_2d.padding_bottom = input_padding_bottom;
214 node->params.pooling_2d.padding_left = input_padding_left;
215 node->params.pooling_2d.pooling_height = pooling_height;
216 node->params.pooling_2d.pooling_width = pooling_width;
217 node->num_inputs = 1;
218 node->inputs[0] = input_id;
219 node->num_outputs = 2;
220 node->outputs[0] = output_value_id;
221 node->outputs[1] = output_index_id;
222 node->flags = flags;
223
224 node->create = create_argmax_pooling_operator;
225 node->setup = setup_argmax_pooling_operator;
226
227 return xnn_status_success;
228 }
229