diff --git a/arraycontext/context.py b/arraycontext/context.py index 04277f1ccefb29f502df6b004b77a385d0a02be5..13ce197ab2f02a62c8fe0a221c618dc3f27555b2 100644 --- a/arraycontext/context.py +++ b/arraycontext/context.py @@ -202,35 +202,6 @@ class ArrayContext(ABC): array understood by the context. """ - @memoize_method - def _get_scalar_func_loopy_program(self, 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(),)) - @abstractmethod def freeze(self, array): """Return a version of the context-defined array *array* that is diff --git a/arraycontext/fake_numpy.py b/arraycontext/fake_numpy.py index 1f208f3dd02f878e7be1523f8586c2caa3f1ddcd..cdb95348c6ca912bb39b01428aa7a0a96ecbfdb2 100644 --- a/arraycontext/fake_numpy.py +++ b/arraycontext/fake_numpy.py @@ -27,6 +27,44 @@ 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 pytools import memoize_in + + +# {{{ _get_scalar_func_loopy_program + +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) + +# }}} # {{{ BaseFakeNumpyNamespace @@ -110,13 +148,11 @@ class BaseFakeNumpyNamespace: def __getattr__(self, name): def loopy_implemented_elwise_func(*args): actx = self._array_context - # FIXME: Maybe involve loopy type inference? - result = actx.empty(args[0].shape, args[0].dtype) - prg = actx._get_scalar_func_loopy_program( + prg = _get_scalar_func_loopy_program(actx, c_name, nargs=len(args), naxes=len(args[0].shape)) - actx.call_loopy(prg, out=result, + outputs = actx.call_loopy(prg, **{"inp%d" % i: arg for i, arg in enumerate(args)}) - return result + return outputs["out"] if name in self._c_to_numpy_arc_functions: from warnings import warn diff --git a/test/test_arraycontext.py b/test/test_arraycontext.py index 4f5a359e8f2c566ad30a3a84fc1ba47514cab283..e1a2c08b4cc63026621b6ceb7b11734015baae53 100644 --- a/test/test_arraycontext.py +++ b/test/test_arraycontext.py @@ -88,6 +88,14 @@ pytest_generate_tests = pytest_generate_tests_for_array_contexts([ _PytatoPyOpenCLArrayContextForTestsFactory, ]) + +def _acf(): + import pyopencl as cl + + context = cl._csc() + queue = cl.CommandQueue(context) + return _PyOpenCLArrayContextForTests(queue, force_device_scalars=True) + # }}} @@ -858,6 +866,17 @@ def test_container_equality(actx_factory): assert isinstance(bcast_dc_of_dofs == bcast_dc_of_dofs_2, MyContainerDOFBcast) +def test_abs_complex(actx_factory): + actx = actx_factory() + a = np.random.randn(2000) + 1j * np.random.randn(2000) + + abs_a_ref = np.abs(a) + abs_a = actx.np.abs(actx.from_numpy(a)) + + assert abs_a.dtype == abs_a_ref.dtype + np.testing.assert_allclose(actx.to_numpy(abs_a), abs_a_ref) + + if __name__ == "__main__": import sys if len(sys.argv) > 1: