Use mkl to accelerate binary ops. (#190)

* Vectorized binary ops with mkl.

* Improve the binary op mkl support.

* Push the support for mkl binary ops.

* Proper vectorization of binary ops.

* Proper mkl'isation when broadcasting binary ops.
This commit is contained in:
Laurent Mazare
2023-07-18 12:04:39 +01:00
committed by GitHub
parent b706f32839
commit ff61a42ad7
4 changed files with 218 additions and 13 deletions

View File

@ -115,9 +115,12 @@ fn main() -> Result<()> {
let token_ids = Tensor::new(&tokens[..], device)?.unsqueeze(0)?;
let token_type_ids = token_ids.zeros_like()?;
println!("Loaded and encoded {:?}", start.elapsed());
for _ in 0..args.n {
for idx in 0..args.n {
let start = std::time::Instant::now();
let _ys = model.forward(&token_ids, &token_type_ids)?;
let ys = model.forward(&token_ids, &token_type_ids)?;
if idx == 0 {
println!("{ys}");
}
println!("Took {:?}", start.elapsed());
}
} else {