Add New Operator

Tip

Before customizing the Op, please refer to the Op list document to avoid unnecessary duplication.

The steps to add an Op are as follows:

1. Add the implementation of the Op for the corresponding Backend (CPU).
2. Add the enumeration of the Op types.
3. Register the implemented Op for the corresponding Backend.
4. Add the frontend representation of the Op.

Add the implementation of the Op for the corresponding Backend

CPU

1. Add CPUAbc.hpp and CPUAbc.cpp to the CPU directory, or you can generate the two files by running new_op.py.
2. Class declaration. Inherit from the base class Op and mainly implement methods like reshape() and execute().
3. Implement CPUOpCreator.

Here is an example:

#ifndef MLLM_CPUABC_H
#define MLLM_CPUABC_H

#include "Op.hpp"
#include "CPUBackend.hpp"

namespace mllm {

class CPUAbc final : public Op {
public:
    CPUAbc(Backend *bn, string opName, int param1, bool param2, int threadCount);
    virtual ~CPUAbc() = default;
    // Calculate the size of the outputs Tensor
    virtual ErrorCode reshape(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
    // Perform the calculation and assign values to outputs
    virtual ErrorCode execute(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
    // Load weight parameters (not required to override)
    virtual ErrorCode load(AbstructLoader &loader) override;
    // Free weight parameters, appearing in pairs with load (not required to override)
    virtual ErrorCode free(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;
    // Manage the memory for outputs(inputs) (not required to override)
    virtual ErrorCode setUp(vector<shared_ptr<Tensor>> inputs, vector<shared_ptr<Tensor>> outputs) override;

private:
    // Parameters of Op
    int param1_;
    bool param2_;
    int thread_count = 4;
};

class CPUAbcCreator : public CPUBackend::Creator {
public:
    virtual Op *create(OpParam op_param, Backend *bn, string name, int threadCount) const {
        // Read the parameters of CPUAbc here
        // OpParam is a vector<float>
        int param1 = (int)op_param["param1"];
        bool param2 = (bool)op_param["param2"];
        return new CPUAbc(bn, name, param1, param2, threadCount);
    }
};

} // namespace mllm

#endif // MLLM_CPUABC_H

Add the enumeration of the Op types

Add the enumeration of the Op types in the OpDefined.hpp file.

#ifndef MLLM_OPDEFINED_H
#define MLLM_OPDEFINED_H

#include <string>
#include <vector>
using std::string;
using std::vector;

namespace mllm {
enum OpType {
    INVALID_VALUE = 0,
    PARAMETER,
    ADD,
    ... ...
    ABC, //<==Add here
    OP_NUM
};

static const vector<string> OpNames = {
    "INVALID_VALUE",
    "Parameter",
    "Add",
    ... ...
    "Abc", //<==Add here
    "OP_NUM"};
} // namespace mllm
#endif

Register the implemented Op for the corresponding Backend

1. Include CPUAbc.hpp in the corresponding Backend.cpp file.
2. Add the addCreator function in the corresponding Backend.cpp file.

#include "CPUBackend.hpp"
#include "CPUView.hpp"
    ... ...
#include "CPUAbc.hpp" //<==Add here

    ... ...
    ... ...

void CPUBackend::registerOps() {
    addCreator(PARAMETER, (CPUBackend::Creator *)(new CPUParameterCreator()));
    ... ...
    addCreator(ABC, (CPUBackend::Creator *)(new CPUAbcCreator())); //<==Add here
}

Add the frontend representation of the Op

Implement the frontend representation of the Op in the file Layer.hpp.

class ABC final : public Layer {
public:
    ABC() = default;
    ABC(std::string name) {
        init(std::move(name), OpType::ABC);
    }
    Tensor operator()(Tensor input) {
        auto ts = run({input}, 1);
        return ts[0];
    }
};