1#!/usr/bin/env python3 2""" 3Code generator for NNAPI wrapper. We can't link directly against 4libneuralnetworks.so because we want PyTorch to work on Android 5devices that don't have it available. Instead, we generate a wrapper 6that opens libneuralnetworks.so with dlopen and finds the functions 7we need with dlsym. We also generate a "check" wrapper that checks 8return values and throws C++ exceptions on errors. 9""" 10 11import re 12import sys 13import textwrap 14from pathlib import Path 15 16 17PREFIX = """\ 18/** 19 * Copyright (c) Facebook, Inc. and its affiliates. 20 * 21 * Licensed under the Apache License, Version 2.0 (the "License"); 22 * you may not use this file except in compliance with the License. 23 * You may obtain a copy of the License at 24 * 25 * http://www.apache.org/licenses/LICENSE-2.0 26 * 27 * Unless required by applicable law or agreed to in writing, software 28 * distributed under the License is distributed on an "AS IS" BASIS, 29 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 30 * See the License for the specific language governing permissions and 31 * limitations under the License. 32 */ 33 34// This file is generated by nnapi/codegen.py 35""" 36 37 38NNAPI_FUNCTIONS = [ 39 ("int", "ANeuralNetworks_getDeviceCount", "uint32_t* numDevices"), # noqa: B950 40 ( 41 "int", 42 "ANeuralNetworks_getDevice", 43 "uint32_t devIndex, ANeuralNetworksDevice** device", 44 ), # noqa: B950 45 ( 46 "int", 47 "ANeuralNetworksDevice_getName", 48 "const ANeuralNetworksDevice* device, const char** name", 49 ), # noqa: B950 50 ( 51 "int", 52 "ANeuralNetworksDevice_getVersion", 53 "const ANeuralNetworksDevice* device, const char** version", 54 ), # noqa: B950 55 ( 56 "int", 57 "ANeuralNetworksDevice_getFeatureLevel", 58 "const ANeuralNetworksDevice* device, int64_t* featureLevel", 59 ), # noqa: B950 60 ( 61 "int", 62 "ANeuralNetworksModel_getSupportedOperationsForDevices", 63 " const ANeuralNetworksModel* model, const ANeuralNetworksDevice* const* devices, uint32_t numDevices, bool* supportedOps", 64 ), # noqa: B950 65 ( 66 "int", 67 "ANeuralNetworksCompilation_createForDevices", 68 "ANeuralNetworksModel* model, const ANeuralNetworksDevice* const* devices, uint32_t numDevices, ANeuralNetworksCompilation** compilation", # noqa: B950 69 ), 70 ( 71 "int", 72 "ANeuralNetworksExecution_compute", 73 "ANeuralNetworksExecution* execution", 74 ), # noqa: B950 75 ( 76 "int", 77 "ANeuralNetworksMemory_createFromFd", 78 "size_t size, int protect, int fd, size_t offset, ANeuralNetworksMemory** memory", 79 ), # noqa: B950 80 ( 81 "void", 82 "ANeuralNetworksMemory_free", 83 "ANeuralNetworksMemory* memory", 84 ), # noqa: B950 85 ( 86 "int", 87 "ANeuralNetworksModel_create", 88 "ANeuralNetworksModel** model", 89 ), # noqa: B950 90 ("void", "ANeuralNetworksModel_free", "ANeuralNetworksModel* model"), # noqa: B950 91 ("int", "ANeuralNetworksModel_finish", "ANeuralNetworksModel* model"), # noqa: B950 92 ( 93 "int", 94 "ANeuralNetworksModel_addOperand", 95 "ANeuralNetworksModel* model, const ANeuralNetworksOperandType* type", 96 ), # noqa: B950 97 ( 98 "int", 99 "ANeuralNetworksModel_setOperandValue", 100 "ANeuralNetworksModel* model, int32_t index, const void* buffer, size_t length", 101 ), # noqa: B950 102 ( 103 "int", 104 "ANeuralNetworksModel_setOperandValueFromMemory", 105 "ANeuralNetworksModel* model, int32_t index, const ANeuralNetworksMemory* memory, size_t offset, size_t length", 106 ), # noqa: B950 107 ( 108 "int", 109 "ANeuralNetworksModel_addOperation", 110 "ANeuralNetworksModel* model, ANeuralNetworksOperationType type, uint32_t inputCount, const uint32_t* inputs, uint32_t outputCount, const uint32_t* outputs", # noqa: B950 111 ), 112 ( 113 "int", 114 "ANeuralNetworksModel_identifyInputsAndOutputs", 115 "ANeuralNetworksModel* model, uint32_t inputCount, const uint32_t* inputs, uint32_t outputCount, const uint32_t* outputs", 116 ), # noqa: B950 117 ( 118 "int", 119 "ANeuralNetworksModel_relaxComputationFloat32toFloat16", 120 "ANeuralNetworksModel* model, bool allow", 121 ), # noqa: B950 122 ( 123 "int", 124 "ANeuralNetworksCompilation_create", 125 "ANeuralNetworksModel* model, ANeuralNetworksCompilation** compilation", 126 ), # noqa: B950 127 ( 128 "void", 129 "ANeuralNetworksCompilation_free", 130 "ANeuralNetworksCompilation* compilation", 131 ), # noqa: B950 132 ( 133 "int", 134 "ANeuralNetworksCompilation_setPreference", 135 "ANeuralNetworksCompilation* compilation, int32_t preference", 136 ), # noqa: B950 137 ( 138 "int", 139 "ANeuralNetworksCompilation_finish", 140 "ANeuralNetworksCompilation* compilation", 141 ), # noqa: B950 142 ( 143 "int", 144 "ANeuralNetworksExecution_create", 145 "ANeuralNetworksCompilation* compilation, ANeuralNetworksExecution** execution", 146 ), # noqa: B950 147 ( 148 "void", 149 "ANeuralNetworksExecution_free", 150 "ANeuralNetworksExecution* execution", 151 ), # noqa: B950 152 ( 153 "int", 154 "ANeuralNetworksExecution_setInput", 155 "ANeuralNetworksExecution* execution, int32_t index, const ANeuralNetworksOperandType* type, const void* buffer, size_t length", # noqa: B950 156 ), 157 ( 158 "int", 159 "ANeuralNetworksExecution_setInputFromMemory", 160 "ANeuralNetworksExecution* execution, int32_t index, const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, size_t offset, size_t length", # noqa: B950 161 ), 162 ( 163 "int", 164 "ANeuralNetworksExecution_setOutput", 165 "ANeuralNetworksExecution* execution, int32_t index, const ANeuralNetworksOperandType* type, void* buffer, size_t length", 166 ), # noqa: B950 167 ( 168 "int", 169 "ANeuralNetworksExecution_setOutputFromMemory", 170 "ANeuralNetworksExecution* execution, int32_t index, const ANeuralNetworksOperandType* type, const ANeuralNetworksMemory* memory, size_t offset, size_t length", # noqa: B950 171 ), 172 ( 173 "int", 174 "ANeuralNetworksExecution_startCompute", 175 "ANeuralNetworksExecution* execution, ANeuralNetworksEvent** event", 176 ), # noqa: B950 177 ("int", "ANeuralNetworksEvent_wait", "ANeuralNetworksEvent* event"), # noqa: B950 178 ("void", "ANeuralNetworksEvent_free", "ANeuralNetworksEvent* event"), # noqa: B950 179 ( 180 "int", 181 "ANeuralNetworksExecution_getOutputOperandRank", 182 "ANeuralNetworksExecution* execution, int32_t index, uint32_t* rank", 183 ), # noqa: B950 184 ( 185 "int", 186 "ANeuralNetworksExecution_getOutputOperandDimensions", 187 "ANeuralNetworksExecution* execution, int32_t index, uint32_t* dimensions", 188 ), # noqa: B950 189] 190 191 192def main(argv): 193 struct_members = [] 194 load_functions = [] 195 define_checks = [] 196 197 for ret, name, args in NNAPI_FUNCTIONS: 198 short_name = name.replace("ANeuralNetworks", "", 1) 199 200 struct_members.append(f" {ret}(*{short_name})({args});") 201 202 load_functions.append( 203 f' *(void**)&nnapi_.{short_name} = dlsym(handle, "{name}");' 204 ) 205 load_functions.append(f" check_nnapi_.{short_name} = check_{short_name};") 206 207 call_args = "".join(re.findall(r"\w+(?:,|$)", args)) 208 if ret == "void": 209 define_checks.append( 210 textwrap.dedent( 211 f"""\ 212 {ret} check_{short_name}({args}) {{ 213 CAFFE_ENFORCE(nnapi_.{short_name}); 214 nnapi_.{short_name}({call_args}); 215 }}""" 216 ) 217 ) 218 if ret == "int": 219 define_checks.append( 220 textwrap.dedent( 221 f"""\ 222 {ret} check_{short_name}({args}) {{ 223 CAFFE_ENFORCE(nnapi_.{short_name}); 224 int ret = nnapi_.{short_name}({call_args}); 225 // TODO: Maybe add better logging here. 226 CAFFE_ENFORCE( 227 ret == ANEURALNETWORKS_NO_ERROR, 228 "{short_name}", "failed with error ", ret 229 ); 230 return ret; 231 }}""" 232 ) 233 ) 234 235 out_dir = Path(__file__).parent 236 237 (out_dir / "nnapi_wrapper.h").write_text( 238 PREFIX 239 + textwrap.dedent( 240 """\ 241 #ifndef NNAPI_WRAPPER_H_ 242 #define NNAPI_WRAPPER_H_ 243 #include <stddef.h> 244 #include <stdint.h> 245 #include <ATen/nnapi/NeuralNetworks.h> 246 struct nnapi_wrapper { 247 __STRUCT_MEMBERS__ 248 }; 249 #ifdef __cplusplus 250 void nnapi_wrapper_load(struct nnapi_wrapper** nnapi, struct nnapi_wrapper** check_nnapi); 251 #endif 252 #endif 253 """ 254 ).replace("__STRUCT_MEMBERS__", "\n".join(struct_members)) 255 ) 256 257 (out_dir / "nnapi_wrapper.cpp").write_text( 258 PREFIX 259 + textwrap.dedent( 260 """\ 261 #ifndef _WIN32 262 #include <dlfcn.h> 263 #endif 264 #include <ATen/nnapi/nnapi_wrapper.h> 265 #include <c10/util/Logging.h> 266 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) 267 static int loaded = 0; 268 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) 269 static struct nnapi_wrapper nnapi_; 270 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables) 271 static struct nnapi_wrapper check_nnapi_; 272 __DEFINE_CHECK_FUNCTIONS__ 273 void nnapi_wrapper_load(struct nnapi_wrapper** nnapi, struct nnapi_wrapper** check_nnapi) { 274 #ifdef _WIN32 275 TORCH_CHECK(false, "Running NNAPI models is not supported on Windows."); 276 #else 277 if (!loaded) { 278 // Clear error flag. 279 dlerror(); 280 void* handle = dlopen("libneuralnetworks.so", RTLD_LAZY | RTLD_LOCAL); 281 CAFFE_ENFORCE(handle, "Failed to load libneuralnetworks.so ", dlerror()); 282 __LOAD_FUNCTIONS__ 283 loaded = 1; 284 } 285 *nnapi = &nnapi_; 286 *check_nnapi = &check_nnapi_; 287 #endif 288 } 289 """ 290 ) 291 .replace("__DEFINE_CHECK_FUNCTIONS__", "\n".join(define_checks)) 292 .replace("__LOAD_FUNCTIONS__", "\n".join(load_functions)) 293 ) 294 295 296if __name__ == "__main__": 297 sys.exit(main(sys.argv)) 298