Skip to content

Support lambda-based parallel operations #346

@IvanGrigorik

Description

@IvanGrigorik

PyKokkos has a way to support the execution of lambda-based expressions execution strictly inside of workloads:

@pk.workload
class someTest:
    def __init__(self, threads: int, value: int):
        self.threads: int = threads
        self.value: int = value
        self.sum: int = 0

    @pk.main
    def run(self):
        self.sum = pk.parallel_reduce(
            "name", self.threads, lambda i, acc: acc + self.value
        )

However, since they will be deprecated soon, we need to port this into workunits or similar semantics, since they are not supported by pykokkos from non-pykokkos functions. So if you have consecutive code that uses parallel_* operation with a lambda function like that:

main():
    value: int = value
    sum: int = 0
    threads: int = 10
    sum = pk.parallel_reduce("name", threads, lambda i, acc: acc + value)

PyKokkos would be unable to translate and execute it, since it is not supported by the framework yet.

I suggest adding a lambda-based expression based on these semantics, since it looks intuitive, and this allows us to move from PyKokkos-specialized syntax to Pythonized syntax itself.

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions