Skip to content
Snippets Groups Projects

Compare revisions

Changes are shown as if the source revision was being merged into the target revision. Learn more about comparing revisions.

Source

Select target project
No results found
Select Git revision

Target

Select target project
  • inducer/arraycontext
  • kaushikcfd/arraycontext
  • fikl2/arraycontext
3 results
Select Git revision
Show changes
Hide whitespace changes
Inline Side-by-side
Showing
with 4431 additions and 140 deletions
arraycontext/impl/pyopencl.py → arraycontext/impl/pyopencl/__init__.py
View file @ 1013ecb2
"""
from __future__ import annotations
__doc__ = """
.. currentmodule:: arraycontext
.. autoclass:: PyOpenCLArrayContext
.. automodule:: arraycontext.impl.pyopencl.taggable_cl_array
"""
__copyright__ = """
Copyright (C) 2020-1 University of Illinois Board of Trustees
"""
......@@ -26,189 +31,27 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
from typing import Sequence, Union
from functools import partial
import operator
from collections.abc import Callable
from typing import TYPE_CHECKING
from warnings import warn
import numpy as np
from pytools import memoize_method
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
from arraycontext.container import serialize_container, is_array_container
from arraycontext.context import ArrayContext
# {{{ 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)
from pytools.tag import ToTagSetConvertible
def not_equal(self, x, y):
return rec_multimap_array_container(operator.ne, x, y)
from arraycontext.container.traversal import rec_map_array_container, with_array_context
from arraycontext.context import (
Array,
ArrayContext,
ArrayOrContainer,
ScalarLike,
UntransformedCodeWarning,
)
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 concatenate(self, arrays, axis=0):
import pyopencl.array as cla
return cla.concatenate(
arrays, axis,
self._array_context.queue,
self._array_context.allocator
)
# }}}
# {{{ fake np.linalg
def _flatten_array(ary):
import pyopencl.array as cl
assert isinstance(ary, cl.Array)
if ary.size == 0:
# Work around https://github.com/inducer/pyopencl/pull/402
return ary._new_with_changes(
data=None, offset=0, shape=(0,), strides=(ary.dtype.itemsize,))
if ary.flags.f_contiguous:
return ary.reshape(-1, order="F")
elif ary.flags.c_contiguous:
return ary.reshape(-1, order="C")
else:
raise ValueError("cannot flatten array "
f"with strides {ary.strides} of {ary.dtype}")
class _PyOpenCLFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
def norm(self, ary, ord=None):
from numbers import Number
if isinstance(ary, Number):
return abs(ary)
if ord is None:
ord = 2
try:
from meshmode.dof_array import DOFArray
except ImportError:
pass
else:
if isinstance(ary, DOFArray):
from warnings import warn
warn("Taking an actx.np.linalg.norm of a DOFArray is deprecated. "
"(DOFArrays use 2D arrays internally, and "
"actx.np.linalg.norm should compute matrix norms of those.) "
"This will stop working in 2022. "
"Use meshmode.dof_array.flat_norm instead.",
DeprecationWarning, stacklevel=2)
import numpy.linalg as la
return la.norm(
[self.norm(_flatten_array(subary), ord=ord)
for _, subary in serialize_container(ary)],
ord=ord)
if is_array_container(ary):
import numpy.linalg as la
return la.norm(
[self.norm(subary, ord=ord)
for _, subary in serialize_container(ary)],
ord=ord)
if len(ary.shape) != 1:
raise NotImplementedError("only vector norms are implemented")
if ary.size == 0:
return 0
if ord == np.inf:
return self._array_context.np.max(abs(ary))
elif isinstance(ord, Number) and ord > 0:
return self._array_context.np.sum(abs(ary)**ord)**(1/ord)
else:
raise NotImplementedError(f"unsupported value of 'ord': {ord}")
# }}}
if TYPE_CHECKING:
import loopy as lp
import pyopencl
# {{{ PyOpenCLArrayContext
......@@ -234,9 +77,15 @@ class PyOpenCLArrayContext(ArrayContext):
of arrays are created (e.g. as results of computation), the associated cost
may become significant. Using e.g. :class:`pyopencl.tools.MemoryPool`
as the allocator can help avoid this cost.
.. automethod:: transform_loopy_program
"""
def __init__(self, queue, allocator=None, wait_event_queue_length=None):
def __init__(self,
queue: pyopencl.CommandQueue,
allocator: pyopencl.tools.AllocatorBase | None = None,
wait_event_queue_length: int | None = None,
force_device_scalars: bool | None = None) -> None:
r"""
:arg wait_event_queue_length: The length of a queue of
:class:`~pyopencl.Event` objects that are maintained by the
......@@ -258,56 +107,153 @@ class PyOpenCLArrayContext(ArrayContext):
For now, *wait_event_queue_length* should be regarded as an
experimental feature that may change or disappear at any minute.
"""
if force_device_scalars is not None:
warn(
"`force_device_scalars` is deprecated and will be removed in 2025.",
DeprecationWarning, stacklevel=2)
if not force_device_scalars:
raise ValueError(
"Passing force_device_scalars=False is not allowed.")
import pyopencl as cl
import pyopencl.array as cl_array
super().__init__()
self.context = queue.context
self.queue = queue
self.allocator = allocator if allocator else None
if wait_event_queue_length is None:
wait_event_queue_length = 10
self._wait_event_queue_length = wait_event_queue_length
self._kernel_name_to_wait_event_queue = {}
self._force_device_scalars = True
# Subclasses might still be using the old
# "force_devices_scalars: bool = False" interface, in which case we need
# to explicitly pass force_device_scalars=True in clone()
self._passed_force_device_scalars = force_device_scalars is not None
import pyopencl as cl
if allocator is None and queue.device.type & cl.device_type.GPU:
from warnings import warn
warn("PyOpenCLArrayContext created without an allocator on a GPU. "
"This can lead to high numbers of memory allocations. "
"Please consider using a pyopencl.tools.MemoryPool. "
"Run with allocator=False to disable this warning.")
self._wait_event_queue_length = wait_event_queue_length
self._kernel_name_to_wait_event_queue: dict[str, list[cl.Event]] = {}
if queue.device.type & cl.device_type.GPU:
if allocator is None:
warn("PyOpenCLArrayContext created without an allocator on a GPU. "
"This can lead to high numbers of memory allocations. "
"Please consider using a pyopencl.tools.MemoryPool. "
"Run with allocator=False to disable this warning.",
stacklevel=2)
if __debug__:
# Use "running on GPU" as a proxy for "they care about speed".
warn("You are using the PyOpenCLArrayContext on a GPU, but you "
"are running Python in debug mode. Use 'python -O' for "
"a noticeable speed improvement.",
stacklevel=2)
self._loopy_transform_cache: \
dict[lp.TranslationUnit, lp.TranslationUnit] = {}
# TODO: Ideally this should only be `(TaggableCLArray,)`, but
# that would break the logic in the downstream users.
self.array_types = (cl_array.Array,)
def _get_fake_numpy_namespace(self):
from arraycontext.impl.pyopencl.fake_numpy import PyOpenCLFakeNumpyNamespace
return PyOpenCLFakeNumpyNamespace(self)
def _rec_map_container(
self, func: Callable[[Array], Array], array: ArrayOrContainer,
allowed_types: tuple[type, ...] | None = None, *,
default_scalar: ScalarLike | None = None,
strict: bool = False) -> ArrayOrContainer:
import arraycontext.impl.pyopencl.taggable_cl_array as tga
if allowed_types is None:
# TODO: replace with 'self.array_types' once `cla.Array` support
# is completely removed
allowed_types = (tga.TaggableCLArray,)
def _wrapper(ary):
if isinstance(ary, allowed_types):
return func(ary)
elif not strict and isinstance(ary, self.array_types):
from warnings import warn
warn(f"Invoking {type(self).__name__}.{func.__name__[1:]} with "
f"{type(ary).__name__} will be unsupported in 2023. Use "
"'to_tagged_cl_array' to convert instances to TaggableCLArray.",
DeprecationWarning, stacklevel=2)
return func(tga.to_tagged_cl_array(ary))
elif np.isscalar(ary):
if default_scalar is None:
return ary
else:
return np.array(ary).dtype.type(default_scalar)
else:
raise TypeError(
f"{type(self).__name__}.{func.__name__[1:]} invoked with "
f"an unsupported array type: got '{type(ary).__name__}', "
f"but expected one of {allowed_types}")
return rec_map_array_container(_wrapper, array)
# {{{ ArrayContext interface
def empty(self, shape, dtype):
import pyopencl.array as cla
return cla.empty(self.queue, shape=shape, dtype=dtype,
allocator=self.allocator)
def from_numpy(self, array):
import arraycontext.impl.pyopencl.taggable_cl_array as tga
def zeros(self, shape, dtype):
import pyopencl.array as cla
return cla.zeros(self.queue, shape=shape, dtype=dtype,
allocator=self.allocator)
def _from_numpy(ary):
return tga.to_device(self.queue, ary, 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)
return with_array_context(
self._rec_map_container(_from_numpy, array, (np.ndarray,), strict=True),
actx=self)
def to_numpy(self, array):
return array.get(queue=self.queue)
def _to_numpy(ary):
return ary.get(queue=self.queue)
return with_array_context(
self._rec_map_container(_to_numpy, array),
actx=None)
def freeze(self, array):
def _freeze(ary):
ary.finish()
return ary.with_queue(None)
return with_array_context(self._rec_map_container(_freeze, array), actx=None)
def thaw(self, array):
def _thaw(ary):
return ary.with_queue(self.queue)
return with_array_context(self._rec_map_container(_thaw, array), actx=self)
def tag(self, tags: ToTagSetConvertible, array):
def _tag(ary):
return ary.tagged(tags)
return self._rec_map_container(_tag, array)
def tag_axis(self, iaxis: int, tags: ToTagSetConvertible, array):
def _tag_axis(ary):
return ary.with_tagged_axis(iaxis, tags)
return self._rec_map_container(_tag_axis, array)
def call_loopy(self, t_unit, **kwargs):
t_unit = self.transform_loopy_program(t_unit)
from arraycontext.loopy import get_default_entrypoint
default_entrypoint = get_default_entrypoint(t_unit)
prg_name = default_entrypoint.name
try:
executor = self._loopy_transform_cache[t_unit]
except KeyError:
orig_t_unit = t_unit
executor = self.transform_loopy_program(t_unit).executor(self.context)
self._loopy_transform_cache[orig_t_unit] = executor
del orig_t_unit
evt, result = t_unit(self.queue, **kwargs, allocator=self.allocator)
evt, result = executor(self.queue, **kwargs, allocator=self.allocator)
if self._wait_event_queue_length is not False:
prg_name = executor.t_unit.default_entrypoint.name
wait_event_queue = self._kernel_name_to_wait_event_queue.setdefault(
prg_name, [])
......@@ -315,24 +261,42 @@ class PyOpenCLArrayContext(ArrayContext):
if len(wait_event_queue) > self._wait_event_queue_length:
wait_event_queue.pop(0).wait()
return result
import arraycontext.impl.pyopencl.taggable_cl_array as tga
def freeze(self, array):
array.finish()
return array.with_queue(None)
# FIXME: Inherit loopy tags for these arrays
return {name: tga.to_tagged_cl_array(ary) for name, ary in result.items()}
def thaw(self, array):
return array.with_queue(self.queue)
def clone(self):
if self._passed_force_device_scalars:
return type(self)(self.queue, self.allocator,
wait_event_queue_length=self._wait_event_queue_length,
force_device_scalars=True)
else:
return type(self)(self.queue, self.allocator,
wait_event_queue_length=self._wait_event_queue_length)
# }}}
@memoize_method
def transform_loopy_program(self, t_unit):
# {{{ transform_loopy_program
def transform_loopy_program(self, t_unit: lp.TranslationUnit) -> lp.TranslationUnit:
from warnings import warn
warn("Using the base "
f"{type(self).__name__}.transform_loopy_program "
"to transform a translation unit. "
"This is largely a no-op and unlikely to result in fast generated "
"code."
f"Instead, subclass {type(self).__name__} and implement "
"the specific transform logic required to transform the program "
"for your package or application. Check higher-level packages "
"(e.g. meshmode), which may already have subclasses you may want "
"to build on.",
UntransformedCodeWarning, stacklevel=2)
# accommodate loopy with and without kernel callables
import loopy as lp
from arraycontext.loopy import get_default_entrypoint
default_entrypoint = get_default_entrypoint(t_unit)
default_entrypoint = t_unit.default_entrypoint
options = default_entrypoint.options
if not (options.return_dict and options.no_numpy):
raise ValueError("Loopy kernel passed to call_loopy must "
......@@ -341,48 +305,41 @@ class PyOpenCLArrayContext(ArrayContext):
"to create this kernel?")
all_inames = default_entrypoint.all_inames()
# FIXME: This could be much smarter.
inner_iname = None
if (len(default_entrypoint.instructions) == 1
and isinstance(default_entrypoint.instructions[0], lp.Assignment)
and any(isinstance(tag, FirstAxisIsElementsTag)
# FIXME: Firedrake branch lacks kernel tags
for tag in getattr(default_entrypoint, "tags", ()))):
stmt, = default_entrypoint.instructions
out_inames = [v.name for v in stmt.assignee.index_tuple]
assert out_inames
outer_iname = out_inames[0]
if len(out_inames) >= 2:
inner_iname = out_inames[1]
elif "iel" in all_inames:
outer_iname = "iel"
if "idof" in all_inames:
inner_iname = "idof"
elif "i0" in all_inames:
if "i0" in all_inames:
outer_iname = "i0"
if "i1" in all_inames:
inner_iname = "i1"
else:
raise RuntimeError(
"Unable to reason what outer_iname and inner_iname "
f"needs to be; all_inames is given as: {all_inames}"
)
return t_unit
if inner_iname is not None:
t_unit = lp.split_iname(t_unit, inner_iname, 16, inner_tag="l.0")
return lp.tag_inames(t_unit, {outer_iname: "g.0"})
t_unit = lp.tag_inames(t_unit, {outer_iname: "g.0"})
return t_unit
# }}}
# {{{ properties
@property
def permits_inplace_modification(self):
return True
def tag(self, tags: Union[Sequence[Tag], Tag], array):
# Sorry, not capable.
return array
@property
def supports_nonscalar_broadcasting(self):
return False
def tag_axis(self, iaxis, tags: Union[Sequence[Tag], Tag], array):
# Sorry, not capable.
return array
@property
def permits_advanced_indexing(self):
return False
# }}}
# }}}
......
arraycontext/impl/pyopencl/fake_numpy.py 0 → 100644
View file @ 1013ecb2
"""
.. currentmodule:: arraycontext
.. autoclass:: PyOpenCLArrayContext
"""
from __future__ import annotations
__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.
"""
import operator
from functools import partial, reduce
import numpy as np
from arraycontext.container import NotAnArrayContainerError, serialize_container
from arraycontext.container.traversal import (
rec_map_array_container,
rec_map_reduce_array_container,
rec_multimap_array_container,
rec_multimap_reduce_array_container,
)
from arraycontext.context import Array, ArrayOrContainer
from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace
from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
from arraycontext.loopy import LoopyBasedFakeNumpyNamespace
try:
import pyopencl as cl # noqa: F401
import pyopencl.array as cl_array
except ImportError:
pass
# {{{ fake numpy
class PyOpenCLFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace):
def _get_fake_numpy_linalg_namespace(self):
return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context)
# NOTE: the order of these follows the order in numpy docs
# NOTE: when adding a function here, also add it to `array_context.rst` docs!
# {{{ array creation routines
def zeros(self, shape, dtype) -> TaggableCLArray:
import arraycontext.impl.pyopencl.taggable_cl_array as tga
return tga.zeros(self._array_context.queue, shape, dtype,
allocator=self._array_context.allocator)
def empty_like(self, ary):
from warnings import warn
warn(f"{type(self._array_context).__name__}.np.empty_like is "
"deprecated and will stop working in 2023. Prefer actx.np.zeros_like "
"instead.",
DeprecationWarning, stacklevel=2)
import arraycontext.impl.pyopencl.taggable_cl_array as tga
actx = self._array_context
def _empty_like(array):
return tga.empty(actx.queue, array.shape, array.dtype,
allocator=actx.allocator, axes=array.axes, tags=array.tags)
return actx._rec_map_container(_empty_like, ary)
def zeros_like(self, ary):
import arraycontext.impl.pyopencl.taggable_cl_array as tga
actx = self._array_context
def _zeros_like(array):
return tga.zeros(
actx.queue, array.shape, array.dtype,
allocator=actx.allocator, axes=array.axes, tags=array.tags)
return actx._rec_map_container(_zeros_like, ary, default_scalar=0)
def ones_like(self, ary):
return self.full_like(ary, 1)
def full_like(self, ary, fill_value):
import arraycontext.impl.pyopencl.taggable_cl_array as tga
actx = self._array_context
def _full_like(subary):
filled = tga.empty(
actx.queue, subary.shape, subary.dtype,
allocator=actx.allocator, axes=subary.axes, tags=subary.tags)
filled.fill(fill_value)
return filled
return actx._rec_map_container(_full_like, ary, default_scalar=fill_value)
def copy(self, ary):
def _copy(subary):
return subary.copy(queue=self._array_context.queue)
return self._array_context._rec_map_container(_copy, ary)
def arange(self, *args, **kwargs):
return cl_array.arange(self._array_context.queue, *args, **kwargs)
# }}}
# {{{ array manipulation routines
def reshape(self, a, newshape, order="C"):
return rec_map_array_container(
lambda ary: ary.reshape(newshape, order=order),
a)
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)
def concatenate(self, arrays, axis=0):
return cl_array.concatenate(
arrays, axis,
self._array_context.queue,
self._array_context.allocator
)
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)
# }}}
# {{{ linear algebra
def vdot(self, x, y, dtype=None):
return rec_multimap_reduce_array_container(
sum,
partial(cl_array.vdot, dtype=dtype, queue=self._array_context.queue),
x, y)
# }}}
# {{{ logic functions
def all(self, a):
queue = self._array_context.queue
def _all(ary):
if np.isscalar(ary):
return np.int8(all([ary]))
return ary.all(queue=queue)
return rec_map_reduce_array_container(
partial(reduce, partial(cl_array.minimum, queue=queue)),
_all,
a)
def any(self, a):
queue = self._array_context.queue
def _any(ary):
if np.isscalar(ary):
return np.int8(any([ary]))
return ary.any(queue=queue)
return rec_map_reduce_array_container(
partial(reduce, partial(cl_array.maximum, queue=queue)),
_any,
a)
def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
actx = self._array_context
queue = actx.queue
# NOTE: pyopencl doesn't like `bool` much, so use `int8` instead
true_ary = actx.from_numpy(np.int8(True))
false_ary = actx.from_numpy(np.int8(False))
def rec_equal(x: ArrayOrContainer, y: ArrayOrContainer) -> cl_array.Array:
if type(x) is not type(y):
return false_ary
try:
serialized_x = serialize_container(x)
serialized_y = serialize_container(y)
except NotAnArrayContainerError:
assert isinstance(x, cl_array.Array)
assert isinstance(y, cl_array.Array)
if x.shape != y.shape:
return false_ary
else:
return (x == y).all()
else:
if len(serialized_x) != len(serialized_y):
return false_ary
return reduce(
partial(cl_array.minimum, queue=queue),
[(true_ary if kx_i == ky_i else false_ary)
and rec_equal(x_i, y_i)
for (kx_i, x_i), (ky_i, y_i)
in zip(serialized_x, serialized_y, strict=True)],
true_ary)
return rec_equal(a, b)
# 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 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 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 logical_or(self, x, y):
return rec_multimap_array_container(cl_array.logical_or, x, y)
def logical_and(self, x, y):
return rec_multimap_array_container(cl_array.logical_and, x, y)
def logical_not(self, x):
return rec_map_array_container(cl_array.logical_not, x)
# }}}
# {{{ mathematical functions
def sum(self, a, axis=None, dtype=None):
if isinstance(axis, int):
axis = axis,
def _rec_sum(ary):
if axis not in [None, tuple(range(ary.ndim))]:
raise NotImplementedError(f"Sum over '{axis}' axes not supported.")
return cl_array.sum(ary, dtype=dtype, queue=self._array_context.queue)
return rec_map_reduce_array_container(sum, _rec_sum, a)
def maximum(self, x, y):
return rec_multimap_array_container(
partial(cl_array.maximum, queue=self._array_context.queue),
x, y)
def amax(self, a, axis=None):
queue = self._array_context.queue
if isinstance(axis, int):
axis = axis,
def _rec_max(ary):
if axis not in [None, tuple(range(ary.ndim))]:
raise NotImplementedError(f"Max. over '{axis}' axes not supported.")
return cl_array.max(ary, queue=queue)
return rec_map_reduce_array_container(
partial(reduce, partial(cl_array.maximum, queue=queue)),
_rec_max,
a)
max = amax
def minimum(self, x, y):
return rec_multimap_array_container(
partial(cl_array.minimum, queue=self._array_context.queue),
x, y)
def amin(self, a, axis=None):
queue = self._array_context.queue
if isinstance(axis, int):
axis = axis,
def _rec_min(ary):
if axis not in [None, tuple(range(ary.ndim))]:
raise NotImplementedError(f"Min. over '{axis}' axes not supported.")
return cl_array.min(ary, queue=queue)
return rec_map_reduce_array_container(
partial(reduce, partial(cl_array.minimum, queue=queue)),
_rec_min,
a)
min = amin
def absolute(self, a):
return self.abs(a)
# }}}
# {{{ sorting, searching, and counting
def where(self, criterion, then, else_):
def where_inner(inner_crit, inner_then, inner_else):
if isinstance(inner_crit, bool | np.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_)
# }}}
# }}}
# {{{ fake np.linalg
class _PyOpenCLFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
pass
# }}}
# vim: foldmethod=marker
arraycontext/impl/pyopencl/taggable_cl_array.py 0 → 100644
View file @ 1013ecb2
"""
.. autoclass:: TaggableCLArray
.. autoclass:: Axis
.. autofunction:: to_tagged_cl_array
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import numpy as np
import pyopencl.array as cla
from pytools import memoize
from pytools.tag import Tag, Taggable, ToTagSetConvertible
# {{{ utils
@dataclass(frozen=True, eq=True)
class Axis(Taggable):
"""
Records the tags corresponding to a dimension of :class:`TaggableCLArray`.
"""
tags: frozenset[Tag]
def _with_new_tags(self, tags: frozenset[Tag]) -> Axis:
from dataclasses import replace
return replace(self, tags=tags)
@memoize
def _construct_untagged_axes(ndim: int) -> tuple[Axis, ...]:
return tuple(Axis(frozenset()) for _ in range(ndim))
def _unwrap_cl_array(ary: cla.Array) -> dict[str, Any]:
return {
"shape": ary.shape,
"dtype": ary.dtype,
"allocator": ary.allocator,
"strides": ary.strides,
"data": ary.base_data,
"offset": ary.offset,
"events": ary.events,
"_context": ary.context,
"_queue": ary.queue,
"_size": ary.size,
"_fast": True,
}
# }}}
# {{{ TaggableCLArray
class TaggableCLArray(cla.Array, Taggable):
"""
A :class:`pyopencl.array.Array` with additional metadata. This is used by
:class:`~arraycontext.PytatoPyOpenCLArrayContext` to preserve tags for data
while frozen, and also in a similar capacity by
:class:`~arraycontext.PyOpenCLArrayContext`.
.. attribute:: axes
A :class:`tuple` of instances of :class:`Axis`, with one :class:`Axis`
for each dimension of the array.
.. attribute:: tags
A :class:`frozenset` of :class:`pytools.tag.Tag`. Typically intended to
record application-specific metadata to drive the optimizations in
:meth:`arraycontext.PyOpenCLArrayContext.transform_loopy_program`.
"""
def __init__(self, cq, shape, dtype, order="C", allocator=None,
data=None, offset=0, strides=None, events=None, _flags=None,
_fast=False, _size=None, _context=None, _queue=None,
axes=None, tags=frozenset()):
super().__init__(cq=cq, shape=shape, dtype=dtype,
order=order, allocator=allocator,
data=data, offset=offset,
strides=strides, events=events,
_flags=_flags, _fast=_fast,
_size=_size, _context=_context,
_queue=_queue)
if __debug__:
if not isinstance(tags, frozenset):
raise TypeError("tags are not a frozenset")
if axes is not None and len(axes) != self.ndim:
raise ValueError("axes length does not match array dimension: "
f"got {len(axes)} axes for {self.ndim}d array")
if axes is None:
axes = _construct_untagged_axes(self.ndim)
self.tags = tags
self.axes = axes
def __repr__(self) -> str:
return (f"{type(self).__name__}(shape={self.shape}, dtype={self.dtype}, "
f"tags={self.tags}, axes={self.axes})")
def copy(self, queue=cla._copy_queue):
ary = super().copy(queue=queue)
return type(self)(None, tags=self.tags, axes=self.axes,
**_unwrap_cl_array(ary))
def _with_new_tags(self, tags: frozenset[Tag]) -> TaggableCLArray:
return type(self)(None, tags=tags, axes=self.axes,
**_unwrap_cl_array(self))
def with_tagged_axis(self, iaxis: int,
tags: ToTagSetConvertible) -> TaggableCLArray:
"""
Returns a copy of *self* with *iaxis*-th axis tagged with *tags*.
"""
new_axes = (*self.axes[:iaxis],
self.axes[iaxis].tagged(tags),
*self.axes[iaxis + 1:])
return type(self)(None, tags=self.tags, axes=new_axes,
**_unwrap_cl_array(self))
def to_tagged_cl_array(ary: cla.Array,
axes: tuple[Axis, ...] | None = None,
tags: frozenset[Tag] = frozenset()) -> TaggableCLArray:
"""
Returns a :class:`TaggableCLArray` that is constructed from the data in
*ary* along with the metadata from *axes* and *tags*. If *ary* is already a
:class:`TaggableCLArray`, the new *tags* and *axes* are added to the
existing ones.
:arg axes: An instance of :class:`Axis` for each dimension of the
array. If passed *None*, then initialized to a :class:`pytato.Axis`
with no tags attached for each dimension.
"""
if axes is not None and len(axes) != ary.ndim:
raise ValueError("axes length does not match array dimension: "
f"got {len(axes)} axes for {ary.ndim}d array")
from pytools.tag import normalize_tags
tags = normalize_tags(tags)
if isinstance(ary, TaggableCLArray):
if axes is not None:
for i, axis in enumerate(axes):
ary = ary.with_tagged_axis(i, axis.tags)
if tags:
ary = ary.tagged(tags)
return ary
elif isinstance(ary, cla.Array):
return TaggableCLArray(None, tags=tags, axes=axes,
**_unwrap_cl_array(ary))
else:
raise TypeError(f"unsupported array type: '{type(ary).__name__}'")
# }}}
# {{{ creation
def empty(queue, shape, dtype=float, *,
axes: tuple[Axis, ...] | None = None,
tags: frozenset[Tag] = frozenset(),
order: str = "C",
allocator=None) -> TaggableCLArray:
if dtype is not None:
dtype = np.dtype(dtype)
return TaggableCLArray(
queue, shape, dtype,
axes=axes, tags=tags,
order=order, allocator=allocator)
def zeros(queue, shape, dtype=float, *,
axes: tuple[Axis, ...] | None = None,
tags: frozenset[Tag] = frozenset(),
order: str = "C",
allocator=None) -> TaggableCLArray:
result = empty(
queue, shape, dtype=dtype, axes=axes, tags=tags,
order=order, allocator=allocator)
result._zero_fill()
return result
def to_device(queue, ary, *,
axes: tuple[Axis, ...] | None = None,
tags: frozenset[Tag] = frozenset(),
allocator=None):
return to_tagged_cl_array(
cla.to_device(queue, ary, allocator=allocator),
axes=axes, tags=tags)
# }}}
arraycontext/impl/pytato/__init__.py 0 → 100644
View file @ 1013ecb2
from __future__ import annotations
__doc__ = """
.. currentmodule:: arraycontext
A :mod:`pytato`-based array context defers the evaluation of an array until it is
frozen. The execution contexts for the evaluations are specific to an
:class:`~arraycontext.ArrayContext` type. For example,
:class:`~arraycontext.PytatoPyOpenCLArrayContext` uses :mod:`pyopencl` to
JIT-compile and execute the array expressions.
The following :mod:`pytato`-based array contexts are provided:
.. autoclass:: PytatoPyOpenCLArrayContext
.. autoclass:: PytatoJAXArrayContext
Compiling a Python callable (Internal)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. automodule:: arraycontext.impl.pytato.compile
Utils
^^^^^
.. automodule:: arraycontext.impl.pytato.utils
"""
__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.
"""
import abc
import sys
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import numpy as np
from pytools import memoize_method
from pytools.tag import Tag, ToTagSetConvertible, normalize_tags
from arraycontext.container.traversal import rec_map_array_container, with_array_context
from arraycontext.context import (
Array,
ArrayContext,
ArrayOrContainer,
ScalarLike,
UntransformedCodeWarning,
)
from arraycontext.metadata import NameHint
if TYPE_CHECKING:
import loopy as lp
import pytato
if getattr(sys, "_BUILDING_SPHINX_DOCS", False):
import pyopencl as cl
import logging
logger = logging.getLogger(__name__)
# {{{ tag conversion
def _preprocess_array_tags(tags: ToTagSetConvertible) -> frozenset[Tag]:
tags = normalize_tags(tags)
name_hints = [tag for tag in tags if isinstance(tag, NameHint)]
if name_hints:
name_hint, = name_hints
from pytato.tags import PrefixNamed
prefix_nameds = [tag for tag in tags if isinstance(tag, PrefixNamed)]
if prefix_nameds:
prefix_named, = prefix_nameds
from warnings import warn
warn("When converting a "
f"arraycontext.metadata.NameHint('{name_hint.name}') "
"to pytato.tags.PrefixNamed, "
f"PrefixNamed('{prefix_named.prefix}') "
"was already present.", stacklevel=1)
tags = (
(tags | frozenset({PrefixNamed(name_hint.name)}))
- {name_hint})
return tags
# }}}
class _NotOnlyDataWrappers(Exception): # noqa: N818
pass
# {{{ _BasePytatoArrayContext
class _BasePytatoArrayContext(ArrayContext, abc.ABC):
"""
An abstract :class:`ArrayContext` that uses :mod:`pytato` data types to
represent.
.. automethod:: __init__
.. automethod:: transform_dag
.. automethod:: compile
"""
def __init__(
self, *,
compile_trace_callback: Callable[[Any, str, Any], None] | None = None
) -> None:
"""
:arg compile_trace_callback: A function of three arguments
*(what, stage, ir)*, where *what* identifies the object
being compiled, *stage* is a string describing the compilation
pass, and *ir* is an object containing the intermediate
representation. This interface should be considered
unstable.
"""
super().__init__()
import pytato as pt
self._freeze_prg_cache: dict[pt.DictOfNamedArrays, lp.TranslationUnit] = {}
self._dag_transform_cache: dict[
pt.DictOfNamedArrays,
tuple[pt.DictOfNamedArrays, str]] = {}
if compile_trace_callback is None:
def _compile_trace_callback(what, stage, ir):
pass
compile_trace_callback = _compile_trace_callback
self._compile_trace_callback = compile_trace_callback
def _get_fake_numpy_namespace(self):
from arraycontext.impl.pytato.fake_numpy import PytatoFakeNumpyNamespace
return PytatoFakeNumpyNamespace(self)
@abc.abstractproperty
def _frozen_array_types(self) -> tuple[type, ...]:
"""
Returns valid frozen array types for the array context.
"""
# {{{ compilation
def transform_dag(self, dag: pytato.DictOfNamedArrays
) -> pytato.DictOfNamedArrays:
"""
Returns a transformed version of *dag*. Sub-classes are supposed to
override this method to implement context-specific transformations on
*dag* (most likely to perform domain-specific optimizations). Every
:mod:`pytato` DAG that is compiled to a GPU-kernel is
passed through this routine.
:arg dag: An instance of :class:`pytato.DictOfNamedArrays`
:returns: A transformed version of *dag*.
"""
return dag
def transform_loopy_program(self, t_unit: lp.TranslationUnit) -> lp.TranslationUnit:
from warnings import warn
warn("Using the base "
f"{type(self).__name__}.transform_loopy_program "
"to transform a translation unit. "
"This is a no-op and will result in unoptimized C code for"
"the requested optimization, all in a single statement."
"This will work, but is unlikely to be performant."
f"Instead, subclass {type(self).__name__} and implement "
"the specific transform logic required to transform the program "
"for your package or application. Check higher-level packages "
"(e.g. meshmode), which may already have subclasses you may want "
"to build on.",
UntransformedCodeWarning, stacklevel=2)
return t_unit
@abc.abstractmethod
def einsum(self, spec, *args, arg_names=None, tagged=()):
pass
# }}}
# {{{ properties
@property
def permits_inplace_modification(self):
return False
@property
def supports_nonscalar_broadcasting(self):
return True
@property
def permits_advanced_indexing(self):
return True
def get_target(self):
return None
# }}}
# }}}
# {{{ PytatoPyOpenCLArrayContext
@dataclass
class ProfileEvent:
"""Holds a profile event that has not been collected by the profiler yet."""
start_cl_event: cl._cl.Event
stop_cl_event: cl._cl.Event
t_unit_name: str
class PytatoPyOpenCLArrayContext(_BasePytatoArrayContext):
"""
An :class:`ArrayContext` that uses :mod:`pytato` data types to represent
the arrays targeting OpenCL for offloading operations.
.. attribute:: queue
A :class:`pyopencl.CommandQueue`.
.. attribute:: allocator
A :mod:`pyopencl` memory allocator. Can also be None (default) or False
to use the default allocator.
.. automethod:: __init__
.. automethod:: transform_dag
.. automethod:: compile
"""
def __init__(
self, queue: cl.CommandQueue, allocator=None, *,
use_memory_pool: bool | None = None,
compile_trace_callback: Callable[[Any, str, Any], None] | None = None,
profile_kernels: bool = False,
# do not use: only for testing
_force_svm_arg_limit: int | None = None,
) -> None:
"""
:arg compile_trace_callback: A function of three arguments
*(what, stage, ir)*, where *what* identifies the object
being compiled, *stage* is a string describing the compilation
pass, and *ir* is an object containing the intermediate
representation. This interface should be considered
unstable.
"""
if allocator is not None and use_memory_pool is not None:
raise TypeError("may not specify both allocator and use_memory_pool")
self.using_svm = None
if allocator is None:
from pyopencl.characterize import has_coarse_grain_buffer_svm
has_svm = has_coarse_grain_buffer_svm(queue.device)
if has_svm:
self.using_svm = True
from pyopencl.tools import SVMAllocator
allocator = SVMAllocator(queue.context, queue=queue)
if use_memory_pool:
from pyopencl.tools import SVMPool
allocator = SVMPool(allocator)
else:
self.using_svm = False
from pyopencl.tools import ImmediateAllocator
allocator = ImmediateAllocator(queue)
if use_memory_pool:
from pyopencl.tools import MemoryPool
allocator = MemoryPool(allocator)
else:
# Check whether the passed allocator allocates SVM
try:
from pyopencl import SVMPointer
mem = allocator(4)
if isinstance(mem, SVMPointer):
self.using_svm = True
else:
self.using_svm = False
except ImportError:
self.using_svm = False
import pyopencl.array as cla
import pytato as pt
super().__init__(compile_trace_callback=compile_trace_callback)
self.queue = queue
self.allocator = allocator
self.array_types = (pt.Array, cla.Array)
# unused, but necessary to keep the context alive
self.context = self.queue.context
self._force_svm_arg_limit = _force_svm_arg_limit
self._enable_profiling(profile_kernels)
# {{{ Profiling functionality
def _enable_profiling(self, enable: bool) -> None:
# List of ProfileEvents that haven't been transferred to profiled
# results yet
self._profile_events: list[ProfileEvent] = []
# Dict of kernel name -> list of kernel execution times
self._profile_results: dict[str, list[int]] = {}
if enable:
import pyopencl as cl
if not self.queue.properties & cl.command_queue_properties.PROFILING_ENABLE:
raise RuntimeError("Profiling was not enabled in the command queue. "
"Please create the queue with "
"cl.command_queue_properties.PROFILING_ENABLE.")
self.profile_kernels = True
else:
self.profile_kernels = False
def _wait_and_transfer_profile_events(self) -> None:
"""Wait for all profiling events to finish and transfer the results
to *self._profile_results*."""
import pyopencl as cl
# First, wait for completion of all events
if self._profile_events:
cl.wait_for_events([p_event.stop_cl_event
for p_event in self._profile_events])
# Then, collect all events and store them
for t in self._profile_events:
name = t.t_unit_name
time = t.stop_cl_event.profile.end - t.start_cl_event.profile.end
self._profile_results.setdefault(name, []).append(time)
self._profile_events = []
def _add_profiling_events(self, start: cl._cl.Event, stop: cl._cl.Event,
t_unit_name: str) -> None:
"""Add profiling events to the list of profiling events."""
self._profile_events.append(ProfileEvent(start, stop, t_unit_name))
def _reset_profiling_data(self) -> None:
"""Reset profiling data."""
self._profile_results = {}
# }}}
@property
def _frozen_array_types(self) -> tuple[type, ...]:
import pyopencl.array as cla
return (cla.Array,)
def _rec_map_container(
self, func: Callable[[Array], Array], array: ArrayOrContainer,
allowed_types: tuple[type, ...] | None = None, *,
default_scalar: ScalarLike | None = None,
strict: bool = False) -> ArrayOrContainer:
import pytato as pt
import arraycontext.impl.pyopencl.taggable_cl_array as tga
if allowed_types is None:
allowed_types = (pt.Array, tga.TaggableCLArray)
def _wrapper(ary):
if isinstance(ary, allowed_types):
return func(ary)
elif np.isscalar(ary):
if default_scalar is None:
return ary
else:
return np.array(ary).dtype.type(default_scalar)
else:
raise TypeError(
f"{func.__qualname__} invoked with "
f"an unsupported array type: got '{type(ary).__name__}', "
f"but expected one of {allowed_types}")
return rec_map_array_container(_wrapper, array)
# {{{ ArrayContext interface
def from_numpy(self, array):
import pytato as pt
import arraycontext.impl.pyopencl.taggable_cl_array as tga
def _from_numpy(ary):
return pt.make_data_wrapper(
tga.to_device(self.queue, ary, allocator=self.allocator)
)
return with_array_context(
self._rec_map_container(_from_numpy, array, (np.ndarray,), strict=True),
actx=self)
def to_numpy(self, array):
def _to_numpy(ary):
return ary.get(queue=self.queue)
return with_array_context(
self._rec_map_container(_to_numpy, self.freeze(array)),
actx=None)
@memoize_method
def get_target(self):
import pyopencl as cl
import pyopencl.characterize as cl_char
dev = self.queue.device
if (
self._force_svm_arg_limit is not None
or (
self.using_svm and dev.type & cl.device_type.GPU
and cl_char.has_coarse_grain_buffer_svm(dev))):
if dev.max_parameter_size == 4352:
# Nvidia devices and PTXAS declare a limit of 4352 bytes,
# which is incorrect. The CUDA documentation at
# https://docs.nvidia.com/cuda/cuda-c-programming-guide/#function-parameters
# mentions a limit of 4KB, which is also incorrect.
# As far as I can tell, the actual limit is around 4080
# bytes, at least on a K40. Reducing the limit further
# in order to be on the safe side.
# Note that the naming convention isn't super consistent
# for Nvidia GPUs, so that we only use the maximum
# parameter size to determine if it is an Nvidia GPU.
limit = 4096-200
from warnings import warn
warn("Running on an Nvidia GPU, reducing the argument "
f"size limit from 4352 to {limit}.", stacklevel=1)
else:
limit = dev.max_parameter_size
if self._force_svm_arg_limit is not None:
limit = self._force_svm_arg_limit
logger.info(
"limiting argument buffer size for %s to %d bytes",
dev, limit)
from arraycontext.impl.pytato.utils import (
ArgSizeLimitingPytatoLoopyPyOpenCLTarget,
)
return ArgSizeLimitingPytatoLoopyPyOpenCLTarget(limit)
else:
return super().get_target()
def freeze(self, array):
if np.isscalar(array):
return array
import pyopencl.array as cla
import pytato as pt
from arraycontext.container.traversal import rec_keyed_map_array_container
from arraycontext.impl.pyopencl.taggable_cl_array import (
TaggableCLArray,
to_tagged_cl_array,
)
from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
from arraycontext.impl.pytato.utils import (
_normalize_pt_expr,
get_cl_axes_from_pt_axes,
)
array_as_dict: dict[str, cla.Array | TaggableCLArray | pt.Array] = {}
key_to_frozen_subary: dict[str, TaggableCLArray] = {}
key_to_pt_arrays: dict[str, pt.Array] = {}
def _record_leaf_ary_in_dict(
key: tuple[Any, ...],
ary: cla.Array | TaggableCLArray | pt.Array) -> None:
key_str = "_ary" + _ary_container_key_stringifier(key)
array_as_dict[key_str] = ary
rec_keyed_map_array_container(_record_leaf_ary_in_dict, array)
# {{{ remove any non pytato arrays from array_as_dict
for key, subary in array_as_dict.items():
if isinstance(subary, TaggableCLArray):
key_to_frozen_subary[key] = subary.with_queue(None)
elif isinstance(subary, self._frozen_array_types):
from warnings import warn
warn(f"Invoking {type(self).__name__}.freeze with"
f" {type(subary).__name__} will be unsupported in 2023. Use"
" `to_tagged_cl_array` to convert instances to TaggableCLArray.",
DeprecationWarning, stacklevel=2)
key_to_frozen_subary[key] = (
to_tagged_cl_array(subary.with_queue(None)))
elif isinstance(subary, pt.DataWrapper):
# trivial freeze.
key_to_frozen_subary[key] = to_tagged_cl_array(
subary.data,
axes=get_cl_axes_from_pt_axes(subary.axes),
tags=subary.tags)
elif isinstance(subary, pt.Array):
# Don't be tempted to take shortcuts here, e.g. for empty
# arrays, as this will inhibit metadata propagation that
# may happen in transform_dag below. See
# https://github.com/inducer/arraycontext/pull/167#issuecomment-1151877480
key_to_pt_arrays[key] = subary
else:
raise TypeError(
f"{type(self).__name__}.freeze invoked with an unsupported "
f"array type: got '{type(subary).__name__}', but expected one "
f"of {self.array_types}")
# }}}
def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
key_str = "_ary" + _ary_container_key_stringifier(key)
return key_to_frozen_subary[key_str]
if not key_to_pt_arrays:
# all cl arrays => no need to perform any codegen
return with_array_context(
rec_keyed_map_array_container(_to_frozen, array),
actx=None)
pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(
key_to_pt_arrays)
normalized_expr, bound_arguments = _normalize_pt_expr(
pt_dict_of_named_arrays)
try:
pt_prg = self._freeze_prg_cache[normalized_expr]
except KeyError:
try:
transformed_dag, function_name = (
self._dag_transform_cache[normalized_expr])
except KeyError:
transformed_dag = self.transform_dag(normalized_expr)
from pytato.tags import PrefixNamed
name_hint_tags = []
for subary in key_to_pt_arrays.values():
name_hint_tags.extend(subary.tags_of_type(PrefixNamed))
from pytools import common_prefix
name_hint = common_prefix([nh.prefix for nh in name_hint_tags])
# All name_hint_tags shared at least some common prefix.
function_name = f"frozen_{name_hint}" if name_hint else "frozen_result"
self._dag_transform_cache[normalized_expr] = (
transformed_dag, function_name)
from arraycontext.loopy import _DEFAULT_LOOPY_OPTIONS
opts = _DEFAULT_LOOPY_OPTIONS
assert opts.return_dict
pt_prg = pt.generate_loopy(transformed_dag,
options=opts,
function_name=function_name,
target=self.get_target()
).bind_to_context(self.context)
pt_prg = pt_prg.with_transformed_translation_unit(
self.transform_loopy_program)
self._freeze_prg_cache[normalized_expr] = pt_prg
else:
transformed_dag, function_name = (
self._dag_transform_cache[normalized_expr])
assert len(pt_prg.bound_arguments) == 0
if self.profile_kernels:
import pyopencl as cl
start_evt = cl.enqueue_marker(self.queue)
evt, out_dict = pt_prg(self.queue,
allocator=self.allocator,
**bound_arguments)
if self.profile_kernels:
self._add_profiling_events(start_evt, evt, pt_prg.program.entrypoint)
assert len(set(out_dict) & set(key_to_frozen_subary)) == 0
key_to_frozen_subary = {
**key_to_frozen_subary,
**{k: to_tagged_cl_array(
v.with_queue(None),
axes=get_cl_axes_from_pt_axes(transformed_dag[k].expr.axes),
tags=transformed_dag[k].expr.tags)
for k, v in out_dict.items()}
}
return with_array_context(
rec_keyed_map_array_container(_to_frozen, array),
actx=None)
def thaw(self, array):
import pytato as pt
import arraycontext.impl.pyopencl.taggable_cl_array as tga
from .utils import get_pt_axes_from_cl_axes
def _thaw(ary):
return pt.make_data_wrapper(ary.with_queue(self.queue),
axes=get_pt_axes_from_cl_axes(ary.axes),
tags=ary.tags)
return with_array_context(
self._rec_map_container(_thaw, array, (tga.TaggableCLArray,)),
actx=self)
def freeze_thaw(self, array):
import pytato as pt
import arraycontext.impl.pyopencl.taggable_cl_array as tga
def _ft(ary):
if isinstance(ary, (pt.DataWrapper, tga.TaggableCLArray)):
return ary
else:
raise _NotOnlyDataWrappers()
try:
return with_array_context(
self._rec_map_container(_ft, array),
actx=self)
except _NotOnlyDataWrappers:
return super().freeze_thaw(array)
def tag(self, tags: ToTagSetConvertible, array):
def _tag(ary):
return ary.tagged(_preprocess_array_tags(tags))
return self._rec_map_container(_tag, array)
def tag_axis(self, iaxis, tags: ToTagSetConvertible, array):
def _tag_axis(ary):
return ary.with_tagged_axis(iaxis, tags)
return self._rec_map_container(_tag_axis, array)
# }}}
# {{{ compilation
def call_loopy(self, program, **kwargs):
import pytato as pt
from pytato.loopy import call_loopy
from pytato.scalar_expr import SCALAR_CLASSES
from arraycontext.impl.pyopencl.taggable_cl_array import TaggableCLArray
entrypoint = program.default_entrypoint.name
# {{{ preprocess args
processed_kwargs = {}
for kw, arg in sorted(kwargs.items()):
if isinstance(arg, (pt.Array, *SCALAR_CLASSES)):
pass
elif isinstance(arg, TaggableCLArray):
arg = self.thaw(arg)
else:
raise ValueError(f"call_loopy argument '{kw}' expected to be an"
" instance of 'pytato.Array', 'Number' or"
f"'TaggableCLArray', got '{type(arg)}'")
processed_kwargs[kw] = arg
# }}}
return call_loopy(program, processed_kwargs, entrypoint)
def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
from .compile import LazilyPyOpenCLCompilingFunctionCaller
return LazilyPyOpenCLCompilingFunctionCaller(self, f)
def transform_dag(self, dag: pytato.DictOfNamedArrays
) -> pytato.DictOfNamedArrays:
import pytato as pt
dag = pt.transform.materialize_with_mpms(dag)
return dag
def einsum(self, spec, *args, arg_names=None, tagged=()):
import pytato as pt
import arraycontext.impl.pyopencl.taggable_cl_array as tga
if arg_names is None:
arg_names = (None,) * len(args)
def preprocess_arg(name, arg):
if isinstance(arg, tga.TaggableCLArray):
ary = self.thaw(arg)
elif isinstance(arg, self._frozen_array_types):
from warnings import warn
warn(f"Invoking {type(self).__name__}.einsum with"
f" {type(arg).__name__} will be unsupported in 2023. Use"
" `to_tagged_cl_array` to convert instances to TaggableCLArray.",
DeprecationWarning, stacklevel=2)
ary = self.thaw(tga.to_tagged_cl_array(arg))
elif isinstance(arg, pt.Array):
ary = arg
else:
raise TypeError(
f"{type(self).__name__}.einsum invoked with an unsupported "
f"array type: got '{type(arg).__name__}', but expected one "
f"of {self.array_types}")
if name is not None: # noqa: SIM102
# Tagging Placeholders with naming-related tags is pointless:
# They already have names. It's also counterproductive, as
# multiple placeholders with the same name that are not
# also the same object are not allowed, and this would produce
# a different Placeholder object of the same name.
if (not isinstance(ary, pt.Placeholder)
and not ary.tags_of_type(NameHint)):
ary = ary.tagged(NameHint(name))
return ary
return pt.einsum(spec, *[
preprocess_arg(name, arg)
for name, arg in zip(arg_names, args, strict=True)
]).tagged(_preprocess_array_tags(tagged))
def clone(self):
return type(self)(self.queue, self.allocator)
# }}}
# }}}
# {{{ PytatoJAXArrayContext
class PytatoJAXArrayContext(_BasePytatoArrayContext):
"""
An arraycontext that uses :mod:`pytato` to represent the thawed state of
the arrays and compiles the expressions using
:class:`pytato.target.python.JAXPythonTarget`.
"""
def __init__(self,
*,
compile_trace_callback: Callable[[Any, str, Any], None] | None = None,
) -> None:
"""
:arg compile_trace_callback: A function of three arguments
*(what, stage, ir)*, where *what* identifies the object
being compiled, *stage* is a string describing the compilation
pass, and *ir* is an object containing the intermediate
representation. This interface should be considered
unstable.
"""
import jax.numpy as jnp
import pytato as pt
super().__init__(compile_trace_callback=compile_trace_callback)
self.array_types = (pt.Array, jnp.ndarray)
@property
def _frozen_array_types(self) -> tuple[type, ...]:
import jax.numpy as jnp
return (jnp.ndarray, )
def _rec_map_container(
self, func: Callable[[Array], Array], array: ArrayOrContainer,
allowed_types: tuple[type, ...] | None = None, *,
default_scalar: ScalarLike | None = None,
strict: bool = False) -> ArrayOrContainer:
if allowed_types is None:
allowed_types = self.array_types
def _wrapper(ary):
if isinstance(ary, allowed_types):
return func(ary)
elif np.isscalar(ary):
if default_scalar is None:
return ary
else:
return np.array(ary).dtype.type(default_scalar)
else:
raise TypeError(
f"{type(self).__name__}.{func.__name__[1:]} invoked with "
f"an unsupported array type: got '{type(ary).__name__}', "
f"but expected one of {allowed_types}")
return rec_map_array_container(_wrapper, array)
# {{{ ArrayContext interface
def from_numpy(self, array):
import jax
import pytato as pt
def _from_numpy(ary):
return pt.make_data_wrapper(jax.device_put(ary))
return with_array_context(
self._rec_map_container(_from_numpy, array, (np.ndarray,)),
actx=self)
def to_numpy(self, array):
import jax
def _to_numpy(ary):
return jax.device_get(ary)
return with_array_context(
self._rec_map_container(_to_numpy, self.freeze(array)),
actx=None)
def freeze(self, array):
if np.isscalar(array):
return array
import jax.numpy as jnp
import pytato as pt
from arraycontext.container.traversal import rec_keyed_map_array_container
from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
array_as_dict: dict[str, jnp.ndarray | pt.Array] = {}
key_to_frozen_subary: dict[str, jnp.ndarray] = {}
key_to_pt_arrays: dict[str, pt.Array] = {}
def _record_leaf_ary_in_dict(key: tuple[Any, ...],
ary: jnp.ndarray | pt.Array) -> None:
key_str = "_ary" + _ary_container_key_stringifier(key)
array_as_dict[key_str] = ary
rec_keyed_map_array_container(_record_leaf_ary_in_dict, array)
# {{{ remove any non pytato arrays from array_as_dict
for key, subary in array_as_dict.items():
if isinstance(subary, jnp.ndarray):
key_to_frozen_subary[key] = subary.block_until_ready()
elif isinstance(subary, pt.DataWrapper):
# trivial freeze.
key_to_frozen_subary[key] = subary.data.block_until_ready()
elif isinstance(subary, pt.Array):
key_to_pt_arrays[key] = subary
else:
raise TypeError(
f"{type(self).__name__}.freeze invoked with an unsupported "
f"array type: got '{type(subary).__name__}', but expected one "
f"of {self.array_types}")
# }}}
def _to_frozen(key: tuple[Any, ...], ary) -> jnp.ndarray:
key_str = "_ary" + _ary_container_key_stringifier(key)
return key_to_frozen_subary[key_str]
if not key_to_pt_arrays:
# all cl arrays => no need to perform any codegen
return with_array_context(
rec_keyed_map_array_container(_to_frozen, array),
actx=None)
pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(key_to_pt_arrays)
transformed_dag = self.transform_dag(pt_dict_of_named_arrays)
pt_prg = pt.generate_jax(transformed_dag, jit=True)
out_dict = pt_prg()
assert len(set(out_dict) & set(key_to_frozen_subary)) == 0
key_to_frozen_subary = {
**key_to_frozen_subary,
**{k: v.block_until_ready()
for k, v in out_dict.items()}
}
return with_array_context(
rec_keyed_map_array_container(_to_frozen, array),
actx=None)
def thaw(self, array):
import pytato as pt
def _thaw(ary):
return pt.make_data_wrapper(ary)
return with_array_context(
self._rec_map_container(_thaw, array, self._frozen_array_types),
actx=self)
def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
from .compile import LazilyJAXCompilingFunctionCaller
return LazilyJAXCompilingFunctionCaller(self, f)
def tag(self, tags: ToTagSetConvertible, array):
def _tag(ary):
import jax.numpy as jnp
if isinstance(ary, jnp.ndarray):
return ary
else:
return ary.tagged(_preprocess_array_tags(tags))
return self._rec_map_container(_tag, array)
def tag_axis(self, iaxis, tags: ToTagSetConvertible, array):
def _tag_axis(ary):
import jax.numpy as jnp
if isinstance(ary, jnp.ndarray):
return ary
else:
return ary.with_tagged_axis(iaxis, tags)
return self._rec_map_container(_tag_axis, array)
# }}}
# {{{ compilation
def call_loopy(self, program, **kwargs):
raise NotImplementedError(
"Calling loopy on JAX arrays is not supported. Maybe rewrite"
" the loopy kernel as numpy-flavored array operations using"
" ArrayContext.np.")
def einsum(self, spec, *args, arg_names=None, tagged=()):
import pytato as pt
if arg_names is None:
arg_names = (None,) * len(args)
def preprocess_arg(name, arg):
import jax.numpy as jnp
if isinstance(arg, jnp.ndarray):
ary = self.thaw(arg)
elif isinstance(arg, pt.Array):
ary = arg
else:
raise TypeError(
f"{type(self).__name__}.einsum invoked with an unsupported "
f"array type: got '{type(arg).__name__}', but expected one "
f"of {self.array_types}")
if name is not None: # noqa: SIM102
# Tagging Placeholders with naming-related tags is pointless:
# They already have names. It's also counterproductive, as
# multiple placeholders with the same name that are not
# also the same object are not allowed, and this would produce
# a different Placeholder object of the same name.
if (not isinstance(ary, pt.Placeholder)
and not ary.tags_of_type(NameHint)):
ary = ary.tagged(NameHint(name))
return ary
return pt.einsum(spec, *[
preprocess_arg(name, arg)
for name, arg in zip(arg_names, args, strict=True)
]).tagged(_preprocess_array_tags(tagged))
def clone(self):
return type(self)()
# }}}
# }}}
# vim: foldmethod=marker
arraycontext/impl/pytato/compile.py 0 → 100644
View file @ 1013ecb2
"""
.. autoclass:: BaseLazilyCompilingFunctionCaller
.. autoclass:: LazilyPyOpenCLCompilingFunctionCaller
.. autoclass:: LazilyJAXCompilingFunctionCaller
.. autoclass:: CompiledFunction
.. autoclass:: FromArrayContextCompile
"""
from __future__ import annotations
__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.
"""
import abc
import itertools
import logging
from collections.abc import Callable, Hashable, Mapping
from dataclasses import dataclass, field
from typing import Any
import numpy as np
from immutabledict import immutabledict
import pytato as pt
from pytools import ProcessLogger, to_identifier
from pytools.tag import Tag
from arraycontext.container import ArrayContainer, is_array_container_type
from arraycontext.container.traversal import rec_keyed_map_array_container
from arraycontext.context import ArrayT
from arraycontext.impl.pytato import (
PytatoJAXArrayContext,
PytatoPyOpenCLArrayContext,
_BasePytatoArrayContext,
)
logger = logging.getLogger(__name__)
def _prg_id_to_kernel_name(f: Any) -> str:
if callable(f):
name = getattr(f, "__name__", "anonymous")
if not name.isidentifier():
return "actx_compiled_" + to_identifier(name)
else:
return name
else:
return to_identifier(str(f))
class FromArrayContextCompile(Tag):
"""
Tagged to the entrypoint kernel of every translation unit that is generated
by :meth:`~arraycontext.PytatoPyOpenCLArrayContext.compile`.
Typically this tag serves as a branch condition in implementing a
specialized transform strategy for kernels compiled by
:meth:`~arraycontext.PytatoPyOpenCLArrayContext.compile`.
"""
# {{{ helper classes: AbstractInputDescriptor
class AbstractInputDescriptor:
"""
Used internally in :class:`BaseLazilyCompilingFunctionCaller` to characterize
an input.
"""
def __eq__(self, other):
raise NotImplementedError
def __hash__(self):
raise NotImplementedError
@dataclass(frozen=True, eq=True)
class ScalarInputDescriptor(AbstractInputDescriptor):
dtype: np.dtype
@dataclass(frozen=True, eq=True)
class LeafArrayDescriptor(AbstractInputDescriptor):
dtype: np.dtype
shape: pt.array.ShapeType
# }}}
# {{{ utilities
def _ary_container_key_stringifier(keys: tuple[Any, ...]) -> str:
"""
Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Stringifies an
array-container's component's key. Goals of this routine:
* No two different keys should have the same stringification
* Stringified key must a valid identifier according to :meth:`str.isidentifier`
* (informal) Shorter identifiers are preferred
"""
def _rec_str(key: Any) -> str:
if isinstance(key, str | int):
return str(key)
elif isinstance(key, tuple):
# t in '_actx_t': stands for tuple
return "_actx_t" + "_".join(_rec_str(k) for k in key) + "_actx_endt"
else:
raise NotImplementedError("Key-stringication unimplemented for "
f"'{type(key).__name__}'.")
return "_".join(_rec_str(key) for key in keys)
def _get_arg_id_to_arg_and_arg_id_to_descr(args: tuple[Any, ...],
kwargs: Mapping[str, Any]
) -> \
tuple[Mapping[tuple[Hashable, ...], Any],
Mapping[tuple[Hashable, ...], AbstractInputDescriptor]]:
"""
Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Extracts
mappings from argument id to argument values and from argument id to
:class:`AbstractInputDescriptor`. See
:attr:`CompiledFunction.input_id_to_name_in_program` for argument-id's
representation.
"""
arg_id_to_arg: dict[tuple[Hashable, ...], Any] = {}
arg_id_to_descr: dict[tuple[Hashable, ...], AbstractInputDescriptor] = {}
for kw, arg in itertools.chain(enumerate(args),
kwargs.items()):
if np.isscalar(arg):
arg_id = (kw,)
arg_id_to_arg[arg_id] = arg
arg_id_to_descr[arg_id] = ScalarInputDescriptor(np.dtype(type(arg)))
elif is_array_container_type(arg.__class__):
def id_collector(keys, ary):
arg_id = (kw, *keys) # noqa: B023
arg_id_to_arg[arg_id] = ary
arg_id_to_descr[arg_id] = LeafArrayDescriptor(
np.dtype(ary.dtype), ary.shape)
return ary
rec_keyed_map_array_container(id_collector, arg)
elif isinstance(arg, pt.Array):
arg_id = (kw,)
arg_id_to_arg[arg_id] = arg
arg_id_to_descr[arg_id] = LeafArrayDescriptor(np.dtype(arg.dtype),
arg.shape)
else:
raise ValueError("Argument to a compiled operator should be"
" either a scalar, pt.Array or an array container. Got"
f" '{arg}'.")
return immutabledict(arg_id_to_arg), immutabledict(arg_id_to_descr)
def _to_input_for_compiled(ary: ArrayT, actx: PytatoPyOpenCLArrayContext):
"""
Preprocess *ary* before turning it into a :class:`pytato.array.Placeholder`
in :meth:`LazilyCompilingFunctionCaller.__call__`.
Preprocessing here refers to:
- Metadata Inference that is supplied via *actx*\'s
:meth:`PytatoPyOpenCLArrayContext.transform_dag`.
"""
import pyopencl.array as cla
from arraycontext.impl.pyopencl.taggable_cl_array import (
TaggableCLArray,
to_tagged_cl_array,
)
if isinstance(ary, pt.Array):
dag = pt.make_dict_of_named_arrays({"_actx_out": ary})
# Transform the DAG to give metadata inference a chance to do its job
return actx.transform_dag(dag)["_actx_out"].expr
elif isinstance(ary, TaggableCLArray):
return ary
elif isinstance(ary, cla.Array):
from warnings import warn
warn("Passing pyopencl.array.Array to a compiled callable"
" is deprecated and will stop working in 2023."
" Use `to_tagged_cl_array` to convert the array to"
" TaggableCLArray", DeprecationWarning, stacklevel=2)
return to_tagged_cl_array(ary,
axes=None,
tags=frozenset())
else:
raise NotImplementedError(type(ary))
def _get_f_placeholder_args(arg, kw, arg_id_to_name, actx):
"""
Helper for :class:`BaseLazilyCompilingFunctionCaller.__call__`. Returns the
placeholder version of an argument to
:attr:`BaseLazilyCompilingFunctionCaller.f`.
"""
if np.isscalar(arg):
from pytato.tags import ForceValueArgTag
name = arg_id_to_name[kw,]
return pt.make_placeholder(name, (), np.dtype(type(arg)),
tags=frozenset({ForceValueArgTag()}))
elif isinstance(arg, pt.Array):
name = arg_id_to_name[kw,]
# Transform the DAG to give metadata inference a chance to do its job
arg = _to_input_for_compiled(arg, actx)
return pt.make_placeholder(name, arg.shape, arg.dtype,
axes=arg.axes,
tags=arg.tags)
elif is_array_container_type(arg.__class__):
def _rec_to_placeholder(keys, ary):
index = (kw, *keys)
name = arg_id_to_name[index]
# Transform the DAG to give metadata inference a chance to do its job
ary = _to_input_for_compiled(ary, actx)
return pt.make_placeholder(name,
ary.shape,
ary.dtype,
axes=ary.axes,
tags=ary.tags)
return rec_keyed_map_array_container(_rec_to_placeholder, arg)
else:
raise NotImplementedError(type(arg))
# }}}
# {{{ BaseLazilyCompilingFunctionCaller
@dataclass
class BaseLazilyCompilingFunctionCaller:
"""
Records a side-effect-free callable :attr:`f` that can be specialized for
the input types with which :meth:`__call__` is invoked.
.. attribute:: f
The callable that will be called to obtain :mod:`pytato` DAGs.
.. automethod:: __call__
"""
actx: _BasePytatoArrayContext
f: Callable[..., Any]
program_cache: dict[Mapping[tuple[Hashable, ...], AbstractInputDescriptor],
CompiledFunction] = field(default_factory=lambda: {})
# {{{ abstract interface
def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
raise NotImplementedError
@property
def compiled_function_returning_array_container_class(
self) -> type[CompiledFunction]:
raise NotImplementedError
@property
def compiled_function_returning_array_class(self) -> type[CompiledFunction]:
raise NotImplementedError
# }}}
def _dag_to_compiled_func(self, ary_or_dict_of_named_arrays,
input_id_to_name_in_program, output_id_to_name_in_program,
output_template):
if isinstance(ary_or_dict_of_named_arrays, pt.Array):
output_id = "_pt_out"
dict_of_named_arrays = pt.make_dict_of_named_arrays(
{output_id: ary_or_dict_of_named_arrays})
pytato_program, name_in_program_to_tags, name_in_program_to_axes = (
self._dag_to_transformed_pytato_prg(dict_of_named_arrays,
prg_id=self.f))
return self.compiled_function_returning_array_class(
self.actx, pytato_program,
input_id_to_name_in_program=input_id_to_name_in_program,
output_tags=name_in_program_to_tags[output_id],
output_axes=name_in_program_to_axes[output_id],
output_name=output_id)
elif isinstance(ary_or_dict_of_named_arrays, pt.DictOfNamedArrays):
pytato_program, name_in_program_to_tags, name_in_program_to_axes = (
self._dag_to_transformed_pytato_prg(ary_or_dict_of_named_arrays,
prg_id=self.f))
return self.compiled_function_returning_array_container_class(
self.actx, pytato_program,
input_id_to_name_in_program=input_id_to_name_in_program,
output_id_to_name_in_program=output_id_to_name_in_program,
name_in_program_to_tags=name_in_program_to_tags,
name_in_program_to_axes=name_in_program_to_axes,
output_template=output_template)
else:
raise NotImplementedError(type(ary_or_dict_of_named_arrays))
def __call__(self, *args: Any, **kwargs: Any) -> Any:
"""
Returns the result of :attr:`~BaseLazilyCompilingFunctionCaller.f`'s
function application on *args*.
Before applying :attr:`~BaseLazilyCompilingFunctionCaller.f`, it is compiled
to a :mod:`pytato` DAG that would apply
:attr:`~BaseLazilyCompilingFunctionCaller.f` with *args* in a lazy-sense.
The intermediary pytato DAG for *args* is memoized in *self*.
"""
arg_id_to_arg, arg_id_to_descr = _get_arg_id_to_arg_and_arg_id_to_descr(
args, kwargs)
try:
compiled_f = self.program_cache[arg_id_to_descr]
except KeyError:
pass
else:
return compiled_f(arg_id_to_arg)
dict_of_named_arrays = {}
output_id_to_name_in_program = {}
input_id_to_name_in_program = {
arg_id: f"_actx_in_{_ary_container_key_stringifier(arg_id)}"
for arg_id in arg_id_to_arg}
output_template = self.f(
*[_get_f_placeholder_args(arg, iarg,
input_id_to_name_in_program, self.actx)
for iarg, arg in enumerate(args)],
**{kw: _get_f_placeholder_args(arg, kw,
input_id_to_name_in_program,
self.actx)
for kw, arg in kwargs.items()})
self.actx._compile_trace_callback(self.f, "post_trace", output_template)
if (not (is_array_container_type(output_template.__class__)
or isinstance(output_template, pt.Array))):
# TODO: We could possibly just short-circuit this interface if the
# returned type is a scalar. Not sure if it's worth it though.
raise NotImplementedError(
f"Function '{self.f.__name__}' to be compiled "
"did not return an array container or pt.Array,"
f" but an instance of '{output_template.__class__}' instead.")
def _as_dict_of_named_arrays(keys, ary):
name = "_pt_out_" + _ary_container_key_stringifier(keys)
output_id_to_name_in_program[keys] = name
dict_of_named_arrays[name] = ary
return ary
rec_keyed_map_array_container(_as_dict_of_named_arrays,
output_template)
compiled_func = self._dag_to_compiled_func(
pt.make_dict_of_named_arrays(dict_of_named_arrays),
input_id_to_name_in_program=input_id_to_name_in_program,
output_id_to_name_in_program=output_id_to_name_in_program,
output_template=output_template)
self.program_cache[arg_id_to_descr] = compiled_func
return compiled_func(arg_id_to_arg)
# }}}
# {{{ LazilyPyOpenCLCompilingFunctionCaller
class LazilyPyOpenCLCompilingFunctionCaller(BaseLazilyCompilingFunctionCaller):
actx: PytatoPyOpenCLArrayContext
@property
def compiled_function_returning_array_container_class(
self) -> type[CompiledFunction]:
return CompiledPyOpenCLFunctionReturningArrayContainer
@property
def compiled_function_returning_array_class(self) -> type[CompiledFunction]:
return CompiledPyOpenCLFunctionReturningArray
def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
if prg_id is None:
prg_id = self.f
from pytato.target.loopy import BoundPyOpenCLExecutable
self.actx._compile_trace_callback(
prg_id, "pre_transform_dag", dict_of_named_arrays)
with ProcessLogger(logger, f"transform_dag for '{prg_id}'"):
pt_dict_of_named_arrays = self.actx.transform_dag(dict_of_named_arrays)
self.actx._compile_trace_callback(
prg_id, "post_transform_dag", pt_dict_of_named_arrays)
name_in_program_to_tags = {
name: out.tags
for name, out in pt_dict_of_named_arrays._data.items()}
name_in_program_to_axes = {
name: out.axes
for name, out in pt_dict_of_named_arrays._data.items()}
self.actx._compile_trace_callback(
prg_id, "pre_generate_loopy", pt_dict_of_named_arrays)
with ProcessLogger(logger, f"generate_loopy for '{prg_id}'"):
from arraycontext.loopy import _DEFAULT_LOOPY_OPTIONS
opts = _DEFAULT_LOOPY_OPTIONS
assert opts.return_dict
pytato_program = pt.generate_loopy(
pt_dict_of_named_arrays,
options=opts,
function_name=_prg_id_to_kernel_name(prg_id),
target=self.actx.get_target(),
).bind_to_context(self.actx.context) # pylint: disable=no-member
assert isinstance(pytato_program, BoundPyOpenCLExecutable)
self.actx._compile_trace_callback(
prg_id, "post_generate_loopy", pytato_program)
self.actx._compile_trace_callback(
prg_id, "pre_transform_loopy_program", pytato_program)
with ProcessLogger(logger, f"transform_loopy_program for '{prg_id}'"):
pytato_program = (pytato_program
.with_transformed_translation_unit(
lambda x: x.with_kernel(
x.default_entrypoint
.tagged(FromArrayContextCompile()))))
pytato_program = (pytato_program
.with_transformed_translation_unit(
self.actx.transform_loopy_program))
self.actx._compile_trace_callback(
prg_id, "post_transform_loopy_program", pytato_program)
self.actx._compile_trace_callback(
prg_id, "final", pytato_program)
return pytato_program, name_in_program_to_tags, name_in_program_to_axes
# }}}
# {{{ preserve back compat
class LazilyCompilingFunctionCaller(LazilyPyOpenCLCompilingFunctionCaller):
def __new__(cls, *args, **kwargs):
from warnings import warn
warn("LazilyCompilingFunctionCaller has been renamed to"
" LazilyPyOpenCLCompilingFunctionCaller. This will be"
" an error in 2023.", DeprecationWarning, stacklevel=2)
return super().__new__(cls)
def _dag_to_transformed_loopy_prg(self, dict_of_named_arrays):
from warnings import warn
warn("_dag_to_transformed_loopy_prg has been renamed to"
" _dag_to_transformed_pytato_prg. This will be"
" an error in 2023.", DeprecationWarning, stacklevel=2)
return super()._dag_to_transformed_pytato_prg(dict_of_named_arrays)
# }}}
# {{{ LazilyJAXCompilingFunctionCaller
class LazilyJAXCompilingFunctionCaller(BaseLazilyCompilingFunctionCaller):
@property
def compiled_function_returning_array_container_class(
self) -> type[CompiledFunction]:
return CompiledJAXFunctionReturningArrayContainer
@property
def compiled_function_returning_array_class(self) -> type[CompiledFunction]:
return CompiledJAXFunctionReturningArray
def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
if prg_id is None:
prg_id = self.f
self.actx._compile_trace_callback(
prg_id, "pre_transform_dag", dict_of_named_arrays)
with ProcessLogger(logger, f"transform_dag for '{prg_id}'"):
pt_dict_of_named_arrays = self.actx.transform_dag(dict_of_named_arrays)
self.actx._compile_trace_callback(
prg_id, "post_transform_dag", pt_dict_of_named_arrays)
name_in_program_to_tags = {
name: out.tags
for name, out in pt_dict_of_named_arrays._data.items()}
name_in_program_to_axes = {
name: out.axes
for name, out in pt_dict_of_named_arrays._data.items()}
self.actx._compile_trace_callback(
prg_id, "pre_generate_jax", pt_dict_of_named_arrays)
with ProcessLogger(logger, f"generate_jax for '{prg_id}'"):
pytato_program = pt.generate_jax(
pt_dict_of_named_arrays,
jit=True,
function_name=_prg_id_to_kernel_name(prg_id))
self.actx._compile_trace_callback(
prg_id, "post_generate_jax", pytato_program)
return pytato_program, name_in_program_to_tags, name_in_program_to_axes
def _args_to_device_buffers(actx, input_id_to_name_in_program, arg_id_to_arg,
fn_name="<unknown>"):
input_kwargs_for_loopy = {}
for arg_id, arg in arg_id_to_arg.items():
if np.isscalar(arg):
if isinstance(actx, PytatoPyOpenCLArrayContext):
# Scalar kernel args are passed as lp.ValueArgs
pass
elif isinstance(actx, PytatoJAXArrayContext):
import jax
arg = jax.device_put(arg)
else:
raise NotImplementedError(type(actx))
elif isinstance(arg, pt.array.DataWrapper):
# got a Datawrapper => simply gets its data
arg = arg.data
elif isinstance(arg, actx._frozen_array_types):
# got a frozen array => do nothing
pass
elif isinstance(arg, pt.Array):
# got an array expression => abort
raise ValueError(
f"Argument '{arg_id}' to the '{fn_name}' compiled function is a"
" pytato array expression. Evaluating it just-in-time"
" potentially causes a significant overhead on each call to the"
" function and is therefore unsupported. "
)
else:
raise NotImplementedError(type(arg))
input_kwargs_for_loopy[input_id_to_name_in_program[arg_id]] = arg
return input_kwargs_for_loopy
# }}}
# {{{ compiled function
class CompiledFunction(abc.ABC):
"""
A callable which captures the :class:`pytato.target.BoundProgram` resulting
from calling :attr:`~BaseLazilyCompilingFunctionCaller.f` with a given set of
input types, and generating :mod:`loopy` IR from it.
.. attribute:: pytato_program
.. attribute:: input_id_to_name_in_program
A mapping from input id to the placeholder name in
:attr:`CompiledFunction.pytato_program`. Input id is represented as the
position of :attr:`~BaseLazilyCompilingFunctionCaller.f`'s argument augmented
with the leaf array's key if the argument is an array container.
.. automethod:: __call__
"""
@abc.abstractmethod
def __call__(self, arg_id_to_arg) -> Any:
"""
:arg arg_id_to_arg: Mapping from input id to the passed argument. See
:attr:`CompiledFunction.input_id_to_name_in_program` for input id's
representation.
"""
pass
# }}}
# {{{ compiled pyopencl function
@dataclass(frozen=True)
class CompiledPyOpenCLFunctionReturningArrayContainer(CompiledFunction):
"""
.. attribute:: output_id_to_name_in_program
A mapping from output id to the name of
:class:`pytato.array.NamedArray` in
:attr:`CompiledFunction.pytato_program`. Output id is represented by
the key of a leaf array in the array container
:attr:`CompiledFunction.output_template`.
.. attribute:: output_template
An instance of :class:`arraycontext.ArrayContainer` that is the return
type of the callable.
"""
actx: PytatoPyOpenCLArrayContext
pytato_program: pt.target.loopy.BoundPyOpenCLExecutable
input_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
output_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
name_in_program_to_tags: Mapping[str, frozenset[Tag]]
name_in_program_to_axes: Mapping[str, tuple[pt.Axis, ...]]
output_template: ArrayContainer
def __call__(self, arg_id_to_arg) -> ArrayContainer:
from .utils import get_cl_axes_from_pt_axes
from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
fn_name = self.pytato_program.program.entrypoint
input_kwargs_for_loopy = _args_to_device_buffers(
self.actx, self.input_id_to_name_in_program, arg_id_to_arg, fn_name)
if self.actx.profile_kernels:
import pyopencl as cl
start_evt = cl.enqueue_marker(self.actx.queue)
evt, out_dict = self.pytato_program(queue=self.actx.queue,
allocator=self.actx.allocator,
**input_kwargs_for_loopy)
if self.actx.profile_kernels:
self.actx._add_profiling_events(start_evt, evt, fn_name)
def to_output_template(keys, _):
name_in_program = self.output_id_to_name_in_program[keys]
return self.actx.thaw(to_tagged_cl_array(
out_dict[name_in_program],
axes=get_cl_axes_from_pt_axes(
self.name_in_program_to_axes[name_in_program]),
tags=self.name_in_program_to_tags[name_in_program]))
return rec_keyed_map_array_container(to_output_template,
self.output_template)
@dataclass(frozen=True)
class CompiledPyOpenCLFunctionReturningArray(CompiledFunction):
"""
.. attribute:: output_name_in_program
Name of the output array in the program.
"""
actx: PytatoPyOpenCLArrayContext
pytato_program: pt.target.loopy.BoundPyOpenCLExecutable
input_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
output_tags: frozenset[Tag]
output_axes: tuple[pt.Axis, ...]
output_name: str
def __call__(self, arg_id_to_arg) -> ArrayContainer:
from .utils import get_cl_axes_from_pt_axes
from arraycontext.impl.pyopencl.taggable_cl_array import to_tagged_cl_array
fn_name = self.pytato_program.program.entrypoint
input_kwargs_for_loopy = _args_to_device_buffers(
self.actx, self.input_id_to_name_in_program, arg_id_to_arg, fn_name)
if self.actx.profile_kernels:
import pyopencl as cl
start_evt = cl.enqueue_marker(self.actx.queue)
evt, out_dict = self.pytato_program(queue=self.actx.queue,
allocator=self.actx.allocator,
**input_kwargs_for_loopy)
if self.actx.profile_kernels:
self.actx._add_profiling_events(start_evt, evt, fn_name)
return self.actx.thaw(to_tagged_cl_array(out_dict[self.output_name],
axes=get_cl_axes_from_pt_axes(
self.output_axes),
tags=self.output_tags))
# }}}
# {{{ compiled jax function
@dataclass(frozen=True)
class CompiledJAXFunctionReturningArrayContainer(CompiledFunction):
"""
.. attribute:: output_id_to_name_in_program
A mapping from output id to the name of
:class:`pytato.array.NamedArray` in
:attr:`CompiledFunction.pytato_program`. Output id is represented by
the key of a leaf array in the array container
:attr:`CompiledFunction.output_template`.
.. attribute:: output_template
An instance of :class:`arraycontext.ArrayContainer` that is the return
type of the callable.
"""
actx: PytatoJAXArrayContext
pytato_program: pt.target.python.BoundJAXPythonProgram
input_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
output_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
name_in_program_to_tags: Mapping[str, frozenset[Tag]]
name_in_program_to_axes: Mapping[str, tuple[pt.Axis, ...]]
output_template: ArrayContainer
def __call__(self, arg_id_to_arg) -> ArrayContainer:
fn_name = self.pytato_program.entrypoint
input_kwargs_for_loopy = _args_to_device_buffers(
self.actx, self.input_id_to_name_in_program, arg_id_to_arg, fn_name)
out_dict = self.pytato_program(**input_kwargs_for_loopy)
def to_output_template(keys, _):
return self.actx.thaw(
out_dict[self.output_id_to_name_in_program[keys]]
.block_until_ready()
)
return rec_keyed_map_array_container(to_output_template,
self.output_template)
@dataclass(frozen=True)
class CompiledJAXFunctionReturningArray(CompiledFunction):
"""
.. attribute:: output_name_in_program
Name of the output array in the program.
"""
actx: PytatoJAXArrayContext
pytato_program: pt.target.python.BoundJAXPythonProgram
input_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
output_tags: frozenset[Tag]
output_axes: tuple[pt.Axis, ...]
output_name: str
def __call__(self, arg_id_to_arg) -> ArrayContainer:
fn_name = self.pytato_program.entrypoint
input_kwargs_for_loopy = _args_to_device_buffers(
self.actx, self.input_id_to_name_in_program, arg_id_to_arg, fn_name)
_evt, out_dict = self.pytato_program(**input_kwargs_for_loopy)
return self.actx.thaw(out_dict[self.output_name])
# }}}
# vim: foldmethod=marker
arraycontext/impl/pytato/fake_numpy.py 0 → 100644
View file @ 1013ecb2
from __future__ import annotations
__copyright__ = """
Copyright (C) 2021 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, reduce
from typing import Any, cast
import numpy as np
import pytato as pt
from arraycontext.container import NotAnArrayContainerError, serialize_container
from arraycontext.container.traversal import (
rec_map_array_container,
rec_map_reduce_array_container,
rec_multimap_array_container,
)
from arraycontext.context import Array, ArrayOrContainer
from arraycontext.fake_numpy import BaseFakeNumpyLinalgNamespace
from arraycontext.loopy import LoopyBasedFakeNumpyNamespace
class PytatoFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace):
# Everything is implemented in the base class for now.
pass
class PytatoFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace):
"""
A :mod:`numpy` mimic for :class:`PytatoPyOpenCLArrayContext`.
.. note::
:mod:`pytato` does not define any memory layout for its arrays. See
:ref:`Pytato docs <pytato:memory-layout>` for more on this.
"""
_pt_unary_funcs = frozenset({
"sin", "cos", "tan", "arcsin", "arccos", "arctan",
"sinh", "cosh", "tanh", "exp", "log", "log10",
"sqrt", "abs", "isnan", "real", "imag", "conj",
"logical_not",
})
_pt_multi_ary_funcs = frozenset({
"arctan2", "equal", "greater", "greater_equal", "less", "less_equal",
"not_equal", "minimum", "maximum", "where", "logical_and", "logical_or",
})
def _get_fake_numpy_linalg_namespace(self):
return PytatoFakeNumpyLinalgNamespace(self._array_context)
def __getattr__(self, name):
if name in self._pt_unary_funcs:
from functools import partial
return partial(rec_map_array_container, getattr(pt, name))
if name in self._pt_multi_ary_funcs:
from functools import partial
return partial(rec_multimap_array_container, getattr(pt, name))
return super().__getattr__(name)
# NOTE: the order of these follows the order in numpy docs
# NOTE: when adding a function here, also add it to `array_context.rst` docs!
# {{{ array creation routines
def zeros(self, shape, dtype):
return pt.zeros(shape, dtype)
def zeros_like(self, ary):
def _zeros_like(array):
return self._array_context.np.zeros(
array.shape, array.dtype).copy(axes=array.axes, tags=array.tags)
return self._array_context._rec_map_container(
_zeros_like, ary, default_scalar=0)
def ones_like(self, ary):
return self.full_like(ary, 1)
def full_like(self, ary, fill_value):
def _full_like(subary):
return pt.full(subary.shape, fill_value, subary.dtype).copy(
axes=subary.axes, tags=subary.tags)
return self._array_context._rec_map_container(
_full_like, ary, default_scalar=fill_value)
def arange(self, *args: Any, **kwargs: Any):
return pt.arange(*args, **kwargs)
def full(self, shape, fill_value, dtype=None):
return pt.full(shape, fill_value, dtype)
# }}}
# {{{ array manipulation routines
def reshape(self, a, newshape, order="C"):
return rec_map_array_container(
lambda ary: pt.reshape(a, newshape, order=order),
a)
def ravel(self, a, order="C"):
"""
:arg order: A :class:`str` describing the order in which the elements
must be traversed while flattening. Can be one of 'F', 'C', 'A' or
'K'. Since, :mod:`pytato` arrays don't have a memory layout, if
*order* is 'A' or 'K', the traversal order while flattening is
undefined.
"""
def _rec_ravel(a):
if order in "FC":
return pt.reshape(a, (-1,), order=order)
elif order in "AK":
# flattening in a C-order
# memory layout is assumed to be "C"
return pt.reshape(a, (-1,), order="C")
else:
raise ValueError("`order` can be one of 'F', 'C', 'A' or 'K'. "
f"(got {order})")
return rec_map_array_container(_rec_ravel, a)
def transpose(self, a, axes=None):
return rec_multimap_array_container(pt.transpose, a, axes)
def broadcast_to(self, array, shape):
return rec_map_array_container(partial(pt.broadcast_to, shape=shape), array)
def concatenate(self, arrays, axis=0):
return rec_multimap_array_container(pt.concatenate, arrays, axis)
def stack(self, arrays, axis=0):
return rec_multimap_array_container(
lambda *args: pt.stack(arrays=args, axis=axis),
*arrays)
# }}}
# {{{ logic functions
def all(self, a):
return rec_map_reduce_array_container(
partial(reduce, pt.logical_and),
lambda subary: pt.all(subary), a)
def any(self, a):
return rec_map_reduce_array_container(
partial(reduce, pt.logical_or),
lambda subary: pt.any(subary), a)
def array_equal(self, a: ArrayOrContainer, b: ArrayOrContainer) -> Array:
actx = self._array_context
# NOTE: not all backends support `bool` properly, so use `int8` instead
true_ary = actx.from_numpy(np.int8(True))
false_ary = actx.from_numpy(np.int8(False))
def rec_equal(x: ArrayOrContainer, y: ArrayOrContainer) -> pt.Array:
if type(x) is not type(y):
return false_ary
try:
serialized_x = serialize_container(x)
serialized_y = serialize_container(y)
except NotAnArrayContainerError:
assert isinstance(x, pt.Array)
assert isinstance(y, pt.Array)
if x.shape != y.shape:
return false_ary
else:
return pt.all(cast(pt.Array, pt.equal(x, y)))
else:
if len(serialized_x) != len(serialized_y):
return false_ary
return reduce(
pt.logical_and,
[(true_ary if kx_i == ky_i else false_ary)
and rec_equal(x_i, y_i)
for (kx_i, x_i), (ky_i, y_i)
in zip(serialized_x, serialized_y, strict=True)],
true_ary)
return cast(Array, rec_equal(a, b))
# }}}
# {{{ mathematical functions
def sum(self, a, axis=None, dtype=None):
def _pt_sum(ary):
if dtype not in [ary.dtype, None]:
raise NotImplementedError
return pt.sum(ary, axis=axis)
return rec_map_reduce_array_container(sum, _pt_sum, a)
def amax(self, a, axis=None):
return rec_map_reduce_array_container(
partial(reduce, pt.maximum), partial(pt.amax, axis=axis), a)
max = amax
def amin(self, a, axis=None):
return rec_map_reduce_array_container(
partial(reduce, pt.minimum), partial(pt.amin, axis=axis), a)
min = amin
def absolute(self, a):
return self.abs(a)
def vdot(self, a: Array, b: Array):
return rec_multimap_array_container(pt.vdot, a, b)
# }}}
arraycontext/impl/pytato/utils.py 0 → 100644
View file @ 1013ecb2
from __future__ import annotations
__doc__ = """
.. autofunction:: transfer_from_numpy
.. autofunction:: transfer_to_numpy
Profiling-related functions
^^^^^^^^^^^^^^^^^^^^^^^^^^^
.. autofunction:: tabulate_profiling_data
"""
__copyright__ = """
Copyright (C) 2021 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 collections.abc import Mapping
from typing import TYPE_CHECKING, Any, cast
import pytools
from pytato.array import (
AbstractResultWithNamedArrays,
Array,
Axis as PtAxis,
DataWrapper,
DictOfNamedArrays,
Placeholder,
SizeParam,
make_placeholder,
)
from pytato.target.loopy import LoopyPyOpenCLTarget
from pytato.transform import ArrayOrNames, CopyMapper
from pytools import UniqueNameGenerator, memoize_method
from arraycontext import ArrayContext
from arraycontext.impl.pyopencl.taggable_cl_array import Axis as ClAxis
from arraycontext.impl.pytato import PytatoPyOpenCLArrayContext
if TYPE_CHECKING:
import loopy as lp
class _DatawrapperToBoundPlaceholderMapper(CopyMapper):
"""
Helper mapper for :func:`normalize_pt_expr`. Every
:class:`pytato.DataWrapper` is replaced with a deterministic copy of
:class:`Placeholder`.
"""
def __init__(self) -> None:
super().__init__()
self.bound_arguments: dict[str, Any] = {}
self.vng = UniqueNameGenerator()
self.seen_inputs: set[str] = set()
def map_data_wrapper(self, expr: DataWrapper) -> Array:
if expr.name is not None:
if expr.name in self.seen_inputs:
raise ValueError("Got multiple inputs with the name"
f"{expr.name} => Illegal.")
self.seen_inputs.add(expr.name)
# Normalizing names so that more arrays can have the same normalized DAG.
from pytato.codegen import _generate_name_for_temp
name = _generate_name_for_temp(expr, self.vng, "_actx_dw")
self.bound_arguments[name] = expr.data
return make_placeholder(
name=name,
shape=tuple(cast(Array, self.rec(s)) if isinstance(s, Array) else s
for s in expr.shape),
dtype=expr.dtype,
axes=expr.axes,
tags=expr.tags)
def map_size_param(self, expr: SizeParam) -> Array:
raise NotImplementedError
def map_placeholder(self, expr: Placeholder) -> Array:
raise ValueError("Placeholders cannot appear in"
" DatawrapperToBoundPlaceholderMapper.")
def _normalize_pt_expr(
expr: DictOfNamedArrays
) -> tuple[Array | AbstractResultWithNamedArrays, Mapping[str, Any]]:
"""
Returns ``(normalized_expr, bound_arguments)``. *normalized_expr* is a
normalized form of *expr*, with all instances of
:class:`pytato.DataWrapper` replaced with instances of :class:`Placeholder`
named in a deterministic manner. The data corresponding to the placeholders
in *normalized_expr* is recorded in the mapping *bound_arguments*.
Deterministic naming of placeholders permits more effective caching of
equivalent graphs.
"""
normalize_mapper = _DatawrapperToBoundPlaceholderMapper()
normalized_expr = normalize_mapper(expr)
assert isinstance(normalized_expr, AbstractResultWithNamedArrays)
return normalized_expr, normalize_mapper.bound_arguments
def get_pt_axes_from_cl_axes(axes: tuple[ClAxis, ...]) -> tuple[PtAxis, ...]:
return tuple(PtAxis(axis.tags) for axis in axes)
def get_cl_axes_from_pt_axes(axes: tuple[PtAxis, ...]) -> tuple[ClAxis, ...]:
return tuple(ClAxis(axis.tags) for axis in axes)
# {{{ arg-size-limiting loopy target
class ArgSizeLimitingPytatoLoopyPyOpenCLTarget(LoopyPyOpenCLTarget):
def __init__(self, limit_arg_size_nbytes: int) -> None:
super().__init__()
self.limit_arg_size_nbytes = limit_arg_size_nbytes
@memoize_method
def get_loopy_target(self) -> lp.PyOpenCLTarget:
from loopy import PyOpenCLTarget
return PyOpenCLTarget(limit_arg_size_nbytes=self.limit_arg_size_nbytes)
# }}}
# {{{ Transfer mappers
class TransferFromNumpyMapper(CopyMapper):
"""A mapper to transfer arrays contained in :class:`~pytato.array.DataWrapper`
instances to be device arrays, using
:meth:`~arraycontext.ArrayContext.from_numpy`.
"""
def __init__(self, actx: ArrayContext) -> None:
super().__init__()
self.actx = actx
def map_data_wrapper(self, expr: DataWrapper) -> Array:
import numpy as np
if not isinstance(expr.data, np.ndarray):
raise ValueError("TransferFromNumpyMapper: tried to transfer data that "
"is already on the device")
# Ideally, this code should just do
# return self.actx.from_numpy(expr.data).tagged(expr.tags),
# but there seems to be no way to transfer the non_equality_tags in that case.
actx_ary = self.actx.from_numpy(expr.data)
assert isinstance(actx_ary, DataWrapper)
# https://github.com/pylint-dev/pylint/issues/3893
# pylint: disable=unexpected-keyword-arg
return DataWrapper(
data=actx_ary.data,
shape=expr.shape,
axes=expr.axes,
tags=expr.tags,
non_equality_tags=expr.non_equality_tags)
class TransferToNumpyMapper(CopyMapper):
"""A mapper to transfer arrays contained in :class:`~pytato.array.DataWrapper`
instances to be :class:`numpy.ndarray` instances, using
:meth:`~arraycontext.ArrayContext.to_numpy`.
"""
def __init__(self, actx: ArrayContext) -> None:
super().__init__()
self.actx = actx
def map_data_wrapper(self, expr: DataWrapper) -> Array:
import numpy as np
import arraycontext.impl.pyopencl.taggable_cl_array as tga
if not isinstance(expr.data, tga.TaggableCLArray):
raise ValueError("TransferToNumpyMapper: tried to transfer data that "
"is already on the host")
np_data = self.actx.to_numpy(expr.data)
assert isinstance(np_data, np.ndarray)
# https://github.com/pylint-dev/pylint/issues/3893
# pylint: disable=unexpected-keyword-arg
# type-ignore: discussed at
# https://github.com/inducer/arraycontext/pull/289#discussion_r1855523967
# possibly related: https://github.com/python/mypy/issues/17375
return DataWrapper( # type: ignore[call-arg]
data=np_data,
shape=expr.shape,
axes=expr.axes,
tags=expr.tags,
non_equality_tags=expr.non_equality_tags)
def transfer_from_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
"""Transfer arrays contained in :class:`~pytato.array.DataWrapper`
instances to be device arrays, using
:meth:`~arraycontext.ArrayContext.from_numpy`.
"""
return TransferFromNumpyMapper(actx)(expr)
def transfer_to_numpy(expr: ArrayOrNames, actx: ArrayContext) -> ArrayOrNames:
"""Transfer arrays contained in :class:`~pytato.array.DataWrapper`
instances to be :class:`numpy.ndarray` instances, using
:meth:`~arraycontext.ArrayContext.to_numpy`.
"""
return TransferToNumpyMapper(actx)(expr)
# }}}
# {{{ Profiling
def tabulate_profiling_data(actx: PytatoPyOpenCLArrayContext) -> pytools.Table:
"""Return a :class:`pytools.Table` with the profiling results."""
actx._wait_and_transfer_profile_events()
tbl = pytools.Table()
# Table header
tbl.add_row(("Kernel", "# Calls", "Time_sum [ns]", "Time_avg [ns]"))
# Precision of results
g = ".5g"
total_calls = 0
total_time = 0.0
for kernel_name, times in actx._profile_results.items():
num_calls = len(times)
total_calls += num_calls
t_sum = sum(times)
t_avg = t_sum / num_calls
if t_sum is not None:
total_time += t_sum
tbl.add_row((kernel_name, num_calls, f"{t_sum:{g}}", f"{t_avg:{g}}"))
tbl.add_row(("", "", "", ""))
tbl.add_row(("Total", total_calls, f"{total_time:{g}}", "--"))
actx._reset_profiling_data()
return tbl
# }}}
# vim: foldmethod=marker
arraycontext/loopy.py
View file @ 1013ecb2
......@@ -2,6 +2,8 @@
.. currentmodule:: arraycontext
.. autofunction:: make_loopy_program
"""
from __future__ import annotations
__copyright__ = """
Copyright (C) 2020-1 University of Illinois Board of Trustees
......@@ -27,8 +29,17 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
from collections.abc import Mapping
from typing import ClassVar
import numpy as np
import loopy as lp
from loopy.version import MOST_RECENT_LANGUAGE_VERSION
from pytools import memoize_in
from arraycontext.container.traversal import multimapped_over_array_containers
from arraycontext.fake_numpy import BaseFakeNumpyNamespace
# {{{ loopy
......@@ -39,7 +50,7 @@ _DEFAULT_LOOPY_OPTIONS = lp.Options(
def make_loopy_program(domains, statements, kernel_data=None,
name="mm_actx_kernel"):
name="mm_actx_kernel", tags=None):
"""Return a :class:`loopy.LoopKernel` suitable for use with
:meth:`ArrayContext.call_loopy`.
"""
......@@ -53,7 +64,8 @@ def make_loopy_program(domains, statements, kernel_data=None,
options=_DEFAULT_LOOPY_OPTIONS,
default_offset=lp.auto,
name=name,
lang_version=MOST_RECENT_LANGUAGE_VERSION)
lang_version=MOST_RECENT_LANGUAGE_VERSION,
tags=tags)
def get_default_entrypoint(t_unit):
......@@ -63,9 +75,96 @@ def get_default_entrypoint(t_unit):
except AttributeError:
try:
return t_unit.root_kernel
except AttributeError:
except AttributeError as err:
raise TypeError("unable to find default entry point for loopy "
"translation unit")
"translation unit") from err
def _get_scalar_func_loopy_program(actx, c_name, nargs, naxes):
@memoize_in(actx, _get_scalar_func_loopy_program)
def get(c_name, nargs, naxes):
from pymbolic.primitives import Subscript, Variable
var_names = [f"i{i}" for i in range(naxes)]
size_names = [f"n{i}" for i in range(naxes)]
subscript = tuple(Variable(vname) for vname in var_names)
from islpy import make_zero_and_vars
v = make_zero_and_vars(var_names, params=size_names)
domain = v[0].domain()
for vname, sname in zip(var_names, size_names, strict=True):
domain = domain & v[0].le_set(v[vname]) & v[vname].lt_set(v[sname])
domain_bset, = domain.get_basic_sets()
import loopy as lp
from arraycontext.transform_metadata import ElementwiseMapKernelTag
def sub(name: str) -> Variable | Subscript:
return Subscript(Variable(name), subscript) if subscript else Variable(name)
return make_loopy_program(
[domain_bset], [
lp.Assignment(
sub("out"),
Variable(c_name)(*[sub(f"inp{i}") for i in range(nargs)]))
], [
lp.GlobalArg("out", dtype=None, shape=lp.auto, offset=lp.auto)
] + [
lp.GlobalArg(f"inp{i}", dtype=None, shape=lp.auto, offset=lp.auto)
for i in range(nargs)
] + [...],
name=f"actx_special_{c_name}",
tags=(ElementwiseMapKernelTag(),))
return get(c_name, nargs, naxes)
class LoopyBasedFakeNumpyNamespace(BaseFakeNumpyNamespace):
_numpy_to_c_arc_functions: ClassVar[Mapping[str, str]] = {
"arcsin": "asin",
"arccos": "acos",
"arctan": "atan",
"arctan2": "atan2",
"arcsinh": "asinh",
"arccosh": "acosh",
"arctanh": "atanh",
}
_c_to_numpy_arc_functions: ClassVar[Mapping[str, str]] = {c_name: numpy_name
for numpy_name, c_name in _numpy_to_c_arc_functions.items()}
def __getattr__(self, name):
def loopy_implemented_elwise_func(*args):
if all(np.isscalar(ary) for ary in args):
return getattr(
np, self._c_to_numpy_arc_functions.get(name, name)
)(*args)
actx = self._array_context
prg = _get_scalar_func_loopy_program(actx,
c_name, nargs=len(args), naxes=len(args[0].shape))
outputs = actx.call_loopy(prg,
**{f"inp{i}": arg for i, arg in enumerate(args)})
return outputs["out"]
if name in self._c_to_numpy_arc_functions:
raise RuntimeError(f"'{name}' in ArrayContext.np has been removed. "
f"Use '{self._c_to_numpy_arc_functions[name]}' as in numpy. ")
# normalize to C names anyway
c_name = self._numpy_to_c_arc_functions.get(name, name)
# limit which functions we try to hand off to loopy
if (name in self._numpy_math_functions
or name in self._c_to_numpy_arc_functions):
return multimapped_over_array_containers(loopy_implemented_elwise_func)
else:
raise AttributeError(
f"'{type(self._array_context).__name__}.np' object "
f"has no attribute '{name}'")
# }}}
......
arraycontext/metadata.py
View file @ 1013ecb2
"""
.. autoclass:: CommonSubexpressionTag
.. autoclass:: FirstAxisIsElementsTag
.. autoclass:: NameHint
"""
from __future__ import annotations
__copyright__ = """
Copyright (C) 2020-1 University of Illinois Board of Trustees
......@@ -27,26 +28,25 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
from pytools.tag import Tag
from dataclasses import dataclass
from pytools.tag import UniqueTag
# {{{ program metadata
class CommonSubexpressionTag(Tag):
"""A tag that is applicable to arrays indicating that this same array
may be evaluated multiple times, and that the implementation should
eliminate those redundant evaluations if possible.
"""
@dataclass(frozen=True)
class NameHint(UniqueTag):
"""A tag acting on arrays or array axes. Express that :attr:`name` is a
useful starting point in forming an identifier for the tagged object.
.. attribute:: name
class FirstAxisIsElementsTag(Tag):
"""A tag that is applicable to array outputs indicating that the
first index corresponds to element indices. This suggests that
the implementation should set element indices as the outermost
loop extent.
A string. Must be a valid Python identifier. Not necessarily unique.
"""
name: str
# }}}
def __post_init__(self):
if not self.name.isidentifier():
raise ValueError("'name' must be an identifier")
# vim: foldmethod=marker
arraycontext/py.typed 0 → 100644
View file @ 1013ecb2
arraycontext/pytest.py
View file @ 1013ecb2
"""
.. currentmodule:: arraycontext
.. autofunction:: pytest_generate_tests_for_pyopencl_array_context
.. autoclass:: PytestArrayContextFactory
.. autoclass:: PytestPyOpenCLArrayContextFactory
.. autofunction:: pytest_generate_tests_for_array_contexts
"""
from __future__ import annotations
__copyright__ = """
......@@ -28,69 +33,371 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
from collections.abc import Callable, Sequence
from typing import Any
from arraycontext import NumpyArrayContext
from arraycontext.context import ArrayContext
# {{{ array context factories
class PytestArrayContextFactory:
@classmethod
def is_available(cls) -> bool:
return True
def __call__(self) -> ArrayContext:
raise NotImplementedError
class PytestPyOpenCLArrayContextFactory(PytestArrayContextFactory):
"""
.. automethod:: __init__
.. automethod:: __call__
"""
def __init__(self, device):
"""
:arg device: a :class:`pyopencl.Device`.
"""
self.device = device
@classmethod
def is_available(cls) -> bool:
try:
import pyopencl # noqa: F401
return True
except ImportError:
return False
def get_command_queue(self):
# Get rid of leftovers from past tests.
# CL implementations are surprisingly limited in how many
# simultaneous contexts they allow...
from pyopencl.tools import clear_first_arg_caches
clear_first_arg_caches()
from gc import collect
collect()
import pyopencl as cl
# On Intel CPU CL, existence of a command queue does not ensure that
# the context survives.
ctx = cl.Context([self.device])
return ctx, cl.CommandQueue(ctx)
class _PytestPyOpenCLArrayContextFactoryWithClass(PytestPyOpenCLArrayContextFactory):
# Deprecated, remove in 2025.
_force_device_scalars = True
@property
def force_device_scalars(self):
from warnings import warn
warn(
"force_device_scalars is deprecated and will be removed in 2025.",
DeprecationWarning, stacklevel=2)
return self._force_device_scalars
@property
def actx_class(self):
from arraycontext import PyOpenCLArrayContext
return PyOpenCLArrayContext
def __call__(self):
# The ostensibly pointless assignment to *ctx* keeps the CL context alive
# long enough to create the array context, which will then start
# holding a reference to the context to keep it alive in turn.
# On some implementations (notably Intel CPU), holding a reference
# to a queue does not keep the context alive.
_ctx, queue = self.get_command_queue()
alloc = None
if queue.device.platform.name == "NVIDIA CUDA":
from pyopencl.tools import ImmediateAllocator
alloc = ImmediateAllocator(queue)
from warnings import warn
warn("Disabling SVM due to memory leak "
"in Nvidia CL when running pytest. "
"See https://github.com/inducer/arraycontext/issues/196",
stacklevel=1)
return self.actx_class(
queue,
allocator=alloc)
def __str__(self):
return (f"<{self.actx_class.__name__} "
f"for <pyopencl.Device '{self.device.name.strip()}' "
f"on '{self.device.platform.name.strip()}'>>")
class _PytestPytatoPyOpenCLArrayContextFactory(PytestPyOpenCLArrayContextFactory):
@classmethod
def is_available(cls) -> bool:
try:
import pyopencl # noqa: F401
import pytato # noqa: F401
return True
except ImportError:
return False
@property
def actx_class(self):
from arraycontext import PytatoPyOpenCLArrayContext
actx_cls = PytatoPyOpenCLArrayContext
return actx_cls
def __call__(self):
# The ostensibly pointless assignment to *ctx* keeps the CL context alive
# long enough to create the array context, which will then start
# holding a reference to the context to keep it alive in turn.
# On some implementations (notably Intel CPU), holding a reference
# to a queue does not keep the context alive.
_ctx, queue = self.get_command_queue()
alloc = None
if queue.device.platform.name == "NVIDIA CUDA":
from pyopencl.tools import ImmediateAllocator
alloc = ImmediateAllocator(queue)
from warnings import warn
warn("Disabling SVM due to memory leak "
"in Nvidia CL when running pytest. "
"See https://github.com/inducer/arraycontext/issues/196",
stacklevel=1)
return self.actx_class(queue, allocator=alloc)
def __str__(self):
return ("<PytatoPyOpenCLArrayContext for "
f"<pyopencl.Device '{self.device.name.strip()}' "
f"on '{self.device.platform.name.strip()}'>>")
class _PytestEagerJaxArrayContextFactory(PytestArrayContextFactory):
def __init__(self, *args, **kwargs):
pass
@classmethod
def is_available(cls) -> bool:
try:
import jax # noqa: F401
return True
except ImportError:
return False
def __call__(self):
from jax import config
from arraycontext import EagerJAXArrayContext
config.update("jax_enable_x64", True)
return EagerJAXArrayContext()
def __str__(self):
return "<EagerJAXArrayContext>"
# {{{ pytest integration
def pytest_generate_tests_for_pyopencl_array_context(metafunc):
"""Parametrize tests for pytest to use a
:class:`~arraycontext.PyOpenCLArrayContext`.
class _PytestPytatoJaxArrayContextFactory(PytestArrayContextFactory):
def __init__(self, *args, **kwargs):
pass
Performs device enumeration analogously to
:func:`pyopencl.tools.pytest_generate_tests_for_pyopencl`.
@classmethod
def is_available(cls) -> bool:
try:
import jax # noqa: F401
import pytato # noqa: F401
return True
except ImportError:
return False
Using the line:
def __call__(self):
from jax import config
from arraycontext import PytatoJAXArrayContext
config.update("jax_enable_x64", True)
return PytatoJAXArrayContext()
def __str__(self):
return "<PytatoJAXArrayContext>"
# {{{ _PytestArrayContextFactory
class _NumpyArrayContextForTests(NumpyArrayContext):
def transform_loopy_program(self, t_unit):
return t_unit
class _PytestNumpyArrayContextFactory(PytestArrayContextFactory):
def __init__(self, *args, **kwargs):
super().__init__()
def __call__(self):
return _NumpyArrayContextForTests()
def __str__(self):
return "<NumpyArrayContext>"
# }}}
_ARRAY_CONTEXT_FACTORY_REGISTRY: dict[str, type[PytestArrayContextFactory]] = {
"pyopencl": _PytestPyOpenCLArrayContextFactoryWithClass,
"pytato:pyopencl": _PytestPytatoPyOpenCLArrayContextFactory,
"pytato:jax": _PytestPytatoJaxArrayContextFactory,
"eagerjax": _PytestEagerJaxArrayContextFactory,
"numpy": _PytestNumpyArrayContextFactory,
}
def register_pytest_array_context_factory(
name: str,
factory: type[PytestArrayContextFactory]) -> None:
if name in _ARRAY_CONTEXT_FACTORY_REGISTRY:
raise ValueError(f"factory '{name}' already exists")
_ARRAY_CONTEXT_FACTORY_REGISTRY[name] = factory
# }}}
# {{{ pytest integration
def pytest_generate_tests_for_array_contexts(
factories: Sequence[str | type[PytestArrayContextFactory]], *,
factory_arg_name: str = "actx_factory",
) -> Callable[[Any], None]:
"""Parametrize tests for pytest to use an :class:`~arraycontext.ArrayContext`.
Using this function in :mod:`pytest` test scripts allows you to use the
argument *factory_arg_name*, which is a callable that returns a
:class:`~arraycontext.ArrayContext`. All test functions will automatically
be run once for each implemented array context. To select specific array
context implementations explicitly define, for example,
.. code-block:: python
from arraycontext import pytest_generate_tests_for_pyopencl
as pytest_generate_tests
pytest_generate_tests = pytest_generate_tests_for_array_context([
"pyopencl",
])
in your pytest test scripts allows you to use the argument ``actx_factory``,
in your test functions, and they will automatically be
run once for each OpenCL device/platform in the system, as appropriate,
with an argument-less function that returns an
:class:`~arraycontext.ArrayContext` when called.
to use the :mod:`pyopencl`-based array context.
It also allows you to specify the ``PYOPENCL_TEST`` environment variable
for device selection.
"""
The environment variable ``ARRAYCONTEXT_TEST`` can also be used to
overwrite any chosen implementations through *factories*. This is a
comma-separated list of known array contexts.
import pyopencl as cl
from pyopencl.tools import _ContextFactory
Current supported implementations include:
class ArrayContextFactory(_ContextFactory):
def __call__(self):
ctx = super().__call__()
from arraycontext.impl.pyopencl import PyOpenCLArrayContext
return PyOpenCLArrayContext(cl.CommandQueue(ctx))
* ``"pyopencl"``, which creates a :class:`~arraycontext.PyOpenCLArrayContext`.
* ``"pytato-pyopencl"``, which creates a
:class:`~arraycontext.PytatoPyOpenCLArrayContext`.
def __str__(self):
return ("<array context factory for <pyopencl.Device '%s' on '%s'>" %
(self.device.name.strip(),
self.device.platform.name.strip()))
:arg factories: a list of identifiers or
:class:`PytestPyOpenCLArrayContextFactory` classes (not instances)
for which to generate test fixtures.
"""
import pyopencl.tools as cl_tools
arg_names = cl_tools.get_pyopencl_fixture_arg_names(
metafunc, extra_arg_names=["actx_factory"])
# {{{ get all requested array context factories
if not arg_names:
return
import os
env_factory_string = os.environ.get("ARRAYCONTEXT_TEST", None)
arg_values, ids = cl_tools.get_pyopencl_fixture_arg_values()
if "actx_factory" in arg_names:
if "ctx_factory" in arg_names or "ctx_getter" in arg_names:
raise RuntimeError("Cannot use both an 'actx_factory' and a "
"'ctx_factory' / 'ctx_getter' as arguments.")
if env_factory_string is not None:
unique_factories = set(env_factory_string.split(","))
else:
unique_factories = set(factories) # type: ignore[arg-type]
for arg_dict in arg_values:
arg_dict["actx_factory"] = ArrayContextFactory(arg_dict["device"])
if not unique_factories:
raise ValueError("no array context factories were selected")
arg_values = [
tuple(arg_dict[name] for name in arg_names)
for arg_dict in arg_values
unknown_factories = [
factory for factory in unique_factories
if (isinstance(factory, str)
and factory not in _ARRAY_CONTEXT_FACTORY_REGISTRY)
]
metafunc.parametrize(arg_names, arg_values, ids=ids)
if unknown_factories:
if env_factory_string is not None:
raise RuntimeError(
"unknown array context factories passed through environment "
f"variable 'ARRAYCONTEXT_TEST': {unknown_factories}")
else:
raise ValueError(f"unknown array contexts: {unknown_factories}")
available_factories = {
factory for key in unique_factories
for factory in [_ARRAY_CONTEXT_FACTORY_REGISTRY.get(key, key)]
if (
not isinstance(factory, str)
and issubclass(factory, PytestArrayContextFactory)
and factory.is_available())
}
from pytools import partition
pyopencl_factories, other_factories = partition(
lambda factory: issubclass(factory, PytestPyOpenCLArrayContextFactory),
available_factories)
# }}}
def inner(metafunc):
# {{{ get pyopencl devices
import pyopencl.tools as cl_tools
arg_names = cl_tools.get_pyopencl_fixture_arg_names(
metafunc, extra_arg_names=[factory_arg_name])
if not arg_names:
return
arg_values, ids = cl_tools.get_pyopencl_fixture_arg_values()
empty_arg_dict = dict.fromkeys(arg_values[0])
# }}}
# {{{ add array context factory to arguments
if factory_arg_name in arg_names:
if "ctx_factory" in arg_names or "ctx_getter" in arg_names:
raise RuntimeError(
f"Cannot use both an '{factory_arg_name}' and a "
"'ctx_factory' / 'ctx_getter' as arguments.")
arg_values_with_actx = []
if pyopencl_factories:
for arg_dict in arg_values:
arg_values_with_actx.extend([
{factory_arg_name: factory(arg_dict["device"]), **arg_dict}
for factory in pyopencl_factories
])
if other_factories:
arg_values_with_actx.extend([
{factory_arg_name: factory(), **empty_arg_dict}
for factory in other_factories
])
else:
arg_values_with_actx = arg_values
# }}}
# NOTE: sorts the args so that parallel pytest works
arg_value_tuples = sorted([
tuple(arg_dict[name] for name in arg_names)
for arg_dict in arg_values_with_actx
], key=lambda x: str(x))
metafunc.parametrize(arg_names, arg_value_tuples, ids=ids)
return inner
# }}}
......
arraycontext/transform_metadata.py 0 → 100644
View file @ 1013ecb2
"""
.. currentmodule:: arraycontext
.. autoclass:: CommonSubexpressionTag
.. autoclass:: ElementwiseMapKernelTag
"""
from __future__ import annotations
__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 pytools.tag import Tag
# {{{ program metadata
class CommonSubexpressionTag(Tag):
"""A tag that is applicable to arrays indicating that this same array
may be evaluated multiple times, and that the implementation should
eliminate those redundant evaluations if possible.
"""
class ElementwiseMapKernelTag(Tag):
"""A tag that applies to :class:`loopy.LoopKernel` indicating that the kernel
is a "map", i.e. that the output array(s) has/have the same shape as the
input array(s), and that each output element only depends on its corresponding
element(s) in the input array(s).
.. note::
"Element" here refers to a scalar element of an array, not an element
in a finite-element discretization.
"""
# }}}
# vim: foldmethod=marker
arraycontext/version.py
View file @ 1013ecb2
VERSION = (2021, 1)
VERSION_TEXT = ".".join(str(i) for i in VERSION)
from __future__ import annotations
from importlib import metadata
def _parse_version(version: str) -> tuple[tuple[int, ...], str]:
import re
m = re.match(r"^([0-9.]+)([a-z0-9]*?)$", VERSION_TEXT)
assert m is not None
return tuple(int(nr) for nr in m.group(1).split(".")), m.group(2)
VERSION_TEXT = metadata.version("arraycontext")
VERSION, VERSION_STATUS = _parse_version(VERSION_TEXT)
doc/Makefile
View file @ 1013ecb2
......@@ -4,7 +4,7 @@
# You can set these variables from the command line, and also
# from the environment for the first two.
SPHINXOPTS ?=
SPHINXBUILD ?= sphinx-build
SPHINXBUILD ?= python $(shell which sphinx-build)
SOURCEDIR = .
BUILDDIR = _build
......
doc/array_context.rst
View file @ 1013ecb2
......@@ -4,11 +4,3 @@ The Array Context Abstraction
.. automodule:: arraycontext
.. automodule:: arraycontext.context
Implementations of the Array Context Abstraction
================================================
Array context based on :mod:`pyopencl.array`
--------------------------------------------
.. automodule:: arraycontext.impl.pyopencl
doc/conf.py
View file @ 1013ecb2
# -- Path setup --------------------------------------------------------------
from importlib import metadata
from urllib.request import urlopen
# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
# sys.path.insert(0, os.path.abspath('.'))
_conf_url = \
"https://raw.githubusercontent.com/inducer/sphinxconfig/main/sphinxconfig.py"
with urlopen(_conf_url) as _inf:
exec(compile(_inf.read(), _conf_url, "exec"), globals())
# -- Project information -----------------------------------------------------
project = "arraycontext"
copyright = "2021, University of Illinois Board of Trustees"
author = "Arraycontext Contributors"
release = metadata.version("arraycontext")
version = ".".join(release.split(".")[:2])
ver_dic = {}
exec(
compile(
open("../arraycontext/version.py").read(),
"../arraycontext/version.py", "exec"),
ver_dic)
version = ".".join(str(x) for x in ver_dic["VERSION"])
# The full version, including alpha/beta/rc tags.
release = ver_dic["VERSION_TEXT"]
# -- General configuration ---------------------------------------------------
# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
"sphinx.ext.autodoc",
"sphinx.ext.doctest",
"sphinx.ext.intersphinx",
"sphinx.ext.mathjax",
"sphinx.ext.graphviz",
"sphinx_copybutton",
]
# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]
# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
intersphinx_mapping = {
"jax": ("https://jax.readthedocs.io/en/latest/", None),
"loopy": ("https://documen.tician.de/loopy", None),
"meshmode": ("https://documen.tician.de/meshmode", None),
"numpy": ("https://numpy.org/doc/stable/", None),
"pymbolic": ("https://documen.tician.de/pymbolic", None),
"pyopencl": ("https://documen.tician.de/pyopencl", None),
"pytato": ("https://documen.tician.de/pytato", None),
"pytest": ("https://docs.pytest.org/en/latest/", None),
"python": ("https://docs.python.org/3/", None),
"pytools": ("https://documen.tician.de/pytools", None),
}
# -- Options for HTML output -------------------------------------------------
# Some modules need to import things just so that sphinx can resolve symbols in
# type annotations. Often, we do not want these imports (e.g. of PyOpenCL) when
# in normal use (because they would introduce unintended side effects or hard
# dependencies). This flag exists so that these imports only occur during doc
# build. Since sphinx appears to resolve type hints lexically (as it should),
# this needs to be cross-module (since, e.g. an inherited arraycontext
# docstring can be read by sphinx when building meshmode, a dependent package),
# this needs a setting of the same name across all packages involved, that's
# why this name is as global-sounding as it is.
import sys
# The theme to use for HTML and HTML Help pages. See the documentation for
# a list of builtin themes.
#
html_theme = "furo"
# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]
sys._BUILDING_SPHINX_DOCS = True
intersphinx_mapping = {
"https://docs.python.org/3/": None,
"https://numpy.org/doc/stable/": None,
"https://documen.tician.de/pytools": None,
"https://documen.tician.de/pyopencl": None,
"https://documen.tician.de/pytato": None,
"https://documen.tician.de/loopy": None,
"https://documen.tician.de/meshmode": None,
"https://docs.pytest.org/en/latest/": None,
}
autoclass_content = "class"
autodoc_typehints = "description"
nitpick_ignore_regex = [
["py:class", r"arraycontext\.context\.ContainerOrScalarT"],
["py:class", r"ArrayOrContainer"],
]
doc/implementations.rst 0 → 100644
View file @ 1013ecb2
Implementations of the Array Context Abstraction
================================================
..
When adding a new array context here, make sure to also add it to and run
```
doc/make_numpy_coverage_table.py
```
to update the coverage table below!
Array context based on :mod:`numpy`
--------------------------------------------
.. automodule:: arraycontext.impl.numpy
Array context based on :mod:`pyopencl.array`
--------------------------------------------
.. automodule:: arraycontext.impl.pyopencl
Lazy/Deferred evaluation array context based on :mod:`pytato`
-------------------------------------------------------------
.. automodule:: arraycontext.impl.pytato
Array context based on :mod:`jax.numpy`
---------------------------------------
.. automodule:: arraycontext.impl.jax
.. _numpy-coverage:
:mod:`numpy` coverage
---------------------
This is a list of functionality implemented by :attr:`arraycontext.ArrayContext.np`.
.. note::
Only functions and methods that have at least one implementation are listed.
.. include:: numpy_coverage.rst
doc/index.rst
View file @ 1013ecb2
......@@ -7,18 +7,53 @@ implementations for:
- :mod:`numpy`
- :mod:`pyopencl`
- :mod:`pytato`
- :mod:`jax.numpy`
- :mod:`pytato` (for lazy/deferred evaluation)
- Debugging
- Profiling
:mod:`arraycontext` started life as an array abstraction for use with the
:mod:`meshmode` unstrucuted discretization package.
Design Guidelines
-----------------
Here are some of the guidelines we aim to follow in :mod:`arraycontext`. There
exist numerous other, related efforts, such as the `Python array API standard
<https://data-apis.org/array-api/latest/purpose_and_scope.html>`__. These
points may aid in clarifying and differentiating our objectives.
- The array context is about exposing the common subset of operations
available in immutable and mutable arrays. As a result, the interface
does *not* seek to support interfaces that provide, enable, or are typically
used only with in-place mutation.
For example: The equivalents of :func:`numpy.empty` were deprecated
and will eventually be removed.
- Each array context offers a specific subset of of :mod:`numpy` under
:attr:`arraycontext.ArrayContext.np`. Functions under this namespace
must be unconditionally :mod:`numpy`-compatible, that is, they may not
offer an interface beyond what numpy offers. Functions that are
incompatible, for example by supporting tag metadata
(cf. :meth:`arraycontext.ArrayContext.einsum`) should live under the
:class:`~arraycontext.ArrayContext` directly.
- Similarly, we strive to minimize redundancy between attributes of
:class:`~arraycontext.ArrayContext` and :attr:`arraycontext.ArrayContext.np`.
For example: ``ArrayContext.empty_like`` was deprecated.
- Array containers are data structures that may contain arrays.
See :mod:`arraycontext.container`. We strive to support these, where sensible,
in :class:`~arraycontext.ArrayContext` and :attr:`arraycontext.ArrayContext.np`.
Contents
--------
.. toctree::
array_context
implementations
container
other
misc
......
doc/make_numpy_coverage_table.py 0 → 100644
View file @ 1013ecb2
"""
Workflow:
1. If a new array context is implemented, it should be added to
:func:`initialize_contexts`.
2. If a new function is implemented, it should be added to the
corresponding ``write_section_name`` function.
3. Once everything is added, regenerate the tables using
.. code::
python make_numpy_support_table.py numpy_coverage.rst
"""
from __future__ import annotations
import pathlib
from mako.template import Template
import arraycontext
# {{{ templating
HEADER = """
.. raw:: html
<style> .red {color:red} </style>
<style> .green {color:green} </style>
.. role:: red
.. role:: green
"""
TABLE_TEMPLATE = Template("""
${title}
${'~' * len(title)}
.. list-table::
:header-rows: 1
* - Function
% for ctx in contexts:
- :class:`~arraycontext.${type(ctx).__name__}`
% endfor
% for name, (directive, in_context) in numpy_functions_for_context.items():
* - :${directive}:`numpy.${name}`
% for ctx in contexts:
<%
flag = in_context.get(type(ctx), "yes").capitalize()
color = "green" if flag == "Yes" else "red"
%> - :${color}:`${flag}`
% endfor
% endfor
""")
def initialize_contexts():
import pyopencl as cl
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
return [
arraycontext.PyOpenCLArrayContext(queue, force_device_scalars=True),
arraycontext.EagerJAXArrayContext(),
arraycontext.PytatoPyOpenCLArrayContext(queue),
arraycontext.PytatoJAXArrayContext(),
]
def build_supported_functions(funcs, contexts):
import numpy as np
numpy_functions_for_context = {}
for directive, name in funcs:
if not hasattr(np, name):
raise ValueError(f"'{name}' not found in numpy namespace")
in_context = {}
for ctx in contexts:
try:
_ = getattr(ctx.np, name)
except AttributeError:
in_context[type(ctx)] = "No"
numpy_functions_for_context[name] = (directive, in_context)
return numpy_functions_for_context
# }}}
# {{{ writing
def write_array_creation_routines(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.array-creation.html
funcs = (
# (sphinx-directive, name)
("func", "empty_like"),
("func", "ones_like"),
("func", "zeros_like"),
("func", "full_like"),
("func", "copy"),
)
r = TABLE_TEMPLATE.render(
title="Array creation routines",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
def write_array_manipulation_routines(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.array-manipulation.html
funcs = (
# (sphinx-directive, name)
("func", "reshape"),
("func", "ravel"),
("func", "transpose"),
("func", "broadcast_to"),
("func", "concatenate"),
("func", "stack"),
)
r = TABLE_TEMPLATE.render(
title="Array manipulation routines",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
def write_linear_algebra(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.linalg.html
funcs = (
# (sphinx-directive, name)
("func", "vdot"),
)
r = TABLE_TEMPLATE.render(
title="Linear algebra",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
def write_logic_functions(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.logic.html
funcs = (
# (sphinx-directive, name)
("func", "all"),
("func", "any"),
("data", "greater"),
("data", "greater_equal"),
("data", "less"),
("data", "less_equal"),
("data", "equal"),
("data", "not_equal"),
)
r = TABLE_TEMPLATE.render(
title="Logic Functions",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
def write_mathematical_functions(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.math.html
funcs = (
("data", "sin"),
("data", "cos"),
("data", "tan"),
("data", "arcsin"),
("data", "arccos"),
("data", "arctan"),
("data", "arctan2"),
("data", "sinh"),
("data", "cosh"),
("data", "tanh"),
("data", "floor"),
("data", "ceil"),
("func", "sum"),
("data", "exp"),
("data", "log"),
("data", "log10"),
("func", "real"),
("func", "imag"),
("data", "conjugate"),
("data", "maximum"),
("func", "amax"),
("data", "minimum"),
("func", "amin"),
("data", "sqrt"),
("data", "absolute"),
("data", "fabs"),
)
r = TABLE_TEMPLATE.render(
title="Mathematical functions",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
def write_searching_sorting_and_counting(outf, contexts):
# https://numpy.org/doc/stable/reference/routines.sort.html
funcs = (
("func", "where"),
)
r = TABLE_TEMPLATE.render(
title="Sorting, searching, and counting",
contexts=contexts,
numpy_functions_for_context=build_supported_functions(funcs, contexts),
)
outf.write(r)
# }}}
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("filename", nargs="?", type=pathlib.Path, default=None)
args = parser.parse_args()
def write(outf):
outf.write(HEADER)
write_array_creation_routines(outf, ctxs)
write_array_manipulation_routines(outf, ctxs)
write_linear_algebra(outf, ctxs)
write_logic_functions(outf, ctxs)
write_mathematical_functions(outf, ctxs)
ctxs = initialize_contexts()
if args.filename:
with open(args.filename, "w") as outf:
write(outf)
else:
import sys
write(sys.stdout)
doc/numpy_coverage.rst 0 → 100644
View file @ 1013ecb2
.. raw:: html
<style> .red {color:red} </style>
<style> .green {color:green} </style>
.. role:: red
.. role:: green
Array creation routines
~~~~~~~~~~~~~~~~~~~~~~~
.. list-table::
:header-rows: 1
* - Function
- :class:`~arraycontext.PyOpenCLArrayContext`
- :class:`~arraycontext.EagerJAXArrayContext`
- :class:`~arraycontext.PytatoPyOpenCLArrayContext`
- :class:`~arraycontext.PytatoJAXArrayContext`
* - :func:`numpy.empty_like`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.ones_like`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.zeros_like`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.full_like`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.copy`
- :green:`Yes`
- :green:`Yes`
- :red:`No`
- :red:`No`
Array manipulation routines
~~~~~~~~~~~~~~~~~~~~~~~~~~~
.. list-table::
:header-rows: 1
* - Function
- :class:`~arraycontext.PyOpenCLArrayContext`
- :class:`~arraycontext.EagerJAXArrayContext`
- :class:`~arraycontext.PytatoPyOpenCLArrayContext`
- :class:`~arraycontext.PytatoJAXArrayContext`
* - :func:`numpy.reshape`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.ravel`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.transpose`
- :red:`No`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.broadcast_to`
- :red:`No`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.concatenate`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.stack`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
Linear algebra
~~~~~~~~~~~~~~
.. list-table::
:header-rows: 1
* - Function
- :class:`~arraycontext.PyOpenCLArrayContext`
- :class:`~arraycontext.EagerJAXArrayContext`
- :class:`~arraycontext.PytatoPyOpenCLArrayContext`
- :class:`~arraycontext.PytatoJAXArrayContext`
* - :func:`numpy.vdot`
- :green:`Yes`
- :green:`Yes`
- :red:`No`
- :red:`No`
Logic Functions
~~~~~~~~~~~~~~~
.. list-table::
:header-rows: 1
* - Function
- :class:`~arraycontext.PyOpenCLArrayContext`
- :class:`~arraycontext.EagerJAXArrayContext`
- :class:`~arraycontext.PytatoPyOpenCLArrayContext`
- :class:`~arraycontext.PytatoJAXArrayContext`
* - :func:`numpy.all`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.any`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.greater`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.greater_equal`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.less`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.less_equal`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.equal`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.not_equal`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
Mathematical functions
~~~~~~~~~~~~~~~~~~~~~~
.. list-table::
:header-rows: 1
* - Function
- :class:`~arraycontext.PyOpenCLArrayContext`
- :class:`~arraycontext.EagerJAXArrayContext`
- :class:`~arraycontext.PytatoPyOpenCLArrayContext`
- :class:`~arraycontext.PytatoJAXArrayContext`
* - :data:`numpy.sin`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.cos`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.tan`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.arcsin`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.arccos`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.arctan`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.arctan2`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.sinh`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.cosh`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.tanh`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.floor`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.ceil`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.sum`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.exp`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.log`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.log10`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.real`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.imag`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.conjugate`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.maximum`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.amax`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.minimum`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :func:`numpy.amin`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.sqrt`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.absolute`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
* - :data:`numpy.fabs`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`
- :green:`Yes`