tvm.dlight¶

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

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

class tvm.dlight.BlockInfo(name: str, iters: List[tvm.dlight.base.analysis.IterInfo], block_rv: tvm.tir.schedule.schedule.BlockRV, reduction_block: bool = False)¶

Information about a TIR block.

dom() → List[Union[int, tvm.ir.expr.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: tvm.tir.schedule.schedule.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: tvm.tir.expr.Var, dom: tvm.ir.expr.PrimExpr, loop_rv: tvm.tir.schedule.schedule.LoopRV)¶

Information about a loop/iter var.

property dom: Union[int, tvm.ir.expr.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: tvm.tir.function.PrimFunc, target: tvm.target.target.Target, tunable: bool) → Union[None, tvm.tir.schedule.schedule.Schedule, List[tvm.tir.schedule.schedule.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[[tvm.tir.function.PrimFunc, tvm.target.target.Target, bool], Union[None, tvm.tir.schedule.schedule.Schedule, List[tvm.tir.schedule.schedule.Schedule]]]], tvm.dlight.base.schedule_rule.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: tvm.target.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