Schedule Primitives

class allo.customize.Schedule(module, top_func, func_args, ip, ext_libs=None, use_def_chain=None, inst_list=None)

Methods:

split(axis, factor)	split will find the loop with loop index axis and tile it with each tile size factor The new inner loop will be named axis.inner and the outer loop will be named axis.outer
reorder(*args)	Reorders nested loops with indices listed in args such that the outermost loop is the first index listed in args, the second is the second outermost, and so on.
unroll(axis[, factor])	Unrolls a loop with loop index axis by factor.
fuse(*args)	Combines loops with indices listed in args into a single loop over a single index.
partition(target[, partition_type, dim, factor])	Partitions a given array, for example if the array is B, this would be <schedule>.B.
buffer_at(target, axis)	Creates a chip buffer to hold the values of target written to in loop with index axis instead of immediately writing them to memory.
reshape(target, shape)	Takes an array in the kernel, target, for example if the array is B, then would be target would be <schedule>.B, and reshapes it to tuple shape.
pipeline(axis[, initiation_interval, rewind])	Pipelines a loop with index axis into initiation_interval stages.
parallel(axis)	Instantiates a loop with index axis to be computed in parallel with the loops it is nested with.
inline([axis])	Inlines a function axis.
dataflow(axis)	Applies a "dataflow" attribute to function axis.
compute_at(from_loop, target_loop)	If from_loop and target_loop are indices over the same range, <schedule>.compute_at(from_loop, target_loop) merges the two loops, taking the body of from_loop and appending it to the body of target_loop.
reuse_at(target, axis)	Takes an array in a kernel, for example if the array is B, this would be <schedule>.B, accessed by index axis and creates a reuse buffer to reuse values from target which are accessed in a sequentially moving window.
to(target, dst[, axis, depth])	Takes an array in the kernel, target, for example if the array is B, this would be target would be <schedule>.B, and converts it into a stream.
unfold(band_name, axes)	Finds a set of nested loops with name band_name and for every <i> in list axes.
compose(schs[, id, instantiate])	Uses schs, a schedule for a kernel called in this kernel, in this kernel.

split(axis, factor)

split will find the loop with loop index axis and tile it with each tile size factor The new inner loop will be named axis.inner and the outer loop will be named axis.outer

Parameters:

• axis (str) – The name of an index in the kernel.

• factor (int) – The size of each tile, e.g. the size of the inner nested loop.

reorder(*args)

Reorders nested loops with indices listed in args such that the outermost loop is the first index listed in args, the second is the second outermost, and so on.

This function is vardic, accepting each index as a separate argument.

unroll(axis, factor=0)

Unrolls a loop with loop index axis by factor.

Parameters:

• axis (str) – The name of an index in the kernel.

• factor (int) – The factor to unroll by, for example a factor of 2 will cause the body to be duplicated once.

fuse(*args)

Combines loops with indices listed in args into a single loop over a single index.

This function is vardic, accepting each index as a separate argument.

partition(target, partition_type=0, dim=0, factor=0)

Partitions a given array, for example if the array is B, this would be <schedule>.B. There are three types, Partition.Complete, Partition.Block, and Partition.cyclic. block: The original array is split into factor equally sized blocks of consecutive elements of the original array cyclic:The original array is split into factor equally sized blocks interleaving the elements of the original array. complete: The original array is split into its individual elements. This corresponds to resolving a memory into registers.

Parameters:

• target (allo.ir.utils.MockBuffer) – The array to partition.

• partition_type (allo.customize.Partition) – The type of partition.

• factor (int) – The number of arrays created by a block or cyclic partition.

• dim (int) – The dimension of target to partition. If dim=0, all dimensions are partitioned.

buffer_at(target, axis)

Creates a chip buffer to hold the values of target written to in loop with index axis instead of immediately writing them to memory.

Parameters:

• target (allo.ir.utils.MockBuffer) – An array written to in a loop.

• axis (str) – The loop index whose body contains writes to target

reshape(target, shape)

Takes an array in the kernel, target, for example if the array is B, then would be target would be <schedule>.B, and reshapes it to tuple shape. As an example, if the desired shape is 32 by 4 by 8, the <shape> would be (32, 4, 8).

Parameters:

• target (allo.ir.utils.MockBuffer) – The array, represented by a memory, to reshape.

• shape (tuple) – The new shape of the memory.

pipeline(axis, initiation_interval=1, rewind=False)

Pipelines a loop with index axis into initiation_interval stages.

Parameters:

• axis (str) – The index of the loop to pipeline.

• initiation_interval (int) – The initiation_interval to be used when pipelining.

• rewind (bool) – If true, rewinding is allowed, allowing continuous loop pipelining. This is only effective for perfect loop nests inside a top level function.

parallel(axis)

Instantiates a loop with index axis to be computed in parallel with the loops it is nested with.

Parameters:

axis (str) – The index of the loop to be computed in parallel.

inline(axis=None)

Inlines a function axis.

Parameters:

axis (str) – The function to inline.

dataflow(axis)

Applies a “dataflow” attribute to function axis. This allows for parallelism if the given function uses streams or the to schedule.

Parameters:

axis (str | allo.ir.LoopWrapper) – The function to add the attribute to.

compute_at(from_loop, target_loop)

If from_loop and target_loop are indices over the same range, <schedule>.compute_at(from_loop, target_loop) merges the two loops, taking the body of from_loop and appending it to the body of target_loop.

Parameters:

• from_loop (str) – The loop whose body is being moved.

• target_loop (str) – The loop whose body is being appended to.

reuse_at(target, axis)

Takes an array in a kernel, for example if the array is B, this would be <schedule>.B, accessed by index axis and creates a reuse buffer to reuse values from target which are accessed in a sequentially moving window.

Parameters:

• target (allo.ir.utils.MockBuffer) – The array being accessed.

• axis (str) – The loop index used to access values in target

to(target, dst, axis=None, depth=-1)

Takes an array in the kernel, target, for example if the array is B, this would be target would be <schedule>.B, and converts it into a stream. dst is the name of the array any value of target is written to. For example if C[i, j] = B[i, j], dst would be specified as “C”. If values of <target> get written to multiple arrays. Multiple calls to <schedule>.to(…) may be needed.

Parameters:

• target (allo.ir.utils.MockBuffer) – The array to convert to a stream.

• dst (str) – An array which a value of target is written to.

• axis (str) – Move axis-th loop body to xcel scope.

• depth (int) – The streaming channel depth.

unfold(band_name, axes)

Finds a set of nested loops with name band_name and for every <i> in list axes. The <i>th nested loop is unfolded into a constant number of copies of it’s loop body.

Parameters:

• band_name (str) – The set of nested loops to unroll.

• axes (list[int]) – A list of the axes to unroll.

compose(schs, id=None, instantiate=None)

Uses schs, a schedule for a kernel called in this kernel, in this kernel.

A kernel, <k1>, may call another kernel, <k2>. This means the output of <k1>.customize() will contain the MLIR for the compiled <k2>, <s2’>. <s2’> will not have any custom schedule. To use a custom schedule, <s2>, the compiled <k2> with some schedule can be created. This is inserted into the schedule for this kernel through self.compose(<s2>).

Parameters:

• schs (allo.customize.Schedule) – The schedule of a kernel used in self.

• id (str) – Identifies the schedule to replace contained in self. This schedule in self must be annotated if id is specified.

• instantiate (list) – This is a list of objects used to instantiate types schs is generic over.

Data Types

allo.ir.types

alias of <module ‘allo.ir.types’ from ‘/__w/allo/allo/allo/ir/types.py’>