From 5218b2d6398d818f517d3a9856a9dabb6076dea0 Mon Sep 17 00:00:00 2001
From: Mit Kotak <53411468+mitkotak@users.noreply.github.com>
Date: Sun, 16 Jan 2022 15:20:19 -0600
Subject: [PATCH] Remove BaseFakeNumpySpace's dependency on call_loopy (#88)
Closes https://github.com/inducer/arraycontext/issues/70.
---
arraycontext/fake_numpy.py | 52 +--------------
arraycontext/impl/pyopencl/fake_numpy.py | 10 ++-
arraycontext/impl/pytato/fake_numpy.py | 7 +-
arraycontext/loopy.py | 84 +++++++++++++++++++++++-
4 files changed, 97 insertions(+), 56 deletions(-)
diff --git a/arraycontext/fake_numpy.py b/arraycontext/fake_numpy.py
index c75295c..25343df 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 4e9d48b..514193f 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 905fd0f..e1730b9 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 f4c9775..728904e 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)
+
# }}}
--
GitLab