[GPU] Introduce USE_PYTORCH_METAL

Pull Request resolved: #46383 The old `USE_METAL` is actually being used by Caffe2. Here we introduce a new macro to enable metal in pytorch. ghstack-source-id: 114384673 Differential Revision: [D24322018](https://our.internmc.facebook.com/intern/diff/D24322018/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D24322018/)!
pytorch · Oct 15, 2020 · 515a581 · 515a581
1 parent 2d6fd22
commit 515a581
Show file tree

Hide file tree

Showing 6 changed files with 18 additions and 12 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -165,7 +165,8 @@ option(USE_GLOG "Use GLOG" OFF)
 option(USE_LEVELDB "Use LEVELDB" OFF)
 option(USE_LITE_PROTO "Use lite protobuf instead of full." OFF)
 option(USE_LMDB "Use LMDB" OFF)
-option(USE_METAL "Use Metal for iOS build" OFF)
+option(USE_METAL "Use Metal for Caffe2 iOS build" ON)
+option(USE_PYTORCH_METAL "Use Metal for PyTorch iOS build" OFF)
 option(USE_NATIVE_ARCH "Use -march=native" OFF)
 cmake_dependent_option(
     USE_NCCL "Use NCCL" ON

diff --git a/aten/src/ATen/CMakeLists.txt b/aten/src/ATen/CMakeLists.txt
@@ -125,7 +125,7 @@ else()
   set(all_cpu_cpp ${all_cpu_cpp} ${vulkan_cpp})
 endif()
 
-if(USE_METAL)
+if(USE_PYTORCH_METAL)
   set(all_cpu_cpp ${all_cpu_cpp} ${metal_cpp} ${native_metal_srcs})
 else()
   set(all_cpu_cpp ${all_cpu_cpp} ${metal_cpp})
@@ -390,7 +390,7 @@ set(INSTALL_HEADERS ${base_h} ${ATen_CORE_HEADERS})
 if(NOT INTERN_BUILD_MOBILE)
   list(APPEND INSTALL_HEADERS ${native_h} ${native_cpu_h} ${native_quantized_h} ${cuda_h} ${native_cuda_h} ${native_hip_h} ${cudnn_h} ${hip_h} ${miopen_h})
 else()
-  if(USE_METAL)
+  if(USE_PYTORCH_METAL)
     list(APPEND INSTALL_HEADERS ${metal_h} ${native_metal_h})
   endif()
 endif()

diff --git a/aten/src/ATen/native/metal/MetalAten.mm b/aten/src/ATen/native/metal/MetalAten.mm
@@ -247,7 +247,7 @@ Tensor reshape(const Tensor& input, IntArrayRef shape) {
 
 struct MetalImpl : public at::metal::MetalInterface {
   bool is_metal_available() const override {
-#if defined(USE_METAL)
+#if defined(USE_PYTORCH_METAL)
     return [[MPSCNNContext sharedInstance] available];
 #else
     return false;
@@ -260,7 +260,7 @@ bool is_metal_available() const override {
     return native::metal::metal_copy_impl_(input, src);
   }
 };
-#if defined(USE_METAL)
+#if defined(USE_PYTORCH_METAL)
 static at::metal::MetalImplRegistrar g_metal_impl(new MetalImpl());
 #endif
 

diff --git a/cmake/Summary.cmake b/cmake/Summary.cmake
@@ -107,6 +107,7 @@ function(caffe2_print_configuration_summary)
     message(STATUS "    LMDB version        : ${LMDB_VERSION}")
   endif()
   message(STATUS "  USE_METAL             : ${USE_METAL}")
+  message(STATUS "  USE_PYTORCH_METAL     : ${USE_PYTORCH_METAL}")
   message(STATUS "  USE_MKL               : ${CAFFE2_USE_MKL}")
   message(STATUS "  USE_MKLDNN            : ${USE_MKLDNN}")
   if(${CAFFE2_USE_MKLDNN})

diff --git a/scripts/build_ios.sh b/scripts/build_ios.sh
@@ -92,7 +92,7 @@ CMAKE_ARGS+=("-DUSE_LEVELDB=OFF")
 CMAKE_ARGS+=("-DUSE_MPI=OFF")
 CMAKE_ARGS+=("-DUSE_NUMPY=OFF")
 CMAKE_ARGS+=("-DUSE_NNPACK=OFF")
-CMAKE_ARGS+=("-DUSE_METAL=OFF")
+CMAKE_ARGS+=("-DUSE_PYTORCH_METAL=OFF")
 CMAKE_ARGS+=("-DUSE_MKLDNN=OFF")
 
 # pthreads

diff --git a/torch/csrc/jit/passes/metal_rewrite.cpp b/torch/csrc/jit/passes/metal_rewrite.cpp
@@ -15,7 +15,7 @@
 namespace torch {
 namespace jit {
 
-#ifdef USE_METAL
+#ifdef USE_PYTORCH_METAL
 
 namespace {
 
@@ -196,27 +196,31 @@ script::Module metalOptimizeForMobile(
 #else
 
 void metalInsertPrePackedOps(std::shared_ptr<Graph>& graph) {
-  TORCH_INTERNAL_ASSERT("metal is not enabled. Please build with USE_METAL=1");
+  TORCH_INTERNAL_ASSERT(
+      "metal is not enabled. Please build with USE_PYTORCH_METAL=1");
 }
 
 void metalInsertPrePackedOps(script::Module& module) {
-  TORCH_INTERNAL_ASSERT("metal is not enabled. Please build with USE_METAL=1");
+  TORCH_INTERNAL_ASSERT(
+      "metal is not enabled. Please build with USE_PYTORCH_METAL=1");
 }
 
 TORCH_API void metalFusePrePackedConvWithClamp(script::Module& module) {
-  TORCH_INTERNAL_ASSERT("metal is not enabled. Please build with USE_METAL=1");
+  TORCH_INTERNAL_ASSERT(
+      "metal is not enabled. Please build with USE_PYTORCH_METAL=1");
 }
 
 TORCH_API void metalFoldPrePackingOps(script::Module& module) {
-  TORCH_INTERNAL_ASSERT("metal is not enabled. Please build with USE_METAL=1");
+  TORCH_INTERNAL_ASSERT(
+      "metal is not enabled. Please build with USE_PYTORCH_METAL=1");
 }
 
 script::Module metalOptimizeForMobile(
     const script::Module& m,
     const std::vector<std::string>& preserved_methods) {
   TORCH_INTERNAL_ASSERT(
       "Mobile optimizaiton only available with metal at the moment. "
-      "metal is not enabled. Please build with USE_METAL=1");
+      "metal is not enabled. Please build with USE_PYTORCH_METAL=1");
   return m;
 }