Documentation

Operator Code Implementation (AI CPU)

Overview

The implementation of an AI CPU operator consists of:

  • Header file (.h) that declares the operator class. The custom operator class inherits the CpuKernel base class.
  • Source file (.cc) that has rewritten the Compute function in the operator class to implement the operator compute logic.

Header File Code Module

Declare the operator class in the cpukernel/impl/reshape_cust_kernel.h file of the operator project. The code module is as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
#ifndef _AICPU_RESHAPE_CUST_KERNELS_H_
#define _AICPU_RESHAPE_CUST_KERNELS_H_

 #include "cpu_kernel.h"        // Defines CpuKernel base class and registration macro.
// Include the following header files as needed:
 #include "cpu_tensor.h"        // Defines tensor and related methods.
#include "cpu_tensor_shape.h"  // Defines tensor shape and related methods.
#include "cpu_types.h"         // Defines data types and formats.
#include "cpu_attr_value.h"    // Defines AttrValue and related methods.

namespace aicpu {        // Defines the namespace aicpu.
class ReshapeCustCpuKernel : public CpuKernel {        // Specifies that the ReshapeCust operator class inherits the CpuKernel base class.
public:
    ~ReshapeCustCpuKernel() = default;
    virtual uint32_t Compute(CpuKernelContext &ctx) override;    // Declares the Compute function, which has been rewritten.
};
} // namespace aicpu
#endif
  • Include related header files.
    • cpu_kernel.h defines the AI CPU operator base class CpuKernel and the kernel registration macro.
    • cpu_tensor.h defines the AI CPU Tensor class and related methods.
    • cpu_tensor_shape.h defines the AI CPU TensorShape class and related methods.
    • cpu_types.h defines the AI CPU data types and formats.
    • cpu_attr_value.h defines the AttrValue class and related methods.
  • Declare the operator class, which is derived from the CpuKernel class. Declare the overloading function Compute, which needs to be implemented in the operator implementation file. For details, see Source File Code Module. The operator class must be defined within the namespace aicpu. The namespace name is fixed to aicpu and cannot be modified.

Source File Code Module

Implement the operator compute logic in the cpukernel/impl/reshape_cust_kernel.cc file of the operator project. The code module is as follows:
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
#include "reshape_cust_kernel.h"     // Includes the header files declaring the ReshapeCust operator class.

namespace {          
const char *RESHAPE_CUST = "ReshapeCust";      //The operator OpType is ReshapeCust class.
}

namespace aicpu {    // Defines the namespace aicpu.
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)  // Implements the Compute function of the custom operator class.
{
    ...
    return 0;
}

REGISTER_CPU_KERNEL(RESHAPE_CUST, ReshapeCustCpuKernel);  // Registers the ReshapeCust Operator.
} // namespace aicpu
  1. Include related header files.

    The header file reshape_cust_kernel.h is declared in Header File Code Module.

    If the following header files have been included in Header File Code Module, you do not need to include them again.

    • cpu_kernel.h defines the AI CPU operator base class CpuKernel and the kernel registration macro.
    • cpu_tensor.h defines the AI CPU Tensor class and related methods.
    • cpu_tensor_shape.h defines the AI CPU TensorShape class and related methods.
    • cpu_types.h defines the AI CPU data types and formats.
    • cpu_attr_value.h defines the AttrValue class and related methods.
  2. Define the namespace and declare that the constant pointer points to the operator OpType.

    See the following example.

    namespace {          
const char *RESHAPE_CUST = "ReshapeCust";
}

    ReshapeCust is the operator OpType, and RESHAPE_CUST is the declared constant pointer that points to the operator OpType.

  3. Define the namespace aicpu, implement the Compute function of custom operator within the namespace aicpu, and define the compute logic of the operator.
    The namespace name is fixed to aicpu. The base class and related definitions are all within the aicpu namespace.
    • Declare the Compute function.
      uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)

      ReshapeCustCpuKernel is the custom operator class defined in the header file. The formal parameter CpuKernelContext is the context of the CPU kernel, including the input and output tensors and attributes of the operator.

    • According to your operator development requirements, edit the Compute function body to obtain the input tensors, organize the compute logic based on the inputs, obtain the output result, and set the output into the output tensor.
  4. Register the kernel implementation of the operator:
    REGISTER_CPU_KERNEL(RESHAPE_CUST, ReshapeCustCpuKernel);
    • RESHAPE_CUST is a string pointer that points to the operator OpType defined in 2.
    • ReshapeCustCpuKernel is the name of the custom operator class.

Compute Function Implementation

The key code for implementing an AI CPU operator is the implementation code of the Compute function. The ReshapeCust operator copies the input tensor data to the output tensor. The shape information of the output tensor is inferred by InferShape defined in the operator prototype.

A sample code of operator implementation is as follows:
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
namespace aicpu {
uint32_t ReshapeCustCpuKernel::Compute(CpuKernelContext &ctx)
{
    Tensor *inputTensor = ctx.Input(0);    // Obtain information such as the input shape and data based on the obtained input_tensor.
    if (inputTensor == nullptr) {
        return -1;
    }

    Tensor *outputTensor = ctx.Output(0);    // Obtain information such as the output shape and data based on the obtained output_tensor.
    if (outputTensor == nullptr) {
        return -1;
    }

    // Obtain the data address of the input tensor.
    auto inputData = inputTensor->GetData();
    if (inputData == nullptr) {
        return -1;
    }

    // Obtain the data address of the output tensor.
    auto outputData = outputTensor->GetData();
    if (outputData == nullptr) {
        return -1;
    }

    // Copy the data of the input tensor to the output tensor.
    uint64_t inputDataSize = inputTensor->GetDataSize();
    memcpy(outputData, inputData, inputDataSize);
    return 0;
}