build/aten/src/ATen/RegisterSchema.cpp

此文件为编译过程中的代码生成，详细参考，模板为aten/src/ATen/templates/RegisterDispatchKey.cpp

【注意】注册的namespace名字（这里为aten）会作为下面Dispatcher注册时函数名的前缀，且Dispatcher查询时在函数名前也要加上相应的namespace名字才能查询

#define TORCH_LIBRARY(ns, m) \
  static void TORCH_LIBRARY_init_ ## ns (torch::Library&); \
  static const torch::detail::TorchLibraryInit TORCH_LIBRARY_static_init_ ## ns ( \
    torch::Library::DEF, \
    &TORCH_LIBRARY_init_ ## ns, \
    #ns, c10::nullopt, __FILE__, __LINE__ \
  ); \
  void TORCH_LIBRARY_init_ ## ns (torch::Library& m)

Library& def(detail::SelectiveStr<true> raw_schema) & {
    return def(raw_schema.operator const char *());
}

template <typename Schema>
  Library& def(Schema&& raw_schema) & {
    c10::FunctionSchema s = schema(std::forward<Schema>(raw_schema));
    return _def(std::move(s));
}

TORCH_LIBRARY的宏定义，包含3步：

1. 声明静态函数，使用namespace名称区分函数名
2. 构造torch::detail::TorchLibraryInit对象过程中，构造了1个kind为torch::Library::DEF的torch::Library对象 并将其作为参数传给1中定义的函数执行
3. 第1步声明函数的具体定义，下面接上m.def等语句

所以如上两部分结合起来，可以看到实际上是构造了1个torch::Library对象，并调用对象中的def方法逐个注册C++ Native Function

需要注意：

1. 因为是多次调用def方法，所以def方法的返回类型为 Library&
2. def函数的输入是代表 C++ Native Function 的字符串，需要调用schema 函数(及torch::jit::parseSchema方法)，将字符串解析为c10::FunctionSchema 对象，再进行注册，FunctionSchema 的主要成员如下，即为函数声明的各个部分：

   • OperatorName name_;

   • std::vector<Argument> arguments_;

   • std::vector<Argument> returns_;

   • c10::optional<AliasAnalysisKind> alias_kind_;

Library& Library::_def(c10::FunctionSchema&& schema, c10::OperatorName* out_name) & {

......

registrars_.emplace_back(
    c10::Dispatcher::singleton().registerDef(
      std::move(schema),
      debugString("", file_, line_)
    )
  );
  return *this;
}

1. 先做一些检查，确保是第一次注册，否则报错
2. 再调用Dispatcher对象的registerDef方法进行注册

C10_ALWAYS_INLINE static Dispatcher& singleton()

C10_EXPORT Dispatcher& Dispatcher::realSingleton() {
  static Dispatcher _singleton;
  return _singleton;
}

RegistrationHandleRAII Dispatcher::registerDef(FunctionSchema schema, std::string debug) {

  ......

1. 调用findOrRegisterName_函数，查找Dispatcher对象中的 operatorLookupTable_ 是否有指定op_name的OperatorHandle。如果没有，则使用op_name新构造一个；如果有，则取出来。需保证def_count为0，否则代表重复定义，报错
2. op.operatorDef_→op表示取到std::list<OperatorDef> operators_上对应此op的OperatorEntry，调用其registerSchema方法进行注册
3. 成功后将 def_count 和 def_and_impl_count 增加计数，并且返回RegistrationHandleRAII对象，其能够撤销注册

附operatorLookupTable_的定义：

• LeftRight<ska::flat_hash_map<OperatorName, OperatorHandle>> operatorLookupTable_

1. 比较要注册的C++ Native Function schema和 成员kernels_字典所有AnnotatedKernel中的inferred_function_schema的一致性，要么缺省要么相等，否则报错
2. 初始化了OperatorEntry中的2个成员，完成注册定义：
   • c10::optional<AnnotatedSchema> schema_。用于和之后注册实现函数的schema进行比较，相同才注册
   • DispatchKeyExtractor dispatchKeyExtractor_。这个对象中包含1个bitset(64位)，如果C++ Native Function中某个位置上的参数是Tensor类型(或Tensor列表)，则对应的位置上（倒序）设置为1。例如位置0上的参数是Tensor类型，则第63 bit设置为1，以此类推。这样做的目的是记录函数的Tensor类型(或Tensor列表)参数的位置，在dispatch时考虑这些位置上的数据类型，选择合适的C++ Native Function执行具体计算逻辑