diff --git a/arraycontext/fake_numpy.py b/arraycontext/fake_numpy.py index cdb95348c6ca912bb39b01428aa7a0a96ecbfdb2..b9784e1f7ab339744c202354826f1d4bc1bf6117 100644 --- a/arraycontext/fake_numpy.py +++ b/arraycontext/fake_numpy.py @@ -25,8 +25,7 @@ THE SOFTWARE. import numpy as np from arraycontext.container import is_array_container, serialize_container -from arraycontext.container.traversal import ( - rec_map_array_container, multimapped_over_array_containers) +from arraycontext.container.traversal import rec_map_array_container from pytools import memoize_in @@ -75,7 +74,7 @@ class BaseFakeNumpyNamespace: self.linalg = self._get_fake_numpy_linalg_namespace() def _get_fake_numpy_linalg_namespace(self): - return BaseFakeNumpyLinalgNamespace(self.array_context) + return BaseFakeNumpyLinalgNamespace(self._array_context) _numpy_math_functions = frozenset({ # https://numpy.org/doc/stable/reference/routines.math.html @@ -145,31 +144,6 @@ class BaseFakeNumpyNamespace: _c_to_numpy_arc_functions = {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): - 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, - **{"inp%d" % i: arg for i, arg in enumerate(args)}) - return outputs["out"] - - if name in self._c_to_numpy_arc_functions: - from warnings import warn - warn(f"'{name}' in ArrayContext.np is deprecated. " - "Use '{c_to_numpy_arc_functions[name]}' as in numpy. " - "The old name will stop working in 2021.", - DeprecationWarning, stacklevel=3) - - # 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: - return multimapped_over_array_containers(loopy_implemented_elwise_func) - else: - raise AttributeError(name) - def _new_like(self, ary, alloc_like): from numbers import Number diff --git a/arraycontext/impl/pyopencl/fake_numpy.py b/arraycontext/impl/pyopencl/fake_numpy.py index 01054bac6b90d2960f3ddc6ee25cd13fc1d91d4d..c82d7824e1b7e42399d16b8ac3024661c074c4e1 100644 --- a/arraycontext/impl/pyopencl/fake_numpy.py +++ b/arraycontext/impl/pyopencl/fake_numpy.py @@ -30,7 +30,9 @@ from functools import partial, reduce import operator from arraycontext.fake_numpy import \ - BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace + BaseFakeNumpyLinalgNamespace +from arraycontext.loopy import \ + LoopyBasedFakeNumpyNamespace from arraycontext.container.traversal import ( rec_multimap_array_container, rec_map_array_container, rec_map_reduce_array_container, @@ -45,7 +47,7 @@ except ImportError: # {{{ fake numpy -class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace): +class PyOpenCLFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace): def _get_fake_numpy_linalg_namespace(self): return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context) @@ -58,6 +60,17 @@ class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace): # These operations provide access to numpy-style comparisons in that # case. + def __getattr__(self, name): + print(name) + cl_funcs = ["abs", "sin", "cos", "tan", "arcsin", "arccos", "arctan", + "sinh", "cosh", "tanh", "exp", "log", "log10", "isnan", + "sqrt", "exp"] + if name in cl_funcs: + from functools import partial + return partial(rec_map_array_container, getattr(cl, name)) + + return super().__getattr__(name) + def equal(self, x, y): return rec_multimap_array_container(operator.eq, x, y) diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py index f17a4abbc13035bf63956f96dafe005afdc6606f..05f87cea4aac9ef72e178943f00fcd5caf8e5d51 100644 --- a/arraycontext/impl/pytato/fake_numpy.py +++ b/arraycontext/impl/pytato/fake_numpy.py @@ -23,9 +23,10 @@ THE SOFTWARE. """ from functools import partial, reduce -from arraycontext.fake_numpy import ( - BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace, - ) +from arraycontext.fake_numpy import \ + BaseFakeNumpyLinalgNamespace +from arraycontext.loopy import \ + LoopyBasedFakeNumpyNamespace from arraycontext.container.traversal import ( rec_multimap_array_container, rec_map_array_container, rec_map_reduce_array_container, @@ -38,7 +39,7 @@ class PytatoFakeNumpyLinalgNamespace(BaseFakeNumpyLinalgNamespace): pass -class PytatoFakeNumpyNamespace(BaseFakeNumpyNamespace): +class PytatoFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace): """ A :mod:`numpy` mimic for :class:`PytatoPyOpenCLArrayContext`. diff --git a/arraycontext/loopy.py b/arraycontext/loopy.py index f4c97754d731961baaaf0191f70dcfeca287b688..37ce9de83d5634da73420b60744a5c014c0fbdee 100644 --- a/arraycontext/loopy.py +++ b/arraycontext/loopy.py @@ -29,7 +29,9 @@ THE SOFTWARE. import loopy as lp from loopy.version import MOST_RECENT_LANGUAGE_VERSION - +from arraycontext.fake_numpy import BaseFakeNumpyNamespace +from arraycontext.container.traversal import multimapped_over_array_containers +from pytools import memoize_in # {{{ loopy @@ -68,6 +70,80 @@ def get_default_entrypoint(t_unit): raise TypeError("unable to find default entry point for loopy " "translation unit") + +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 import var + + var_names = ["i%d" % i for i in range(naxes)] + size_names = ["n%d" % i for i in range(naxes)] + subscript = tuple(var(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): + 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 .loopy import make_loopy_program + from arraycontext.transform_metadata import ElementwiseMapKernelTag + return make_loopy_program( + [domain_bset], + [ + lp.Assignment( + var("out")[subscript], + var(c_name)(*[ + var("inp%d" % i)[subscript] for i in range(nargs)])) + ], + name="actx_special_%s" % c_name, + tags=(ElementwiseMapKernelTag(),)) + + return get(c_name, nargs, naxes) + + +class LoopyBasedFakeNumpyNamespace(BaseFakeNumpyNamespace): + _numpy_to_c_arc_functions = { + "arcsin": "asin", + "arccos": "acos", + "arctan": "atan", + "arctan2": "atan2", + + "arcsinh": "asinh", + "arccosh": "acosh", + "arctanh": "atanh", + } + + _c_to_numpy_arc_functions = {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): + 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, + **{"inp%d" % i: arg for i, arg in enumerate(args)}) + return outputs["out"] + + if name in self._c_to_numpy_arc_functions: + from warnings import warn + warn(f"'{name}' in ArrayContext.np is deprecated. " + "Use '{c_to_numpy_arc_functions[name]}' as in numpy. " + "The old name will stop working in 2022.", + DeprecationWarning, stacklevel=3) + + # 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: + return multimapped_over_array_containers(loopy_implemented_elwise_func) + else: + raise AttributeError(name) + # }}}