1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
//! # Gemm
//!
//! General Matrix Multiplication: Y = alpha * A' * B' + beta * C
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Gemm.html>
//!
//! ## Description
//! Computes Y = alpha * A' * B' + beta * C, where:
//! - Input tensor A has shape (M, K) or (K, M) if transposed
//! - Input tensor B has shape (K, N) or (N, K) if transposed
//! - Input tensor C is broadcastable to shape (M, N)
//! - Output tensor Y has shape (M, N)
//!
//! ## Type Constraints
//! - T: tensor(float), tensor(double), tensor(float16), tensor(bfloat16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)
//!
//! ## Opset Versions
//! - **Opset 1**: Initial version with basic GEMM operation.
//! - **Opset 6**: Clarified broadcasting behavior for input C.
//! - **Opset 7**: Added support for additional data types (int32, int64, uint32, uint64).
//! - **Opset 11**: Changed attribute types to more specific types; clarified unidirectional broadcasting for C.
//! - **Opset 13**: Added bfloat16 and float16 support; updated type constraints.
//!
//! **Implementation Note**: This implementation validates opset 11+.
//!
//! ## Notes
//! In the node conversion phase, Gemm nodes are converted to Linear nodes when:
//! - `alpha` = 1.0
//! - `beta` = 1.0
//! - `transB` = 1
//!
//! This optimization allows the use of optimized Linear layer implementations in Burn.
use derive_new::new;
use onnx_ir_derive::NodeBuilder;
use crate::ir::{ArgType, Argument, Node, RawNode, TensorType};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
use core::cmp::max;
/// Configuration for Gemm operation
#[derive(Debug, Clone, new)]
pub struct GemmConfig {
pub alpha: f32,
pub beta: f32,
pub trans_a: i64,
pub trans_b: i64,
}
/// Node representation for Gemm operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct GemmNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: GemmConfig,
}
pub(crate) struct GemmProcessor;
impl NodeProcessor for GemmProcessor {
type Config = GemmConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 11,
max_opset: None,
inputs: InputSpec::AtLeast(2),
outputs: OutputSpec::Exact(1),
}
}
fn infer_types(
&self,
node: &mut RawNode,
_opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
// TODO: Validate A and B tensor ranks are exactly 2 per ONNX spec - GEMM is defined for 2D matrices only - Missing rank validation
// TODO: Validate C tensor is broadcastable to output shape (M, N) per spec - Missing broadcasting validation
// TODO: Validate compatible dimensions for matrix multiplication - After transpositions, need K dimension to match - Missing dimension validation
// Extract input A tensor type
let a_rank = match &node.inputs[0].ty {
ArgType::Tensor(tensor) => tensor.rank,
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs[0].ty),
});
}
};
// Extract input B tensor type
let b_rank = match &node.inputs[1].ty {
ArgType::Tensor(tensor) => tensor.rank,
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs[1].ty),
});
}
};
let output_rank = max(a_rank, b_rank);
let elem_type = match &node.inputs[0].ty {
ArgType::Tensor(t) => t.dtype,
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs[0].ty),
});
}
};
node.outputs[0].ty = ArgType::Tensor(TensorType {
rank: output_rank,
static_shape: None,
dtype: elem_type,
});
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
let mut alpha: f32 = 1.0;
let mut beta: f32 = 1.0;
let mut trans_a: i64 = 0;
let mut trans_b: i64 = 0;
for (key, value) in node.attrs.iter() {
match key.as_str() {
"alpha" => alpha = value.clone().into_f32(),
"beta" => beta = value.clone().into_f32(),
"transA" => trans_a = value.clone().into_i64(),
"transB" => trans_b = value.clone().into_i64(),
_ => {
return Err(ProcessError::InvalidAttribute {
name: key.clone(),
reason: format!("Unexpected attribute for Gemm: {}", key),
});
}
}
}
let config = GemmConfig::new(alpha, beta, trans_a, trans_b);
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::Gemm(GemmNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::NodeType;
use crate::node::test_utils::TestNodeBuilder;
fn create_test_node(
alpha: Option<f32>,
beta: Option<f32>,
trans_a: Option<i64>,
trans_b: Option<i64>,
) -> RawNode {
let mut builder = TestNodeBuilder::new(NodeType::Gemm, "test_gemm")
.input_tensor_f32("A", 2, None)
.input_tensor_f32("B", 2, None)
.input_tensor_f32("C", 2, None)
.output_tensor_f32("Y", 2, None);
if let Some(alpha_val) = alpha {
builder = builder.attr_float("alpha", alpha_val);
}
if let Some(beta_val) = beta {
builder = builder.attr_float("beta", beta_val);
}
if let Some(trans_a_val) = trans_a {
builder = builder.attr_int("transA", trans_a_val);
}
if let Some(trans_b_val) = trans_b {
builder = builder.attr_int("transB", trans_b_val);
}
builder.build()
}
#[test]
fn test_gemm_config_defaults() {
let node = create_test_node(None, None, None, None);
let mut node = node;
let processor = GemmProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.alpha, 1.0);
assert_eq!(config.beta, 1.0);
assert_eq!(config.trans_a, 0);
assert_eq!(config.trans_b, 0);
}
#[test]
fn test_gemm_config_with_attrs() {
let node = create_test_node(Some(2.0), Some(3.0), Some(1), Some(1));
let mut node = node;
let processor = GemmProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.alpha, 2.0);
assert_eq!(config.beta, 3.0);
assert_eq!(config.trans_a, 1);
assert_eq!(config.trans_b, 1);
}
#[test]
fn test_gemm_config_partial_attrs() {
let node = create_test_node(Some(0.5), None, Some(1), None);
let mut node = node;
let processor = GemmProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.alpha, 0.5);
assert_eq!(config.beta, 1.0); // default
assert_eq!(config.trans_a, 1);
assert_eq!(config.trans_b, 0); // default
}
// TODO: Add test for non-2D tensors - GEMM requires 2D tensors, should error for rank != 2 - Missing rank validation test
// TODO: Add test for incompatible matrix dimensions - Test where K dimension doesn't match after transposition - Missing dimension validation test
// TODO: Add test for C broadcasting validation - Test various C shapes that should/shouldn't broadcast to (M, N) - Missing broadcasting test
// TODO: Add test for alpha=0 edge case - When alpha=0, should output beta*C regardless of A*B - Missing edge case test
// TODO: Add test for beta=0 edge case - When beta=0, should output alpha*A*B regardless of C - Missing edge case test
// TODO: Add test for static shape computation - When A, B have static shapes, output should compute static shape - Missing shape inference test
// TODO: Add test for different data types - Spec supports float, double, float16, bfloat16, int types - Only testing f32
// TODO: Add test for opset < 11 - Should fail per implementation requirement - Missing opset validation test
// TODO: Add test for missing C input - C is optional per spec, test with 2 inputs only - Missing optional input test
}