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

Faster gather and scatter. (#682)

Browse Source 
Reduce unnecessary integer ops, especially since
there kernels are integer bound.

Increase number of iterations for benchmarks for
better smoothing.

Github Issue #506

Co-authored-by: Vijay Krishnamoorthy <vijay_krish@apple.com>
This commit is contained in:
Vijay Krish
2024-02-13 17:47:41 -08:00
committed by GitHub
parent be6e9d6a9f
commit 2fdc2462c3
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/metal/kernels/utils.h
View File

@@ -71,7 +71,7 @@ inline size_t elem_to_loc(
device const size_t* strides,
int ndim) {
size_t loc = 0;
for (int i = ndim - 1; i >= 0; --i) {
for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
loc += (elem % shape[i]) * strides[i];
elem /= shape[i];
}
@@ -84,7 +84,7 @@ inline size_t elem_to_loc(
constant const size_t* strides,
int ndim) {
size_t loc = 0;
for (int i = ndim - 1; i >= 0; --i) {
for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
loc += (elem % shape[i]) * strides[i];
elem /= shape[i];
}