Home | History | Annotate | Download | only in stream_executor
      1 /* Copyright 2015 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 // Common DSO loading functionality: exposes callables that dlopen DSOs
     17 // in either the runfiles directories
     18 
     19 #ifndef TENSORFLOW_STREAM_EXECUTOR_DSO_LOADER_H_
     20 #define TENSORFLOW_STREAM_EXECUTOR_DSO_LOADER_H_
     21 
     22 #include "tensorflow/stream_executor/platform/port.h"
     23 #include <vector>
     24 
     25 #include "tensorflow/stream_executor/lib/status.h"
     26 #include "tensorflow/stream_executor/lib/statusor.h"
     27 #include "tensorflow/stream_executor/lib/stringpiece.h"
     28 #include "tensorflow/stream_executor/platform.h"
     29 #include "tensorflow/stream_executor/platform/mutex.h"
     30 
     31 namespace perftools {
     32 namespace gputools {
     33 namespace internal {
     34 
     35 // Permits StreamExecutor code to dynamically load a pre-determined set of
     36 // relevant DSOs via dlopen.
     37 //
     38 // Thread-safe.
     39 class DsoLoader {
     40  public:
     41   // The following methods either load the DSO of interest and return a dlopen
     42   // handle or error status in the canonical namespace.
     43 
     44   static port::Status GetCublasDsoHandle(void** dso_handle);
     45   static port::Status GetCudnnDsoHandle(void** dso_handle);
     46   static port::Status GetCufftDsoHandle(void** dso_handle);
     47   static port::Status GetCurandDsoHandle(void** dso_handle);
     48   static port::Status GetLibcudaDsoHandle(void** dso_handle);
     49   static port::Status GetLibcuptiDsoHandle(void** dso_handle);
     50 
     51   // Registers a new binary-relative path to use as a dlopen search path.
     52   static void RegisterRpath(port::StringPiece path);
     53 
     54  private:
     55   // Registered rpaths (singleton vector) and a mutex that guards it.
     56   static std::vector<string>* GetRpaths();
     57 
     58   // Descriptive boolean wrapper to indicate whether symbols are made available
     59   // to resolve in later-loaded libraries.
     60   enum class LoadKind { kLocal, kGlobal };
     61 
     62   // Loads a DSO from the given "path" (which can technically be any dlopen-able
     63   // name). If the load kind is global, the symbols in the loaded DSO are
     64   // visible to subsequent DSO loading operations.
     65   static port::Status GetDsoHandle(port::StringPiece path, void** dso_handle,
     66                                    LoadKind load_kind = LoadKind::kLocal);
     67 
     68 
     69   // Returns the binary directory (or binary path) associated with the currently
     70   // executing program. If strip_executable_name is true, the executable file is
     71   // stripped off of the path.
     72   static string GetBinaryDirectory(bool strip_executable_name);
     73 
     74   // Invokes realpath on the original path; updates candidate and returns true
     75   // if it succeeds (i.e. a file exists at the path); otherwise, returns false.
     76   static bool TrySymbolicDereference(string* candidate);
     77 
     78   // Attempts to find a path to the DSO of interest, otherwise returns the
     79   // bare library name:
     80   // Arguments:
     81   //   library_name: the filename in tree; e.g. libOpenCL.so.1.0.0
     82   //   runfiles_relpath: where to look for the library relative to the runfiles
     83   //      root; e.g. third_party/gpus/cuda/lib64
     84   static string FindDsoPath(port::StringPiece library_name,
     85                             port::StringPiece runfiles_relpath);
     86 
     87   // Return platform dependent paths for DSOs
     88   static string GetCudaLibraryDirPath();
     89   static string GetCudaDriverLibraryPath();
     90   static string GetCudaCuptiLibraryPath();
     91 
     92   SE_DISALLOW_COPY_AND_ASSIGN(DsoLoader);
     93 };
     94 
     95 // Wrapper around the DsoLoader that prevents us from dlopen'ing any of the DSOs
     96 // more than once.
     97 class CachedDsoLoader {
     98  public:
     99   // Cached versions of the corresponding DsoLoader methods above.
    100   static port::StatusOr<void*> GetCublasDsoHandle();
    101   static port::StatusOr<void*> GetCudnnDsoHandle();
    102   static port::StatusOr<void*> GetCufftDsoHandle();
    103   static port::StatusOr<void*> GetCurandDsoHandle();
    104   static port::StatusOr<void*> GetLibcudaDsoHandle();
    105   static port::StatusOr<void*> GetLibcuptiDsoHandle();
    106 
    107  private:
    108   // Fetches a DSO handle via "load_dso" and returns the StatusOr form of the
    109   // result.
    110   static port::StatusOr<void*> FetchHandleResult(
    111       std::function<port::Status(void**)> load_dso);
    112 
    113   SE_DISALLOW_COPY_AND_ASSIGN(CachedDsoLoader);
    114 };
    115 
    116 }  // namespace internal
    117 }  // namespace gputools
    118 }  // namespace perftools
    119 
    120 #endif  // TENSORFLOW_STREAM_EXECUTOR_DSO_LOADER_H_
    121