diff --git a/arraycontext/impl/pyopencl.py b/arraycontext/impl/pyopencl/__init__.py
similarity index 63%
rename from arraycontext/impl/pyopencl.py
rename to arraycontext/impl/pyopencl/__init__.py
index fa2aa7493b1d4bb074c75a6b9e54a8c95acdfe11..ca5b2dd725c3a121e87256859d5de37f68e01a11 100644
--- a/arraycontext/impl/pyopencl.py
+++ b/arraycontext/impl/pyopencl/__init__.py
@@ -28,153 +28,16 @@ THE SOFTWARE.
from warnings import warn
from typing import Sequence, Union
-from functools import partial
-import operator
import numpy as np
from pytools.tag import Tag
from arraycontext.metadata import FirstAxisIsElementsTag
-from arraycontext.fake_numpy import \
- BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace
-from arraycontext.container.traversal import (rec_multimap_array_container,
- rec_map_array_container)
from arraycontext.context import ArrayContext
from numbers import Number
-# {{{ fake numpy
-
-class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace):
- def _get_fake_numpy_linalg_namespace(self):
- return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context)
-
- # {{{ comparisons
-
- # FIXME: This should be documentation, not a comment.
- # These are here mainly because some arrays may choose to interpret
- # equality comparison as a binary predicate of structural identity,
- # i.e. more like "are you two equal", and not like numpy semantics.
- # These operations provide access to numpy-style comparisons in that
- # case.
-
- def equal(self, x, y):
- return rec_multimap_array_container(operator.eq, x, y)
-
- def not_equal(self, x, y):
- return rec_multimap_array_container(operator.ne, x, y)
-
- def greater(self, x, y):
- return rec_multimap_array_container(operator.gt, x, y)
-
- def greater_equal(self, x, y):
- return rec_multimap_array_container(operator.ge, x, y)
-
- def less(self, x, y):
- return rec_multimap_array_container(operator.lt, x, y)
-
- def less_equal(self, x, y):
- return rec_multimap_array_container(operator.le, x, y)
-
- # }}}
-
- def ones_like(self, ary):
- def _ones_like(subary):
- ones = self._array_context.empty_like(subary)
- ones.fill(1)
- return ones
-
- return self._new_like(ary, _ones_like)
-
- def maximum(self, x, y):
- import pyopencl.array as cl_array
- return rec_multimap_array_container(
- partial(cl_array.maximum, queue=self._array_context.queue),
- x, y)
-
- def minimum(self, x, y):
- import pyopencl.array as cl_array
- return rec_multimap_array_container(
- partial(cl_array.minimum, queue=self._array_context.queue),
- x, y)
-
- def where(self, criterion, then, else_):
- import pyopencl.array as cl_array
-
- def where_inner(inner_crit, inner_then, inner_else):
- if isinstance(inner_crit, bool):
- return inner_then if inner_crit else inner_else
- return cl_array.if_positive(inner_crit != 0, inner_then, inner_else,
- queue=self._array_context.queue)
-
- return rec_multimap_array_container(where_inner, criterion, then, else_)
-
- def sum(self, a, dtype=None):
- import pyopencl.array as cl_array
- return cl_array.sum(
- a, dtype=dtype, queue=self._array_context.queue).get()[()]
-
- def min(self, a):
- import pyopencl.array as cl_array
- return cl_array.min(a, queue=self._array_context.queue).get()[()]
-
- def max(self, a):
- import pyopencl.array as cl_array
- return cl_array.max(a, queue=self._array_context.queue).get()[()]
-
- def stack(self, arrays, axis=0):
- import pyopencl.array as cla
- return rec_multimap_array_container(
- lambda *args: cla.stack(arrays=args, axis=axis,
- queue=self._array_context.queue),
- *arrays)
-
- def reshape(self, a, newshape):
- import pyopencl.array as cla
- return cla.reshape(a, newshape)
-
- def concatenate(self, arrays, axis=0):
- import pyopencl.array as cla
- return cla.concatenate(
- arrays, axis,
- self._array_context.queue,
- self._array_context.allocator
- )
-
- def ravel(self, a, order="C"):
- def _rec_ravel(a):
- if order in "FC":
- return a.reshape(-1, order=order)
- elif order == "A":
- # TODO: upstream this to pyopencl.array
- if a.flags.f_contiguous:
- return a.reshape(-1, order="F")
- elif a.flags.c_contiguous:
- return a.reshape(-1, order="C")
- else:
- raise ValueError("For `order='A'`, array should be either"
- " F-contiguous or C-contiguous.")
- elif order == "K":
- raise NotImplementedError("PyOpenCLArrayContext.np.ravel not "
- "implemented for 'order=K'")
- else:
- raise ValueError("`order` can be one of 'F', 'C', 'A' or 'K'. "
- f"(got {order})")
-
- return rec_map_array_container(_rec_ravel, a)
-
-# }}}
-
-
-# {{{ fake np.linalg
-
-class _PyOpenCLFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
- pass
-
-# }}}
-
-
# {{{ PyOpenCLArrayContext
class PyOpenCLArrayContext(ArrayContext):
@@ -222,6 +85,8 @@ class PyOpenCLArrayContext(ArrayContext):
For now, *wait_event_queue_length* should be regarded as an
experimental feature that may change or disappear at any minute.
"""
+ import pyopencl as cl
+
super().__init__()
self.context = queue.context
self.queue = queue
@@ -233,7 +98,6 @@ class PyOpenCLArrayContext(ArrayContext):
self._wait_event_queue_length = wait_event_queue_length
self._kernel_name_to_wait_event_queue = {}
- import pyopencl as cl
if queue.device.type & cl.device_type.GPU:
if allocator is None:
warn("PyOpenCLArrayContext created without an allocator on a GPU. "
@@ -250,23 +114,24 @@ class PyOpenCLArrayContext(ArrayContext):
self._loopy_transform_cache = {}
def _get_fake_numpy_namespace(self):
+ from arraycontext.impl.pyopencl.fake_numpy import PyOpenCLFakeNumpyNamespace
return PyOpenCLFakeNumpyNamespace(self)
# {{{ ArrayContext interface
def empty(self, shape, dtype):
- import pyopencl.array as cla
- return cla.empty(self.queue, shape=shape, dtype=dtype,
+ import pyopencl.array as cl_array
+ return cl_array.empty(self.queue, shape=shape, dtype=dtype,
allocator=self.allocator)
def zeros(self, shape, dtype):
- import pyopencl.array as cla
- return cla.zeros(self.queue, shape=shape, dtype=dtype,
+ import pyopencl.array as cl_array
+ return cl_array.zeros(self.queue, shape=shape, dtype=dtype,
allocator=self.allocator)
def from_numpy(self, array: np.ndarray):
- import pyopencl.array as cla
- return cla.to_device(self.queue, array, allocator=self.allocator)
+ import pyopencl.array as cl_array
+ return cl_array.to_device(self.queue, array, allocator=self.allocator)
def to_numpy(self, array):
if isinstance(array, Number):
diff --git a/arraycontext/impl/pyopencl/fake_numpy.py b/arraycontext/impl/pyopencl/fake_numpy.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2bfe00df269b62b435f1d8d3e8a6c57bd8c4da0
--- /dev/null
+++ b/arraycontext/impl/pyopencl/fake_numpy.py
@@ -0,0 +1,165 @@
+"""
+.. currentmodule:: arraycontext
+.. autoclass:: PyOpenCLArrayContext
+"""
+__copyright__ = """
+Copyright (C) 2020-1 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+from functools import partial
+import operator
+
+from arraycontext.fake_numpy import \
+ BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace
+from arraycontext.container.traversal import (rec_multimap_array_container,
+ rec_map_array_container)
+
+try:
+ import pyopencl as cl # noqa: F401
+ import pyopencl.array as cl_array
+except ImportError:
+ pass
+
+
+# {{{ fake numpy
+
+class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace):
+ def _get_fake_numpy_linalg_namespace(self):
+ return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context)
+
+ # {{{ comparisons
+
+ # FIXME: This should be documentation, not a comment.
+ # These are here mainly because some arrays may choose to interpret
+ # equality comparison as a binary predicate of structural identity,
+ # i.e. more like "are you two equal", and not like numpy semantics.
+ # These operations provide access to numpy-style comparisons in that
+ # case.
+
+ def equal(self, x, y):
+ return rec_multimap_array_container(operator.eq, x, y)
+
+ def not_equal(self, x, y):
+ return rec_multimap_array_container(operator.ne, x, y)
+
+ def greater(self, x, y):
+ return rec_multimap_array_container(operator.gt, x, y)
+
+ def greater_equal(self, x, y):
+ return rec_multimap_array_container(operator.ge, x, y)
+
+ def less(self, x, y):
+ return rec_multimap_array_container(operator.lt, x, y)
+
+ def less_equal(self, x, y):
+ return rec_multimap_array_container(operator.le, x, y)
+
+ # }}}
+
+ def ones_like(self, ary):
+ def _ones_like(subary):
+ ones = self._array_context.empty_like(subary)
+ ones.fill(1)
+ return ones
+
+ return self._new_like(ary, _ones_like)
+
+ def maximum(self, x, y):
+ return rec_multimap_array_container(
+ partial(cl_array.maximum, queue=self._array_context.queue),
+ x, y)
+
+ def minimum(self, x, y):
+ return rec_multimap_array_container(
+ partial(cl_array.minimum, queue=self._array_context.queue),
+ x, y)
+
+ def where(self, criterion, then, else_):
+ def where_inner(inner_crit, inner_then, inner_else):
+ if isinstance(inner_crit, bool):
+ return inner_then if inner_crit else inner_else
+ return cl_array.if_positive(inner_crit != 0, inner_then, inner_else,
+ queue=self._array_context.queue)
+
+ return rec_multimap_array_container(where_inner, criterion, then, else_)
+
+ def sum(self, a, dtype=None):
+ return cl_array.sum(
+ a, dtype=dtype, queue=self._array_context.queue).get()[()]
+
+ def min(self, a):
+ return cl_array.min(a, queue=self._array_context.queue).get()[()]
+
+ def max(self, a):
+ return cl_array.max(a, queue=self._array_context.queue).get()[()]
+
+ def stack(self, arrays, axis=0):
+ return rec_multimap_array_container(
+ lambda *args: cl_array.stack(arrays=args, axis=axis,
+ queue=self._array_context.queue),
+ *arrays)
+
+ def reshape(self, a, newshape):
+ return cl_array.reshape(a, newshape)
+
+ def concatenate(self, arrays, axis=0):
+ return cl_array.concatenate(
+ arrays, axis,
+ self._array_context.queue,
+ self._array_context.allocator
+ )
+
+ def ravel(self, a, order="C"):
+ def _rec_ravel(a):
+ if order in "FC":
+ return a.reshape(-1, order=order)
+ elif order == "A":
+ # TODO: upstream this to pyopencl.array
+ if a.flags.f_contiguous:
+ return a.reshape(-1, order="F")
+ elif a.flags.c_contiguous:
+ return a.reshape(-1, order="C")
+ else:
+ raise ValueError("For `order='A'`, array should be either"
+ " F-contiguous or C-contiguous.")
+ elif order == "K":
+ raise NotImplementedError("PyOpenCLArrayContext.np.ravel not "
+ "implemented for 'order=K'")
+ else:
+ raise ValueError("`order` can be one of 'F', 'C', 'A' or 'K'. "
+ f"(got {order})")
+
+ return rec_map_array_container(_rec_ravel, a)
+
+# }}}
+
+
+# {{{ fake np.linalg
+
+class _PyOpenCLFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
+ pass
+
+# }}}
+
+
+# vim: foldmethod=marker