Symbolization Overview

Introduction

Symbolization is to express the shape of an operator through symbols and provide the simplification, inference, and guard functions. The prerequisite for automatic fusion is symbolization. After symbolization, partial operator semantics can be retained to provide key information for subsequent loop axis fusion and memory optimization.

Basis of Symbolization

Currently, the open-source C++ SymEngine library is encapsulated to support symbol expression, computation, and simplification. The advantage of encapsulating this layer is to ensure the stability of use and avoid the binding with open-source modules. In this way, this layer can be flexibly replaced with other third-party libraries or self-developed libraries in the future.

Symbol Generation

In the build phase, before automatic fusion starts, all inputs on a graph are symbolized to generate symbols. After this phase, the shape inputs such as Data and RefData on the graph will contain the generalized symbol shape input. Specific symbolization logic: If an input is greater than 0, it is symbolized as a constant symbol. If an input is -1, it is symbolized as a dynamic symbol. The first -1 is symbolized as s0, the second -1 is symbolized as s1, and so on.

As shown in Figure 1, there are three inputs. The shapes before and after symbolization are as follows:

data0: [3, -1] -> [3, s0];
data1: [-1, 5] -> [s1, 5];
data2: [-1, -1] -> [s2, s3].

Figure 1 Graph structure

Symbol Propagation

Symbol propagation means symbol inference. Each node derives output symbol shapes based on the input symbol shapes. Symbols are generated by data nodes, and then the symbolic output shape of each node is derived along the nodes in the graph until the entire graph is traversed. Symbol propagation is used in the following scenarios:

InferSymbolShape inference (class-1 operator processing): The output symbolic shape is inferred based on the input symbolic shape of the operator.
For example, in the following figure, Concat has two inputs, whose symbolic shapes are data0[s0, s1] and data1[s0, s2], and concat_dim is 1. According to the function of the Concat operator, it can be inferred that the output shape is [s0, (s1 + s2)].

Figure 2 Class-1 operator processing
InferValue inference (class-2 operator processing): symbol folding on the graph.
A class-2 operator can be used to infer output symbols during build. As the output inference is value-dependent, output symbols cannot be inferred only based on the input shape symbols. The output needs to be inferred based on the input values in addition to the input symbols. In the example shown in Figure 3, to infer the symbolic shape of the BroadcastTo operator, you need to know both the output shape and output value of the Pack node.

Figure 3 Class-2 operator processing

Symbols are generated during build. Therefore, for class-2 operators (such as ReShape), constant folding can be extended. By registering the kernel for symbol folding, symbol-based folding operations can be implemented during build to complete symbol-based kernel computation and InferValue logic implementation. The logic of constant folding is as follows:
1. Register the operator's kernel function that supports symbol computations, such as Add, Sub, Mul, and Shape.
2. Traverse the entire graph by referring to constant folding. If specific values can be obtained for all input nodes of a node, call the symbol kernel function corresponding to the current node_type to perform a round of folding computation.
3. After the current node is folded, the symbol value is set to the peer output tensor of the node. After the symbol folding is complete, such nodes can be optimized during build to reduce scheduling and delivery times during execution because these expressions are computed by the host for the shape.
4. Continue the folding computation.
Processing of class-3 and class-4 operators:
Class-3 and class-4 operators refer to operators whose output symbols cannot be derived based on the input shape or symbol value during build. The specific shape can be output only after the execution is complete. In automatic fusion scenarios, graph partitioning and building can be enabled for class-3 and class-4 operators. For details, see the --experimental_enable_jit_executor_v2 control point in AutoFuse Enabling Method.
Processing of customized operators:
1. If the symbol inference process of customized operators is not supplemented, the processing is similar to that of the preceding class-3 and class-4 operators.
2. If the symbol inference process is supplemented for customized operators, operators of class 1 and class 2 are inferred based on typical operators of class 1 and class 2, and the fallback operation is performed for operators of class 3 and class 4.
After build, a symbol expression is generated. In the execution phase, the GE needs to complete shape inference and memory allocation based on the actual Dim value corresponding to the symbol. Therefore, how to compute the actual shape based on the symbol is a key point in the execution phase.

Symbolization includes symbol inference and symbol computation at the function implementation layer. Symbol computation is performed prior to symbol inference. If neither of them exists, static inference is performed. The sequence is as follows:

Symbol computation
The output shape and output value of an operator can be inferred. For example, for the Shape operator, if the input data shape is [s0, 4], the output shape is [2], and the output value is {s0, 4}.
Symbol inference
If the output value of an operator cannot be inferred, the system attempts to infer only the output shape.
Static inference
If the preceding methods cannot be used for inference, check whether the output shape of the operator is static. If it is static, use the static operator as the symbol inference result. If it is not static, no inference can be performed for the operator. In this case, you need to supplement symbol inference and symbol computation functions for the operator.

Parent topic: Symbolization