Repository: cupy/cupy
Pull request: Fix cuTENSOR contraction descriptor cache
Close #7318. Close #7812.

The mode labels encapsulated by the `Mode` object should be cached, not just the pointer to it as the same memory could be garbage collected and reused, leading to erroneous cache hits. 

The added tests would fail without the fix.

Problem statement:
Wrong answer got from `contraction` function in `cutensor` lib
### Description

In some case, the `contraction` function in `cutensor` lib will give the wrong answer. This problem is very likely to show up when call several `contraction` function sequencely and with similar contraction order. See the code bellow for example. There are several ways to correct the code:

1. Change `b = contraction('lijk', a, 'jm', t2, 'likm', buf=buf)` to `b = contraction( 'jm', t2, 'lijk', a,'likm', buf=buf)`. This only swap the contraction parameter `a` and `t1`, which should not affect the answer but make the contraction paremeters to get `b` and `c` no longer similar.
2. Swap the contraction parameter of `c` similarly.
3. Do not use the buf, i.e. Change `b = contraction('lijk', a, 'jm', t2, 'likm', buf=buf)` to `b = contraction('lijk', a, 'jm', t2, 'likm')`.
4. Comment out the 443 and 473 line in `cupy/cupyx/cutensor.pyx`. `443: # _contraction_descriptors[key] = desc`, `473: # _contraction_plans[key] = plan`. This disable the 'cache' mechanism for `contraction_descriptors` and `contraction_plans`.

 After several months' investagation, i think this is very likely related to the 'cache' mechanism in the `cupy/cupyx/cutensor.pyx`, i.e. the following dicts:
```
29: cdef dict _handles = {}
30: cdef dict _tensor_descriptors = {}
31: cdef dict _contraction_descriptors = {}
32: cdef dict _contraction_finds = {}
33: cdef dict _contraction_plans = {}
34: cdef dict _modes = {}
35: cdef dict _scalars = {}
```
Perhaps, `cutensor` lib automaticly destroy the object cached in these dicts causeing this problem. I am not sure which specific dicts cause the problem (disableign all of them can solve the problem of course but less efficient).

### To Reproduce

```py
import numpy as np
import cupy
from cupyx import cutensor

def contraction(inda, a, indb, b, indc, buf=None):
    modea = tuple(inda)
    modeb = tuple(indb)
    modec = tuple(indc)
    dtype = np.result_type(a, b)
    cshape = [a.shape[modea.index(
        c)] if c in modea else b.shape[modeb.index(c)] for c in modec]
    if buf is None:
        c = cupy.empty(cshape, dtype=dtype)
    else:
        assert np.prod(buf.shape) >= np.prod(cshape), 'buf is not enough'
        c = cupy.ndarray(cshape, dtype=dtype, memptr=buf.data)
    desca = cutensor.create_tensor_descriptor(a)
    descb = cutensor.create_tensor_descriptor(b)
    descc = cutensor.create_tensor_descriptor(c)
    c = cutensor.contraction(1.0, a, desca, cutensor.create_mode(
        *modea), b, descb, cutensor.create_mode(*modeb), 0.0, c, descc, cutensor.create_mode(*modec))
    return c

if __name__ == "__main__":
    n = 10
    t1 = cupy.random.rand(n, n, n, n)
    t2 = cupy.random.rand(n, n)

    # cutensor
    buf = cupy.empty(n**4)
    a = contraction('ijkl', t1, 'km', t2, 'lijm')
    b = contraction('lijk', a, 'jm', t2, 'likm', buf=buf)
    c = contraction('likj', b, 'im', t2, 'ljkm')
    print('a:', a.max(), 'b:', b.max(), 'c:', c.max())
    buf = a = b = c = None

    #cupy dot
    a = cupy.dot(t1.transpose((3, 0, 1, 2)), t2)
    b = cupy.dot(a.transpose((0, 1, 3, 2)), t2)
    c = cupy.dot(b.transpose((0, 3, 2, 1)), t2)
    print('a:', a.max(), 'b:', b.max(), 'c:', c.max())
```


### Installation

Built from GitHub source

### Environment

```
OS                           : Linux-5.4.143.bsk.7-amd64-x86_64-with-glibc2.28
Python Version               : 3.9.15
CuPy Version                 : 12.0.0b2
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.21.5
SciPy Version                : 1.7.3
Cython Build Version         : 0.29.32
Cython Runtime Version       : 0.29.32
CUDA Root                    : /usr/local/cuda
nvcc PATH                    : /usr/local/cuda/bin/nvcc
CUDA Build Version           : 11030
CUDA Driver Version          : 11030
CUDA Runtime Version         : 11030
cuBLAS Version               : (available)
cuFFT Version                : 10402
cuRAND Version               : 10204
cuSOLVER Version             : (11, 1, 2)
cuSPARSE Version             : (available)
NVRTC Version                : (11, 3)
Thrust Version               : 101100
CUB Build Version            : 101100
Jitify Build Version         : 4a37de0
cuDNN Build Version          : 8700
cuDNN Version                : 8700
NCCL Build Version           : 21104
NCCL Runtime Version         : 21104
cuTENSOR Version             : 10602
cuSPARSELt Build Version     : None
Device 0 Name                : A100-SXM-80GB
Device 0 Compute Capability  : 80
Device 0 PCI Bus ID          : 0000:4E:00.0
```


### Additional Information


cutensor crash in the second run
### Description

The first run of cutensor.contraction gives the identical result with cupy.einsum. However, the second run gives the incorrect result. If I comment out the first run and run the second contraction, it gives the identical result with cupy.einsum

### To Reproduce

```py
import cupy
from cupyx import cutensor

def contract(pattern, a, b):
    pattern = pattern.replace(" ", "")
    str_a, rest = pattern.split(',')
    str_b, str_c = rest.split('->')
    key = str_a + str_b
    val = list(a.shape) + list(b.shape)
    shape = {k:v for k, v in zip(key, val)}
    c = cupy.zeros([shape[k] for k in str_c], order='C')

    mode_a = list(str_a)
    mode_b = list(str_b)
    mode_c = list(str_c)

    desc_a = cutensor.create_tensor_descriptor(a)
    desc_b = cutensor.create_tensor_descriptor(b)
    desc_c = cutensor.create_tensor_descriptor(c)
    
    mode_a = cutensor.create_mode(*mode_a)
    mode_b = cutensor.create_mode(*mode_b)
    mode_c = cutensor.create_mode(*mode_c)
    alpha = 1.0
    beta = 0.0
    
    cutensor.contraction(alpha, a, desc_a, mode_a, b, desc_b, mode_b, beta, c, desc_c, mode_c)
    return c

a = cupy.random.rand(20,20,20,20)
b = cupy.random.rand(20,20)

c = contract('ijkl,ik->jl', a, b)
c_einsum = cupy.einsum('ijkl,ik->jl', a, b)
print(cupy.linalg.norm(c - c_einsum))

c = contract('ijkl,jl->ik', a, b)
c_einsum = cupy.einsum('ijkl,jl->ik', a, b)
print(cupy.linalg.norm(c - c_einsum))
```


### Installation

Wheel (`pip install cupy-***`)

### Environment

```
OS                           : Linux-5.4.143.bsk.8-amd64-x86_64-with-glibc2.31
Python Version               : 3.9.2
CuPy Version                 : 12.2.0
CuPy Platform                : NVIDIA CUDA
NumPy Version                : 1.24.2
SciPy Version                : 1.10.0
Cython Build Version         : 0.29.36
Cython Runtime Version       : 0.29.33
CUDA Root                    : /usr/local/cuda
nvcc PATH                    : /usr/local/cuda/bin/nvcc
CUDA Build Version           : 11080
CUDA Driver Version          : 11070
CUDA Runtime Version         : 11070
cuBLAS Version               : (available)
cuFFT Version                : 10702
cuRAND Version               : 10210
cuSOLVER Version             : (11, 4, 0)
cuSPARSE Version             : (available)
NVRTC Version                : (11, 7)
Thrust Version               : 101501
CUB Build Version            : 101501
Jitify Build Version         : <unknown>
cuDNN Build Version          : 8801
cuDNN Version                : 8500
NCCL Build Version           : 21602
NCCL Runtime Version         : 21304
cuTENSOR Version             : 10602
cuSPARSELt Build Version     : None
Device 0 Name                : Tesla V100-SXM2-32GB
Device 0 Compute Capability  : 70
Device 0 PCI Bus ID          : 0000:1B:00.0
```


### Additional Information

_No response_

Hints from issue discussion:
I think this bug is related to a similar caching scheme in cuda called "Plan Cache" (https://docs.nvidia.com/cuda/cutensor/plan_cache.html) that is automatically enabled. The cupy caching scheme may conflict with the plan cache. I propose to disable all caching schemes in cupy.cutensor. 
@jemiryguo so sorry for late response, and thank you for bringing this to our attention and your time on investigating the issue for us. 

As I pointed out in another bug report (https://github.com/cupy/cupy/issues/7812#issuecomment-1686254600), you were right that it's due to erroneous cache hits. CuPy only caches the pointer addresses of the mode labels, not their entire contents, so depending on the contraction problem & gc the pointer addresses could be reused but with different mode labels. This will be fixed by #7814.
I reproduced the case, ~and the second `c_einsum` seems to have the wrong value as you pointed out.~

cc/ @leofang 
I noticed that in the 2nd call `cutensorInitContractionDescriptor` is not called, which is wrong because the two contractions are distinct and should not reuse the same descriptor created/cached from the first call. Will look into the cache issue later today or tomorrow.
OK I figured out. It has to do with the interplay with garbage collection. Consider this code:
https://github.com/cupy/cupy/blob/db0be4f6be5f849dcba0a9ca2cb5f840d8028923/cupyx/cutensor.pyx#L425-L431
`mode_A` & co are created, their pointers cached, and then destroyed due to out of scope after the 1st call is done. Then, in the 2nd call `mode_B` and `mode_C` are swapped, but they are still created in order, and the pointer address are _identical_, leading to a *cache hit* (but should have been a *miss*):
```
$ CUTENSOR_LOG_LEVEL=0 python ../cupy_cutensor.py 
['i', 'j', 'k', 'l'] ['i', 'k'] ['j', 'l']
mode(105, 106, 107, 108) 94562718015312
mode(105, 107) 94562634127728
mode(106, 108) 94562633847120
key (94562711502672, 16, 94562718497632, 94562718015312, 16, 94562706419168, 94562634127728, 16, 94562706419168, 94562633847120, 16) is not found
3.410605131648481e-13

['i', 'j', 'k', 'l'] ['j', 'l'] ['i', 'k']
mode(105, 106, 107, 108) 94562718015312
mode(106, 108) 94562634127728
mode(105, 107) 94562633847120
key (94562711502672, 16, 94562718497632, 94562718015312, 16, 94562706419168, 94562634127728, 16, 94562706419168, 94562633847120, 16) is found
99.923464935204
```  
@kmaehashi @takagi feel free to assign to me if you'd like me to to send a patch 🙂
@asi1024 I suspect #7318 has the same root cause (I didn't check carefully, but it's plausible).
(permission tested)
> I suspect https://github.com/cupy/cupy/issues/7318 has the same root cause (I didn't check carefully, but it's plausible).

Confirmed.
Fantastic! @leofang Thank you for the prompt response.

Commit message: add tests

Patch to review:
diff --git a/cupyx/cutensor.pyx b/cupyx/cutensor.pyx
index a672a8eb914..e157ae3f286 100644
--- a/cupyx/cutensor.pyx
+++ b/cupyx/cutensor.pyx
@@ -122,6 +122,14 @@ cdef class Mode:
     def __repr__(self):
         return 'mode(' + ', '.join([str(x) for x in self._array]) + ')'
 
+    def __eq__(self, other):
+        if not isinstance(other, Mode):
+            return False
+        return (self._array == (<Mode>other)._array).all()
+
+    def __hash__(self):
+        return hash(bytes(self._array))
+
 
 cdef class _Scalar:
 
@@ -423,9 +431,9 @@ cdef inline ContractionDescriptor _create_contraction_descriptor(
     cdef ContractionDescriptor desc
 
     key = (handle.ptr, cutensor_compute_type,
-           desc_A.ptr, mode_A.data, alignment_req_A,
-           desc_B.ptr, mode_B.data, alignment_req_B,
-           desc_C.ptr, mode_C.data, alignment_req_C)
+           desc_A.ptr, mode_A, alignment_req_A,
+           desc_B.ptr, mode_B, alignment_req_B,
+           desc_C.ptr, mode_C, alignment_req_C)
     if key in _contraction_descriptors:
         desc = _contraction_descriptors[key]
         return desc
diff --git a/tests/cupyx_tests/test_cutensor.py b/tests/cupyx_tests/test_cutensor.py
index 0520490feca..a573c8ea3dc 100644
--- a/tests/cupyx_tests/test_cutensor.py
+++ b/tests/cupyx_tests/test_cutensor.py
@@ -1,3 +1,5 @@
+import gc
+
 import numpy
 import pytest
 
@@ -181,6 +183,18 @@ def test_create_mode_ascii(self):
         assert m.ndim == 2
         assert repr(m) == 'mode(120, 121)'
 
+    def test_mode_compare(self):
+        m1 = cutensor.create_mode(10, 11, 12)
+        m2 = cutensor.create_mode(10, 11, 12)
+        assert m1 == m2
+        assert hash(m1) == hash(m2)
+        assert m1.data != m2.data
+
+        m2 = cutensor.create_mode(12, 11, 10)
+        assert m1 != m2
+        assert hash(m1) != hash(m2)
+        assert m1.data != m2.data
+
 
 @pytest.mark.skipif(not ct.available, reason='cuTensor is unavailable')
 class TestScalar:
@@ -273,7 +287,7 @@ def test_reduction(self):
 @testing.parameterize(*testing.product({
     'dtype_combo': ['eee', 'fff', 'ddd', 'FFF', 'DDD', 'dDD', 'DdD'],
     'compute_type_hint': [None, 'down-convert', 'TF32'],
-    'shape': [(40, 30, 20)],
+    'shape': [(40, 20, 20)],  # let last two dim be the same for testing cache
     'alpha': [1.0],
     'beta': [0.0, 1.0],
 }))
@@ -342,3 +356,17 @@ def test_contraction(self):
                              self.c, desc_c, mode_c)
         cupy.testing.assert_allclose(self.c, self.c_ref,
                                      rtol=self.tol, atol=self.tol)
+
+        # test the contraction descriptor cache (issues #7318, #7812)
+        del mode_b
+        gc.collect()
+        mode_b = cutensor.create_mode('n', 'k')  # flipped
+        self.c_ref = self.alpha * cupy.matmul(self.a, self.b.T)
+        self.c_ref += self.beta * self.c
+        cutensor.contraction(self.alpha,
+                             self.a, desc_a, mode_a,
+                             self.b, desc_b, mode_b,
+                             self.beta,
+                             self.c, desc_c, mode_c)
+        cupy.testing.assert_allclose(self.c, self.c_ref,
+                                     rtol=self.tol, atol=self.tol)


Code review comments:
diff_hunk: @@ -423,9 +431,9 @@ cdef inline ContractionDescriptor _create_contraction_descriptor(
     cdef ContractionDescriptor desc
 
     key = (handle.ptr, cutensor_compute_type,
-           desc_A.ptr, mode_A.data, alignment_req_A,
-           desc_B.ptr, mode_B.data, alignment_req_B,
-           desc_C.ptr, mode_C.data, alignment_req_C)
+           desc_A.ptr, mode_A, alignment_req_A,
+           desc_B.ptr, mode_B, alignment_req_B,
+           desc_C.ptr, mode_C, alignment_req_C)
path: cupyx/cutensor.pyx
text: @author:
I feel it could be error prone to cache pointers that CuPy does not own (just like the modes), so perhaps we shouldn't use `desc.ptr` here either? Let me look into that.

@author:
For tensor descriptors, `create_tensor_descriptor()` would cache the created descriptor, so `desc.ptr` would not be invalidated. Perhaps we should do the same thing in `create_modes()` and cache the `Mode` object, then `mode.data` would stay valid too.

@author:
Done 8af1d6f.