Merge branch 'update_0.1.0.1' into 'main'

Support pip wheel release + Fix samples Closes CUQNT-571 See merge request cuda-hpc-libraries/cuquantum-sdk/cuquantum-public!3
NVIDIA · Mar 12, 2022 · 38dae54 · 38dae54
2 parents 4e63ab0 + d65a6d0
commit 38dae54
Show file tree

Hide file tree

Showing 10 changed files with 335 additions and 198 deletions.
diff --git a/README.md b/README.md
@@ -1,9 +1,14 @@
+<div align="center"><img src="https://developer.nvidia.com/sites/default/files/akamai/nvidia-cuquantum-icon.svg" width="250"/></div>
+
 # Welcome to the cuQuantum repository!
 
 This public repository contains two sets of files related to the [NVIDIA cuQuantum SDK](https://developer.nvidia.com/cuquantum-sdk):
 
 - `samples`: All C/C++ sample codes for the cuQuantum SDK.
 - `python`: The open-sourced cuQuantum Python project.
+  - Available for download on
+    - conda-forge: `cuquantum` [![Conda Version](https://img.shields.io/conda/vn/conda-forge/cuquantum.svg)](https://anaconda.org/conda-forge/cuquantum), `cuquantum-python` [![Conda Version](https://img.shields.io/conda/vn/conda-forge/cuquantum-python.svg)](https://anaconda.org/conda-forge/cuquantum-python)
+    - PyPI: *coming soon!*
 
 Other components of the cuQuantum SDK can be accessed following the instruction given in the documentation.
 

diff --git a/python/README.md b/python/README.md
@@ -2,7 +2,7 @@
 
 ## Documentation
 
-Please visit the [NVIDIA cuQuantum SDK documentation](https://docs.nvidia.com/cuda/cuquantum/).
+Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com/cuda/cuquantum/python).
 
 ## Building
 
@@ -22,7 +22,7 @@ If you already have a Conda environment set up, it is the easiest to install cuQ
 ```
 conda install -c conda-forge cuquantum-python
 ```
-The Conda solver will address all dependencies for you.
+The Conda solver will install all required dependencies for you.
 
 ### Install cuQuantum Python from source
 
@@ -57,6 +57,8 @@ Runtime dependencies of the cuQuantum Python package include:
 
 If you install everything from conda-forge, the dependencies are taken care for you (except for the driver).
 
+If you install the pip wheels, cuTENSOR and cuQuantum are installed for you.
+
 If you build cuQuantum Python from source, please make sure the paths to the cuQuantum and cuTENSOR libraries are added
 to your `LD_LIBRARY_PATH` environment variable.
 

diff --git a/python/cuquantum/__init__.py b/python/cuquantum/__init__.py
@@ -4,6 +4,7 @@
     contract, contract_path, einsum, einsum_path, Network,
     NetworkOptions, OptimizerInfo, OptimizerOptions, PathFinderOptions, ReconfigOptions, SlicerOptions)
 from cuquantum.utils import ComputeType, cudaDataType, libraryPropertyType
+from cuquantum._version import __version__
 
 
 # We patch all enum values so that they have the correct docstrings

diff --git a/python/cuquantum/_version.py b/python/cuquantum/_version.py
@@ -0,0 +1,2 @@
+# TODO: find a better approach to sync the version string with the C libs
+__version__ = '.'.join(['0.1.0', '1'])  # the last digit is for cuQuantum Python only
diff --git a/python/setup.py b/python/setup.py
@@ -7,23 +7,32 @@
 from Cython.Build import cythonize
 
 
+# Get __version__ variable
+source_root = os.path.abspath(os.path.dirname(__file__))
+with open(os.path.join(source_root, 'cuquantum', '_version.py')) as f:
+    exec(f.read())
+
+
 # search order:
 # 1. installed "cuquantum" package
 # 2. env var
 for path in site.getsitepackages():
     path = os.path.join(path, 'cuquantum')
     if os.path.isdir(path):
         cuquantum_root = path
+        using_cuquantum_wheel = True
         break
 else:
     cuquantum_root = os.environ.get('CUQUANTUM_ROOT')
+    using_cuquantum_wheel = False
 
 
 # We allow setting CUSTATEVEC_ROOT and CUTENSORNET_ROOT separately for the ease
 # of development, but users are encouraged to either install cuquantum from PyPI
 # or conda, or set CUQUANTUM_ROOT to the existing installation.
 try:
     custatevec_root = os.environ['CUSTATEVEC_ROOT']
+    using_cuquantum_wheel = False
 except KeyError as e:
     if cuquantum_root is None:
         raise RuntimeError('cuStateVec is not found, please install "cuquantum" '
@@ -32,6 +41,7 @@
         custatevec_root = cuquantum_root
 try:
     cutensornet_root = os.environ['CUTENSORNET_ROOT']
+    using_cuquantum_wheel = False
 except KeyError as e:
     if cuquantum_root is None:
         raise RuntimeError('cuTensorNet is not found, please install "cuquantum" '
@@ -47,9 +57,11 @@
     path = os.path.join(path, 'cutensor')
     if os.path.isdir(path):
         cutensor_root = path
+        assert using_cuquantum_wheel  # if this raises, the env is corrupted
         break
 else:
     cutensor_root = os.environ.get('CUTENSOR_ROOT')
+    assert not using_cuquantum_wheel
 if cutensor_root is None:
     raise RuntimeError('cuTENSOR is not found, please install "cutensor" '
                        'or set $CUTENSOR_ROOT')
@@ -73,10 +85,12 @@
     ]
 ignore_cuquantum_dep = bool(os.environ.get('CUQUANTUM_IGNORE_SOLVER', False))
 if not ignore_cuquantum_dep:
-    setup_requires.append('cuquantum==0.0.1.*')
-    setup_requires.append('cutensor>=1.4.*')
-    install_requires.append('cuquantum==0.0.1.*')
-    install_requires.append('cutensor>=1.4.*')
+    assert using_cuquantum_wheel  # if this raises, the env is corrupted
+    # cuTENSOR version is constrained in the cuquantum package, so we don't
+    # need to list it
+    ver = '.'.join(__version__.split('.')[:3])  # remove the Python patch number
+    setup_requires.append('cuquantum=='+ver+'.*')
+    install_requires.append('cuquantum=='+ver+'.*')
 
 
 def check_cuda_version():
@@ -110,12 +124,40 @@ def check_cuda_version():
     raise RuntimeError(f"Unsupported CUDA version: {cuda_ver}")
 
 
-print()
-print("****************************************************************")
+def prepare_libs_and_rpaths():
+    global cusv_lib_dir, cutn_lib_dir
+    # we include both lib64 and lib to accommodate all possible sources
+    cusv_lib_dir = [os.path.join(custatevec_root, 'lib'),
+                    os.path.join(custatevec_root, 'lib64')]
+    cutn_lib_dir = [os.path.join(cutensornet_root, 'lib'),
+                    os.path.join(cutensornet_root, 'lib64'),
+                    os.path.join(cutensor_root, 'lib', cutensor_ver)]
+
+    global cusv_lib, cutn_lib, extra_linker_flags
+    if using_cuquantum_wheel:
+        cusv_lib = [':libcustatevec.so.0']
+        cutn_lib = [':libcutensornet.so.0', ':libcutensor.so.1']
+        # The rpaths must be adjusted given the following full-wheel installation:
+        #   cuquantum-python: site-packages/cuquantum/{custatevec, cutensornet}/  [=$ORIGIN]
+        #   cusv & cutn:      site-packages/cuquantum/lib/
+        #   cutensor:         site-packages/cutensor/lib/CUDA_VER/
+        ldflag = "-Wl,--disable-new-dtags,"
+        ldflag += "-rpath,$ORIGIN/../lib,"
+        ldflag += f"-rpath,$ORIGIN/../../cutensor/lib/{cutensor_ver}"
+        extra_linker_flags = [ldflag]
+    else:
+        cusv_lib = ['custatevec']
+        cutn_lib = ['cutensornet', 'cutensor']
+        extra_linker_flags = []
+
+
+prepare_libs_and_rpaths()
+print("\n****************************************************************")
 print("CUDA version:", cuda_ver)
 print("CUDA path:", cuda_path)
 print("cuStateVec path:", custatevec_root)
 print("cuTensorNet path:", cutensornet_root)
+print("cuTENSOR path:", cutensor_root)
 print("****************************************************************\n")
 
 
@@ -124,8 +166,9 @@ def check_cuda_version():
     sources=["cuquantum/custatevec/custatevec.pyx"],
     include_dirs=[os.path.join(cuda_path, 'include'),
                   os.path.join(custatevec_root, 'include')],
-    library_dirs=[os.path.join(custatevec_root, 'lib64')],
-    libraries=['custatevec'],
+    library_dirs=cusv_lib_dir,
+    libraries=cusv_lib,
+    extra_link_args=extra_linker_flags,
 )
 
 
@@ -134,15 +177,15 @@ def check_cuda_version():
     sources=["cuquantum/cutensornet/cutensornet.pyx"],
     include_dirs=[os.path.join(cuda_path, 'include'),
                   os.path.join(cutensornet_root, 'include')],
-    library_dirs=[os.path.join(cutensornet_root, 'lib64'),
-                  os.path.join(cutensor_root, 'lib', cutensor_ver)],
-    libraries=['cutensornet', 'cutensor'],
+    library_dirs=cutn_lib_dir,
+    libraries=cutn_lib,
+    extra_link_args=extra_linker_flags,
 )
 
 
 setup(
     name="cuquantum-python",
-    version='0.1.0.0',  # the last digit is dedicated to cuQuantum Python
+    version=__version__,
     description="Python APIs for cuQuantum",
     url="https://github.com/NVIDIA/cuQuantum",
     author="NVIDIA Corporation",

diff --git a/python/tests/cuquantum_tests/cutensornet_tests/test_cutensornet.py b/python/tests/cuquantum_tests/cutensornet_tests/test_cutensornet.py
@@ -105,8 +105,10 @@ def test_get_version(self):
         assert ver == cutensornet.VERSION
 
     def test_get_cudart_version(self):
+        # CUDA runtime is statically linked, so we can't compare
+        # with the "runtime" version
         ver = cutensornet.get_cudart_version()
-        assert ver == cupy.cuda.runtime.runtimeGetVersion()
+        assert isinstance(ver, int)
 
 
 class TestHandle:

diff --git a/samples/custatevec/CMakeLists.txt b/samples/custatevec/CMakeLists.txt
@@ -1,6 +1,6 @@
-#  
-# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
 # 
+# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
 # 
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are
@@ -13,7 +13,7 @@
 #  - Neither the name(s) of the copyright holder(s) nor the names of its
 #    contributors may be used to endorse or promote products derived
 #    from this software without specific prior written permission.
-#
+# 
 # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 # "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 # LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 
@@ -24,44 +24,62 @@
 # DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 # THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+# ---[ Check cmake version.
 cmake_minimum_required(VERSION 3.13.0 FATAL_ERROR)
-if(CMAKE_VERSION VERSION_GREATER 3.18)
-  cmake_policy(SET CMP0104 OLD)
-  set(CUDA_ARCHITECTURES FALSE)
-endif()
 
 # ---[ Project specification.
 project(custatevec_example LANGUAGES C CXX CUDA)
 
 include(GNUInstallDirs)
 
 # ##########################################
-# cuStateVec_example build mode
+# custatevec_example build mode
 # ##########################################
 
-if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
-    message(STATUS "Setting build type to 'Debug' as none was specified.")
-    set(CMAKE_BUILD_TYPE "Debug" CACHE STRING "Choose the type of build." FORCE)
-    set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "" "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
+if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  message(STATUS "Setting build type to 'Debug' as none was specified.")
+  set(CMAKE_BUILD_TYPE "Debug" CACHE STRING "Choose the type of build.")
+  set_property(
+    CACHE
+    CMAKE_BUILD_TYPE
+    PROPERTY
+      STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo"
+  )
 else()
-    message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
+  message(STATUS "Build type: ${CMAKE_BUILD_TYPE}")
 endif()
 
 # ##########################################
-# cuStateVec_example dependencies
+# custatevec_example check-env utilities
 # ##########################################
 
-if (NOT CUSTATEVEC_ROOT)
-  set(CUSTATEVEC_ROOT ${CUDA_TOOLKIT_ROOT_DIR})
-endif()
+function(set_with_fallback VARIABLE FALLBACK)
+  if (NOT DEFINED ${VARIABLE} OR ${VARIABLE} STREQUAL "")
+    set(${VARIABLE} $ENV{${VARIABLE}} CACHE INTERNAL ${VARIABLE})
+    if (${VARIABLE} STREQUAL "")
+      if (NOT ${FALLBACK} STREQUAL "")
+        set(${VARIABLE} $ENV{${FALLBACK}} CACHE INTERNAL ${VARIABLE})
+      endif ()
+    endif ()
+  endif ()
+endfunction()
 
-include_directories(${CUSTATEVEC_ROOT}/include)
-link_directories(${CUSTATEVEC_ROOT}/lib64)
+# ##########################################
+# custatevec_example dependencies
+# ##########################################
+
+set_with_fallback(CUSTATEVEC_ROOT CUQUANTUM_ROOT)
+
+if (CUSTATEVEC_ROOT STREQUAL "")
+    message(FATAL_ERROR "Please set the environment variables CUSTATEVEC_ROOT or CUQUANTUM_ROOT to the path of the cuQuantum installation.")
+endif ()
+
+message(STATUS "Using CUSTATEVEC_ROOT = ${CUSTATEVEC_ROOT}")
 
 # ##########################################
-# cuStateVec_example building flags
+# custatevec_example global flags
 # ##########################################
 
 # Global CXX flags/options
@@ -75,39 +93,54 @@ set(CMAKE_CUDA_STANDARD 11)
 set(CMAKE_CUDA_STANDARD_REQUIRED ON)
 set(CMAKE_CUDA_EXTENSIONS OFF)
 
-set(CMAKE_CUDA_FLAGS_ARCH_SM60 "-gencode arch=compute_60,code=sm_60")
-set(CMAKE_CUDA_FLAGS_ARCH_SM70 "-gencode arch=compute_70,code=sm_70")
-set(CMAKE_CUDA_FLAGS_ARCH_SM75 "-gencode arch=compute_75,code=sm_75")
-set(CMAKE_CUDA_FLAGS_ARCH_SM80 "-gencode arch=compute_80,code=sm_80 -gencode arch=compute_80,code=compute_80")
-set(CMAKE_CUDA_FLAGS_ARCH "${CMAKE_CUDA_FLAGS_ARCH_SM60} ${CMAKE_CUDA_FLAGS_ARCH_SM70} ${CMAKE_CUDA_FLAGS_ARCH_SM75} ${CMAKE_CUDA_FLAGS_ARCH_SM80}")
-set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} ${CMAKE_CUDA_FLAGS_ARCH}")
-
 # ##########################################
-# cuStateVec_example target
+# custatevec_example utility function
 # ##########################################
 
 function(add_custatevec_example GROUP_TARGET EXAMPLE_NAME EXAMPLE_SOURCES)
-    list(GET EXAMPLE_SOURCES 0 EXAMPLE_MAIN_SOURCE)
-    get_filename_component(EXAMPLE_TARGET ${EXAMPLE_MAIN_SOURCE} NAME_WE)
-    add_executable(${EXAMPLE_TARGET} ${EXAMPLE_SOURCES})
-    target_include_directories(${EXAMPLE_TARGET}
-        PUBLIC
-            ${CUDA_INCLUDE_DIRS}
-            ${CUSTATEVEC_ROOT}/include
-    )
-    target_link_libraries(${EXAMPLE_TARGET}
-        PUBLIC custatevec cudart cublas cublasLt
-    )
-    # Install example
-    install(
-        TARGETS ${EXAMPLE_TARGET}
-        RUNTIME
-        DESTINATION ${CUSTATEVEC_EXAMPLE_BINARY_INSTALL_DIR}
-        PERMISSIONS OWNER_EXECUTE OWNER_WRITE OWNER_READ GROUP_EXECUTE GROUP_READ WORLD_EXECUTE WORLD_READ
-    )
-    add_dependencies(${GROUP_TARGET} ${EXAMPLE_TARGET})
+  list(GET EXAMPLE_SOURCES 0 EXAMPLE_MAIN_SOURCE)
+  get_filename_component(EXAMPLE_TARGET ${EXAMPLE_MAIN_SOURCE} NAME_WE)
+  add_executable(${EXAMPLE_TARGET} ${EXAMPLE_SOURCES})
+  target_include_directories(
+    ${EXAMPLE_TARGET}
+      PUBLIC
+        ${CUDA_INCLUDE_DIRS}
+        ${CUSTATEVEC_ROOT}/include
+  )
+  target_link_directories(
+    ${EXAMPLE_TARGET}
+      PUBLIC
+        ${CUSTATEVEC_ROOT}/lib
+        ${CUSTATEVEC_ROOT}/lib64
+  )
+  target_link_libraries(
+    ${EXAMPLE_TARGET}
+      PUBLIC
+        custatevec
+        cudart
+        cublas
+        cublasLt
+  )
+  set_target_properties(
+    ${EXAMPLE_TARGET}
+    PROPERTIES
+      CUDA_ARCHITECTURES
+        "70;75;80"
+  )
+  # Install example
+  install(
+      TARGETS ${EXAMPLE_TARGET}
+      RUNTIME
+      DESTINATION ${CUSTATEVEC_EXAMPLE_BINARY_INSTALL_DIR}
+      PERMISSIONS OWNER_EXECUTE OWNER_WRITE OWNER_READ GROUP_EXECUTE GROUP_READ WORLD_EXECUTE WORLD_READ
+  )
+  add_dependencies(${GROUP_TARGET} ${EXAMPLE_TARGET})
 endfunction()
 
+# ##########################################
+# custatevec_example add all targets
+# ##########################################
+
 add_custom_target(custatevec_examples)
 
 add_custatevec_example(custatevec_examples "cuStateVec.example.gate_application"   gate_application.cu)