Repository: webonnx/wonnx (language: python)

Issue:
AveragePool fails when output width * height is not a multiple of 4
**Describe the bug**

When you feed AveragePool an NxCxWxH tensor where the output WxH (which depends on the kernel size) is not divisible by four, the following error occurs:

````
Shader error:
error: expected ')', found 'u'
   ┌─ wgsl:25:23
   │
25 │         if (gidx < 4.5u) {
   │                       ^ expected ')'
````

The relevant part of the template (templates/pool/aggregate.wgsl):

````wgsl
if (gidx < {{ o_lens[0] / 4 }}u) {
````

I ran into this when implementing `GlobalAveragePool` which is basically `AveragePool` with the kernel size equal to the image size, i.e. simply averaging to a single number per channel (NxCxWxH -> NxCx1x1).

**To Reproduce**

Use AveragePool in a way that sets `o_lens[0]` to something not divisible by 4.

**Expected behavior**

This should just work - the output may of course be a vector of length rounded up to the next multiple of 4, that is automatically chopped of if it's the output vector (or not relevant when this is an input to the next op).



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index d715690f..f1dce73e 100644
--- a/README.md
+++ b/README.md
@@ -197,7 +197,7 @@ fn test_matmul_square_matrix() {
 |<a href="#GatherElements">GatherElements</a>|<a href="Changelog.md#GatherElements-13">13</a>, <a href="Changelog.md#GatherElements-11">11</a>|
 |<a href="#GatherND">GatherND</a>|<a href="Changelog.md#GatherND-13">13</a>, <a href="Changelog.md#GatherND-12">12</a>, <a href="Changelog.md#GatherND-11">11</a>|
 |<a href="#Gemm">Gemm</a>|<a href="Changelog.md#Gemm-13">13</a>, <a href="Changelog.md#Gemm-11">11</a>, <a href="Changelog.md#Gemm-9">9</a>, <a href="Changelog.md#Gemm-7">7</a>, <a href="Changelog.md#Gemm-6">6</a>, <a href="Changelog.md#Gemm-1">1</a>|✅|
-|<a href="#GlobalAveragePool">GlobalAveragePool</a>|<a href="Changelog.md#GlobalAveragePool-1">1</a>|
+|<a href="#GlobalAveragePool">GlobalAveragePool</a>|<a href="Changelog.md#GlobalAveragePool-1">1</a>|✅|
 |<a href="#GlobalLpPool">GlobalLpPool</a>|<a href="Changelog.md#GlobalLpPool-2">2</a>, <a href="Changelog.md#GlobalLpPool-1">1</a>|
 |<a href="#GlobalMaxPool">GlobalMaxPool</a>|<a href="Changelog.md#GlobalMaxPool-1">1</a>|
 |<a href="#Greater">Greater</a>|<a href="Changelog.md#Greater-13">13</a>, <a href="Changelog.md#Greater-9">9</a>, <a href="Changelog.md#Greater-7">7</a>, <a href="Changelog.md#Greater-1">1</a>|✅|
diff --git a/src/compiler.rs b/src/compiler.rs
index cef3ca27..214b5663 100644
--- a/src/compiler.rs
+++ b/src/compiler.rs
@@ -319,13 +319,29 @@ pub fn compile(
                 1,
             )
         }
-        op @ ("MaxPool" | "AveragePool" | "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish") => {
+        op
+        @
+        ("MaxPool" | "AveragePool" | "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish"
+        | "GlobalAveragePool") => {
             // TODO: Conv only support NxCxHxW for the moment.
             debug_assert!(input_shape[0].rank() == 4);
 
+            // GlobalAveragePool is equivalent to AveragePool, with the kernel shape set to the size of the input tensor
+            // See https://github.com/onnx/onnx/blob/main/docs/Operators.md#globalaveragepool
+            // Other attributes are not supported and also not relevant, and are simply ignored
+            let is_global_average_pool = op == "GlobalAveragePool";
+            if is_global_average_pool {
+                // Generate shader code as if this were a regular AveragePool
+                context.insert("op_type", "AveragePool");
+            }
+
             let auto_pad = get_attribute("auto_pad", Some("NOTSET".to_string()), node)?;
             let dilations = get_attribute("dilations", Some(vec![1, 1]), node)?;
-            let kernel_shape = get_attribute::<Vec<i64>>("kernel_shape", None, node)?;
+            let kernel_shape = if is_global_average_pool {
+                vec![input_shape[0].dim(2) as i64, input_shape[0].dim(3) as i64]
+            } else {
+                get_attribute::<Vec<i64>>("kernel_shape", None, node)?
+            };
             let strides = get_attribute("strides", Some(vec![1, 1]), node)?;
             let pads = get_attribute("pads", Some(vec![0, 0, 0, 0]), node)?;
 
@@ -388,7 +404,7 @@ pub fn compile(
 
             // GLSL shader for convolution computation
             match op {
-                "MaxPool" | "AveragePool" => (
+                "MaxPool" | "AveragePool" | "GlobalAveragePool" => (
                     "pool/aggregate.wgsl".to_string(),
                     ceil(output_lengths[0], 1024) as _,
                     1,


Test patch:
diff --git a/tests/globalaveragepool.rs b/tests/globalaveragepool.rs
new file mode 100644
index 00000000..6480d156
--- /dev/null
+++ b/tests/globalaveragepool.rs
@@ -0,0 +1,60 @@
+use approx::assert_ulps_eq;
+use std::collections::HashMap;
+use wonnx::utils::{graph, model, node, tensor};
+
+fn abs_eq_vector(xs: &[f32], ys: &[f32]) {
+    assert_eq!(xs.len(), ys.len());
+    for i in 0..xs.len() {
+        assert_ulps_eq!(xs[i], ys[i], max_ulps = 2);
+    }
+}
+
+#[test]
+fn global_average_pool() {
+    let mut input_data = HashMap::new();
+
+    let batches = 1;
+    let width_height: usize = 2;
+    let channels: usize = 4;
+    // FIXME: we are testing with 4 channels because the AveragePool op doesn't support output tensors with total length non divisible by 4
+    let data: Vec<f32> = (0..(batches * width_height * width_height * channels))
+        .map(|x| x as f32)
+        .collect();
+    let shape = vec![
+        batches as i64,
+        channels as i64,
+        width_height as i64,
+        width_height as i64,
+    ];
+    let output_shape = vec![batches as i64, channels as i64, 1, 1];
+    input_data.insert("X".to_string(), data.as_slice());
+
+    let bn_model = model(graph(
+        vec![tensor("X", &shape)],
+        vec![tensor("Y", &output_shape)],
+        vec![],
+        vec![],
+        vec![node(
+            vec!["X"],
+            vec!["Y"],
+            "gap",
+            "GlobalAveragePool",
+            vec![],
+        )],
+    ));
+
+    // LOGIC
+    let session =
+        pollster::block_on(wonnx::Session::from_model(bn_model)).expect("Session did not create");
+
+    let result = pollster::block_on(session.run(input_data)).unwrap();
+    let out_y = &result["Y"];
+
+    // The GlobalAveragePool op simply averages all pixels in an image (NxCxWxH becomes NxCx1x1). In our test data pixels
+    // range from 0..16, so the test data and output are given as:
+    // Channel 1: [[0,1], [2,3]] => average is 1,5
+    // Channel 2: [[4,5], [6,7]] => average is 5,5
+    // Channel 3: [[8,9], [10, 11]] => average is 9,5
+    // Channel 4: [[12,13], [14, 15]] => average is 13,5
+    abs_eq_vector(out_y.as_slice(), &[1.5, 5.5, 9.5, 13.5]);
+}