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

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

Hide file tree

Showing 5 changed files with 29 additions and 1 deletion.
diff --git a/.circleci/config.yml b/.circleci/config.yml
diff --git a/.circleci/verbatim-sources/job-specs/job-specs-custom.yml b/.circleci/verbatim-sources/job-specs/job-specs-custom.yml
@@ -574,6 +574,8 @@
             instruments -s -devices
             if [ $(BUILD_LITE_INTERPRETER) == 1 ]; then
               fastlane scan --only_testing TestAppTests/TestAppTests/testLiteInterpreter
+            elif [ $(USE_PYTORCH_METAL) == 1 ]; then
+              fastlane scan --only_testing TestAppTests/TestAppTests/testMetal
             else
               fastlane scan --only_testing TestAppTests/TestAppTests/testFullJIT
             fi

diff --git a/ios/TestApp/TestAppTests/TestAppTests.mm b/ios/TestApp/TestAppTests/TestAppTests.mm
@@ -31,4 +31,14 @@ - (void)testFullJIT {
   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
@@ -67,6 +67,7 @@
 end
 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']
 targets.each do |target|
     target.frameworks_build_phases.clear
     for lib in libs do
@@ -78,5 +79,16 @@
     end
 end
 
+# link system frameworks
+if 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
 puts "Done."
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")