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

Function array_copy_strided_to_linear

mlx/data/Array.cpp:406–463  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

404
405template <class T>
406void array_copy_strided_to_linear(
407 void* dst,
408 int64_t offset,
409 void* src,
410 const std::vector<int64_t>& shape,
411 const std::vector<int64_t>& stride) {
412 auto dst_t = reinterpret_cast<T*>(dst);
413 auto src_t = reinterpret_cast<T*>(src);
414 int ndim = shape.size();
415 while (ndim > 1 && (shape[ndim - 1] == 1)) {
416 ndim--;
417 }
418 // note: isize/istride considers the istride-contiguous portion
419 // of the last dimensions (not only the last one)
420 int64_t isize = shape[ndim - 1];
421 int64_t istride = stride[ndim - 1];
422 while ((ndim > 2) && (stride[ndim - 2] == isize * istride)) {
423 isize *= shape[ndim - 2];
424 ndim--;
425 }
426 int64_t osize = 1;
427 for (int dim = 0; dim < ndim - 1; dim++) {
428 osize *= shape[dim];
429 }
430 int64_t oprod = 1;
431 for (int dim = 1; dim < ndim - 1; dim++) {
432 oprod *= shape[dim];
433 }
434 if (istride == 1) {
435 for (int64_t idx = 0; idx < osize; idx++) {
436 int64_t roffset = offset;
437 int64_t r = idx;
438 int64_t rprod = oprod;
439 for (int dim = 0; dim < ndim - 1; dim++) {
440 int64_t rshape = r / rprod;
441 r = r % rprod;
442 rprod /= shape[dim + 1];
443 roffset += stride[dim] * rshape;
444 }
445 std::memcpy(dst_t + idx * isize, src_t + roffset, isize * sizeof(T));
446 }
447 } else {
448 for (int64_t idx = 0; idx < osize; idx++) {
449 int64_t roffset = offset;
450 int64_t r = idx;
451 int64_t rprod = oprod;
452 for (int dim = 0; dim < ndim - 1; dim++) {
453 int64_t rshape = r / rprod;
454 r = r % rprod;
455 rprod /= shape[dim + 1];
456 roffset += stride[dim] * rshape;
457 }
458 for (int64_t k = 0; k < isize; k++) {
459 dst_t[idx * isize + k] = src_t[roffset + k * istride];
460 }
461 }
462 }
463}

Callers

nothing calls this directly

Calls 1

sizeMethod · 0.45

Tested by

no test coverage detected