1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
|
#pragma once
#include <ATen/core/builtin_function.h>
#include <ATen/core/stack.h>
#include <torch/csrc/jit/backends/backend_detail.h>
#include <torch/csrc/jit/backends/backend_interface.h>
#include <torch/custom_class.h>
namespace torch {
namespace jit {
// Static registration API for backends.
template <class TBackendInterface>
class backend {
static_assert(
std::is_base_of<PyTorchBackendInterface, TBackendInterface>::value,
"torch::jit::backend_<T> requires T to inherit from PyTorchBackendInterface");
std::string backend_name_;
public:
explicit backend(const std::string& name) : backend_name_(name) {
static auto cls =
torch::class_<TBackendInterface>(detail::kBackendsNamespace, name)
.def(torch::init<>())
._def_unboxed(
"preprocess",
detail::getPreprocessFunc<TBackendInterface>(),
detail::getPreprocessSchema())
._def_unboxed(
"compile",
detail::getCompileFunc<TBackendInterface>(),
detail::getCompileSchema())
._def_unboxed(
"execute",
detail::getExecuteFunc<TBackendInterface>(),
detail::getExecuteSchema());
}
};
} // namespace jit
} // namespace torch
|
