Skip to content

Commit 16dc42b

Browse files
committed
Convert remaining yscv-onnx std HashMaps to FxHashMap
The last std::collections::HashMap holdouts: the public `ShapeMap` type alias and `infer_shapes_from_tensors`' input map, the graph_cost diagnostic map (usize-keyed — Fx is a clean win there), and the shape/ optimizer test fixtures. Move them all to FxHashMap so the workspace has zero std HashMap. Tests use FxHashMap::from_iter (no inherent `from`).
1 parent 8f85e7a commit 16dc42b

4 files changed

Lines changed: 16 additions & 12 deletions

File tree

crates/yscv-onnx/src/optimizer/graph_cost.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ pub fn graph_cost(model: &OnnxModel, shapes: &ShapeInference) -> GraphCost {
5252
let mut estimated_element_ops = 0u64;
5353
let mut estimated_bytes_read = 0u64;
5454
let mut estimated_bytes_written = 0u64;
55-
let diagnostic_by_index: std::collections::HashMap<usize, String> = shapes
55+
let diagnostic_by_index: rustc_hash::FxHashMap<usize, String> = shapes
5656
.diagnostics
5757
.iter()
5858
.map(|d| (d.node_index, d.error.to_string()))

crates/yscv-onnx/src/shape_infer.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::collections::HashMap;
1+
use rustc_hash::FxHashMap;
22

33
use thiserror::Error;
44
use yscv_tensor::Tensor;
@@ -62,7 +62,7 @@ impl TensorShape {
6262
}
6363
}
6464

65-
pub type ShapeMap = HashMap<String, TensorShape>;
65+
pub type ShapeMap = FxHashMap<String, TensorShape>;
6666

6767
#[derive(Debug, Clone, PartialEq, Eq)]
6868
pub struct ShapeDiagnostic {
@@ -157,7 +157,7 @@ pub fn infer_shapes(model: &OnnxModel, input_shapes: &ShapeMap) -> ShapeInferenc
157157

158158
pub fn infer_shapes_from_tensors(
159159
model: &OnnxModel,
160-
inputs: &HashMap<String, Tensor>,
160+
inputs: &FxHashMap<String, Tensor>,
161161
) -> ShapeInference {
162162
let input_shapes: ShapeMap = inputs
163163
.iter()

crates/yscv-onnx/src/tests/optimizer.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ fn graph_cost_report_is_deterministic_and_sorted_by_stable_key() {
173173
let model = load_onnx_model(&bytes).unwrap();
174174
let shapes = infer_shapes(
175175
&model,
176-
&std::collections::HashMap::from([
176+
&rustc_hash::FxHashMap::from_iter([
177177
("mid_a".to_string(), TensorShape::known(vec![1, 4])),
178178
("mid_b".to_string(), TensorShape::known(vec![1, 4])),
179179
]),
@@ -229,7 +229,7 @@ fn graph_cost_report_shows_lighter_graph_after_optimization() {
229229
let bytes = build_minimal_onnx_model(nodes, vec![w], vec!["x"], vec!["y"]);
230230
let mut model = load_onnx_model(&bytes).unwrap();
231231
let input_shapes =
232-
std::collections::HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]);
232+
rustc_hash::FxHashMap::from_iter([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]);
233233

234234
let before_shapes = infer_shapes(&model, &input_shapes);
235235
let before = graph_cost(&model, &before_shapes);

crates/yscv-onnx/src/tests/shape_infer.rs

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use std::collections::HashMap;
1+
use rustc_hash::FxHashMap;
22

33
use yscv_tensor::Tensor;
44

@@ -40,7 +40,8 @@ fn infers_conv_shape_and_cost() {
4040
vec!["y"],
4141
))
4242
.unwrap();
43-
let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 32, 32]))]);
43+
let input_shapes =
44+
FxHashMap::from_iter([("x".to_string(), TensorShape::known(vec![1, 3, 32, 32]))]);
4445

4546
let inferred = infer_shapes(&model, &input_shapes);
4647
assert!(
@@ -84,7 +85,8 @@ fn infers_reshape_and_matmul_shapes() {
8485
vec!["y"],
8586
))
8687
.unwrap();
87-
let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 4, 4]))]);
88+
let input_shapes =
89+
FxHashMap::from_iter([("x".to_string(), TensorShape::known(vec![1, 3, 4, 4]))]);
8890

8991
let inferred = infer_shapes(&model, &input_shapes);
9092
assert!(
@@ -116,7 +118,8 @@ fn reports_unknowns_without_guessing() {
116118
vec!["y"],
117119
))
118120
.unwrap();
119-
let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]);
121+
let input_shapes =
122+
FxHashMap::from_iter([("x".to_string(), TensorShape::known(vec![1, 3, 8, 8]))]);
120123

121124
let inferred = infer_shapes(&model, &input_shapes);
122125
assert_eq!(inferred.diagnostics.len(), 1);
@@ -144,7 +147,7 @@ fn infers_shapes_from_initializer_inputs() {
144147
))
145148
.unwrap();
146149

147-
let inferred = infer_shapes(&model, &HashMap::new());
150+
let inferred = infer_shapes(&model, &FxHashMap::default());
148151
assert_eq!(inferred.shapes["y"].as_known_dims(), Some(vec![2]));
149152
}
150153

@@ -167,7 +170,8 @@ fn global_average_pool_cost_reads_spatial_from_inferred_shape() {
167170
vec!["y"],
168171
))
169172
.unwrap();
170-
let input_shapes = HashMap::from([("x".to_string(), TensorShape::known(vec![1, 8, 4, 4]))]);
173+
let input_shapes =
174+
FxHashMap::from_iter([("x".to_string(), TensorShape::known(vec![1, 8, 4, 4]))]);
171175

172176
let inferred = infer_shapes(&model, &input_shapes);
173177
assert!(

0 commit comments

Comments
 (0)