1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ 16 #define TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ 17 18 // This file is depended on by kernels that have to build for mobile devices. 19 // For this reason, we avoid relying on TensorFlow and instead only use the 20 // standard C++ library. 21 22 #include <mutex> // NOLINT 23 #include <string> 24 #include <unordered_map> 25 26 namespace xla { 27 namespace cpu { 28 29 // The CPU JIT compiler uses this registry to resolve symbolic CustomCall 30 // targets; so when using the CPU JIT, CustomCall targets need to be registered 31 // here with the symbol name used in the CustomCall. 32 // 33 // The XLA AOT compiler links using a standard offline linker; so when compiling 34 // in AOT mode, you *also* need to make sure the name of the callee (presumably 35 // implemented in C++) matches up with the symbolic name used in the CustomCall. 36 // 37 // We maintain the registry in both the JIT and the AOT cases for simplicity, 38 // but we only use it when running in JIT mode. 39 class CustomCallTargetRegistry { 40 public: 41 static CustomCallTargetRegistry* Global(); 42 43 void Register(const std::string& symbol, void* address); 44 void* Lookup(const std::string& symbol) const; 45 46 private: 47 std::unordered_map<std::string, void*> registered_symbols_; 48 mutable std::mutex mu_; 49 }; 50 51 class RegisterCustomCallTarget { 52 public: RegisterCustomCallTarget(const std::string & name,void * address)53 explicit RegisterCustomCallTarget(const std::string& name, void* address) { 54 CustomCallTargetRegistry::Global()->Register(name, address); 55 } 56 }; 57 58 #define REGISTER_CUSTOM_CALL_CONCAT(a, b) a##b 59 60 #define REGISTER_CUSTOM_CALL_TARGET_WITH_SYM_HELPER(symbol, address, counter) \ 61 static ::xla::cpu::RegisterCustomCallTarget REGISTER_CUSTOM_CALL_CONCAT( \ 62 custom_call_target_register, counter)(symbol, \ 63 reinterpret_cast<void*>(address)) 64 65 #define REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(symbol, address) \ 66 REGISTER_CUSTOM_CALL_TARGET_WITH_SYM_HELPER(symbol, address, __COUNTER__) 67 68 #define REGISTER_CUSTOM_CALL_TARGET(function) \ 69 REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(#function, function) 70 71 } // namespace cpu 72 } // namespace xla 73 74 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CPU_CUSTOM_CALL_TARGET_REGISTRY_H_ 75