diff --git a/arraycontext/fake_numpy.py b/arraycontext/fake_numpy.py
index c75295c1b582527c57655b7ec278ce1a2110fbce..25343dfec164c461a7bdd1b6839dc2ae748938fa 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 NotAnArrayContainerError, 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
@@ -68,14 +67,13 @@ def _get_scalar_func_loopy_program(actx, c_name, nargs, naxes):
# {{{ BaseFakeNumpyNamespace
-
class BaseFakeNumpyNamespace:
def __init__(self, array_context):
self._array_context = array_context
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
@@ -128,54 +126,8 @@ class BaseFakeNumpyNamespace:
# FIXME:
# "interp",
-
})
- _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):
- 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,
- **{"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. "
- f"Use '{self._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
- or name in self._c_to_numpy_arc_functions):
- return multimapped_over_array_containers(loopy_implemented_elwise_func)
- else:
- raise AttributeError(name)
-
def _new_like(self, ary, alloc_like):
if np.isscalar(ary):
# NOTE: `np.zeros_like(x)` returns `array(x, shape=())`, which
diff --git a/arraycontext/impl/pyopencl/fake_numpy.py b/arraycontext/impl/pyopencl/fake_numpy.py
index 4e9d48beca3de0b505591300b2f17df1e1de31ac..514193ff4e448dde3c441899411f4f7ee6e34124 100644
--- a/arraycontext/impl/pyopencl/fake_numpy.py
+++ b/arraycontext/impl/pyopencl/fake_numpy.py
@@ -31,8 +31,12 @@ import operator
import numpy as np
-from arraycontext.fake_numpy import \
- BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace
+from arraycontext.fake_numpy import (
+ BaseFakeNumpyLinalgNamespace
+ )
+from arraycontext.loopy import (
+ LoopyBasedFakeNumpyNamespace
+ )
from arraycontext.container import NotAnArrayContainerError, serialize_container
from arraycontext.container.traversal import (
rec_map_array_container,
@@ -50,7 +54,7 @@ except ImportError:
# {{{ fake numpy
-class PyOpenCLFakeNumpyNamespace(BaseFakeNumpyNamespace):
+class PyOpenCLFakeNumpyNamespace(LoopyBasedFakeNumpyNamespace):
def _get_fake_numpy_linalg_namespace(self):
return _PyOpenCLFakeNumpyLinalgNamespace(self._array_context)
diff --git a/arraycontext/impl/pytato/fake_numpy.py b/arraycontext/impl/pytato/fake_numpy.py
index 905fd0f873ba86597eed79ac0387ca9e8a494ac9..e1730b93c6834f09f69a19958429ddd8ec40fa4a 100644
--- a/arraycontext/impl/pytato/fake_numpy.py
+++ b/arraycontext/impl/pytato/fake_numpy.py
@@ -26,7 +26,10 @@ from functools import partial, reduce
import numpy as np
from arraycontext.fake_numpy import (
- BaseFakeNumpyNamespace, BaseFakeNumpyLinalgNamespace,
+ BaseFakeNumpyLinalgNamespace
+ )
+from arraycontext.loopy import (
+ LoopyBasedFakeNumpyNamespace
)
from arraycontext.container import NotAnArrayContainerError, serialize_container
from arraycontext.container.traversal import (
@@ -42,7 +45,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..728904e44a17423b8d9ae38883bf73ec23a34550 100644
--- a/arraycontext/loopy.py
+++ b/arraycontext/loopy.py
@@ -27,9 +27,12 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""
+import numpy as np
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 +71,85 @@ 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):
+ 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,
+ **{"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. "
+ f"Use '{self._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
+ or name in self._c_to_numpy_arc_functions):
+ return multimapped_over_array_containers(loopy_implemented_elwise_func)
+ else:
+ raise AttributeError(name)
+
# }}}