Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

single_version_only for actx.compile #228

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion arraycontext/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -504,7 +504,8 @@ def clone(self: SelfType) -> SelfType:
"setup-only" array context "leaks" into the application.
"""

def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
def compile(self, f: Callable[..., Any],
single_version_only: bool = False) -> Callable[..., Any]:
inducer marked this conversation as resolved.
Show resolved Hide resolved
"""Compiles *f* for repeated use on this array context. *f* is expected
to be a `pure function <https://en.wikipedia.org/wiki/Pure_function>`__
performing an array computation.
Expand Down
6 changes: 4 additions & 2 deletions arraycontext/impl/pytato/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -612,9 +612,11 @@ def call_loopy(self, program, **kwargs):

return call_loopy(program, processed_kwargs, entrypoint)

def compile(self, f: Callable[..., Any]) -> Callable[..., Any]:
def compile(self, f: Callable[..., Any],
single_version_only: bool = False) -> Callable[..., Any]:
from .compile import LazilyPyOpenCLCompilingFunctionCaller
return LazilyPyOpenCLCompilingFunctionCaller(self, f)
return LazilyPyOpenCLCompilingFunctionCaller(self,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe also handle the other sub-classes of BaseLazilyCompilingFunctionCaller?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added in 5ed25a5

f, single_version_only)

def transform_dag(self, dag: "pytato.DictOfNamedArrays"
) -> "pytato.DictOfNamedArrays":
Expand Down
7 changes: 6 additions & 1 deletion arraycontext/impl/pytato/compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,8 @@ class BaseLazilyCompilingFunctionCaller:
program_cache: Dict["PMap[Tuple[Any, ...], AbstractInputDescriptor]",
"CompiledFunction"] = field(default_factory=lambda: {})

single_version_only: bool = False
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm worried this might lead to developer errors. I would prefer if we don't default and ensure that the current version is passed during instantiation.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I removed the default in 5ed25a5

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fails in grudge now:

  File "wave-mpi-lazy.py", line 201, in main
    compiled_rhs = actx.compile(rhs)
                   ^^^^^^^^^^^^^^^^^
  File "/home/runner/work/arraycontext/arraycontext/mirgecom/src/grudge/grudge/array_context.py", line 416, in compile
    return _DistributedLazilyPyOpenCLCompilingFunctionCaller(self, f)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: BaseLazilyCompilingFunctionCaller.__init__() missing 1 required positional argument: 'single_version_only'

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm... Maybe then we introduce the default as you suggested, but have a deprecation period attached to it?


# {{{ abstract interface

def _dag_to_transformed_pytato_prg(self, dict_of_named_arrays, *, prg_id=None):
Expand Down Expand Up @@ -324,7 +326,10 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any:
try:
compiled_f = self.program_cache[arg_id_to_descr]
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we save ourselves having to find arg_id_to_descr?

except KeyError:
pass
if self.single_version_only and self.program_cache:
raise ValueError(
f"Function '{self.f.__name__}' to be compiled "
"was already compiled previously with different arguments.")
else:
return compiled_f(arg_id_to_arg)

Expand Down