Add C++ support for raw CUDA pointer backed state-vector

.github/workflows/build_wheel_manylinux2014.yml

@@ -23,7 +23,7 @@ env:
   CIBW_ENVIRONMENT: PATH=$PATH:/usr/local/cuda/bin LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:$(python -c "import site; print( f'{site.getsitepackages()[0]}/cuquantum/lib')")
 
   CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
-  CIBW_REPAIR_WHEEL_COMMAND_LINUX: "./docker/auditwheel repair -w {dest_dir} {wheel}"
+  CIBW_REPAIR_WHEEL_COMMAND_LINUX: "python ./docker/auditwheel repair -w {dest_dir} {wheel}"
 
 jobs:
   linux-wheels-x86-64:
@@ -41,7 +41,12 @@ jobs:
         with:
           access_token: ${{ github.token }}
 
-      - uses: actions/checkout@v2
+      - name: Remove unused SDKs for extra storage during build
+        run: |
+          sudo rm -rf /usr/share/dotnet       # DotNet
+          sudo rm -rf /usr/local/lib/android  # Android SDK
+
+      - uses: actions/checkout@v3
 
       - uses: actions/setup-python@v2
         name: Install Python
@@ -63,7 +68,7 @@ jobs:
           regex: '.*[0-9]+.[0-9]+.[0-9]+[-_]?rc[0-9]+'
 
       - uses: actions/upload-artifact@v2
-        if: ${{ github.event_name == 'release' || github.ref == 'refs/heads/main' || steps.rc_build.outputs.match != ''}}
+        #if: ${{ github.event_name == 'release' || github.ref == 'refs/heads/main' || steps.rc_build.outputs.match != ''}}
         with:
           name: ${{ runner.os }}-wheels
           path: ./wheelhouse/*.whl

CMakeLists.txt

@@ -52,6 +52,7 @@ option(ENABLE_CLANG_TIDY "Enable clang-tidy build checks" OFF)
 option(DISABLE_CUDA_SAFETY "Build without CUDA call safety checks" OFF)
 option(ENABLE_PYTHON "Build Python bindings" ON)
 option(ENABLE_SANITIZER "Enable address sanitizer" OFF)
+option(ENABLE_INTERFACES "Enable ML framework interfaces" ON)
 
 # Build options
 if(NOT CMAKE_BUILD_TYPE)
@@ -98,7 +99,7 @@ set(ENABLE_OPENMP ON)
 FetchContent_Declare(
     pennylane_lightning
     GIT_REPOSITORY https://github.com/PennyLaneAI/pennylane-lightning.git
-    GIT_TAG        latest
+    GIT_TAG        master
 )
 FetchContent_MakeAvailable(pennylane_lightning)
 
@@ -108,7 +109,7 @@ if(ENABLE_OPENMP)
     find_package(OpenMP REQUIRED)
 endif()
 
-find_package (Python COMPONENTS Interpreter Development)
+find_package(Python ${PY_VERSION} EXACT REQUIRED)
 
 find_library(CUSTATEVEC_LIB
     NAMES   libcustatevec.so.1 custatevec.so.1
@@ -124,7 +125,7 @@ find_library(CUSTATEVEC_LIB
             ${CUDAToolkit_LIBRARY_DIR}
             ${CUDA_TOOLKIT_ROOT_DIR}/lib
             ${CUDA_TOOLKIT_ROOT_DIR}/lib64
-            ${Python_SITELIB}/cuquantum/lib
+            ${Python_SITEARCH}/cuquantum/lib
             ENV LD_LIBRARY_PATH
 )
 
@@ -139,7 +140,7 @@ find_file( CUSTATEVEC_INC
             ${CUQUANTUM_SDK}/include
             ${CUDAToolkit_INCLUDE_DIRS}
             ${CUDA_TOOLKIT_ROOT_DIR}/include
-            ${Python_SITELIB}/cuquantum/include
+            ${Python_SITEARCH}/cuquantum/include
             ENV CPATH
 )
 

docker/Dockerfile

@@ -11,14 +11,20 @@ RUN mkdir -p -m 0700 ~/.ssh && ssh-keyscan github.com >> ~/.ssh/known_hosts
 
 COPY ./  /pennylane-lightning-gpu
 
-# Create venv for each required Python version
 RUN cd /pennylane-lightning-gpu \
     && export PATH=$PATH:/usr/local/cuda/bin \
     && export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64 \
-    && for i in {7..10}; do python3.${i} -m venv pyenv3.${i}; source pyenv3.${i}/bin/activate && python3 -m pip install auditwheel ninja wheel && python -m pip install --no-deps cuquantum && python3 setup.py build_ext --define=CMAKE_CXX_COMPILER=/opt/rh/devtoolset-9/root/usr/bin/g++ --define=ENABLE_CLANG_TIDY=0  && python3 setup.py bdist_wheel && deactivate && rm -rf ./build ; done
+    && for i in `seq 7 10`; do \
+        python3.${i} -m pip install cmake auditwheel ninja wheel \
+        && python3.${i} -m pip install --no-deps cuquantum \
+        && python3.${i} setup.py build_ext \
+            --define="CMAKE_CXX_COMPILER=/opt/rh/devtoolset-10/root/usr/bin/g++" \
+            --define="ENABLE_CLANG_TIDY=0" \
+            --define="PY_VERSION=3.${i}" \
+        && python3.${i} setup.py bdist_wheel && rm -rf ./build ; \
+    done
 
 RUN cd /pennylane-lightning-gpu \
     && export PATH=$PATH:/usr/local/cuda/bin \
-    && source pyenv3.10/bin/activate \
-    && export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:$(python -c "import site; print(site.getsitepackages()[0])")/cuquantum/lib \
-    && for i in $(ls ./dist); do /pennylane-lightning-gpu/docker/auditwheel repair -w /wheelhouse /pennylane-lightning-gpu/dist/$i; done
+    && export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:$(python3.7 -c "import site; print(site.getsitepackages()[0])")/cuquantum/lib \
+    && for i in $(ls ./dist); do python3.7 /pennylane-lightning-gpu/docker/auditwheel repair -w /wheelhouse /pennylane-lightning-gpu/dist/$i; done

docker/auditwheel

@@ -1,5 +1,3 @@
-#!/usr/bin/env python3
-
 # Patch to not ship CUDA system libraries
 # Follows https://github.com/DIPlib/diplib/tree/master/tools/travis
 import sys

pennylane_lightning_gpu/_serialize.py

@@ -33,9 +33,9 @@
 import pennylane as qml
 
 try:
-    from pennylane_lightning_gpu.lightning_gpu_qubit_ops import (
-        LightningGPU_C128,
-        LightningGPU_C64,
+    from pennylane_lightning_gpu.module_bindings import (
+        LightningGPU_StateVectorCudaManaged_C128,
+        LightningGPU_StateVectorCudaManaged_C64,
         ObsStructGPU_C64,
         ObsStructGPU_C128,
     )
@@ -141,7 +141,11 @@ def _serialize_ops(
 
     uses_stateprep = False
 
-    sv_py = LightningGPU_C64 if use_csingle else LightningGPU_C128
+    sv_py = (
+        LightningGPU_StateVectorCudaManaged_C64
+        if use_csingle
+        else LightningGPU_StateVectorCudaManaged_C128
+    )
 
     for o in tape.operations:
         if isinstance(o, (BasisState, QubitStateVector)):

pennylane_lightning_gpu/lightning_gpu.py

@@ -40,20 +40,7 @@
 from ._version import __version__
 
 try:
-    from .lightning_gpu_qubit_ops import (
-        LightningGPU_C128,
-        LightningGPU_C64,
-        AdjointJacobianGPU_C128,
-        AdjointJacobianGPU_C64,
-        device_reset,
-        is_gpu_supported,
-        get_gpu_arch,
-        DevPool,
-        ObsStructGPU_C128,
-        ObsStructGPU_C64,
-        OpsStructGPU_C128,
-        OpsStructGPU_C64,
-    )
+    from .module_bindings import *
 
     from ._serialize import _serialize_obs, _serialize_ops
     from ctypes.util import find_library
@@ -87,10 +74,22 @@
 )
 
 
-def _gpu_dtype(dtype):
+def _gpu_dtype(dtype, use_managed=True):
     if dtype not in [np.complex128, np.complex64]:
         raise ValueError(f"Data type is not supported for state-vector computation: {dtype}")
-    return LightningGPU_C128 if dtype == np.complex128 else LightningGPU_C64
+    if use_managed:
+        ctor = (
+            LightningGPU_StateVectorCudaManaged_C128
+            if dtype == np.complex128
+            else LightningGPU_StateVectorCudaManaged_C64
+        )
+    else:
+        ctor = (
+            LightningGPU_StateVectorCudaRaw_C128
+            if dtype == np.complex128
+            else LightningGPU_StateVectorCudaRaw_C64
+        )
+    return ctor
 
 
 class LightningGPU(LightningQubit):
@@ -103,7 +102,7 @@ class LightningGPU(LightningQubit):
 
     name = "PennyLane plugin for GPU-backed Lightning device using NVIDIA cuQuantum SDK"
     short_name = "lightning.gpu"
-    pennylane_requires = ">=0.22"
+    pennylane_requires = ">=0.23"
     version = __version__
     author = "Xanadu Inc."
     _CPP_BINARY_AVAILABLE = True

pennylane_lightning_gpu/module_bindings.py

@@ -0,0 +1,16 @@
+from .lightning_gpu_qubit_ops import (
+    LightningGPU_StateVectorCudaManaged_C128,
+    LightningGPU_StateVectorCudaManaged_C64,
+    LightningGPU_StateVectorCudaRaw_C128,
+    LightningGPU_StateVectorCudaRaw_C64,
+    AdjointJacobianGPU_C128,
+    AdjointJacobianGPU_C64,
+    ObsStructGPU_C128,
+    ObsStructGPU_C64,
+    OpsStructGPU_C128,
+    OpsStructGPU_C64,
+    device_reset,
+    is_gpu_supported,
+    get_gpu_arch,
+    DevPool,
+)