1 /* Copyright 2021 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/compiler/xla/service/gpu/nvptx_helper.h"
17
18 #include <string>
19
20 #include "absl/strings/str_join.h"
21 #include "tensorflow/core/lib/io/path.h"
22 #include "tensorflow/core/platform/cuda_libdevice_path.h"
23
24 namespace xla {
25 namespace gpu {
26
27 namespace {
28
CandidateCudaRoots(const HloModuleConfig & config)29 std::vector<std::string> CandidateCudaRoots(const HloModuleConfig& config) {
30 return tensorflow::CandidateCudaRoots(
31 config.debug_options().xla_gpu_cuda_data_dir());
32 }
33
34 } // namespace
35
CantFindCudaMessage(absl::string_view msg,const HloModuleConfig & hlo_module_config)36 std::string CantFindCudaMessage(absl::string_view msg,
37 const HloModuleConfig& hlo_module_config) {
38 return absl::StrCat(
39 msg, "\nSearched for CUDA in the following directories:\n ",
40 absl::StrJoin(CandidateCudaRoots(hlo_module_config), "\n "),
41 "\nYou can choose the search directory by setting xla_gpu_cuda_data_dir "
42 "in HloModule's DebugOptions. For most apps, setting the environment "
43 "variable XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda will work.");
44 }
45
GetLibdeviceDir(const HloModuleConfig & hlo_module_config)46 std::string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) {
47 for (const std::string& cuda_root : CandidateCudaRoots(hlo_module_config)) {
48 std::string libdevice_dir =
49 tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice");
50 VLOG(2) << "Looking for libdevice at " << libdevice_dir;
51 if (tensorflow::Env::Default()->IsDirectory(libdevice_dir).ok()) {
52 VLOG(2) << "Found libdevice dir " << libdevice_dir;
53 return libdevice_dir;
54 }
55 }
56 LOG(WARNING) << CantFindCudaMessage(
57 "Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice. This may "
58 "result in compilation or runtime failures, if the program we try to run "
59 "uses routines from libdevice.",
60 hlo_module_config);
61
62 // GetCudaRootCandidates always includes ".", but if everything fails, we
63 // return it anyway. Better than returning the empty string.
64 return ".";
65 }
66
67 } // namespace gpu
68 } // namespace xla
69