xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/mobile/prim_ops_registery.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <torch/csrc/jit/mobile/prim_ops_registery.h>
2 
3 namespace torch::jit::mobile {
4 
5 static std::unordered_map<std::string, std::function<void(Stack&)>>&
primOpsFnTable()6 primOpsFnTable() {
7   static std::unordered_map<std::string, std::function<void(Stack&)>>
8       prim_ops_fn;
9   return prim_ops_fn;
10 }
11 
registerPrimOpsFunction(const std::string & name,const std::function<void (Stack &)> & fn)12 void registerPrimOpsFunction(
13     const std::string& name,
14     const std::function<void(Stack&)>& fn) {
15   primOpsFnTable()[name] = fn;
16 }
17 
hasPrimOpsFn(const std::string & name)18 bool hasPrimOpsFn(const std::string& name) {
19   return primOpsFnTable().count(name);
20 }
21 
getPrimOpsFn(const std::string & name)22 std::function<void(Stack&)>& getPrimOpsFn(const std::string& name) {
23   TORCH_CHECK(
24       hasPrimOpsFn(name),
25       "Prim Ops Function for ",
26       name,
27       " is not promoted yet.");
28   return primOpsFnTable()[name];
29 }
30 
31 } // namespace torch::jit::mobile
32