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

fix large ops (#1620)

Browse Source
This commit is contained in:
Awni Hannun
2024-11-24 09:17:10 -08:00
committed by GitHub
parent bb303c45a5
commit 211411faf2
12 changed files with 37 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -40,6 +40,9 @@ void ternary_op_gpu_inplace(
auto ndim = shape.size();
int work_per_thread;
if (topt == TernaryOpType::General) {
large |=
(a.data_size() > UINT32_MAX || b.data_size() > UINT32_MAX ||
c.data_size() > UINT32_MAX);
work_per_thread = large ? 4 : 2;
} else {
work_per_thread = 1;