diff options
Diffstat (limited to 'tensorflow/stream_executor/dso_loader.cc')
| -rw-r--r-- | tensorflow/stream_executor/dso_loader.cc | 60 |
1 files changed, 30 insertions, 30 deletions
diff --git a/tensorflow/stream_executor/dso_loader.cc b/tensorflow/stream_executor/dso_loader.cc index cce31ef4dc..83f2cadfad 100644 --- a/tensorflow/stream_executor/dso_loader.cc +++ b/tensorflow/stream_executor/dso_loader.cc @@ -13,6 +13,9 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +// TODO(jhen): Replace hardcoded, platform specific path strings in GetXXXPath() +// with a function in e.g. cuda.h. + #include "tensorflow/stream_executor/dso_loader.h" #include <dlfcn.h> @@ -32,19 +35,17 @@ limitations under the License. #include "tensorflow/stream_executor/lib/stringprintf.h" #include "tensorflow/stream_executor/platform/logging.h" #include "tensorflow/stream_executor/platform/port.h" -#include "tensorflow/stream_executor/lib/str_util.h" namespace perftools { namespace gputools { namespace internal { -// TensorFlow OSS configure uses the following lines to configure versions. For -// any modifications of the format, please make sure the script still works. -string GetCudaVersion() { return ""; } -string GetCudnnVersion() { return ""; } +string GetCudaVersion() { return TF_CUDA_VERSION; } +string GetCudnnVersion() { return TF_CUDNN_VERSION; } /* static */ port::Status DsoLoader::GetCublasDsoHandle(void** dso_handle) { - return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName("cublas", GetCudaVersion()), + return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName( + "cublas", GetCudaVersion()), GetCudaLibraryDirPath()), dso_handle); } @@ -53,35 +54,38 @@ string GetCudnnVersion() { return ""; } // libcudnn is versioned differently than the other libraries and may have a // different version number than other CUDA libraries. See b/22397368 for // some details about the complications surrounding this. - return GetDsoHandle( - FindDsoPath(tensorflow::internal::FormatLibraryFileName("cudnn", GetCudnnVersion()), - GetCudaLibraryDirPath()), + return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName( + "cudnn", GetCudnnVersion()), + GetCudaLibraryDirPath()), dso_handle); } /* static */ port::Status DsoLoader::GetCufftDsoHandle(void** dso_handle) { - return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName("cufft", GetCudaVersion()), + return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName( + "cufft", GetCudaVersion()), GetCudaLibraryDirPath()), dso_handle); } /* static */ port::Status DsoLoader::GetCurandDsoHandle(void** dso_handle) { - return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName("curand", GetCudaVersion()), + return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName( + "curand", GetCudaVersion()), GetCudaLibraryDirPath()), dso_handle); } /* static */ port::Status DsoLoader::GetLibcudaDsoHandle(void** dso_handle) { - return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName("cuda", "1"), - GetCudaDriverLibraryPath()), - dso_handle); + return GetDsoHandle( + FindDsoPath(tensorflow::internal::FormatLibraryFileName("cuda", "1"), + GetCudaDriverLibraryPath()), + dso_handle); } /* static */ port::Status DsoLoader::GetLibcuptiDsoHandle(void** dso_handle) { - return GetDsoHandle( - FindDsoPath(tensorflow::internal::FormatLibraryFileName("cupti", GetCudaVersion()), - GetCudaCuptiLibraryPath()), - dso_handle); + return GetDsoHandle(FindDsoPath(tensorflow::internal::FormatLibraryFileName( + "cupti", GetCudaVersion()), + GetCudaCuptiLibraryPath()), + dso_handle); } /* static */ void DsoLoader::RegisterRpath(port::StringPiece path) { @@ -89,11 +93,9 @@ string GetCudnnVersion() { return ""; } GetRpaths()->push_back(path.ToString()); } - /* static */ port::Status DsoLoader::GetDsoHandle(port::StringPiece path, void** dso_handle, LoadKind load_kind) { - int dynload_flags = RTLD_LAZY | (load_kind == LoadKind::kLocal ? RTLD_LOCAL : RTLD_GLOBAL); string path_string = path.ToString(); @@ -138,9 +140,9 @@ string GetCudnnVersion() { return ""; } static std::vector<string>* CreatePrimordialRpaths() { auto rpaths = new std::vector<string>; #if defined(__APPLE__) - rpaths->push_back("driver/driver_sh.runfiles/org_tensorflow/third_party/gpus/cuda/lib"); + rpaths->push_back("driver/driver_sh.runfiles/local_config_cuda/cuda/lib"); #else - rpaths->push_back("driver/driver_sh.runfiles/org_tensorflow/third_party/gpus/cuda/lib64"); + rpaths->push_back("driver/driver_sh.runfiles/local_config_cuda/cuda/lib64"); #endif return rpaths; } @@ -165,7 +167,6 @@ static std::vector<string>* CreatePrimordialRpaths() { /* static */ string DsoLoader::FindDsoPath(port::StringPiece library_name, port::StringPiece runfiles_relpath) { - // Keep a record of the paths we attempted so we can dump out meaningful // diagnostics if no path is found. std::vector<string> attempted; @@ -191,29 +192,28 @@ static std::vector<string>* CreatePrimordialRpaths() { /* static */ string DsoLoader::GetCudaLibraryDirPath() { #if defined(__APPLE__) - return "third_party/gpus/cuda/lib"; + return "external/local_config_cuda/cuda/lib"; #else - return "third_party/gpus/cuda/lib64"; + return "external/local_config_cuda/cuda/lib64"; #endif } /* static */ string DsoLoader::GetCudaDriverLibraryPath() { #if defined(__APPLE__) - return "third_party/gpus/cuda/driver/lib"; + return "external/local_config_cuda/cuda/driver/lib"; #else - return "third_party/gpus/cuda/driver/lib64"; + return "external/local_config_cuda/cuda/driver/lib64"; #endif } /* static */ string DsoLoader::GetCudaCuptiLibraryPath() { #if defined(__APPLE__) - return "third_party/gpus/cuda/extras/CUPTI/lib"; + return "external/local_config_cuda/cuda/extras/CUPTI/lib"; #else - return "third_party/gpus/cuda/extras/CUPTI/lib64"; + return "external/local_config_cuda/cuda/extras/CUPTI/lib64"; #endif } - // -- CachedDsoLoader /* static */ port::StatusOr<void*> CachedDsoLoader::GetCublasDsoHandle() { |