MCPcopy Create free account
hub / github.com/ml-explore/mlx-data / array_copy_linear_to_strided

Function array_copy_linear_to_strided

mlx/data/Array.cpp:339–403  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

337
338template <class T>
339void array_copy_linear_to_strided(
340 void* dst,
341 int64_t offset,
342 void* src,
343 const std::vector<int64_t>& shape,
344 const std::vector<int64_t>& stride) {
345 auto dst_t = reinterpret_cast<T*>(dst);
346 auto src_t = reinterpret_cast<T*>(src);
347
348 // Deal with scalar arrays
349 if (shape.size() == 0) {
350 dst_t[offset] = src_t[0];
351 return;
352 }
353
354 int ndim = shape.size();
355 while (ndim > 1 && (shape[ndim - 1] == 1)) {
356 ndim--;
357 }
358 // note: isize/istride considers the istride-contiguous portion
359 // of the last dimensions (not only the last one)
360 int64_t isize = shape[ndim - 1];
361 int64_t istride = stride[ndim - 1];
362 while ((ndim > 2) && (stride[ndim - 2] == isize * istride)) {
363 isize *= shape[ndim - 2];
364 ndim--;
365 }
366 int64_t osize = 1;
367 for (int dim = 0; dim < ndim - 1; dim++) {
368 osize *= shape[dim];
369 }
370 int64_t oprod = 1;
371 for (int dim = 1; dim < ndim - 1; dim++) {
372 oprod *= shape[dim];
373 }
374 if (istride == 1) {
375 for (int64_t idx = 0; idx < osize; idx++) {
376 int64_t roffset = offset;
377 int64_t r = idx;
378 int64_t rprod = oprod;
379 for (int dim = 0; dim < ndim - 1; dim++) {
380 int64_t rshape = r / rprod;
381 r = r % rprod;
382 rprod /= shape[dim + 1];
383 roffset += stride[dim] * rshape;
384 }
385 std::memcpy(dst_t + roffset, src_t + idx * isize, isize * sizeof(T));
386 }
387 } else {
388 for (int64_t idx = 0; idx < osize; idx++) {
389 int64_t roffset = offset;
390 int64_t r = idx;
391 int64_t rprod = oprod;
392 for (int dim = 0; dim < ndim - 1; dim++) {
393 int64_t rshape = r / rprod;
394 r = r % rprod;
395 rprod /= shape[dim + 1];
396 roffset += stride[dim] * rshape;

Callers

nothing calls this directly

Calls 1

sizeMethod · 0.45

Tested by

no test coverage detected