Add a pass to outline computations in a function

arraycontext/context.py

@@ -167,7 +167,7 @@
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Hashable, Mapping
 from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, TypeVar, Union, overload
 from warnings import warn
 
@@ -328,6 +328,7 @@ class ArrayContext(ABC):
     .. automethod:: tag
     .. automethod:: tag_axis
     .. automethod:: compile
+    .. automethod:: outline
     """
 
     array_types: tuple[type, ...] = ()
@@ -577,6 +578,34 @@ def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
         """
         return f
 
+    # FIXME: Think about making this a standalone function? Would make it easier to
+    # pass arguments when used as a decorator, e.g.:
+    #     @outline(actx, id=...)
+    #     def func(...):
+    # vs.
+    #     outline = partial(actx.outline, id=...)
+    #
+    #     @outline
+    #     def func(...):
+    def outline(self,
+                f: Callable[..., Any],
+                *,
+                id: Hashable | None = None) -> Callable[..., Any]:
+        """
+        Returns a drop-in-replacement for *f*. The behavior of the returned
+        callable is specific to the derived class.
+
+        The reason for the existence of such a routine is mainly for
+        arraycontexts that allow a lazy mode of execution. In such
+        arraycontexts, the computations within *f* maybe staged to potentially
+        enable additional compiler transformations. See
+        :func:`pytato.trace_call` or :func:`jax.named_call` for examples.
+
+        :arg f: the function executing the computation to be staged.
+        :return: a function with the same signature as *f*.
+        """
+        return f
+
     # undocumented for now
     @property
     @abstractmethod

arraycontext/impl/pytato/__init__.py

@@ -53,7 +53,7 @@
 
 import abc
 import sys
-from collections.abc import Callable
+from collections.abc import Callable, Hashable
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -226,6 +226,21 @@ def get_target(self):
 
     # }}}
 
+    def outline(self,
+                f: Callable[..., Any],
+                *,
+                id: Hashable | None = None,
+                tags: frozenset[Tag] = frozenset()
+                ) -> Callable[..., Any]:
+        from pytato.tags import FunctionIdentifier
+
+        from .outline import OutlinedCall
+        id = id or getattr(f, "__name__", None)
+        if id is not None:
+            tags = tags | {FunctionIdentifier(id)}
+
+        return OutlinedCall(self, f, tags)
+
 # }}}
 
 
@@ -443,8 +458,8 @@ def freeze(self, array):
             TaggableCLArray,
             to_tagged_cl_array,
         )
-        from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
         from arraycontext.impl.pytato.utils import (
+            _ary_container_key_stringifier,
             _normalize_pt_expr,
             get_cl_axes_from_pt_axes,
         )
@@ -507,6 +522,12 @@ def _to_frozen(key: tuple[Any, ...], ary) -> TaggableCLArray:
 
         pt_dict_of_named_arrays = pt.make_dict_of_named_arrays(
                 key_to_pt_arrays)
+
+        # FIXME: Remove this if/when _normalize_pt_expr gets support for functions
+        pt_dict_of_named_arrays = pt.tag_all_calls_to_be_inlined(
+            pt_dict_of_named_arrays)
+        pt_dict_of_named_arrays = pt.inline_calls(pt_dict_of_named_arrays)
+
         normalized_expr, bound_arguments = _normalize_pt_expr(
                 pt_dict_of_named_arrays)
 
@@ -674,7 +695,7 @@ def preprocess_arg(name, arg):
                 # multiple placeholders with the same name that are not
                 # also the same object are not allowed, and this would produce
                 # a different Placeholder object of the same name.
-                if (not isinstance(ary, pt.Placeholder)
+                if (not isinstance(ary, pt.Placeholder | pt.NamedArray)
                         and not ary.tags_of_type(NameHint)):
                     ary = ary.tagged(NameHint(name))
 
@@ -779,7 +800,7 @@ def freeze(self, array):
         import pytato as pt
 
         from arraycontext.container.traversal import rec_keyed_map_array_container
-        from arraycontext.impl.pytato.compile import _ary_container_key_stringifier
+        from arraycontext.impl.pytato.utils import _ary_container_key_stringifier
 
         array_as_dict: dict[str, jnp.ndarray | pt.Array] = {}
         key_to_frozen_subary: dict[str, jnp.ndarray] = {}

arraycontext/impl/pytato/compile.py

@@ -110,28 +110,6 @@ class LeafArrayDescriptor(AbstractInputDescriptor):
 
 # {{{ utilities
 
-def _ary_container_key_stringifier(keys: tuple[Any, ...]) -> str:
-    """
-    Helper for :meth:`BaseLazilyCompilingFunctionCaller.__call__`. Stringifies an
-    array-container's component's key. Goals of this routine:
-
-    * No two different keys should have the same stringification
-    * Stringified key must a valid identifier according to :meth:`str.isidentifier`
-    * (informal) Shorter identifiers are preferred
-    """
-    def _rec_str(key: Any) -> str:
-        if isinstance(key, str | int):
-            return str(key)
-        elif isinstance(key, tuple):
-            # t in '_actx_t': stands for tuple
-            return "_actx_t" + "_".join(_rec_str(k) for k in key) + "_actx_endt"
-        else:
-            raise NotImplementedError("Key-stringication unimplemented for "
-                                      f"'{type(key).__name__}'.")
-
-    return "_".join(_rec_str(key) for key in keys)
-
-
 def _get_arg_id_to_arg_and_arg_id_to_descr(args: tuple[Any, ...],
                                            kwargs: Mapping[str, Any]
                                            ) -> \
@@ -322,6 +300,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
         :attr:`~BaseLazilyCompilingFunctionCaller.f` with *args* in a lazy-sense.
         The intermediary pytato DAG for *args* is memoized in *self*.
         """
+        from arraycontext.impl.pytato.utils import _ary_container_key_stringifier
         arg_id_to_arg, arg_id_to_descr = _get_arg_id_to_arg_and_arg_id_to_descr(
             args, kwargs)