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.
- is_elementwise(sch: tvm.tir.schedule.schedule.Schedule) bool ¶
Whether the block is elementwise, i.e. trivial mapping between read/write region
- 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.
- tvm.dlight.normalize_prim_func(sch: tvm.tir.schedule.schedule.Schedule) Optional[List[tvm.dlight.base.analysis.BlockInfo]] ¶
Normalize the primfunc to normal form
- tvm.dlight.try_inline(sch: tvm.tir.schedule.schedule.Schedule, blocks: List[tvm.dlight.base.analysis.BlockInfo]) List[tvm.dlight.base.analysis.BlockInfo] ¶
Try to inline as many blocks as possible, and return the remaining blocks.
- tvm.dlight.try_inline_contiguous_spatial(sch: tvm.tir.schedule.schedule.Schedule, block_infos: List[tvm.dlight.base.analysis.BlockInfo]) List[tvm.dlight.base.analysis.BlockInfo] ¶
Try to inline contiguous spatial blocks in a schedule