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

fix copy for large arrays (#1953)

Browse Source
This commit is contained in:
Awni Hannun 2025-03-10 15:04:25 -07:00 committed by GitHub
parent 3c3e558c60
commit 117e1355a2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/backend/metal/jit_kernels.cpp
View File

@ -176,8 +176,13 @@ MTL::ComputePipelineState* get_copy_kernel(
auto out_type = get_type_string(out.dtype());
kernel_source +=
get_template_definition("s_" + lib_name, "copy_s", in_type, out_type);
kernel_source +=
get_template_definition("s2_" + lib_name, "copy_s2", in_type, out_type);
kernel_source +=
get_template_definition("v_" + lib_name, "copy_v", in_type, out_type);
kernel_source +=
get_template_definition("v2_" + lib_name, "copy_v2", in_type, out_type);
kernel_source += get_template_definition(
"g1_" + lib_name, "copy_g_nd1", in_type, out_type, "int");
kernel_source += get_template_definition(