Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Support dipu device #1127

Merged
merged 4 commits into from
May 25, 2023
Merged

Support dipu device #1127

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
14f539f
Support dipu debice
CokeDong May 5, 2023
4f65be7
Fix lint
CokeDong May 10, 2023
6a9f104
Merge branch 'main' into test
CokeDong May 10, 2023
6008a66
Fix merge lint
CokeDong May 10, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
6 changes: 3 additions & 3 deletions mmengine/device/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
# Copyright (c) OpenMMLab. All rights reserved.
from .utils import (get_device, get_max_cuda_memory, is_cuda_available,
is_mlu_available, is_mps_available, is_npu_available,
is_npu_support_full_precision)
is_dipu_available, is_mlu_available, is_mps_available,
is_npu_available, is_npu_support_full_precision)

__all__ = [
'get_max_cuda_memory', 'get_device', 'is_cuda_available',
'is_mlu_available', 'is_mps_available', 'is_npu_available',
'is_npu_support_full_precision'
'is_dipu_available', 'is_npu_support_full_precision'
]
12 changes: 12 additions & 0 deletions mmengine/device/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@
except Exception:
IS_NPU_AVAILABLE = False

try:
import torch_dipu # noqa: F401
IS_DIPU_AVAILABLE = True
except Exception:
IS_DIPU_AVAILABLE = False


def get_max_cuda_memory(device: Optional[torch.device] = None) -> int:
"""Returns the maximum GPU memory occupied by tensors in megabytes (MB) for
Expand Down Expand Up @@ -63,6 +69,10 @@ def is_mps_available() -> bool:
return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()


def is_dipu_available() -> bool:
return IS_DIPU_AVAILABLE


def is_npu_support_full_precision() -> bool:
"""Returns True if npu devices support full precision training."""
version_of_support_full_precision = 220
Expand All @@ -79,6 +89,8 @@ def is_npu_support_full_precision() -> bool:
DEVICE = 'mlu'
elif is_mps_available():
DEVICE = 'mps'
elif is_dipu_available():
DEVICE = 'dipu'


def get_device() -> str:
Expand Down