Repository: pytorch
Issue: BuildExtension does not choose correct CUDA installation
Labels: module: cpp-extensions,module: cuda,triaged

### 🐛 Describe the bug

When trying to build CUDA extensions via the `BuildExtension` class (`torch.utils.cpp_extension.BuildExtension`), if multiple CUDA versions are installed, the build might fail. This is specifically due to [`_find_cuda_home()`](https://github.com/pytorch/pytorch/blob/619db911658e7f06c34e4198f7d0c9d4d629dd4a/torch/utils/cpp_extension.py#L81) choosing the first CUDA installation it can find, regardless of its compatibility with the current PyTorch version. This then causes the check at [`_check_cuda_version()`](https://github.com/pytorch/pytorch/blob/619db911658e7f06c34e4198f7d0c9d4d629dd4a/torch/utils/cpp_extension.py#L800) to error out, if the wrong path was selected automatically, regardless of whether a compatible CUDA version is available on the system.

In my specific case, there are installations of both CUDA 10.1 and various CUDA 11.x versions present on a Ubuntu machine. `nvcc` points to the 10.1 install, while the one at `/usr/local/cuda` is a compatible version. The logic of `_find_cuda_home()` checks the location of `nvcc` first and only uses `usr/local/cuda` as a fallback, giving me the following error:
```
Traceback (most recent call last):
  File "setup.py", line 20, in <module>
    setup(
  File "[...]/env/lib/python3.8/site-packages/setuptools/__init__.py", line 87, in setup
    return distutils.core.setup(**attrs)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/core.py", line 148, in setup
    return run_commands(dist)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/core.py", line 163, in run_commands
    dist.run_commands()
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 967, in run_commands
    self.run_command(cmd)
  File "[...]/env/lib/python3.8/site-packages/setuptools/dist.py", line 1214, in run_command
    super().run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 986, in run_command
    cmd_obj.run()
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/install.py", line 74, in run
    self.do_egg_install()
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/install.py", line 123, in do_egg_install
    self.run_command('bdist_egg')
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/dist.py", line 1214, in run_command
    super().run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 986, in run_command
    cmd_obj.run()
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 165, in run
    cmd = self.call_command('install_lib', warn_dir=0)
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 151, in call_command
    self.run_command(cmdname)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/dist.py", line 1214, in run_command
    super().run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 986, in run_command
    cmd_obj.run()
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/install_lib.py", line 11, in run
    self.build()
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/command/install_lib.py", line 107, in build
    self.run_command('build_ext')
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 313, in run_command
    self.distribution.run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/dist.py", line 1214, in run_command
    super().run_command(command)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 986, in run_command
    cmd_obj.run()
  File "[...]/env/lib/python3.8/site-packages/setuptools/command/build_ext.py", line 79, in run
    _build_ext.run(self)
  File "[...]/env/lib/python3.8/site-packages/setuptools/_distutils/command/build_ext.py", line 339, in run
    self.build_extensions()
  File "[...]/env/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 410, in build_extensions
    self._check_cuda_version()
  File "[...]/env/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 787, in _check_cuda_version
    raise RuntimeError(CUDA_MISMATCH_MESSAGE.format(cuda_str_version, torch.version.cuda))
RuntimeError: 
The detected CUDA version (10.1) mismatches the version that was used to compile
PyTorch (11.3). Please make sure to use the same CUDA versions.
```

This can of course be worked around by manually setting `CUDA_HOME` to the correct path, but in my mind, the fact that PyTorch detects & uses the correct CUDA installation for training & inference and only fails to use the correct one when building extensions does sound unintended to me.

Ideally in my mind, [`_find_cuda_home()`](https://github.com/pytorch/pytorch/blob/619db911658e7f06c34e4198f7d0c9d4d629dd4a/torch/utils/cpp_extension.py#L81) should be adjusted to detect the same CUDA version as used by PyTorch normally or even automatically verify version compatibility for each candidate install location instead of just choosing the first one and verifying compatibility later.

Quick addendum: this error occurs with every CUDA extension I try to build and also happens on multiple systems with different hardware & CUDA configurations, as long as there are multiple CUDA installations and the one used by `nvcc` is an incompatible one.

### Versions

```
PyTorch version: 1.11.0+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.4 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.4
Libc version: glibc-2.31

Python version: 3.8.10 (default, Mar 15 2022, 12:22:08)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-107-generic-x86_64-with-glibc2.29
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000
GPU 4: NVIDIA RTX A6000
GPU 5: NVIDIA RTX A6000
GPU 6: NVIDIA RTX A6000
GPU 7: NVIDIA RTX A6000

Nvidia driver version: 470.103.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.21.0
[pip3] pytorch-lightning==1.6.3
[pip3] torch==1.11.0+cu113
[pip3] torchaudio==0.11.0+cu113
[pip3] torchmetrics==0.8.2
[pip3] torchtext==0.12.0
[pip3] torchvision==0.12.0+cu113
[conda] Could not collect
```

cc @malfet @zou3519 @ngimel