Adding a new Op

You can create a custom op if it is not supported yet.

To add a custom op, you need to finish the following steps:

Define the Op class

Define the new Op class in mace/ops/my_custom_op.h.

#ifndef MACE_OPS_MY_CUSTOM_OP_H_
#define MACE_OPS_MY_CUSTOM_OP_H_

#include "mace/core/operator.h"
#include "mace/kernels/my_custom_op.h"

namespace mace {
namespace ops {

template <DeviceType D, typename T>
class MyCustomOp : public Operator<D, T> {
 public:
  MyCustomOp(const OperatorDef &op_def, Workspace *ws)
      : Operator<D, T>(op_def, ws),
        functor_() {}

  bool Run(StatsFuture *future) override {
    const Tensor *input = this->Input(INPUT);
    Tensor *output = this->Output(OUTPUT);
   
    functor_(input, output, future);
    return true;
  }

 protected:
  OP_INPUT_TAGS(INPUT);
  OP_OUTPUT_TAGS(OUTPUT);

 private:
  kernels::MyCustomOpFunctor<D, T> functor_;
};

}  // namespace ops
}  // namespace mace

#endif  // MACE_OPS_MY_CUSTOM_OP_H_

Register the new Op

Define the Ops registering function in mace/ops/my_custom_op.cc.

#include "mace/ops/my_custom_op.h"

namespace mace {
namespace ops {

void Register_My_Custom_Op(OperatorRegistry *op_registry) {
  REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
                                     .Device(DeviceType::CPU)
                                     .TypeConstraint<float>("T")
                                     .Build(),
                    Custom_Op<DeviceType::CPU, float>);

  REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
                                     .Device(DeviceType::OPENCL)
                                     .TypeConstraint<float>("T")
                                     .Build(),
                    Custom_Op<DeviceType::OPENCL, float>);

  REGISTER_OPERATOR(op_registry, OpKeyBuilder("my_custom_op")
                                     .Device(DeviceType::OPENCL)
                                     .TypeConstraint<half>("T")
                                     .Build(),
                    Custom_Op<DeviceType::OPENCL, half>);
}

}  // namespace ops
}  // namespace mace

And then register the new Op in mace/core/operator.cc.

Implement the Op kernel code

You need to implement the CPU kernel in a mace/kernels/my_custom_op.h and optionally OpenCL kernel in mace/kernels/kernels/my_custom_op_opencl.cc and mace/kernels/kernels/cl/my_custom_op.cl. You can also optimize the CPU kernel with NEON.

Add test and benchmark

It's strongly recommended to add unit test and micro benchmark for your new Op. If you wish to contribute back, it's required.

Document the new Op

Finally, add an entry in operator table in the document.