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_api_dlsym_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
31 // Reminder: Update tpu_library_loader_windows.cc if you are adding new publicly
32 // visible methods.
33
34 namespace tensorflow {
35 namespace tpu {
36
37 #if defined(PLATFORM_GOOGLE)
InitializeTpuLibrary(void * library_handle)38 Status InitializeTpuLibrary(void* library_handle) {
39 return errors::Unimplemented("You must statically link in a TPU library.");
40 }
41 #else // PLATFORM_GOOGLE
42 #include "tensorflow/core/tpu/tpu_library_init_fns.inc"
43
44 Status InitializeTpuLibrary(void* library_handle) {
45 Status s = InitializeTpuStructFns(library_handle);
46
47 // Retrieve arguments from environment if applicable
48 std::pair<std::vector<std::string>, std::vector<const char*> > args =
49 GetLibTpuInitArguments();
50
51 // TPU platform registration must only be performed after the library is
52 // loaded. We do not want to register a TPU platform in XLA without the
53 // supporting library providing the necessary APIs.
54 if (s.ok()) {
55 void (*initialize_fn)(bool init_library, int num_args, const char** args);
56 initialize_fn = reinterpret_cast<decltype(initialize_fn)>(
57 dlsym(library_handle, "TfTpu_Initialize"));
58 (*initialize_fn)(/*init_library=*/true, args.second.size(),
59 args.second.data());
60
61 RegisterTpuPlatform();
62 }
63
64 return s;
65 }
66
67 bool FindAndLoadTpuLibrary() {
68 if (!TryAcquireTpuLock()) return false;
69
70 void* library = dlopen("libtpu.so", RTLD_NOW);
71 if (library) {
72 InitializeTpuLibrary(library);
73 }
74 return true;
75 }
76
77 static bool tpu_library_finder = FindAndLoadTpuLibrary();
78 #endif // PLATFORM_GOOGLE
79
80 } // namespace tpu
81 } // namespace tensorflow
82