|
7 | 7 | #include <memory> |
8 | 8 | #include <string> |
9 | 9 | #include <vector> |
| 10 | +#include <span> |
10 | 11 |
|
11 | 12 | #ifdef __OBJC__ |
12 | 13 | #import <Foundation/Foundation.h> |
13 | 14 | #endif |
14 | 15 |
|
15 | | -template <typename T> NSArray *arrayToNSArray(const void *array, ssize_t numel); |
| 16 | +using namespace ::executorch::extension; |
| 17 | +using namespace ::torch::executor; |
16 | 18 |
|
17 | | -std::vector<int> NSArrayToIntVector(NSArray *inputArray); |
| 19 | +template <typename T> T getValueFromNSNumber(NSNumber *number) { |
| 20 | + if constexpr (std::is_same<T, int8_t>::value) { |
| 21 | + return static_cast<T>([number charValue]); // `charValue` for 8-bit integers |
| 22 | + } else if constexpr (std::is_same<T, int32_t>::value) { |
| 23 | + return static_cast<T>([number intValue]); // `intValue` for 32-bit integers |
| 24 | + } else if constexpr (std::is_same<T, int64_t>::value) { |
| 25 | + return static_cast<T>( |
| 26 | + [number longLongValue]); // `longLongValue` for 64-bit integers |
| 27 | + } else if constexpr (std::is_same<T, float>::value) { |
| 28 | + return static_cast<T>([number floatValue]); |
| 29 | + } else if constexpr (std::is_same<T, double>::value) { |
| 30 | + return static_cast<T>([number doubleValue]); |
| 31 | + } |
| 32 | +} |
18 | 33 |
|
19 | 34 | template <typename T> |
20 | | -std::unique_ptr<T[]> NSArrayToTypedArray(NSArray *nsArray); |
| 35 | +std::unique_ptr<T[]> NSArrayToTypedArray(NSArray *nsArray) { |
| 36 | + size_t arraySize = [nsArray count]; |
21 | 37 |
|
22 | | -template <typename T> T getValueFromNSNumber(NSNumber *number); |
| 38 | + std::unique_ptr<T[]> typedArray(new T[arraySize]); |
| 39 | + |
| 40 | + for (NSUInteger i = 0; i < arraySize; ++i) { |
| 41 | + NSNumber *number = [nsArray objectAtIndex:i]; |
| 42 | + if ([number isKindOfClass:[NSNumber class]]) { |
| 43 | + typedArray[i] = getValueFromNSNumber<T>(number); |
| 44 | + } else { |
| 45 | + typedArray[i] = T(); |
| 46 | + } |
| 47 | + } |
| 48 | + return typedArray; |
| 49 | +} |
| 50 | + |
| 51 | +template <typename T> |
| 52 | +NSArray *arrayToNSArray(const void *array, ssize_t numel) { |
| 53 | + const T *typedArray = static_cast<const T *>(array); |
| 54 | + NSMutableArray *nsArray = [NSMutableArray arrayWithCapacity:numel]; |
| 55 | + |
| 56 | + for (int i = 0; i < numel; ++i) { |
| 57 | + [nsArray addObject:@(typedArray[i])]; |
| 58 | + } |
| 59 | + |
| 60 | + return [nsArray copy]; |
| 61 | +} |
23 | 62 |
|
24 | 63 | template <typename T> |
25 | | -const T* |
26 | | -runForwardFromNSArray(NSArray *inputArray, ssize_t& numel, std::vector<int> shapes, |
27 | | - std::unique_ptr<executorch::extension::Module> &model); |
| 64 | +NSArray *arrayToNSArray(const std::vector<std::span<const T>> &dataPtrVec) { |
| 65 | + NSMutableArray *nsArray = [NSMutableArray array]; |
| 66 | + for (const auto &span : dataPtrVec) { |
| 67 | + NSMutableArray *innerArray = [NSMutableArray arrayWithCapacity:span.size()]; |
| 68 | + for(auto x : span) { |
| 69 | + [innerArray addObject:@(x)]; |
| 70 | + } |
| 71 | + [nsArray addObject:[innerArray copy]]; |
| 72 | + } |
| 73 | + return [nsArray copy]; |
| 74 | +} |
| 75 | + |
| 76 | +std::vector<int> NSArrayToIntVector(NSArray *inputArray) { |
| 77 | + std::vector<int> output; |
| 78 | + for (NSUInteger i = 0; i < [inputArray count]; ++i) { |
| 79 | + NSNumber *number = [inputArray objectAtIndex:i]; |
| 80 | + if (number) { |
| 81 | + output.push_back([number intValue]); |
| 82 | + } else { |
| 83 | + output.push_back(0); |
| 84 | + } |
| 85 | + } |
| 86 | + return output; |
| 87 | +} |
| 88 | + |
| 89 | +template <typename T> |
| 90 | +std::vector<std::span<const T>> |
| 91 | +runForwardFromNSArray(NSArray *inputArray, std::vector<int> shapes, |
| 92 | + std::unique_ptr<Module> &model) { |
| 93 | + std::unique_ptr<T[]> inputPtr = NSArrayToTypedArray<T>(inputArray); |
| 94 | + |
| 95 | + TensorPtr inputTensor = from_blob(inputPtr.get(), shapes); |
| 96 | + Result result = model->forward(inputTensor); |
| 97 | + |
| 98 | + if (result.ok()) { |
| 99 | + std::vector<std::span<const T>> outputVec; |
| 100 | + |
| 101 | + for (const auto ¤tResult : *result) { |
| 102 | + Tensor currentTensor = currentResult.toTensor(); |
| 103 | + std::span<const T> currentSpan(currentTensor.const_data_ptr<T>(), currentTensor.numel()); |
| 104 | + outputVec.push_back(std::move(currentSpan)); |
| 105 | + } |
| 106 | + return outputVec; |
| 107 | + } |
28 | 108 |
|
| 109 | + @throw [NSException |
| 110 | + exceptionWithName:@"forward_error" |
| 111 | + reason:[NSString stringWithFormat:@"%d", (int)result.error()] |
| 112 | + userInfo:nil]; |
| 113 | +} |
29 | 114 |
|
30 | 115 | #endif // Utils_hpp |
0 commit comments