From 81fd3a2b0e328a425c0ee8bf5dd82281c181c25f Mon Sep 17 00:00:00 2001
From: peterjc123 <peterghost86@gmail.com>
Date: Fri, 19 Feb 2021 12:20:42 -0800
Subject: [PATCH] Automatically set BUILD_SPLIT_CUDA for cpp exts (#52503)

Summary:
Fixes https://github.com/pytorch/vision/pull/3418#issuecomment-781673110

Pull Request resolved: https://github.com/pytorch/pytorch/pull/52503

Reviewed By: malfet

Differential Revision: D26546857

Pulled By: janeyx99

fbshipit-source-id: a100b408e7cd28695145a1dda7f2fa081bb7f21f
---
 torch/utils/cpp_extension.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/torch/utils/cpp_extension.py b/torch/utils/cpp_extension.py
index ec650a5a149c2a..ca673033e1767f 100644
--- a/torch/utils/cpp_extension.py
+++ b/torch/utils/cpp_extension.py
@@ -22,16 +22,21 @@
 from setuptools.command.build_ext import build_ext
 from pkg_resources import packaging  # type: ignore
 
-BUILD_SPLIT_CUDA = os.getenv('BUILD_SPLIT_CUDA')
 IS_WINDOWS = sys.platform == 'win32'
 LIB_EXT = '.pyd' if IS_WINDOWS else '.so'
 EXEC_EXT = '.exe' if IS_WINDOWS else ''
+CLIB_PREFIX = '' if IS_WINDOWS else 'lib'
+CLIB_EXT = '.dll' if IS_WINDOWS else '.so'
 SHARED_FLAG = '/DLL' if IS_WINDOWS else '-shared'
 
 _HERE = os.path.abspath(__file__)
 _TORCH_PATH = os.path.dirname(os.path.dirname(_HERE))
 TORCH_LIB_PATH = os.path.join(_TORCH_PATH, 'lib')
 
+
+BUILD_SPLIT_CUDA = os.getenv('BUILD_SPLIT_CUDA') or (os.path.exists(os.path.join(
+    TORCH_LIB_PATH, f'{CLIB_PREFIX}torch_cuda_cu{CLIB_EXT}')) and os.path.exists(os.path.join(TORCH_LIB_PATH, f'{CLIB_PREFIX}torch_cuda_cpp{CLIB_EXT}')))
+
 # Taken directly from python stdlib < 3.9
 # See https://github.com/pytorch/pytorch/issues/48617
 def _nt_quote_args(args: Optional[List[str]]) -> List[str]: