frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

migrating some straggler pytorch ops in fbcode to the new registration API (#48954)

Browse Source 
Summary:
I already migrated the majority of fbcode ops to the new registration API, but there are a few stragglers (mostly new files that were created in the last two weeks).

The goal is mostly to stamp out as much of the legacy registration API usage as possible, so that people only see the new API when they look around the code for examples of how to register their own ops.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/48954

ghstack-source-id: 118140663

Test Plan: Ran buck targets for each file that I migrated

Reviewed By: ezyang

Differential Revision: D25380422

fbshipit-source-id: 268139a1d7b9ef14c07befdf9e5a31f15b96a48c
This commit is contained in:
Brian Hirsh
2020-12-09 14:40:53 -08:00
committed by Facebook GitHub Bot
parent 67d12c9582
commit c7cc8a48c0
2 changed files with 8 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
benchmarks/operator_benchmark/pt_extension/extension.cpp
View File

@ -17,9 +17,10 @@ List<Tensor> consume_list(List<Tensor> a) {
// That caused an issue for our op benchmark which needs to run an op
// in a loop and report the execution time. This diff resolves that issue by
// registering this consume op with correct alias information which is DEFAULT.
auto reg = torch::RegisterOperators()
.op("operator_benchmark::_consume", &consume)
.op("operator_benchmark::_consume.list", &consume_list);
TORCH_LIBRARY_FRAGMENT(operator_benchmark, m) {
m.def("_consume", &consume);
m.def("_consume.list", &consume_list);
}
PYBIND11_MODULE(cpp_extension, m) {
m.def("_consume", &consume, "consume");

8
torch/csrc/autograd/record_function_ops.cpp
View File

@ -65,10 +65,10 @@ c10::intrusive_ptr<c10::ivalue::Future> _call_end_callbacks_on_fut(
}
// Internal only, do not use directly, use Python's record_function()
static auto registry =
RegisterOperators()
.op("profiler::_record_function_enter", &record_function_enter)
.op("profiler::_record_function_exit", &record_function_exit);
TORCH_LIBRARY_FRAGMENT(profiler, m) {
m.def("_record_function_enter", &record_function_enter);
m.def("_record_function_exit", &record_function_exit);
}
// Needed to register JIT operator in operator registry below
c10::AliasAnalysisKind aliasAnalysisFromSchema() {