Work Units

Work units are the third key concept (out of three) in PyKokkos (Concepts). PyKokkos work units are specified as one of the arguments to patterns (Patterns). Each work unit is a Python function decorated with @pk.workunit. At the time a pattern is executed (e.g., parallel_for), the work unit function is invoked for each unique id depending on the policy used.

Note

PyKokkos work units are written in a subset of Python. These functions are translated to C++ and compiled to various target architectures. Thus, work units should be lightweight and free from use of various Python libraries.

Signatures

Depending on the pattern used to invoke a work unit (e.g., parallel_for), the signature of the work unit function differs.

In case of parallel_for, the work unit function is required to accept, as the first argument, unique id for one unit of work.

@pk.workunit
def work(wid, [keyword arguments])

In case of parallel_reduce, the work unit function is required to accept two arguments: (1) unique id for one unit of work, and (2) an accumulator.

@pk.workunit
def work(wid, acc, [keyword arguments])

Finally, in case of parallel_scan, the work unit function is requires to accept three arguments: (1) unique id for one unit of work, (2) an accumulator, and (3) a boolean flag to indicate if the scan for that unique id is complete.

@pk.workunit
def work(wid, acc, final, [keyword arguments])

Keyword Arguments

Each work unit function can accept an arbitrary number of arguments after those required in the signature. At the time of a pattern execution, the arguments have to be given as keyword arguments.

Type Annotations

Type annotations for the arguments are optional, but recommended. We use type annotations heavily in the examples in the PyKokkos repository.

Note

Type annotations for variables defined inside a work unit are currently required.