[iOS][OSS][BE] Add simulator tests for Metal

Pull Request resolved: #64852 ghstack-source-id: 137970230 Differential Revision: [D30877961](https://our.internmc.facebook.com/intern/diff/D30877961/)
pytorch · Sep 13, 2021 · ab884c2 · ab884c2
1 parent 1b2f4f7
commit ab884c2
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 14 deletions.
diff --git a/.circleci/cimodel/data/simple/ios_definitions.py b/.circleci/cimodel/data/simple/ios_definitions.py
@@ -70,6 +70,9 @@ def gen_tree(self):
     IOSJob(XCODE_VERSION, ArchVariant("arm64", "metal"), extra_props={
         "use_metal": miniutils.quote(str(int(True))),
         "lite_interpreter": miniutils.quote(str(int(True)))}),
+    IOSJob(XCODE_VERSION, ArchVariant("x86_64", "metal"), extra_props={
+        "use_metal": miniutils.quote(str(int(True))),
+        "lite_interpreter": miniutils.quote(str(int(True)))}),
     IOSJob(XCODE_VERSION, ArchVariant("arm64", "full_jit"), extra_props={
         "lite_interpreter": miniutils.quote(str(int(False)))}),
     IOSJob(XCODE_VERSION, ArchVariant("arm64", "custom"), extra_props={

diff --git a/.circleci/config.yml b/.circleci/config.yml
diff --git a/.circleci/scripts/binary_ios_test.sh b/.circleci/scripts/binary_ios_test.sh
@@ -26,4 +26,4 @@ if ! [ -x "$(command -v xcodebuild)" ]; then
     exit 1
 fi
 PROFILE=PyTorch_CI_2021
-ruby ${PROJ_ROOT}/scripts/xcode_build.rb -i ${PROJ_ROOT}/build_ios/install -x ${PROJ_ROOT}/ios/TestApp/TestApp.xcodeproj -p ${IOS_PLATFORM} -c ${PROFILE} -t ${IOS_DEV_TEAM_ID} -f Accelerate,MetalPerformanceShaders,CoreML
+ruby ${PROJ_ROOT}/scripts/xcode_build.rb -i ${PROJ_ROOT}/build_ios/install -x ${PROJ_ROOT}/ios/TestApp/TestApp.xcodeproj -p ${IOS_PLATFORM} -c ${PROFILE} -t ${IOS_DEV_TEAM_ID}
diff --git a/.circleci/verbatim-sources/job-specs/job-specs-custom.yml b/.circleci/verbatim-sources/job-specs/job-specs-custom.yml
@@ -525,9 +525,8 @@
             fi
             export IOS_ARCH=${IOS_ARCH}
             export IOS_PLATFORM=${IOS_PLATFORM}
-            if [ ${IOS_PLATFORM} != "SIMULATOR" ]; then
-              export USE_PYTORCH_METAL=${USE_METAL}
-            fi
+            export USE_PYTORCH_METAL=${USE_METAL}
+            
             unbuffer ${PROJ_ROOT}/scripts/build_ios.sh 2>&1 | ts
       - run:
           name: Run Build Test
@@ -577,7 +576,11 @@
             cd ${PROJ_ROOT}/ios/TestApp
             instruments -s -devices
             if [ ${BUILD_LITE_INTERPRETER} == 1 ]; then
-              fastlane scan --only_testing TestAppTests/TestAppTests/testLiteInterpreter
+              if [ ${USE_PYTORCH_METAL} == 1 ]; then
+                fastlane scan --only_testing TestAppTests/TestAppTests/testMetal
+              else
+                fastlane scan --only_testing TestAppTests/TestAppTests/testLiteInterpreter
+              fi
             else
               fastlane scan --only_testing TestAppTests/TestAppTests/testFullJIT
             fi

diff --git a/ios/TestApp/TestAppTests/TestLiteInterpreter.mm b/ios/TestApp/TestAppTests/TestLiteInterpreter.mm
@@ -21,4 +21,14 @@ - (void)testLiteInterpreter {
   XCTAssertTrue(outputTensor.numel() == 1000);
 }
 
+- (void)testMetal {
+  NSString* modelPath = [[NSBundle bundleForClass:[self class]] pathForResource:@"model_metal"
+                                                                         ofType:@"ptl"];
+  auto module = torch::jit::_load_for_mobile(modelPath.UTF8String);
+  c10::InferenceMode mode;
+  auto input = torch::ones({1, 3, 224, 224}, at::kFloat).metal();
+  auto outputTensor = module.forward({input}).toTensor().cpu();
+  XCTAssertTrue(outputTensor.numel() == 1000);
+}
+
 @end
diff --git a/ios/TestApp/benchmark/setup.rb b/ios/TestApp/benchmark/setup.rb
@@ -93,6 +93,7 @@
 
 puts "Linking static libraries..."
 libs = ['libc10.a', 'libclog.a', 'libpthreadpool.a', 'libXNNPACK.a', 'libeigen_blas.a', 'libcpuinfo.a', 'libpytorch_qnnpack.a', 'libtorch_cpu.a', 'libtorch.a']
+frameworks = ['CoreML', 'Metal', 'MetalPerformanceShaders', 'Accelerate', 'UIKit']
 targets.each do |target|
     target.frameworks_build_phases.clear
     for lib in libs do
@@ -102,6 +103,14 @@
             target.frameworks_build_phases.add_file_reference(libref)
         end
     end
+    # link system frameworks
+    frameworks.each do |framework|
+        path = "System/Library/Frameworks/#{framework}.framework"
+        framework_ref = project.frameworks_group.new_reference(path)
+        framework_ref.name = "#{framework}.framework"
+        framework_ref.source_tree = 'SDKROOT'
+        target.frameworks_build_phases.add_file_reference(framework_ref)
+    end
 end
 
 project.save

diff --git a/ios/TestApp/benchmark/trace_model.py b/ios/TestApp/benchmark/trace_model.py
@@ -8,4 +8,6 @@
 traced_script_module = torch.jit.trace(model, example)
 optimized_scripted_module = optimize_for_mobile(traced_script_module)
 torch.jit.save(optimized_scripted_module, '../models/model.pt')
-exported_optimized_scripted_module = optimized_scripted_module._save_for_lite_interpreter("../models/model_lite.ptl")
+optimized_scripted_module._save_for_lite_interpreter("../models/model_lite.ptl")
+optimized_scripted_module = optimize_for_mobile(traced_script_module, backend='metal')
+optimized_scripted_module._save_for_lite_interpreter("../models/model_metal.ptl")
diff --git a/ios/TestApp/fastlane/Scanfile b/ios/TestApp/fastlane/Scanfile
@@ -4,3 +4,4 @@ clean(true)
 suppress_xcode_output(true)
 force_quit_simulator(true)
 include_simulator_logs(false)
+deployment_target_version("14.0")
diff --git a/scripts/xcode_build.rb b/scripts/xcode_build.rb
@@ -19,9 +19,6 @@
  opts.on('-t', '--team_id ', 'development team ID') { |value|
     options[:team_id] = value
  }
- opts.on('-f', '--framework ', 'system frameworks') { |value|
-    options[:framework] = value.split(',')
-}
 end.parse!
 puts options.inspect
 
@@ -63,7 +60,7 @@
     end
 end
 # link system frameworks
-frameworks = options[:framework]
+frameworks = ['CoreML', 'Metal', 'MetalPerformanceShaders', 'Accelerate', 'UIKit']
 if frameworks
     frameworks.each do |framework|
         path = "System/Library/Frameworks/#{framework}.framework"