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
|
#include <torch/csrc/jit/mobile/prim_ops_registery.h>
namespace torch::jit::mobile {
static std::unordered_map<std::string, std::function<void(Stack&)>>&
primOpsFnTable() {
static std::unordered_map<std::string, std::function<void(Stack&)>>
prim_ops_fn;
return prim_ops_fn;
}
void registerPrimOpsFunction(
const std::string& name,
const std::function<void(Stack&)>& fn) {
primOpsFnTable()[name] = fn;
}
bool hasPrimOpsFn(const std::string& name) {
return primOpsFnTable().count(name);
}
std::function<void(Stack&)>& getPrimOpsFn(const std::string& name) {
TORCH_CHECK(
hasPrimOpsFn(name),
"Prim Ops Function for ",
name,
" is not promoted yet.");
return primOpsFnTable()[name];
}
} // namespace torch::jit::mobile
|
