feat[next]: Compiled variant for field operators #2368
Conversation
|
|
||
| foast_to_itir: workflow.Workflow[AOT_FOP, itir.FunctionDefinition] | ||
|
|
||
| def __call__(self, inp: AOT_FOP) -> AOT_PRG: |
There was a problem hiding this comment.
This function used to be wrong. The out argument was previously passed in kwarg_types, but this function just ignored all values in kwarg_types and just added out manually again. With the CompiledProgramPool the out argument is just the third argument and not a keyword argument anymore (consistent with the program definition below). Since this function should be reworked or replaced I didn't bother and just dropped the manual addition and added some sanity checks.
| ) | ||
|
|
||
|
|
||
| @dataclasses.dataclass(frozen=True) |
There was a problem hiding this comment.
This is not used anywhere including in icon4py and PMAP so I removed it instead of investing time into making it compatible with _CommonProgramLike.
|
|
||
| # TODO(tehrengruber): This class does not follow the Liskov-Substitution principle as it doesn't | ||
| # have a program definition. Revisit. | ||
| @dataclasses.dataclass(frozen=True) |
There was a problem hiding this comment.
This is not used anywhere including in icon4py and PMAP so I removed it instead of investing time into making it compatible with _CommonProgramLike. Since Backend.__call__ was removed it couldn't stay.
| @@ -1,400 +0,0 @@ | |||
| ```python | |||
There was a problem hiding this comment.
This file has become mostly obsolete with the introduction of CompiledProgram, hence I removed it here.
egparedes
left a comment
egparedes
left a comment
There was a problem hiding this comment.
Looks good to me. Don't have any serious objection, mostly naming and code style issues, many of them optional.
src/gt4py/next/ffront/decorator.py
Outdated
|
|
||
| DEFAULT_BACKEND: next_backend.Backend | None = None | ||
|
|
||
| DefinitionT = TypeVar( |
There was a problem hiding this comment.
Not criticial but this typevar name is too generic. What about DSLDefT or something in that direction?
There was a problem hiding this comment.
I went for ProgramLikeDefinitionT for now. If we really want to do DSLCallable we should also rename CompiledProgram to CompiledDSLCallable etc. This doesn't feel like a good fit for this PR.
There was a problem hiding this comment.
I think ProgramLike is meaningless in this context and that the right solution would be exactly what you propose as non-fitting this PR: use DSLCallable here and rename CompiledProgram to CompiledDSLCallable. Anyway, unless @havogt has a strong opinion towards one of the options, this should not be a blocker for this PR.
src/gt4py/next/ffront/decorator.py
Outdated
| class _ProgramLikeMixin(Generic[DefinitionT]): | ||
| """ | ||
| Mixing used by program and program-like objects. | ||
|
|
||
| Contains functionality and configuration options common to all kinds of program-likes. | ||
| """ |
There was a problem hiding this comment.
I would suggest a different naming, in the same spirit as the previous suggestion:
| class _ProgramLikeMixin(Generic[DefinitionT]): | |
| """ | |
| Mixing used by program and program-like objects. | |
| Contains functionality and configuration options common to all kinds of program-likes. | |
| """ | |
| class _DSLCallableMixin(Generic[DSLDefT]): | |
| """ | |
| Mix-in used by directly callable GT4Py DSL artifacts like programs and operators. | |
| Contains functionality and configuration options common to all kinds of callable artifacts. | |
| """ |
There was a problem hiding this comment.
See comment above.
3 participants
In #1968 we introduced pre- and jit-compilation of static arguments. In the first version this was only support for programs, but not field operators which were directly called. This PR extends to functionality to also support field operators and scans.
A good amount of technical debt which accumulated around the
Backend.__call__method was also removed as it became obsolete with the unified approach.Since scans are the only type of callable with generic arguments and there is a good amount of technical debt surrounding them that was beyond the scope of this PR, the first version barely implements something working for them and prints a warning as performance is likely not great. If performance is needed a simple workaround is to wrap the scan into a field operator or program.