Concepts
In the current operator implementation model, the complete computation logic covers the host part (usually called the tiling function) and the device part (usually called the kernel). When an operator is being executed, the tiling function on the host is also executed based on the input shape information to generate the tiling data, which is used as the input for the kernel to continue the execution of the device part. On the device, a piece of kernel code is executed concurrently on multiple hardware units (also called cores or blocks). Therefore, each core is responsible for a part of the entire operator computation. Logically, the tiling function is responsible for selecting the most efficient kernel implementation for the current shape and computing a proper core distribution policy for the selected kernel. The result is transferred to the kernel along with the tiling data to guide the kernel to run properly.
After automatic fusion, an AscBackend or FusedAscBackend node is generated in the GE graph. (FusedAscBackend also carries a subgraph object of the ComputeGraph type, and ComputeGraph contains one or more AscBackend nodes.) It is only a shell node generated by the automatic fusion software stack. The real computation logic is stored in the AscGraph attribute of the node. AscGraph uses the directed acyclic graph (DAG) as the basic data structure. AscGraph can completely express the computation logic (including the core distribution policy) of an operator based on the symbolic capability. The Codegen module generates single-core kernel code based on the intra-core computation part in AscGraph. The Auto Tiling module generates the tiling function based on the complete AscGraph. The TilingData definition and InferShape inference function are generated based on the symbol information in AscGraph.
AscGraph has two forms:
- HintGraph: output of the fusion framework, which is used as the backend input and describes only the computation logic of an operator.
- ImplGraph: result after Auto Schedule processing, which is used as the input of the Codegen and Auto Tiling blocks and describes the complete implementation logic of an operator, including all information such as the schedule policies, memory management, and pipeline parallelism.
Although HintGraph and ScheduleResult differ in expression layers and purposes, they share one basic logic for graph building, reuse some fields, and define certain specific fields. The semantics of the same field remain consistent across different forms and do not change with the form.
AscGraph Structure
AscGraph expresses multi-step computation in a multi-layer loop. Each step corresponds to a node in the graph. The node is instantiated by AscIR. The directed edges between nodes indicate the data transfer relationship. Therefore, AscGraph can be considered as a DAG. Each node also contains attributes, which are used to specify the loop layer of the computation.
Instantiation refers to the process of converting AscIR into AscGraph nodes. This is similar to the process of creating a class and then an object. In this process, the input, output, and attributes of AscIR are assigned specific values so that the nodes are executable. For example, if the AscIR type is Cast, the dst_dtype attribute is set during instantiation to specify the target data type that the input data will be converted to.
Attributes are fields defined by AscIR specifications and can be on graphs or nodes. They can be simple values or containers (similar to the struct in the C language) that contain more subfields.
Node
A node is an entity after AscIR instantiation. Its input, output, and attributes comply with the AscIR definition. In AscGraph, each node represents a computation operation and is nested in a loop structure for multiple executions.
- Computation semantics of a node
The complete computation semantics of a node consists of the following parts:
- Node type: indicates the computation type of a node. For example, Add indicates the addition operation.
- Node attribute: indicates the necessary information for supplementing the computation logic. For example, the dst_type attribute of the Cast node specifies the target data type that the input data will be converted to.
- Input and output description: Some nodes require the preceding input and output information to describe the computation logic. For example, the Broadcast node needs to determine the axis on which the broadcast occurs based on the repeats and axis information in the input and output, and the Cast node needs to determine the target data type that the input data will be converted to based on the dtype in the output.
- Node attributes
Each AscGraph node contains the following common attributes that take effect in all forms (HintGraph and ImplGraph):
- name: node name, which is unique in AscGraph. If there are multiple AscGraphs (for example, multiple ImplGraphs in the ScheduleResult phase or multiple AscGraphs in FusedAscBackend in the HintGraph phase), duplicate node names are allowed.
- type: node type, indicating the AscIR type of a node.
- sched: execution approach of the node, including the sched.axis attribute, indicating the loop layer to which the node belongs. For details, see Axis.
- Other dedicated attributes defined by AscIR: A node also contains the IR-specific attributes defined in AscIR.
- Input and output attributes
Each node has certain input and output attributes. AscGraph requires that the input and output attribute values at both ends of one edge are the same. Therefore, in most cases, you only need to pay attention to the output attributes of the node, as the input attributes of the connected nodes must be the same. Below are common attributes applicable to all forms:
- axis: number of axes contained in the output. All axes must reference the axes defined globally in AscGraph.
- repeats: number of data repetitions on each axis, which is the size of the output on each axis.
- stride: index stride of the output on each axis.
- Node execution approach
In AscGraph, nodes are usually nested in a multi-layer loop and are executed multiple times. Assume that the scheduling information about the Foo node is as follows:
1Foo.sched.axis = [z0, z1, z2, z3];
This means that the Foo node will be executed in s0 * s1 * s2 * s3 iterations of the loop. Each loop variable [z0, z1, z2, z3] traverses its corresponding size [s0, s1, s2, s3] in sequence, which is equivalent to the following C++ code:
1 2 3 4 5 6 7 8 9
for (int64_t i0 = 0; i0 < s0; ++i0) { // Traverse the z0 axis. for (int64_t i1 = 0; i1 < s1; ++i1) { // Traverse the z1 axis. for (int64_t i2 = 0; i2 < s2; ++i2) { // Traverse the z2 axis. for (int64_t i3 = 0; i3 < s3; ++i3) { // Traverse the z3 axis. // Perform the Foo computation. } } } }
Axis
Axis is an important concept in AscGraph. AscGraph defines loops and data expressions based on axes. An axis represents a dimension of data and can be identified by name or ID. The core attribute size is used to describe the length of the dimension. The axis and its size symbol are stored as attributes on AscGraph and are defined at the graph level. Therefore, an axis can be referenced by any element (such as a node) in the graph by using the axis ID to implement shared cognition and reference of the same dimension.
Create an axis on the graph. For example, create an axis whose length is s0.
1 2 3 4 5 6 7 | AscGraph graph; // The axis length can be a variable whose name is s0. Expression s0 = graph.CreateSizeVar("s0"); // Create an axis named z0 with a length of s0. The axis ID is assigned during creation and can be obtained through z0.id. Axis &z0 = graph.CreateAxis("z0", s0); |
In the naming convention of automatic fusion, prefix s usually indicates the size variable. For example, s0 and s1 indicate the 0th and 1st size variables, respectively. s is the first letter of symbol, indicating that this is a symbolically expressed variable.
Axis names are usually prefixed with letter z, for example, z0 and z1. The size of z0 is represented by s0, the size of z1 is represented by s1, and so on. This naming method ensures that the mapping between the axis (z) and its size (s) is clear and readable.
Use axes to express a loop. For example, the current AscGraph has four axes:
1 2 3 4 5 6 7 8 9 10 | AscGraph graph; Expression s0 = graph.CreateSizeVar("s0"); Expression s1 = graph.CreateSizeVar("s1"); Expression s2 = graph.CreateSizeVar("s2"); Expression s3 = graph.CreateSizeVar("s3"); Axis &z0 = graph.CreateAxis("z0", s0); Axis &z1 = graph.CreateAxis("z1", s1); Axis &z2 = graph.CreateAxis("z2", s2); Axis &z3 = graph.CreateAxis("z3", s3); |
If the loop consists of z0, z1, z2, and z3, the four axes are traversed fully and sequentially. The equivalent C language expression is as follows:
1 2 3 4 5 6 7 8 | for (int64_t i0 = 0; i0 < s0; ++i0) { // Traverse the z0 axis. The traversal length is the length s0 of the z0 axis. for (int64_t i1 = 0; i1 < s1; ++i1) { // Traverse the z1 axis. for (int64_t i2 = 0; i2 < s2; ++i2) { // Traverse the z2 axis. for (int64_t i3 = 0; i3 < s3; ++i3) { // Traverse the z3 axis. } } } } |
Graph Input and Output
AscGraph uses nodes of specific types to express inputs and outputs. The corresponding AscIRs are as follows:
- Data AscIR indicates the input of a graph.
- Data AscIR involves no input and has only one output, which corresponds to an input of the graph.
- The int32_t index attribute indicates the sequence number of the input in the graph.
- Each data node corresponds to a graph input, and the node connected to the data node is the operator that reads the input.
- Output AscIR indicates the output of a graph.
- Output AscIR involves no output and has only one input, which corresponds to an output of the graph.
- The int index attribute indicates the sequence number of the output in the graph.
- Each output node corresponds to a graph output, and the node connected to the output node is the final output operator of the entire graph.
Unlike the input and output of a common node, the output of Data and the input of Output represent the input and output of AscGraph, respectively. They are the commitments of AscGraph to external systems and cannot be modified during the entire backend running process.