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

Use load_vector in arg_reduce (#2439)

Browse Source
This commit is contained in:
Cheng
2025-07-30 17:40:26 +09:00
committed by GitHub
parent a0ae49d397
commit 3628e5d497
2 changed files with 12 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
mlx/backend/cuda/device/utils.cuh
View File

@@ -131,20 +131,6 @@ inline __device__ void store_vector(
}
}
// Helper for accessing strided data.
template <typename T>
struct StridedIterator {
T it;
int64_t stride;
__host__ __device__ StridedIterator(T it, int64_t stride)
: it(it), stride(stride) {}
__host__ __device__ auto operator[](int i) const {
return it[i * stride];
}
};
///////////////////////////////////////////////////////////////////////////////
// Type limits utils
///////////////////////////////////////////////////////////////////////////////