Use eGPU if available for CoreML models in image feature extraction (#3063)

Author: nickjong

src/toolkits/image_deep_feature_extractor/CMakeLists.txt

@@ -13,6 +13,13 @@ if(HAS_CORE_ML)
   find_library(core_video NAMES CoreVideo)
   message("CoreVideo found at ${core_video}.")
 
+  if(HAS_MACOS_10_15 AND NOT TC_BUILD_IOS)
+    # On macOS 10.15+, we can use TCMPSDeviceManager to pass eGPUs to CoreML.
+    set(additional_image_deep_feature_extractor_requirements tcmps)
+  else()
+    set(additional_image_deep_feature_extractor_requirements "")
+  endif()
+
   make_library(image_deep_feature_extractor OBJECT
     SOURCES
       mlmodel_image_feature_extractor.mm
@@ -25,6 +32,7 @@ if(HAS_CORE_ML)
       ${CORE_ML}
       ${FOUNDATION}
       ${CORE_VIDEO}
+      ${additional_image_deep_feature_extractor_requirements}
     EXTERNAL_VISIBILITY
   )
 

src/toolkits/image_deep_feature_extractor/mlmodel_image_feature_extractor.mm

@@ -20,6 +20,10 @@
 
 #include <coremltools/mlmodel/src/Format.hpp>
 
+#ifdef HAS_MACOS_10_15
+#import <ml/neural_net/mps_device_manager.h>
+#endif
+
 namespace turi {
 namespace image_deep_feature_extractor {
 
@@ -158,6 +162,27 @@ void build_vision_feature_print_scene_spec(const std::string& model_path) {
 
 }
 
+API_AVAILABLE(macos(10.13), ios(11.0))
+MLModel* create_model(NSURL* url, NSError* _Nullable* error) {
+#if defined(HAS_MACOS_10_15) && !defined(TC_BUILD_IOS)
+  if (@available(macos 10.15, ios 10.13, *)) {
+    MLModelConfiguration* config = [[MLModelConfiguration alloc] init];
+    config.preferredMetalDevice = [TCMPSDeviceManager sharedInstance].preferredDevice;
+
+    if (config.preferredMetalDevice != nil) {
+      logprogress_stream << "Using GPU (" << config.preferredMetalDevice.name.UTF8String
+                         << ") to extract features.";
+    } else {
+      // Assume that CoreML will fall back to CPU if no Metal devices are available.
+      logprogress_stream << "Using CPU to extract features.";
+    }
+
+    return [MLModel modelWithContentsOfURL:url configuration:config error:error];
+  }
+#endif  // defined(HAS_MACOS_10_15) && !defined(TC_BUILD_IOS)
+  return [MLModel modelWithContentsOfURL:url error:error];
+}
+
 API_AVAILABLE(macos(10.13),ios(11.0))
 static MLModel *create_model(const std::string& download_path,
 			     const std::string& model_name) {
@@ -173,7 +198,7 @@ void build_vision_feature_print_scene_spec(const std::string& model_path) {
   if (boost::filesystem::exists(compiled_modified_model_path)) {
 
     NSError* error = nil;
-    result = [MLModel modelWithContentsOfURL:compiledModelURL error:&error];
+    result = create_model(compiledModelURL, &error);
 
     if (error || !result) {
 
@@ -240,7 +265,7 @@ void build_vision_feature_print_scene_spec(const std::string& model_path) {
 
     // Load the compiled modified model
     NSError* error = nil;
-    result = [MLModel modelWithContentsOfURL:compiledModelURL error:&error];
+    result = create_model(compiledModelURL, &error);
     checkNSError(error);
   }