[RUNTIME][CONTRIB] CoreML Runtime (#5283)

* [RUNTIME][CONTRIB] CoreML Runtime

* fix lint

* fix CI

* use xcrun to compile coreml model

CMakeLists.txt

@@ -69,6 +69,7 @@ tvm_option(USE_ANTLR "Build with ANTLR for Relay parsing" OFF)
 tvm_option(USE_CPP_RPC "Build CPP RPC" OFF)
 tvm_option(USE_TFLITE "Build with tflite support" OFF)
 tvm_option(USE_TENSORFLOW_PATH "TensorFlow root path when use TFLite" none)
+tvm_option(USE_COREML "Build with coreml support" OFF)
 
 if(USE_CPP_RPC AND UNIX)
   message(FATAL_ERROR "USE_CPP_RPC is only supported with WIN32. Use the Makefile for non-Windows.")
@@ -301,6 +302,7 @@ include(cmake/modules/contrib/NNPack.cmake)
 include(cmake/modules/contrib/HybridDump.cmake)
 include(cmake/modules/contrib/TFLite.cmake)
 include(cmake/modules/contrib/TF_TVMDSOOP.cmake)
+include(cmake/modules/contrib/CoreML.cmake)
 
 if(NOT MSVC)
   include(CheckCXXCompilerFlag)

apps/ios_rpc/tvmrpc.xcodeproj/project.pbxproj

@@ -249,7 +249,7 @@
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 			shellPath = /bin/sh;
-			shellScript = "libpath=${CONFIGURATION_BUILD_DIR}/${CONTENTS_FOLDER_PATH}/Frameworks/tvm\nmkdir -p ${libpath}\nrm -rf ${libpath}/*\n                  \nif [ -f ${SRCROOT}/rpc_config.txt ]; then\n    head -n 1 ${SRCROOT}/rpc_config.txt > ${libpath}/rpc_config.txt\n    tail -n +2 ${SRCROOT}/rpc_config.txt | xargs -J % cp % ${libpath}\nfi\n\n";
+			shellScript = "libpath=${CONFIGURATION_BUILD_DIR}/${CONTENTS_FOLDER_PATH}/Frameworks/tvm\nmkdir -p ${libpath}\nrm -rf ${libpath}/*\n                  \nif [ -f ${SRCROOT}/rpc_config.txt ]; then\n    head -n 1 ${SRCROOT}/rpc_config.txt > ${libpath}/rpc_config.txt\n    tail -n +2 ${SRCROOT}/rpc_config.txt | xargs -J % cp -r % ${libpath}\nfi\n\n";
 		};
 /* End PBXShellScriptBuildPhase section */
 

apps/ios_rpc/tvmrpc/TVMRuntime.mm

@@ -46,6 +46,8 @@
 // Metal
 #include "../../../src/runtime/metal/metal_module.mm"
 #include "../../../src/runtime/metal/metal_device_api.mm"
+// CoreML
+#include "../../../src/runtime/contrib/coreml/coreml_runtime.mm"
 
 namespace dmlc {
 // Override logging mechanism

cmake/modules/contrib/CoreML.cmake

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+if(USE_COREML)
+  message(STATUS "Build with contrib.coreml")
+  find_library(FOUNDATION_LIB Foundation)
+  find_library(COREML_LIB Coreml)
+  file(GLOB COREML_CONTRIB_SRC src/runtime/contrib/coreml/*.mm)
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${FOUNDATION_LIB} ${COREML_LIB})
+  list(APPEND RUNTIME_SRCS ${COREML_CONTRIB_SRC})
+endif(USE_COREML)

python/tvm/contrib/coreml_runtime.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""CoreML runtime that load and run coreml models."""
+import tvm._ffi
+from ..rpc import base as rpc_base
+
+def create(compiled_model_path, output_names, ctx):
+    """Create a runtime executor module given a coreml model and context.
+    Parameters
+    ----------
+    compiled_model_path : str
+        The path of the compiled model to be deployed.
+    output_names : list of str
+        The output names of the model.
+    ctx : TVMContext
+        The context to deploy the module. It can be local or remote when there
+        is only one TVMContext.
+    Returns
+    -------
+    coreml_runtime : CoreMLModule
+        Runtime coreml module that can be used to execute the coreml model.
+    """
+    device_type = ctx.device_type
+    runtime_func = "tvm.coreml_runtime.create"
+
+    if device_type >= rpc_base.RPC_SESS_MASK:
+        fcreate = ctx._rpc_sess.get_function(runtime_func)
+    else:
+        fcreate = tvm._ffi.get_global_func(runtime_func)
+
+    return CoreMLModule(fcreate(compiled_model_path, ctx, *output_names))
+
+
+class CoreMLModule(object):
+    """Wrapper runtime module.
+
+    This is a thin wrapper of the underlying TVM module.
+    you can also directly call set_input, run, and get_output
+    of underlying module functions
+
+    Parameters
+    ----------
+    module : Module
+        The internal tvm module that holds the actual coreml functions.
+
+    Attributes
+    ----------
+    module : Module
+        The internal tvm module that holds the actual coreml functions.
+    """
+
+    def __init__(self, module):
+        self.module = module
+        self.invoke = module["invoke"]
+        self.set_input = module["set_input"]
+        self.get_output = module["get_output"]
+        self.get_num_outputs = module["get_num_outputs"]

python/tvm/contrib/xcode.py

@@ -170,6 +170,17 @@ def compile_metal(code, path_target=None, sdk="macosx"):
     return libbin
 
 
+def compile_coreml(model, out_dir="."):
+    """Compile coreml model and return the compiled model path.
+    """
+    mlmodel_path = os.path.join(out_dir, "tmp.mlmodel")
+    model.save(mlmodel_path)
+
+    xcrun(["coremlcompiler", "compile", mlmodel_path, out_dir])
+
+    return os.path.join(out_dir, "tmp.mlmodelc")
+
+
 class XCodeRPCServer(object):
     """Wrapper for RPC server
 

src/runtime/contrib/coreml/coreml_runtime.h

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \brief CoreML runtime that can run coreml model
+ *        containing only tvm PackedFunc.
+ * \file coreml_runtime.h
+ */
+#ifndef TVM_RUNTIME_CONTRIB_COREML_COREML_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_COREML_COREML_RUNTIME_H_
+
+#import <Foundation/Foundation.h>
+#import <CoreML/CoreML.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/runtime/ndarray.h>
+#include <tvm/runtime/packed_func.h>
+
+#include <vector>
+#include <string>
+#include <memory>
+
+namespace tvm {
+namespace runtime {
+
+/*!
+ * \brief CoreML runtime.
+ *
+ *  This runtime can be accessed in various language via
+ *  TVM runtime PackedFunc API.
+ */
+class CoreMLRuntime : public ModuleNode {
+ public:
+  /*!
+   * \brief Get member function to front-end.
+   * \param name The name of the function.
+   * \param sptr_to_self The pointer to the module node.
+   * \return The corresponding member function.
+   */
+  virtual PackedFunc GetFunction(const std::string& name,
+                                 const ObjectPtr<Object>& sptr_to_self);
+
+  /*!
+   * \return The type key of the executor.
+   */
+  const char* type_key() const {
+    return "CoreMLRuntime";
+  }
+
+  /*!
+   * \brief Invoke the coreml prediction.
+   */
+  void Invoke();
+
+  /*!
+   * \brief Initialize the coreml runtime with coreml model and context.
+   * \param model_path The compiled model path.
+   * \param ctx The context where the coreml model will be executed on.
+   * \param output_names The output names of the model.
+   */
+  void Init(const std::string& model_path,
+            TVMContext ctx,
+            const std::vector<NSString *>& output_names);
+
+  /*!
+   * \brief set input to the model.
+   * \param key The input name.
+   * \param data_in The input data.
+   */
+  void SetInput(const std::string& key, DLTensor* data_in);
+  /*!
+   * \brief Return NDArray for given output index.
+   * \param index The output index.
+   *
+   * \return NDArray corresponding to given output node index.
+   */
+  NDArray GetOutput(int index) const;
+  /*!
+   * \brief Return the number of outputs
+   *
+   * \return The number of outputs
+   */
+  int GetNumOutputs() const;
+
+  // CoreML model
+  MLModel *model_;
+  // CoreML model input dictionary
+  NSMutableDictionary<NSString *, id> *input_dict_;
+  // CoreML model output
+  id<MLFeatureProvider> output_;
+  // List of output names
+  std::vector<NSString *> output_names_;
+  // TVM context
+  TVMContext ctx_;
+};
+
+}  // namespace runtime
+}  // namespace tvm
+
+#endif  // TVM_RUNTIME_CONTRIB_COREML_COREML_RUNTIME_H_

src/runtime/contrib/coreml/coreml_runtime.mm

+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file coreml_runtime.cc
+ */
+#include <tvm/runtime/registry.h>
+
+#include "coreml_runtime.h"
+
+namespace tvm {
+namespace runtime {
+
+MLModel *load_coreml_model(const std::string& model_path) {
+  NSBundle* bundle = [NSBundle mainBundle];
+  NSString* base = [bundle privateFrameworksPath];
+  NSString* fname = [NSString stringWithUTF8String:("tvm/" + model_path).c_str()];
+  NSString* assetPath = [base stringByAppendingPathComponent: fname];
+
+  if (![[NSFileManager defaultManager] fileExistsAtPath:assetPath]) {
+    assetPath = [NSString stringWithCString: model_path.c_str() encoding:NSUTF8StringEncoding];
+  }
+
+  NSURL *url = [NSURL fileURLWithPath:assetPath];
+
+  MLModel *model = [MLModel modelWithContentsOfURL:url error:nil];
+  if (model == nil) {
+    NSLog(@"modelc %@ not found", url);
+  }
+  return model;
+}
+
+void CoreMLRuntime::Init(const std::string& model_path,
+                         TVMContext ctx,
+                         const std::vector<NSString *>& output_names) {
+  model_ = load_coreml_model(model_path);
+  ctx_ = ctx;
+  input_dict_ = [NSMutableDictionary dictionary];
+  output_names_ = output_names;
+}
+
+void CoreMLRuntime::Invoke() {
+  id<MLFeatureProvider> input = [[MLDictionaryFeatureProvider alloc] initWithDictionary:input_dict_ error:nil];
+  output_ = [model_ predictionFromFeatures:input error:nil];
+}
+
+void CoreMLRuntime::SetInput(const std::string& key, DLTensor* data_in) {
+  int64_t size = 1;
+  NSMutableArray *shape = [[NSMutableArray alloc] init];
+  for (int64_t i = 0; i < data_in->ndim; ++i) {
+    size *= data_in->shape[i];
+    [shape addObject:[NSNumber numberWithInteger:data_in->shape[i]]];
+  }
+
+  DataType dtype(data_in->dtype);
+  MLMultiArrayDataType dataType;
+  if (dtype == DataType::Float(64)) {
+    dataType = MLMultiArrayDataTypeDouble;
+    size *= sizeof(double);
+  } else if (dtype == DataType::Float(32)) {
+    dataType = MLMultiArrayDataTypeFloat32;
+    size *= sizeof(float);
+  } else {
+    LOG(FATAL) << "unsupported data type " << dtype;
+    return;
+  }
+
+  MLMultiArray *dest = [[MLMultiArray alloc] initWithShape:shape
+                        dataType:dataType error:nil];
+
+  CHECK(data_in->strides == NULL);
+  memcpy(dest.dataPointer, data_in->data, size);
+
+  NSString *nsKey = [NSString stringWithUTF8String:key.c_str()];
+  [input_dict_ setObject:dest forKey:nsKey];
+}
+
+NDArray CoreMLRuntime::GetOutput(int index) const {
+  NSString *name = output_names_[index];
+  MLModelDescription *model_desc = model_.modelDescription;
+  MLFeatureDescription *output_desc = model_desc.outputDescriptionsByName[name];
+  MLMultiArrayConstraint *data_desc = output_desc.multiArrayConstraint;
+  std::vector<int64_t> shape;
+  int64_t size = 1;
+  for (int64_t i = 0; i < data_desc.shape.count; ++i) {
+    int n = data_desc.shape[i].intValue;
+    size *= n;
+    shape.push_back(n);
+  }
+
+  DataType dtype;
+  if (data_desc.dataType == MLMultiArrayDataTypeDouble) {
+    dtype = DataType::Float(64);
+    size *= sizeof(double);
+  } else if (data_desc.dataType == MLMultiArrayDataTypeFloat32) {
+    dtype = DataType::Float(32);
+    size *= sizeof(float);
+  } else {
+    LOG(FATAL) << "unexpected data type " << data_desc.dataType;
+  }
+  MLMultiArray *src = [output_ featureValueForName:name].multiArrayValue;
+  NDArray ret = NDArray::Empty(shape, dtype, ctx_);
+  ret.CopyFromBytes(src.dataPointer, size);
+
+  return ret;
+}
+
+int CoreMLRuntime::GetNumOutputs() const {
+  return output_names_.size();
+}
+
+PackedFunc CoreMLRuntime::GetFunction(
+    const std::string& name,
+    const ObjectPtr<Object>& sptr_to_self) {
+  // Return member functions during query.
+  if (name == "invoke") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        this->Invoke();
+      });
+  } else if (name == "set_input") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        const auto& input_name = args[0].operator std::string();
+        this->SetInput(input_name, args[1]);
+      });
+  } else if (name == "get_output") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        *rv = this->GetOutput(args[0]);
+      });
+  } else if (name == "get_num_outputs") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        *rv = this->GetNumOutputs();
+      });
+  } else {
+    return PackedFunc();
+  }
+}
+
+Module CoreMLRuntimeCreate(const std::string& model_path,
+                           TVMContext ctx,
+                           const std::vector<NSString *>& output_names) {
+  auto exec = make_object<CoreMLRuntime>();
+  exec->Init(model_path, ctx, output_names);
+  return Module(exec);
+}
+
+TVM_REGISTER_GLOBAL("tvm.coreml_runtime.create")
+  .set_body([](TVMArgs args, TVMRetValue* rv) {
+      std::vector<NSString *> output_names;
+      for (size_t i = 2; i < args.size(); i++) {
+        const std::string& name = args[i];
+        output_names.push_back([NSString stringWithUTF8String:name.c_str()]);
+      }
+      *rv = CoreMLRuntimeCreate(args[0], args[1], output_names);
+  });
+}  // namespace runtime
+}  // namespace tvm

tests/python/contrib/test_coreml_runtime.py

+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import tvm
+from tvm import te
+import numpy as np
+from tvm import rpc
+from tvm.contrib import util, xcode, coreml_runtime
+
+import os
+
+proxy_host = os.environ.get("TVM_IOS_RPC_PROXY_HOST", "localhost")
+proxy_port = os.environ.get("TVM_IOS_RPC_PROXY_PORT", 9090)
+destination = os.environ.get("TVM_IOS_RPC_DESTINATION", "")
+key = "iphone"
+
+def skipped_test_coreml_runtime():
+
+    import coremltools
+    from coremltools.models.neural_network import NeuralNetworkBuilder
+
+    def create_coreml_model():
+        shape = (2,)
+        alpha = 2
+
+        inputs = [
+            ('input0', coremltools.models.datatypes.Array(*shape)),
+            ('input1', coremltools.models.datatypes.Array(*shape))
+        ]
+        outputs = [
+            ('output0', coremltools.models.datatypes.Array(*shape)),
+            ('output1', coremltools.models.datatypes.Array(*shape)),
+        ]
+        builder = NeuralNetworkBuilder(inputs, outputs)
+        builder.add_elementwise(name='Add',
+                                input_names=['input0', 'input1'],
+                                output_name='output0',
+                                mode='ADD')
+        builder.add_elementwise(name='Mul',
+                                alpha=alpha,
+                                input_names=['input0'],
+                                output_name='output1',
+                                mode='MULTIPLY')
+        return coremltools.models.MLModel(builder.spec)
+
+    def verify(coreml_model, compiled_model_path, ctx):
+        coreml_model = create_coreml_model()
+
+        out_spec = coreml_model.output_description._fd_spec
+        out_names = [spec.name for spec in out_spec]
+
+        # inference via coremltools
+        inputs = {}
+        for in_spec in coreml_model.input_description._fd_spec:
+            name = in_spec.name
+            shape = in_spec.type.multiArrayType.shape
+            inputs[name] = np.random.random_sample(shape)
+
+        coreml_outputs = [coreml_model.predict(inputs)[name] for name in out_names]
+
+        # inference via tvm coreml runtime
+        runtime = coreml_runtime.create(compiled_model_path, out_names, ctx)
+        for name in inputs:
+            runtime.set_input(name, tvm.nd.array(inputs[name], ctx))
+        runtime.invoke()
+        tvm_outputs = [runtime.get_output(i).asnumpy() for i in range(runtime.get_num_outputs())]
+
+        for c_out, t_out in zip(coreml_outputs, tvm_outputs):
+            np.testing.assert_almost_equal(c_out, t_out, 3)
+
+    def check_remote(coreml_model):
+        temp = util.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        xcode.popen_test_rpc(proxy_host, proxy_port, key, destination=destination,
+                             libs=[compiled_model])
+        compiled_model = os.path.basename(compiled_model)
+        remote = rpc.connect(proxy_host, proxy_port, key=key)
+        ctx = remote.cpu(0)
+        verify(coreml_model, compiled_model, ctx)
+
+    def check_local(coreml_model):
+        temp = util.tempdir()
+        compiled_model = xcode.compile_coreml(coreml_model, out_dir=temp.temp_dir)
+        ctx = tvm.cpu(0)
+        verify(coreml_model, compiled_model, ctx)
+
+    coreml_model = create_coreml_model()
+    check_remote(coreml_model)
+    check_local(coreml_model)
+
+
+if __name__ == "__main__":
+    # skipped_test_coreml_runtime()
+    pass