1 // Copyright (c) Facebook, Inc. and its affiliates.
2 // All rights reserved.
3 //
4 // Copyright 2019 Google LLC
5 //
6 // This source code is licensed under the BSD-style license found in the
7 // LICENSE file in the root directory of this source tree.
8 
9 #include <assert.h>
10 #include <math.h>
11 #include <stdbool.h>
12 #include <stddef.h>
13 #include <stdint.h>
14 #include <stdlib.h>
15 #include <string.h>
16 
17 #include <xnnpack.h>
18 #include <xnnpack/allocator.h>
19 #include <xnnpack/operator.h>
20 #include <xnnpack/common.h>
21 #include <xnnpack/log.h>
22 #include <xnnpack/math.h>
23 #include <xnnpack/params-init.h>
24 #include <xnnpack/params.h>
25 #include <xnnpack/indirection.h>
26 
27 
compute_output_dimension(size_t padded_input_dimension,size_t pooling_dimension,size_t stride_dimension)28 static inline size_t compute_output_dimension(
29     size_t padded_input_dimension,
30     size_t pooling_dimension,
31     size_t stride_dimension)
32 {
33   return (padded_input_dimension - pooling_dimension) / stride_dimension + 1;
34 }
35 
compute_output_dimension_with_tf_same_padding(size_t input_dimension,size_t stride_dimension)36 static inline size_t compute_output_dimension_with_tf_same_padding(
37     size_t input_dimension,
38     size_t stride_dimension)
39 {
40   return divide_round_up(input_dimension, stride_dimension);
41 }
42 
xnn_create_average_pooling2d_nhwc_qu8(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 stride_height,uint32_t stride_width,size_t channels,size_t input_pixel_stride,size_t output_pixel_stride,uint8_t input_zero_point,float input_scale,uint8_t output_zero_point,float output_scale,uint8_t output_min,uint8_t output_max,uint32_t flags,xnn_operator_t * average_pooling_op_out)43 enum xnn_status xnn_create_average_pooling2d_nhwc_qu8(
44     uint32_t input_padding_top,
45     uint32_t input_padding_right,
46     uint32_t input_padding_bottom,
47     uint32_t input_padding_left,
48     uint32_t pooling_height,
49     uint32_t pooling_width,
50     uint32_t stride_height,
51     uint32_t stride_width,
52     size_t channels,
53     size_t input_pixel_stride,
54     size_t output_pixel_stride,
55     uint8_t input_zero_point,
56     float input_scale,
57     uint8_t output_zero_point,
58     float output_scale,
59     uint8_t output_min,
60     uint8_t output_max,
61     uint32_t flags,
62     xnn_operator_t* average_pooling_op_out)
63 {
64   xnn_operator_t average_pooling_op = NULL;
65   enum xnn_status status = xnn_status_uninitialized;
66 
67   if ((xnn_params.init_flags & XNN_INIT_FLAG_XNNPACK) == 0) {
68     xnn_log_error("failed to create %s operator: XNNPACK is not initialized",
69       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8));
70     goto error;
71   }
72 
73   status = xnn_status_invalid_parameter;
74 
75   const uint32_t pooling_size = pooling_height * pooling_width;
76   if (pooling_size == 0) {
77     xnn_log_error(
78       "failed to create %s operator with %" PRIu32 "x%" PRIu32 " pooling size: "
79       "pooling size dimensions must be non-zero",
80       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), pooling_width, pooling_height);
81     goto error;
82   }
83 
84   if (pooling_size == 1) {
85     xnn_log_error(
86       "failed to create %s operator with 1 pooling element: 1x1 pooling is meaningless",
87       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8));
88     goto error;
89   }
90 
91   if (stride_height == 0 || stride_width == 0) {
92     xnn_log_error(
93       "failed to create %s operator with %" PRIu32 "x%" PRIu32 " stride: stride dimensions must be non-zero",
94       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), stride_width, stride_height);
95     goto error;
96   }
97 
98   if (channels == 0) {
99     xnn_log_error(
100       "failed to create %s operator with %zu channels: number of channels must be non-zero",
101       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), channels);
102     goto error;
103   }
104 
105   if (input_pixel_stride < channels) {
106     xnn_log_error(
107       "failed to create %s operator with input pixel stride of %zu: "
108       "stride must be at least as large as the number of channels (%zu)",
109       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), input_pixel_stride, channels);
110     goto error;
111   }
112 
113   if (output_pixel_stride < channels) {
114     xnn_log_error(
115       "failed to create %s operator with output pixel stride of %zu: "
116       "stride must be at least as large as the number of channels (%zu)",
117       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), output_pixel_stride, channels);
118     goto error;
119   }
120 
121   if (input_scale <= 0.0f || !isnormal(input_scale)) {
122     xnn_log_error(
123       "failed to create %s operator with %.7g input scale: scale must be finite, normalized, and positive",
124       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), input_scale);
125     goto error;
126   }
127 
128   if (output_scale <= 0.0f || !isnormal(output_scale)) {
129     xnn_log_error(
130       "failed to create %s operator with %.7g output scale: scale must be finite, normalized, and positive",
131       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), output_scale);
132     goto error;
133   }
134 
135   if (output_min >= output_max) {
136     xnn_log_error(
137       "failed to create %s operator with [%" PRIu8 ", %" PRIu8 "] output range: range min must be below range max",
138       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8), output_min, output_max);
139     goto error;
140   }
141 
142   const bool any_padding = (input_padding_left | input_padding_top | input_padding_right | input_padding_bottom) != 0;
143   if ((flags & XNN_FLAG_TENSORFLOW_SAME_PADDING) != 0) {
144     if (any_padding) {
145       xnn_log_error(
146         "failed to create %s operator with %" PRIu32 "+%" PRIu32 "x%" PRIu32 "+%" PRIu32" padding: "
147         "TensorFlow SAME padding can't be combined with explicit padding specification",
148         xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8),
149         input_padding_top, input_padding_left, input_padding_bottom, input_padding_right);
150       goto error;
151     }
152   }
153 
154   status = xnn_status_unsupported_parameter;
155 
156   const float input_output_scale = input_scale / output_scale;
157   if (input_output_scale < 0x1.0p-8f || input_output_scale >= 0x1.0p+8f) {
158     xnn_log_error(
159       "failed to create %s operator with %.7g input scale and %.7g output scale: "
160       "input-to-output scale ratio (%.7f) must be in [2**-8, 2**8) range",
161       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8),
162       input_scale, output_scale, input_output_scale);
163     goto error;
164   }
165 
166   if (pooling_size >= 16777216) {
167     xnn_log_error(
168       "failed to create %s operator with %"PRIu32" (%" PRIu32 "x%" PRIu32 ") pooling elements: "
169       "the number of elements in the pooling area must be below 2**24",
170       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8),
171       pooling_size, pooling_width, pooling_height);
172     goto error;
173   }
174 
175   status = xnn_status_out_of_memory;
176 
177   average_pooling_op = xnn_allocate_zero_simd_memory(sizeof(struct xnn_operator));
178   if (average_pooling_op == NULL) {
179     xnn_log_error(
180       "failed to allocate %zu bytes for %s operator descriptor",
181       sizeof(struct xnn_operator), xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8));
182     goto error;
183   }
184 
185   const size_t zero_bytes = channels * sizeof(uint8_t) + XNN_EXTRA_BYTES;
186   void* zero_buffer = xnn_allocate_simd_memory(zero_bytes);
187   if (zero_buffer == NULL) {
188     xnn_log_error(
189       "failed to allocate %zu bytes for %s operator zero padding",
190       zero_bytes, xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8));
191     goto error;
192   }
193   memset(zero_buffer, input_zero_point, channels * sizeof(uint8_t));
194   average_pooling_op->zero_buffer = zero_buffer;
195 
196   average_pooling_op->padding_top = input_padding_top;
197   average_pooling_op->padding_right = input_padding_right;
198   average_pooling_op->padding_bottom = input_padding_bottom;
199   average_pooling_op->padding_left = input_padding_left;
200 
201   average_pooling_op->kernel_height = pooling_height;
202   average_pooling_op->kernel_width = pooling_width;
203   average_pooling_op->stride_height = stride_height;
204   average_pooling_op->stride_width = stride_width;
205   average_pooling_op->dilation_height = 1;
206   average_pooling_op->dilation_width = 1;
207   average_pooling_op->channels = channels;
208   average_pooling_op->input_pixel_stride = input_pixel_stride;
209   average_pooling_op->output_pixel_stride = output_pixel_stride;
210 
211   average_pooling_op->input_zero_point = (int32_t) (uint32_t) input_zero_point;
212   average_pooling_op->output_zero_point = output_zero_point;
213   average_pooling_op->input_scale = input_scale;
214   average_pooling_op->output_scale = output_scale;
215   average_pooling_op->output_min = output_min;
216   average_pooling_op->output_max = output_max;
217 
218   // Number of rows read in the AVGPOOL micro-kernel.
219   const size_t avgpool_nrows =
220     round_up(doz(pooling_size, xnn_params.qu8.avgpool.mr), xnn_params.qu8.avgpool.qr) + xnn_params.qu8.avgpool.mr;
221   average_pooling_op->params.qu8_avgpool =
222     xnn_init_qu8_avgpool_params(
223       (int32_t) -((uint32_t) input_zero_point * (uint32_t) avgpool_nrows),
224       input_scale / (output_scale * (float) pooling_size),
225       output_zero_point, output_min, output_max);
226 
227   average_pooling_op->type = xnn_operator_type_average_pooling_nhwc_qu8;
228   average_pooling_op->ukernel.type = xnn_ukernel_type_average_pooling;
229   average_pooling_op->flags = flags;
230 
231   *average_pooling_op_out = average_pooling_op;
232   return xnn_status_success;
233 
234 error:
235   xnn_delete_operator(average_pooling_op);
236   return status;
237 }
238 
xnn_create_average_pooling2d_nhwc_f32(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 stride_height,uint32_t stride_width,size_t channels,size_t input_pixel_stride,size_t output_pixel_stride,float output_min,float output_max,uint32_t flags,xnn_operator_t * average_pooling_op_out)239 enum xnn_status xnn_create_average_pooling2d_nhwc_f32(
240     uint32_t input_padding_top,
241     uint32_t input_padding_right,
242     uint32_t input_padding_bottom,
243     uint32_t input_padding_left,
244     uint32_t pooling_height,
245     uint32_t pooling_width,
246     uint32_t stride_height,
247     uint32_t stride_width,
248     size_t channels,
249     size_t input_pixel_stride,
250     size_t output_pixel_stride,
251     float output_min,
252     float output_max,
253     uint32_t flags,
254     xnn_operator_t* average_pooling_op_out)
255 {
256   xnn_operator_t average_pooling_op = NULL;
257   enum xnn_status status = xnn_status_uninitialized;
258 
259   if ((xnn_params.init_flags & XNN_INIT_FLAG_XNNPACK) == 0) {
260     xnn_log_error("failed to create %s operator: XNNPACK is not initialized",
261       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
262     goto error;
263   }
264 
265   status = xnn_status_invalid_parameter;
266 
267   const uint32_t pooling_size = pooling_height * pooling_width;
268   if (pooling_size == 0) {
269     xnn_log_error(
270       "failed to create %s operator with %" PRIu32 "x%" PRIu32 " pooling size: "
271       "pooling size dimensions must be non-zero",
272       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), pooling_width, pooling_height);
273     goto error;
274   }
275 
276   if (pooling_size == 1) {
277     xnn_log_error(
278       "failed to create %s operator with 1 pooling element: 1x1 pooling is meaningless",
279       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
280     goto error;
281   }
282 
283   if (stride_height == 0 || stride_width == 0) {
284     xnn_log_error(
285       "failed to create %s operator with %" PRIu32 "x%" PRIu32 " stride: stride dimensions must be non-zero",
286       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), stride_width, stride_height);
287     goto error;
288   }
289 
290   if (channels == 0) {
291     xnn_log_error(
292       "failed to create %s operator with %zu channels: number of channels must be non-zero",
293       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), channels);
294     goto error;
295   }
296 
297   if (input_pixel_stride < channels) {
298     xnn_log_error(
299       "failed to create %s operator with input pixel stride of %zu: "
300       "stride must be at least as large as the number of channels (%zu)",
301       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), input_pixel_stride, channels);
302     goto error;
303   }
304 
305   if (output_pixel_stride < channels) {
306     xnn_log_error(
307       "failed to create %s operator with output pixel stride of %zu: "
308       "stride must be at least as large as the number of channels (%zu)",
309       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), output_pixel_stride, channels);
310     goto error;
311   }
312 
313   if (isnan(output_min)) {
314     xnn_log_error(
315       "failed to create %s operator with NaN output lower bound: lower bound must be non-NaN",
316       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
317     goto error;
318   }
319 
320   if (isnan(output_max)) {
321     xnn_log_error(
322       "failed to create %s operator with NaN output upper bound: upper bound must be non-NaN",
323       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
324     goto error;
325   }
326 
327   if (output_min >= output_max) {
328     xnn_log_error(
329       "failed to create %s operator with [%.7g, %.7g] output range: lower bound must be below upper bound",
330       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32), output_min, output_max);
331     goto error;
332   }
333 
334   const bool any_padding = (input_padding_left | input_padding_top | input_padding_right | input_padding_bottom) != 0;
335   if ((flags & XNN_FLAG_TENSORFLOW_SAME_PADDING) != 0) {
336     if (any_padding) {
337       xnn_log_error(
338         "failed to create %s operator with %" PRIu32 "+%" PRIu32 "x%" PRIu32 "+%" PRIu32" padding: "
339         "TensorFlow SAME padding can't be combined with explicit padding specification",
340         xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32),
341         input_padding_top, input_padding_left, input_padding_bottom, input_padding_right);
342       goto error;
343     }
344   }
345 
346   status = xnn_status_out_of_memory;
347 
348   average_pooling_op = xnn_allocate_zero_simd_memory(sizeof(struct xnn_operator));
349   if (average_pooling_op == NULL) {
350     xnn_log_error(
351       "failed to allocate %zu bytes for %s operator descriptor",
352       sizeof(struct xnn_operator), xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
353     goto error;
354   }
355 
356   const size_t zero_bytes = channels * sizeof(float) + XNN_EXTRA_BYTES;
357   void* zero_buffer = xnn_allocate_zero_simd_memory(zero_bytes);
358   if (zero_buffer == NULL) {
359     xnn_log_error(
360       "failed to allocate %zu bytes for %s operator zero padding",
361       zero_bytes, xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32));
362     goto error;
363   }
364   average_pooling_op->zero_buffer = zero_buffer;
365 
366   average_pooling_op->padding_top = input_padding_top;
367   average_pooling_op->padding_right = input_padding_right;
368   average_pooling_op->padding_bottom = input_padding_bottom;
369   average_pooling_op->padding_left = input_padding_left;
370 
371   average_pooling_op->kernel_height = pooling_height;
372   average_pooling_op->kernel_width = pooling_width;
373   average_pooling_op->stride_height = stride_height;
374   average_pooling_op->stride_width = stride_width;
375   average_pooling_op->dilation_height = 1;
376   average_pooling_op->dilation_width = 1;
377   average_pooling_op->channels = channels;
378   average_pooling_op->input_pixel_stride = input_pixel_stride;
379   average_pooling_op->output_pixel_stride = output_pixel_stride;
380 
381   average_pooling_op->type = xnn_operator_type_average_pooling_nhwc_f32;
382   average_pooling_op->params.f32_scaleminmax =
383     xnn_init_f32_scaleminmax_params(1.0f / (float) pooling_size, output_min, output_max);
384   const bool tf_same_padding = (flags & XNN_FLAG_TENSORFLOW_SAME_PADDING) != 0;
385   if (any_padding || tf_same_padding) {
386     average_pooling_op->params.f32_minmax =
387       xnn_init_f32_minmax_params(output_min, output_max);
388     average_pooling_op->ukernel.type = xnn_ukernel_type_pixelwise_average_pooling;
389   } else {
390     average_pooling_op->ukernel.type = xnn_ukernel_type_average_pooling;
391   }
392   average_pooling_op->flags = flags;
393 
394   *average_pooling_op_out = average_pooling_op;
395   return xnn_status_success;
396 
397 error:
398   xnn_delete_operator(average_pooling_op);
399   return status;
400 }
401 
setup_average_pooling2d(xnn_operator_t average_pooling_op,size_t batch_size,size_t input_height,size_t input_width,const void * input,void * output,uint32_t log2_input_element_size,uint32_t log2_output_element_size,struct avgpool_parameters avgpool[restrict XNN_MIN_ELEMENTS (1)],struct pavgpool_parameters pavgpool[restrict1],struct gavgpool_parameters gavgpool[restrict XNN_MIN_ELEMENTS (1)],const void * params,size_t params_size,const void * global_params,size_t global_params_size,size_t num_threads,bool is_pixelwise)402 static enum xnn_status setup_average_pooling2d(
403   xnn_operator_t average_pooling_op,
404   size_t batch_size,
405   size_t input_height,
406   size_t input_width,
407   const void* input,
408   void* output,
409   uint32_t log2_input_element_size,
410   uint32_t log2_output_element_size,
411   struct avgpool_parameters avgpool[restrict XNN_MIN_ELEMENTS(1)],
412   struct pavgpool_parameters pavgpool[restrict 1],
413   struct gavgpool_parameters gavgpool[restrict XNN_MIN_ELEMENTS(1)],
414   const void* params,
415   size_t params_size,
416   const void* global_params,
417   size_t global_params_size,
418   size_t num_threads,
419   bool is_pixelwise)
420 {
421   assert(!is_pixelwise || pavgpool != NULL);
422 
423   average_pooling_op->state = xnn_run_state_invalid;
424 
425   if ((xnn_params.init_flags & XNN_INIT_FLAG_XNNPACK) == 0) {
426     xnn_log_error("failed to setup %s operator: XNNPACK is not initialized",
427       xnn_operator_type_to_string(average_pooling_op->type));
428     return xnn_status_uninitialized;
429   }
430 
431   if (input_width == 0 || input_height == 0) {
432     xnn_log_error(
433       "failed to setup %s operator with %zux%zu input: input dimensions must be non-zero",
434       xnn_operator_type_to_string(average_pooling_op->type), input_width, input_height);
435     return xnn_status_invalid_parameter;
436   }
437 
438   if (batch_size == 0) {
439     average_pooling_op->state = xnn_run_state_skip;
440     return xnn_status_success;
441   }
442 
443   average_pooling_op->input_height = input_height;
444   average_pooling_op->input_width = input_width;
445   average_pooling_op->input = input;
446 
447   const bool tf_same_padding = (average_pooling_op->flags & XNN_FLAG_TENSORFLOW_SAME_PADDING) != 0;
448   if (tf_same_padding) {
449     average_pooling_op->output_height = compute_output_dimension_with_tf_same_padding(
450         input_height, average_pooling_op->stride_height);
451     average_pooling_op->output_width = compute_output_dimension_with_tf_same_padding(
452         input_width, average_pooling_op->stride_width);
453 
454     const uint32_t kernel_height = average_pooling_op->kernel_height;
455     const uint32_t kernel_width = average_pooling_op->kernel_width;
456     const uint32_t total_padding_height =
457       (average_pooling_op->output_height - 1) * average_pooling_op->stride_height + kernel_height - input_height;
458     const uint32_t total_padding_width =
459       (average_pooling_op->output_width - 1) * average_pooling_op->stride_width + kernel_width - input_width;
460     average_pooling_op->padding_top = total_padding_height / 2;
461     average_pooling_op->padding_left = total_padding_width / 2;
462     average_pooling_op->padding_bottom = total_padding_height - average_pooling_op->padding_top;
463     average_pooling_op->padding_right = total_padding_width - average_pooling_op->padding_left;
464   } else {
465     average_pooling_op->output_height = compute_output_dimension(
466         average_pooling_op->padding_top + input_height + average_pooling_op->padding_bottom,
467         average_pooling_op->kernel_height,
468         average_pooling_op->stride_height);
469     average_pooling_op->output_width = compute_output_dimension(
470         average_pooling_op->padding_left + input_width + average_pooling_op->padding_right,
471         average_pooling_op->kernel_width,
472         average_pooling_op->stride_width);
473   }
474   average_pooling_op->output = output;
475 
476   const size_t output_height = average_pooling_op->output_height;
477   const size_t output_width = average_pooling_op->output_width;
478   const size_t padded_input_width = average_pooling_op->padding_left + input_width + average_pooling_op->padding_right;
479   const size_t padded_input_height = average_pooling_op->padding_top + input_height + average_pooling_op->padding_bottom;
480   if (padded_input_width == average_pooling_op->kernel_width && padded_input_height == average_pooling_op->kernel_height) {
481     // Global average pooling
482     const size_t input_elements = input_height * input_width;
483     const size_t input_stride_in_bytes = average_pooling_op->input_pixel_stride << log2_input_element_size;
484     const size_t channels = average_pooling_op->channels;
485     average_pooling_op->context.global_average_pooling_nwc = (struct global_average_pooling_nwc_context) {
486         .input = input,
487         .zero = average_pooling_op->zero_buffer,
488         .input_pixel_stride = input_stride_in_bytes,
489         .input_batch_stride = input_stride_in_bytes * input_elements,
490         .input_elements = input_elements,
491         .channels = channels,
492         .output = output,
493         .output_batch_stride = average_pooling_op->output_pixel_stride << log2_output_element_size,
494     };
495     memcpy(&average_pooling_op->context.global_average_pooling_nwc.params, global_params, global_params_size);
496     average_pooling_op->compute.type = xnn_parallelization_type_1d;
497     average_pooling_op->compute.range[0] = batch_size;
498 
499     if (input_elements <= gavgpool->mr) {
500       average_pooling_op->compute.task_1d = (pthreadpool_task_1d_t) xnn_compute_global_average_pooling_nwc_unipass;
501       average_pooling_op->context.global_average_pooling_nwc.unipass_ukernel = gavgpool->up;
502     } else {
503       average_pooling_op->compute.task_1d = (pthreadpool_task_1d_t) xnn_compute_global_average_pooling_nwc_multipass;
504       average_pooling_op->context.global_average_pooling_nwc.multipass_ukernel = gavgpool->mp;
505     }
506   } else {
507     // Non-global average pooling
508     const size_t pooling_height = average_pooling_op->kernel_height;
509     const size_t pooling_width = average_pooling_op->kernel_width;
510     const size_t pooling_size = pooling_height * pooling_width;
511 
512     const uint32_t mr = is_pixelwise ? pavgpool->mr : avgpool->mr;
513 
514     const size_t step_width = min(average_pooling_op->stride_width, pooling_width);
515     const size_t step_height = pooling_size + (output_width - 1) * step_width * pooling_height;
516 
517     const size_t last_input_height = average_pooling_op->last_input_height;
518     const size_t last_input_width = average_pooling_op->last_input_width;
519     if (input_height != last_input_height || input_width != last_input_width) {
520       // Micro-kernel may read up to (mr - 1) elements after the end of indirection buffer.
521       const size_t indirection_buffer_size = sizeof(void*) * ((mr - 1) + output_height * step_height);
522 
523       const void** indirection_buffer =
524         (const void**) xnn_reallocate_memory(average_pooling_op->indirection_buffer, indirection_buffer_size);
525       if (indirection_buffer == NULL) {
526         xnn_log_error("failed to allocate %zu bytes for %s operator indirection buffer",
527           indirection_buffer_size, xnn_operator_type_to_string(average_pooling_op->type));
528         return xnn_status_out_of_memory;
529       }
530       average_pooling_op->indirection_buffer = indirection_buffer;
531 
532       xnn_indirection_init_dwconv2d(average_pooling_op, step_height, step_width, log2_input_element_size);
533 
534       average_pooling_op->last_input = input;
535       average_pooling_op->last_input_height = input_height;
536       average_pooling_op->last_input_width = input_width;
537     }
538 
539     const size_t channels = average_pooling_op->channels;
540 
541     const size_t indirect_input_height_stride = step_height * sizeof(void*);
542     const size_t output_width_stride = average_pooling_op->output_pixel_stride << log2_output_element_size;
543     const size_t output_height_stride = output_width * output_width_stride;
544 
545     if (is_pixelwise) {
546       /* This part is specific to FP32, needs revision if another data types get a PAVGPOOL micro-kernel */
547       if (input_height != last_input_height || input_width != last_input_width) {
548         const size_t pixelwise_buffer_size = output_height * output_width * sizeof(float);
549         float* pixelwise_buffer =
550           (float*) xnn_reallocate_memory(average_pooling_op->pixelwise_buffer, pixelwise_buffer_size);
551         if (pixelwise_buffer == NULL) {
552           xnn_log_error("failed to allocate %zu bytes for %s operator pixelwise buffer",
553             pixelwise_buffer_size, xnn_operator_type_to_string(average_pooling_op->type));
554           return xnn_status_out_of_memory;
555         }
556         average_pooling_op->pixelwise_buffer = pixelwise_buffer;
557 
558         float* pixelwise_pointer = pixelwise_buffer;
559         for (size_t output_y = 0; output_y < output_height; output_y++) {
560           const size_t input_y_start = doz(output_y * average_pooling_op->stride_height, average_pooling_op->padding_top);
561           const size_t input_y_end =
562             min(doz(output_y * average_pooling_op->stride_height + average_pooling_op->kernel_height, average_pooling_op->padding_top), input_height);
563           const uint32_t input_y_range = (uint32_t) (input_y_end - input_y_start);
564           for (size_t output_x = 0; output_x < output_width; output_x++) {
565             const size_t input_x_start = doz(output_x * average_pooling_op->stride_width, average_pooling_op->padding_left);
566             const size_t input_x_end =
567               min(doz(output_x * average_pooling_op->stride_width + average_pooling_op->kernel_width, average_pooling_op->padding_left), input_width);
568             const uint32_t input_x_range = (uint32_t) (input_x_end - input_x_start);
569             *pixelwise_pointer++ = 1.0f / ((float) (int32_t) (input_y_range * input_x_range));
570           }
571         }
572       }
573 
574       const uint32_t qr = pavgpool->qr;
575       const size_t multipass_adjustment =
576         pooling_size > mr ? round_up(pooling_size - mr, qr) + mr - qr : 0;
577       average_pooling_op->context.pixelwise_average_pooling = (struct pixelwise_average_pooling_context) {
578         .indirect_input = average_pooling_op->indirection_buffer,
579         .indirect_input_height_stride = indirect_input_height_stride,
580         .input_batch_stride = input_height * input_width * average_pooling_op->input_pixel_stride << log2_input_element_size,
581         .input_offset = (size_t) ((uintptr_t) input - (uintptr_t) average_pooling_op->last_input),
582         .pixelwise_buffer = average_pooling_op->pixelwise_buffer,
583         .pixelwise_buffer_height_stride = output_width * sizeof(float),
584         .output = output,
585         .output_batch_stride = output_height * output_height_stride,
586         .output_height_stride = output_height_stride,
587         .output_width = output_width,
588         .pooling_size = pooling_size,
589         .channels = channels,
590         .zero = average_pooling_op->zero_buffer,
591         .input_increment = (pooling_height * step_width - multipass_adjustment) * sizeof(void*),
592         .output_increment = output_width_stride - (channels << log2_output_element_size),
593       };
594       memcpy(&average_pooling_op->context.pixelwise_average_pooling.params, params, params_size);
595       if (pooling_size <= mr) {
596         average_pooling_op->context.pixelwise_average_pooling.unipass_ukernel = pavgpool->up;
597         average_pooling_op->compute.task_2d = (pthreadpool_task_2d_t) xnn_compute_pixelwise_average_pooling_unipass;
598       } else {
599         average_pooling_op->context.pixelwise_average_pooling.multipass_ukernel = pavgpool->mp;
600         average_pooling_op->compute.task_2d = (pthreadpool_task_2d_t) xnn_compute_pixelwise_average_pooling_multipass;
601       }
602     } else {
603       const uint32_t qr = avgpool->qr;
604       const size_t multipass_adjustment =
605         pooling_size > mr ? round_up(pooling_size - mr, qr) + mr - qr : 0;
606       average_pooling_op->context.average_pooling = (struct average_pooling_context) {
607         .indirect_input = average_pooling_op->indirection_buffer,
608         .indirect_input_height_stride = indirect_input_height_stride,
609         .input_offset = (size_t) ((uintptr_t) input - (uintptr_t) average_pooling_op->last_input),
610         .input_batch_stride = input_height * input_width * average_pooling_op->input_pixel_stride << log2_input_element_size,
611         .output = output,
612         .output_batch_stride = output_height * output_height_stride,
613         .output_height_stride = output_height_stride,
614         .output_width = output_width,
615         .pooling_size = pooling_size,
616         .channels = channels,
617         .zero = average_pooling_op->zero_buffer,
618         .input_increment = (pooling_height * step_width - multipass_adjustment) * sizeof(void*),
619         .output_increment = output_width_stride - (channels << log2_output_element_size),
620         .params.f32 = average_pooling_op->params.f32_scaleminmax,
621       };
622       memcpy(&average_pooling_op->context.average_pooling.params, params, params_size);
623       if (pooling_size <= mr) {
624         average_pooling_op->context.average_pooling.unipass_ukernel = avgpool->up;
625         average_pooling_op->compute.task_2d = (pthreadpool_task_2d_t) xnn_compute_average_pooling_unipass;
626       } else {
627         average_pooling_op->context.average_pooling.multipass_ukernel = avgpool->mp;
628         average_pooling_op->compute.task_2d = (pthreadpool_task_2d_t) xnn_compute_average_pooling_multipass;
629       }
630     }
631     average_pooling_op->compute.type = xnn_parallelization_type_2d;
632     average_pooling_op->compute.range[0] = batch_size;
633     average_pooling_op->compute.range[1] = output_height;
634   }
635   average_pooling_op->state = xnn_run_state_ready;
636 
637   return xnn_status_success;
638 }
639 
xnn_setup_average_pooling2d_nhwc_qu8(xnn_operator_t average_pooling_op,size_t batch_size,size_t input_height,size_t input_width,const uint8_t * input,uint8_t * output,pthreadpool_t threadpool)640 enum xnn_status xnn_setup_average_pooling2d_nhwc_qu8(
641     xnn_operator_t average_pooling_op,
642     size_t batch_size,
643     size_t input_height,
644     size_t input_width,
645     const uint8_t* input,
646     uint8_t* output,
647     pthreadpool_t threadpool)
648 {
649   if (average_pooling_op->type != xnn_operator_type_average_pooling_nhwc_qu8) {
650     xnn_log_error("failed to setup operator: operator type mismatch (expected %s, got %s)",
651       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_qu8),
652       xnn_operator_type_to_string(average_pooling_op->type));
653     return xnn_status_invalid_parameter;
654   }
655 
656   assert(average_pooling_op->ukernel.type == xnn_ukernel_type_average_pooling);
657 
658   // Number of rows read in the GAVGPOOL micro-kernel.
659   const size_t input_size = input_height * input_width;
660   const size_t pooling_size = average_pooling_op->kernel_height * average_pooling_op->kernel_width;
661   const size_t gavgpool_nrows = round_up(input_size, xnn_params.qu8.gavgpool.mr);
662   average_pooling_op->params.qu8_gavgpool =
663     xnn_init_qu8_avgpool_params(
664       -(average_pooling_op->input_zero_point * (int32_t) gavgpool_nrows),
665       average_pooling_op->input_scale / (average_pooling_op->output_scale * (float) pooling_size),
666       average_pooling_op->output_zero_point,
667       average_pooling_op->output_min,
668       average_pooling_op->output_max);
669 
670   return setup_average_pooling2d(
671     average_pooling_op,
672     batch_size, input_height, input_width,
673     input, output,
674     0 /* log2(sizeof(input element)) = log2(sizeof(uint8_t)) */,
675     0 /* log2(sizeof(output element)) = log2(sizeof(uint8_t)) */,
676     &xnn_params.qu8.avgpool,
677     NULL /* no PAVGPOOL micro-kernel */,
678     &xnn_params.qu8.gavgpool,
679     &average_pooling_op->params.qu8_avgpool,
680     sizeof(average_pooling_op->params.qu8_avgpool),
681     &average_pooling_op->params.qu8_gavgpool,
682     sizeof(average_pooling_op->params.qu8_gavgpool),
683     pthreadpool_get_threads_count(threadpool),
684     false /* pixelwise not supported */);
685 }
686 
xnn_setup_average_pooling2d_nhwc_f32(xnn_operator_t average_pooling_op,size_t batch_size,size_t input_height,size_t input_width,const float * input,float * output,pthreadpool_t threadpool)687 enum xnn_status xnn_setup_average_pooling2d_nhwc_f32(
688     xnn_operator_t average_pooling_op,
689     size_t batch_size,
690     size_t input_height,
691     size_t input_width,
692     const float* input,
693     float* output,
694     pthreadpool_t threadpool)
695 {
696   if (average_pooling_op->type != xnn_operator_type_average_pooling_nhwc_f32) {
697     xnn_log_error("failed to setup operator: operator type mismatch (expected %s, got %s)",
698       xnn_operator_type_to_string(xnn_operator_type_average_pooling_nhwc_f32),
699       xnn_operator_type_to_string(average_pooling_op->type));
700     return xnn_status_invalid_parameter;
701   }
702 
703   assert(average_pooling_op->ukernel.type == xnn_ukernel_type_average_pooling ||
704          average_pooling_op->ukernel.type == xnn_ukernel_type_pixelwise_average_pooling);
705 
706   const bool is_pixelwise = average_pooling_op->ukernel.type == xnn_ukernel_type_pixelwise_average_pooling;
707   if (is_pixelwise) {
708     const size_t input_size = input_height * input_width;
709     xnn_update_f32_scaleminmax_params(&average_pooling_op->params.f32_scaleminmax, 1.0f / (float) input_size);
710   }
711 
712   return setup_average_pooling2d(
713     average_pooling_op,
714     batch_size, input_height, input_width,
715     input, output,
716     2 /* log2(sizeof(input element)) = log2(sizeof(float)) */,
717     2 /* log2(sizeof(output element)) = log2(sizeof(float)) */,
718     &xnn_params.f32.avgpool,
719     &xnn_params.f32.pavgpool,
720     &xnn_params.f32.gavgpool,
721     is_pixelwise ? (const void*) &average_pooling_op->params.f32_minmax : (const void*) &average_pooling_op->params.f32_scaleminmax,
722     is_pixelwise ? sizeof(average_pooling_op->params.f32_minmax) : sizeof(average_pooling_op->params.f32_scaleminmax),
723     &average_pooling_op->params.f32_scaleminmax,
724     sizeof(average_pooling_op->params.f32_scaleminmax),
725     pthreadpool_get_threads_count(threadpool),
726     is_pixelwise);
727 }
728