Revert "upgrade 'unevaluated array as argument' warning to error (inducer#305)"

This reverts commit 4aeaed4.

Author: MTCam

arraycontext/impl/jax/fake_numpy.py

@@ -80,7 +80,7 @@ def _empty_like(array):
 
     def zeros_like(self, ary):
         def _zeros_like(array):
-            return self._array_context.np.zeros(array.shape, array.dtype)
+            return self._array_context.zeros(array.shape, array.dtype)
 
         return self._array_context._rec_map_container(
             _zeros_like, ary, default_scalar=0)

arraycontext/impl/pytato/__init__.py

@@ -535,6 +535,7 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
 
             pt_prg = pt.generate_loopy(transformed_dag,
                                        options=opts,
+                                       cl_device=self.queue.device,
                                        function_name=function_name,
                                        target=self.get_target()
                                        ).bind_to_context(self.context)

arraycontext/impl/pytato/compile.py

@@ -529,8 +529,7 @@ def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
         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>"):
+def _args_to_device_buffers(actx, input_id_to_name_in_program, arg_id_to_arg):
     input_kwargs_for_loopy = {}
 
     for arg_id, arg in arg_id_to_arg.items():
@@ -551,20 +550,32 @@ def _args_to_device_buffers(actx, input_id_to_name_in_program, arg_id_to_arg,
             # 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. "
-            )
+            # got an array expression => evaluate it
+            from warnings import warn
+            warn(f"Argument array '{arg_id}' to a compiled function is "
+                    "unevaluated. Evaluating just-in-time, at "
+                    "considerable expense. This is deprecated and will stop "
+                    "working in 2023. To avoid this warning, force evaluation "
+                    "of all arguments via freeze/thaw.",
+                    DeprecationWarning, stacklevel=4)
+
+            arg = actx.freeze(arg)
         else:
             raise NotImplementedError(type(arg))
 
         input_kwargs_for_loopy[input_id_to_name_in_program[arg_id]] = arg
 
     return input_kwargs_for_loopy
 
+
+def _args_to_cl_buffers(actx, input_id_to_name_in_program, arg_id_to_arg):
+    from warnings import warn
+    warn("_args_to_cl_buffer has been renamed to"
+         " _args_to_device_buffers. This will be"
+         " an error in 2023.", DeprecationWarning, stacklevel=2)
+    return _args_to_device_buffers(actx, input_id_to_name_in_program,
+                                   arg_id_to_arg)
+
 # }}}
 
 
@@ -620,7 +631,7 @@ class CompiledPyOpenCLFunctionReturningArrayContainer(CompiledFunction):
        type of the callable.
     """
     actx: PytatoPyOpenCLArrayContext
-    pytato_program: pt.target.loopy.BoundPyOpenCLExecutable
+    pytato_program: pt.target.BoundProgram
     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]]
@@ -631,10 +642,8 @@ 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)
+                self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
         evt, out_dict = self.pytato_program(queue=self.actx.queue,
                                             allocator=self.actx.allocator,
@@ -665,7 +674,7 @@ class CompiledPyOpenCLFunctionReturningArray(CompiledFunction):
         Name of the output array in the program.
     """
     actx: PytatoPyOpenCLArrayContext
-    pytato_program: pt.target.loopy.BoundPyOpenCLExecutable
+    pytato_program: pt.target.BoundProgram
     input_id_to_name_in_program: Mapping[tuple[Hashable, ...], str]
     output_tags: frozenset[Tag]
     output_axes: tuple[pt.Axis, ...]
@@ -675,10 +684,8 @@ 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)
+                self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
         evt, out_dict = self.pytato_program(queue=self.actx.queue,
                                             allocator=self.actx.allocator,
@@ -716,18 +723,16 @@ class CompiledJAXFunctionReturningArrayContainer(CompiledFunction):
        type of the callable.
     """
     actx: PytatoJAXArrayContext
-    pytato_program: pt.target.python.BoundJAXPythonProgram
+    pytato_program: pt.target.BoundProgram
     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)
+                self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
         out_dict = self.pytato_program(**input_kwargs_for_loopy)
 
@@ -749,17 +754,15 @@ class CompiledJAXFunctionReturningArray(CompiledFunction):
         Name of the output array in the program.
     """
     actx: PytatoJAXArrayContext
-    pytato_program: pt.target.python.BoundJAXPythonProgram
+    pytato_program: pt.target.BoundProgram
     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)
+                self.actx, self.input_id_to_name_in_program, arg_id_to_arg)
 
         _evt, out_dict = self.pytato_program(**input_kwargs_for_loopy)
 

test/test_arraycontext.py

@@ -1149,6 +1149,7 @@ def test_actx_compile_kwargs(actx_factory):
 def test_actx_compile_with_tuple_output_keys(actx_factory):
     # arraycontext.git<=3c9aee68 would fail due to a bug in output
     # key stringification logic.
+    from arraycontext import from_numpy, to_numpy
     actx = actx_factory()
     rng = np.random.default_rng()
 
@@ -1162,11 +1163,11 @@ def my_rhs(scale, vel):
     v_x = rng.uniform(size=10)
     v_y = rng.uniform(size=10)
 
-    vel = actx.from_numpy(Velocity2D(v_x, v_y, actx))
+    vel = from_numpy(Velocity2D(v_x, v_y, actx), actx)
 
     scaled_speed = compiled_rhs(3.14, vel=vel)
 
-    result = actx.to_numpy(scaled_speed)[0, 0]
+    result = to_numpy(scaled_speed, actx)[0, 0]
     np.testing.assert_allclose(result.u, -3.14*v_y)
     np.testing.assert_allclose(result.v, 3.14*v_x)
 
@@ -1292,8 +1293,6 @@ class ArrayContainerWithNumpy:
     u: np.ndarray
     v: DOFArray
 
-    __array_ufunc__ = None
-
 
 def test_array_container_with_numpy(actx_factory):
     actx = actx_factory()
@@ -1412,16 +1411,14 @@ def test_compile_anonymous_function(actx_factory):
 
     # See https://github.com/inducer/grudge/issues/287
     actx = actx_factory()
-
-    ones = actx.thaw(actx.freeze(
-        actx.np.zeros(shape=(10, 4), dtype=np.float64) + 1
-        ))
-
     f = actx.compile(lambda x: 2*x+40)
-    np.testing.assert_allclose(actx.to_numpy(f(ones)), 42)
-
+    np.testing.assert_allclose(
+        actx.to_numpy(f(1+actx.np.zeros((10, 4), "float64"))),
+        42)
     f = actx.compile(partial(lambda x: 2*x+40))
-    np.testing.assert_allclose(actx.to_numpy(f(ones)), 42)
+    np.testing.assert_allclose(
+        actx.to_numpy(f(1+actx.np.zeros((10, 4), "float64"))),
+        42)
 
 
 @pytest.mark.parametrize(

test/testlib.py

@@ -160,7 +160,7 @@ def array_context(self):
 
 @with_container_arithmetic(
         bcasts_across_obj_array=False,
-        container_types_bcast_across=(DOFArray, np.ndarray),
+        bcast_container_types=(DOFArray, np.ndarray),
         matmul=True,
         rel_comparison=True,
         _cls_has_array_context_attr=True,
@@ -173,8 +173,6 @@ class MyContainerDOFBcast:
     momentum: np.ndarray
     enthalpy: DOFArray | np.ndarray
 
-    __array_ufunc__ = None
-
     @property
     def array_context(self):
         if isinstance(self.mass, np.ndarray):
@@ -211,8 +209,6 @@ class Velocity2D:
     v: ArrayContainer
     array_context: ArrayContext
 
-    __array_ufunc__ = None
-
 
 @with_array_context.register(Velocity2D)
 # https://github.com/python/mypy/issues/13040