1 /* Copyright 2020 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
16 #include "tensorflow/core/tpu/tpu_model_server_initializer.h"
17
18 #include <dlfcn.h>
19
20 #include "tensorflow/core/platform/errors.h"
21 #include "tensorflow/core/platform/status.h"
22 #include "tensorflow/core/tpu/tpu_api_dlsym_set_fn.h"
23
24 #if !defined(PLATFORM_GOOGLE)
25 #include "tensorflow/core/tpu/tpu_api.h"
26 #include "tensorflow/core/tpu/tpu_initializer_helper.h"
27 #include "tensorflow/stream_executor/tpu/tpu_platform.h"
28 #endif
29
30 namespace tensorflow {
31 namespace tpu {
32
33
34 #if defined(PLATFORM_GOOGLE)
InitializeTpuModelServer(void * library_handle)35 Status InitializeTpuModelServer(void* library_handle) {
36 return errors::Unimplemented("You must statically link in a TPU library.");
37 }
38 #else // PLATFORM_GOOGLE
39 #include "tensorflow/core/tpu/tpu_library_init_fns.inc"
40
41 Status InitializeTpuModelServer(void* library_handle) {
42 Status s = InitializeTpuStructFns(library_handle);
43
44 // Retrieve arguments from environment if applicable
45 std::pair<std::vector<std::string>, std::vector<const char*> > args =
46 GetLibTpuInitArguments();
47
48 // TPU platform registration must only be performed after the library is
49 // loaded. We do not want to register a TPU platform in XLA without the
50 // supporting library providing the necessary APIs.
51 if (s.ok()) {
52 void (*initialize_fn)(bool init_library, int num_args, const char** args);
53 initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
54 dlsym(library_handle, "TfTpu_Initialize"));
55 (*initialize_fn)(/*init_library=*/true, args.second.size(),
56 args.second.data());
57
58 RegisterTpuPlatform();
59 }
60
61 OpsApiFn()->TfTpu_InitializeTpuModelServerFn();
62 return s;
63 }
64
65 bool FindAndLoadTpuModelServer() {
66 if (!TryAcquireTpuLock()) return false;
67 void* library = dlopen("libtpu.so", RTLD_NOW);
68 if (library) {
69 InitializeTpuModelServer(library);
70 }
71 return true;
72 }
73
74 static bool tpu_library_finder = FindAndLoadTpuModelServer();
75 #endif // PLATFORM_GOOGLE
76
77 } // namespace tpu
78 } // namespace tensorflow
79