Skip to content

Commit dd2b7c4

Browse files
author
Holger Kohr
committed
MAINT: improve import of cupy
1 parent d028dd5 commit dd2b7c4

File tree

1 file changed

+173
-157
lines changed

1 file changed

+173
-157
lines changed

‎odl/space/cupy_tensors.py

Lines changed: 173 additions & 157 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,15 @@
66
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
77
# obtain one at https://mozilla.org/MPL/2.0/.
88

9-
"""Implementation of tensor spaces using ``pygpu``."""
9+
"""Implementation of tensor spaces using CuPy.
10+
11+
See https://cupy.chainer.org/ or https://github.com/cupy/cupy for details
12+
on the backend.
13+
"""
1014

1115
from __future__ import print_function, division, absolute_import
1216
import numpy as np
17+
import warnings
1318

1419
from odl.set import RealNumbers
1520
from odl.space.base_tensors import TensorSpace, Tensor
@@ -23,9 +28,11 @@
2328
except ImportError:
2429
CUPY_AVAILABLE = False
2530
else:
26-
from pkg_resources import parse_version
27-
if parse_version(cupy.__version__) < parse_version('2.0.0rc1'):
28-
raise ImportError('cupy <2.0.0rc1 not supported')
31+
_maj = int(cupy.__version__.split('.')[0])
32+
if _maj < 2:
33+
raise warnings.warn(
34+
'your version {} of CuPy is not supported; please upgrade to '
35+
'version 2.0.0 or higher'.format(cupy.__version__), RuntimeWarning)
2936
CUPY_AVAILABLE = True
3037

3138

@@ -35,10 +42,11 @@
3542
# --- Space method implementations --- #
3643

3744

38-
lico = cupy.ElementwiseKernel(in_params='T a, T x, T b, T y',
39-
out_params='T z',
40-
operation='z = a * x + b * y;',
41-
name='lico')
45+
if CUPY_AVAILABLE:
46+
lico = cupy.ElementwiseKernel(in_params='T a, T x, T b, T y',
47+
out_params='T z',
48+
operation='z = a * x + b * y;',
49+
name='lico')
4250

4351

4452
def fallback_scal(a, x):
@@ -77,12 +85,17 @@ def _cublas_func(name, dtype):
7785
ValueError :
7886
If the data type is not supported by cuBLAS.
7987
"""
80-
if np.dtype(dtype) == 'float32':
88+
dtype, dtype_in = np.dtype(dtype), dtype
89+
if dtype == 'float32':
8190
prefix = 's'
82-
elif np.dtype(dtype) == 'float64':
91+
elif dtype == 'float64':
8392
prefix = 'd'
93+
elif dtype == 'complex64':
94+
prefix = 'c'
95+
elif dtype == 'complex128':
96+
prefix = 'z'
8497
else:
85-
raise ValueError('dtype {!r} not supported by cuBLAS'.format(dtype))
98+
raise ValueError('dtype {!r} not supported by cuBLAS'.format(dtype_in))
8699

87100
return getattr(cupy.cuda.cublas, prefix + name)
88101

@@ -197,13 +210,15 @@ def _lincomb_impl(a, x1, b, x2, out):
197210

198211
class CupyTensorSpace(TensorSpace):
199212

200-
"""Tensor space implemented with GPU arrays.
213+
"""Tensor space implemented with CUDA arrays using the CuPy library.
201214
202215
This space implements tensors of arbitrary rank over a `Field` ``F``,
203216
which is either the real or complex numbers.
204217
205218
Its elements are represented as instances of the
206219
`CupyTensor` class.
220+
221+
See https://github.com/cupy/cupy for details on the backend.
207222
"""
208223

209224
def __init__(self, shape, dtype='float64', device=None, **kwargs):
@@ -1799,157 +1814,158 @@ def _weighting(weights, exponent):
17991814

18001815
# Kernels for space functions
18011816

1802-
dotw = cupy.ReductionKernel(in_params='T x, T y, W w',
1803-
out_params='T res',
1804-
map_expr='x * y * w',
1805-
reduce_expr='a + b',
1806-
post_map_expr='res = a',
1807-
identity='0',
1808-
name='dotw')
1809-
1810-
nrm0 = cupy.ReductionKernel(in_params='T x',
1811-
out_params='int64 res',
1812-
map_expr='x != 0',
1813-
reduce_expr='a + b',
1814-
post_map_expr='res = a',
1815-
identity='0',
1816-
name='nrm0')
1817-
1818-
nrm1 = cupy.ReductionKernel(in_params='T x',
1819-
out_params='T res',
1820-
map_expr='abs(x)',
1821-
reduce_expr='a + b',
1822-
post_map_expr='res = a',
1823-
identity='0',
1824-
name='nrm1w')
1825-
1826-
nrm1w = cupy.ReductionKernel(in_params='T x, W w',
1827-
out_params='T res',
1828-
map_expr='abs(x) * w',
1829-
reduce_expr='a + b',
1830-
post_map_expr='res = a',
1831-
identity='0',
1832-
name='nrm1w')
1833-
1834-
nrm2 = cupy.ReductionKernel(in_params='T x',
1835-
out_params='T res',
1836-
map_expr='x * x',
1837-
reduce_expr='a + b',
1838-
post_map_expr='res = sqrt(a)',
1839-
identity='0',
1840-
name='nrm2')
1841-
1842-
nrm2w = cupy.ReductionKernel(in_params='T x, W w',
1843-
out_params='T res',
1844-
map_expr='x * x * w',
1845-
reduce_expr='a + b',
1846-
post_map_expr='res = sqrt(a)',
1847-
identity='0',
1848-
name='nrm2w')
1849-
1850-
nrminf = cupy.ReductionKernel(in_params='T x',
1851-
out_params='T res',
1852-
map_expr='abs(x)',
1853-
reduce_expr='a > b ? a : b',
1854-
post_map_expr='res = a',
1855-
identity='0',
1856-
name='nrminf')
1857-
1858-
nrmneginf = cupy.ReductionKernel(in_params='T x',
1817+
if CUPY_AVAILABLE:
1818+
dotw = cupy.ReductionKernel(in_params='T x, T y, W w',
1819+
out_params='T res',
1820+
map_expr='x * y * w',
1821+
reduce_expr='a + b',
1822+
post_map_expr='res = a',
1823+
identity='0',
1824+
name='dotw')
1825+
1826+
nrm0 = cupy.ReductionKernel(in_params='T x',
1827+
out_params='int64 res',
1828+
map_expr='x != 0',
1829+
reduce_expr='a + b',
1830+
post_map_expr='res = a',
1831+
identity='0',
1832+
name='nrm0')
1833+
1834+
nrm1 = cupy.ReductionKernel(in_params='T x',
1835+
out_params='T res',
1836+
map_expr='abs(x)',
1837+
reduce_expr='a + b',
1838+
post_map_expr='res = a',
1839+
identity='0',
1840+
name='nrm1w')
1841+
1842+
nrm1w = cupy.ReductionKernel(in_params='T x, W w',
18591843
out_params='T res',
1860-
map_expr='abs(x)',
1861-
reduce_expr='a > b ? b : a',
1844+
map_expr='abs(x) * w',
1845+
reduce_expr='a + b',
18621846
post_map_expr='res = a',
18631847
identity='0',
1864-
name='nrmneginf')
1865-
1866-
nrmp = cupy.ReductionKernel(in_params='T x, T p',
1867-
out_params='T res',
1868-
map_expr='pow(abs(x), p)',
1869-
reduce_expr='a + b',
1870-
post_map_expr='res = pow(a, 1 / p)',
1871-
identity='0',
1872-
name='nrmp')
1873-
1874-
nrmpw = cupy.ReductionKernel(in_params='T x, T p, W w',
1875-
out_params='T res',
1876-
map_expr='pow(abs(x), p) * w',
1877-
reduce_expr='a + b',
1878-
post_map_expr='res = pow(a, 1 / p)',
1879-
identity='0',
1880-
name='nrmpw')
1881-
1882-
dist0 = cupy.ReductionKernel(in_params='T x, T y',
1883-
out_params='int64 res',
1884-
map_expr='x != y',
1885-
reduce_expr='a + b',
1886-
post_map_expr='res = a',
1887-
identity='0',
1888-
name='dist0')
1889-
1890-
dist1 = cupy.ReductionKernel(in_params='T x, T y',
1891-
out_params='T res',
1892-
map_expr='abs(x - y)',
1893-
reduce_expr='a + b',
1894-
post_map_expr='res = a',
1895-
identity='0',
1896-
name='dist1')
1897-
1898-
dist1w = cupy.ReductionKernel(in_params='T x, T y, W w',
1899-
out_params='T res',
1900-
map_expr='abs(x - y) * w',
1901-
reduce_expr='a + b',
1902-
post_map_expr='res = a',
1903-
identity='0',
1904-
name='dist1w')
1905-
1906-
dist2 = cupy.ReductionKernel(in_params='T x, T y',
1907-
out_params='T res',
1908-
map_expr='(x - y) * (x - y)',
1909-
reduce_expr='a + b',
1910-
post_map_expr='res = sqrt(a)',
1911-
identity='0',
1912-
name='dist2')
1913-
1914-
dist2w = cupy.ReductionKernel(in_params='T x, T y, W w',
1915-
out_params='T res',
1916-
map_expr='(x - y) * (x - y) * w',
1917-
reduce_expr='a + b',
1918-
post_map_expr='res = sqrt(a)',
1919-
identity='0',
1920-
name='dist2w')
1921-
1922-
distinf = cupy.ReductionKernel(in_params='T x, T y',
1923-
out_params='T res',
1924-
map_expr='abs(x - y)',
1925-
reduce_expr='a > b ? a : b',
1926-
post_map_expr='res = a',
1927-
identity='0',
1928-
name='distinf')
1929-
1930-
distneginf = cupy.ReductionKernel(in_params='T x, T y',
1848+
name='nrm1w')
1849+
1850+
nrm2 = cupy.ReductionKernel(in_params='T x',
1851+
out_params='T res',
1852+
map_expr='x * x',
1853+
reduce_expr='a + b',
1854+
post_map_expr='res = sqrt(a)',
1855+
identity='0',
1856+
name='nrm2')
1857+
1858+
nrm2w = cupy.ReductionKernel(in_params='T x, W w',
1859+
out_params='T res',
1860+
map_expr='x * x * w',
1861+
reduce_expr='a + b',
1862+
post_map_expr='res = sqrt(a)',
1863+
identity='0',
1864+
name='nrm2w')
1865+
1866+
nrminf = cupy.ReductionKernel(in_params='T x',
19311867
out_params='T res',
1932-
map_expr='abs(x - y)',
1933-
reduce_expr='a > b ? b : a',
1868+
map_expr='abs(x)',
1869+
reduce_expr='a > b ? a : b',
19341870
post_map_expr='res = a',
19351871
identity='0',
1936-
name='distneginf')
1937-
1938-
distp = cupy.ReductionKernel(in_params='T x, T y, T p',
1939-
out_params='T res',
1940-
map_expr='pow(abs(x - y), p)',
1941-
reduce_expr='a + b',
1942-
post_map_expr='res = pow(a, 1 / p)',
1943-
identity='0',
1944-
name='distp')
1945-
1946-
distpw = cupy.ReductionKernel(in_params='T x, T y, T p, W w',
1947-
out_params='T res',
1948-
map_expr='pow(abs(x - y), p) * w',
1949-
reduce_expr='a + b',
1950-
post_map_expr='res = pow(a, 1 / p)',
1951-
identity='0',
1952-
name='distpw')
1872+
name='nrminf')
1873+
1874+
nrmneginf = cupy.ReductionKernel(in_params='T x',
1875+
out_params='T res',
1876+
map_expr='abs(x)',
1877+
reduce_expr='a > b ? b : a',
1878+
post_map_expr='res = a',
1879+
identity='0',
1880+
name='nrmneginf')
1881+
1882+
nrmp = cupy.ReductionKernel(in_params='T x, T p',
1883+
out_params='T res',
1884+
map_expr='pow(abs(x), p)',
1885+
reduce_expr='a + b',
1886+
post_map_expr='res = pow(a, 1 / p)',
1887+
identity='0',
1888+
name='nrmp')
1889+
1890+
nrmpw = cupy.ReductionKernel(in_params='T x, T p, W w',
1891+
out_params='T res',
1892+
map_expr='pow(abs(x), p) * w',
1893+
reduce_expr='a + b',
1894+
post_map_expr='res = pow(a, 1 / p)',
1895+
identity='0',
1896+
name='nrmpw')
1897+
1898+
dist0 = cupy.ReductionKernel(in_params='T x, T y',
1899+
out_params='int64 res',
1900+
map_expr='x != y',
1901+
reduce_expr='a + b',
1902+
post_map_expr='res = a',
1903+
identity='0',
1904+
name='dist0')
1905+
1906+
dist1 = cupy.ReductionKernel(in_params='T x, T y',
1907+
out_params='T res',
1908+
map_expr='abs(x - y)',
1909+
reduce_expr='a + b',
1910+
post_map_expr='res = a',
1911+
identity='0',
1912+
name='dist1')
1913+
1914+
dist1w = cupy.ReductionKernel(in_params='T x, T y, W w',
1915+
out_params='T res',
1916+
map_expr='abs(x - y) * w',
1917+
reduce_expr='a + b',
1918+
post_map_expr='res = a',
1919+
identity='0',
1920+
name='dist1w')
1921+
1922+
dist2 = cupy.ReductionKernel(in_params='T x, T y',
1923+
out_params='T res',
1924+
map_expr='(x - y) * (x - y)',
1925+
reduce_expr='a + b',
1926+
post_map_expr='res = sqrt(a)',
1927+
identity='0',
1928+
name='dist2')
1929+
1930+
dist2w = cupy.ReductionKernel(in_params='T x, T y, W w',
1931+
out_params='T res',
1932+
map_expr='(x - y) * (x - y) * w',
1933+
reduce_expr='a + b',
1934+
post_map_expr='res = sqrt(a)',
1935+
identity='0',
1936+
name='dist2w')
1937+
1938+
distinf = cupy.ReductionKernel(in_params='T x, T y',
1939+
out_params='T res',
1940+
map_expr='abs(x - y)',
1941+
reduce_expr='a > b ? a : b',
1942+
post_map_expr='res = a',
1943+
identity='0',
1944+
name='distinf')
1945+
1946+
distneginf = cupy.ReductionKernel(in_params='T x, T y',
1947+
out_params='T res',
1948+
map_expr='abs(x - y)',
1949+
reduce_expr='a > b ? b : a',
1950+
post_map_expr='res = a',
1951+
identity='0',
1952+
name='distneginf')
1953+
1954+
distp = cupy.ReductionKernel(in_params='T x, T y, T p',
1955+
out_params='T res',
1956+
map_expr='pow(abs(x - y), p)',
1957+
reduce_expr='a + b',
1958+
post_map_expr='res = pow(a, 1 / p)',
1959+
identity='0',
1960+
name='distp')
1961+
1962+
distpw = cupy.ReductionKernel(in_params='T x, T y, T p, W w',
1963+
out_params='T res',
1964+
map_expr='pow(abs(x - y), p) * w',
1965+
reduce_expr='a + b',
1966+
post_map_expr='res = pow(a, 1 / p)',
1967+
identity='0',
1968+
name='distpw')
19531969

19541970

19551971
class CupyTensorSpaceArrayWeighting(ArrayWeighting):

0 commit comments

Comments
 (0)