xref: /aosp_15_r20/external/pytorch/aten/src/ATen/nnapi/nnapi_register.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <ATen/nnapi/nnapi_bind.h>
2 
3 // Set flag if running on ios
4 #ifdef __APPLE__
5   #include <TargetConditionals.h>
6   #if TARGET_OS_IPHONE
7     #define IS_IOS_NNAPI_BIND
8   #endif
9 #endif
10 
11 #ifndef IS_IOS_NNAPI_BIND
TORCH_LIBRARY(_nnapi,m)12 TORCH_LIBRARY(_nnapi, m) {
13   m.class_<torch::nnapi::bind::NnapiCompilation>("Compilation")
14     .def(torch::jit::init<>())
15     .def("init", &torch::nnapi::bind::NnapiCompilation::init)
16     .def("init2", &torch::nnapi::bind::NnapiCompilation::init2)
17     .def("run", &torch::nnapi::bind::NnapiCompilation::run)
18     ;
19 }
20 #else
21   #undef IS_IOS_NNAPI_BIND
22 #endif
23