huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Performance improvement. (#181)

Browse Source
This commit is contained in:
Laurent Mazare
2023-07-17 11:07:14 +01:00
committed by GitHub
parent 28e1c07304
commit 5b1c0bc9be
1 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
candle-core/src/cpu_backend.rs
View File

@ -66,6 +66,7 @@ struct WCond<'a>(&'a [u32], &'a Layout);
impl<'a> Map2 for WCond<'a> {
const OP: &'static str = "where";
#[inline(always)]
fn f<T: WithDType>(&self, t: &[T], t_l: &Layout, f: &[T], f_l: &Layout) -> Result<Vec<T>> {
let vs = match (
self.1.contiguous_offsets(),
@ -116,18 +117,18 @@ impl<'a> Map1 for Sum<'a> {
}
fn unary_map<T: Copy, U: Copy, F: FnMut(T) -> U>(vs: &[T], layout: &Layout, mut f: F) -> Vec<U> {
let mut result = vec![];
result.reserve(layout.shape().elem_count());
match layout.strided_blocks() {
crate::StridedBlocks::SingleBlock { start_offset, len } => {
for &v in vs[start_offset..start_offset + len].iter() {
result.push(f(v))
}
}
crate::StridedBlocks::SingleBlock { start_offset, len } => vs
[start_offset..start_offset + len]
.iter()
.map(|&v| f(v))
.collect(),
crate::StridedBlocks::MultipleBlocks {
block_start_index,
block_len,
} => {
let mut result = vec![];
result.reserve(layout.shape().elem_count());
// Specialize the case where block_len is one to avoid the second loop.
if block_len == 1 {
for index in block_start_index {
@ -142,9 +143,9 @@ fn unary_map<T: Copy, U: Copy, F: FnMut(T) -> U>(vs: &[T], layout: &Layout, mut
}
}
}
result
}
}
result
}
// This function maps over two strided index sequences.