[WIP] Include GPU in wheels

.github/workflows/testing_wheels.yml

@@ -32,32 +32,47 @@ jobs:
               arch: x86_64
               build: "cp36* cp37* cp38* cp39*"
               manylinux_image: manylinux2014
+          - os: ubuntu-20.04
+            name: manylinux2014-gpu
+            cibw:
+              arch: x86_64
+              build: "cp36*"
+              manylinux_image: manylinux2014
           - os: windows-2019
             name: win_amd64
             architecture: x64
             cibw:
               build: "cp36-win_amd64 cp37-win_amd64 cp38-win_amd64 cp39-win_amd64"
     env:
+      QSIMCIRQ_PROJECT: "${{ matrix.name == 'manylinux2014-gpu' && 'qsimcirq-gpu' || 'qsimcirq'}}"
       CIBW_BUILD: "${{ matrix.cibw.build || '*' }}"
       CIBW_ARCHS: "${{ matrix.cibw.arch || 'auto' }}"
       CIBW_MANYLINUX_X86_64_IMAGE: "${{ matrix.cibw.manylinux_image }}"
+      CIBW_BEFORE_ALL_LINUX: "bash build_tools/cuda_install.sh && echo $PATH"
+      CIBW_ENVIRONMENT_LINUX: PATH=/usr/local/cuda-11/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-11/lib64:$LD_LIBRARY_PATH
+      CIBW_ENVIRONMENT_PASS_LINUX: PATH LD_LIBRARY_PATH
       CIBW_BEFORE_BUILD_MACOS: "brew install libomp"
       CIBW_REPAIR_WHEEL_COMMAND_MACOS: ""
       # due to package and module name conflict have to temporarily move it away to run tests
       CIBW_BEFORE_TEST: "mv {package}/qsimcirq /tmp"
-      CIBW_TEST_COMMAND: "pytest {package}/qsimcirq_tests/qsimcirq_test.py && mv /tmp/qsimcirq {package}"
+      # prevent GPU tests (no GPU on GHA runners)
+      CIBW_TEST_COMMAND: "pytest {package}/qsimcirq_tests/qsimcirq_test.py -k \"not gpu\" && mv /tmp/qsimcirq {package}"
     steps:
       - uses: actions/checkout@v2
 
       # Used to host cibuildwheel
       - uses: actions/setup-python@v2
 
+      # Install NVCC on Ubuntu to include GPU in the wheel.
+      - name: cuda-toolkit
+        if: ${{ matrix.os == 'ubuntu-20.04' }}
+        uses: Jimver/[email protected]
+        with:
+          cuda: "11.5.1"
+
       - name: Install cibuildwheel and twine
         run: python -m pip install cibuildwheel==1.11.0
 
-      - name: Run C++ tests
-        run: bash build_tools/test_libs.sh
-
       - name: Build wheels
         run: python -m cibuildwheel --output-dir wheelhouse
 

CMakeLists.txt

@@ -1,14 +1,13 @@
 cmake_minimum_required(VERSION 3.11)
 
-execute_process(COMMAND which nvcc OUTPUT_VARIABLE has_nvcc)
-if(has_nvcc STREQUAL "")
-    project(qsim)
-else()
+if(ENV{QSIMCIRQ_PROJECT} STREQUAL "qsimcirq-gpu")
     project(qsim LANGUAGES CXX CUDA)
     ADD_SUBDIRECTORY(pybind_interface/cuda)
     if(DEFINED ENV{CUQUANTUM_DIR})
         ADD_SUBDIRECTORY(pybind_interface/custatevec)
     endif()
+else()
+    project(qsim)
 endif()
 
 ADD_SUBDIRECTORY(pybind_interface/sse)

build_tools/cuda_install.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+set -x
+
+# install dependencies
+kernel_version=$(uname -r)
+yum clean all
+yum install -y kernel
+yum install -y kernel-devel-${kernel_version} kernel-headers-${kernel_version} pciutils
+
+# install python and CUDA toolkit
+# TODO: need to pin version
+
+yum install -y yum-utils epel-release
+yum install -y python3-devel.x86_64
+yum-config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/cuda-rhel7.repo
+yum clean all
+yum -y install nvidia-driver-latest-dkms cuda
+yum -y install cuda-drivers
+
+# post-install actions
+
+# commented to allow outer PATH-setting
+# export PATH=/usr/local/cuda-11/bin:$PATH
+# export LD_LIBRARY_PATH=/usr/local/cuda-11/lib64:$LD_LIBRARY_PATH
+
+# works here, does not escape to build process
+nvcc --version
+which nvcc
+nvidia-smi
+
+set +x

pybind_interface/cuda/CMakeLists.txt

@@ -27,7 +27,7 @@ if(NOT pybind11_POPULATED)
   add_subdirectory(${pybind11_SOURCE_DIR} ${pybind11_BINARY_DIR})
 endif()
 
-find_package(PythonLibs 3.6 REQUIRED)
+find_package(Python 3.6 REQUIRED COMPONENTS Interpreter Development)
 find_package(CUDA REQUIRED)
 
 include_directories(${PYTHON_INCLUDE_DIRS} ${pybind11_SOURCE_DIR}/include)

pybind_interface/decide/CMakeLists.txt

@@ -1,10 +1,9 @@
 cmake_minimum_required(VERSION 3.11)
 
-execute_process(COMMAND which nvcc OUTPUT_VARIABLE has_nvcc)
-if(has_nvcc STREQUAL "")
-  project(qsim)
-else()
+if(ENV{QSIMCIRQ_PROJECT} STREQUAL "qsimcirq-gpu")
   project(qsim LANGUAGES CXX CUDA)
+else()
+  project(qsim)
 endif()
 
 IF (WIN32)
@@ -32,9 +31,7 @@ if(NOT pybind11_POPULATED)
   add_subdirectory(${pybind11_SOURCE_DIR} ${pybind11_BINARY_DIR})
 endif()
 
-if(has_nvcc STREQUAL "")
-  pybind11_add_module(qsim_decide decide.cpp)
-else()
+if(ENV{QSIMCIRQ_PROJECT} STREQUAL "qsimcirq-gpu")
   find_package(PythonLibs 3.6 REQUIRED)
   find_package(CUDA REQUIRED)
 
@@ -53,4 +50,6 @@ else()
         SUFFIX "${PYTHON_MODULE_EXTENSION}"
   )
   set_source_files_properties(decide.cpp PROPERTIES LANGUAGE CUDA)
+else()
+  pybind11_add_module(qsim_decide decide.cpp)
 endif()

setup.py

@@ -9,6 +9,9 @@
 from distutils.version import LooseVersion
 
 
+project_name = os.environ.get("QSIMCIRQ_PROJECT", "qsimcirq")
+
+
 class CMakeExtension(Extension):
     def __init__(self, name, sourcedir=""):
         Extension.__init__(self, name, sources=[])
@@ -87,27 +90,32 @@ def build_extension(self, ext):
 __version__ = ""
 exec(open("qsimcirq/_version.py").read())
 
+cmake_exts = [
+    CMakeExtension("qsimcirq/qsim_avx512"),
+    CMakeExtension("qsimcirq/qsim_avx2"),
+    CMakeExtension("qsimcirq/qsim_sse"),
+    CMakeExtension("qsimcirq/qsim_basic"),
+    CMakeExtension("qsimcirq/qsim_decide"),
+]
+if "gpu" in project_name:
+    cmake_exts += [
+        CMakeExtension("qsimcirq/qsim_cuda"),
+        CMakeExtension("qsimcirq/qsim_custatevec"),
+    ]
+
 setup(
-    name="qsimcirq",
+    name=project_name,
     version=__version__,
     url="https://github.com/quantumlib/qsim",
-    author="Vamsi Krishna Devabathini",
-    author_email="devabathini92@gmail.com",
-    python_requires=">=3.3.0",
+    author="The qsim Developers",
+    author_email="qsim-qsimh-dev@googlegroups.com",
+    python_requires=">=3.6.0",
     install_requires=requirements,
     license="Apache 2",
     description=description,
     long_description=long_description,
     long_description_content_type="text/markdown",
-    ext_modules=[
-        CMakeExtension("qsimcirq/qsim_avx512"),
-        CMakeExtension("qsimcirq/qsim_avx2"),
-        CMakeExtension("qsimcirq/qsim_sse"),
-        CMakeExtension("qsimcirq/qsim_basic"),
-        CMakeExtension("qsimcirq/qsim_cuda"),
-        CMakeExtension("qsimcirq/qsim_custatevec"),
-        CMakeExtension("qsimcirq/qsim_decide"),
-    ],
+    ext_modules=cmake_exts,
     cmdclass=dict(build_ext=CMakeBuild),
     zip_safe=False,
     packages=["qsimcirq"],