Repository: webonnx/wonnx (language: python)

Issue:
Gemm does not appear to work properly when input dim is Nx2
**Describe the bug**

In BertSQuAD, there is this Gemm operation:

<img width="161" alt="image" src="https://user-images.githubusercontent.com/181337/164986528-13fd7953-0f42-4009-8d58-c4175b270e0d.png">

Executing this leads to all zeroes even though the inputs are all non-zero. Looking through the code it seems the shader assumes the second dimension of input B to be at least 4 (it multiplies blocks of 4x4). 

**To Reproduce**

Perform Gemm with an input B of size NxM where M <4, e.g. 768x2 as in my example. Output will be all zeroes.

**Expected behavior**

Output should be non-zero.

**Screenshots**
n/a

**Desktop (please complete the following information):**
 - OS: macOS

Gemm does not appear to work properly when input dim is Nx2
**Describe the bug**

In BertSQuAD, there is this Gemm operation:

<img width="161" alt="image" src="https://user-images.githubusercontent.com/181337/164986528-13fd7953-0f42-4009-8d58-c4175b270e0d.png">

Executing this leads to all zeroes even though the inputs are all non-zero. Looking through the code it seems the shader assumes the second dimension of input B to be at least 4 (it multiplies blocks of 4x4). 

**To Reproduce**

Perform Gemm with an input B of size NxM where M <4, e.g. 768x2 as in my example. Output will be all zeroes.

**Expected behavior**

Output should be non-zero.

**Screenshots**
n/a

**Desktop (please complete the following information):**
 - OS: macOS



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 8fafd530..acdcee70 100644
--- a/README.md
+++ b/README.md
@@ -260,7 +260,7 @@ fn test_matmul_square_matrix() {
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gather">Gather</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gather-1">1</a>|✅ (axis=0)|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherElements">GatherElements</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GatherElements-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GatherElements-11">11</a>|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GatherND">GatherND</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GatherND-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GatherND-12">12</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GatherND-11">11</a>|
-|<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gemm">Gemm</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-9">9</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-7">7</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-6">6</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-1">1</a>|✅|
+|<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#Gemm">Gemm</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-13">13</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-11">11</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-9">9</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-7">7</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-6">6</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#Gemm-1">1</a>|✅*|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GlobalAveragePool">GlobalAveragePool</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GlobalAveragePool-1">1</a>|✅|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GlobalLpPool">GlobalLpPool</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GlobalLpPool-2">2</a>, <a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GlobalLpPool-1">1</a>|
 |<a href="https://github.com/onnx/onnx/blob/main/docs/Operators.md#GlobalMaxPool">GlobalMaxPool</a>|<a href="https://github.com/onnx/onnx/blob/main/docs/Changelog.md#GlobalMaxPool-1">1</a>|
@@ -395,3 +395,6 @@ fn test_matmul_square_matrix() {
 
 * Internally 64-bit integers are not supported (the reason is they are not supported in the current version of WGSL); 
   inputs and initializers with 64-bit scalars are converted to 32-bit values (possibly overflowing).
+
+* For `MatMul` and `Gemm`, the matrix dimensions must be divisible by 2, or the output matrix must be of size (1, N). Matrix 
+  multiplication only supports floats, not integers (this is a WebGPU/WGSL limitation).
diff --git a/wonnx-cli/src/gpu.rs b/wonnx-cli/src/gpu.rs
index e53810b2..44a4dc86 100644
--- a/wonnx-cli/src/gpu.rs
+++ b/wonnx-cli/src/gpu.rs
@@ -49,7 +49,6 @@ impl Inferer for GPUInferer {
             None => result.values().next().unwrap().clone(),
         };
 
-        log::info!("gpu result: {:?}", result);
         Ok(result)
     }
 }
diff --git a/wonnx/src/compiler.rs b/wonnx/src/compiler.rs
index aabfca7b..27cd99f2 100644
--- a/wonnx/src/compiler.rs
+++ b/wonnx/src/compiler.rs
@@ -1,6 +1,7 @@
 use crate::utils::{
     ceil, get_attribute, AttributeNotFoundError, DataTypeError, MultiType, ScalarType, Shape,
 };
+use num::integer::gcd;
 use tera::{Context, Tera};
 use thiserror::Error;
 
@@ -903,55 +904,219 @@ pub fn compile(
             }
         }
         op @ ("Gemm" | "MatMul") => {
-            let alpha = get_attribute("alpha", Some(1.0), node)?;
-            let beta = get_attribute("beta", Some(1.0), node)?;
+            // Generic matrix multiplication; outputs an M*N matrix from inputs A (size M*K) and B (size K*N)
+
+            // MatMul behaves "like numpy.matmul" (https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.matmul.html)
+            // If both arguments are 2-D they are multiplied like conventional matrices. If they are not, special rules are
+            // applied, that may (among others) lead to 'stacking' (basically multiple matrix multiplications are performed
+            // in parallel). Below, these rules are applied. These rewrite the shapes for input and output and in some cases
+            // set a 'stack count' and 'stack stride'.
+            let mut stack_count = 1;
+            let mut input_left_shape = input_shapes[0].clone();
+            let mut input_right_shape = input_shapes[1].clone();
+            let mut output_shape = output_shapes[0].clone();
+            let mut stack_left_stride: u64 = 0;
+            let mut stack_right_stride: u64 = 0;
+            let mut stack_output_stride: u64 = 0;
+
+            if op == "MatMul" {
+                // - If either argument is N-D, N > 2, it is treated as a stack of matrices residing in the last two indexes
+                //   and broadcast accordingly.
+                if input_left_shape.rank() > 2 || input_right_shape.rank() > 2 {
+                    if input_left_shape.rank() != input_right_shape.rank()
+                        || output_shape.rank() != input_left_shape.rank()
+                        || input_left_shape.dims[0..(input_left_shape.dims.len() - 2)]
+                            != input_right_shape.dims[0..(input_right_shape.dims.len() - 2)]
+                        || input_left_shape.dims[0..(input_left_shape.dims.len() - 2)]
+                            != output_shape.dims[0..(output_shape.dims.len() - 2)]
+                    {
+                        return Err(CompileError::UnimplementedVariant {
+                            variant: format!("broadcasting for two stacks of matrixes (left side has shape {}, right side has shape {})", input_left_shape, input_right_shape),
+                            op: op.to_string(),
+                        });
+                    }
+
+                    let stack_dims = input_left_shape.dims.len() - 2;
+                    stack_count = input_left_shape.dims[0..stack_dims].iter().product();
+                    input_left_shape.dims.drain(0..stack_dims);
+                    input_right_shape.dims.drain(0..stack_dims);
+                    output_shape.dims.drain(0..stack_dims);
+                    stack_left_stride = input_left_shape.dims.iter().product();
+                    stack_right_stride = input_right_shape.dims.iter().product();
+                    stack_output_stride = output_shape.dims.iter().product();
+
+                    log::debug!(
+                        "MatMul stacking: left {} right {} stack_dims={} stack_count={} stack_left_stride={} stack_right_stride={} stack_output_stride={}",
+                        input_left_shape,
+                        input_right_shape,
+                        stack_dims,
+                        stack_count,
+                        stack_left_stride,
+                        stack_right_stride,
+                        stack_output_stride
+                    );
+                }
+
+                // If the first argument is 1-D, it is promoted to a matrix by prepending a 1 to its dimensions. After
+                // matrix multiplication the prepended 1 is removed.
+                if input_left_shape.rank() == 1 {
+                    input_left_shape.dims.insert(0, 1);
+                }
+
+                // If the second argument is 1-D, it is promoted to a matrix by appending a 1 to its dimensions. After
+                // matrix multiplication the appended 1 is removed.
+                if input_right_shape.rank() == 1 {
+                    input_left_shape.dims.push(1);
+                }
+            }
+
+            context.insert("stack_left_stride", &stack_left_stride);
+            context.insert("stack_right_stride", &stack_right_stride);
+            context.insert("stack_output_stride", &stack_output_stride);
+            context.insert("left_shape", &input_left_shape.dims);
+            context.insert("right_shape", &input_right_shape.dims);
+            context.insert("output_shape", &output_shape.dims);
+
+            // Check dimensions
+            let dim_m = output_shape.dim(0);
+            let dim_n = output_shape.dim(1);
+            let dim_k = input_left_shape.dim(1);
+            if dim_m != input_left_shape.dim(0) {
+                return Err(CompileError::InvalidInputShape {
+                    input_index: 0,
+                    input_shape: input_left_shape.clone(),
+                });
+            }
+
+            if dim_n != input_right_shape.dim(1) || dim_k != input_right_shape.dim(0) {
+                return Err(CompileError::InvalidInputShape {
+                    input_index: 1,
+                    input_shape: input_right_shape.clone(),
+                });
+            }
 
-            // Whether A resp. B should be transposed, or C should be broadcast (default: 0 = false)
             if op == "Gemm" {
+                // Check if A resp. B should be transposed, or C should be broadcast (default: 0 = false)
                 let transpose_a = get_attribute("transA", Some(0), node)?;
                 let transpose_b = get_attribute("transB", Some(0), node)?;
                 let broadcast = get_attribute("broadcast", Some(0), node)?;
 
                 if transpose_a != 0 || transpose_b != 0 || broadcast != 0 {
                     return Err(CompileError::UnimplementedVariant {
-                        variant: "Gemm with transA/transB/broadcast not equal to zero".to_string(),
+                        variant: "with transA/transB/broadcast not equal to zero".to_string(),
                         op: op.to_string(),
                     });
                 }
-            }
 
-            context.insert("alpha", &alpha);
-            context.insert("beta", &beta);
+                // If there is a bias input, it should be "unidirectionally broadcastable to M*N". Currently we only support a bias ofM*N though
+                if input_shapes.len() > 2 {
+                    let mut bias_shape = input_shapes[2].clone();
 
-            // Whether A resp. B should be transposed, or C should be broadcast (default: 0 = false)
-            if op == "Gemm" {
-                let transpose_a = get_attribute("transA", Some(0), node)?;
-                let transpose_b = get_attribute("transB", Some(0), node)?;
-                let broadcast = get_attribute("broadcast", Some(0), node)?;
+                    // A shape of higher rank than 2 can never be broadcasted
+                    if bias_shape.rank() > 2 || bias_shape.rank() == 0 {
+                        return Err(CompileError::InvalidInputShape {
+                            input_index: 2,
+                            input_shape: bias_shape.clone(),
+                        });
+                    }
 
-                if transpose_a != 0 || transpose_b != 0 || broadcast != 0 {
+                    // For bias shape of rank 1, prepend a 1
+                    if bias_shape.rank() == 1 {
+                        bias_shape.dims.insert(0, 1);
+                    }
+
+                    // First dimension of bias must be either 1 or M
+                    if bias_shape.dim(0) != dim_m && bias_shape.dim(0) != 1 {
+                        return Err(CompileError::InvalidInputShape {
+                            input_index: 2,
+                            input_shape: bias_shape.clone(),
+                        });
+                    }
+
+                    // Second dimension of bias must be either 1 or N
+                    if bias_shape.dim(1) != dim_n && bias_shape.dim(1) != 1 {
+                        return Err(CompileError::InvalidInputShape {
+                            input_index: 2,
+                            input_shape: bias_shape.clone(),
+                        });
+                    }
+
+                    context.insert("bias_shape", &bias_shape.dims);
+                    context.insert("bias_broadcast_rows", &(bias_shape.dim(0) == 1));
+                    context.insert("bias_broadcast_columns", &(bias_shape.dim(1) == 1));
+                }
+            }
+
+            // Due to a limitation in WGSL, we currently only support floating-point matrix multiplication
+            // See https://github.com/gfx-rs/naga/issues/1896
+            let scalar_type = agreed_type(input_shapes, output_shapes)?;
+            match scalar_type {
+                ScalarType::I32 | ScalarType::I64 => {
                     return Err(CompileError::UnimplementedVariant {
-                        variant: "Gemm with transA/transB/broadcast not equal to zero".to_string(),
+                        variant: "with integers".to_string(),
                         op: op.to_string(),
-                    });
+                    })
                 }
+                ScalarType::F32 => (),
             }
 
-            if input_shapes[0].dim(0) == 1 {
+            // Obtain alpha and beta coefficients
+            let alpha = get_attribute("alpha", Some(1.0), node)?;
+            let beta = get_attribute("beta", Some(1.0), node)?;
+            context.insert("alpha", &alpha);
+            context.insert("beta", &beta);
+
+            // Determine and set thread count/workgroup size when stacking (shader y dimension)
+            let (y_threads, workgroup_size_y) = workgroup_size(
+                stack_count,
+                MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                MAX_WORKGROUP_SIZE_Y,
+            )?;
+            context.insert("workgroup_size_y", &workgroup_size_y);
+
+            if dim_m == 1 {
+                let n_elements = output_shapes[0].dim(1);
+                let (x_threads, workgroup_size_x) = workgroup_size(
+                    n_elements,
+                    MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                    MAX_WORKGROUP_SIZE_X,
+                )?;
+
+                context.insert("workgroup_size_x", &workgroup_size_x);
                 NodeTemplate {
                     scalar_type: agreed_type(input_shapes, output_shapes)?,
                     template: "matrix/gemm_1.wgsl",
-                    threads: (output_shapes[0].dim(1) as _, 1, 1),
+                    threads: (x_threads as _, y_threads, 1),
                 }
             } else {
+                // Matrix multiplication is performed (by the gemm.wgsl shader) in blocks of 4x4, except when the output matrix
+                // or any of the inputs has a dimension smaller than 4, in which case we can do 3x3 or 2x2
+                let kernel_size = gcd(dim_m, gcd(dim_k, dim_n)).min(4).max(1);
+                if kernel_size == 1 {
+                    return Err(CompileError::UnimplementedVariant {
+                        variant: String::from(
+                            "usage with input matrixes whose dimensions are not divisible by 2",
+                        ),
+                        op: op.to_string(),
+                    });
+                }
+
+                let n_blocks = ceil(dim_m * dim_n, kernel_size * kernel_size);
+                let (x_threads, workgroup_size_x) = workgroup_size(
+                    n_blocks,
+                    MAX_COMPUTE_WORKGROUPS_PER_DIMENSION,
+                    MAX_WORKGROUP_SIZE_X,
+                )?;
+
+                context.insert("m_chunks", &(dim_m / kernel_size).max(1));
+                context.insert("n_chunks", &(dim_n / kernel_size).max(1));
+                context.insert("k_chunks", &(dim_k / kernel_size).max(1));
+                context.insert("kernel_size", &kernel_size);
+                context.insert("workgroup_size_x", &workgroup_size_x);
                 NodeTemplate {
-                    scalar_type: agreed_type(input_shapes, output_shapes)?,
+                    scalar_type,
                     template: "matrix/gemm.wgsl",
-                    threads: (
-                        (input_shapes[0].dim(0) * input_shapes[1].dim(1) / 16) as _,
-                        1,
-                        1,
-                    ),
+                    threads: (x_threads as _, y_threads, 1),
                 }
             }
         }
diff --git a/wonnx/templates/matrix/gemm.wgsl b/wonnx/templates/matrix/gemm.wgsl
index 27901360..6bce19b2 100644
--- a/wonnx/templates/matrix/gemm.wgsl
+++ b/wonnx/templates/matrix/gemm.wgsl
@@ -1,71 +1,115 @@
-{%- include "structs.wgsl" -%}
+type Scalar = {{ scalar_type }};
+type GemmVec = vec{{ kernel_size }}<{{ scalar_type }}>;
+type GemmMat = mat{{ kernel_size }}x{{ kernel_size }}<{{ scalar_type }}>;
+
+struct GemmArrayVector {
+	data: array<GemmVec>;
+};
 
 [[group(0), binding(0)]]
-var<storage, read> input_0: ArrayVector;
+var<storage, read> input_left: GemmArrayVector;
 
 [[group(0), binding(1)]]
-var<storage, read> input_1: ArrayVector;
+var<storage, read> input_right: GemmArrayVector;
 
-{%- if i_lens | length == 3 -%} // Bias
-[[group(0), binding(2)]]
-var<storage, read> input_2: ArrayVector;
+{% if i_lens | length == 3 %} // Bias
+	[[group(0), binding(2)]]
+	var<storage, read> input_bias: GemmArrayVector;
 
-[[group(0), binding(3)]]
-var<storage, write> output_0: ArrayVector;
-{%- else -%}
-[[group(0), binding(2)]]
-var<storage, write> output_0: ArrayVector;
-{%- endif -%}  
+	[[group(0), binding(3)]]
+	var<storage, write> output_0: GemmArrayVector;
+{% else %}
+	[[group(0), binding(2)]]
+	var<storage, write> output_0: GemmArrayVector;
+{% endif %}
 
-[[stage(compute), workgroup_size(1)]]
+[[stage(compute), workgroup_size({{ workgroup_size_x }}, {{ workgroup_size_y }})]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
-	let y = global_id.x % {{ i_shape[1][1] / 4 | int }}u;
-	let x = global_id.x / {{ i_shape[1][1] / 4 | int }}u;
-	let index = x * {{ i_shape[1][1] }}u + y;
+	let y = global_id.x % {{ n_chunks }}u;
+	let x = global_id.x / {{ n_chunks }}u;
+
+	{# Calculate stacking offsets #}
+	let stack_index = global_id.y;
+	let left_offset = stack_index * {{ stack_left_stride / kernel_size }}u;
+	let right_offset = stack_index * {{ stack_right_stride / kernel_size }}u;
+	let output_offset = stack_index * {{ stack_output_stride / kernel_size }}u;
+
+	let index = output_offset + (x * {{ right_shape[1] }}u) + y;
+
+	{# Create zero vector and matrix of the correct size for later use #}
+	let zero_vec = GemmVec(
+		{% for i in range(end = kernel_size) %}
+			Scalar(0) {%-if not loop.last -%},{%- endif -%}
+		{% endfor %}
+	);
 
-	let zero = Vec4(Scalar(0), Scalar(0), Scalar(0), Scalar(0));
-	var tmpsum = Mat4x4(zero, zero, zero, zero);
-	var product = Mat4x4(zero, zero, zero, zero);
+	let zero_matrix = GemmMat(
+		{% for i in range(end = kernel_size) %}
+			zero_vec {%-if not loop.last -%},{%- endif -%}
+		{% endfor %}
+	);
 
-	for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
-		let index_left = x * {{ i_shape[0][1] }}u + k; 
-		let index_right = k * {{ i_shape[0][1] }}u + y; 
+	var tmpsum = zero_matrix;
+	var product = zero_matrix;
 
-		let mat_left = Mat4x4(
-			input_0.data[index_left], 
-			input_0.data[index_left + {{ i_shape[0][1] / 4 | int }}u],
-			input_0.data[index_left + {{ 2 * i_shape[0][1] / 4 | int }}u],
-			input_0.data[index_left + {{ 3 * i_shape[0][1] / 4 | int }}u],
+	for(var k: u32 = 0u; k < {{ k_chunks }}u; k = k + 1u) {
+		let index_left = left_offset + (x * {{ left_shape[1] }}u) + k;
+		let index_right = right_offset + (k * {{ right_shape[1] }}u) + y;
+
+		let mat_left = GemmMat(
+			{% for i in range(end = kernel_size) %}
+				input_left.data[index_left + {{ i * k_chunks }}u] {%-if not loop.last -%},{%- endif -%}
+			{% endfor %}
 		);
 		
-		let mat_right = Mat4x4(
-			input_1.data[index_right], 
-			input_1.data[index_right + {{ i_shape[1][1] / 4 | int }}u],
-			input_1.data[index_right + {{ 2 * i_shape[1][1] / 4 | int }}u],
-			input_1.data[index_right + {{ 3 * i_shape[1][1] / 4 | int }}u],
+		let mat_right = GemmMat(
+			{% for i in range(end = kernel_size) %}
+				input_right.data[index_right + ({{ i * n_chunks }}u)] {%-if not loop.last -%},{%- endif -%}
+			{% endfor %}
 		);
 	
 		product = mat_right * mat_left;
 	
-		for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
+		for(var index_mat: u32 = 0u; index_mat < {{ kernel_size }}u; index_mat = index_mat + 1u) {
 			tmpsum[index_mat] = tmpsum[index_mat] + product[index_mat];
 		}
 	}
 	
-	{%- if i_lens | length == 3 -%}
-		let bias_row = input_2.data[x]; 
-		var bias = transpose(Mat4x4(bias_row, bias_row, bias_row, bias_row));
-		for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
-			output_0.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = 
+	{% if i_lens | length == 3 %}
+		let bias_index =
+			{% if not bias_broadcast_rows %} (x * {{ bias_shape[1] }}u) + {% endif %} 
+			{% if not bias_broadcast_columns %} y {% else  %} 0u {% endif %};
+
+		for(var index_mat: u32 = 0u; index_mat < {{ kernel_size }}u; index_mat = index_mat + 1u) {
+			{% if bias_broadcast_columns and bias_broadcast_rows %}
+				{# Bias is just a single number, broadcast over both dimensions #}
+				let bias_number = input_bias.data[bias_index][0];
+				let bias = GemmVec(
+					{% for k in range(end = kernel_size) %}
+						bias_number {%-if not loop.last -%},{%- endif -%}
+					{% endfor %}
+				);
+			{% elif bias_broadcast_columns %}
+				let bias_number = input_bias.data[bias_index][index_mat];
+				let bias = GemmVec(
+					{% for k in range(end = kernel_size) %}
+						bias_number {%-if not loop.last -%},{%- endif -%}
+					{% endfor %}
+				);
+			{% else  %}
+				{# Bias is a matrix, no broadcasting, one bias scalar for each output element #}
+				let bias = input_bias.data[bias_index {% if not bias_broadcast_rows %} + (index_mat * {{ n_chunks }}u) {% endif %}];
+			{% endif %}
+
+			output_0.data[index + (index_mat * {{ n_chunks }}u)] = 
 				{%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%} 
 				tmpsum[index_mat] + 
 				{%- if beta != 1 -%} {{ beta | float }} * {%- endif -%} 
-				bias[index_mat]
-			;
+				bias;
 		}
-	{%- else -%}
-		for(var index_mat: u32 = 0u; index_mat < 4u; index_mat = index_mat + 1u) {
-			output_0.data[index + index_mat * {{ i_shape[1][1] / 4 | int }}u] = {%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%}tmpsum[index_mat];
+	{% else %}
+		for(var index_mat: u32 = 0u; index_mat < {{ kernel_size }}u; index_mat = index_mat + 1u) {
+			output_0.data[index + (index_mat * {{ n_chunks }}u)] = {% if alpha != 1 -%} {{ alpha | float }} * {%- endif -%} tmpsum[index_mat];
 		}
-	{%- endif -%}
+	{% endif %}
 }
\ No newline at end of file
diff --git a/wonnx/templates/matrix/gemm_1.wgsl b/wonnx/templates/matrix/gemm_1.wgsl
index 8b2b1521..87ca41de 100644
--- a/wonnx/templates/matrix/gemm_1.wgsl
+++ b/wonnx/templates/matrix/gemm_1.wgsl
@@ -7,43 +7,50 @@ var<storage, read> input_0: ArrayVector;
 var<storage, read> input_1: Array;
 
 {%- if i_lens | length == 3 -%} // Bias
-[[group(0), binding(2)]]
-var<storage, read> input_2: Array;
 
-[[group(0), binding(3)]]
-var<storage, write> output_0: Array;
+	[[group(0), binding(2)]]
+	var<storage, read> input_2: Array;
+
+	[[group(0), binding(3)]]
+	var<storage, write> output_0: Array;
 
 {%- else -%}
-[[group(0), binding(2)]]
-var<storage, write> output_0: Array;
+
+	[[group(0), binding(2)]]
+	var<storage, write> output_0: Array;
 
 {%- endif -%}  
 
-[[stage(compute), workgroup_size(1)]]
+[[stage(compute), workgroup_size(1, {{ workgroup_size_y }})]]
 fn main([[builtin(global_invocation_id)]] global_id: vec3<u32>) {
 	let gidx = global_id.x;
 
+	{# Calculate stacking offsets #}
+	let left_offset = global_id.y * {{ stack_left_stride / 4 }}u;
+	let right_offset = global_id.y * {{ stack_right_stride }}u;
+	let output_offset = global_id.y * {{ stack_output_stride }}u;
+
 	var tmpsum = Scalar(0);
 	var product = Scalar(0);
 
-	for(var k: u32 = 0u; k < {{ i_shape[0][1] / 4 | int }}u; k = k + 1u) {
-		let index_left = k; 
-		let index_right = k * {{ i_shape[1][1] * 4 }}u + gidx; 
+	for(var k: u32 = 0u; k < {{ left_shape[1] / 4 | int }}u; k = k + 1u) {
+		let index_left = left_offset + k; 
+		let index_right = right_offset + (k * {{ right_shape[1] * 4 }}u) + gidx; 
 
 		let vec_left = input_0.data[index_left];
 
 		let vec_right = Vec4(
 			input_1.data[index_right], 
-			input_1.data[index_right + {{ i_shape[1][1] }}u],
-			input_1.data[index_right + {{ 2 * i_shape[1][1] }}u],
-			input_1.data[index_right + {{ 3 * i_shape[1][1] }}u],
+			input_1.data[index_right + {{ right_shape[1] }}u],
+			input_1.data[index_right + {{ 2 * right_shape[1] }}u],
+			input_1.data[index_right + {{ 3 * right_shape[1] }}u],
 		);
 	
 		product = dot(vec_left, vec_right);
 		tmpsum = tmpsum + product;
 	}
 	
-	output_0.data[gidx] = 
+	output_0.data[output_offset + gidx] = 
 		{%- if alpha != 1 -%} {{ alpha | float }} * {%- endif -%} 
 		tmpsum
 		{%- if i_lens | length == 3 -%}


Test patch:
diff --git a/wonnx/tests/matrix.rs b/wonnx/tests/matrix.rs
index 6fd0f7cc..80a37bb2 100644
--- a/wonnx/tests/matrix.rs
+++ b/wonnx/tests/matrix.rs
@@ -4,6 +4,7 @@ mod common;
 
 #[test]
 fn test_matmul_square_matrix() {
+    let _ = env_logger::builder().is_test(true).try_init();
     let n = 16;
     let mut input_data = HashMap::new();
 
@@ -79,6 +80,7 @@ fn test_two_transposes() {
 
 #[test]
 fn test_split() {
+    let _ = env_logger::builder().is_test(true).try_init();
     let mut input_data = HashMap::new();
     let data = (1..=2 * 6).map(|x| x as f32).collect::<Vec<f32>>();
     input_data.insert("X".to_string(), data.as_slice().into());
@@ -144,7 +146,7 @@ fn test_pad_example() {
 #[test]
 fn test_pad_complex() {
     let mut input_data = HashMap::new();
-    let data = (1..=1 * 3 * 2).map(|x| x as f32).collect::<Vec<f32>>();
+    let data = (1..=3 * 2).map(|x| x as f32).collect::<Vec<f32>>();
     input_data.insert("X".to_string(), data.as_slice().into());
 
     let kv = 0.5;
@@ -179,24 +181,16 @@ fn test_pad_complex() {
         pollster::block_on(wonnx::Session::from_model(model)).expect("session did not create");
     let result = pollster::block_on(session.run(&input_data)).unwrap();
 
-    /*
-    [
-    [[1,2],
-     [3,4],
-     [5,6]],
-    ]
-
-    ->
-
+    /* [[[1,2], [3,4], [5,6]]] ->
     [
-    [[0,0,0,0,0],
-     [0,0,1,2,0],
-     [0,0,3,4,0],
-     [0,0,5,6,0]],
-    [[0,0,0,0,0],
-     [0,0,0,0,0],
-     [0,0,0,0,0],
-     [0,0,0,0,0]]
+        [[0,0,0,0,0],
+        [0,0,1,2,0],
+        [0,0,3,4,0],
+        [0,0,5,6,0]],
+        [[0,0,0,0,0],
+        [0,0,0,0,0],
+        [0,0,0,0,0],
+        [0,0,0,0,0]]
     ]
     */
 
@@ -213,6 +207,7 @@ fn test_pad_complex() {
         kv, kv, kv, kv, kv,
     ];
     let actual: &[_] = (&result["Y"]).try_into().unwrap();
+
     // No arithmetic is done, so `assert_eq!` can be used.
     assert_eq!(actual, &test_y);
 }
@@ -273,3 +268,296 @@ fn test_resize() {
     //];
     //assert_eq!(result["Y"], test_y);
 }
+
+// Multiply a 2x2 matrix with an identity matrix of size 2x2.
+#[test]
+fn test_matmul_square_matrix_small() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let n = 2;
+    let mut input_data = HashMap::new();
+
+    let data_a = ndarray::Array2::eye(n);
+    let mut data_b = ndarray::Array2::<f32>::zeros((n, n));
+    data_b[[0, 0]] = 0.0;
+    data_b[[0, 1]] = 0.5;
+    data_b[[1, 0]] = 1.0;
+    data_b[[1, 1]] = 1.5;
+
+    let sum = data_a.dot(&data_b);
+
+    input_data.insert("A".to_string(), data_a.as_slice().unwrap().into());
+    input_data.insert("B".to_string(), data_b.as_slice().unwrap().into());
+
+    let n = n as i64;
+    let model = model(graph(
+        vec![tensor("A", &[n, n]), tensor("B", &[n, n])],
+        vec![tensor("C", &[n, n])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B"], vec!["C"], "MatMul", "MatMul", 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();
+
+    common::assert_eq_vector((&result["C"]).try_into().unwrap(), sum.as_slice().unwrap());
+}
+
+// Multiply a 4x4 matrix with a ones matrix of size 4x2.
+// a = np.arange(0,16).reshape((4,4))
+// b = np.arange(-8,0).reshape((4,2))
+// c = np.matmul(a, b)
+//array([[ -20,  -14], [-100,  -78], [-180, -142], [-260, -206]])
+#[test]
+fn test_matmul_nonsquare_matrix_small() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..16).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..8).map(|x| -8.0 + (x as f32)).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+
+    let model = model(graph(
+        vec![tensor("A", &[4, 4]), tensor("B", &[4, 2])],
+        vec![tensor("C", &[4, 2])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B"], vec!["C"], "MatMul", "MatMul", 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 out = &[-20., -14., -100., -78., -180., -142., -260., -206.];
+    common::assert_eq_vector((&result["C"]).try_into().unwrap(), out);
+}
+
+// Multiply two stacks of 2x2 matrixes
+// a = np.arange(0,16).reshape((4,2,2))
+// b = np.arange(16,32).reshape((4,2,2))
+// c = np.matmul(a, b)
+#[test]
+fn test_matmul_stacks() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..16).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (16..32).map(|x| x as f32).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+
+    let model = model(graph(
+        vec![tensor("A", &[4, 2, 2]), tensor("B", &[4, 2, 2])],
+        vec![tensor("C", &[4, 2, 2])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B"], vec!["C"], "MatMul", "MatMul", 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 out: Vec<f32> = vec![
+        18, 19, 86, 91, 190, 199, 274, 287, 426, 443, 526, 547, 726, 751, 842, 871,
+    ]
+    .iter()
+    .map(|x| *x as f32)
+    .collect();
+    common::assert_eq_vector((&result["C"]).try_into().unwrap(), &out);
+}
+
+// Multiply a 1x4 matrix with a ones matrix of size 4x2.
+// a = np.arange(0,4).reshape((1,4))
+// b = np.arange(-8,0).reshape((4,2))
+// c = np.matmul(a, b)
+//array([[ -20,  -14]])
+#[test]
+fn test_matmul_1d() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..4).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..8).map(|x| -8.0 + (x as f32)).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+
+    let model = model(graph(
+        vec![tensor("A", &[1, 4]), tensor("B", &[4, 2])],
+        vec![tensor("C", &[1, 2])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B"], vec!["C"], "MatMul", "MatMul", 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 out = &[-20., -14.];
+    common::assert_eq_vector((&result["C"]).try_into().unwrap(), out);
+}
+
+// Test Gemm with matrix bias
+// a = np.arange(24).reshape((4,6))
+// b = np.arange(24).reshape((6,4))
+// c = np.arange(16).reshape((4,4))
+// d = np.dot(a,b) + c
+// d = array([[ 220,  236,  252,  268], [ 584,  636,  688,  740], [ 948, 1036, 1124, 1212], [1312, 1436, 1560, 1684]])
+#[test]
+fn test_gemm_matrix_bias() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let c_data: Vec<f32> = (0..16).map(|x| x as f32).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+    input_data.insert("C".to_string(), c_data.as_slice().into());
+
+    let model = model(graph(
+        vec![
+            tensor("A", &[4, 6]),
+            tensor("B", &[6, 4]),
+            tensor("C", &[4, 4]),
+        ],
+        vec![tensor("D", &[4, 4])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B", "C"], vec!["D"], "Gemm", "Gemm", 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 out = &[
+        220., 236., 252., 268., 584., 636., 688., 740., 948., 1036., 1124., 1212., 1312., 1436.,
+        1560., 1684.,
+    ];
+    common::assert_eq_vector((&result["D"]).try_into().unwrap(), out);
+}
+
+// Test Gemm with broadcasting bias
+// a = np.arange(24).reshape((4,6))
+// b = np.arange(24).reshape((6,4))
+// c = np.arange(4).reshape((1,4))
+// d = np.dot(a,b) + c
+// d = array([[ 220,  236,  252,  268], [ 580,  632,  684,  736], [ 940, 1028, 1116, 1204], [1300, 1424, 1548, 1672]])
+#[test]
+fn test_gemm_broadcasting_bias() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let c_data: Vec<f32> = (0..4).map(|x| x as f32).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+    input_data.insert("C".to_string(), c_data.as_slice().into());
+
+    let model = model(graph(
+        vec![
+            tensor("A", &[4, 6]),
+            tensor("B", &[6, 4]),
+            tensor("C", &[1, 4]),
+        ],
+        vec![tensor("D", &[4, 4])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B", "C"], vec!["D"], "Gemm", "Gemm", 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 out = &[
+        220., 236., 252., 268., 580., 632., 684., 736., 940., 1028., 1116., 1204., 1300., 1424.,
+        1548., 1672.,
+    ];
+    common::assert_eq_vector((&result["D"]).try_into().unwrap(), out);
+}
+
+// Test Gemm with broadcasting bias (same as above but bias is shaped (4,1) instead of (1,4))
+// a = np.arange(24).reshape((4,6))
+// b = np.arange(24).reshape((6,4))
+// c = np.arange(4).reshape((1,4))
+// d = np.dot(a,b) + c
+// d = array([[ 220,  235,  250,  265], [ 581,  632,  683,  734], [ 942, 1029, 1116, 1203], [1303, 1426, 1549, 1672]])
+#[test]
+fn test_gemm_broadcasting_second_bias() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let c_data: Vec<f32> = (0..4).map(|x| x as f32).collect();
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+    input_data.insert("C".to_string(), c_data.as_slice().into());
+
+    let model = model(graph(
+        vec![
+            tensor("A", &[4, 6]),
+            tensor("B", &[6, 4]),
+            tensor("C", &[4, 1]),
+        ],
+        vec![tensor("D", &[4, 4])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B", "C"], vec!["D"], "Gemm", "Gemm", 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 out = &[
+        220., 235., 250., 265., 581., 632., 683., 734., 942., 1029., 1116., 1203., 1303., 1426.,
+        1549., 1672.,
+    ];
+    common::assert_eq_vector((&result["D"]).try_into().unwrap(), out);
+}
+
+// Test Gemm with broadcasting bias
+// a = np.arange(24).reshape((4,6))
+// b = np.arange(24).reshape((6,4))
+// c = np.array(-1000)
+// d = np.dot(a,b) + c
+// d = array([[-780, -765, -750, -735], [-420, -369, -318, -267], [ -60,   27,  114,  201], [ 300,  423,  546,  669]])
+#[test]
+fn test_gemm_scalar_bias() {
+    let _ = env_logger::builder().is_test(true).try_init();
+    let a_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let b_data: Vec<f32> = (0..24).map(|x| x as f32).collect();
+    let c_data: Vec<f32> = vec![-1000.0];
+
+    let mut input_data = HashMap::new();
+    input_data.insert("A".to_string(), a_data.as_slice().into());
+    input_data.insert("B".to_string(), b_data.as_slice().into());
+    input_data.insert("C".to_string(), c_data.as_slice().into());
+
+    let model = model(graph(
+        vec![
+            tensor("A", &[4, 6]),
+            tensor("B", &[6, 4]),
+            tensor("C", &[1]),
+        ],
+        vec![tensor("D", &[4, 4])],
+        vec![],
+        vec![],
+        vec![node(vec!["A", "B", "C"], vec!["D"], "Gemm", "Gemm", 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 out = &[
+        -780., -765., -750., -735., -420., -369., -318., -267., -60., 27., 114., 201., 300., 423.,
+        546., 669.,
+    ];
+    common::assert_eq_vector((&result["D"]).try_into().unwrap(), out);
+}