1 #include <torch/extension.h>
2 #include <torch/script.h>
3
4 using torch::List;
5 using torch::Tensor;
6
consume(Tensor a)7 Tensor consume(Tensor a) {
8 return a;
9 }
10
consume_list(List<Tensor> a)11 List<Tensor> consume_list(List<Tensor> a) {
12 return a;
13 }
14
15 // When JIT tracing is used on function with constant for loop,
16 // the for loop is optimized away because of dead code elimination.
17 // That caused an issue for our op benchmark which needs to run an op
18 // in a loop and report the execution time. This diff resolves that issue by
19 // registering this consume op with correct alias information which is DEFAULT.
TORCH_LIBRARY_FRAGMENT(operator_benchmark,m)20 TORCH_LIBRARY_FRAGMENT(operator_benchmark, m) {
21 m.def("_consume", &consume);
22 m.def("_consume.list", &consume_list);
23 }
24
PYBIND11_MODULE(benchmark_cpp_extension,m)25 PYBIND11_MODULE(benchmark_cpp_extension, m) {
26 m.def("_consume", &consume, "consume");
27 m.def("_consume_list", &consume_list, "consume_list");
28 }
29