LeakyRelu Operator (TBE DSL)

Operator Code Implementation

1. Only four data types (float16, float32, int32, and int8) are supported by the LeakyRelu operator. First, verify the operator type, create a placeholder for the input tensor, call the operator compute function, and perform auto scheduling.

   def leaky_relu_demo(x, y, negative_slope=0, kernel_name="leaky_relu"):
       # check input tensor shape
       shape = x.get("shape")
       dtype = x.get("dtype")
   
       # check input tensor data_type
       check_list = ["float16", "float32", "int32", "int8"]
       if dtype.lower() not in check_list:
           raise RuntimeError(
               "leaky relu only support %s while dtype is %s"
               % (",".join(check_list), dtype))
   
       inp_dtype = dtype.lower()
       input_data_x = tvm.placeholder(shape, name="input_data_x", dtype=inp_dtype)
   
       with tvm.target.cce():
           res = leaky_relu_compute(input_data_x, y, negative_slope, kernel_name)
           sch = generic.auto_schedule(res)

2. Implement the operator compute APIs in the following logic:
   1. When negative_slope = 0, the output y is the larger value between the input x and 0.
      • If the input data type is float16 or int8, the te.lang.cce.vrelu(x) API can be called for compute process.
      • If the input data type is float32 or int32, precision loss may occur if the te.lang.cce.vrelu(x) API is directly called. This is because the te.lang.cce.vrelu(x) API converts float32 and int32 to float16 for compute process. Therefore, you need to compare the input data with tensor_zero (with the same shape as x and 0 elements) by using the te.lang.cce.vmax() API.
   2. When negative_slope is not 0, you need to distinguish the scenarios where negative_slope is less than 1 and negative_slope is large than 1 according to Function Description.

      When negative_slope is less than or equals to 1, the larger value between x and x*negative_slope is used. When negative_slope is large than 1, the smaller value between x and x*negative_slope is used.

      • As required by the te.lang.cce.vmuls() API, the input tensor should be of the same data type of the scalar. Therefore, you need to convert negative_slope to the data type of input x by using the tvm.const API.
      • When the te.lang.cce.vmuls() API is used, data types int8 and int32 will be converted to float16. Therefore, after the te.lang.cce.vmuls() operation is performed, you need to convert the result to the original data type by calling the te.lang.cce.cast_to() API and compares the value of negative_slope with x.

   The compute implementation is as follows:

   @fusion_manager.register("leaky_relu")
   def leaky_relu_demo_compute(x, y, negative_slope=0, kernel_name="leaky_relu"):
       """
       compute for caffe_relu_layer_cce
       """
       inp_dtype = x.dtype.lower()         # Convert all uppercase letters to lowercase letters.
       shape = x.shape                     # Obtain the input shape.
   
       # The original relu logic remains unchanged.
       if negative_slope == 0:
           if inp_dtype in ("float32", "int32"):   # Perform the vmax operation to avoid precision loss.
               tensor_zero = te.lang.cce.broadcast(tvm.const(0, inp_dtype), shape) # Each element is a tensor of inp_dtype and has the value of 0.
               data_res = te.lang.cce.vmax(x, tensor_zero)  # Obtain the larger value between x and tensor_zero.
           else:  # When the data type is float16 or int8,
               data_res = te.lang.cce.vrelu(x)      # Perform the ReLU operation when the data type is float16 or int8.
   
           data_res = te.lang.cce.cast_to(data_res, inp_dtype) # Convert the data type of data_res to inp_dtype.
   
           return data_res
   
       # negative_slope != 0
       if inp_dtype in ("float16", "float32"):
           slope_tmp = tvm.const(negative_slope, dtype=inp_dtype)  # Convert the constant negative_slope to the data type Inp_dtype.
           tmp = te.lang.cce.vmuls(x, slope_tmp)  # Perform the x*slop_tmp operation.
           if negative_slope <= 1:
               res = te.lang.cce.vmax(x, tmp)   # Obtain the larger value between x and tmp when negative_slope ≤ 1.
           else:
               res = te.lang.cce.vmin(x, tmp)   # Obtain the smaller value between x and tmp when negative_slope > 1.
       else:
           # inp_dtype in ("int32", "int8")
           slope_tmp = tvm.const(negative_slope, dtype=inp_dtype)
           tmp = te.lang.cce.vmuls(x, slope_tmp)      # Perform the vmuls operation to convert int8 and int32 to float16.
           tmp_oritype = te.lang.cce.cast_to(tmp, inp_dtype)     # Perform the cast operation to convert the tmp data type to the original data type.
           if negative_slope <= 1:
               res = te.lang.cce.vmax(x, tmp_oritype)
           else:
               res = te.lang.cce.vmin(x, tmp_oritype)
   
           res = te.lang.cce.cast_to(res, inp_dtype)
   
       return res

3. Call the te.lang.cce.cce_build_code() API to perform auto scheduling and build.

      config = {"name": kernel_name,
                 "tensor_list": [input_data_x, res]}
       te.lang.cce.cce_build_code(sch, config)

Operator Plug-in Implementation

You need to customize the ParseParamsLeakyRelu function to implement the attribute mapping from the Caffe LeakyReLUDemo operator to the LeakyReluDemo operator adapted to the Ascend AI Processor.

The ParseParamsLeakyRelu function is implemented as follows:

Status ParseParamsLeakyRelu(const ge::Operator& op_src, ge::Operator& op_dest)
{
  // trans op_src to op_dest
  float negative_slope;
if (ge::GRAPH_SUCCESS == op_src.GetAttr("negative_slope", negative_slope)){// Obtains the value of the negative_slope attribute in the original model.
The op_dest.SetAttr("negative_slope", negative_slope); // Assigns the obtained value of the negative_slope attribute to the negative_slope attribute of the op_dest object.
    }else{
        op_dest.SetAttr("negative_slope", float(0));
    }
  return SUCCESS;
}

Operator Prototype Definition

The key point of prototype definition is to infer the shape and dtype of the output tensor. The LeakyReluDemo operator directly assigns the shape and dtype of the input tensor to the output tensor, as shown in the following:

IMPLEMT_INFERFUNC(LeakyReluDemo, LeakyReluDemoInferShape) {
  auto x_shape = op.GetInputDesc("x").GetShape().GetDims();
  DataType x_dtype = op.GetInputDesc("x").GetDataType();
  TensorDesc y_desc = op.GetOutputDesc("y");
  y_desc.SetShape(ge::Shape(x_shape));
  y_desc.SetDataType(x_dtype);
  (void)op.UpdateOutputDesc("y", y_desc);
  return GRAPH_SUCCESS;
}

Operator Information Definition

For details about the information definition file of the LeakyReluDemo operator, see tbe/op_info_cfg/ai_core/<soc_version>/leaky_relu.ini. The name of the Python file of the operator implementation code (opFile.vaule) and name of the operator definition function (opInterface.vaule) are not configured in the information definition file. Therefore, FE converts uppercase letters of the operator type into underscored lowercase letters to match the operator implementation file and operator definition function. For details about the mapping rules, see 4.