Operator Lowering Implementation

Operator lowering uses the loop IR to express the computation logic of an operator. For example, the computation process of the Add operation is as follows:

Load data from the input tensor.
Broadcast the input data.
Determine the dtype used for computation and convert the input to the required type.
Perform the Add operation on the loaded data.
Write the result to the output tensor of the node.

During graph build on the GE, the node anchor is used to express the tensor input at run time. The anchor can be considered as the placeholder of the tensor input at build time. Below is an example of the pseudocode for Add computation using the Loop IR:

graphStatus LoweringAdd(const ge::Node &node) {
     auto x = loop::Load(node.GetInDataAnchor(0)); // Load data from input anchor 0.
     auto y = loop::Load(node.GetInDataAnchor(1)); // Load data from input anchor 1.
     vector<Expression> broadcasted_shape = xxx; // Compute the output after broadcasting.
     ge::DataType compute_dtype = xxx; // Compute the dtype after type promotion. Note that this step is different from the proactive type promotion during code generation.
     x = loop::Broadcast(x, broadcasted_shape);
     x = loop::Cast(x, compute_dtype);
     y = loop::Broadcast(y, broadcasted_shape);
     y = loop::Cast(y, compute_dtype);
     auto result = loop::Add(x, y); // Express the computation.
     loop::Store(node.GetOutDataAnchor(0), result); // Save the computation result to the output anchor.
 }

The operations of loading inputs, saving outputs, computing broadcast, and type promotion are highly similar. After the common implementation is extracted, the lowering of an operator can be simplified to only the computation part.

Var LoopAdd(const vector<LoopVar> &inputs) { // Var is the type of the intermediate result output by the loop computation.
     return loop::Add(inputs[0], inputs[1]);
 }

Finally, LoopAdd is called during lowering.

graphStatus LoweringAdd(const ge::Node &node) {
     // x, y = common logic, which processes input loading, broadcasting, and type conversion.
     auto result = LoopAdd(x, y); // Express the computation.
     // Common logic, which saves the output result.
 }

Parent topic: Lowering