zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-22 02:58:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

More buffer donation with no-ops (#1591)

Browse Source 
* more donation

* fix test

* fix build
This commit is contained in:
Awni Hannun
2024-11-18 08:35:41 -08:00
committed by GitHub
parent 6931f84412
commit 9bd03dd9b4
7 changed files with 82 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/metal/primitives.cpp
View File

@@ -5,6 +5,7 @@
#include <sstream>
#include "mlx/backend/common/load.h"
#include "mlx/backend/common/utils.h"
#include "mlx/backend/metal/copy.h"
#include "mlx/backend/metal/device.h"
#include "mlx/backend/metal/kernels.h"
@@ -343,7 +344,7 @@ void SliceUpdate::eval_gpu(const std::vector<array>& inputs, array& out) {
auto& upd = inputs[1];
if (upd.size() == 0) {
out.copy_shared_buffer(in);
move_or_copy(in, out);
return;
}
@@ -420,8 +421,8 @@ void View::eval_gpu(const std::vector<array>& inputs, array& out) {
strides[i] *= ibytes;
strides[i] /= obytes;
}
out.copy_shared_buffer(
in, strides, in.flags(), in.data_size() * ibytes / obytes);
move_or_copy(
in, out, strides, in.flags(), in.data_size() * ibytes / obytes);
} else {
auto tmp = array(in.shape(), in.dtype(), nullptr, {});
tmp.set_data(allocator::malloc_or_wait(tmp.nbytes()));

2
mlx/backend/metal/unary.cpp
View File

@@ -161,7 +161,7 @@ void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
unary_op_gpu(inputs, out, get_primitive_string(this));
} else {
// No-op integer types
out.copy_shared_buffer(in);
move_or_copy(in, out);
}
}