From dc44d6ee25f280318da4f7a2f16c2932719d9c24 Mon Sep 17 00:00:00 2001 From: Dieter Weber Date: Tue, 10 Jan 2023 11:01:06 +0100 Subject: [PATCH] Introduce `result_type()` function * Introduce `result_type()` to find the smallest NumPy dtype that accomodates all parameters. * Support `cupyx.scipy.sparse.csr_matrix` with `dtype=bool`. --- README.md | 14 +- src/sparseconverter/__init__.py | 280 ++++++++++++++++++++++---------- tests/test_sparseconverters.py | 47 +++++- 3 files changed, 255 insertions(+), 86 deletions(-) diff --git a/README.md b/README.md index d0602f0..fd0165c 100644 --- a/README.md +++ b/README.md @@ -70,11 +70,23 @@ and to perform efficient conversion to supported formats as needed. ## Still TODO -* cupyx.sparse formats with dtype `bool` * PyTorch arrays * SciPy sparse arrays as opposed to SciPy sparse matrices. * More detailed cost metric based on more real-world use cases and parameters. +## Changelog + +### 0.2.0 + +* Introduce `result_type()` to find the smallest NumPy dtype that accomodates + all parameters. Allowed as parameters are all valid arguments to + `numpy.result_type(...)` plus backend specifiers. +* Support `cupyx.scipy.sparse.csr_matrix` with `dtype=bool`. + +### 0.1.1 + +Initial release + ## Known issues * `conda install -c conda-forge cupy` on Python 3.7 and Windows 11 may install `cudatoolkit` 10.1 and `cupy` 8.3, which have sporadically produced invalid data structures for `cupyx.sparse.csc_matrix` for unknown reasons. This doesn't happen with current versions. Running the benchmark function `benchmark_conversions()` can help to debug such issues since it performs all pairwise conversions and checks for correctness. diff --git a/src/sparseconverter/__init__.py b/src/sparseconverter/__init__.py index a5e396e..d2d117f 100644 --- a/src/sparseconverter/__init__.py +++ b/src/sparseconverter/__init__.py @@ -3,7 +3,6 @@ import time from typing import TYPE_CHECKING, Callable, Dict, Iterable, Optional, Tuple, Union, Hashable from typing_extensions import Literal -import warnings import numpy as np import scipy.sparse as sp @@ -146,6 +145,44 @@ def _get_lazy(self, item): CUPY_SCIPY_CSR: np.float32, } +# Exceptions for the above matrix +_special_mappings = { + (CUPY_SCIPY_CSR, np.dtype(bool)): bool, +} + + +def result_type(*args) -> np.dtype: + ''' + Find a dtype that fulfills the following properties: + + * Can contain :code:`numpy.result_type(...)` of all items in :code:`args` + that are not a backend specifier. + * Supported by all items in :code:`args` that are backend specifiers + ''' + backends = [] + others = [] + + for a in args: + try: + if a in BACKENDS: + backends.append(a) + else: + others.append(a) + except TypeError: + others.append(a) + if others: + result_dtype = np.result_type(*others) + else: + result_dtype = np.dtype(bool) + stable = False + while not stable: + prev = result_dtype + for b in backends: + typ = _special_mappings.get((b, result_dtype), _base_dtypes[b]) + result_dtype = np.result_type(result_dtype, typ) + stable = prev == result_dtype + return result_dtype + def prod(shape: Tuple) -> int: ''' @@ -339,6 +376,32 @@ def fail_coo_complex128(arr): np.float32, np.float64, np.complex64, np.complex128 } + CUPY_SPARSE_CSR_DTYPES = { + np.dtype(bool), np.float32, np.float64, np.complex64, np.complex128 + } + + def _adjust_dtype_cupy_sparse(array_backend: ArrayBackend): + ''' + return dtype upconversion function for the specified cupyx.scipy.sparse backend. + ''' + if array_backend == CUPY_SCIPY_CSR: + allowed = CUPY_SPARSE_CSR_DTYPES + elif array_backend in (CUPY_SCIPY_COO, CUPY_SCIPY_CSC): + allowed = CUPY_SPARSE_DTYPES + else: + allowed = BACKENDS + + def adjust(arr): + if arr.dtype in allowed: + res = arr + # print('_adjust_dtype_cupy_sparse passthrough', arr.dtype, res.dtype) + else: + # Base dtype is the same for all cupyx.scipy.sparse matrices + res = arr.astype(result_type(arr, array_backend)) + return res + + return adjust + def _GCXS_to_cupy_coo(arr: sparse.GCXS): reshaped = arr.reshape((arr.shape[0], -1)) return cupyx.scipy.sparse.coo_matrix(reshaped.to_scipy_sparse()) @@ -378,7 +441,7 @@ def _CUPY_to_scipy_csr(arr: cupy.ndarray): Use GPU for sparsification if dtype allows, otherwise CPU ''' reshaped = arr.reshape((arr.shape[0], -1)) - if arr.dtype in CUPY_SPARSE_DTYPES: + if arr.dtype in CUPY_SPARSE_CSR_DTYPES: intermediate = cupyx.scipy.sparse.csr_matrix(reshaped) return intermediate.get() else: @@ -413,7 +476,7 @@ def _CUPY_to_sparse_gcxs(arr: cupy.ndarray): ''' Use GPU for sparsification if dtype allows, otherwise CPU ''' - if arr.dtype in CUPY_SPARSE_DTYPES: + if arr.dtype in CUPY_SPARSE_CSR_DTYPES: reshaped = arr.reshape((arr.shape[0], -1)) intermediate = cupyx.scipy.sparse.csr_matrix(reshaped) return sparse.GCXS(intermediate.get()).reshape(arr.shape) @@ -433,6 +496,62 @@ def _CUPY_to_sparse_dok(arr: cupy.ndarray): intermediate = cupy.asnumpy(arr) return sparse.DOK.from_numpy(intermediate) + _adjust_dtype_for_cupy_scipy_coo = _adjust_dtype_cupy_sparse(CUPY_SCIPY_COO) + _adjust_dtype_for_cupy_scipy_csr = _adjust_dtype_cupy_sparse(CUPY_SCIPY_CSR) + _adjust_dtype_for_cupy_scipy_csc = _adjust_dtype_cupy_sparse(CUPY_SCIPY_CSC) + + def _CUPY_to_cupy_scipy_coo(arr: cupy.array): + if has_cupy_coo_construction: + return cupyx.scipy.sparse.coo_matrix( + _adjust_dtype_for_cupy_scipy_coo( + _flatsig(arr) + ) + ) + elif not has_csr_complex128_bug or arr.dtype != np.complex128: + return cupyx.scipy.sparse.coo_matrix( + cupyx.scipy.sparse.csr_matrix( + _adjust_dtype_for_cupy_scipy_coo( + _flatsig(arr) + ) + ) + ) + else: + fail_coo_complex128(arr) + + def _CUPY_to_cupy_scipy_csr(arr: cupy.array): + if has_csr_complex128_bug and arr.dtype == np.complex128: + if has_cupy_coo_construction: + return cupyx.scipy.sparse.csr_matrix( + cupyx.scipy.sparse.coo_matrix( + _flatsig(arr) + ) + ) + else: + return fail_coo_complex128(arr) + else: + return cupyx.scipy.sparse.csr_matrix( + _adjust_dtype_for_cupy_scipy_csr( + _flatsig(arr) + ) + ) + + def _CUPY_to_cupy_scipy_csc(arr: cupy.array): + if has_csr_complex128_bug and arr.dtype == np.complex128: + if has_cupy_coo_construction: + return cupyx.scipy.sparse.csc_matrix( + cupyx.scipy.sparse.coo_matrix( + _flatsig(arr) + ) + ) + else: + return fail_coo_complex128(arr) + else: + return cupyx.scipy.sparse.csc_matrix( + _adjust_dtype_for_cupy_scipy_csc( + _flatsig(arr) + ) + ) + def _sparse_coo_to_CUPY(arr: sparse.COO): ''' Use GPU for densification if dtype allows, otherwise CPU @@ -449,12 +568,13 @@ def _sparse_gcxs_to_CUPY(arr: sparse.GCXS): ''' Use GPU for densification if dtype allows, otherwise CPU ''' - if arr.dtype in CUPY_SPARSE_DTYPES: + if arr.compressed_axes == (0, ) and arr.dtype in CUPY_SPARSE_CSR_DTYPES: reshaped = arr.reshape((arr.shape[0], -1)) - if arr.compressed_axes == (0, ): - intermediate = cupyx.scipy.sparse.csr_matrix(reshaped.to_scipy_sparse()) - elif arr.compressed_axes == (1, ): - intermediate = cupyx.scipy.sparse.csc_matrix(reshaped.to_scipy_sparse()) + intermediate = cupyx.scipy.sparse.csr_matrix(reshaped.to_scipy_sparse()) + return intermediate.toarray().reshape(arr.shape) + elif arr.compressed_axes == (1, ) and arr.dtype in CUPY_SPARSE_DTYPES: + reshaped = arr.reshape((arr.shape[0], -1)) + intermediate = cupyx.scipy.sparse.csc_matrix(reshaped.to_scipy_sparse()) return intermediate.toarray().reshape(arr.shape) else: intermediate = arr.todense() @@ -472,25 +592,68 @@ def _sparse_dok_to_CUPY(arr: sparse.DOK): intermediate = arr.todense() return cupy.array(intermediate) - def _adjust_dtype_cupy_sparse(arr): + def _scipy_csr_to_CUPY(arr: sp.csr_matrix): ''' - dtype upconversion to cupyx.scipy.sparse. + Use GPU for densification if dtype allows, otherwise CPU + ''' + if arr.dtype in CUPY_SPARSE_CSR_DTYPES: + intermediate = cupyx.scipy.sparse.csr_matrix(arr) + return intermediate.toarray() + else: + intermediate = arr.toarray() + return cupy.array(intermediate) - FIXME add support for bool + def _scipy_coo_to_CUPY(arr: sp.coo_matrix): + ''' + Use GPU for densification if dtype allows, otherwise CPU ''' if arr.dtype in CUPY_SPARSE_DTYPES: - res = arr - # print('_adjust_dtype_cupy_sparse passthrough', arr.dtype, res.dtype) + intermediate = cupyx.scipy.sparse.coo_matrix(arr) + return intermediate.toarray() else: - # Base dtype is the same for all cupyx.scipy.sparse matrices - res = arr.astype(np.result_type(arr, _base_dtypes[CUPY_SCIPY_COO])) - # print('_adjust_dtype_cupy_sparse convert', arr.dtype, res.dtype) - return res + intermediate = arr.toarray() + return cupy.array(intermediate) + + def _scipy_csc_to_CUPY(arr: sp.coo_matrix): + ''' + Use GPU for densification if dtype allows, otherwise CPU + ''' + if arr.dtype in CUPY_SPARSE_DTYPES: + intermediate = cupyx.scipy.sparse.csc_matrix(arr) + return intermediate.toarray() + else: + intermediate = arr.toarray() + return cupy.array(intermediate) + + def _cupy_scipy_csr_to_scipy_coo(arr: cupyx.scipy.sparse.csr_matrix): + if arr.dtype in CUPY_SPARSE_DTYPES: + return cupyx.scipy.sparse.coo_matrix(arr).get() + else: + return sp.coo_matrix(arr.get()) + + def _cupy_scipy_csr_to_scipy_csc(arr: cupyx.scipy.sparse.csr_matrix): + if arr.dtype in CUPY_SPARSE_DTYPES: + return cupyx.scipy.sparse.csc_matrix(arr).get() + else: + return sp.csc_matrix(arr.get()) self._converters[(NUMPY, CUPY)] = cupy.array self._converters[(CUDA, CUPY)] = cupy.array self._converters[(CUPY, NUMPY)] = cupy.asnumpy self._converters[(CUPY, CUDA)] = cupy.asnumpy + + for left in CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC: + if (left, CUPY) not in self._converters: + self._converters[(left, CUPY)] = _classes[left].toarray + + for (left, right) in [ + (CUPY_SCIPY_COO, SCIPY_COO), + (CUPY_SCIPY_CSR, SCIPY_CSR), + (CUPY_SCIPY_CSC, SCIPY_CSC) + ]: + if (left, right) not in self._converters: + self._converters[(left, right)] = _classes[left].get + # Accepted by constructor of target class for left in (SCIPY_COO, SCIPY_CSR, SCIPY_CSC, CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC): @@ -498,84 +661,33 @@ def _adjust_dtype_cupy_sparse(arr): if (left, right) not in self._converters: if left in ND_BACKENDS: self._converters[(left, right)] = chain( - _flatsig, _adjust_dtype_cupy_sparse, _classes[right] + _flatsig, _adjust_dtype_cupy_sparse(right), _classes[right] ) else: self._converters[(left, right)] = chain( - _adjust_dtype_cupy_sparse, _classes[right] + _adjust_dtype_cupy_sparse(right), _classes[right] ) # Work around https://github.com/cupy/cupy/issues/7035 # Otherwise CUPY -> {CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC} # could be handled by the block above. - left, right = CUPY, CUPY_SCIPY_COO - if (left, right) not in self._converters: - if has_cupy_coo_construction: - self._converters[(left, right)] = chain( - _flatsig, _adjust_dtype_cupy_sparse, _classes[right] - ) - elif not has_csr_complex128_bug: - self._converters[(left, right)] = chain( - _flatsig, _adjust_dtype_cupy_sparse, _classes[CUPY_SCIPY_CSR], _classes[right] - ) - else: - self._converters[(left, right)] = fail_coo_complex128 - warnings.warn( - "Please upgrade CuPy and/or CUDA: " - f"Constructing {CUPY_SCIPY_COO} from {CUPY} doesn't work (old CuPy version)" - "and installation is affected by this bug: " - "https://github.com/cupy/cupy/issues/7035 " - "(old cuSPARSE version)." - ) - for right in CUPY_SCIPY_CSR, CUPY_SCIPY_CSC: - if (left, right) not in self._converters: - if has_csr_complex128_bug and has_cupy_coo_construction: - self._converters[(left, right)] = chain( - # First convert to COO which is not affected - # Fortunately the overhead is not too bad. - _flatsig, _adjust_dtype_cupy_sparse, - _classes[CUPY_SCIPY_COO], - _classes[right] - ) - elif not has_csr_complex128_bug: - self._converters[(left, right)] = chain( - - _flatsig, _adjust_dtype_cupy_sparse, _classes[right] - ) - else: - self._converters[(left, right)] = fail_coo_complex128 - warnings.warn( - "Please upgrade CuPy and/or CUDA: " - f"Constructing {CUPY_SCIPY_COO} from {CUPY} doesn't work (old CuPy version)" - "and installation is affected by this bug: " - "https://github.com/cupy/cupy/issues/7035 " - "(old cuSPARSE version)." - ) + self._converters[CUPY, CUPY_SCIPY_COO] = _CUPY_to_cupy_scipy_coo + self._converters[CUPY, CUPY_SCIPY_CSR] = _CUPY_to_cupy_scipy_csr + self._converters[CUPY, CUPY_SCIPY_CSC] = _CUPY_to_cupy_scipy_csc for left in NUMPY, CUDA: for right in CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC: if (left, right) not in self._converters: c1 = self._converters[(left, CUPY)] c2 = self._converters[(CUPY, right)] self._converters[(left, right)] = chain(c1, c2) - for left in CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC: - if (left, CUPY) not in self._converters: - self._converters[(left, CUPY)] = _classes[left].toarray - - for (left, right) in [ - (CUPY_SCIPY_COO, SCIPY_COO), - (CUPY_SCIPY_CSR, SCIPY_CSR), - (CUPY_SCIPY_CSC, SCIPY_CSC) - ]: - if (left, right) not in self._converters: - self._converters[(left, right)] = _classes[left].get self._converters[(SPARSE_GCXS, CUPY_SCIPY_COO)] = chain( - _adjust_dtype_cupy_sparse, _GCXS_to_cupy_coo + _adjust_dtype_cupy_sparse(CUPY_SCIPY_COO), _GCXS_to_cupy_coo ) self._converters[(SPARSE_GCXS, CUPY_SCIPY_CSR)] = chain( - _adjust_dtype_cupy_sparse, _GCXS_to_cupy_csr + _adjust_dtype_cupy_sparse(CUPY_SCIPY_CSR), _GCXS_to_cupy_csr ) self._converters[(SPARSE_GCXS, CUPY_SCIPY_CSC)] = chain( - _adjust_dtype_cupy_sparse, _GCXS_to_cupy_csc + _adjust_dtype_cupy_sparse(CUPY_SCIPY_CSC), _GCXS_to_cupy_csc ) self._converters[(SPARSE_GCXS, CUPY)] = _GCXS_to_cupy @@ -591,6 +703,13 @@ def _adjust_dtype_cupy_sparse(arr): self._converters[(SPARSE_GCXS, CUPY)] = _sparse_gcxs_to_CUPY self._converters[(SPARSE_DOK, CUPY)] = _sparse_dok_to_CUPY + self._converters[(SCIPY_COO, CUPY)] = _scipy_coo_to_CUPY + self._converters[(SCIPY_CSR, CUPY)] = _scipy_csr_to_CUPY + self._converters[(SCIPY_CSC, CUPY)] = _scipy_csc_to_CUPY + + self._converters[(CUPY_SCIPY_CSR, SCIPY_COO)] = _cupy_scipy_csr_to_scipy_coo + self._converters[(CUPY_SCIPY_CSR, SCIPY_CSC)] = _cupy_scipy_csr_to_scipy_csc + proxies = [ (CUPY_SCIPY_COO, SCIPY_COO, NUMPY), (CUPY_SCIPY_COO, SCIPY_COO, CUDA), @@ -602,9 +721,8 @@ def _adjust_dtype_cupy_sparse(arr): (CUPY_SCIPY_CSR, SCIPY_CSR, NUMPY), (CUPY_SCIPY_CSR, SCIPY_CSR, CUDA), - (CUPY_SCIPY_CSR, CUPY_SCIPY_COO, SCIPY_COO), - (CUPY_SCIPY_CSR, CUPY_SCIPY_CSC, SCIPY_CSC), - (CUPY_SCIPY_CSR, CUPY_SCIPY_COO, SPARSE_COO), + + (CUPY_SCIPY_CSR, SCIPY_COO, SPARSE_COO), (CUPY_SCIPY_CSR, SCIPY_CSR, SPARSE_GCXS), (CUPY_SCIPY_CSR, SCIPY_CSR, SPARSE_DOK), @@ -616,10 +734,6 @@ def _adjust_dtype_cupy_sparse(arr): (CUPY_SCIPY_CSC, SCIPY_CSC, SPARSE_GCXS), (CUPY_SCIPY_CSC, SCIPY_CSC, SPARSE_DOK), - (SCIPY_COO, CUPY_SCIPY_COO, CUPY), - (SCIPY_CSR, CUPY_SCIPY_CSR, CUPY), - (SCIPY_CSC, CUPY_SCIPY_CSC, CUPY), - (SPARSE_COO, SCIPY_COO, CUPY_SCIPY_COO), (SPARSE_COO, SCIPY_COO, CUPY_SCIPY_CSR), (SPARSE_COO, SCIPY_COO, CUPY_SCIPY_CSC), diff --git a/tests/test_sparseconverters.py b/tests/test_sparseconverters.py index f01ad6a..0961f88 100644 --- a/tests/test_sparseconverters.py +++ b/tests/test_sparseconverters.py @@ -6,7 +6,7 @@ CPU_BACKENDS, CUDA, CUPY_BACKENDS, CUPY_SCIPY_COO, CUPY_SCIPY_CSC, CUPY_SCIPY_CSR, DENSE_BACKENDS, NUMPY, BACKENDS, CUDA_BACKENDS, ND_BACKENDS, SCIPY_COO, SPARSE_BACKENDS, SPARSE_COO, SPARSE_DOK, SPARSE_GCXS, cheapest_pair, check_shape, for_backend, - get_backend, get_device_class, make_like, prod, benchmark_conversions + get_backend, get_device_class, make_like, prod, benchmark_conversions, result_type ) @@ -75,13 +75,18 @@ def test_for_backend(left, right, dtype): CUPY_SPARSE_DTYPES = { np.float32, np.float64, np.complex64, np.complex128 } + CUPY_SPARSE_CSR_DTYPES = { + bool, np.float32, np.float64, np.complex64, np.complex128 + } CUPY_SPARSE_FORMATS = { CUPY_SCIPY_COO, CUPY_SCIPY_CSR, CUPY_SCIPY_CSC } print(left, right, dtype) if cupy is None and (left in CUDA_BACKENDS or right in CUDA_BACKENDS): pytest.skip("No CuPy, skipping CuPy test") - if left in CUPY_SPARSE_FORMATS and dtype not in CUPY_SPARSE_DTYPES: + if left == CUPY_SCIPY_CSR and dtype in CUPY_SPARSE_CSR_DTYPES: + pass + elif left in CUPY_SPARSE_FORMATS and dtype not in CUPY_SPARSE_DTYPES: pytest.skip(f"Dtype {dtype} not supported for left format {left}, skipping.") shape = (7, 11, 13, 17) left_ref = _mk_random(shape, dtype=dtype, array_backend=NUMPY) @@ -124,8 +129,23 @@ def test_for_backend(left, right, dtype): assert converted.shape == target_shape assert converted_back.shape == target_shape + assert np.allclose(left_ref.reshape(target_shape), converted_back) + keep_dtype = left not in CUPY_SPARSE_FORMATS and right not in CUPY_SPARSE_FORMATS + keep_dtype = keep_dtype or dtype in CUPY_SPARSE_DTYPES + no_bool = (CUPY_SCIPY_COO, CUPY_SCIPY_CSC) + keep_dtype = keep_dtype or dtype == bool and left not in no_bool and right not in no_bool + + target_dtype = result_type(left, right, dtype) + + if (keep_dtype): + assert left_ref.dtype == data.dtype + assert left_ref.dtype == converted.dtype + assert left_ref.dtype == converted_back.dtype + else: + assert converted.dtype == target_dtype + def test_unknown_format(): not_an_array = "I am not a supported array type" @@ -209,3 +229,26 @@ def test_graceful_no_cupy(): repeats=1, warmup=False ) + + +@pytest.mark.parametrize( + 'args,expected', [ + ((bool, ), bool), + ((np.uint8, ), np.uint8), + ((bool, CUPY_SCIPY_COO), np.float32), + ((np.uint8, CUPY_SCIPY_COO), np.float32), + ((np.uint8, CUPY_SCIPY_CSR, NUMPY), np.float32), + ((np.complex64, CUPY_SCIPY_COO, np.int16, NUMPY), np.complex64), + ((np.uint32, NUMPY, CUPY_SCIPY_COO, bool), np.float64), + ((np.uint64, NUMPY, bool), np.uint64), + ((CUPY_SCIPY_CSC, np.uint64, NUMPY, bool), np.float64), + (DENSE_BACKENDS.union( + SPARSE_BACKENDS.intersection(CPU_BACKENDS), (CUPY_SCIPY_CSR, ) + ), bool), + (BACKENDS, np.float32), + ] +) +def test_result_type(args, expected): + print(args, expected) + res = result_type(*args) + assert res == expected