legate.core.task.task#

legate.core.task.task( func: object = None, *, tuple variants: tuple[VariantCode, ...] = DEFAULT_VARIANT_LIST, constraints: Sequence[DeferredConstraint | Sequence[DeferredConstraint]] | None = None, options: TaskConfig | VariantOptions | None = None, ) → Callable[[UserFunction], PyTask] | PyTask#

Convert a Python function to a Legate task.

The user may pass either a TaskConfig or VariantOptions to select the task options. If the former, the user can select the task ID of the task themselves, while additionally setting variant options via the TaskConfig.variant_options property. If the latter, then a new, unique task ID will be automatically generated.

Parameters:

func (UserFunction) – The base user function to invoke in the task.
variants (tuple[VariantCode, ...], optional) – The list of variants for which func is applicable. Defaults to (<VariantCode.CPU: 1>,).
constraints (Sequence[DeferredConstraint | Sequence[DeferredConstraint]], optional) – The list of constraints which are to be applied to the arguments of func, if any. Defaults to no constraints.
options (TaskConfig | VariantOptions, optional) – Either a TaskConfig or VariantOptions describing the task configuration.

Returns:

The task object.

Return type:

PyTask

Example

from legate.core import broadcast, align, VariantCode, VariantOptions
from legate.core.task import task, InputArray, OutputArray

@task
def my_basic_task(
    x: InputArray,
    y: OutputArray,
    z: tuple[int, ...] = (1, 2, 3)
 ) -> None:
    ...

@task(
    variants=(VariantCode.CPU, VariantCode.GPU),
    constraints=(align("x", "y"), broadcast("x")),
    options=VariantOptions(may_throw_exception=True)
)
def my_task_with_options(
    x: InputArray,
    y: OutputArray,
    z: tuple[int, ...] = (1, 2, 3)
) -> None:
    raise RuntimeError("Exceptional!")

By default, the @task decorator registers the wrapped callable for all variants. It is possible, however, to select additional functions to act as specific variants. Variants must have identical calling signatures to the original function and inherit the same constraints and options.

The original function is used for any unspecified variants.

from legate.core.task import task

# Acts as the CPU and OMP variant as these were not specified
# separately
@task
def my_task() -> None:
    ...

@my_task.gpu_variant
def gpu_task_variant() -> None:
    ...