sync with internal repo (commit d92ae646)

Author: leofang

LICENSE

@@ -1,6 +1,6 @@
 BSD-3-Clause
 
-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 met:

python/.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+*.egg-info
+*.eggs
+*.cpp
+*.so

python/LICENSE

@@ -1,6 +1,6 @@
 BSD-3-Clause
 
-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 met:

python/README.md

@@ -11,10 +11,11 @@ Please visit the [NVIDIA cuQuantum Python documentation](https://docs.nvidia.com
 Build-time dependencies of the cuQuantum Python package and some versions that
 are known to work are as follows:
 
-* CUDA Toolkit 11.4+
-* cuQuantum 0.1.0
-* cuTENSOR 1.4.0+
+* CUDA Toolkit 11.x
+* cuQuantum 22.03
+* cuTENSOR 1.5.0+
 * Cython - e.g. 0.29.21
+* [packaging](https://packaging.pypa.io/en/latest/)
 
 ### Install cuQuantum Python from conda-forge
 
@@ -24,6 +25,18 @@ conda install -c conda-forge cuquantum-python
 ```
 The Conda solver will install all required dependencies for you.
 
+### Install cuQuantum Python from PyPI
+
+Alternatively, assuming you already have a Python environment set up (it doesn't matter if it's a Conda env or not),
+you can also install cuQuantum Python this way:
+
+```
+pip install cuquantum-python
+```
+The `pip` solver will also install both cuTENSOR and cuQuantum for you.
+
+Note: To properly install the wheels the environment variable `CUQUANTUM_ROOT` must not be set.
+
 ### Install cuQuantum Python from source
 
 To compile and install cuQuantum Python from source, please follow the steps below:
@@ -48,9 +61,9 @@ Runtime dependencies of the cuQuantum Python package include:
 
 * An NVIDIA GPU with compute capability 7.0+
 * Driver: Linux (450.80.02+)
-* CUDA Toolkit 11.4+
-* cuQuantum 0.1.0
-* cuTENSOR 1.4.0+
+* CUDA Toolkit 11.x
+* cuQuantum 22.03
+* cuTENSOR 1.5.0+
 * NumPy v1.17+
 * CuPy v9.5.0+
 * PyTorch v1.10+ (optional)
@@ -77,4 +90,5 @@ library in Python.
 ## Testing
 
 If pytest is installed, run `pytest tests` in the Python source root directory would
-run all tests.
+run all tests. Some tests would be skipped if `cffi` is not installed or if the environment
+variable `CUDA_PATH` is not set.

python/cuquantum/__init__.py

@@ -1,7 +1,7 @@
 from cuquantum import custatevec
 from cuquantum import cutensornet
 from cuquantum.cutensornet import (
-    contract, contract_path, einsum, einsum_path, Network,
+    contract, contract_path, einsum, einsum_path, Network, BaseCUDAMemoryManager, MemoryPointer,
     NetworkOptions, OptimizerInfo, OptimizerOptions, PathFinderOptions, ReconfigOptions, SlicerOptions)
 from cuquantum.utils import ComputeType, cudaDataType, libraryPropertyType
 from cuquantum._version import __version__
@@ -17,14 +17,11 @@
         cutensornet.ContractionOptimizerInfoAttribute,
         cutensornet.ContractionOptimizerConfigAttribute,
         cutensornet.ContractionAutotunePreferenceAttribute,
+        cutensornet.WorksizePref,
+        cutensornet.Memspace,
+        cutensornet.GraphAlgo,
+        cutensornet.MemoryModel,
         ):
     cutensornet._internal.enum_utils.add_enum_class_doc(enum, chomp="_ATTRIBUTE|_PREFERENCE_ATTRIBUTE")
-# these have yet another convention...
-for v in cutensornet.GraphAlgorithm:
-    v.__doc__ = f"See `CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_GRAPH_ALGORITHM_{v.name}`."
-cutensornet.MemoryModel.SLICER_HEURISTIC.__doc__ = \
-    f"See `CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_MODEL_HEURISTIC`."
-cutensornet.MemoryModel.SLICER_CUTENSOR.__doc__ = \
-    f"See `CUTENSORNET_CONTRACTION_OPTIMIZER_CONFIG_SLICER_MEMORY_MODEL_CUTENSOR`."
 
-del enum, utils, v
+del enum, utils

python/cuquantum/_version.py

@@ -1,2 +1,4 @@
-# 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
+# Note: cuQuantum Python follows the cuQuantum SDK version, which is now
+# switched to YY.MM and is different from individual libraries' (semantic)
+# versioning scheme.
+__version__ = '22.03.0'  # the last digit is for cuQuantum Python only

python/cuquantum/custatevec/custatevec.pxd

@@ -3,28 +3,42 @@
 
 from libc.stdint cimport intptr_t, int32_t, uint32_t, int64_t
 
+from cuquantum.utils cimport Stream
+
 
 # The C types are prefixed with an underscore because we are not
 # yet protected by the module namespaces as done in CUDA Python.
 # Once we switch over the names would be prettier (in the Cython
 # layer).
 
+
+# Cython limitation: need standalone typedef if we wanna use it for casting
+ctypedef int (*DeviceAllocType)(void*, void**, size_t, Stream)
+ctypedef int (*DeviceFreeType)(void*, void*, size_t, Stream)
+
+
 cdef extern from '<custatevec.h>' nogil:
     # cuStateVec types
     ctypedef void* _Handle 'custatevecHandle_t'
     ctypedef int64_t _Index 'custatevecIndex_t'
     ctypedef int _Status 'custatevecStatus_t'
-    ctypedef struct _SamplerDescriptor 'custatevecSamplerDescriptor_t':
-        pass
-    ctypedef struct _AccessorDescriptor 'custatevecAccessorDescriptor':
-        pass
+    ctypedef void* _SamplerDescriptor 'custatevecSamplerDescriptor_t'
+    ctypedef void* _AccessorDescriptor 'custatevecAccessorDescriptor_t'
     ctypedef enum _ComputeType 'custatevecComputeType_t':
         pass
-    # ctypedef void(*custatevecLoggerCallback_t)(
-    #     int32_t logLevel,
-    #     const char* functionName,
-    #     const char* message)
-    # ctypedef custatevecLoggerCallback_t LoggerCallback
+    ctypedef struct _DeviceMemHandler 'custatevecDeviceMemHandler_t':
+        void* ctx
+        DeviceAllocType device_alloc
+        DeviceFreeType device_free
+
+        # Cython limitation: cannot use C defines in declaring a static array,
+        # so we just have to hard-code CUSTATEVEC_ALLOCATOR_NAME_LEN here...
+        char name[64]
+    ctypedef void(*LoggerCallbackData 'custatevecLoggerCallbackData_t')(
+        int32_t logLevel,
+        const char* functionName,
+        const char* message,
+        void* userData)
 
     # cuStateVec enums
     ctypedef enum _Pauli 'custatevecPauli_t':
@@ -55,3 +69,4 @@ cdef extern from '<custatevec.h>' nogil:
     int CUSTATEVEC_VER_MINOR
     int CUSTATEVEC_VER_PATCH
     int CUSTATEVEC_VERSION
+    int CUSTATEVEC_ALLOCATOR_NAME_LEN