Fix the caffe2_gpu linkage with torch on Windows (#16071)

Summary:
Fixes https://github.com/pytorch/pytorch/issues/15992.
Inspired by https://docs.microsoft.com/en-us/cpp/build/reference/optimization-best-practices?view=vs-2017. But this PR needs to be tested.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16071

Differential Revision: D13693006

Pulled By: soumith

fbshipit-source-id: e83e9ae2591fa4da01d2b1b593558dba3bdc3cf7

diff --git a/torch/CMakeLists.txt b/torch/CMakeLists.txt
index 02d6abc..210b3bc 100644 (file)
--- a/torch/CMakeLists.txt
+++ b/torch/CMakeLists.txt
@@ -419,6 +419,12 @@ if(NOT ${CMAKE_VERSION} VERSION_LESS "3.1")
   set_property(TARGET torch PROPERTY CXX_STANDARD 11)
 endif()
 
+# Prevent the unused functions being optimized away
+# Otherwise torch.dll will be linked without caffe2_gpu.dll
+if (MSVC)
+  set_target_properties(torch PROPERTIES LINK_FLAGS "/OPT:NOREF")
+endif(MSVC)
+
 install(DIRECTORY "${TORCH_SRC_DIR}/csrc"
         DESTINATION ${TORCH_INSTALL_INCLUDE_DIR}/torch
         FILES_MATCHING PATTERN "*.h")

diff --git a/torch/csrc/cuda/comm.cpp b/torch/csrc/cuda/comm.cpp
index d4045b8..a4f5067 100644 (file)
--- a/torch/csrc/cuda/comm.cpp
+++ b/torch/csrc/cuda/comm.cpp
@@ -16,6 +16,18 @@
 #include <cstddef>
 #include <vector>
 
+
+// The following code is used to ensure torch is linked against caffe2_gpu.
+#ifdef _MSC_VER
+namespace {
+#pragma optimize("", off)
+  int warp_size() {
+    return at::cuda::warp_size();
+  }
+#pragma optimize("", on)
+}
+#endif
+
 namespace torch { namespace cuda {
 using namespace at;
 using namespace torch::autograd;