zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-21 18:28:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

CUDA backend: unary ops (#2158)

Browse Source
This commit is contained in:
Cheng
2025-06-09 22:45:08 +09:00
committed by GitHub
parent 5866b3857b
commit f8bad60609
13 changed files with 1074 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,5 +1,6 @@
// Copyright © 2024 Apple Inc.
#include "mlx/backend/common/utils.h"
#include "mlx/backend/common/unary.h"
#include "mlx/backend/metal/device.h"
#include "mlx/backend/metal/kernels.h"
#include "mlx/backend/metal/utils.h"
@@ -99,21 +100,7 @@ void unary_op_gpu(
array& out,
const std::string op,
const Stream& s) {
auto& in = inputs[0];
bool contig = in.flags().contiguous;
if (contig) {
if (in.is_donatable() && in.itemsize() == out.itemsize()) {
out.copy_shared_buffer(in);
} else {
out.set_data(
allocator::malloc(in.data_size() * out.itemsize()),
in.data_size(),
in.strides(),
in.flags());
}
} else {
out.set_data(allocator::malloc(out.nbytes()));
}
set_unary_output_data(inputs[0], out);
unary_op_gpu_inplace(inputs, out, op, s);
}