fix: Change iOS native bindings so they return an array of outputs (#43)

This PR fixes the issue in iOS ExecuTorch bindings where if the model
returned multiple output arrays, only the first one was considered.

## Description
<!-- Provide a concise and descriptive summary of the changes
implemented in this PR. -->

### Type of change
- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] Documentation update (improves or adds clarity to existing
documentation)

### Tested on
- [x] iOS
- [ ] Android

### Testing instructions
<!-- Provide step-by-step instructions on how to test your changes.
Include setup details if necessary. -->

### Screenshots
<!-- Add screenshots here, if applicable -->

### Related issues
<!-- Link related issues here using #issue-number -->

### Checklist
- [x] I have performed a self-review of my code
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have updated the documentation accordingly
- [ ] My changes generate no new warnings

### Additional notes
<!-- Include any additional information, assumptions, or context that
reviewers might need to understand this PR. -->

Author: chmjkb

ios/ExecutorchLib.xcframework/ios-arm64-simulator/ExecutorchLib.framework/ExecutorchLib

@@ -21,7 +21,7 @@ export interface ExecutorchModule {
   error: string | null;
   isModelLoading: boolean;
   isModelRunning: boolean;
-  forward: (input: ETInput, shape: number[]) => Promise<number[]>;
+  forward: (input: ETInput, shape: number[]) => Promise<number[][]>;
   loadMethod: (methodName: string) => Promise<void>;
   loadForward: () => Promise<void>;
 }

ios/ExecutorchLib.xcframework/ios-arm64/ExecutorchLib.framework/ExecutorchLib

@@ -42,7 +42,7 @@
 		55EA2C572CB90E7D004315B3 /* Accelerate.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 55EA2C562CB90E7D004315B3 /* Accelerate.framework */; };
 		55EA2C592CB90E80004315B3 /* CoreML.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 55EA2C582CB90E80004315B3 /* CoreML.framework */; };
 		55EA2C5B2CB90E85004315B3 /* libsqlite3.tbd in Frameworks */ = {isa = PBXBuildFile; fileRef = 55EA2C5A2CB90E85004315B3 /* libsqlite3.tbd */; };
-		A851C4062CF9F1B600424E93 /* Utils.mm in Sources */ = {isa = PBXBuildFile; fileRef = A851C4052CF9F1B600424E93 /* Utils.mm */; };
+		A84198842D02DF29006D4D5E /* InputType.h in Headers */ = {isa = PBXBuildFile; fileRef = A84198832D02DF29006D4D5E /* InputType.h */; };
 		A851C4072CF9F1B600424E93 /* Utils.hpp in Headers */ = {isa = PBXBuildFile; fileRef = A851C4042CF9F1B600424E93 /* Utils.hpp */; };
 /* End PBXBuildFile section */
 
@@ -83,8 +83,8 @@
 		55EA2C562CB90E7D004315B3 /* Accelerate.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Accelerate.framework; path = System/Library/Frameworks/Accelerate.framework; sourceTree = SDKROOT; };
 		55EA2C582CB90E80004315B3 /* CoreML.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreML.framework; path = System/Library/Frameworks/CoreML.framework; sourceTree = SDKROOT; };
 		55EA2C5A2CB90E85004315B3 /* libsqlite3.tbd */ = {isa = PBXFileReference; lastKnownFileType = "sourcecode.text-based-dylib-definition"; name = libsqlite3.tbd; path = usr/lib/libsqlite3.tbd; sourceTree = SDKROOT; };
+		A84198832D02DF29006D4D5E /* InputType.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = InputType.h; sourceTree = "<group>"; };
 		A851C4042CF9F1B600424E93 /* Utils.hpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.h; path = Utils.hpp; sourceTree = "<group>"; };
-		A851C4052CF9F1B600424E93 /* Utils.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = Utils.mm; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -135,7 +135,7 @@
 				55EA2C322CB90C7A004315B3 /* sampler */,
 				55EA2C3E2CB90C7A004315B3 /* tokenizer */,
 				A851C4042CF9F1B600424E93 /* Utils.hpp */,
-				A851C4052CF9F1B600424E93 /* Utils.mm */,
+				A84198832D02DF29006D4D5E /* InputType.h */,
 			);
 			path = ExecutorchLib;
 			sourceTree = "<group>";
@@ -220,6 +220,7 @@
 				55EA2C542CB90E70004315B3 /* LLaMARunner.h in Headers */,
 				5576B4B72CEF9709005027B7 /* ETModel.h in Headers */,
 				55EA2C532CB90C7A004315B3 /* tokenizer.h in Headers */,
+				A84198842D02DF29006D4D5E /* InputType.h in Headers */,
 				55EA2C412CB90C7A004315B3 /* stats.h in Headers */,
 				55EA2C4E2CB90C7A004315B3 /* bpe_tokenizer.h in Headers */,
 				55EA2C402CB90C7A004315B3 /* runner.h in Headers */,
@@ -308,7 +309,6 @@
 			buildActionMask = 2147483647;
 			files = (
 				55EA2C482CB90C7A004315B3 /* sampler.cpp in Sources */,
-				A851C4062CF9F1B600424E93 /* Utils.mm in Sources */,
 				55EA2C3F2CB90C7A004315B3 /* runner.cpp in Sources */,
 				55EA2C422CB90C7A004315B3 /* text_decoder_runner.cpp in Sources */,
 				55EA2C4D2CB90C7A004315B3 /* bpe_tokenizer.cpp in Sources */,

src/types.ts

@@ -1,5 +1,6 @@
 #import "ETModel.h"
 #include "Utils.hpp"
+#include <InputType.h>
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/core/error.h>
@@ -34,38 +35,32 @@ - (NSArray *)forward:(NSArray *)input
            inputType:(NSNumber *)inputType {
   int inputTypeIntValue = [inputType intValue];
   std::vector<int> shapes = NSArrayToIntVector(shape);
-  ssize_t outputNumel = 0;
   @try {
     switch (inputTypeIntValue) {
-    case 0: {
-      // Int8Array
-      const int8_t *output =
-          runForwardFromNSArray<int8_t>(input, outputNumel, shapes, _model);
-      return arrayToNSArray<int8_t>(output, outputNumel);
+    case InputTypeInt8: {
+      std::vector<std::span<const int8_t>> output =
+          runForwardFromNSArray<int8_t>(input, shapes, _model);
+      return arrayToNSArray<int8_t>(output);
     }
-    case 1: {
-      // Int32Array
-      const int32_t *output =
-          runForwardFromNSArray<int32_t>(input, outputNumel, shapes, _model);
-      return arrayToNSArray<int32_t>(output, outputNumel);
+    case InputTypeInt32: {
+      std::vector<std::span<const int32_t>> output =
+          runForwardFromNSArray<int32_t>(input, shapes, _model);
+      return arrayToNSArray<int32_t>(output);
     }
-    case 2: {
-      // BigInt64Array
-      const int64_t *output =
-          runForwardFromNSArray<int64_t>(input, outputNumel, shapes, _model);
-      return arrayToNSArray<int64_t>(output, outputNumel);
+    case InputTypeInt64: {
+      std::vector<std::span<const int64_t>> output =
+          runForwardFromNSArray<int64_t>(input, shapes, _model);
+      return arrayToNSArray<int64_t>(output);
     }
-    case 3: {
-      // Float32Array
-      const float *output =
-          runForwardFromNSArray<float>(input, outputNumel, shapes, _model);
-      return arrayToNSArray<float>(output, outputNumel);
+    case InputTypeFloat32: {
+      std::vector<std::span<const float>> output =
+          runForwardFromNSArray<float>(input, shapes, _model);
+      return arrayToNSArray<float>(output);
     }
-    case 4: {
-      // Float64Array
-      const double *output =
-          runForwardFromNSArray<double>(input, outputNumel, shapes, _model);
-      return arrayToNSArray<double>(output, outputNumel);
+    case InputTypeFloat64: {
+      std::vector<std::span<const double>> output =
+          runForwardFromNSArray<double>(input, shapes, _model);
+      return arrayToNSArray<double>(output);
     }
     }
   } @catch (NSException *exception) {

third-party/ios/ExecutorchLib/ExecutorchLib.xcodeproj/project.pbxproj

@@ -0,0 +1,12 @@
+#ifndef InputType_h
+#define InputType_h
+
+typedef NS_ENUM(NSInteger, InputType) {
+    InputTypeInt8 = 0,
+    InputTypeInt32 = 1,
+    InputTypeInt64 = 2,
+    InputTypeFloat32 = 3,
+    InputTypeFloat64 = 4,
+};
+
+#endif /* InputType_h */

third-party/ios/ExecutorchLib/ExecutorchLib/Exported/ETModel.mm

@@ -7,24 +7,109 @@
 #include <memory>
 #include <string>
 #include <vector>
+#include <span>
 
 #ifdef __OBJC__
 #import <Foundation/Foundation.h>
 #endif
 
-template <typename T> NSArray *arrayToNSArray(const void *array, ssize_t numel);
+using namespace ::executorch::extension;
+using namespace ::torch::executor;
 
-std::vector<int> NSArrayToIntVector(NSArray *inputArray);
+template <typename T> T getValueFromNSNumber(NSNumber *number) {
+  if constexpr (std::is_same<T, int8_t>::value) {
+    return static_cast<T>([number charValue]); // `charValue` for 8-bit integers
+  } else if constexpr (std::is_same<T, int32_t>::value) {
+    return static_cast<T>([number intValue]); // `intValue` for 32-bit integers
+  } else if constexpr (std::is_same<T, int64_t>::value) {
+    return static_cast<T>(
+        [number longLongValue]); // `longLongValue` for 64-bit integers
+  } else if constexpr (std::is_same<T, float>::value) {
+    return static_cast<T>([number floatValue]);
+  } else if constexpr (std::is_same<T, double>::value) {
+    return static_cast<T>([number doubleValue]);
+  }
+}
 
 template <typename T>
-std::unique_ptr<T[]> NSArrayToTypedArray(NSArray *nsArray);
+std::unique_ptr<T[]> NSArrayToTypedArray(NSArray *nsArray) {
+  size_t arraySize = [nsArray count];
 
-template <typename T> T getValueFromNSNumber(NSNumber *number);
+  std::unique_ptr<T[]> typedArray(new T[arraySize]);
+
+  for (NSUInteger i = 0; i < arraySize; ++i) {
+    NSNumber *number = [nsArray objectAtIndex:i];
+    if ([number isKindOfClass:[NSNumber class]]) {
+      typedArray[i] = getValueFromNSNumber<T>(number);
+    } else {
+      typedArray[i] = T();
+    }
+  }
+  return typedArray;
+}
+
+template <typename T>
+NSArray *arrayToNSArray(const void *array, ssize_t numel) {
+  const T *typedArray = static_cast<const T *>(array);
+  NSMutableArray *nsArray = [NSMutableArray arrayWithCapacity:numel];
+
+  for (int i = 0; i < numel; ++i) {
+    [nsArray addObject:@(typedArray[i])];
+  }
+
+  return [nsArray copy];
+}
 
 template <typename T>
-const T*
-runForwardFromNSArray(NSArray *inputArray, ssize_t& numel, std::vector<int> shapes,
-                      std::unique_ptr<executorch::extension::Module> &model);
+NSArray *arrayToNSArray(const std::vector<std::span<const T>> &dataPtrVec) {
+    NSMutableArray *nsArray = [NSMutableArray array];
+    for (const auto &span : dataPtrVec) {
+        NSMutableArray *innerArray = [NSMutableArray arrayWithCapacity:span.size()];
+        for(auto x : span) {
+            [innerArray addObject:@(x)];
+        }
+        [nsArray addObject:[innerArray copy]];
+    }
+    return [nsArray copy];
+}
+
+std::vector<int> NSArrayToIntVector(NSArray *inputArray) {
+  std::vector<int> output;
+  for (NSUInteger i = 0; i < [inputArray count]; ++i) {
+    NSNumber *number = [inputArray objectAtIndex:i];
+    if (number) {
+      output.push_back([number intValue]);
+    } else {
+      output.push_back(0);
+    }
+  }
+  return output;
+}
+
+template <typename T>
+std::vector<std::span<const T>>
+runForwardFromNSArray(NSArray *inputArray, std::vector<int> shapes,
+                      std::unique_ptr<Module> &model) {
+  std::unique_ptr<T[]> inputPtr = NSArrayToTypedArray<T>(inputArray);
+
+  TensorPtr inputTensor = from_blob(inputPtr.get(), shapes);
+  Result result = model->forward(inputTensor);
+
+  if (result.ok()) {
+    std::vector<std::span<const T>> outputVec;
+
+    for (const auto &currentResult : *result) {
+      Tensor currentTensor = currentResult.toTensor();
+      std::span<const T> currentSpan(currentTensor.const_data_ptr<T>(), currentTensor.numel());
+      outputVec.push_back(std::move(currentSpan));
+    }
+    return outputVec;
+  }
 
+  @throw [NSException
+      exceptionWithName:@"forward_error"
+                 reason:[NSString stringWithFormat:@"%d", (int)result.error()]
+               userInfo:nil];
+}
 
 #endif // Utils_hpp