Replace pycuda with cuda-python in detect_target (facebookincubator#391)

Summary:
NVIDIA provides official python API for CUDA - [NVIDIA/cuda-python](https://github.com/NVIDIA/cuda-python).  We can use it instead of [inducer/pycuda](https://github.com/inducer/pycuda).

Benefits of `cuda-python`:
- Pre-installed in many NVIDIA containers (e.g. [tensorrt](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt) and [pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) containers)
- `pip install` is fast. **No compilation** during the installation.
- Versioning is more clear - it corresponds to cuda version (e.g. 11.6.1). In contrast `pycuda` uses date based versioning  - e.g. 2022.2.2 which does not correspond to cuda version.
- Official NVIDIA package

Pull Request resolved: facebookincubator#391

Reviewed By: chenyang78

Differential Revision: D43962921

Pulled By: tenpercent

fbshipit-source-id: a2a7a5c8fc627ba3f122f208ba38edc60afda396

.circleci/config.yml

@@ -27,7 +27,7 @@ setup_env: &setup_env
           python3.8 setup.py bdist_wheel &&
           sudo python3.8 -m pip install --no-input dist/*.whl &&
           cd /home/circleci/project &&
-          python3.8 -m pip install pycuda &&
+          python3.8 -m pip install 'cuda-python<12.0.0' &&
           python3.8 -m pip install pytest &&
           python3.8 -m pip install torch &&
           python3.8 -m pip install numpy &&

docker/Dockerfile.cuda

@@ -40,8 +40,8 @@ RUN bash /Install/install_doc_dep.sh
 # install Pytorch
 RUN pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
 
-# install Pycuda
-RUN pip3 install pycuda
+# install NVIDIA cuda-python
+RUN pip3 install 'cuda-python<12.0.0'
 
 # for detection
 RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata

python/aitemplate/testing/detect_target.py

@@ -54,10 +54,16 @@ def _detect_cuda_with_nvidia_smi():
 
 def _detect_cuda():
     try:
-        import pycuda.driver as drv
+        from cuda import cuda
 
-        drv.init()
-        major, minor = drv.Device(0).compute_capability()
+        def assert_cuda(res):
+            if res[0].value != 0:
+                raise RuntimeError(f"CUDA error code={res[0].value}")
+            return res[1:]
+
+        assert_cuda(cuda.cuInit(0))
+        # Get Compute Capability of the first Visible device
+        major, minor = assert_cuda(cuda.cuDeviceComputeCapability(0))
         comp_cap = major * 10 + minor
         if comp_cap >= 90:
             return "90"