experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Remove XLA::GetCudaRootCandidates, extend tensorflow::CandidateCudaRoots

Browse Source 
instead.

PiperOrigin-RevId: 248247226
This commit is contained in:
George Karpenkov 2019-05-14 17:50:33 -07:00 committed by TensorFlower Gardener
parent 36773d3037
commit cd134aa61d
4 changed files with 26 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc
View File

@ -120,13 +120,18 @@ namespace {
namespace tracing = tensorflow::tracing;
static std::vector<std::string> CandidateCudaRoots(
const HloModuleConfig& config) {
return tensorflow::CandidateCudaRoots(
config.debug_options().xla_gpu_cuda_data_dir());
}
void PrintCantFindCudaMessage(absl::string_view msg,
const HloModuleConfig& hlo_module_config) {
LOG(WARNING) << msg;
LOG(WARNING) << "Searched in the following directories:";
for (const auto& dir :
GetCudaRootCandidates(PtxCompilationOptions(hlo_module_config))) {
for (const auto& dir : CandidateCudaRoots(hlo_module_config)) {
LOG(WARNING) << " " << dir;
}
LOG(WARNING)
@ -137,9 +142,7 @@ void PrintCantFindCudaMessage(absl::string_view msg,
// Returns the directory containing nvvm libdevice files.
string GetLibdeviceDir(const HloModuleConfig& hlo_module_config) {
const auto& candidate_dirs =
GetCudaRootCandidates(PtxCompilationOptions(hlo_module_config));
for (const string& cuda_root : candidate_dirs) {
for (const string& cuda_root : CandidateCudaRoots(hlo_module_config)) {
string libdevice_dir =
tensorflow::io::JoinPath(cuda_root, "nvvm", "libdevice");
VLOG(2) << "Looking for libdevice at " << libdevice_dir;

23
tensorflow/compiler/xla/service/gpu/stream_executor_util.cc
View File

@ -300,26 +300,6 @@ void WarnIfBadPtxasVersion(const string& ptxas_path) {
}
}
// Returns a vector of potential locations of the CUDA root directory.
// Searches through tensorflow CUDA locations AND through the CUDA location
// specified in HLO configuration.
std::vector<string> GetCudaRootCandidates(
PtxCompilationOptions compile_ptx_options) {
std::vector<string> potential_cuda_roots = tensorflow::CandidateCudaRoots();
// "." is our last resort, even though it probably won't work.
potential_cuda_roots.push_back(".");
// CUDA location explicitly specified by user via --xla_gpu_cuda_data_dir has
// highest priority.
string xla_gpu_cuda_data_dir = compile_ptx_options.xla_gpu_cuda_data_dir;
if (!xla_gpu_cuda_data_dir.empty()) {
potential_cuda_roots.insert(potential_cuda_roots.begin(),
xla_gpu_cuda_data_dir);
}
return potential_cuda_roots;
}
StatusOr<absl::Span<const uint8>> CompilePtxOrGetCached(
se::StreamExecutor* executor, absl::string_view ptx,
PtxCompilationOptions compilation_options) {
@ -360,7 +340,8 @@ StatusOr<std::vector<uint8>> CompilePtx(
"Compile PTX", tensorflow::profiler::TraceMeLevel::kInfo);
auto env = tensorflow::Env::Default();
string ptxas_path;
for (const string& cuda_root : GetCudaRootCandidates(compile_ptx_options)) {
for (const string& cuda_root : tensorflow::CandidateCudaRoots(
/*preferred_location=*/compile_ptx_options.xla_gpu_cuda_data_dir)) {
ptxas_path = tensorflow::io::JoinPath(cuda_root, "bin", "ptxas");
VLOG(2) << "Looking for ptxas at " << ptxas_path;
if (env->FileExists(ptxas_path).ok()) {

6
tensorflow/compiler/xla/service/gpu/stream_executor_util.h
View File

@ -145,12 +145,6 @@ StatusOr<absl::Span<const uint8>> CompilePtxOrGetCached(
se::StreamExecutor* executor, absl::string_view ptx,
PtxCompilationOptions compilation_options);
// Returns a vector of potential locations of the CUDA root directory.
// Searches through tensorflow CUDA locations AND through the CUDA location
// specified in compile_ptx_options (can be constructed from HloModuleConfig).
std::vector<string> GetCudaRootCandidates(
PtxCompilationOptions compile_ptx_options);
} // namespace gpu
} // namespace xla

16
tensorflow/core/platform/cuda_libdevice_path.h
View File

@ -25,6 +25,22 @@ namespace tensorflow {
// the CUDA SDK, which contains sub-folders such as bin, lib64, and nvvm.
std::vector<string> CandidateCudaRoots();
// A convenient wrapper for CandidateCudaRoots, which allows supplying a
// preferred location (inserted first in the output vector), and a flag whether
// the current working directory should be searched (inserted last).
inline std::vector<string> CandidateCudaRoots(
string preferred_location, bool use_working_directory = true) {
std::vector<string> candidates = CandidateCudaRoots();
if (!preferred_location.empty()) {
candidates.insert(candidates.begin(), preferred_location);
}
// "." is our last resort, even though it probably won't work.
candidates.push_back(".");
return candidates;
}
} // namespace tensorflow
#endif // TENSORFLOW_CORE_PLATFORM_CUDA_LIBDEVICE_PATH_H_