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_CUSTOM_CALL_TARGET_REGISTRY_H_ 16 #define TENSORFLOW_COMPILER_XLA_SERVICE_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 <map> 23 #include <mutex> // NOLINT 24 #include <string> 25 26 namespace xla { 27 28 // XLA JIT compilers use this registry to resolve symbolic CustomCall targets; 29 // so when using XLA as a JIT, CustomCall targets need to be registered here 30 // with the symbol name used in the CustomCall. 31 // 32 // The XLA:CPU ahead-of-time (AOT) compiler links using a standard offline 33 // linker; so when compiling in CPU AOT mode, you *also* need to make sure the 34 // name of the callee (presumably implemented in C++) matches up with the 35 // 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 const std::string& platform); 45 void* Lookup(const std::string& symbol, const std::string& platform) const; 46 47 private: 48 // Maps the pair (symbol, platform) to a C function implementing a custom call 49 // named `symbol` for StreamExecutor platform `platform`. 50 // 51 // Different platforms have different ABIs. TODO(jlebar): Describe them! 52 // 53 // (We std::map rather than std::unordered_map because the STL doesn't provide 54 // a default hasher for pair<std::string, std::string>, and we want to avoid 55 // pulling in dependencies that might define this.) 56 std::map<std::pair<std::string, std::string>, void*> registered_symbols_; 57 mutable std::mutex mu_; 58 }; 59 60 class RegisterCustomCallTarget { 61 public: RegisterCustomCallTarget(const std::string & name,void * address,const std::string & platform)62 explicit RegisterCustomCallTarget(const std::string& name, void* address, 63 const std::string& platform) { 64 CustomCallTargetRegistry::Global()->Register(name, address, platform); 65 } 66 }; 67 68 #define XLA_REGISTER_CUSTOM_CALL_CONCAT(a, b) a##b 69 70 #define XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM_HELPER(symbol, address, \ 71 platform, counter) \ 72 static ::xla::RegisterCustomCallTarget XLA_REGISTER_CUSTOM_CALL_CONCAT( \ 73 custom_call_target_register, counter)( \ 74 symbol, reinterpret_cast<void*>(address), platform) 75 76 #define XLA_REGISTER_CUSTOM_CALL_TARGET(function, platform) \ 77 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(#function, function, platform) 78 79 #define XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(symbol, address, platform) \ 80 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM_HELPER(symbol, address, platform, \ 81 __COUNTER__) 82 83 // Convenience overloads for registering custom-call targets on the CPU. 84 #define XLA_CPU_REGISTER_CUSTOM_CALL_TARGET(function) \ 85 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(#function, function, "Host") 86 87 #define XLA_CPU_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(symbol, address) \ 88 XLA_REGISTER_CUSTOM_CALL_TARGET_WITH_SYM(symbol, address, "Host") 89 90 } // namespace xla 91 92 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_CUSTOM_CALL_TARGET_REGISTRY_H_ 93