Skip to content

Commit

Permalink
Updating all call-sites of the legacy dispatcher registration API in …
Browse files Browse the repository at this point in the history
…fbcode to the new API. (#48178)

Summary:
Pull Request resolved: pytorch/pytorch#48178

I migrated all call sites that used the legacy dispatcher registration API (RegisterOperators()) to use the new API (TORCH_LIBRARY...). I found all call-sites by running `fbgs RegisterOperators()`. This includes several places, including other OSS code (nestedtensor, torchtext, torchvision). A few things to call out:

For simple ops that only had one registered kernel without a dispatch key, I replaced them with:
```
TORCH_LIBRARY_FRAGMENT(ns, m) {
   m.def("opName", fn_name);
}
```

For ops that registered to a specific dispatch key / had multiple kernels registered, I registered the common kernel (math/cpu) directly inside a `TORCH_LIBRARY_FRAGMENT` block, and registered any additional kernels from other files (e.g. cuda) in a separate `TORCH_LIBRARY_IMPL` block.

```
// cpu file
TORCH_LIBRARY_FRAGMENT(ns, m) {
  m.def("opName(schema_inputs) -> schema_outputs");
  m.impl("opName", torch::dispatch(c10::DispatchKey::CPU, TORCH_FN(cpu_kernel)));
}

// cuda file
TORCH_LIBRARY_IMPL(ns, CUDA, m) {
  m.impl("opName", torch::dispatch(c10::DispatchKey::CUDA, TORCH_FN(cuda_kernel)));
}
```
Special cases:

I found a few ops that used a (legacy) `CPUTensorId`/`CUDATensorId` dispatch key. Updated those to use CPU/CUDA- this seems safe because the keys are aliased to one another in `DispatchKey.h`

There were a handful of ops that registered a functor (function class) to the legacy API. As far as I could tell we don't allow this case in the new API, mainly because you can accomplish the same thing more cleanly with lambdas. Rather than delete the class I wrote a wrapper function on top of the class, which I passed to the new API.

There were a handful of ops that were registered only to a CUDA dispatch key. I put them inside a TORCH_LIBRARY_FRAGMENT block, and used a `def()` and `impl()` call like in case two above.

Test Plan: Imported from OSS

Reviewed By: ezyang

Differential Revision: D25056090

Pulled By: bdhirsh

fbshipit-source-id: 8f868b45f545e5da2f21924046e786850eba70d9
  • Loading branch information
bdhirsh authored and facebook-github-bot committed Dec 2, 2020
1 parent 522e59b commit d1686a9
Showing 1 changed file with 11 additions and 12 deletions.
23 changes: 11 additions & 12 deletions torchtext/csrc/register_bindings.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
#include <sentencepiece.h> // @manual
#include <torch/csrc/utils/pybind.h> // @manual
#include <torch/script.h>
#include <torch/library.h> // @manual
#include <vectors.h> // @manual
#include <vocab.h> // @manual

Expand Down Expand Up @@ -174,16 +175,14 @@ static auto vectors =
});

// Registers our custom op with torch.
static auto registry =
torch::RegisterOperators()
.op("torchtext::generate_sp_model", &generate_sp_model)
.op(torch::RegisterOperators::options()
.schema("torchtext::load_sp_model(str path) -> "
"__torch__.torch.classes.torchtext.SentencePiece model")
.catchAllKernel<decltype(load_sp_model), &load_sp_model>())
.op(torch::RegisterOperators::options()
.schema("torchtext::load_sp_model_string(str content) -> "
"__torch__.torch.classes.torchtext.SentencePiece model")
.catchAllKernel<decltype(load_sp_model_string),
&load_sp_model_string>());
TORCH_LIBRARY_FRAGMENT(torchtext, m) {
m.def("generate_sp_model", generate_sp_model);
m.def("load_sp_model(str path) -> "
"__torch__.torch.classes.torchtext.SentencePiece model",
load_sp_model);
m.def("load_sp_model_string(str content) -> "
"__torch__.torch.classes.torchtext.SentencePiece model",
load_sp_model_string);
}

} // namespace torchtext

0 comments on commit d1686a9

Please sign in to comment.