Add New Tensor Function

Tip

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

The steps to add an tensor function are as follows:

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

Add the implementation of the tensor function for the corresponding Backend

CPU

Under the CPU directory, add the subclass of TensorFunction in CPUTensorFunction.hpp.

Here is an example:

#ifndef CPUTENSORFUNCTION_HPP
#define CPUTENSORFUNCTION_HPP
    ... ...

namespace mllm {
class Tensor;

class CPUabcFunction: public TensorFunction {
public:
    void setup(vector<Tensor*> outputs, vector<Tensor*> inputs, vector<float> args) override {
        ... ...
    }
    void execute(vector<Tensor*> outputs, vector<Tensor*> inputs, vector<float> args) override {
        ... ...
    }
};

} // namespace mllm
#endif // CPUTENSORFUNCTION_HPP

Add the enumeration of the tensor function types

Add the enumeration of the tensor function types in the OpDefined.hpp file.

#ifndef MLLM_OPDEFINED_H
#define MLLM_OPDEFINED_H
    ... ...

namespace mllm {
    ... ...

enum TensorFuncType {
    FUNC_ADD,
    FUNC_SUB,
    ... ...
    FUNC_ABC, //<==Add Here
};

}
#endif

Register the implemented tensor function for the corresponding Backend

Add the CPUabcFunction in the corresponding Backend.cpp file.

    ... ...

void CPUBackend::registerFuncs() {
    map_function_[TensorFuncType::FUNC_ADD] = new CPUaddFunction();
    map_function_[TensorFuncType::FUNC_SUB] = new CPUsubFunction();
    ... ...
    map_function_[TensorFuncType::FUNC_ABC] = new CPUabcFunction();//<==Add Here
}

Add the frontend representation of the tensor function

1. Register the corresponding front-end representation of the Tensor function in Tensor.hpp.

Tensor abc(... ...);

2. Implement the corresponding front-end representation of the Tensor function in Tensor.cpp.

Tensor Tensor::abc(... ...);{
    ... ...
};