xref: /aosp_15_r20/external/executorch/backends/apple/coreml/runtime/delegate/MLMultiArray_Copy.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1//
2// MLMultiArray+Copy.mm
3//
4// Copyright © 2024 Apple Inc. All rights reserved.
5//
6// Please refer to the license found in the LICENSE file in the root directory of the source tree.
7
8#import <MLMultiArray_Copy.h>
9
10#import <objc_array_util.h>
11#import <multiarray.h>
12
13namespace {
14using namespace executorchcoreml;
15
16MultiArray to_multi_array(void *data,
17                          MLMultiArrayDataType dataType,
18                          NSArray<NSNumber *> *shape,
19                          NSArray<NSNumber *> *strides) {
20    auto layout = MultiArray::MemoryLayout(to_multiarray_data_type(dataType).value(),
21                                           to_vector<size_t>(shape),
22                                           to_vector<ssize_t>(strides));
23    return MultiArray(data, std::move(layout));
24}
25} //namespace
26
27@implementation MLMultiArray (Copy)
28
29- (void)copyInto:(MLMultiArray *)dstMultiArray {
30    [self getBytesWithHandler:^(const void *srcBytes, __unused NSInteger srcSize) {
31        [dstMultiArray getMutableBytesWithHandler:^(void *dstBytes, __unused NSInteger size, NSArray<NSNumber *> * strides) {
32            auto src = ::to_multi_array(const_cast<void *>(srcBytes), self.dataType, self.shape, self.strides);
33            auto dst = ::to_multi_array(dstBytes, dstMultiArray.dataType, dstMultiArray.shape, strides);
34            src.copy(dst);
35        }];
36    }];
37}
38
39@end
40