Add Operator (TBE DSL)

Operator Code Implementation

The Add operator supports only three data types: float16, float32, and int32. Therefore, the input data type needs to be verified. The two inputs may have different shapes. This scenario is supported by the Add operator, but not supported by the operator compute API te.lang.cce.vadd(). As a result, the two input shapes need to be broadcast and verified. The operator implementation code is as follows:

tbe/impl/add.py

SHAPE_SIZE_LIMIT = 2147483648
# Compare each dimension of the two input shapes and assign the larger value of each dimension to generate out_shape.
def _produce_shapes(shape1, shape2):
    shape1 = list(shape1)
    shape2 = list(shape2)
    flag = 0
    if len(shape1) < len(shape2):
        shape1, shape2 = shape2, shape1
        flag = 1

    output_shape_len = len(shape1)
    dec = output_shape_len - len(shape2)
    for i in range(dec):
        shape2 = [1] + shape2

    out_shape = []
    for i in range(output_shape_len):
        if (shape1[i] != shape2[i]) and (shape1[i] != 1) and (shape2[i] != 1):
            raise RuntimeError("input shapes not match!")
        out_shape.append(shape1[i] if shape1[i] > shape2[i] else shape2[i])

    if flag == 1:
        shape1, shape2 = shape2, shape1

    return shape1, shape2, out_shape

# Convert the shape to a list.
def _shape_to_list(shape):
    result = []
    for i in shape:
        if isinstance(i, tvm.expr.Var):
            result.append(i)
        else:
            result.append(i.value)
    return result

# Implement the computation logic of the Add operator.
@fusion_manager.register("add")
def add_compute(input_x, input_y, output_z, kernel_name="add"):
    shape_x = _shape_to_list(input_x.shape)
    shape_y = _shape_to_list(input_y.shape)
    shape_x, shape_y, shape_max = _produce_shapes(shape_x, shape_y)   # Assign the larger value of each dimension of shape_x and shape_y to shape_max.
    shape_size = reduce(lambda x, y: x * y, shape_max[:])      
    if shape_size > SHAPE_SIZE_LIMIT:
        raise RuntimeError("the shape is too large to calculate")

    input_x = te.lang.cce.broadcast(input_x, shape_max)       # Broadcast the shape of input_x as shape_max.
    input_y = te.lang.cce.broadcast(input_y, shape_max)       # Broadcast the shape of input_y as shape_max.
    res = te.lang.cce.vadd(input_x, input_y)        # Execute input_x + input_y.

    return res          # Return the output tensor.

# Operator definition function
def add(input_x, input_y, output_z, kernel_name="add"):
    # Obtain the shape and data type of the operator input tensor.
    shape_x = input_x.get("shape")      
    shape_y = input_y.get("shape")
    check_tuple = ("float16", "float32", "int32")
    input_data_type = input_x.get("dtype").lower()
    if input_data_type not in check_tuple:
        raise RuntimeError("only support %s while dtype is %s" %
                           (",".join(check_tuple), input_data_type))
    
# Broadcast shape_x and shape_y, which is a preparation for tensor placeholder.
    shape_x, shape_y, shape_max = _produce_shapes(shape_x, shape_y)  
    if shape_x[-1] == 1 and shape_y[-1] == 1 and shape_max[-1] == 1: 
       # If the shape length is 1, assign a value directly. If the shape length is not 1 and the last dimension is 1, reduce the last dimension. For the scenario where the last dimension = 1 and where the last dimension does not exist, the format is the same. For example, 2*3 = 2*3*1. Reducing the last dimension = 1 can improve the subsequent scheduling efficiency.
        shape_x = shape_x if len(shape_x) == 1 else shape_x[:-1]   
        shape_y = shape_y if len(shape_y) == 1 else shape_y[:-1]
        shape_max = shape_max if len(shape_max) == 1 else shape_max[:-1]

  
    # Call the placeholder API of TVM to place the first input tensor, returning a tensor object.
    data_x = tvm.placeholder(shape_x, name="data_1", dtype=input_data_type)
    # Call the placeholder API of TVM to place the second input tensor, returning a tensor object.
    data_y = tvm.placeholder(shape_y, name="data_2", dtype=input_data_type)

    # Call the compute implementation function.
    res = add_compute(data_x, data_y, output_z, kernel_name)  
    # Auto scheduling
    with tvm.target.cce():
        schedule = generic.auto_schedule(res)
    # Build configuration
    config = {"name": kernel_name,
              "tensor_list": (data_x, data_y, res)}
    te.lang.cce.cce_build_code(schedule, config)

Operator Plug-In Implementation

Parse and map the Add operator developed under the TensorFlow framework to the Add operator supported by the Ascend AI Processor. You can implement the operator attribute mapping by calling AutoMappingFn(). For details about the complete code, see the framework/tf_plugin/add_plugin.cpp folder in the sample.

Operator Prototype Definition

The key point of prototype definition is to infer the shape of the output tensor and verify the internal association of the operator inputs.

The principle of inferring the output shape is as follows: Obtain the two input shapes, broadcast them to the same shape by assigning the larger size of each dimension of the two inputs. The implementation code is as follows:

op_proto/add.cpp

bool InferShapeAndTypeAdd(Operator& op, const string& input_name1, const string& input_name2, const string& output_name) {
  // vOutputDesc.push_back(op.GetInputDesc(0));
  TensorDesc vOutputDesc = op.GetOutputDesc(output_name);

  DataType input_dtype = op.GetInputDesc(input_name1).GetDataType();
  Format input_format = op.GetInputDesc(input_name1).GetFormat();
  // Exchange the shape dimensions.
  ge::Shape shapeX = op.GetInputDesc(input_name1).GetShape();
  ge::Shape shapeY = op.GetInputDesc(input_name2).GetShape();
  std::vector<int64_t> dimsX = shapeX.GetDims();
  std::vector<int64_t> dimsY = shapeY.GetDims();
  if (dimsX.size() < dimsY.size()) {
    std::vector<int64_t> dimsTmp = dimsX;
    dimsX = dimsY;
    dimsY = dimsTmp;
  }

  // Pad the smaller shape dimension with 1.
  if (dimsX.size() != dimsY.size()) {
    int dec = dimsX.size() - dimsY.size();
    for (int i = 0; i < dec; i++) {
      dimsY.insert(dimsY.begin(), (int64_t)1);
    }
  }

  // Use the larger value of each dimension in the two input shapes to form the output shape dimensions.
  std::vector<int64_t> dimVec;
  for (size_t i = 0; i < dimsX.size(); i++) {
    if ((dimsX[i] != dimsY[i]) && (dimsX[i] != 1) && (dimsY[i] != 1)) {
      return false;
    }

    int64_t dims = dimsX[i] > dimsY[i] ? dimsX[i] : dimsY[i];
    dimVec.push_back(dims);
  }
  ge::Shape outputShape = ge::Shape(dimVec);

  vOutputDesc.SetShape(outputShape);
  vOutputDesc.SetDataType(input_dtype);
  vOutputDesc.SetFormat(input_format);
  op.UpdateOutputDesc(output_name, vOutputDesc);

  return true;
}

Operator Information Definition

For details about the information definition file of the Add operator, see tbe/op_info_cfg/ai_core/<soc_version>/add.ini.