hidet.ir.task

Classes:

Task(name, inputs, outputs, *[, ...])

A task defines the operator computation.

class hidet.ir.task.Task(name, inputs, outputs, *, inverse_map=None, attributes=None, share_map=None)

A task defines the operator computation.

name

The name of the task.

Type:

str

inputs

The input tensors of this task.

Type:

List[TensorInput]

outputs

The output tensors of this task. They are derived by applying computation on the input tensors.

Type:

List[TensorNode]

inverse_map

The inverse map. It records how the input tensors are derived from the output tensor (when there is only one output tensor). This is used to support epilogue fusion.

Type:

Dict[TensorInput, InverseMap]

attrs

The attributes of this task.

Type:

Dict[str, Union[str, float, int, bool]]

assertions

The assertions of this task. Each assertion is a tuple of an expression and an optional message. The expression is evaluated at runtime. If the expression is not true, the program will abort and print the message.

Type:

List[Tuple[Expr, Optional[str]]]

share_map

The share map. If one output tensor shares memory with one input tensor, it is specified in this map. For example, share_map = {0: 0, 1: 2} means that the output tensor 0 shares the memory with input tensor 0, and output tensor 1 shares the memory with input tensor 2.

Type:

Dict[int, int]

symbols

The list of symbols used in this task.

Type:

List[SymbolVar]

Methods:

generate_arguments(inputs, outputs)	Generate arguments for the compiled function of this task given the tensor parameters.
dummy_arguments(device)	Generate dummy arguments for the compiled function of this task.
build(target[, load])	Build the task for the given target to a callable function.

generate_arguments(inputs, outputs)

Generate arguments for the compiled function of this task given the tensor parameters.

Parameters:

• inputs (Sequence[Tensor]) – The input tensors.

• outputs (Sequence[Tensor]) – The output tensors.

Returns:

args – The arguments for the compiled function.

Return type:

Sequence[Tensor or int]

dummy_arguments(device)

Generate dummy arguments for the compiled function of this task.

Parameters:

device (str) – The target device.

Returns:

args – The arguments for the compiled function.

Return type:

Sequence[Tensor or int]

build(target, load=True)

Build the task for the given target to a callable function.

Parameters:

• target (Union[str, Target]) – The target device.

• load (bool) – Whether to load the task

Returns:

func – The compiled module.

Return type:

hidet.runtime.CompiledTask