saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[cpp_builder] refactor to reduce libcudart_static logs (#146394)

Browse source 
Want to reduce logs from `log_msg = f'"libcudart_static.a" not found under {path}'`, which was added in https://github.com/pytorch/pytorch/pull/142175

Differential Revision: D69096354

Pull Request resolved: https://github.com/pytorch/pytorch/pull/146394
Approved by: https://github.com/benjaminglass1, https://github.com/chenyang78
This commit is contained in:
Henry Tsang 2025-02-05 00:41:30 +00:00 committed by PyTorch MergeBot
parent 658e22d495
commit 9d5bf38dec
1 changed files with 12 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
torch/_inductor/cpp_builder.py
View file

@ -1151,19 +1151,25 @@ def _set_gpu_runtime_env() -> None:
os.environ["CUDA_HOME"] = build_paths.sdk_home
@functools.lru_cache(8)
def _find_libcudart_static(path: str) -> Optional[Path]:
lib_dirs = list(Path(path).rglob("libcudart_static.a"))
if lib_dirs:
return lib_dirs[0].resolve().parent
log_msg = f'"libcudart_static.a" not found under {path}'
log.info(log_msg)
return None
def _transform_cuda_paths(lpaths: list[str]) -> None:
# This handles two cases:
# 1. Cases where libs are in (e.g.) lib/cuda-12 and lib/cuda-12/stubs
# 2. Linux machines may have CUDA installed under either lib64/ or lib/
for i, path in enumerate(lpaths):
if "CUDA_HOME" in os.environ and path.startswith(os.environ["CUDA_HOME"]):
try:
lib_dir = next(Path(path).rglob("libcudart_static.a")).resolve().parent
except StopIteration:
log_msg = f'"libcudart_static.a" not found under {path}'
log.info(log_msg)
lib_dir: Optional[Path] = _find_libcudart_static(path)
if lib_dir is None:
continue
lpaths[i] = str(lib_dir)
stub_dir = lib_dir / "stubs"
if stub_dir.exists():