From 29390688cb1aa425523103ee0fa0c8986d5790a0 Mon Sep 17 00:00:00 2001
From: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
Date: Fri, 21 Oct 2022 14:02:26 +0800
Subject: [PATCH] [Feats]: Add ROCm info when collecting env (#633)
* Add ROCm info when collecting env
* minor fix
---
mmengine/utils/dl_utils/collect_env.py | 30 ++++++++++++++++++--------
1 file changed, 21 insertions(+), 9 deletions(-)
diff --git a/mmengine/utils/dl_utils/collect_env.py b/mmengine/utils/dl_utils/collect_env.py
index a6e87dce..c04b8df1 100644
--- a/mmengine/utils/dl_utils/collect_env.py
+++ b/mmengine/utils/dl_utils/collect_env.py
@@ -67,15 +67,27 @@ def collect_env():
env_info['CUDA_HOME'] = CUDA_HOME
if CUDA_HOME is not None and osp.isdir(CUDA_HOME):
- try:
- nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
- nvcc = subprocess.check_output(f'"{nvcc}" -V', shell=True)
- nvcc = nvcc.decode('utf-8').strip()
- release = nvcc.rfind('Cuda compilation tools')
- build = nvcc.rfind('Build ')
- nvcc = nvcc[release:build].strip()
- except subprocess.SubprocessError:
- nvcc = 'Not Available'
+ if CUDA_HOME == '/opt/rocm':
+ try:
+ nvcc = osp.join(CUDA_HOME, 'hip/bin/hipcc')
+ nvcc = subprocess.check_output(
+ f'"{nvcc}" --version', shell=True)
+ nvcc = nvcc.decode('utf-8').strip()
+ release = nvcc.rfind('HIP version:')
+ build = nvcc.rfind('')
+ nvcc = nvcc[release:build].strip()
+ except subprocess.SubprocessError:
+ nvcc = 'Not Available'
+ else:
+ try:
+ nvcc = osp.join(CUDA_HOME, 'bin/nvcc')
+ nvcc = subprocess.check_output(f'"{nvcc}" -V', shell=True)
+ nvcc = nvcc.decode('utf-8').strip()
+ release = nvcc.rfind('Cuda compilation tools')
+ build = nvcc.rfind('Build ')
+ nvcc = nvcc[release:build].strip()
+ except subprocess.SubprocessError:
+ nvcc = 'Not Available'
env_info['NVCC'] = nvcc
try:
--
GitLab