use onnx_runtime_ir::{
static_shape, Attribute, DataType, Dim, Graph, Node, NodeId, Shape, TensorData, ValueId,
WeightRef,
};
use onnx_runtime_session::{InferenceSession, SessionError, Tensor, WarmupShape};
fn f32_bytes(data: &[f32]) -> Vec<u8> {
data.iter().flat_map(|v| v.to_le_bytes()).collect()
}
fn f32_init(g: &mut Graph, name: &str, dims: &[usize], data: &[f32]) -> ValueId {
let vid = g.create_named_value(name, DataType::Float32, static_shape(dims.iter().copied()));
g.set_initializer(
vid,
WeightRef::Inline(TensorData::from_raw(
DataType::Float32,
dims.to_vec(),
f32_bytes(data),
)),
);
vid
}
fn input(g: &mut Graph, name: &str, dtype: DataType, dims: &[usize]) -> ValueId {
let vid = g.create_named_value(name, dtype, static_shape(dims.iter().copied()));
g.add_input(vid);
vid
}
fn op(
g: &mut Graph,
op_type: &str,
inputs: &[ValueId],
out_dtype: DataType,
out_dims: &[usize],
attrs: &[(&str, Attribute)],
) -> ValueId {
g.opset_imports.entry(String::new()).or_insert(17);
let out = g.create_value(out_dtype, static_shape(out_dims.iter().copied()));
let mut node = Node::new(
NodeId(0),
op_type,
inputs.iter().map(|&v| Some(v)).collect(),
vec![out],
);
for (k, v) in attrs {
node.attributes.insert((*k).to_string(), v.clone());
}
g.insert_node(node);
out
}
fn input_shaped(g: &mut Graph, name: &str, dtype: DataType, shape: Shape) -> ValueId {
let vid = g.create_named_value(name, dtype, shape);
g.add_input(vid);
vid
}
fn op_shaped(
g: &mut Graph,
op_type: &str,
inputs: &[ValueId],
out_dtype: DataType,
out_shape: Shape,
attrs: &[(&str, Attribute)],
) -> ValueId {
g.opset_imports.entry(String::new()).or_insert(17);
let out = g.create_value(out_dtype, out_shape);
let mut node = Node::new(
NodeId(0),
op_type,
inputs.iter().map(|&v| Some(v)).collect(),
vec![out],
);
for (k, v) in attrs {
node.attributes.insert((*k).to_string(), v.clone());
}
g.insert_node(node);
out
}
#[test]
fn unsupported_op_error_is_actionable() {
let mut graph = Graph::new();
graph.opset_imports.insert(String::new(), 17);
let x = input(&mut graph, "x", DataType::Float32, &[1]);
let y = graph.create_named_value("y", DataType::Float32, static_shape([1]));
let mut node = Node::new(NodeId(0), "Conv", vec![Some(x)], vec![y]);
node.name = "unsupported_activation".to_string();
graph.insert_node(node);
graph.add_output(y);
let message = match InferenceSession::from_graph(graph) {
Err(err) => err.to_string(),
Ok(_) => panic!("unsupported operator unexpectedly built"),
};
assert!(message.contains("Conv"), "{message}");
assert!(message.contains("ai.onnx"), "{message}");
assert!(message.contains("unsupported_activation"), "{message}");
assert!(message.contains("opset 17"), "{message}");
assert!(message.contains("cpu_ep"), "{message}");
assert!(message.contains("To fix:"), "{message}");
}
#[test]
fn unsupported_op_error_formats_unnamed_node_gracefully() {
let mut graph = Graph::new();
graph.opset_imports.insert(String::new(), 0);
let x = input(&mut graph, "x", DataType::Float32, &[1]);
let y = graph.create_named_value("y", DataType::Float32, static_shape([1]));
graph.insert_node(Node::new(
NodeId(0),
"Conv",
vec![Some(x)],
vec![y],
));
graph.add_output(y);
let message = match InferenceSession::from_graph(graph) {
Err(err) => err.to_string(),
Ok(_) => panic!("unsupported operator unexpectedly built"),
};
assert!(
message.contains("node <unnamed node #0>, opset 0"),
"{message}"
);
assert!(!message.contains("node \"\""), "{message}");
}
#[test]
fn from_graph_rejects_missing_opset_import_at_load_time() {
let mut graph = Graph::new();
let x = input(&mut graph, "x", DataType::Float32, &[1]);
let y = graph.create_named_value("y", DataType::Float32, static_shape([1]));
let mut node = Node::new(NodeId(0), "Sigmoid", vec![Some(x)], vec![y]);
node.name = "missing_opset_import".to_string();
graph.insert_node(node);
graph.add_output(y);
let message = match InferenceSession::from_graph(graph) {
Err(err) => err.to_string(),
Ok(_) => panic!("illegal graph unexpectedly built"),
};
assert_eq!(
message,
"illegal ONNX model: operator ai.onnx::Sigmoid at node \"missing_opset_import\" uses \
domain 'ai.onnx' but no corresponding opset_import is declared. RULES #1: the model must \
declare an opset_import for domain 'ai.onnx'; if you built this graph programmatically, \
add it before loading; if this is a file, the model is malformed/invalid per the ONNX spec"
);
assert!(message.contains("Sigmoid"), "{message}");
assert!(message.contains("ai.onnx"), "{message}");
assert!(message.contains("RULES #1"), "{message}");
assert!(!message.contains("18446744073709551615"), "{message}");
}
fn ref_matmul(a: &[f32], m: usize, k: usize, b: &[f32], n: usize) -> Vec<f32> {
let mut out = vec![0.0f32; m * n];
for i in 0..m {
for j in 0..n {
let mut acc = 0.0f32;
for p in 0..k {
acc += a[i * k + p] * b[p * n + j];
}
out[i * n + j] = acc;
}
}
out
}
fn ref_add_rowvec(m: &[f32], rows: usize, cols: usize, bias: &[f32]) -> Vec<f32> {
let mut out = vec![0.0f32; rows * cols];
for r in 0..rows {
for c in 0..cols {
out[r * cols + c] = m[r * cols + c] + bias[c];
}
}
out
}
fn ref_layernorm_last(
x: &[f32],
rows: usize,
cols: usize,
scale: &[f32],
bias: &[f32],
eps: f32,
) -> Vec<f32> {
let mut out = vec![0.0f32; rows * cols];
for r in 0..rows {
let row = &x[r * cols..r * cols + cols];
let mean = row.iter().sum::<f32>() / cols as f32;
let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / cols as f32;
let inv = 1.0 / (var + eps).sqrt();
for c in 0..cols {
out[r * cols + c] = (row[c] - mean) * inv * scale[c] + bias[c];
}
}
out
}
fn ref_relu(x: &[f32]) -> Vec<f32> {
x.iter().map(|&v| v.max(0.0)).collect()
}
fn assert_close(got: &[f32], want: &[f32]) {
assert_eq!(got.len(), want.len(), "length mismatch");
for (i, (g, w)) in got.iter().zip(want).enumerate() {
assert!((g - w).abs() < 1e-4, "element {i}: got {g}, want {w}");
}
}
#[test]
fn matmul_add_layernorm_relu_chain_matches_reference() {
let x_data = [0.5f32, -1.0, 2.0, 1.5, 0.0, -0.5];
let w_data = [
0.1f32, 0.2, -0.3, 0.4, -0.5, 0.6, 0.7, -0.8, 0.9, -1.0, 0.2, 0.3,
];
let bias = [0.1f32, -0.2, 0.3, 0.05];
let scale = [1.2f32, 0.8, 1.0, 0.5];
let ln_bias = [0.0f32, 0.1, -0.1, 0.2];
let mut g = Graph::new();
let x = input(&mut g, "X", DataType::Float32, &[2, 3]);
let w = f32_init(&mut g, "W", &[3, 4], &w_data);
let m = op(&mut g, "MatMul", &[x, w], DataType::Float32, &[2, 4], &[]);
let b = f32_init(&mut g, "B", &[4], &bias);
let a = op(&mut g, "Add", &[m, b], DataType::Float32, &[2, 4], &[]);
let s = f32_init(&mut g, "Scale", &[4], &scale);
let bn = f32_init(&mut g, "LnBias", &[4], &ln_bias);
let l = op(
&mut g,
"LayerNormalization",
&[a, s, bn],
DataType::Float32,
&[2, 4],
&[("axis", Attribute::Int(-1))],
);
let y = op(&mut g, "Relu", &[l], DataType::Float32, &[2, 4], &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build session");
let x_tensor = Tensor::from_f32(&[2, 3], &x_data).unwrap();
let outputs = session.run(&[("X", &x_tensor)]).expect("run");
assert_eq!(outputs.len(), 1);
let m_ref = ref_matmul(&x_data, 2, 3, &w_data, 4);
let a_ref = ref_add_rowvec(&m_ref, 2, 4, &bias);
let l_ref = ref_layernorm_last(&a_ref, 2, 4, &scale, &ln_bias, 1e-5);
let y_ref = ref_relu(&l_ref);
assert_close(&outputs[0].to_vec_f32(), &y_ref);
assert_eq!(outputs[0].shape, vec![2, 4]);
}
#[test]
fn gather_then_transpose_matches_reference() {
let table = [
0.0f32, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0,
];
let idx = [2i64, 0, 3];
let mut g = Graph::new();
let data = f32_init(&mut g, "Table", &[4, 3], &table);
let indices = input(&mut g, "Idx", DataType::Int64, &[3]);
let gathered = op(
&mut g,
"Gather",
&[data, indices],
DataType::Float32,
&[3, 3],
&[("axis", Attribute::Int(0))],
);
let transposed = op(
&mut g,
"Transpose",
&[gathered],
DataType::Float32,
&[3, 3],
&[("perm", Attribute::Ints(vec![1, 0]))],
);
g.add_output(transposed);
let mut session = InferenceSession::from_graph(g).expect("build session");
let idx_tensor = Tensor::from_i64(&[3], &idx).unwrap();
let outputs = session.run(&[("Idx", &idx_tensor)]).expect("run");
let mut gathered_ref = Vec::new();
for &i in &idx {
let base = i as usize * 3;
gathered_ref.extend_from_slice(&table[base..base + 3]);
}
let mut want = vec![0.0f32; 9];
for r in 0..3 {
for c in 0..3 {
want[c * 3 + r] = gathered_ref[r * 3 + c];
}
}
assert_close(&outputs[0].to_vec_f32(), &want);
}
#[test]
fn shape_keyed_cache_is_reused_across_runs() {
let mut g = Graph::new();
let x = input(&mut g, "X", DataType::Float32, &[2, 2]);
let w = f32_init(&mut g, "W", &[2, 2], &[1.0, 0.0, 0.0, 1.0]);
let m = op(&mut g, "MatMul", &[x, w], DataType::Float32, &[2, 2], &[]);
let y = op(&mut g, "Relu", &[m], DataType::Float32, &[2, 2], &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build");
let after_build = session.cache_stats();
assert_eq!(after_build.entries, 2, "two nodes compiled");
assert_eq!(after_build.misses, 2);
assert_eq!(after_build.hits, 0);
let x_tensor = Tensor::from_f32(&[2, 2], &[1.0, 2.0, 3.0, 4.0]).unwrap();
let out1 = session.run(&[("X", &x_tensor)]).unwrap();
let after_run1 = session.cache_stats();
assert_eq!(after_run1.entries, 2, "no new entries on run");
assert_eq!(after_run1.misses, 2, "no recompilation");
assert_eq!(after_run1.hits, 2, "each node served from cache");
let out2 = session.run(&[("X", &x_tensor)]).unwrap();
let after_run2 = session.cache_stats();
assert_eq!(after_run2.entries, 2);
assert_eq!(after_run2.misses, 2);
assert_eq!(after_run2.hits, 4, "second run hit the cache again");
assert_close(&out1[0].to_vec_f32(), &[1.0, 2.0, 3.0, 4.0]);
assert_close(&out2[0].to_vec_f32(), &[1.0, 2.0, 3.0, 4.0]);
}
#[test]
fn warmup_validates_input_names() {
let mut g = Graph::new();
let x = input(&mut g, "X", DataType::Float32, &[1, 2]);
let y = op(&mut g, "Relu", &[x], DataType::Float32, &[1, 2], &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).unwrap();
assert!(session
.warmup(&[WarmupShape {
input_name: "nope".into(),
shape: vec![1, 2],
}])
.is_err());
assert!(session
.warmup(&[WarmupShape {
input_name: "X".into(),
shape: vec![1, 2],
}])
.is_ok());
}
#[test]
fn missing_input_is_rejected() {
let mut g = Graph::new();
let x = input(&mut g, "X", DataType::Float32, &[1, 2]);
let y = op(&mut g, "Relu", &[x], DataType::Float32, &[1, 2], &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).unwrap();
let err = session.run(&[]).unwrap_err();
assert!(matches!(
err,
onnx_runtime_session::SessionError::InputNotFound { .. }
));
}
#[test]
fn input_shape_mismatch_is_rejected() {
let mut g = Graph::new();
let x = input(&mut g, "X", DataType::Float32, &[2, 2]);
let y = op(&mut g, "Relu", &[x], DataType::Float32, &[2, 2], &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).unwrap();
let wrong = Tensor::from_f32(&[3, 2], &[0.0; 6]).unwrap();
let err = session.run(&[("X", &wrong)]).unwrap_err();
assert!(matches!(
err,
onnx_runtime_session::SessionError::ShapeMismatch { .. }
));
}
#[test]
fn symbolic_batch_matmul_chain_runs_for_multiple_shapes() {
let w_data = [
1.0f32, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0,
];
let bias = [0.5f32, -0.5, 1.0, -1.0];
let mut g = Graph::new();
let batch = g.intern_symbol("batch");
let sym_row = || vec![Dim::Symbolic(batch), Dim::Static(4)];
let x = input_shaped(&mut g, "X", DataType::Float32, sym_row());
let w = f32_init(&mut g, "W", &[4, 4], &w_data);
let m = op_shaped(&mut g, "MatMul", &[x, w], DataType::Float32, sym_row(), &[]);
let b = f32_init(&mut g, "B", &[4], &bias);
let a = op_shaped(&mut g, "Add", &[m, b], DataType::Float32, sym_row(), &[]);
let y = op_shaped(&mut g, "Relu", &[a], DataType::Float32, sym_row(), &[]);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build symbolic session");
let after_build = session.cache_stats();
assert_eq!(after_build.entries, 0, "no kernels compiled before first run");
assert_eq!(after_build.misses, 0);
let run_batch = |session: &mut InferenceSession, rows: usize, fill: f32| -> Vec<f32> {
let data: Vec<f32> = (0..rows * 4).map(|i| fill + i as f32).collect();
let x_tensor = Tensor::from_f32(&[rows, 4], &data).unwrap();
let out = session.run(&[("X", &x_tensor)]).expect("run");
assert_eq!(out[0].shape, vec![rows, 4]);
let m_ref = ref_matmul(&data, rows, 4, &w_data, 4);
let a_ref = ref_add_rowvec(&m_ref, rows, 4, &bias);
let y_ref = ref_relu(&a_ref);
assert_close(&out[0].to_vec_f32(), &y_ref);
out[0].to_vec_f32()
};
run_batch(&mut session, 2, 0.0);
let s2 = session.cache_stats();
assert_eq!(s2.entries, 3, "three nodes compiled for batch=2");
assert_eq!(s2.misses, 3);
assert_eq!(s2.hits, 0);
run_batch(&mut session, 3, 10.0);
let s3 = session.cache_stats();
assert_eq!(s3.entries, 6, "batch=3 adds three distinct shape-keyed entries");
assert_eq!(s3.misses, 6);
assert_eq!(s3.hits, 0);
run_batch(&mut session, 2, 100.0);
let s2b = session.cache_stats();
assert_eq!(s2b.entries, 6, "no new entries: batch=2 plan reused");
assert_eq!(s2b.misses, 6);
assert_eq!(s2b.hits, 3, "each node served from the batch=2 cache");
}
#[test]
fn symbol_conflict_across_inputs_is_rejected() {
let mut g = Graph::new();
let batch = g.intern_symbol("batch");
let sym_row = || vec![Dim::Symbolic(batch), Dim::Static(4)];
let a = input_shaped(&mut g, "A", DataType::Float32, sym_row());
let b = input_shaped(&mut g, "B", DataType::Float32, sym_row());
let s = op_shaped(&mut g, "Add", &[a, b], DataType::Float32, sym_row(), &[]);
g.add_output(s);
let mut session = InferenceSession::from_graph(g).expect("build");
let a_t = Tensor::from_f32(&[2, 4], &[0.0; 8]).unwrap();
let b_t = Tensor::from_f32(&[3, 4], &[0.0; 12]).unwrap();
let err = session.run(&[("A", &a_t), ("B", &b_t)]).unwrap_err();
assert!(
matches!(err, SessionError::SymbolConflict { .. }),
"expected SymbolConflict, got {err:?}"
);
let a_ok = Tensor::from_f32(&[2, 4], &[1.0; 8]).unwrap();
let b_ok = Tensor::from_f32(&[2, 4], &[2.0; 8]).unwrap();
let out = session.run(&[("A", &a_ok), ("B", &b_ok)]).expect("run");
assert_close(&out[0].to_vec_f32(), &[3.0; 8]);
assert_eq!(out[0].shape, vec![2, 4]);
}
#[test]
fn unresolved_symbol_reports_uninferred_shape() {
let mut g = Graph::new();
let batch = g.intern_symbol("batch");
let ghost = g.intern_symbol("ghost");
let x = input_shaped(
&mut g,
"X",
DataType::Float32,
vec![Dim::Symbolic(batch), Dim::Static(4)],
);
let y = op_shaped(
&mut g,
"Relu",
&[x],
DataType::Float32,
vec![Dim::Symbolic(ghost), Dim::Static(4)],
&[],
);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build");
let x_t = Tensor::from_f32(&[2, 4], &[0.0; 8]).unwrap();
let err = session.run(&[("X", &x_t)]).unwrap_err();
assert!(
matches!(err, SessionError::UnresolvedShape { ref op, .. } if op == "Relu"),
"expected UnresolvedShape naming the producing op, got {err:?}"
);
}
#[test]
fn symbolic_input_rank_mismatch_is_rejected() {
let mut g = Graph::new();
let batch = g.intern_symbol("batch");
let x = input_shaped(
&mut g,
"X",
DataType::Float32,
vec![Dim::Symbolic(batch), Dim::Static(4)],
);
let y = op_shaped(
&mut g,
"Relu",
&[x],
DataType::Float32,
vec![Dim::Symbolic(batch), Dim::Static(4)],
&[],
);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build");
let wrong = Tensor::from_f32(&[2, 2, 4], &[0.0; 16]).unwrap();
let err = session.run(&[("X", &wrong)]).unwrap_err();
assert!(
matches!(err, SessionError::RankMismatch { .. }),
"expected RankMismatch, got {err:?}"
);
}
#[test]
fn symbolic_input_static_dim_mismatch_is_rejected() {
let mut g = Graph::new();
let batch = g.intern_symbol("batch");
let x = input_shaped(
&mut g,
"X",
DataType::Float32,
vec![Dim::Symbolic(batch), Dim::Static(4)],
);
let y = op_shaped(
&mut g,
"Relu",
&[x],
DataType::Float32,
vec![Dim::Symbolic(batch), Dim::Static(4)],
&[],
);
g.add_output(y);
let mut session = InferenceSession::from_graph(g).expect("build");
let wrong = Tensor::from_f32(&[2, 5], &[0.0; 10]).unwrap();
let err = session.run(&[("X", &wrong)]).unwrap_err();
assert!(
matches!(err, SessionError::ShapeMismatch { .. }),
"expected ShapeMismatch on the static dim, got {err:?}"
);
}
#[test]
fn from_graph_rejects_unimplemented_control_flow_subgraph_at_build() {
let mut graph = Graph::new();
graph.opset_imports.insert(String::new(), 17);
let x = input(&mut graph, "x", DataType::Float32, &[1]);
let y = graph.create_named_value("y", DataType::Float32, static_shape([1]));
let mut node = Node::new(NodeId(0), "SequenceMap", vec![Some(x)], vec![y]);
node.name = "control_flow_seqmap".to_string();
node.attributes
.insert("body".to_string(), Attribute::Graph(Box::new(Graph::new())));
graph.insert_node(node);
graph.add_output(y);
let message = match InferenceSession::from_graph(graph) {
Err(err) => err.to_string(),
Ok(_) => panic!("unimplemented control-flow subgraph unexpectedly built"),
};
assert!(message.contains("SequenceMap"), "{message}");
assert!(message.contains("body"), "{message}");
assert!(message.contains("control-flow"), "{message}");
assert!(message.contains("RULES #1"), "{message}");
}
#[test]
fn from_graph_rejects_dangling_tensor_reference_at_build() {
let mut graph = Graph::new();
graph.opset_imports.insert(String::new(), 17);
let x = input(&mut graph, "x", DataType::Float32, &[2]);
let z = graph.create_named_value("z", DataType::Float32, static_shape([2]));
let y = graph.create_named_value("y", DataType::Float32, static_shape([2]));
let mut node = Node::new(NodeId(0), "Add", vec![Some(x), Some(z)], vec![y]);
node.name = "dangling_add".to_string();
graph.insert_node(node);
graph.add_output(y);
let message = match InferenceSession::from_graph(graph) {
Err(err) => err.to_string(),
Ok(_) => panic!("dangling reference unexpectedly built"),
};
assert!(message.contains("'z'"), "{message}");
assert!(message.contains("Add"), "{message}");
assert!(message.contains("RULES #1"), "{message}");
}