Make kernel instance on class level.
@@ -15,6 +15,11 @@ from cupy.sparse import util
class _compressed_sparse_matrix(sparse_data._data_matrix):
+ _compress_getitem_kern = cupy.ElementwiseKernel(
+ 'T d, S ind, int32 minor', 'raw T answer',
+ 'if (ind == minor) atomicAdd(&answer[0], d);',
+ 'compress_getitem', preamble=util._preamble_atomic_add)
+
def __init__(self, arg1, shape=None, dtype=None, copy=False):
if shape is not None and len(shape) != 2:
raise ValueError(
@@ -201,11 +206,8 @@ class _compressed_sparse_matrix(sparse_data._data_matrix):
start = self.indptr[major]
end = self.indptr[major + 1]
answer = cupy.zeros((), self.dtype)
- kern = cupy.ElementwiseKernel(
- 'T d, S ind, int32 minor', 'raw T answer',
- 'if (ind == minor) atomicAdd(&answer[0], d);',
- 'compress_getitem', preamble=util._preamble_atomic_add)
- kern(self.data[start:end], self.indices[start:end], minor, answer)
+ self._compress_getitem_kern(
+ self.data[start:end], self.indices[start:end], minor, answer)
return answer[()]
def _get_major_slice(self, major):
