zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Allow querying the allocator for the buffer size (#1404)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-09-11 21:02:16 -07:00
committed by GitHub
parent 8b30acd7eb
commit 881f09b2e2
10 changed files with 42 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
mlx/backend/common/binary.h
View File

@@ -43,13 +43,15 @@ void set_binary_op_output_data(
array& out,
BinaryOpType bopt,
bool donate_with_move = false) {
bool b_donatable = is_donatable(b, out);
bool a_donatable = is_donatable(a, out);
switch (bopt) {
case BinaryOpType::ScalarScalar:
out.set_data(
allocator::malloc_or_wait(out.itemsize()), 1, a.strides(), a.flags());
break;
case BinaryOpType::ScalarVector:
if (b.is_donatable() && b.itemsize() == out.itemsize()) {
if (b_donatable) {
if (donate_with_move) {
out.move_shared_buffer(b);
} else {
@@ -64,7 +66,7 @@ void set_binary_op_output_data(
}
break;
case BinaryOpType::VectorScalar:
if (a.is_donatable() && a.itemsize() == out.itemsize()) {
if (a_donatable) {
if (donate_with_move) {
out.move_shared_buffer(a);
} else {
@@ -79,13 +81,13 @@ void set_binary_op_output_data(
}
break;
case BinaryOpType::VectorVector:
if (a.is_donatable() && a.itemsize() == out.itemsize()) {
if (a_donatable) {
if (donate_with_move) {
out.move_shared_buffer(a);
} else {
out.copy_shared_buffer(a);
}
} else if (b.is_donatable() && b.itemsize() == out.itemsize()) {
} else if (b_donatable) {
if (donate_with_move) {
out.move_shared_buffer(b);
} else {
@@ -100,16 +102,14 @@ void set_binary_op_output_data(
}
break;
case BinaryOpType::General:
if (a.is_donatable() && a.flags().row_contiguous &&
a.itemsize() == out.itemsize() && a.size() == out.size()) {
if (a_donatable && a.flags().row_contiguous && a.size() == out.size()) {
if (donate_with_move) {
out.move_shared_buffer(a);
} else {
out.copy_shared_buffer(a);
}
} else if (
b.is_donatable() && b.flags().row_contiguous &&
b.itemsize() == out.itemsize() && b.size() == out.size()) {
b_donatable && b.flags().row_contiguous && b.size() == out.size()) {
if (donate_with_move) {
out.move_shared_buffer(b);
} else {