Add dynamic shape support to CoreML

Differential Revision: D70904915

Pull Request resolved: pytorch#9094

Author: metascroy

CMakeLists.txt

@@ -922,6 +922,10 @@ if(EXECUTORCH_BUILD_EXECUTOR_RUNNER)
     endif()
   endif()
 
+  if(EXECUTORCH_BUILD_COREML)
+    list(APPEND _executor_runner_libs coremldelegate)
+  endif()
+
   add_executable(executor_runner ${_executor_runner__srcs})
   if(CMAKE_BUILD_TYPE STREQUAL "Release")
     if(APPLE)

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

@@ -71,7 +71,7 @@ __attribute__((objc_subclassing_restricted))
 /// @param error   On failure, error is filled with the failure information.
 /// @retval `YES` if the execution succeeded otherwise `NO`.
 - (BOOL)executeModelWithHandle:(ModelHandle*)handle
-                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                       argsVec:(std::vector<executorchcoreml::MultiArray>&)argsVec
                 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
                    eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
                          error:(NSError* __autoreleasing*)error;

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

@@ -734,7 +734,7 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle
 }
 
 - (BOOL)executeModelWithHandle:(ModelHandle *)handle
-                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                       argsVec:(std::vector<executorchcoreml::MultiArray>&)argsVec
                 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
                    eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
                          error:(NSError * __autoreleasing *)error {
@@ -785,6 +785,12 @@ - (BOOL)executeModelWithHandle:(ModelHandle *)handle
             return NO;
         }
 
+        // Resize for dynamic shapes
+        for (int i = 0; i < outputArgs.size(); i++) {
+            auto new_size = to_vector<size_t>(modelOutputs[i].shape);
+            outputArgs[i].resize(new_size);
+            argsVec[model.orderedInputNames.count + i].resize(new_size);
+        }
         ::set_outputs(outputArgs, modelOutputs);
         return YES;
     }

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

@@ -89,7 +89,7 @@ class BackendDelegate {
     /// @param error   On failure, error is filled with the failure information.
     /// @retval `true` if the execution succeeded otherwise `false`.
     virtual bool execute(Handle* handle,
-                         const std::vector<MultiArray>& args,
+                         std::vector<MultiArray>& args,
                          const ModelLoggingOptions& logging_options,
                          ModelEventLogger* event_logger,
                          std::error_code& error) const noexcept = 0;

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

@@ -104,7 +104,7 @@ - (ModelHandle*)loadModelFromAOTData:(NSData*)data
                                error:(NSError* __autoreleasing*)error;
 
 - (BOOL)executeModelWithHandle:(ModelHandle*)handle
-                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                       argsVec:(std::vector<executorchcoreml::MultiArray>&)argsVec
                 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
                    eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
                          error:(NSError* __autoreleasing*)error;
@@ -199,7 +199,7 @@ - (ModelHandle*)loadModelFromAOTData:(NSData*)data
 }
 
 - (BOOL)executeModelWithHandle:(ModelHandle*)handle
-                       argsVec:(const std::vector<executorchcoreml::MultiArray>&)argsVec
+                       argsVec:(std::vector<executorchcoreml::MultiArray>&)argsVec
                 loggingOptions:(const executorchcoreml::ModelLoggingOptions&)loggingOptions
                    eventLogger:(const executorchcoreml::ModelEventLogger* _Nullable)eventLogger
                          error:(NSError* __autoreleasing*)error {
@@ -286,7 +286,7 @@ explicit BackendDelegateImpl(const Config& config) noexcept
     }
 
     bool execute(Handle* handle,
-                 const std::vector<MultiArray>& args,
+                 std::vector<MultiArray>& args,
                  const ModelLoggingOptions& logging_options,
                  ModelEventLogger *event_logger,
                  std::error_code& ec) const noexcept override {

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

@@ -12,13 +12,15 @@
 #import <coreml_backend/delegate.h>
 #import <executorch/runtime/core/evalue.h>
 #import <executorch/runtime/platform/log.h>
+#import <executorch/runtime/kernel/kernel_includes.h>
 #import <memory>
 #import <model_event_logger.h>
 #import <model_logging_options.h>
 #import <multiarray.h>
 #import <objc_safe_cast.h>
 #import <unordered_map>
 #import <vector>
+#include <array>
 
 #ifdef ET_EVENT_TRACER_ENABLED
 #import <model_event_logger_impl.h>
@@ -40,6 +42,9 @@
 using executorch::runtime::FreeableBuffer;
 using executorch::runtime::get_backend_class;
 using executorch::runtime::Result;
+using executorch::aten::SizesType;
+using executorch::aten::Tensor;
+using executorch::runtime::kTensorDimensionLimit;
 
 std::optional<MultiArray::DataType> get_data_type(ScalarType scalar_type) {
     switch (scalar_type) {
@@ -221,6 +226,21 @@ ModelLoggingOptions get_logging_options(BackendExecutionContext& context) {
                              ETCoreMLStrings.delegateIdentifier.UTF8String);
 #endif
 
+    // Resize for dynamic shape
+    std::array<SizesType, kTensorDimensionLimit> new_shape;
+    for (size_t i = nInputs; i < nInputs + nOutputs; i++) {
+        Tensor& t = args[i]->toTensor();
+        int rank = delegate_args[i].layout().rank();
+        assert (rank <= new_shape.size());
+        for (int d = 0; d < rank; d++) {
+            new_shape[d] = delegate_args[i].layout().shape()[d];
+        }
+        ET_CHECK_OR_RETURN_ERROR(
+            resize_tensor(t, ArrayRef(new_shape.data(), rank)) == Error::Ok,
+            DelegateInvalidHandle,
+            "%s: Failed to resize delegate output %zu",  ETCoreMLStrings.delegateIdentifier.UTF8String, i);
+    }
+
     return Error::Ok;
 }
 

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

@@ -84,6 +84,11 @@ class MultiArray final {
         /// Returns `true` if the memory layout is packed otherwise `false`.
         bool is_packed() const noexcept;
 
+        // Resizes memory layout
+        // New shape must be the same dimension and no larger than current shape in all dimensions
+        // New format is contiguous
+        void resize(const std::vector<size_t>& shape);
+
     private:
         DataType dataType_;
         std::vector<size_t> shape_;
@@ -126,6 +131,8 @@ class MultiArray final {
         *ptr = value;
     }
 
+    void resize(const std::vector<size_t>& shape) { layout_.resize(shape); }
+
 private:
     void* data(const std::vector<size_t>& indices) const noexcept;
 

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

@@ -512,6 +512,24 @@ ssize_t get_data_offset(size_t index, const std::vector<size_t>& shape, const st
 
 namespace executorchcoreml {
 
+void MultiArray::MemoryLayout::resize(const std::vector<size_t>& shape) {
+    assert(shape.size() == shape_.size());
+    for (int i = 0; i < shape.size(); ++i) {
+        assert (shape[i] >= 1);
+        assert(shape[i] <= shape_[i]);
+    }
+    int stride = 1;
+    for (int i = shape.size() - 1; i >= 0; --i) {
+        shape_[i] = shape[i];
+        strides_[i] = stride;
+        if (shape[i] > 1) {
+            stride *= shape[i];
+        }
+    }
+}
+
+
+
 size_t MultiArray::MemoryLayout::num_elements() const noexcept {
     if (shape_.size() == 0) {
         return 0;

backends/apple/coreml/runtime/test/BackendDelegateTests.mm

@@ -162,8 +162,9 @@ - (void)testAddModelExecution {
     MLMultiArray *output = [ETCoreMLTestUtils filledMultiArrayWithShape:inputs[0].shape dataType:inputs[0].dataType repeatedValue:@(0) error:&localError];
     NSArray<MLMultiArray *> *args = [inputs arrayByAddingObject:output];
     std::error_code errorCode;
+    auto argsVec = to_multiarrays(args);
     XCTAssertTrue(_delegate->execute(handle,
-                                     to_multiarrays(args),
+                                     argsVec,
                                      ModelLoggingOptions(),
                                      nullptr,
                                      errorCode));
@@ -187,8 +188,9 @@ - (void)testMulModelExecution {
     MLMultiArray *output = [ETCoreMLTestUtils filledMultiArrayWithShape:inputs[0].shape dataType:inputs[0].dataType repeatedValue:@(0) error:&localError];
     NSArray<MLMultiArray *> *args = [inputs arrayByAddingObject:output];
     std::error_code errorCode;
-    XCTAssertTrue(_delegate->execute(handle, 
-                                     to_multiarrays(args),
+    auto argsVec = to_multiarrays(args);
+    XCTAssertTrue(_delegate->execute(handle,
+                                     argsVec,
                                      ModelLoggingOptions(),
                                      nullptr,
                                      errorCode));

backends/apple/coreml/runtime/test/MultiArrayTests.mm

@@ -130,4 +130,20 @@ - (void)testNonAdjacentDataCopy {
     [self verifyDataCopyWithShape:shape srcStrides:srcStrides dstStrides:dstStrides];
 }
 
+- (void)testResize {
+    std::vector<size_t> shape = {3, 1, 2, 5};
+    std::vector<ssize_t> strides = {1*2*5, 2*5, 5, 1};
+    std::vector<uint8_t> storage;
+    std::vector<size_t> newShape = {3, 1, 1, 1};
+    
+    auto array = make_multi_array_and_fill<int>(shape, strides, storage);
+    for (size_t i = 0; i < array.layout().rank(); ++i) {
+        XCTAssertEqual(array.layout().shape()[i], shape[i]);
+    }
+    array.resize(newShape);
+    for (size_t i = 0; i < array.layout().rank(); ++i) {
+        XCTAssertEqual(array.layout().shape()[i], newShape[i]);
+    }
+}
+
 @end

backends/apple/coreml/runtime/workspace/executorchcoreml.xcodeproj/xcshareddata/xcschemes/executorchcoreml_tests.xcscheme

@@ -23,6 +23,11 @@
                BlueprintName = "executorchcoreml_tests"
                ReferencedContainer = "container:executorchcoreml.xcodeproj">
             </BuildableReference>
+            <SkippedTests>
+               <Test
+                  Identifier = "ETCoreMLModelDebuggerTests/testMV3ProgramDebugging">
+               </Test>
+            </SkippedTests>
          </TestableReference>
       </Testables>
    </TestAction>

examples/apple/coreml/scripts/export.py

@@ -76,6 +76,12 @@ def parse_args() -> argparse.ArgumentParser:
     parser.add_argument("--use_partitioner", action=argparse.BooleanOptionalAction)
     parser.add_argument("--generate_etrecord", action=argparse.BooleanOptionalAction)
     parser.add_argument("--save_processed_bytes", action=argparse.BooleanOptionalAction)
+    parser.add_argument(
+        "--dynamic_shapes",
+        action=argparse.BooleanOptionalAction,
+        required=False,
+        default=False,
+    )
 
     args = parser.parse_args()
     # pyre-fixme[7]: Expected `ArgumentParser` but got `Namespace`.
@@ -164,16 +170,20 @@ def main():
             f"Valid compute units are {valid_compute_units}."
         )
 
-    model, example_inputs, _, _ = EagerModelFactory.create_model(
+    model, example_inputs, _, dynamic_shapes = EagerModelFactory.create_model(
         *MODEL_NAME_TO_MODEL[args.model_name]
     )
+    if not args.dynamic_shapes:
+        dynamic_shapes = None
 
     compile_specs = generate_compile_specs_from_args(args)
     lowered_module = None
 
     if args.use_partitioner:
         model.eval()
-        exir_program_aten = torch.export.export(model, example_inputs, strict=True)
+        exir_program_aten = torch.export.export(
+            model, example_inputs, dynamic_shapes=dynamic_shapes, strict=True
+        )
 
         edge_program_manager = exir.to_edge(exir_program_aten)
         edge_copy = copy.deepcopy(edge_program_manager)