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

Fusion: allow fusion of a single parallel_reduce with a sequence of parallel_fors #257

Merged
merged 5 commits into from
Feb 24, 2024

Conversation

NaderAlAwar
Copy link
Contributor

This PR enables fusion of a single parallel_reduce that with a sequence of parallel_fors that it follows.

@NaderAlAwar NaderAlAwar merged commit 9e98896 into kokkos:main Feb 24, 2024
5 checks passed
) -> Tuple[Dict[str, Any], List[ast.arg]]:
"""
Fuse the parameters and runtime arguments of a list of workunits and rename them as necessary

:param workunits_trees: the list of workunit trees (ASTs) being merged
:param kwargs: the keyword arguments passed to the call
:param operation: they type of parallel operation ("parallel_for", "parallel_reduce", or "parallel_scan")
Copy link
Contributor

Choose a reason for hiding this comment

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

they?

:returns: a tuple of the fused kwargs and the combined inspected parameters
"""

if operation == "parallel_scan":
Copy link
Contributor

Choose a reason for hiding this comment

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

would be nice to put supported operations in an enum or just constants

@@ -140,6 +140,9 @@ def __setitem__(self, key: Union[int, TeamMember], value: Union[int, float]) ->
:param value: the new value at the index.
"""

if "PK_FUSION" in os.environ:
Copy link
Contributor

Choose a reason for hiding this comment

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

I love env variables

Copy link
Contributor

@gliga gliga left a comment

Choose a reason for hiding this comment

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

LG

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants