ImplGraph

The Auto Schedule module generates one or more ScheduleResults according to multiple schedule policies when processing the HintGraph. Each ScheduleResult completely expresses the implementation logic of an operator. The overall computation of a ScheduleResult is split into one or more sub-computations. Each sub-computation step is called a ScheduleGroup, which means that the schedule policy is independently applied. In a ScheduleGroup, each schedule policy is used to generate a new AscGraph, which is called an ImplGraph.

Each ImplGraph contains the schedule policy, inter-core and intra-core tiling modes, memory, and vectorization.

Axis Transformation

An axis can be split into two sub-axises, the inner axis and outer axis. For example, after the z0 axis with repeats value s0 is split, two axes z0Outer and z0Inner are generated. The product of the repeats of the two axes is equal to s0. Therefore, the repeats of the two axes can be expressed as ceil(s0/s0I) and s0I.

Axis splitting is classified into block partitioning and tile partitioning.

Block partitioning indicates that data is split and then allocated to multiple cores for parallel execution. Generally, the outer axis and inner axis of block partitioning are named by adding B/b to the end of the original axis name. An ImplGraph can have only one outer axis of the block. For example, if z0 is split by block, the outer axis and inner axis after splitting are named z0B and z0b, respectively.
Tile partitioning is common partitioning and does not involve special semantics.

Two or more consecutive loop axes can be merged into one. Consecutive loop axes refer to consecutive axes in sched.axis. For example, if sched.axis = [z0, z1, z2, z3], [z1, z2, z3] are consecutive axes. If the axis is split, for example, into [z0, z1T, z1t, z2T, z2t, z3], [z1t, z2T] are consecutive axes.

If the loop axis has been reordered, for example, the sequence is changed to [z0, z2, z1, z3], [z0, z2] are consecutive axes. If the axis is reordered after splitting, the reordered axis cannot be further merged. In the preceding example, z1 is split and reordered, and the sequence is changed to [z0, z1T, z2, z1t, z3]. In this case, [z2, z1t] and [z1t, z3] are not consecutive axes and cannot be merged. The sequence of [z0, z1T] is not affected by the reorder and can still be merged.

The following fields are used to express the axis transformation relationship:

axis.type: indicates the axis transformation type. If an axis undergoes multiple transformations, only the last transformation type is saved.
- kAxisTypeOriginal: original axis, which has not been split or merged.
- kAxisTypeBlockOuter: outer axis after block partitioning.
- kAxisTypeBlockInner: inner axis after block partitioning.
- kAxisTypeTileOuter: outer axis after tile partitioning.
- kAxisTypeTileInner: inner axis after tile partitioning.
- kAxisTypeMerged: axis after multiple axes are merged.
from: If the axis is split, the split axis is saved using the from field. If the current axis is merged from multiple axes, the axis before merging is saved using the from field.
split_pair: If the current axis is split, the other axis generated during the splitting is saved using the split_pair field.

Vectorization and Loop Axis

In the ImplGraph, the definition of axis is extended to express three meanings. Consider calling an API:

// There are four axes: z0, z1, z2, z3, with lengths s0, s1, s2, s3 respectively. The index is represented as q + index.
for q0 in range(s0):
    buffer[s1 * s2 * s3];
    for q1 in range(s1):
        // inputs contains four axes: q0, q1, q2, q3. In this call, the data of the q2 and q3 axes is computed.
        buffer[index] = CalcApi(inputs[q0][q1])

The preceding call contains the following information:

Number of loop layers: q0 and q1
Length of data computed at a time: s2 * s3
Length of the data storage buffer after computation: In the example, the buffer contains the z1, z2, and z3 axes. Therefore, the length is the product of the lengths of the three axes.

It is expressed by the following fields.

Location	Attribute Key	Description	Value in This Example
Node	sched.axis	All axes on the node	[z0, z1, z2, z3]
Node	sched.loop_axis	Loop stop axis. The stop axis and the axes before it are the loop part. The part after the stop axis is the data volume computed at a time.	z1
Output	vectorized_axis	Length of the data storage buffer	[z1, z2, z3]
Output	vectorized_stride	Stride for storing data	[s2*s3, s3, 1]

Parent topic: AscGraph