huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-17 02:58:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Handle the contiguous case in an optimized way when copying cpu memory.

Browse Source
This commit is contained in:
laurent
2023-06-23 12:20:16 +01:00
parent d4054ab500
commit 52c503ba8f
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/cpu_backend.rs
View File

@ -115,19 +115,27 @@ impl CpuStorage {
src_stride: &[usize], src_stride: &[usize],
dst_offset: usize, dst_offset: usize,
) -> Result<()> { ) -> Result<()> {
// TODO: Optimize the contiguous case.
let src_indexes = StridedIndex::new(src_shape.dims(), src_stride);
match (self, dst) { match (self, dst) {
(Self::F32(src), Self::F32(dst)) => { (Self::F32(src), Self::F32(dst)) => {
if src_shape.is_contiguous(src_stride) {
dst[dst_offset..].copy_from_slice(src)
} else {
let src_indexes = StridedIndex::new(src_shape.dims(), src_stride);
for (dst_index, src_index) in src_indexes.enumerate() { for (dst_index, src_index) in src_indexes.enumerate() {
dst[dst_index + dst_offset] = src[src_index] dst[dst_index + dst_offset] = src[src_index]
} }
} }
}
(Self::F64(src), Self::F64(dst)) => { (Self::F64(src), Self::F64(dst)) => {
if src_shape.is_contiguous(src_stride) {
dst[dst_offset..].copy_from_slice(src)
} else {
let src_indexes = StridedIndex::new(src_shape.dims(), src_stride);
for (dst_index, src_index) in src_indexes.enumerate() { for (dst_index, src_index) in src_indexes.enumerate() {
dst[dst_index + dst_offset] = src[src_index] dst[dst_index + dst_offset] = src[src_index]
} }
} }
}
(_, dst) => { (_, dst) => {
// This should be covered by the dtype check above. // This should be covered by the dtype check above.
return Err(Error::DTypeMismatchBinaryOp { return Err(Error::DTypeMismatchBinaryOp {