如何让CUDA XLA知道TensorFlow路径

Question

我通过命令pip install tf-nightly-gpu --prefix=/tf/install/path安装了TensorFlow夜间构建版本

当我尝试运行任何XLA示例时，TensorFlow出现错误"Unable to find libdevice dir. Using‘.“”无法将ptx编译为cubin。将尝试让GPU驱动程序编译ptx。找不到: /usr/local/cuda-10.0/bin/ptxas not find“。

所以很明显TensorFlow找不到我的CUDA路径。在我的系统中，CUDA安装在/cm/shared/apps/cuda/toolkit/10.0.130中。因为我不是从源代码构建TensorFlow的，所以在默认情况下，XLA搜索文件夹/user/local/cuda-*。但是因为我没有这个文件夹，所以它会发出一个错误。

目前，我的解决方法是创建一个符号链接。我检查了tensorflow/compiler/xla/service/gpu/nvptx_compiler.cc.中的TensorFlow源代码文件"//用户通过--xla_gpu_ CUDA _data_dir显式指定的cuda位置具有最高优先级“中有一个注释。那么如何将值传递给这个标志呢？我尝试了以下两个环境变量，但它们都不起作用：

export XLA_FLAGS="--xla_gpu_cuda_data_dir=/cm/shared/apps/cuda10.0/toolkit/10.0.130/"
export TF_XLA_FLAGS="--xla_gpu_cuda_data_dir=/cm/shared/apps/cuda10.0/toolkit/10.0.130/"

那么如何使用标志"--xla_gpu_cuda_data_dir“呢？谢谢。

Answer 1

针对此问题有代码更改，但不清楚如何使用。检查此处https://github.com/tensorflow/tensorflow/issues/23783

Answer 2

您可以在终端中运行export XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda