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

Remove thrust iterators (#2396)

Browse Source
This commit is contained in:
Cheng
2025-07-21 23:30:27 +09:00
committed by GitHub
parent 93d70419e7
commit f55c4ed1d6
7 changed files with 23 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -49,6 +49,20 @@ store_vector(T* ptr, uint32_t offset, const AlignedVector<T, N>& vec) {
to[offset] = vec;
}
// 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
///////////////////////////////////////////////////////////////////////////////