[GPU] Enable Metal on macosx

Pull Request resolved: pytorch/pytorch#47635 Add macosx support for metal. The supported os version is 10.13 and above. ghstack-source-id: 116607205 Differential Revision: [D24825088](https://our.internmc.facebook.com/intern/diff/D24825088/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D24825088/)!
jasperzhong · Nov 13, 2020 · ae26e86 · ae26e86
1 parent 6422a78
commit ae26e86
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 11 deletions.
diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt
@@ -135,7 +135,7 @@ endif()
 
 # Metal
 if(USE_PYTORCH_METAL)
-  if(IOS)
+  if(APPLE)
     set(all_cpu_cpp ${all_cpu_cpp} ${metal_cpp} ${native_metal_srcs})
   else()
     # Add files needed from optimized_for_mobile

diff --git a/aten/src/ATen/native/metal/MetalPrepackOpRegister.cpp b/aten/src/ATen/native/metal/MetalPrepackOpRegister.cpp
@@ -3,7 +3,7 @@
 #include <ATen/ATen.h>
 
 
-#if defined(C10_IOS)
+#if (C10_IOS || TARGET_OS_MAC)
 #import <ATen/native/metal/mpscnn/MPSCNNOps.h>
 #endif
 
@@ -94,19 +94,19 @@ c10::intrusive_ptr<Conv2dOpContext> conv2d_prepack(
 Tensor conv2d_prepack_run(
     const Tensor& input,
     const c10::intrusive_ptr<Conv2dOpContext>& op_context) {
-#if defined(C10_IOS)
+#if (C10_IOS || TARGET_OS_MAC)
   return mpscnn::conv2d(input, *op_context);
 #else
-  TORCH_CHECK(false, "conv2d_prepack_run can only be invoked on iOS");
+  TORCH_CHECK(false, "conv2d_prepack_run can only be invoked on iOS and MacOS");
   return input;
 #endif
 }
 
 Tensor copy_to_host(const Tensor& input) {
-#if defined(C10_IOS)
+#if (C10_IOS || TARGET_OS_MAC)
   return mpscnn::copy_to_host(input);
 #else
-  TORCH_CHECK(false, "copy_to_host can only be invoked on iOS");
+  TORCH_CHECK(false, "copy_to_host can only be invoked on iOS and MacOS");
   return input;
 #endif
 }

diff --git a/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.mm b/aten/src/ATen/native/metal/mpscnn/MPSCNNContext.mm
@@ -4,8 +4,10 @@
 #include <torch/script.h>
 #include <mutex>
 
-#if defined(C10_IOS)
+#if C10_IOS
 #import <UIKit/UIKit.h>
+#elif TARGET_OS_MAC
+#import <Foundation/NSProcessInfo.h>
 #endif
 
 @implementation MPSCNNContext {
@@ -31,10 +33,11 @@ + (instancetype)sharedInstance {
 }
 
 - (BOOL)available {
-#if defined(C10_IOS)
-#if TARGET_IPHONE_SIMULATOR
+#if !defined(__APPLE__)
   return false;
-#else
+#elif TARGET_IPHONE_SIMULATOR
+  return false;
+#elif TARGET_OS_IPHONE
   if (!MPSSupportsMTLDevice(_device)) {
     return false;
   }
@@ -45,8 +48,23 @@ - (BOOL)available {
           supportsFeatureSet:MTLFeatureSet_iOS_GPUFamily3_v2]) {
     return false;
   }
+#elif TARGET_OS_MAC
+  if (!MPSSupportsMTLDevice(_device)) {
+    return false;
+  }
+  NSOperatingSystemVersion supportedVer = {10, 13, 0};
+  if (![[NSProcessInfo processInfo]
+          isOperatingSystemAtLeastVersion:supportedVer]) {
+    return false;
+  }
+  if (![MTLCreateSystemDefaultDevice()
+          supportsFeatureSet:MTLFeatureSet_macOS_GPUFamily1_v3]) {
+    return false;
+  }
+#else
+  return false;
 #endif
-#endif
+
   return _device && _library && _commandQueue;
 }