Home | History | Annotate | Download | only in cpu 
      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:
     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