From f81df0f341330951005ce102b61cd280b10c51e1 Mon Sep 17 00:00:00 2001 From: Lars Gleim 00627932 Date: Tue, 14 Jan 2025 12:07:04 +0100 Subject: [PATCH 1/2] Reworked backend & device selection logic Fixes Vulkan on non-cuda GPUs & adds Pytorch 2.5+ Intel XPU support Fixes https://github.com/Genesis-Embodied-AI/Genesis/issues/274 Merges the logic of get_gpu_device() and get_cpu_device() to handle device selection in a single place. If a specific backend is requested, this request is honored. The vulkan backend will use intel's pytorch xps device if available, reducing cross-device data transfer. The `gpu` backend will - default to `cuda` when available, - resolve to `mps` on macOS - or fallback to `vulkan`. --- genesis/__init__.py | 18 +++++------ genesis/utils/misc.py | 69 +++++++++++++++++++++++++++---------------- 2 files changed, 50 insertions(+), 37 deletions(-) diff --git a/genesis/__init__.py b/genesis/__init__.py index 37a81f1e..6b21d0de 100644 --- a/genesis/__init__.py +++ b/genesis/__init__.py @@ -23,7 +23,7 @@ def fake_print(*args, **kwargs): from .constants import backend as gs_backend from .logging import Logger from .version import __version__ -from .utils import set_random_seed, get_platform, get_cpu_device, get_gpu_device +from .utils import set_random_seed, get_platform, get_device _initialized = False backend = None @@ -43,7 +43,6 @@ def init( theme="dark", logger_verbose_time=False, ): - # genesis._initialized global _initialized if _initialized: @@ -72,19 +71,16 @@ def init( first_init = False - # get default device and compute total device memory + # genesis.backend global platform - global device platform = get_platform() - if backend == gs_backend.cpu: - device, device_name, total_mem = get_cpu_device() - else: - device, device_name, total_mem = get_gpu_device() - - # genesis.backend if backend not in GS_ARCH[platform]: raise_exception(f"backend ~~<{backend}>~~ not supported for platform ~~<{platform}>~~") - backend = GS_ARCH[platform][backend] + + # get default device and compute total device memory + global device + device, device_name, total_mem, backend = get_device(backend) + _globalize_backend(backend) logger.info( diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index 7ec3843f..fc5211f7 100644 --- a/genesis/utils/misc.py +++ b/genesis/utils/misc.py @@ -11,6 +11,7 @@ import torch import genesis as gs +from genesis.constants import backend as gs_backend def raise_exception(msg="Something went wrong."): @@ -96,32 +97,48 @@ def get_cpu_name(): return platform.processor() -def get_cpu_device(): - device_name = get_cpu_name() - total_mem = psutil.virtual_memory().total / 1024**3 - device = torch.device("cpu") - return device, device_name, total_mem - - -def get_gpu_device(): - if get_platform() == "macOS": - if not torch.backends.mps.is_available(): - gs.raise_exception("metal device not available") - - # on mac, cpu and gpu are in the same device - _, device_name, total_mem = get_cpu_device() - device = torch.device("mps") - - else: - if not torch.cuda.is_available(): - gs.raise_exception("cuda device not available") - - device = torch.device("cuda") - device_property = torch.cuda.get_device_properties(0) - device_name = device_property.name - total_mem = device_property.total_memory / 1024**3 - - return device, device_name, total_mem +def get_device(backend: gs_backend): + match backend: + case gs_backend.cuda: + if not torch.cuda.is_available(): + gs.raise_exception("cuda device not available") + + device = torch.device("cuda") + device_property = torch.cuda.get_device_properties(0) + device_name = device_property.name + total_mem = device_property.total_memory / 1024**3 + + case gs_backend.metal: + if not torch.backends.mps.is_available(): + gs.raise_exception("metal device not available") + + # on mac, cpu and gpu are in the same device + _, device_name, total_mem, _ = get_device(gs_backend.cpu) + device = torch.device("mps") + + case gs_backend.vulkan: + if torch.xpu.is_available(): # pytorch 2.5+ Intel XPU device + device = torch.device("xpu") + device_property = torch.xpu.get_device_properties(0) + device_name = device_property.name + total_mem = device_property.total_memory / 1024**3 + else: # pytorch tensors on cpu + device, device_name, total_mem, _ = get_device(gs_backend.cpu) + + case gs_backend.gpu: + if torch.cuda.is_available(): + return get_device(gs_backend.cuda) + elif get_platform() == "macOS": + return get_device(gs_backend.metal) + else: + return get_device(gs_backend.vulkan) + + case _: + device_name = get_cpu_name() + total_mem = psutil.virtual_memory().total / 1024**3 + device = torch.device("cpu") + + return device, device_name, total_mem, backend def get_src_dir(): From 50ebafd166c314813b1022e9edb11c2002ce6bea Mon Sep 17 00:00:00 2001 From: Lars Gleim 00627932 Date: Thu, 16 Jan 2025 11:24:28 +0100 Subject: [PATCH 2/2] Python 3.9 compatibility Switched from match / case to if / else syntax --- genesis/utils/misc.py | 71 +++++++++++++++++++++---------------------- 1 file changed, 35 insertions(+), 36 deletions(-) diff --git a/genesis/utils/misc.py b/genesis/utils/misc.py index fc5211f7..1c0082f3 100644 --- a/genesis/utils/misc.py +++ b/genesis/utils/misc.py @@ -98,45 +98,44 @@ def get_cpu_name(): def get_device(backend: gs_backend): - match backend: - case gs_backend.cuda: - if not torch.cuda.is_available(): - gs.raise_exception("cuda device not available") - - device = torch.device("cuda") - device_property = torch.cuda.get_device_properties(0) + if backend == gs_backend.cuda: + if not torch.cuda.is_available(): + gs.raise_exception("cuda device not available") + + device = torch.device("cuda") + device_property = torch.cuda.get_device_properties(0) + device_name = device_property.name + total_mem = device_property.total_memory / 1024**3 + + elif backend == gs_backend.metal: + if not torch.backends.mps.is_available(): + gs.raise_exception("metal device not available") + + # on mac, cpu and gpu are in the same device + _, device_name, total_mem, _ = get_device(gs_backend.cpu) + device = torch.device("mps") + + elif backend == gs_backend.vulkan: + if torch.xpu.is_available(): # pytorch 2.5+ Intel XPU device + device = torch.device("xpu") + device_property = torch.xpu.get_device_properties(0) device_name = device_property.name total_mem = device_property.total_memory / 1024**3 + else: # pytorch tensors on cpu + device, device_name, total_mem, _ = get_device(gs_backend.cpu) + + elif backend == gs_backend.gpu: + if torch.cuda.is_available(): + return get_device(gs_backend.cuda) + elif get_platform() == "macOS": + return get_device(gs_backend.metal) + else: + return get_device(gs_backend.vulkan) - case gs_backend.metal: - if not torch.backends.mps.is_available(): - gs.raise_exception("metal device not available") - - # on mac, cpu and gpu are in the same device - _, device_name, total_mem, _ = get_device(gs_backend.cpu) - device = torch.device("mps") - - case gs_backend.vulkan: - if torch.xpu.is_available(): # pytorch 2.5+ Intel XPU device - device = torch.device("xpu") - device_property = torch.xpu.get_device_properties(0) - device_name = device_property.name - total_mem = device_property.total_memory / 1024**3 - else: # pytorch tensors on cpu - device, device_name, total_mem, _ = get_device(gs_backend.cpu) - - case gs_backend.gpu: - if torch.cuda.is_available(): - return get_device(gs_backend.cuda) - elif get_platform() == "macOS": - return get_device(gs_backend.metal) - else: - return get_device(gs_backend.vulkan) - - case _: - device_name = get_cpu_name() - total_mem = psutil.virtual_memory().total / 1024**3 - device = torch.device("cpu") + else: + device_name = get_cpu_name() + total_mem = psutil.virtual_memory().total / 1024**3 + device = torch.device("cpu") return device, device_name, total_mem, backend