tensorflow源码阅读-opkernel注册

//tensorflow/core/framework/kernel_def_builder.h
#define REGISTER_KERNEL_BUILDER(kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ_HELPER(__COUNTER__, kernel_builder, __VA_ARGS__)

#define REGISTER_KERNEL_BUILDER_UNIQ_HELPER(ctr, kernel_builder, ...) \
  REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, __VA_ARGS__)

#define REGISTER_KERNEL_BUILDER_UNIQ(ctr, kernel_builder, ...)        \
  constexpr bool should_register_##ctr##__flag =                      \
      SHOULD_REGISTER_OP_KERNEL(#__VA_ARGS__);                        \
  static ::tensorflow::kernel_factory::OpKernelRegistrar              \
      registrar__body__##ctr##__object(                               \
          should_register_##ctr##__flag                               \
              ? ::tensorflow::register_kernel::kernel_builder.Build() \
              : nullptr,                                              \
          #__VA_ARGS__,                                               \
          [](::tensorflow::OpKernelConstruction* context)             \
              -> ::tensorflow::OpKernel* {                            \
            return new __VA_ARGS__(context);                          \
          });

opkernel的注册， REGISTER_KERNEL_BUILDER（）里面有两个参数。__VA_ARGS__表示不定参数。
REGISTER_KERNEL_BUILDER宏会调用OpKernelRegistrar，它有三个参数。第一个为kernel_builder，返回的是一个KernelDef指针，第二个是opkernel名字，#__VA_ARGS__表示把传入的参数变成字符串，第三个参数是创建opkernel的对象（被包装成Factory ）

//tensorflow/core/framework/op_kernel.h
class OpKernelRegistrar {
public:
 typedef OpKernel* (*Factory)(OpKernelConstruction*);

 OpKernelRegistrar(const KernelDef* kernel_def, StringPiece kernel_class_name,
                   Factory factory) {
   if (kernel_def != nullptr) {
     InitInternal(kernel_def, kernel_class_name, factory);
   }
 }
};

OpKernelRegistrar的第一个参数

kernel_def从REGISTER_KERNEL_BUILDER传入：

//tensorflow/core/kernels/matmul_op.cc
#define REGISTER_CPU_EIGEN(T)                                                  \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), \
      MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>);

Name(“MatMul”).Device(DEVICE_CPU).TypeConstraint(“T”).Label(“eigen”)，它会先调用Name(“MatMul”)，详细如下：

//tensorflow/core/framework/op_kernel.h
class Name : public KernelDefBuilder {
 public:
  explicit Name(const char* op)
      : KernelDefBuilder(SHOULD_REGISTER_OP(op) ? op : "_no_register") {}
};

类Name继承KernelDefBuilder 并调用KernelDefBuilder构造方法:

//tensorflow/core/framework/kernel_def_builder.cc
KernelDefBuilder::KernelDefBuilder(const char* op_name) {
  kernel_def_ = new KernelDef;
  kernel_def_->set_op(op_name);
}

KernelDefBuilder构造方法创建一个KernelDef对象，并设置op_name。这里的KernelDef是通过brotobuf（tensorflow/core/framework/kernel_def.proto）生成的类,里面定义了opkernel的属性。

既然Name继承了KernelDefBuilder，那Name对象具有KernelDefBuilder的所有方法，因此Name(“MatMul”).Device(DEVICE_CPU)则是设置设备的类型,并返回对象本身。接着串联地设置opkernel属性。

//tensorflow/core/framework/kernel_def_builder.cc
KernelDefBuilder& KernelDefBuilder::Device(const char* device_type) {
  kernel_def_->set_device_type(device_type);
  return *this;
}

最后调用Build()方法，将KernelDef指针返回：

//tensorflow/core/framework/kernel_def_builder.cc
const KernelDef* KernelDefBuilder::Build() {
  KernelDef* r = kernel_def_;
  kernel_def_ = nullptr;
  return r;
}

OpKernelRegistrar的第二个参数

//tensorflow/core/kernels/matmul_op.cc
#define REGISTER_CPU_EIGEN(T)                                                  \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("MatMul").Device(DEVICE_CPU).TypeConstraint<T>("T").Label("eigen"), \
      MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>);

REGISTER_KERNEL_BUILDER的第二个参数MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU */>)传入到OpKernelRegistrar，这里的#__VA_ARGS__表示把传入的参数变成字符串，也就是说，第二个参数是opkernel的名字。

OpKernelRegistrar的第三个参数

它是将 MatMulOp<CPUDevice, T, false /* cublas, ignored for CPU />通过lamda表达式，创建一个opkernel对象并返回。 typedef OpKernel (Factory)(OpKernelConstruction);被包装成Factory

opkernel注册

//tensorflow/core/framework/op_kernel.h
class OpKernelRegistrar {
public:
 typedef OpKernel* (*Factory)(OpKernelConstruction*);

 OpKernelRegistrar(const KernelDef* kernel_def, StringPiece kernel_class_name,
                   Factory factory) {
   if (kernel_def != nullptr) {
     InitInternal(kernel_def, kernel_class_name, factory);
   }
 }
};

OpKernelRegistrar 会调用InitInternal，具体如下：

//tensorflow/core/framework/op_kernel.cc
void OpKernelRegistrar::InitInternal(const KernelDef* kernel_def,
                                    StringPiece kernel_class_name,
                                    Factory factory) {
 // See comments in register_kernel::Name in header for info on _no_register.
 if (kernel_def->op() != "_no_register") {
   const string key =
       Key(kernel_def->op(), DeviceType(kernel_def->device_type()),
           kernel_def->label());
   GlobalKernelRegistryTyped()->insert(std::make_pair(
       key, KernelRegistration(*kernel_def, kernel_class_name, factory)));
 }
 delete kernel_def;
}

struct KernelRegistration {
  KernelRegistration(const KernelDef& d, StringPiece c,
                     kernel_factory::OpKernelRegistrar::Factory f)
      : def(d), kernel_class_name(c), factory(f) {}
  const KernelDef def;
  const string kernel_class_name;
  const kernel_factory::OpKernelRegistrar::Factory factory;
};

可以看到OpKernelRegistrar这个类主要是负责根据传进来的KernelDef和KernelFactory，首先依据一定规则生成一个适当的key，并插入到一个全局唯一的Kernel注册表里，注册表当然是一个map但是值得注意的是它是multimap因此支持一个键对应多个kernel副本。

typedef std::unordered_multimap<string, KernelRegistration> KernelRegistry;