zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-26 02:33:21 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

don't duplicate malloc with custom kernel init (#1830)

Browse Source
This commit is contained in:
Awni Hannun 2025-02-04 13:20:57 -08:00 committed by GitHub
parent f6c0499b8d
commit a229c8cef0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/metal/custom_kernel.cpp
View File

@ -15,10 +15,11 @@ void CustomKernel::eval_gpu(
std::vector<array> copies; std::vector<array> copies;
for (auto& out : outputs) { for (auto& out : outputs) {
out.set_data(allocator::malloc_or_wait(out.nbytes()));
if (init_value_) { if (init_value_) {
copies.emplace_back(init_value_.value(), out.dtype()); copies.emplace_back(init_value_.value(), out.dtype());
fill_gpu(copies.back(), out, s); fill_gpu(copies.back(), out, s);
} else {
out.set_data(allocator::malloc_or_wait(out.nbytes()));
} }
} }