Repository: cupy/cupy

Issue:
RFC: Support performant `einsum()` and `einsum_path()` using cuQuantum
### Description

Related to #6078.

The new [cuQuantum SDK](https://developer.nvidia.com/cuquantum-sdk) provides two libraries, one of which is [cuTensorNet](https://docs.nvidia.com/cuda/cuquantum/cutensornet/index.html) that accelerates tensor network contraction (backed by cuTENSOR). Specifically, cuTensorNet currently handles two major challenges from problems that have a huge Einsum expression:
1. Finding the optimal contraction path
2. Executing the actual pairwise contraction for a given path in an optimal fashion

Those two functionalities map very nicely to NumPy's `einsum_path()` and `einsum()` APIs, respectively. As a result, in [cuQuantum Python](https://docs.nvidia.com/cuda/cuquantum/python/), the Python binding for cuQuantum libraries, we provide pythonic APIs [`cuquantum.einsum_path()`](https://docs.nvidia.com/cuda/cuquantum/python/api/generated/cuquantum.einsum_path.html) and [`cuquantum.einsum()`](https://docs.nvidia.com/cuda/cuquantum/python/api/generated/cuquantum.einsum.html) among other things. We currently provide two conda packages, `cuquantum` and `cuquantum-python`, and a pip wheel release is in our near-term roadmap.

I would like to propose using cuQuantum Python to back CuPy's `einsum_path()` and `einsum()` when available. Performance improvement data can be shared if needed. However, there are a few issues to be addressed:
- Because CuPy is currently a required dependency of cuQuantum Python, we need to break any potential circular import.
- For the initial release, `cuquantum.einsum_path()` and `cuquantum.einsum()` are not yet a 100% drop-in replacement of their NumPy counterparts (we mark most optional arguments as unsupported).
- Directly using our drop-in replacement APIs has a small performance penalty (due to the library handle not reused).
- We currently only support classical einsum; ellipsis, broadcasting, etc are not yet supported 

Below is my 4-step proposal, involving a small refactoring of the existing codebase:
1. Add `cupy.einsum_path()` as a fallback: The functionalities are already in `cupy/linalg/_einsum.py`, we just need to put everything together.
2. Introduce `ACCELERATOR_CUQUANTUM` as a optional routine accelerator backend: This way, we allow users to set `CUPY_ACCELERATORS=cub,cutensor,cuquantum`. This is desired to be able to test things separately (ex: [see here](https://github.com/NVIDIA/cuQuantum/blob/4e63ab07d0d1bce2fb53e8eebcb3818d9444dbc5/python/tests/cuquantum_tests/cutensornet_tests/__init__.py#L4-L10)).
3. Use the core functionalities of cuTensorNet to back `cupy.einsum_path()` and `cupy.einsum()`: I don't think I'd like to use `cuquantum.einsum_path()` and `cuquantum.einsum()` directly due to the aforementioned reasons, so most likely we'll need to add some helper functions to check
   - If `ACCELERATOR_CUQUANTUM` is requested
   - If `cuquantum` can be imported
   - If all user inputs can be handled by cuTensorNet
   - If so, forward to cuQuantum Python, otherwise fall back 
4. Add some mock tests to ensure the backend can be used when available

### Additional Information

_No response_


Discussion:
The alternative of this proposal is to advice users to install cuQuantum Python and use it as is:
```python
#from cupy import einsum_path, einsum
from cuquantum import einsum_path, einsum

path, info = einsum_path(...)
out = einsum(..., optimize=path)
```
But I am not sure if this is in line with CuPy's interest 😅
> * Add `cupy.einsum_path()` as a fallback: The functionalities are already in `cupy/linalg/_einsum.py`, we just need to put everything together.

Looking deeper at https://github.com/cupy/cupy/blob/master/cupy/linalg/_einsum_opt.py I think this file should just be removed from the CuPy codebase and `cupy.einsum_path()` should just delegate to `numpy.einsum_path()` for this step because the code is simply copied from NumPy. This implementation is CPU based and we shouldn't waste any bandwidth in maintaining the duplication. (In fact, our experience shows it's complicated to implement a GPU-based path finder.)
@kmaehashi Did the team get a chance to discuss the plan? 🙂 In any case we should get part of the plan done to support `einsum_path`.
@leofang Sorry to get back to you late! We discussed and we are positive to introduce cuQuantum based `einsum` 😃 
The only concern is how to set up unit tests. It seems the easiest way to install cuQuantum is to use conda?
> It seems the easiest way to install cuQuantum is to use conda?

For today, yes, but wheels should be available soon... 😅
> For today, yes, but wheels should be available soon... 😅

@kmaehashi Well, we hit unexpected delay in the wheel release so we decided to skip cuQuantum v0.1.0 + cuTENSOR 1.4.0, and do cuQuantum v1.0.0 + cuTENSOR 1.5.0 directly. cuTENSOR wheel is up: https://pypi.org/project/cutensor/

(btw, given the lack of symlink support in wheels it's a bit awkward to build a project against the cuTENSOR wheel (do pip install yourself to see what I mean), so I am not sure if it'd be useful for CuPy... Also notice the wheel size...)
> (btw, given the lack of symlink support in wheels it's a bit awkward to build a project against the cuTENSOR wheel (do pip install yourself to see what I mean), so I am not sure if it'd be useful for CuPy... Also notice the wheel size...)

Can't try because on macOS right now, but interested to learn a bit more. Do you need more than the include and library directories for building against cuTENSOR? Those should at least be supportable by wheels. 
Hi @rgommers sorry I left the note in a rush last night. Happy to expand because this is one big challenge for distributing any binary repackaged wheels, not just NVIDIA products. This only happens with wheels not conda packages. 

If you do `pip install cutensor` and check what's installed, you'd see this:
```
$ tree $CONDA_PREFIX/lib/python3.8/site-packages/cutensor/
/home/leof/miniforge3/envs/opt_einsum_dev/lib/python3.8/site-packages/cutensor/
├── LICENSE
├── include
│   ├── cutensor
│   │   └── types.h
│   ├── cutensor.h
│   └── cutensorMg.h
└── lib
    ├── 10.2
    │   ├── libcutensor.so.1
    │   └── libcutensorMg.so.1
    ├── 11
    │   ├── libcutensor.so.1
    │   └── libcutensorMg.so.1
    └── 11.0
        ├── libcutensor.so.1
        └── libcutensorMg.so.1

6 directories, 10 files
```
Note that those `*.so.1` files are actual shared libraries, not symlinks. But this is a hack, because in other packaging formats we offer, there should exist these files
```
libcutensor.so        # symlink
libcutensor.so.1      # symlink
libcutensor.so.1.5.0  # actual file
```
In particular the first file is of utter importance as it allows passing `-lcutensor` to the linker. 

However, because when generating a wheel all symlinks are "materialized", we either triple the file size to keep all these (see also this comment https://github.com/pypa/wheel/issues/203#issuecomment-622598330) or trim the symlinks to lock the wheel size. The latter is what we do and our automation picked the `*.so.1` variant (for a good reason I can't recall now). As a result, in cuQuantum Python we had to build in a special way if cuTENSOR is coming from wheel: https://github.com/NVIDIA/cuQuantum/blob/79113afc3a52e82c5e681252e971452d8960983d/python/setup.py#L148
which is equivalent of passing `-l:libcutensor.so.1` to the linker.

This is awful, and is considered a temporary solution because it prevents downstream (non-NVIDIA, in this case) projects from naively consuming it, but until wheel can support symlinks (which was rejected back in 2017 straight away, see https://github.com/pypa/wheel/issues/203, and there's a WIP community effort https://github.com/karellen/wheel-axle-runtime & https://github.com/karellen/wheel-axle to the rescue) some hacks have to be done one way or another. I see no other way out...
Thanks for the detailed explanation @leofang! I wasn't aware of this symlink issue - that seems like a pretty serious design problem with wheels indeed. 

> and there's a WIP community effort https://github.com/karellen/wheel-axle-runtime & https://github.com/karellen/wheel-axle to the rescue) some hacks have to be done one way or another. I see no other way out...

Ingenious, but not something that anyone should be using right now for production-quality releases it looks like.
> Below is my 4-step proposal, involving a small refactoring of the existing codebase:
> 
> 1. Add `cupy.einsum_path()` as a fallback: The functionalities are already in `cupy/linalg/_einsum.py`, we just need to put everything together.

Spent some time trying to work on the proposed refactoring to add `einsum_path()` but I think for the time being we still need copy/paste from NumPy. I reported the reason and change request to NumPy here: https://github.com/numpy/numpy/issues/21379.

Resolving patch:
diff --git a/cupy/linalg/_einsum.py b/cupy/linalg/_einsum.py
--- a/cupy/linalg/_einsum.py
+++ b/cupy/linalg/_einsum.py
@@ -9,6 +9,7 @@
 from cupy import _util
 from cupy.linalg._einsum_opt import _greedy_path
 from cupy.linalg._einsum_opt import _optimal_path
+from cupy.linalg._einsum_cutn import _try_use_cutensornet
 
 
 try:
@@ -444,7 +445,7 @@ def _tuple_sorted_by_0(zs):
 
 
 def einsum(*operands, **kwargs):
-    """einsum(subscripts, *operands, dtype=False)
+    """einsum(subscripts, *operands, dtype=None, optimize=False)
 
     Evaluates the Einstein summation convention on the operands.
     Using the Einstein summation convention, many common multi-dimensional
@@ -452,21 +453,46 @@ def einsum(*operands, **kwargs):
     provides a way to compute such summations.
 
     .. note::
-       Memory contiguity of calculation result is not always compatible with
-       `numpy.einsum`.
-       ``out``, ``order``, and ``casting`` options are not supported.
+
+       - Memory contiguity of the returned array is not always compatible with
+         that of :func:`numpy.einsum`.
+       - ``out``, ``order``, and ``casting`` options are not supported.
+       - If :envvar:`CUPY_ACCELERATORS` includes ``cutensornet``, the `einsum`
+         calculation will be performed by the cuTensorNet backend if possible.
+
+           - The support of the ``optimize`` option is limited (currently, only
+             `False`, 'cutensornet', or a custom path for pairwise contraction
+             is supported, and the maximum intermediate size is ignored). If
+             you need finer control for path optimization, consider replacing
+             :func:`cupy.einsum` by :func:`cuquantum.contract` instead.
+           - Requires `cuQuantum Python`_ (v22.03+).
+
+       - If :envvar:`CUPY_ACCELERATORS` includes ``cutensor``, `einsum` will be
+         accelerated by the cuTENSOR backend whenever possible.
 
     Args:
         subscripts (str): Specifies the subscripts for summation.
         operands (sequence of arrays): These are the arrays for the operation.
+        dtype: If provided, forces the calculation to use the data type
+            specified. Default is None.
+        optimize: Valid options include {`False`, `True`, 'greedy', 'optimal'}.
+            Controls if intermediate optimization should occur. No optimization
+            will occur if `False`, and `True` will default to the 'greedy'
+            algorithm. Also accepts an explicit contraction list from
+            :func:`numpy.einsum_path`. Defaults to `False`. If a pair is
+            supplied, the second argument is assumed to be the maximum
+            intermediate size created.
 
     Returns:
         cupy.ndarray:
             The calculation based on the Einstein summation convention.
 
     .. seealso:: :func:`numpy.einsum`
-
+    .. _cuQuantum Python: https://docs.nvidia.com/cuda/cuquantum/python/
     """
+    out = _try_use_cutensornet(*operands, **kwargs)
+    if out is not None:
+        return out
 
     input_subscripts, output_subscript, operands = \
         _parse_einsum_input(operands)
diff --git a/cupy/linalg/_einsum_cutn.py b/cupy/linalg/_einsum_cutn.py
new file mode 100644
--- /dev/null
+++ b/cupy/linalg/_einsum_cutn.py
@@ -0,0 +1,141 @@
+import warnings
+
+try:
+    from cuquantum import cutensornet
+    cutn_handle_cache = {}  # type: ignore  # noqa
+except ImportError:
+    cutensornet = None
+
+import cupy
+from cupy._core import _accelerator
+
+
+def _get_einsum_operands(args):
+    """Parse & retrieve einsum operands, assuming ``args`` is in either
+    "subscript" or "interleaved" format.
+    """
+
+    if len(args) == 0:
+        raise ValueError(
+            'must specify the einstein sum subscripts string and at least one '
+            'operand, or at least one operand and its corresponding '
+            'subscripts list')
+
+    if isinstance(args[0], str):
+        expr = args[0]
+        operands = list(args[1:])
+        return expr, operands
+    else:
+        args = list(args)
+        operands = []
+        inputs = []
+        output = None
+        while len(args) >= 2:
+            operands.append(args.pop(0))
+            inputs.append(args.pop(0))
+        if len(args) == 1:
+            output = args.pop(0)
+        assert not args
+        return inputs, operands, output
+
+
+def _try_use_cutensornet(*args, **kwargs):
+    if cupy.cuda.runtime.is_hip:
+        return None
+
+    if (_accelerator.ACCELERATOR_CUTENSORNET not in
+            _accelerator.get_routine_accelerators()):
+        return None
+
+    if cutensornet is None:
+        warnings.warn(
+            'using the cuTensorNet backend was requested but it cannot be '
+            'imported -- maybe you forgot to install cuQuantum Python? '
+            'Please do "pip install cuquantum-python" or "conda install '
+            '-c conda-forge cuquantum-python" and retry',
+            stacklevel=2)
+        return None
+
+    # cannot pop as we might still need kwargs later
+    dtype = kwargs.get('dtype', None)
+    path = kwargs.get('optimize', False)
+    if path is True:
+        path = 'greedy'
+
+    # we do very lightweight pre-processing here just to inspect the
+    # operands; the actual input verification is deferred to cuTensorNet
+    # which can generate far better diagonostic messages
+    args = _get_einsum_operands(args)
+    operands = [cupy.asarray(op) for op in args[1]]
+
+    if len(operands) == 1:
+        # As of cuTENSOR 1.5.0 it still chokes with some common operations
+        # like trace ("ii->") so it's easier to just skip all single-operand
+        # cases instead of whitelisting what could be done explicitly
+        return None
+
+    if (any(op.size == 0 for op in operands) or
+            any(len(op.shape) == 0 for op in operands)):
+        # To cuTensorNet the shape is invalid
+        return None
+
+    # all input dtypes must be identical (to a numerical dtype)
+    result_dtype = cupy.result_type(*operands) if dtype is None else dtype
+    if result_dtype not in (
+            cupy.float32, cupy.float64, cupy.complex64, cupy.complex128):
+        return None
+    operands = [op.astype(result_dtype, copy=False) for op in operands]
+
+    # prepare cutn inputs
+    device = cupy.cuda.runtime.getDevice()
+    handle = cutn_handle_cache.get(
+        device, cutensornet.create())
+    cutn_options = {'device_id': device, 'handle': handle,
+                    'memory_limit': 4**31}  # TODO(leofang): fix?
+
+    # TODO(leofang): support all valid combinations:
+    # - path from user, contract with cutn (done)
+    # - path from cupy, contract with cutn (not yet)
+    # - path from cutn, contract with cutn (done)
+    # - path from cutn, contract with cupy (not yet)
+    raise_warning = False
+    if path is False:
+        # following the same convention (contracting from the right) as would
+        # be produced by _iter_path_pairs(), but converting to a list of pairs
+        # due to cuTensorNet's requirement
+        path = [(i-1, i-2) for i in range(len(operands), 1, -1)]
+    elif len(path) and path[0] == 'einsum_path':
+        # let cuTensorNet check if the format is correct
+        path = path[1:]
+    elif len(path) == 2:
+        if isinstance(path[1], (int, float)):
+            raise_warning = True
+        if path[0] != 'cutensornet':
+            raise_warning = True
+        path = None
+    else:  # path is a string
+        if path != 'cutensornet':
+            raise_warning = True
+        path = None
+    if raise_warning:
+        warnings.warn(
+            'the cuTensorNet backend ignores the "optimize" option '
+            'except when an explicit contraction path is provided '
+            'or when optimize=False (disable optimization); also, '
+            'the maximum intermediate size, if set, is ignored',
+            stacklevel=2)
+    cutn_optimizer = {'path': path} if path else None
+
+    if len(args) == 2:
+        out = cutensornet.contract(
+            args[0], *operands, options=cutn_options, optimize=cutn_optimizer)
+    elif len(args) == 3:
+        inputs = [i for pair in zip(operands, args[0]) for i in pair]
+        if args[2] is not None:
+            inputs.append(args[2])
+        out = cutensornet.contract(
+            *inputs, options=cutn_options, optimize=cutn_optimizer)
+    else:
+        assert False
+
+    return out
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -368,6 +368,7 @@
     'numpy': ('https://numpy.org/doc/stable/', None),
     'scipy': ('https://docs.scipy.org/doc/scipy/', None),
     'numba': ('https://numba.readthedocs.io/en/stable', None),
+    'cuquantum': ('https://docs.nvidia.com/cuda/cuquantum/', None),
     # blocked by data-apis/array-api#428
     #'array-api': ('https://data-apis.org/array-api/2021.12/', None),
 }


Test patch: