Skip to content

Commit 6b4b013

Browse files
authored
Use Memcpy in copy_utils
Differential Revision: D76061894 Pull Request resolved: #11430
1 parent f3a6c18 commit 6b4b013

File tree

1 file changed

+11
-4
lines changed

1 file changed

+11
-4
lines changed

kernels/portable/cpu/util/copy_ops_util.h

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,17 @@ void _as_strided_copy(
2727
ArrayRef<int64_t> stride,
2828
int64_t dim) {
2929
// the last dimension, copy data
30+
const int64_t stride_dim = stride.at(dim);
3031
if (dim == static_cast<int64_t>(size.size()) - 1) {
31-
for (const auto i : c10::irange(size.at(dim))) {
32-
output_data[i] = *input_data;
33-
input_data += stride.at(dim);
32+
const size_t num_elements = size.at(dim);
33+
// use memcpy for contiguous memory
34+
if (stride_dim == 1) {
35+
memcpy(output_data, input_data, num_elements * sizeof(CTYPE));
36+
} else {
37+
for (const auto i : c10::irange(num_elements)) {
38+
output_data[i] = *input_data;
39+
input_data += stride_dim;
40+
}
3441
}
3542
return;
3643
}
@@ -39,7 +46,7 @@ void _as_strided_copy(
3946
for ([[maybe_unused]] const auto i : c10::irange(size.at(dim))) {
4047
_as_strided_copy<CTYPE>(
4148
input_data, output_data, out, size, stride, dim + 1);
42-
input_data += stride.at(dim);
49+
input_data += stride_dim;
4350
output_data += trailing_dims;
4451
}
4552
}

0 commit comments

Comments
 (0)