pytorch

1
#include <c10/util/Exception.h>
2
#include <operator_registry.h>
3

4
namespace torch {
5
namespace executor {
6

7
KernelRegistry& getKernelRegistry() {
8
  static KernelRegistry kernel_registry;
9
  return kernel_registry;
10
}
11

12
bool register_kernels(const ArrayRef<Kernel>& kernels) {
13
  return getKernelRegistry().register_kernels(kernels);
14
}
15

16
bool KernelRegistry::register_kernels(
17
    const ArrayRef<Kernel>& kernels) {
18
  for (const auto& kernel : kernels) {
19
    this->kernels_map_[kernel.name_] = kernel.kernel_;
20
  }
21
  return true;
22
}
23

24
bool hasKernelFn(const char* name) {
25
  return getKernelRegistry().hasKernelFn(name);
26
}
27

28
bool KernelRegistry::hasKernelFn(const char* name) {
29
  auto kernel = this->kernels_map_.find(name);
30
  return kernel != this->kernels_map_.end();
31
}
32

33
KernelFunction& getKernelFn(const char* name) {
34
  return getKernelRegistry().getKernelFn(name);
35
}
36

37
KernelFunction& KernelRegistry::getKernelFn(const char* name) {
38
  auto kernel = this->kernels_map_.find(name);
39
  TORCH_CHECK_MSG(kernel != this->kernels_map_.end(), "Kernel not found!");
40
  return kernel->second;
41
}
42

43

44
} // namespace executor
45
} // namespace torch
46