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

Remove redundant simd_sum in logsumexp (#2210)

Browse Source
This commit is contained in:
Cheng 2025-05-21 23:25:03 +09:00 committed by GitHub
parent 35c87741cf
commit 7774b87cbd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
mlx/backend/metal/kernels/logsumexp.h
View File

@ -134,10 +134,7 @@ template <typename T, typename AccT = float, int N_READS = 4>
threadgroup_barrier(mem_flags::mem_threadgroup); threadgroup_barrier(mem_flags::mem_threadgroup);
normalizer = simd_sum(local_normalizer[simd_lane_id]); normalizer = simd_sum(local_normalizer[simd_lane_id]);
if (simd_group_id == 0) { if (lid == 0) {
normalizer = simd_sum(local_normalizer[simd_lane_id]); out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
if (simd_lane_id == 0) {
out[gid] = isinf(maxval) ? T(maxval) : T(log(normalizer) + maxval);
}
} }
} }