Switch to pybind11 style registration function API. (#36258)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/36258

Previous we had a && chaining style API.  There are some downsides to
this API:

- It's easy to forget the 'static' qualifier in front, leading to
  subtle ODR bugs.
- It is not compatible with torchbind class_ definitions, as these
  need multiple levels of chaining.  So in practice people end
  up having to define multiple static initializers, one per class.
- It's not like pybind11.
- There's no way to conveniently get the file and line number of
  the registration, as there is no macro point in the API.
- The old API doesn't really encourage people to put all of their
  definitions for a library in one place, and to give a custom
  namespace for it.  Similarly, the old API wasn't very DRY, because
  you had to keep repeating the namespace/dispatch key you
  were writing implementations for.

The new API is modeled exactly off of the PYBIND11_MODULE macro:
you write:

```
TORCH_LIBRARY(aten, m) {
  m.def("aten::add(Tensor self, Tensor other) -> Tensor");
  ...
}
```

in a non-chaining fashion, and under the hood the macro expands to
define a function, and define a static initializer that allocates
c10::Library (previously called c10::Module, but we renamed it
to avoid confusion with the existing NN module concept), passes
it to your function, and then retains it for the rest of the lifetime
of the program.  Specification of the namespace is mandatory,
and in later commit I plan to make it a hard error to TORCH_LIBRARY
the same library name twice.

If you are specifying an implementation for an existing operator
(e.g., you're the XLA backend, or even if you're just putting
registrations for implementations at the implementation site),
you should use TORCH_LIBRARY_IMPL, which instead takes a backend
argument (instead of namespace) and can be used to specify an
implementation for a backend.  Unlike TORCH_LIBRARY, you can do
as many of these as you want for a backend.

This needs updates to the mobile code analyzer.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>

Test Plan: Imported from OSS

Differential Revision: D20929257

Pulled By: ezyang

fbshipit-source-id: ba04d78492e8c93ae7190165fb936f6872896ada

torch/csrc/utils/python_dispatch.cpp

@@ -62,16 +62,16 @@ void initDispatchBindings(PyObject* module) {
     .def("schema", &c10::OperatorHandle::schema);
 
   // TODO: figure out how to do chaining
-  py::class_<c10::Module>(m, "_DispatchModule")
+  py::class_<c10::Library>(m, "_DispatchModule")
     .def("def_", [](py::object self, const char* schema, const char* alias) {
-      self.cast<c10::Module&>().def(torch::schema(schema, parseAliasAnalysisKind(alias)));
+      self.cast<c10::Library&>().def(torch::schema(schema, parseAliasAnalysisKind(alias)));
       return self;
     }, "", py::arg("schema"), py::arg("alias") = "")
     // Simulated "legacy" def where alias analysis kind is not set.
     // Ordinarily this can only be exercised from RegisterOperators() API
     // but I am not going to bind that here
     .def("def_legacy", [](py::object self, const char* schema) {
-      self.cast<c10::Module&>().def(torch::jit::parseSchema(schema));
+      self.cast<c10::Library&>().def(torch::jit::parseSchema(schema));
       return self;
     }, "", py::arg("schema"))
     // We can't conveniently turn Python functions into valid functions
@@ -83,7 +83,7 @@ void initDispatchBindings(PyObject* module) {
     // Mangling scheme: args_rets.  One character per.
     //  t = Tensor
     .def("def_name_t_t", [](py::object self, const char* name, const char* dispatch, const char* debug) {
-      self.cast<c10::Module&>().def(
+      self.cast<c10::Library&>().def(
         name,
         dispatch_str(dispatch, [](const at::Tensor& a) {
           return a;
@@ -94,7 +94,7 @@ void initDispatchBindings(PyObject* module) {
            py::arg("dispatch") = "",
            py::arg("debug") = "default_def_name_t_t")
     .def("def_schema_t_t", [](py::object self, const char* schema, const char* dispatch, const char* alias, const char* debug) {
-      self.cast<c10::Module&>().def(
+      self.cast<c10::Library&>().def(
         torch::schema(schema, parseAliasAnalysisKind(alias)),
         dispatch_str(dispatch, [](const at::Tensor& a) {
           return a;
@@ -108,7 +108,7 @@ void initDispatchBindings(PyObject* module) {
     // TODO: maybe consider deduplicating the definitions here, it's getting
     // pretty long
     .def("impl_t_t", [](py::object self, const char* name, const char* dispatch, const char* debug) {
-      self.cast<c10::Module&>().impl(
+      self.cast<c10::Library&>().impl(
         name,
         dispatch_str(dispatch, [](const at::Tensor& a) {
           return a;
@@ -119,7 +119,7 @@ void initDispatchBindings(PyObject* module) {
            py::arg("dispatch") = "",
            py::arg("debug") = "impl_t_t")
     .def("impl_tt_t", [](py::object self, const char* name, const char* dispatch, const char* debug) {
-      self.cast<c10::Module&>().impl(
+      self.cast<c10::Library&>().impl(
         name,
         dispatch_str(dispatch, [](const at::Tensor& a, const at::Tensor& b) {
           return a;
@@ -130,10 +130,12 @@ void initDispatchBindings(PyObject* module) {
   ;
 
   m.def("_dispatch_import", [](std::string name) {
+    // This is a wee bit dodgy right now, but the "underlying" API is much
+    // easier to test than the high level (using TORCH_LIBRARY, e.g.)
     if (name.empty()) {
-      return std::make_unique<c10::Module>(torch::import());
+      return std::make_unique<c10::Library>(c10::Library::FRAGMENT, "_", c10::DispatchKey::CatchAll, "/dev/null", 0);
     } else {
-      return std::make_unique<c10::Module>(c10::_import_DOES_NOT_WORK_WITH_MOBILE_CUSTOM_BUILD(name));
+      return std::make_unique<c10::Library>(c10::Library::FRAGMENT, name, c10::nullopt, "/dev/null", 0);
     }
   });