Automatically set BUILD_SPLIT_CUDA for cpp exts (pytorch#52503)

Summary:
Fixes pytorch/vision#3418 (comment)

Pull Request resolved: pytorch#52503

Reviewed By: malfet

Differential Revision: D26546857

Pulled By: janeyx99

fbshipit-source-id: a100b408e7cd28695145a1dda7f2fa081bb7f21f

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]: