tvm.dlight

DLight package provides efficient schedules out-of-box for deep learning workloads.

class tvm.dlight.ApplyDefaultSchedule(*rules: ScheduleRule): A IRModule pass that applies a list of ScheduleRules to all PrimFuncs in the module.

class tvm.dlight.BlockInfo(name: str, iters: List[IterInfo], block_rv: BlockRV, reduction_block: bool = False)

Information about a TIR block.

dom() → List[int | PrimExpr]: The iteration domain of the block.

dom_kind() → str: The iteration domain kind of the block, for example, SSSS, SSSR.

is_injective() → bool: Whether the block is injective, i.e. all its iteration domains are injective.

is_elementwise(sch: Schedule) → bool: Whether the block is elementwise, i.e. trivial mapping between read/write region

is_reduction() → bool: Whether the block is a reduction workload.

is_gemv() → bool: Whether the block is a GEMV workload.

is_gemm() → bool: Whether the block is a GEMM workload.

class tvm.dlight.IterInfo(kind: typing_extensions.Literal[S, R, O], var: Var, dom: PrimExpr, loop_rv: LoopRV)

Information about a loop/iter var.

property dom: int | PrimExpr: The iteration domain of the loop.

class tvm.dlight.ScheduleRule

A thin wrapper on an arbitrary function that can be used to schedule a TIR PrimFunc.

Given a PrimFunc, a target, and a tunable flag, the apply method of a ScheduleRule returns either a Schedule, a list of Schedules, or None, where None means that the rule is not applicable to the given PrimFunc. If the tunable flag is True, the ScheduleRule is allowed to return either a Schedule or a list of Schedules, and the Schedules are allowed to contain tunable instructions. If the tunable flag is False, the ScheduleRule is only allowed to return a Schedule, and the Schedule is not allowed to contain tunable instructions.

apply(func: PrimFunc, target: Target, tunable: bool) → None | Schedule | List[Schedule]

Apply the ScheduleRule to the given PrimFunc.

Parameters:

func (tir.PrimFunc) – The PrimFunc to apply the ScheduleRule to.
target (Target) – The compilation target the schedule is supposed to be built for.
tunable (bool) – Whether the schedule is allowed to contain tunable instructions.

Returns:

results – Either a Schedule, a list of Schedules, or None, where None means that the rule is not applicable to the given PrimFunc.

Return type:

Union[None, tir.Schedule, List[tir.Schedule]]

static from_callable(name) → Callable[[Callable[[PrimFunc, Target, bool], None | Schedule | List[Schedule]]], ScheduleRule]

Create a ScheduleRule from a callable.

Parameters:: name (str)
Returns:: decorator – A decorator that takes a callable and returns a ScheduleRule.
Return type:: Callable

Examples

@ScheduleRule.from_callable("MyRule")
def my_rule(func: tir.PrimFunc, target: Target, tunable: bool) -> Union[None, Schedule]
    # Do something with func and target

is_target_available(target: Target) → bool

Check whether the rule is available for the given target.

Parameters:: target (Target) – The compilation target the schedule is supposed to be built for.
Returns:: available – Whether the rule is available for the given target.
Return type:: bool

tvm.dlight.normalize_prim_func(sch: Schedule) → List[BlockInfo] | None: Normalize the primfunc to normal form

tvm.dlight.try_inline(sch: Schedule, blocks: List[BlockInfo]) → List[BlockInfo]

Try to inline as many blocks as possible, and return the remaining blocks.

Parameters:

sch (tir.Schedule) – The TIR schedule used to inline blocks.
blocks (List[BlockInfo]) – The blocks to be inlined.

Returns:

remaining – The remaining blocks that cannot be inlined.

Return type:

List[BlockInfo]

tvm.dlight.try_inline_contiguous_spatial(sch: Schedule, block_infos: List[BlockInfo]) → List[BlockInfo]

Try to inline contiguous spatial blocks in a schedule

Parameters:

sch (tir.Schedule) – The TIR schedule used to inline blocks.
block_infos (List[BlockInfo]) – The blocks to be try.

Returns:

remaining – The remaining blocks that cannot be inlined.

Return type:

List[BlockInfo]