Let a Core ML model read and write bound pixel buffers in place (#24261)

Package.swift

@@ -48,6 +48,7 @@ let products = deliverables([
     "frameworks": [
       "Accelerate",
       "CoreML",
+      "CoreVideo",
     ],
     "libraries": [
       "sqlite3",

backends/apple/coreml/CMakeLists.txt

@@ -9,6 +9,7 @@ if(APPLE)
   find_library(FOUNDATION_FRAMEWORK Foundation REQUIRED)
   find_library(ACCELERATE_FRAMEWORK Accelerate REQUIRED)
   find_library(COREML_FRAMEWORK CoreML REQUIRED)
+  find_library(COREVIDEO_FRAMEWORK CoreVideo REQUIRED)
   find_library(SQLITE_LIBRARY sqlite3 REQUIRED)
 endif()
 
@@ -257,8 +258,9 @@ if(APPLE)
   endif()
 
   target_link_libraries(
-    coremldelegate PRIVATE ${ACCELERATE_FRAMEWORK} ${COREML_FRAMEWORK}
-                           ${FOUNDATION_FRAMEWORK} ${SQLITE_LIBRARY}
+    coremldelegate
+    PRIVATE ${ACCELERATE_FRAMEWORK} ${COREML_FRAMEWORK} ${COREVIDEO_FRAMEWORK}
+            ${FOUNDATION_FRAMEWORK} ${SQLITE_LIBRARY}
   )
   if(EXECUTORCH_BUILD_SHARED)
     set_target_properties(

backends/apple/coreml/runtime/delegate/ETCoreMLModel.h

@@ -6,6 +6,7 @@
 // Please refer to the license found in the LICENSE file in the root directory of the source tree.
 
 #import <CoreML/CoreML.h>
+#import <CoreVideo/CoreVideo.h>
 #import <vector>
 
 NS_ASSUME_NONNULL_BEGIN
@@ -83,6 +84,24 @@ __attribute__((objc_subclassing_restricted))
 
 - (BOOL)prewarmAndReturnError:(NSError* __autoreleasing*)error;
 
+/// Binds an IOSurface-backed `kCVPixelFormatType_OneComponent16Half` pixel buffer to this model, so that an input or
+/// output whose data pointer is the buffer's base address is passed to Core ML as the pixel buffer: Core ML reads or
+/// writes it in place at the buffer's row stride instead of copying it. Must be called before the first prediction.
+///
+/// @param pixelBuffer The pixel buffer, retained for the model's lifetime.
+/// @param error On failure, error is filled with the failure information.
+/// @retval `YES` if the buffer was bound, `NO` if it is not an IOSurface-backed fp16 buffer, the OS is older than
+/// macOS 12 / iOS 16, locking fails, or the buffer has no base address.
+- (BOOL)bindPixelBuffer:(CVPixelBufferRef)pixelBuffer error:(NSError* __autoreleasing*)error;
+
+/// Sets `multiArray` to an array backed by the pixel buffer bound at `arg`'s data pointer, or to `nil` if no buffer is
+/// bound there. While `arg`'s shape stays the same, every call returns the same array.
+///
+/// @retval `NO` if the bound pixel buffer doesn't have `arg`'s shape and data type.
+- (BOOL)getPixelBufferMultiArray:(MLMultiArray* _Nullable __autoreleasing* _Nonnull)multiArray
+                          forArg:(const executorchcoreml::MultiArray&)arg
+                           error:(NSError* __autoreleasing*)error;
+
 @end
 
 NS_ASSUME_NONNULL_END

backends/apple/coreml/runtime/delegate/ETCoreMLModel.mm

@@ -172,8 +172,20 @@ void reset_state_for_feature_name(NSString *feature_name, MLState *state) {
 
 }
 
+namespace {
+/// A pixel buffer bound to the model, with the array Core ML reads or writes it through.
+struct BoundPixelBuffer {
+    void *baseAddress = nullptr;
+    id pixelBuffer = nil; // CVPixelBufferRef
+    MLMultiArray *multiArray = nil;
+    std::vector<size_t> multiArrayShape;
+};
+} // namespace
+
 #pragma mark - ETCoreMLModel
-@interface ETCoreMLModel ()
+@interface ETCoreMLModel () {
+    std::vector<BoundPixelBuffer> _boundPixelBuffers;
+}
 
 @property (strong, readonly, nonatomic) NSCache<ETCoreMLMultiArrayDescriptor *, NSMutableData *> *cache;
 @property (copy, readonly, nonatomic) NSDictionary<NSString *, MLMultiArrayConstraint *> *inputConstraintsByName;
@@ -318,8 +330,24 @@ - (nullable instancetype)initWithCompiledModelURL:(NSURL *)compiledModelURL
 
         const auto& layout = arg.layout();
         auto dataType = to_ml_multiarray_data_type(layout.dataType());
+        MLMultiArray *pixelBufferArray = nil;
+        if (![self getPixelBufferMultiArray:&pixelBufferArray forArg:arg error:error]) {
+            return nil;
+        }
+
         MLMultiArray *multiArrayArg = nil;
-        if (dataType == constraint.dataType) {
+        if (pixelBufferArray) {
+            if (pixelBufferArray.dataType != constraint.dataType) {
+                ETCoreMLLogErrorAndSetNSError(error,
+                                              ETCoreMLErrorInternalError,
+                                              "The pixel buffer bound to arg '%@' does not match the model feature's data type.",
+                                              argName);
+                return nil;
+            }
+            multiArrayArg = pixelBufferArray;
+            // Core ML reads and writes the pixel buffer at its own row stride, not the arg's strides.
+            lCopyData = NO;
+        } else if (dataType == constraint.dataType) {
             // We can use the same data storage.
             multiArrayArg = [[MLMultiArray alloc] initWithDataPointer:arg.data()
                                                                 shape:to_array(layout.shape())
@@ -430,4 +458,83 @@ - (BOOL)prewarmAndReturnError:(NSError* __autoreleasing*)error {
     return result;
 }
 
+- (BOOL)bindPixelBuffer:(CVPixelBufferRef)pixelBuffer error:(NSError * __autoreleasing *)error {
+    BOOL bindable = NO;
+    if (@available(macOS 12.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)) {
+        bindable = CVPixelBufferGetIOSurface(pixelBuffer) != NULL &&
+            CVPixelBufferGetPixelFormatType(pixelBuffer) == kCVPixelFormatType_OneComponent16Half;
+    }
+    if (!bindable) {
+        ETCoreMLLogErrorAndSetNSError(error,
+                                      ETCoreMLErrorInternalError,
+                                      "Only an IOSurface-backed fp16 pixel buffer can be bound to a model, on macOS 12 / iOS 16 or later.");
+        return NO;
+    }
+
+    const CVReturn lockStatus = CVPixelBufferLockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
+    if (lockStatus != kCVReturnSuccess) {
+        ETCoreMLLogErrorAndSetNSError(error,
+                                      ETCoreMLErrorInternalError,
+                                      "Failed to lock the pixel buffer for binding (status=%d).",
+                                      static_cast<int>(lockStatus));
+        return NO;
+    }
+    void *baseAddress = CVPixelBufferGetBaseAddress(pixelBuffer);
+    CVPixelBufferUnlockBaseAddress(pixelBuffer, kCVPixelBufferLock_ReadOnly);
+    if (baseAddress == nullptr) {
+        ETCoreMLLogErrorAndSetNSError(error,
+                                      ETCoreMLErrorInternalError,
+                                      "The pixel buffer has no base address and cannot be bound.");
+        return NO;
+    }
+
+    BoundPixelBuffer bound;
+    bound.baseAddress = baseAddress;
+    bound.pixelBuffer = (__bridge id)pixelBuffer;
+    for (auto& existing : _boundPixelBuffers) {
+        if (existing.baseAddress == bound.baseAddress) {
+            existing = std::move(bound);
+            return YES;
+        }
+    }
+    _boundPixelBuffers.push_back(std::move(bound));
+    return YES;
+}
+
+- (BOOL)getPixelBufferMultiArray:(MLMultiArray * _Nullable __autoreleasing *)multiArray
+                          forArg:(const executorchcoreml::MultiArray&)arg
+                           error:(NSError * __autoreleasing *)error {
+    *multiArray = nil;
+    for (auto& bound : _boundPixelBuffers) {
+        if (bound.baseAddress != arg.data()) {
+            continue;
+        }
+
+        // Wrapping the buffer costs Core ML a setup per array, so the array is made once per shape.
+        const auto& shape = arg.layout().shape();
+        if (bound.multiArray == nil || bound.multiArrayShape != shape) {
+            bound.multiArray = nil;
+            bound.multiArrayShape = shape;
+            CVPixelBufferRef buffer = (__bridge CVPixelBufferRef)bound.pixelBuffer;
+            if (@available(macOS 12.0, iOS 16.0, tvOS 16.0, watchOS 9.0, *)) {
+                if (!shape.empty() && shape.back() == CVPixelBufferGetWidth(buffer) &&
+                    std::accumulate(shape.cbegin(), shape.cend() - 1, size_t(1), std::multiplies<size_t>{}) == CVPixelBufferGetHeight(buffer)) {
+                    bound.multiArray = [[MLMultiArray alloc] initWithPixelBuffer:buffer shape:to_array(shape)];
+                }
+            }
+        }
+        if (bound.multiArray == nil || arg.layout().dataType() != MultiArray::DataType::Float16) {
+            ETCoreMLLogErrorAndSetNSError(error,
+                                          ETCoreMLErrorInternalError,
+                                          "The pixel buffer bound at an arg's data pointer doesn't match the arg's shape and data type.");
+            return NO;
+        }
+
+        *multiArray = bound.multiArray;
+        return YES;
+    }
+
+    return YES;
+}
+
 @end

backends/apple/coreml/runtime/delegate/ETCoreMLModelManager.mm

@@ -154,7 +154,24 @@ void set_outputs(NSArray<MLMultiArray *> *outputs, NSArray<MLMultiArray *> *mode
     }
 }
 
-void copy(MLMultiArray *src, executorchcoreml::MultiArray& dst) {
+BOOL copy(ETCoreMLModel *model,
+          MLMultiArray *src,
+          executorchcoreml::MultiArray& dst,
+          NSError * __autoreleasing *error) {
+    // When Core ML didn't write into a bound pixel buffer itself, the copy has to follow the
+    // buffer's row stride rather than dst's strides.
+    MLMultiArray *pixel_buffer_dst = nil;
+    if (![model getPixelBufferMultiArray:&pixel_buffer_dst forArg:dst error:error]) {
+        return NO;
+    }
+    if (pixel_buffer_dst) {
+        // Core ML returns the output backing itself when it wrote into it.
+        if (src != pixel_buffer_dst) {
+            ::copy(src, pixel_buffer_dst);
+        }
+        return YES;
+    }
+
     void (^copy_data)(void *) = ^(void *bytes) {
         if (bytes == dst.data()) {
             return;
@@ -174,15 +191,23 @@ void copy(MLMultiArray *src, executorchcoreml::MultiArray& dst) {
     } else {
         copy_data(src.dataPointer);
     }
+
+    return YES;
 }
 
-void set_outputs(std::vector<executorchcoreml::MultiArray>& outputs,
-                 NSArray<MLMultiArray *> *model_outputs) {
+BOOL set_outputs(ETCoreMLModel *model,
+                 std::vector<executorchcoreml::MultiArray>& outputs,
+                 NSArray<MLMultiArray *> *model_outputs,
+                 NSError * __autoreleasing *error) {
     NSEnumerator<MLMultiArray *> *enumerator = [model_outputs objectEnumerator];
     for (auto& output : outputs) {
         MLMultiArray *model_output = [enumerator nextObject];
-        ::copy(model_output, output);
+        if (!::copy(model, model_output, output, error)) {
+            return NO;
+        }
     }
+
+    return YES;
 }
 
 NSData * _Nullable get_file_data(const inmemoryfs::InMemoryFileSystem *inMemoryFS,
@@ -1176,8 +1201,7 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle
                         outputArgs[i].resize(new_size);
                         argsVec[model.orderedInputNames.count + i].resize(new_size);
                     }
-                    ::set_outputs(outputArgs, modelOutputs);
-                    result = YES;
+                    result = ::set_outputs(model, outputArgs, modelOutputs, &localError);
                 }
             }
         }