Merge pull request #42674 from levinxo:master

PiperOrigin-RevId: 328672559
Change-Id: I50de6ce97b230cbb59ababf5c9c83b61f45e6eb9
This commit is contained in:
TensorFlower Gardener 2020-08-26 22:05:05 -07:00
commit a8529f2c5d

View File

@ -176,6 +176,7 @@ def _header_paths():
"include/*-linux-gnu", "include/*-linux-gnu",
"extras/CUPTI/include", "extras/CUPTI/include",
"include/cuda/CUPTI", "include/cuda/CUPTI",
"local/cuda/extras/CUPTI/include",
] ]
@ -188,6 +189,8 @@ def _library_paths():
"lib/*-linux-gnu", "lib/*-linux-gnu",
"lib/x64", "lib/x64",
"extras/CUPTI/*", "extras/CUPTI/*",
"local/cuda/lib64",
"local/cuda/extras/CUPTI/lib64",
] ]
@ -268,12 +271,14 @@ def _find_cuda_config(base_paths, required_version):
nvcc_path, nvcc_version = _find_versioned_file(base_paths, [ nvcc_path, nvcc_version = _find_versioned_file(base_paths, [
"", "",
"bin", "bin",
"local/cuda/bin",
], nvcc_name, cuda_version, get_nvcc_version) ], nvcc_name, cuda_version, get_nvcc_version)
nvvm_path = _find_file(base_paths, [ nvvm_path = _find_file(base_paths, [
"nvvm/libdevice", "nvvm/libdevice",
"share/cuda", "share/cuda",
"lib/nvidia-cuda-toolkit/libdevice", "lib/nvidia-cuda-toolkit/libdevice",
"local/cuda/nvvm/libdevice",
], "libdevice*.10.bc") ], "libdevice*.10.bc")
cupti_header_path = _find_file(base_paths, _header_paths(), "cupti.h") cupti_header_path = _find_file(base_paths, _header_paths(), "cupti.h")