huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-18 11:37:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add the source offset when copying the data around.

Browse Source
This commit is contained in:
laurent
2023-06-24 08:35:49 +01:00
parent 4db972781f
commit d6cb4f1c53
5 changed files with 35 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/cpu_backend.rs
View File

@ -91,16 +91,23 @@ fn binary_map<T: Copy, F: FnMut(T, T) -> T>(
fn copy_strided_src_<T: Copy>(
src: &[T],
dst: &mut [T],
dst_offset: usize,
src_shape: &Shape,
src_stride: &[usize],
dst_offset: usize,
src_offset: usize,
) {
let src = &src[src_offset..];
if src_shape.is_contiguous(src_stride) {
dst[dst_offset..].copy_from_slice(src)
let elem_to_copy = dst.len() - dst_offset;
dst[dst_offset..].copy_from_slice(&src[..elem_to_copy])
} else {
let src_indexes = StridedIndex::new(src_shape.dims(), src_stride);
for (dst_index, src_index) in src_indexes.enumerate() {
dst[dst_index + dst_offset] = src[src_index]
let dst_index = dst_index + dst_offset;
if dst_index >= dst.len() {
break;
}
dst[dst_index] = src[src_index]
}
}
}
@ -289,22 +296,23 @@ impl CpuStorage {
pub(crate) fn copy_strided_src(
&self,
dst: &mut Self,
dst_offset: usize,
src_shape: &Shape,
src_stride: &[usize],
dst_offset: usize,
src_offset: usize,
) -> Result<()> {
if src_shape.rank() != src_stride.len() {
panic!("incoherent shape and strides {src_shape:?} {src_stride:?}")
}
match (self, dst) {
(Self::U32(src), Self::U32(dst)) => {
copy_strided_src_(src, dst, src_shape, src_stride, dst_offset)
copy_strided_src_(src, dst, dst_offset, src_shape, src_stride, src_offset)
}
(Self::F32(src), Self::F32(dst)) => {
copy_strided_src_(src, dst, src_shape, src_stride, dst_offset)
copy_strided_src_(src, dst, dst_offset, src_shape, src_stride, src_offset)
}
(Self::F64(src), Self::F64(dst)) => {
copy_strided_src_(src, dst, src_shape, src_stride, dst_offset)
copy_strided_src_(src, dst, dst_offset, src_shape, src_stride, src_offset)
}
(_, dst) => {
// This should be covered by the dtype check above.