zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-01 21:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix unintuitive metal kernel caching (#2242)

Browse Source 
* Fix unintuitive metal kernel caching

* alternative solution
This commit is contained in:
Awni Hannun
2025-06-06 20:08:15 -07:00
committed by GitHub
parent 2e8cf0b450
commit 1ca616844b
13 changed files with 713 additions and 593 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
examples/extensions/axpby/axpby.cpp
View File

@@ -172,11 +172,11 @@ void Axpby::eval_gpu(
kname << (contiguous_kernel ? "contiguous_" : "general_");
kname << type_to_name(out);
// Make sure the metal library is available
d.register_library("mlx_ext");
// Load the metal library
auto lib = d.get_library("mlx_ext");
// Make a kernel from this metal library
auto kernel = d.get_kernel(kname.str(), "mlx_ext");
auto kernel = d.get_kernel(kname.str(), lib);
// Prepare to encode kernel
auto& compute_encoder = d.get_command_encoder(s.index);