Repository: webonnx/wonnx (language: python)

Issue:
ComputeLimitExceeded("X threads", 96800, 65535)
**Describe the bug**
I get the error `ComputeLimitExceeded("X threads", 96800, 65535)` when loading a network that concatenates two tensors.
The inputs are two tensors with shape [1, 64, 440, 440].

**To Reproduce**
Load a network that has a similar input size.
My ONNX sample input wouldn't be helpful here, because it uses a ConvTranspose operation which I am working on adding to wonnx in my fork.

**Expected behavior**
The concatenation should work.

**Screenshots**
I added some output to print the input and output shapes in the compiler:
```
[2023-07-23T19:54:42Z DEBUG wonnx::compiler] [
        Shape {
            dims: [
                1,
                64,
                440,
                440,
            ],
            data_type: F32,
        },
        Shape {
            dims: [
                1,
                64,
                440,
                440,
            ],
            data_type: F32,
        },
    ]
[2023-07-23T19:54:42Z DEBUG wonnx::compiler] [
        Shape {
            dims: [
                1,
                128,
                440,
                440,
            ],
            data_type: F32,
        },
    ]
```

Is this a general limitation and an indicator that the network might be too complex for wonnx or would it make sense to consider bumping the workgroup size or maybe making it dynamic depending on the input size?



Discussion / hints:
WebGPU (and hence `wgpu`) impose limits on the number of workgroups (see [here](https://www.w3.org/TR/webgpu/#dom-supported-limits-maxcomputeworkgroupsizex)). The error occurs when invoking `dispatch` with workgroup counts that exceed these limits.

Many ops in `wonnx` determine the workgroup count dynamically as to stay within these limits. I am not sure if you are doing this in your op implementation but I would suggest looking at other ops code in `compiler.rs`.
I see, thanks for the link, I missed that part in the compiler source code.

But this issue is not caused by my current implementation, it's the already existing `Concat` implementation.
I think it could be possible to distribute the workgroups differently by changing `@workgroup_size(256,1,1)` to `@workgroup_size(16,16,1)` and calculating the workgroup counts accordingly (since the limits are applied per dimension, so getting each dimension to a lower value should be helpful here).

I'm going to give that a go and try to refactor the `Concat`implementation.
:+1: 

Most ops call `workgroup_size` to determine the appropriate workgroup size dynamically. It appears the implementation of `Concat` (actually one of the first ops to be implemented in wonnx) does not do that, so overflowing the limits is expected. Happy to review a PR that fixes this (shouldn't be that difficult!).
Sure, I'll take a look at other operations and try to port the logic to `Concat` :+1: 
I also got my network to run using the approach I described above - just need to actually implement a proper shader for `ConvTranspose`now.

I think I'm also going to add a unit test for `Concat` as it seems like there is no test for that right now. 

Fix patch:
diff --git a/wonnx/src/compiler.rs b/wonnx/src/compiler.rs
index bed2b865..ba757c20 100644
--- a/wonnx/src/compiler.rs
+++ b/wonnx/src/compiler.rs
@@ -2,7 +2,7 @@
 use crate::utils::{
     ceil, AttributeNotFoundError, DataTypeError, MultiType, NodeAttributes, ScalarType, Shape,
 };
-use num::integer::gcd;
+use num::integer::{gcd, Roots};
 use tera::{Context, Tera};
 use thiserror::Error;
 
@@ -792,10 +792,13 @@ pub fn compile(
             }
             context.insert("cum_len", &input_cumulative_len);
 
+            let root = output_lengths[0].sqrt() + 1;
+            let per_dim = ceil(root, 16) + 1;
+
             NodeTemplate {
                 scalar_type: agreed_type(input_shapes, output_shapes)?,
                 template: "matrix/concat.wgsl",
-                threads: (ceil(output_lengths[0], 256) as u32, 1, 1),
+                threads: (per_dim as u32, per_dim as u32, 1),
             }
         }
         op @ ("MaxPool" | "AveragePool" | "Conv" | "ConvRelu" | "ConvLeakyRelu" | "ConvMish"
diff --git a/wonnx/templates/matrix/concat.wgsl b/wonnx/templates/matrix/concat.wgsl
index 76872cf7..4c282723 100644
--- a/wonnx/templates/matrix/concat.wgsl
+++ b/wonnx/templates/matrix/concat.wgsl
@@ -12,19 +12,24 @@ var<storage, read> input_{{ loop.index0 }}: Array;
 @group({{ binding_len  / 4 | int }}) @binding({{ binding_len % 4 }})
 var<storage, read_write> output_0: Array;
 
-@compute @workgroup_size(256, 1, 1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+@compute @workgroup_size(16, 16, 1)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>, @builtin(num_workgroups) num_workgroups: vec3<u32>) {
 	let gidx = global_id.x;
+        let gidy = global_id.y;
+
+        let nx = num_workgroups.x;
+
+        let actual_idx = gidx + gidy * nx;
 	
 	{% for input in i_lens %}
 		{% if loop.first %}
-			if (gidx < {{ i_lens[0] }}u) {
-				output_0.data[gidx] = input_0.data[gidx];
+			if (actual_idx < {{ i_lens[0] }}u) {
+				output_0.data[actual_idx] = input_0.data[actual_idx];
 			}
 
 		{% else %}
-			if ((gidx >= {{ cum_len | nth(n=loop.index0 -1) }}u) && (gidx < {{ cum_len | nth(n=loop.index0)}}u)) {
-				output_0.data[gidx] = input_{{ loop.index0 }}.data[gidx - {{ cum_len | nth(n=loop.index0 -1) }}u];
+			if ((actual_idx >= {{ cum_len | nth(n=loop.index0 -1) }}u) && (actual_idx < {{ cum_len | nth(n=loop.index0)}}u)) {
+				output_0.data[actual_idx] = input_{{ loop.index0 }}.data[actual_idx - {{ cum_len | nth(n=loop.index0 -1) }}u];
 			}
 			
 		{% endif %}


Test patch:
diff --git a/wonnx/tests/concat.rs b/wonnx/tests/concat.rs
new file mode 100644
index 00000000..980ad039
--- /dev/null
+++ b/wonnx/tests/concat.rs
@@ -0,0 +1,80 @@
+use std::{collections::HashMap, convert::TryInto};
+use wonnx::utils::{graph, model, node, tensor};
+mod common;
+
+#[test]
+fn test_concat() {
+    let n: usize = 16;
+
+    let xdata: Vec<f32> = (0..n).map(|x| x as f32).collect();
+    let mut ydata: Vec<f32> = (n..2 * n).map(|x| x as f32).collect();
+    let input_dims = vec![n as i64];
+    let output_dims = vec![(n * 2) as i64];
+
+    let input_data = HashMap::from([
+        ("X".into(), xdata.as_slice().into()),
+        ("Y".into(), ydata.as_slice().into()),
+    ]);
+
+    let model = model(graph(
+        vec![tensor("X", &input_dims), tensor("Y", &input_dims)],
+        vec![tensor("Z", &output_dims)],
+        vec![],
+        vec![],
+        vec![node(vec!["X", "Y"], vec!["Z"], "a", "Concat", vec![])],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
+
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+
+    let mut expected_result = xdata.clone();
+    expected_result.append(&mut ydata);
+
+    common::assert_eq_vector((&result["Z"]).try_into().unwrap(), &expected_result);
+}
+
+#[test]
+fn test_concat4() {
+    let n: usize = 13;
+
+    let xdata: Vec<f32> = (0..n).map(|x| x as f32).collect();
+    let mut ydata: Vec<f32> = (n..2 * n).map(|x| x as f32).collect();
+    let mut zdata: Vec<f32> = (n * 2..3 * n).map(|x| x as f32).collect();
+    let mut wdata: Vec<f32> = (n * 3..4 * n).map(|x| x as f32).collect();
+    let input_dims = vec![n as i64];
+    let output_dims = vec![(n * 4) as i64];
+
+    let input_data = HashMap::from([
+        ("X".into(), xdata.as_slice().into()),
+        ("Y".into(), ydata.as_slice().into()),
+        ("Z".into(), zdata.as_slice().into()),
+        ("W".into(), wdata.as_slice().into()),
+    ]);
+
+    let model = model(graph(
+        vec![
+            tensor("X", &input_dims),
+            tensor("Y", &input_dims),
+            tensor("Z", &input_dims),
+            tensor("W", &input_dims),
+        ],
+        vec![tensor("O", &output_dims)],
+        vec![],
+        vec![],
+        vec![node(vec!["X", "Y", "Z", "W"], vec!["O"], "a", "Concat", vec![])],
+    ));
+
+    let session =
+        pollster::block_on(wonnx::Session::from_model(model)).expect("Session did not create");
+
+    let result = pollster::block_on(session.run(&input_data)).unwrap();
+
+    let mut expected_result = xdata.clone();
+    expected_result.append(&mut ydata);
+    expected_result.append(&mut zdata);
+    expected_result.append(&mut wdata);
+
+    common::assert_eq_vector((&result["O"]).try_into().unwrap(), &expected_result);
+}