Skip to content

Commit 187ee53

Browse files
committed
Use AHash for borrow checking hash tables.
1 parent 5e321dc commit 187ee53

File tree

3 files changed

+11
-9
lines changed

3 files changed

+11
-9
lines changed

Cargo.toml

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ keywords = ["python", "numpy", "ffi", "pyo3"]
1515
license = "BSD-2-Clause"
1616

1717
[dependencies]
18+
ahash = "0.7"
1819
libc = "0.2"
1920
num-complex = ">= 0.2, < 0.5"
2021
num-integer = "0.1"
@@ -23,7 +24,7 @@ ndarray = ">= 0.13, < 0.16"
2324
pyo3 = { version = "0.16", default-features = false, features = ["macros"] }
2425

2526
[dev-dependencies]
26-
pyo3 = { version = "0.16", features = ["auto-initialize"] }
27+
pyo3 = { version = "0.16", default-features = false, features = ["auto-initialize"] }
2728

2829
[workspace]
2930
members = ["examples/*"]

benches/borrow.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ use pyo3::Python;
99
#[bench]
1010
fn initial_shared_borrow(bencher: &mut Bencher) {
1111
Python::with_gil(|py| {
12-
let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
12+
let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
1313

1414
bencher.iter(|| {
1515
let array = black_box(array);
@@ -22,7 +22,7 @@ fn initial_shared_borrow(bencher: &mut Bencher) {
2222
#[bench]
2323
fn additional_shared_borrow(bencher: &mut Bencher) {
2424
Python::with_gil(|py| {
25-
let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
25+
let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
2626

2727
let _shared = (0..128).map(|_| array.readonly()).collect::<Vec<_>>();
2828

@@ -37,7 +37,7 @@ fn additional_shared_borrow(bencher: &mut Bencher) {
3737
#[bench]
3838
fn exclusive_borrow(bencher: &mut Bencher) {
3939
Python::with_gil(|py| {
40-
let array = PyArray::<f64, _>::zeros(py, (1, 2, 3), false);
40+
let array = PyArray::<f64, _>::zeros(py, (6, 5, 4, 3, 2, 1), false);
4141

4242
bencher.iter(|| {
4343
let array = black_box(array);

src/borrow.rs

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -163,11 +163,12 @@
163163

164164
use std::any::type_name;
165165
use std::cell::UnsafeCell;
166-
use std::collections::hash_map::{Entry, HashMap};
166+
use std::collections::hash_map::Entry;
167167
use std::fmt;
168168
use std::mem::size_of;
169169
use std::ops::Deref;
170170

171+
use ahash::AHashMap;
171172
use ndarray::{ArrayView, ArrayViewMut, Dimension, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, IxDyn};
172173
use num_integer::gcd;
173174
use pyo3::{FromPyObject, PyAny, PyResult};
@@ -235,7 +236,7 @@ impl BorrowKey {
235236
}
236237
}
237238

238-
type BorrowFlagsInner = HashMap<usize, HashMap<BorrowKey, isize>>;
239+
type BorrowFlagsInner = AHashMap<usize, AHashMap<BorrowKey, isize>>;
239240

240241
struct BorrowFlags(UnsafeCell<Option<BorrowFlagsInner>>);
241242

@@ -248,7 +249,7 @@ impl BorrowFlags {
248249

249250
#[allow(clippy::mut_from_ref)]
250251
unsafe fn get(&self) -> &mut BorrowFlagsInner {
251-
(*self.0.get()).get_or_insert_with(HashMap::new)
252+
(*self.0.get()).get_or_insert_with(AHashMap::new)
252253
}
253254

254255
fn acquire<T, D>(&self, array: &PyArray<T, D>) -> Result<(), BorrowError>
@@ -292,7 +293,7 @@ impl BorrowFlags {
292293
}
293294
}
294295
Entry::Vacant(entry) => {
295-
let mut same_base_arrays = HashMap::with_capacity(1);
296+
let mut same_base_arrays = AHashMap::with_capacity(1);
296297
same_base_arrays.insert(key, 1);
297298
entry.insert(same_base_arrays);
298299
}
@@ -363,7 +364,7 @@ impl BorrowFlags {
363364
}
364365
}
365366
Entry::Vacant(entry) => {
366-
let mut same_base_arrays = HashMap::with_capacity(1);
367+
let mut same_base_arrays = AHashMap::with_capacity(1);
367368
same_base_arrays.insert(key, -1);
368369
entry.insert(same_base_arrays);
369370
}

0 commit comments

Comments
 (0)