wehub-resource-sync
2269 lines
70 KiB
Python
2269 lines
70 KiB
Python
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
|
||
|
||
from __future__ import annotations
|
||
|
||
import ctypes
|
||
import importlib
|
||
import os
|
||
import re
|
||
import sys
|
||
import types
|
||
from contextlib import contextmanager
|
||
from typing import TYPE_CHECKING, Any, TypeAlias, overload
|
||
|
||
import paddle
|
||
from paddle.amp import autocast as _autocast
|
||
from paddle.amp.grad_scaler import GradScaler as _GradScaler
|
||
from paddle.base import core, framework
|
||
from paddle.base.framework import (
|
||
is_compiled_with_cinn,
|
||
is_compiled_with_cuda,
|
||
is_compiled_with_distribute,
|
||
is_compiled_with_rocm,
|
||
)
|
||
from paddle.tensor.creation import (
|
||
BFloat16Tensor,
|
||
BoolTensor,
|
||
ByteTensor,
|
||
CharTensor,
|
||
DoubleTensor,
|
||
FloatTensor,
|
||
HalfTensor,
|
||
IntTensor,
|
||
LongTensor,
|
||
ShortTensor,
|
||
)
|
||
|
||
from . import ( # noqa: F401
|
||
cuda,
|
||
xpu,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
from collections.abc import Generator
|
||
from contextlib import AbstractContextManager
|
||
from types import TracebackType
|
||
|
||
from paddle import IPUPlace as _IPUPlace, XPUPlace as _XPUPlace
|
||
from paddle._typing.device_like import PlaceLike
|
||
from paddle.base.core import Place
|
||
|
||
_InitStreamBase = core.CUDAStream | core.CustomDeviceStream | core.XPUStream
|
||
_InitEventBase = core.CUDAEvent | core.CustomDeviceEvent | core.XPUEvent
|
||
|
||
from paddle import CUDAPlace, CustomPlace
|
||
from paddle.base.libpaddle import _customDeviceProperties
|
||
|
||
_CustomPlaceLike: TypeAlias = (
|
||
CUDAPlace
|
||
| CustomPlace
|
||
| str # some string like "iluvatar_gpu" "metax_gpu:0", etc.
|
||
| int # some int like 0, 1, etc.
|
||
)
|
||
|
||
# Dynamically import device functions based on available devices
|
||
current_device_is_cpu = 0
|
||
if core.is_compiled_with_cuda():
|
||
from .cuda import (
|
||
create_event as _create_event_base,
|
||
create_stream as _create_stream_base,
|
||
device_count,
|
||
empty_cache,
|
||
get_device_properties as _get_device_properties,
|
||
get_rng_state,
|
||
manual_seed,
|
||
max_memory_allocated,
|
||
max_memory_reserved,
|
||
memory_allocated,
|
||
memory_reserved,
|
||
reset_max_memory_allocated,
|
||
reset_max_memory_reserved,
|
||
set_rng_state,
|
||
)
|
||
elif core.is_compiled_with_xpu():
|
||
from .xpu import (
|
||
create_event as _create_event_base,
|
||
create_stream as _create_stream_base,
|
||
device_count,
|
||
empty_cache,
|
||
get_device_properties as _get_device_properties,
|
||
get_rng_state,
|
||
manual_seed,
|
||
max_memory_allocated,
|
||
max_memory_reserved,
|
||
memory_allocated,
|
||
memory_reserved,
|
||
reset_max_memory_allocated,
|
||
reset_max_memory_reserved,
|
||
set_rng_state,
|
||
)
|
||
else:
|
||
if hasattr(core, 'get_all_custom_device_type'):
|
||
dev_types = core.get_all_custom_device_type()
|
||
else:
|
||
dev_types = []
|
||
if dev_types and core.is_compiled_with_custom_device(dev_types[0]):
|
||
from .custom_device import (
|
||
create_event as _create_event_base,
|
||
create_stream as _create_stream_base,
|
||
device_count,
|
||
empty_cache,
|
||
get_device_properties as _get_device_properties,
|
||
get_rng_state,
|
||
manual_seed,
|
||
max_memory_allocated,
|
||
max_memory_reserved,
|
||
memory_allocated,
|
||
memory_reserved,
|
||
reset_max_memory_allocated,
|
||
reset_max_memory_reserved,
|
||
set_rng_state,
|
||
)
|
||
else:
|
||
current_device_is_cpu = 1
|
||
from .cpu import (
|
||
device_count,
|
||
get_rng_state,
|
||
manual_seed,
|
||
max_memory_allocated,
|
||
max_memory_reserved,
|
||
reset_max_memory_allocated,
|
||
reset_max_memory_reserved,
|
||
set_rng_state,
|
||
)
|
||
|
||
|
||
__all__ = [
|
||
'get_cudnn_version',
|
||
'set_device',
|
||
'get_device',
|
||
'XPUPlace',
|
||
'IPUPlace',
|
||
'is_compiled_with_xpu',
|
||
'is_compiled_with_ipu',
|
||
'is_compiled_with_cinn',
|
||
'is_compiled_with_cuda',
|
||
'is_compiled_with_rocm',
|
||
'is_compiled_with_distribute',
|
||
'is_compiled_with_custom_device',
|
||
'get_all_device_type',
|
||
'get_all_custom_device_type',
|
||
'get_available_device',
|
||
'get_available_custom_device',
|
||
'get_device_properties',
|
||
'Stream',
|
||
'Event',
|
||
'current_stream',
|
||
'set_stream',
|
||
'stream_guard',
|
||
'device_guard',
|
||
'synchronize',
|
||
'device_count',
|
||
'empty_cache',
|
||
'max_memory_allocated',
|
||
'max_memory_reserved',
|
||
'reset_max_memory_allocated',
|
||
'reset_max_memory_reserved',
|
||
'memory_allocated',
|
||
'memory_reserved',
|
||
'is_available',
|
||
'is_current_stream_capturing',
|
||
'get_device_name',
|
||
'get_device_capability',
|
||
'get_rng_state',
|
||
'set_rng_state',
|
||
'FloatTensor',
|
||
'DoubleTensor',
|
||
'HalfTensor',
|
||
'BFloat16Tensor',
|
||
'ByteTensor',
|
||
'CharTensor',
|
||
'ShortTensor',
|
||
'IntTensor',
|
||
'LongTensor',
|
||
'BoolTensor',
|
||
'device',
|
||
'is_bf16_supported',
|
||
'manual_seed',
|
||
'reset_peak_memory_stats',
|
||
'ipc_collect',
|
||
'get_stream_from_external',
|
||
'StreamContext',
|
||
]
|
||
|
||
_cudnn_version = None
|
||
|
||
|
||
def is_compiled_with_custom_device(device_type: str) -> bool:
|
||
"""
|
||
|
||
Whether paddle was built with Paddle_CUSTOM_DEVICE .
|
||
|
||
Args:
|
||
device_type (str): the registered device type, like "npu".
|
||
|
||
Return:
|
||
bool, ``True`` if CustomDevice is supported, otherwise ``False``.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> support_npu = paddle.device.is_compiled_with_custom_device("npu")
|
||
|
||
"""
|
||
return core.is_compiled_with_custom_device(device_type)
|
||
|
||
|
||
def is_compiled_with_ipu() -> bool:
|
||
"""
|
||
|
||
Whether paddle was built with WITH_IPU=ON to support Graphcore IPU.
|
||
|
||
Returns (bool): `True` if IPU is supported, otherwise `False`.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> support_ipu = paddle.is_compiled_with_ipu()
|
||
|
||
"""
|
||
return core.is_compiled_with_ipu()
|
||
|
||
|
||
def IPUPlace() -> _IPUPlace:
|
||
"""
|
||
|
||
Return a Graphcore IPU Place
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:IPU)
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.set_device('ipu')
|
||
>>> place = paddle.device.IPUPlace()
|
||
|
||
"""
|
||
return core.IPUPlace()
|
||
|
||
|
||
def is_compiled_with_xpu() -> bool:
|
||
"""
|
||
|
||
Whether paddle was built with WITH_XPU=ON to support Baidu Kunlun
|
||
|
||
Returns (bool): whether paddle was built with WITH_XPU=ON
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> support_xpu = paddle.device.is_compiled_with_xpu()
|
||
|
||
"""
|
||
return core.is_compiled_with_xpu()
|
||
|
||
|
||
def XPUPlace(dev_id: int) -> _XPUPlace:
|
||
"""
|
||
|
||
Return a Baidu Kunlun Place
|
||
|
||
Args:
|
||
dev_id(int): Baidu Kunlun device id
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:XPU)
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.set_device('xpu')
|
||
>>> place = paddle.device.XPUPlace(0)
|
||
|
||
"""
|
||
return core.XPUPlace(dev_id)
|
||
|
||
|
||
def is_available() -> bool:
|
||
"""
|
||
Check whether **any supported device** is available in the current environment.
|
||
|
||
This function checks whether Paddle is built with support for at least one
|
||
type of accelerator (e.g., CUDA, XPU, CustomDevice) and whether there is
|
||
at least one device of that type available.
|
||
|
||
If any supported device is available, this function returns True. Otherwise,
|
||
it returns False.
|
||
|
||
Returns:
|
||
bool: True if there is at least one available device (GPU/XPU/CustomDevice),
|
||
False otherwise.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
>>> if paddle.device.is_available():
|
||
... print("At least one device is available")
|
||
... else:
|
||
... print("No supported devices available")
|
||
"""
|
||
return device_count() >= 1
|
||
|
||
|
||
def is_current_stream_capturing() -> bool:
|
||
"""
|
||
Check whether the current stream is in CUDA graph capturing state.
|
||
|
||
Returns:
|
||
bool: True if the current stream is capturing, False otherwise.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> if paddle.device.is_available():
|
||
... graph = paddle.device.cuda.graphs.CUDAGraph()
|
||
... graph.capture_begin()
|
||
... print(paddle.device.is_current_stream_capturing()) # True
|
||
... graph.capture_end()
|
||
"""
|
||
return core.is_cuda_graph_capturing()
|
||
|
||
|
||
def get_cudnn_version() -> int | None:
|
||
"""
|
||
|
||
This function return the version of cudnn. the return value is int which represents the
|
||
cudnn version. For example, if it return 7600, it represents the version of cudnn is 7.6.
|
||
|
||
Returns:
|
||
int: A int value which represents the cudnn version. If cudnn version is not installed, it return None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
>>> cudnn_version = paddle.device.get_cudnn_version()
|
||
|
||
|
||
|
||
"""
|
||
global _cudnn_version
|
||
if not core.is_compiled_with_cuda():
|
||
return None
|
||
if _cudnn_version is None:
|
||
cudnn_version = int(core.cudnn_version())
|
||
_cudnn_version = cudnn_version
|
||
if _cudnn_version < 0:
|
||
return None
|
||
else:
|
||
return cudnn_version
|
||
else:
|
||
return _cudnn_version
|
||
|
||
|
||
def device_to_place(device: Place | int | str | None = None) -> Place:
|
||
"""
|
||
Convert input device(Place | int | str | None) into corresponding Place object.
|
||
"""
|
||
device = _device_to_paddle(device)
|
||
device = _convert_to_place(device)
|
||
return device
|
||
|
||
|
||
def _convert_to_place(device: PlaceLike) -> Place:
|
||
if not isinstance(device, str):
|
||
if type(device) is core.Place:
|
||
if device.is_gpu_place():
|
||
return core.CUDAPlace(device.gpu_device_id())
|
||
elif device.is_cpu_place():
|
||
return core.CPUPlace()
|
||
elif device.is_xpu_place():
|
||
return core.XPUPlace(device.xpu_device_id())
|
||
elif device.is_custom_place():
|
||
return core.CustomPlace(
|
||
device.custom_device_type(), device.custom_device_id()
|
||
)
|
||
return device
|
||
|
||
lower_device = device.lower()
|
||
if lower_device.startswith("cuda"):
|
||
lower_device = lower_device.replace("cuda", "gpu")
|
||
if device in core.get_all_custom_device_type():
|
||
selected_devices = os.getenv(f"FLAGS_selected_{device}s", "0").split(
|
||
","
|
||
)
|
||
device_id = int(selected_devices[0])
|
||
place = core.CustomPlace(device, device_id)
|
||
elif lower_device == 'cpu':
|
||
place = core.CPUPlace()
|
||
elif lower_device == 'gpu' or lower_device == 'dcu':
|
||
if not core.is_compiled_with_cuda():
|
||
raise ValueError(
|
||
"The device should not be 'gpu', "
|
||
"since PaddlePaddle is not compiled with CUDA"
|
||
)
|
||
place = core.CUDAPlace(paddle.distributed.ParallelEnv().dev_id)
|
||
elif lower_device == 'xpu':
|
||
if not core.is_compiled_with_xpu():
|
||
raise ValueError(
|
||
"The device should not be 'xpu', "
|
||
"since PaddlePaddle is not compiled with XPU"
|
||
)
|
||
selected_xpus = os.getenv("FLAGS_selected_xpus", "0").split(",")
|
||
device_id = int(selected_xpus[0])
|
||
place = core.XPUPlace(device_id)
|
||
elif lower_device == 'ipu':
|
||
if not core.is_compiled_with_ipu():
|
||
raise ValueError(
|
||
"The device should not be 'ipu', "
|
||
"since PaddlePaddle is not compiled with IPU"
|
||
)
|
||
place = core.IPUPlace()
|
||
else:
|
||
available_gpu_device = re.match(r'gpu:\d+', lower_device) or re.match(
|
||
r'dcu:\d+', lower_device
|
||
)
|
||
available_xpu_device = re.match(r'xpu:\d+', lower_device)
|
||
if available_gpu_device:
|
||
if not core.is_compiled_with_cuda():
|
||
raise ValueError(
|
||
f"The device should not be {available_gpu_device}, since PaddlePaddle is "
|
||
"not compiled with CUDA"
|
||
)
|
||
device_info_list = device.split(':', 1)
|
||
device_id = device_info_list[1]
|
||
device_id = int(device_id)
|
||
place = core.CUDAPlace(device_id)
|
||
if available_xpu_device:
|
||
if not core.is_compiled_with_xpu():
|
||
raise ValueError(
|
||
f"The device should not be {available_xpu_device}, since PaddlePaddle is "
|
||
"not compiled with XPU"
|
||
)
|
||
device_info_list = device.split(':', 1)
|
||
device_id = device_info_list[1]
|
||
device_id = int(device_id)
|
||
place = core.XPUPlace(device_id)
|
||
if not available_gpu_device and not available_xpu_device:
|
||
device_info_list = device.split(':', 1)
|
||
device_type = device_info_list[0]
|
||
if device_type in core.get_all_custom_device_type():
|
||
device_id = device_info_list[1]
|
||
device_id = int(device_id)
|
||
place = core.CustomPlace(device_type, device_id)
|
||
else:
|
||
raise ValueError(
|
||
"The device must be a string which is like 'cpu', {}".format(
|
||
', '.join(
|
||
f"'{x}', '{x}:x'"
|
||
for x in [
|
||
'gpu',
|
||
'dcu',
|
||
'xpu',
|
||
'npu',
|
||
*core.get_all_custom_device_type(),
|
||
]
|
||
)
|
||
)
|
||
)
|
||
return place
|
||
|
||
|
||
class device:
|
||
r"""Context-manager that changes the selected device.
|
||
|
||
Args:
|
||
device (paddle.Place, int or str): device index to select.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
>>> print(paddle.device.get_device()) # gpu:0
|
||
>>> with paddle.device.device("cpu"):
|
||
... print(paddle.device.get_device()) # cpu
|
||
|
||
>>> # paddle.cuda.device is an alias of paddle.device.device
|
||
>>> with paddle.cuda.device("cpu"):
|
||
... print(paddle.device.get_device()) # cpu
|
||
>>> print(paddle.device.get_device())
|
||
"""
|
||
|
||
def __init__(self, device: Place | int | str | None = None):
|
||
self.place = device_to_place(device)
|
||
self.prev_place_str = "-1"
|
||
|
||
def __enter__(self):
|
||
self.prev_place_str = get_device()
|
||
set_device(self.place)
|
||
|
||
def __exit__(
|
||
self,
|
||
exc_type: type[BaseException] | None,
|
||
exc_value: BaseException | None,
|
||
traceback: types.TracebackType | None,
|
||
) -> bool | None:
|
||
set_device(self.prev_place_str)
|
||
return False
|
||
|
||
|
||
def current_device() -> int:
|
||
"""
|
||
Return the index of a currently selected device.
|
||
|
||
Returns:
|
||
int: The index of the currently selected device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
>>> device_id = paddle.device.current_device() # this is equivalent to paddle.cuda.current_device()
|
||
>>> print(f"Current device index: {device_id}")
|
||
"""
|
||
# Use paddle.device.get_device() to get the current device string
|
||
device_str = get_device()
|
||
|
||
# Parse the device string to extract the device index
|
||
# Format examples: 'gpu:0', 'xpu:0', 'custom_device:0'
|
||
if ':' in device_str:
|
||
device_id = int(device_str.split(':')[1])
|
||
else:
|
||
# If no device index is specified, default to 0
|
||
device_id = 0
|
||
|
||
return device_id
|
||
|
||
|
||
def is_bf16_supported(including_emulation: bool = True) -> bool:
|
||
"""
|
||
Return a bool indicating if the current CUDA/ROCm device supports dtype bfloat16.
|
||
|
||
Args:
|
||
including_emulation (bool = True): Whether to treat software-emulated BF16 as supported; if False, only native hardware BF16 support is considered.
|
||
|
||
Returns:
|
||
bool: A boolean value which indicates whether the current CUDA/ROCm device supports dtype bfloat16.
|
||
|
||
Examples:
|
||
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
>>> paddle.device.is_bf16_supported()
|
||
>>> # paddle.cuda.is_bf16_supported() is an alias of paddle.device.is_bf16_supported()
|
||
>>> paddle.cuda.is_bf16_supported()
|
||
|
||
"""
|
||
# including_emulation is not used here, but kept for compatibility with the original implementation
|
||
if core.is_bfloat16_supported(paddle.framework._current_expected_place_()):
|
||
return True
|
||
|
||
# If CUDA is not available, than it does not support bf16 either
|
||
if not is_available():
|
||
return False
|
||
|
||
device = get_device()
|
||
|
||
# Check for CUDA version and device compute capability.
|
||
# This is a fast way to check for it.
|
||
if not including_emulation:
|
||
return False
|
||
|
||
# Finally try to create a bfloat16 device.
|
||
try:
|
||
paddle.tensor([1.0], dtype=paddle.bfloat16, device=device)
|
||
return True
|
||
except:
|
||
return False
|
||
|
||
|
||
def set_device(device: PlaceLike | int) -> PlaceLike:
|
||
"""
|
||
|
||
Paddle supports running calculations on various types of devices, including CPU, GPU, XPU, NPU and IPU.
|
||
They are represented by string identifiers. This function can specify the global device
|
||
which the OP will run.
|
||
|
||
Args:
|
||
device(str, Place or int): This parameter determines the specific running device.
|
||
It can be ``cpu``, ``gpu``, ``xpu``, ``npu``, ``gpu:x``, ``xpu:x``, ``npu:x`` and ``ipu``,
|
||
where ``x`` is the index of the GPUs, XPUs or NPUs.
|
||
|
||
Returns:
|
||
Place,the Place to set.
|
||
|
||
Examples:
|
||
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
|
||
>>> paddle.device.set_device("cpu")
|
||
>>> x1 = paddle.ones(name='x1', shape=[1, 2], dtype='int32')
|
||
>>> print(x1.place)
|
||
Place(cpu)
|
||
|
||
>>> paddle.device.set_device("gpu:0")
|
||
>>> x2 = paddle.zeros(name='x2', shape=[1, 2], dtype='int32')
|
||
>>> print(x2.place)
|
||
Place(gpu:0)
|
||
>>> # x1 is still on cpu
|
||
>>> print(x1.place)
|
||
Place(cpu)
|
||
|
||
"""
|
||
place = device_to_place(device)
|
||
framework._set_expected_place(place)
|
||
if framework.in_dygraph_mode():
|
||
core.eager_set_device_id()
|
||
return place
|
||
|
||
|
||
@overload
|
||
def get_device(input: None = None) -> str: ...
|
||
|
||
|
||
@overload
|
||
def get_device(input: paddle.Tensor) -> int: ...
|
||
|
||
|
||
def get_device(input: paddle.Tensor | None = None) -> str | int:
|
||
"""
|
||
|
||
This function can get the current global device of the program is running.
|
||
It's a string which is like 'cpu', 'gpu:x', 'xpu:x' and 'npu:x'. if the global device is not
|
||
set, it will return a string which is 'gpu:x' when cuda is available or it
|
||
will return a string which is 'cpu' when cuda is not available.
|
||
|
||
Returns:
|
||
if input is Tensor, this function will return the device ID where the given Tensor is located.
|
||
int:
|
||
- -1, if the Tensor is on CPU.
|
||
- The device ID (e.g., 0, 1, ...) if the Tensor is on GPU.
|
||
|
||
if input is not Tensor, this function will return the device name where the program is running.
|
||
str:
|
||
- 'cpu': If the program is running on CPU.
|
||
- 'gpu:x': If the program is running on GPU, where `x` is the index of the GPU.
|
||
- 'xpu:x': If the program is running on XPU, where `x` is the index of the XPU.
|
||
- 'npu:x': If the program is running on NPU, where `x` is the index of
|
||
Examples:
|
||
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> device = paddle.device.get_device()
|
||
|
||
>>> x_cpu = paddle.to_tensor([1, 2, 3], place=paddle.CPUPlace())
|
||
>>> id = paddle.get_device(x_cpu) # -1
|
||
|
||
|
||
|
||
"""
|
||
if isinstance(input, paddle.Tensor):
|
||
if 'cpu' in str(input.place):
|
||
return -1
|
||
return input.place.gpu_device_id()
|
||
device = ''
|
||
place = framework._current_expected_place_()
|
||
if isinstance(place, core.CPUPlace):
|
||
device = 'cpu'
|
||
elif isinstance(place, core.CUDAPlace):
|
||
device_id = place.get_device_id()
|
||
device = 'gpu:' + str(device_id)
|
||
elif isinstance(place, core.XPUPlace):
|
||
device_id = place.get_device_id()
|
||
device = 'xpu:' + str(device_id)
|
||
elif isinstance(place, core.IPUPlace):
|
||
num_devices = core.get_ipu_device_count()
|
||
device = f"ipus:{{0-{num_devices - 1}}}"
|
||
elif isinstance(place, core.CustomPlace):
|
||
device_id = place.get_device_id()
|
||
device_type = place.get_device_type()
|
||
device = device_type + ':' + str(device_id)
|
||
else:
|
||
raise ValueError(f"The device specification {place} is invalid")
|
||
|
||
return device
|
||
|
||
|
||
def get_default_device() -> paddle.device:
|
||
"""
|
||
Returns:
|
||
str: The default device for PaddlePaddle.
|
||
Example:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
>>> print(paddle.get_default_device())
|
||
"""
|
||
dev = get_device()
|
||
# Only replace exact "gpu" device type, not substrings in custom device names
|
||
if dev.startswith("gpu"):
|
||
dev = "cuda" + dev[3:]
|
||
return paddle.device(dev)
|
||
|
||
|
||
def set_default_device(device: PlaceLike | int | None = None) -> None:
|
||
"""
|
||
Paddle supports running calculations on various types of devices, including CPU, GPU, XPU, NPU and IPU.
|
||
This function can specify the global device which the OP will run.
|
||
|
||
Args:
|
||
device(str | Place | paddle.device | int, optional): This parameter determines the specific running device.
|
||
It can be ``cpu``, ``gpu``, ``xpu``, ``npu``, ``gpu:x``, ``xpu:x``, ``npu:x`` and ``ipu``,
|
||
where ``x`` is the index of the GPUs, XPUs or NPUs. Defaults is ``None``, which means current device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.set_device("cpu")
|
||
"""
|
||
set_device(device)
|
||
|
||
|
||
def get_all_device_type() -> list[str]:
|
||
"""
|
||
|
||
Get all available device types.
|
||
|
||
Returns:
|
||
A list of all available device types.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.get_all_device_type()
|
||
|
||
>>> # Case 1: paddlepaddle-cpu package installed, and no custom device registered.
|
||
>>> # Output: []
|
||
|
||
>>> # Case 2: paddlepaddle-gpu package installed, and no custom device registered.
|
||
>>> # Output: ['gpu']
|
||
|
||
>>> # Case 3: paddlepaddle-cpu package installed, and custom device 'CustomCPU' is registered.
|
||
>>> # Output: ['CustomCPU']
|
||
|
||
>>> # Case 4: paddlepaddle-gpu package installed, and custom device 'CustomCPU' and 'CustomGPU' is registered.
|
||
>>> # Output: ['gpu', 'CustomCPU', 'CustomGPU']
|
||
|
||
"""
|
||
return core.get_all_device_type()
|
||
|
||
|
||
def get_all_custom_device_type() -> list[str] | None:
|
||
"""
|
||
|
||
Get all available custom device types.
|
||
|
||
Returns:
|
||
A list of all available custom device types.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.get_all_custom_device_type()
|
||
|
||
>>> # Case 1: paddlepaddle-gpu package installed, and no custom device registered.
|
||
>>> # Output: None
|
||
|
||
>>> # Case 2: paddlepaddle-gpu package installed, and custom device 'CustomCPU' and 'CustomGPU' is registered.
|
||
>>> # Output: ['CustomCPU', 'CustomGPU']
|
||
|
||
"""
|
||
return core.get_all_custom_device_type()
|
||
|
||
|
||
def get_available_device() -> list[str]:
|
||
"""
|
||
|
||
Get all available devices.
|
||
|
||
Returns:
|
||
A list of all available devices.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.get_available_device()
|
||
|
||
>>> # Case 1: paddlepaddle-cpu package installed, and no custom device registered.
|
||
>>> # Output: []
|
||
|
||
>>> # Case 2: paddlepaddle-gpu package installed, and no custom device registered.
|
||
>>> # Output: ['gpu:0', 'gpu:1']
|
||
|
||
>>> # Case 3: paddlepaddle-cpu package installed, and custom device 'CustomCPU' is registered.
|
||
>>> # Output: ['CustomCPU']
|
||
|
||
>>> # Case 4: paddlepaddle-gpu package installed, and custom device 'CustomCPU' and 'CustomGPU' is registered.
|
||
>>> # Output: ['gpu:0', 'gpu:1', 'CustomCPU', 'CustomGPU:0', 'CustomGPU:1']
|
||
|
||
"""
|
||
return core.get_available_device()
|
||
|
||
|
||
def get_available_custom_device() -> list[str] | None:
|
||
"""
|
||
|
||
Get all available custom devices.
|
||
|
||
Returns:
|
||
A list of all available custom devices.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.get_available_custom_device()
|
||
|
||
>>> # Case 1: paddlepaddle-gpu package installed, and no custom device registered.
|
||
>>> # Output: None
|
||
|
||
>>> # Case 2: paddlepaddle-gpu package installed, and custom device 'CustomCPU' and 'CustomGPU' is registered.
|
||
>>> # Output: ['CustomCPU', 'CustomGPU:0', 'CustomGPU:1']
|
||
|
||
"""
|
||
return core.get_available_custom_device()
|
||
|
||
|
||
def get_device_properties(
|
||
device: _CustomPlaceLike | None = None,
|
||
) -> _customDeviceProperties:
|
||
"""
|
||
|
||
Return the properties of given device.
|
||
|
||
Args:
|
||
device(|paddle.CustomPlace|int|str|None, optional): The device, the id of the device or
|
||
the string name of device like npu:x' which to get the properties of the
|
||
device from. If device is None, the device is the current device.
|
||
Default: None.
|
||
|
||
Returns:
|
||
_customDeviceProperties: The properties of the device which include ASCII string
|
||
identifying device, major compute capability, minor compute capability, global
|
||
memory available and the number of multiprocessors on the device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # import paddle
|
||
>>> # paddle.device.set_device('npu')
|
||
>>> # paddle.device.get_device_properties('npu:0')
|
||
>>> # _customDeviceProperties(name='', major=0, minor=0, total_memory=0MB, multi_processor_count=0)
|
||
|
||
>>> # paddle.device.get_device_properties('npu')
|
||
>>> # _customDeviceProperties(name='', major=0, minor=0, total_memory=0MB, multi_processor_count=0)
|
||
"""
|
||
device = device_to_place(device)
|
||
return _get_device_properties(device)
|
||
|
||
|
||
def get_device_module(device: _CustomPlaceLike = None):
|
||
"""
|
||
Returns the Paddle module associated with a given device.
|
||
|
||
Args:
|
||
device (_CustomPlaceLike, optional): The device to query.
|
||
Can be one of the following:
|
||
- paddle.Place object (e.g., paddle.CUDAPlace(0))
|
||
- str (e.g., "gpu:0", "xpu", "npu")
|
||
- int (device index, e.g., 0 -> "gpu:0")
|
||
- None (use current expected place)
|
||
|
||
Returns:
|
||
module: The corresponding Paddle device module (e.g., paddle.cuda, paddle.device.xpu)
|
||
|
||
Raises:
|
||
RuntimeError: If the device type is CPU (Paddle does not expose `paddle.cpu`)
|
||
or if no matching device module is found.
|
||
|
||
Example:
|
||
.. code-block:: pycon
|
||
|
||
>>> paddle.get_device_module("gpu:0")
|
||
<module 'paddle.cuda' ...>
|
||
|
||
>>> # paddle.get_device_module(paddle.XPUPlace(0))
|
||
>>> # <module 'paddle.device.xpu' ...>
|
||
"""
|
||
device = _device_to_paddle(device)
|
||
if isinstance(device, str):
|
||
device = device.lower().split(':')[0]
|
||
custom_device_types = {
|
||
"metax_gpu",
|
||
"biren_gpu",
|
||
"custom_cpu",
|
||
"gcu",
|
||
"iluvatar_gpu",
|
||
"intel_gpu",
|
||
"intel_hpu",
|
||
"mlu",
|
||
"mps",
|
||
"npu",
|
||
"sdaa",
|
||
}
|
||
if device in ("cuda", "gpu"):
|
||
return paddle.cuda
|
||
elif device == "xpu":
|
||
return paddle.device.xpu
|
||
elif device in custom_device_types:
|
||
return paddle.device.custom_device
|
||
elif device == "cpu":
|
||
return paddle.device.cpu
|
||
else:
|
||
raise RuntimeError(f"Unsupported device type: {device}")
|
||
|
||
place = (
|
||
paddle.framework._current_expected_place_()
|
||
if device is None
|
||
else _convert_to_place(device)
|
||
)
|
||
|
||
place_to_module = {
|
||
core.CUDAPlace: paddle.cuda,
|
||
core.CustomPlace: paddle.device.custom_device,
|
||
core.XPUPlace: paddle.device.xpu,
|
||
core.CPUPlace: paddle.device,
|
||
}
|
||
|
||
for place_type, module in place_to_module.items():
|
||
if isinstance(place, place_type):
|
||
return module
|
||
|
||
|
||
def get_device_name(
|
||
device: _CustomPlaceLike | None = None,
|
||
) -> str:
|
||
"""
|
||
|
||
Return the properties of given device.
|
||
|
||
Args:
|
||
device(|paddle.CustomPlace|int|str|None, optional): The device, the id of the device or
|
||
the string name of device like npu:x' which to get the properties of the
|
||
device from. If device is None, the device is the current device.
|
||
Default: None.
|
||
|
||
Returns:
|
||
str: The name of the CUDA device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
>>> name = paddle.device.get_device_name()
|
||
>>> print(name)
|
||
"""
|
||
return get_device_properties(device).name
|
||
|
||
|
||
def get_device_capability(
|
||
device: _CustomPlaceLike | None = None,
|
||
) -> tuple[int, int]:
|
||
"""
|
||
|
||
Return the device_capability of given device.
|
||
|
||
Args:
|
||
device(|paddle.CustomPlace|int|str|None, optional): The device, the id of the device or
|
||
the string name of device like npu:x' which to get the properties of the
|
||
device from. If device is None, the device is the current device.
|
||
Default: None.
|
||
|
||
Returns:
|
||
str: The device_capability of given device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
>>> cap = paddle.device.get_device_capability()
|
||
>>> print(cap)
|
||
"""
|
||
prop = get_device_properties(device)
|
||
return prop.major, prop.minor
|
||
|
||
|
||
def extract_device_id(device: _CustomPlaceLike, op_name: str) -> int:
|
||
'''
|
||
Return the id of the given device. It is just a utility that will not be exposed to users.
|
||
|
||
Args:
|
||
device(paddle.CUDAPlace|paddle.CustomPlace|int|str): The device, the id of the device or
|
||
the string name of device like 'gpu:x' or 'custom_device:x'.
|
||
Default: None.
|
||
|
||
Return:
|
||
int: The id of the given device. If device is None, return the id of current device.
|
||
'''
|
||
if device is None:
|
||
return core.get_cuda_current_device_id()
|
||
|
||
if isinstance(device, int):
|
||
device_id = device
|
||
if core.is_compiled_with_cuda():
|
||
device_type = 'gpu'
|
||
else:
|
||
device_type = None
|
||
available_custom_devices = core.get_available_custom_device()
|
||
if len(available_custom_devices) == 1:
|
||
if device == 0:
|
||
device_type = available_custom_devices[0]
|
||
else:
|
||
raise ValueError(
|
||
f"Device id {device} not found in available_custom_devices: [{available_custom_devices[0]}:0]"
|
||
)
|
||
else:
|
||
for d in available_custom_devices:
|
||
dev_type, dev_id = d.split(':')
|
||
if int(dev_id) == device:
|
||
device_type = dev_type
|
||
if device_type is None:
|
||
raise ValueError(
|
||
f"Device id {device} not found in available_custom_devices: {available_custom_devices}"
|
||
)
|
||
elif isinstance(device, core.CUDAPlace):
|
||
device_type = 'gpu'
|
||
device_id = device.get_device_id()
|
||
elif isinstance(device, core.CustomPlace):
|
||
device_type = device.get_device_type()
|
||
device_id = device.get_device_id()
|
||
elif isinstance(device, str):
|
||
if device.startswith('gpu:'):
|
||
device_id = int(device[4:])
|
||
elif (
|
||
':' in device
|
||
): # handle custom device formats like npu:0, metax_gpu:1
|
||
device_type, device_id_str = device.split(':', 1)
|
||
device_id = int(device_id_str)
|
||
else:
|
||
raise ValueError(
|
||
f"The current string {device} is not expected. Because {op_name} only support string which is like 'gpu:x' or '<custom_device>:x'. "
|
||
"Please input appropriate string again!"
|
||
)
|
||
else:
|
||
raise ValueError(
|
||
f"The device type {device} is not expected. Because {op_name} only support int, str (format 'gpu:x' or '<custom_device>:x'), paddle.CUDAPlace or paddle.CustomPlace. "
|
||
"Please input appropriate device again!"
|
||
)
|
||
|
||
assert device_id >= 0, (
|
||
f"The device id must be not less than 0, but got id = {device_id}."
|
||
)
|
||
|
||
if core.is_compiled_with_cuda():
|
||
assert device_id < device_count(), (
|
||
f"The device id {device_id} exceeds gpu card number {device_count()}"
|
||
)
|
||
else:
|
||
assert device_id < core.get_custom_device_count(device_type), (
|
||
f"The device id {device_id} exceeds {device_type} device card number {core.get_custom_device_count(device_type)}"
|
||
)
|
||
return device_id
|
||
|
||
|
||
class Event:
|
||
'''
|
||
|
||
A device event wrapper around StreamBase.
|
||
|
||
Args:
|
||
enable_timing (bool, optional): indicates if the event should measure time, default is False
|
||
blocking (bool, optional): if True, ``wait`` will be blocking, default is False
|
||
interprocess (bool): if True, the event can be shared between processes, default is False
|
||
|
||
Returns:
|
||
Event: The event.
|
||
|
||
Note:
|
||
The `device` parameter has been removed in the latest version. The event will always use the current device context.
|
||
Previously, you could specify the device like:
|
||
```python
|
||
# Old usage (no longer supported)
|
||
e = paddle.device.Event(device="gpu:0")
|
||
```
|
||
Now it will automatically use the current device:
|
||
```python
|
||
# New usage
|
||
paddle.set_device("gpu:0") # Set device first
|
||
e = paddle.device.Event() # Will use gpu:0
|
||
```
|
||
|
||
paddle.device.Event is equivalent to paddle.cuda.Event.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> e1 = paddle.device.Event() # Uses current device (custom_cpu)
|
||
>>>
|
||
>>> # Old usage (no longer supported):
|
||
>>> # e2 = paddle.device.Event('custom_cpu')
|
||
>>> # e3 = paddle.device.Event('custom_cpu:0')
|
||
>>> # e4 = paddle.device.Event(paddle.CustomPlace('custom_cpu', 0))
|
||
>>>
|
||
>>> # New equivalent usage:
|
||
>>> paddle.set_device('custom_cpu:0')
|
||
>>> e5 = paddle.device.Event() # Uses custom_cpu:0
|
||
|
||
'''
|
||
|
||
device: PlaceLike | None
|
||
enable_timing: bool
|
||
event_base: _InitEventBase
|
||
|
||
def __init__(
|
||
self,
|
||
enable_timing: bool = False,
|
||
blocking: bool = False,
|
||
interprocess: bool = False,
|
||
) -> None:
|
||
self.device = paddle.framework._current_expected_place_()
|
||
|
||
device_id = (
|
||
self.device.get_device_id()
|
||
if hasattr(self.device, 'get_device_id')
|
||
else None
|
||
)
|
||
device_type = (
|
||
self.device.get_device_type()
|
||
if hasattr(self.device, 'get_device_type')
|
||
else None
|
||
)
|
||
|
||
self.event_base = _create_event_base(
|
||
enable_timing=enable_timing,
|
||
blocking=blocking,
|
||
interprocess=interprocess,
|
||
device_type=device_type,
|
||
device_id=device_id,
|
||
)
|
||
|
||
def record(self, stream: Stream | None = None) -> None:
|
||
'''
|
||
|
||
Records the event in a given stream.
|
||
|
||
Args:
|
||
stream(Stream, optional): The given stream. By default, stream is None,
|
||
event will be recorded in current_stream.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> e = paddle.device.Event()
|
||
>>> e.record()
|
||
|
||
>>> s = paddle.device.Stream()
|
||
>>> e.record(s)
|
||
|
||
'''
|
||
if stream is None:
|
||
stream = current_stream(self.device)
|
||
|
||
self.event_base.record(stream.stream_base)
|
||
|
||
def query(self) -> bool:
|
||
'''
|
||
|
||
Checks if all work currently captured by event has completed.
|
||
|
||
Returns:
|
||
bool: Whether all work currently captured by event has completed.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> e = paddle.device.Event()
|
||
>>> e.record()
|
||
>>> e.query()
|
||
|
||
'''
|
||
return self.event_base.query()
|
||
|
||
def elapsed_time(self, end_event: Event) -> int:
|
||
'''
|
||
|
||
Returns the time elapsed in milliseconds after the event was
|
||
recorded and before the end_event was recorded.
|
||
|
||
Returns:
|
||
int: The time.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> e1 = paddle.device.Event()
|
||
>>> e1.record()
|
||
|
||
>>> e2 = paddle.device.Event()
|
||
>>> e2.record()
|
||
>>> e1.elapsed_time(e2)
|
||
|
||
'''
|
||
return self.event_base.elapsed_time(end_event.event_base)
|
||
|
||
def synchronize(self) -> None:
|
||
'''
|
||
|
||
Waits for the event to complete.
|
||
Waits until the completion of all work currently captured in this event.
|
||
This prevents the CPU thread from proceeding until the event completes.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> e = paddle.device.Event()
|
||
>>> e.record()
|
||
>>> e.synchronize()
|
||
|
||
'''
|
||
self.event_base.synchronize()
|
||
|
||
def __repr__(self) -> str:
|
||
return f"Event(device={self.device}, event_base={self.event_base})"
|
||
|
||
|
||
class Stream:
|
||
'''
|
||
|
||
A device stream wrapper around StreamBase.
|
||
paddle.cuda.Stream() is equivalent to paddle.device.Stream().
|
||
|
||
Args:
|
||
device(str|paddle.CUDAPlace(n)|paddle.CustomPlace(n)|None): Which device the stream run on. If device is None, the device is the current device. Default: None.
|
||
It can be ``gpu``, ``gpu:x``, ``custom_device``, ``custom_device:x``, where ``custom_device`` is the name of CustomDevice,
|
||
where ``x`` is the index of the GPUs, XPUs. And it can be paddle.CUDAPlace(n) or paddle.CustomPlace(n).
|
||
priority(int, optional): priority of the CUDA stream. Can be either
|
||
1 or -1 (high priority) or 0 or 2 (low priority). By default, streams have
|
||
priority 2.
|
||
|
||
Returns:
|
||
Stream: The stream.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s1 = paddle.device.Stream()
|
||
>>> s2 = paddle.device.Stream('custom_cpu')
|
||
>>> s3 = paddle.device.Stream('custom_cpu:0')
|
||
>>> s4 = paddle.device.Stream(paddle.CustomPlace('custom_cpu', 0))
|
||
|
||
'''
|
||
|
||
stream_base: _InitStreamBase
|
||
device: PlaceLike | int
|
||
_priority_map: dict[int, int] = {-1: 1, 0: 2, 1: 1, 2: 2}
|
||
|
||
def __init__(
|
||
self,
|
||
device: PlaceLike | int | None = None,
|
||
priority: int = 2,
|
||
stream_base: _InitStreamBase | None = None,
|
||
) -> None:
|
||
if stream_base is not None:
|
||
if isinstance(
|
||
stream_base,
|
||
(core.CUDAStream, core.CustomDeviceStream, core.XPUStream),
|
||
):
|
||
self.stream_base = stream_base
|
||
self.device = stream_base.place
|
||
else:
|
||
raise TypeError(
|
||
"stream_base should be CUDAStream, XPUStream, CustomDeviceStream"
|
||
)
|
||
return
|
||
self.device = device_to_place(device)
|
||
|
||
device_id = (
|
||
self.device.get_device_id()
|
||
if hasattr(self.device, 'get_device_id')
|
||
else None
|
||
)
|
||
device_type = (
|
||
self.device.get_device_type()
|
||
if hasattr(self.device, 'get_device_type')
|
||
else None
|
||
)
|
||
priority = self._priority_map.get(priority, 2)
|
||
self.stream_base = _create_stream_base(
|
||
device_id=device_id,
|
||
priority=priority,
|
||
blocking=False,
|
||
device_type=device_type,
|
||
)
|
||
|
||
def wait_event(self, event: Event) -> None:
|
||
'''
|
||
|
||
Makes all future work submitted to the stream wait for an event.
|
||
|
||
Args:
|
||
event (Event): an event to wait for.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s1 = paddle.device.Stream()
|
||
>>> s2 = paddle.device.Stream()
|
||
>>> e = paddle.device.Event()
|
||
>>> e.record(s1)
|
||
>>> s2.wait_event(e)
|
||
|
||
'''
|
||
self.stream_base.wait_event(event.event_base)
|
||
|
||
def wait_stream(self, stream: Stream) -> None:
|
||
'''
|
||
|
||
Synchronizes with another stream.
|
||
All future work submitted to this stream will wait until all kernels
|
||
submitted to a given stream at the time of call complete.
|
||
|
||
Args:
|
||
stream (Stream): a stream to synchronize.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s1 = paddle.device.Stream()
|
||
>>> s2 = paddle.device.Stream()
|
||
>>> s1.wait_stream(s2)
|
||
|
||
'''
|
||
self.stream_base.wait_stream(stream.stream_base)
|
||
|
||
def record_event(self, event: Event | None = None) -> Event:
|
||
'''
|
||
|
||
Records an event.
|
||
|
||
Args:
|
||
event (Event, optional): event to record. If not given, a new one
|
||
will be allocated.
|
||
|
||
Returns:
|
||
Event: Recorded event.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s = paddle.device.Stream()
|
||
>>> e1 = s.record_event()
|
||
|
||
>>> e2 = paddle.device.Event()
|
||
>>> s.record_event(e2)
|
||
|
||
'''
|
||
if event is None:
|
||
event = Event()
|
||
event.record(self)
|
||
return event
|
||
|
||
def query(self) -> bool:
|
||
'''
|
||
|
||
Checks if all the work submitted has been completed.
|
||
|
||
Returns:
|
||
bool: Whether all kernels in this stream are completed.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s = paddle.device.Stream()
|
||
>>> s.query()
|
||
|
||
'''
|
||
return self.stream_base.query()
|
||
|
||
def synchronize(self) -> None:
|
||
'''
|
||
|
||
Wait for all the kernels in this stream to complete.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s = paddle.device.Stream()
|
||
>>> s.synchronize()
|
||
|
||
'''
|
||
self.stream_base.synchronize()
|
||
|
||
@property
|
||
def _as_parameter_(self):
|
||
if isinstance(self.stream_base, core.CUDAStream):
|
||
return ctypes.c_void_p(self.stream_base.cuda_stream)
|
||
elif isinstance(self.stream_base, core.XPUStream):
|
||
return ctypes.c_void_p(self.stream_base.xpu_stream)
|
||
else:
|
||
return ctypes.c_void_p(self.stream_base.raw_stream)
|
||
|
||
@property
|
||
def cuda_stream(self) -> int:
|
||
assert isinstance(self.stream_base, core.CUDAStream), (
|
||
"cuda_stream is only available for CUDA streams"
|
||
)
|
||
return self.stream_base.cuda_stream
|
||
|
||
def __cuda_stream__(self):
|
||
"""
|
||
CUDA Stream protocol described at
|
||
https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol
|
||
|
||
Returns a tuple of (protocol_version, cudaStream_t)
|
||
"""
|
||
return (0, self.stream_base.raw_stream)
|
||
|
||
def __eq__(self, o: Stream | None) -> bool:
|
||
if isinstance(o, Stream):
|
||
return super().__eq__(o)
|
||
return False
|
||
|
||
def __hash__(self) -> int:
|
||
return hash((self.stream_base, self.device))
|
||
|
||
def __repr__(self) -> str:
|
||
return f'<paddle.device.Stream device={self.device} stream={self._as_parameter_.value:#x}>'
|
||
|
||
|
||
def _device_to_paddle(
|
||
dev: Place | int | str | None = None,
|
||
):
|
||
if isinstance(dev, int):
|
||
if dev < 0:
|
||
raise ValueError(f"Device index must be non-negative, got {dev}")
|
||
current_place = get_device() # e.g. "gpu:0", "cpu"
|
||
if current_place == "cpu":
|
||
if dev != 0:
|
||
raise ValueError(f"CPU device only supports index 0, got {dev}")
|
||
return "cpu"
|
||
device_type = current_place.split(":")[0]
|
||
return f"{device_type}:{dev}"
|
||
elif isinstance(dev, str):
|
||
cleaned_device = dev.strip()
|
||
return (
|
||
cleaned_device.replace("cuda", "gpu")
|
||
if "cuda" in cleaned_device
|
||
else cleaned_device
|
||
)
|
||
elif dev is None:
|
||
return get_device()
|
||
else:
|
||
return dev
|
||
|
||
|
||
def current_stream(device: PlaceLike | None = None) -> Stream:
|
||
'''
|
||
|
||
Return the current stream by the device.
|
||
|
||
Args:
|
||
device(str|paddle.CUDAPlace(n)|paddle.CustomPlace(n)): The device which want to get stream from. If device is None, the device is the current device. Default: None.
|
||
It can be ``gpu``, ``gpu:x``, ``custom_device``, ``custom_device:x``, where ``custom_device`` is the name of CustomDevice,
|
||
where ``x`` is the index of the GPUs, CustomDevices. And it can be paddle.CUDAPlace(n) or paddle.CustomPlace(n).
|
||
|
||
Returns:
|
||
Stream: The stream to the device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s1 = paddle.device.current_stream()
|
||
>>> s2 = paddle.device.current_stream("custom_cpu:0")
|
||
>>> place = paddle.CustomPlace('custom_cpu', 0)
|
||
>>> s3 = paddle.device.current_stream(place)
|
||
|
||
'''
|
||
if device is None:
|
||
place = paddle.framework._current_expected_place_()
|
||
elif isinstance(device, str):
|
||
place = paddle.device._convert_to_place(device)
|
||
else:
|
||
place = device
|
||
|
||
if paddle.is_compiled_with_cuda() and isinstance(place, paddle.CUDAPlace):
|
||
return Stream(
|
||
stream_base=core._get_current_stream(place.get_device_id())
|
||
)
|
||
elif paddle.is_compiled_with_xpu() and isinstance(place, paddle.XPUPlace):
|
||
return Stream(
|
||
stream_base=core._xpu_get_current_stream(place.get_device_id())
|
||
)
|
||
elif isinstance(place, paddle.CustomPlace):
|
||
return Stream(
|
||
stream_base=core._get_current_custom_device_stream(
|
||
place.get_device_type(), place.get_device_id()
|
||
)
|
||
)
|
||
else:
|
||
raise TypeError(
|
||
"device should be gpu, xpu, {}".format(
|
||
",".join(paddle.device.get_all_custom_device_type())
|
||
)
|
||
)
|
||
|
||
|
||
def set_stream(stream: Stream) -> Stream:
|
||
'''
|
||
|
||
Set the current stream.
|
||
|
||
Args:
|
||
stream(Stream): The selected stream.
|
||
|
||
Returns:
|
||
Stream: The previous stream.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s = paddle.device.Stream()
|
||
>>> # paddle.cuda.set_stream(s) is equivalent to paddle.device.set_stream(s)
|
||
>>> paddle.device.set_stream(s)
|
||
|
||
'''
|
||
|
||
prev_stream = current_stream(stream.stream_base.place)
|
||
|
||
if paddle.is_compiled_with_cuda() and isinstance(
|
||
stream.stream_base.place, paddle.CUDAPlace
|
||
):
|
||
core._set_current_stream(stream.stream_base)
|
||
elif paddle.is_compiled_with_xpu() and isinstance(
|
||
stream.stream_base.place, paddle.XPUPlace
|
||
):
|
||
core._xpu_set_current_stream(stream.stream_base.idx)
|
||
elif isinstance(stream.stream_base.place, paddle.CustomPlace):
|
||
core._set_current_custom_device_stream(
|
||
stream.stream_base.place.get_device_type(),
|
||
stream.stream_base.place.get_device_id(),
|
||
stream.stream_base,
|
||
)
|
||
else:
|
||
raise TypeError(
|
||
"device should be gpu, xpu, {}".format(
|
||
",".join(paddle.device.get_all_custom_device_type())
|
||
)
|
||
)
|
||
|
||
return prev_stream
|
||
|
||
|
||
class stream_guard:
|
||
'''
|
||
|
||
Notes:
|
||
This API only supports dynamic graph mode currently.
|
||
A context manager that specifies the current stream context by the given stream.
|
||
|
||
Args:
|
||
stream(Stream, optional): the selected stream. If stream is None, just yield.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> s = paddle.device.Stream()
|
||
>>> data1 = paddle.ones(shape=[20])
|
||
>>> data2 = paddle.ones(shape=[20])
|
||
>>> data3 = data1 + data2
|
||
>>> # This is equivalent to paddle.cuda.StreamContext(s) and paddle.device.StreamContext(s)
|
||
>>> with paddle.device.stream_guard(s):
|
||
... s.wait_stream(paddle.device.default_stream()) # type: ignore[attr-defined]
|
||
... data4 = data1 + data3
|
||
|
||
'''
|
||
|
||
stream: Stream | None
|
||
|
||
def __init__(self, stream: Stream | None = None) -> None:
|
||
self.stream = stream
|
||
|
||
def __enter__(self) -> None:
|
||
cur_stream = self.stream
|
||
if cur_stream is None:
|
||
return
|
||
|
||
self.src_prev_stream = current_stream(cur_stream.device)
|
||
if self.src_prev_stream.device != cur_stream.device:
|
||
self.tmp_place = paddle.base.framework._current_expected_place_()
|
||
paddle.base.framework._set_expected_place(cur_stream.device)
|
||
self.dst_prev_stream = current_stream(cur_stream.device)
|
||
set_stream(cur_stream)
|
||
else:
|
||
set_stream(cur_stream)
|
||
|
||
def __exit__(
|
||
self,
|
||
exc_type: type[BaseException] | None,
|
||
exc_val: BaseException | None,
|
||
exc_tb: TracebackType | None,
|
||
) -> None:
|
||
cur_stream = self.stream
|
||
if cur_stream is None:
|
||
return
|
||
|
||
if self.src_prev_stream.device != cur_stream.device:
|
||
set_stream(self.dst_prev_stream)
|
||
paddle.base.framework._set_expected_place(self.tmp_place)
|
||
set_stream(self.src_prev_stream)
|
||
else:
|
||
set_stream(self.src_prev_stream)
|
||
|
||
|
||
StreamContext = stream_guard
|
||
|
||
|
||
def stream(stream: Stream | None) -> stream_guard:
|
||
'''
|
||
|
||
Notes:
|
||
This API only supports dynamic graph mode currently.
|
||
A context manager that specifies the current stream context by the given stream.
|
||
|
||
Args:
|
||
stream(Stream, optional): the selected stream. If stream is None, just yield.
|
||
|
||
Returns:
|
||
None.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
|
||
>>> paddle.set_device('cuda')
|
||
>>> s = paddle.device.Stream()
|
||
>>> data1 = paddle.ones(shape=[20])
|
||
>>> data2 = paddle.ones(shape=[20])
|
||
>>> data3 = data1 + data2
|
||
|
||
>>> with paddle.device.stream(s): # this is equivalent to paddle.cuda.stream(s)
|
||
... s.wait_stream(paddle.cuda.current_stream())
|
||
... data4 = data1 + data3
|
||
>>> print(data4)
|
||
|
||
'''
|
||
return StreamContext(stream)
|
||
|
||
|
||
class device_guard:
|
||
'''
|
||
|
||
Notes:
|
||
This API only supports dynamic graph mode currently.
|
||
|
||
A context manager that specifies the current device context by the given device.
|
||
|
||
Args:
|
||
device(PlaceLike): The specified device.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
|
||
>>> # Set the global default device to CPU
|
||
>>> paddle.set_device("cpu")
|
||
>>> # Temporarily switch to GPU:0 using device_guard with string input
|
||
>>> with paddle.device.device_guard("gpu:0"):
|
||
... x = paddle.randn([4, 4]) # Create a Tensor on GPU:0
|
||
... x = x.tanh() * 2 # Perform computation on GPU:0
|
||
... print(x.place) # Check the device of the Tensor
|
||
Place(gpu:0)
|
||
|
||
>>> # Set the global default device to GPU:0
|
||
>>> paddle.set_device("gpu:0")
|
||
>>> # Temporarily switch to CPU using device_guard with Place object (CPUPlace)
|
||
>>> cpu_place = paddle.CPUPlace()
|
||
>>> with paddle.device.device_guard(cpu_place):
|
||
... x = paddle.randn([4, 4]) # Create a Tensor on CPU
|
||
... x = x.tanh() * 2 # Perform computation on CPU
|
||
... print(x.place)
|
||
Place(cpu)
|
||
'''
|
||
|
||
_target_place: Place
|
||
_original_place: Place
|
||
|
||
def __init__(self, device: PlaceLike) -> None:
|
||
if isinstance(device, str):
|
||
self._target_place = paddle.device._convert_to_place(device)
|
||
elif isinstance(device, paddle.base.libpaddle.Place):
|
||
self._target_place = device
|
||
else:
|
||
raise ValueError(
|
||
"'device' must be a string or an instance of a subclass of "
|
||
f"paddle.base.libpaddle.Place, but got {type(device)}"
|
||
)
|
||
|
||
def __enter__(self) -> None:
|
||
self._original_place = paddle.framework._current_expected_place_()
|
||
if self._original_place != self._target_place:
|
||
paddle.framework._set_expected_place(self._target_place)
|
||
|
||
def __exit__(
|
||
self,
|
||
exc_type: type[BaseException] | None,
|
||
exc_val: BaseException | None,
|
||
exc_tb: TracebackType | None,
|
||
) -> None:
|
||
if self._original_place != self._target_place:
|
||
paddle.framework._set_expected_place(self._original_place)
|
||
|
||
|
||
def synchronize(device: PlaceLike | None = None) -> None:
|
||
"""
|
||
Wait for the compute on the given device to finish.
|
||
|
||
Args:
|
||
device(str|paddle.CUDAPlace(n)|paddle.XPUPlace(n)|paddle.CustomPlace(n)): The device which want to wait for. If device is None, the device is the current device. Default: None.
|
||
It can be ``gpu``, ``gpu:x``, ``xpu``, ``xpu:x``, ``custom_device``, ``custom_device:x``, where ``custom_device`` is the name of CustomDevice,
|
||
where ``x`` is the index of the GPUs, XPUs. And it can be paddle.CUDAPlace(n) or paddle.XPUPlace(n) or paddle.CustomPlace(n).
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:CUSTOM_DEVICE)
|
||
>>> import paddle
|
||
>>> paddle.set_device('custom_cpu')
|
||
>>> paddle.device.synchronize()
|
||
>>> paddle.device.synchronize("custom_cpu:0")
|
||
>>> place = paddle.CustomPlace('custom_cpu', 0)
|
||
>>> paddle.device.synchronize(place)
|
||
"""
|
||
|
||
if device is None:
|
||
place = paddle.framework._current_expected_place_()
|
||
elif isinstance(device, str):
|
||
place = paddle.device._convert_to_place(device)
|
||
else:
|
||
place = device
|
||
|
||
if paddle.is_compiled_with_cuda() and isinstance(place, paddle.CUDAPlace):
|
||
core._device_synchronize(place.get_device_id())
|
||
elif paddle.is_compiled_with_xpu() and isinstance(place, paddle.XPUPlace):
|
||
core._xpu_device_synchronize(place.get_device_id())
|
||
elif isinstance(place, paddle.CustomPlace):
|
||
core._synchronize_custom_device(
|
||
place.get_device_type(), place.get_device_id()
|
||
)
|
||
else:
|
||
raise TypeError(
|
||
"device should be gpu, xpu, {}".format(
|
||
",".join(paddle.device.get_all_custom_device_type())
|
||
)
|
||
)
|
||
|
||
|
||
def ipc_collect() -> None:
|
||
"""
|
||
Force collects GPU memory after it has been released by CUDA IPC.
|
||
This function checks if any sent CUDA tensors could be cleaned from the memory.
|
||
Force closes shared memory file used for reference counting if there is no active counters.
|
||
Useful when the producer process stopped actively sending tensors and want to release unused memory.
|
||
Returns:
|
||
None
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
>>> # Force collect expired IPC memory
|
||
>>> paddle.device.ipc_collect() # this is equivalent to paddle.cuda.ipc_collect()
|
||
"""
|
||
paddle.base.libpaddle._ipc_collect()
|
||
|
||
|
||
def get_stream_from_external(
|
||
data_ptr: int, device: PlaceLike | None = None
|
||
) -> Stream:
|
||
r'''
|
||
Return a :class:`Stream` from an externally allocated CUDA stream.
|
||
|
||
This function is used to wrap streams allocated in other libraries in order
|
||
to facilitate data exchange and multi-library interactions.
|
||
|
||
.. note::
|
||
This function doesn't manage the stream life-cycle, it is the user
|
||
responsibility to keep the referenced stream alive while this returned
|
||
stream is being used.
|
||
|
||
Args:
|
||
data_ptr(int): Integer representation of the CUDA stream handle (``cudaStream_t``)
|
||
that is allocated externally.
|
||
device(str|paddle.CUDAPlace(n), optional):
|
||
The CUDA device where the stream was originally allocated.
|
||
If device is None, the current CUDA device is used.
|
||
It can be ``gpu``, ``gpu:x``, or ``paddle.CUDAPlace(n)``.
|
||
|
||
Returns:
|
||
Stream: The wrapped CUDA stream corresponding to the given external pointer.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> # doctest: +SKIP('original_raw_ptr not exist')
|
||
>>> original_raw_ptr = 77777
|
||
>>> external_stream = paddle.device.get_stream_from_external(original_raw_ptr, "cuda:0")
|
||
'''
|
||
if device is None:
|
||
place = paddle.framework._current_expected_place_()
|
||
elif isinstance(device, str):
|
||
place = paddle.device._convert_to_place(device)
|
||
else:
|
||
place = device
|
||
|
||
return Stream(
|
||
stream_base=core._get_stream_from_external(
|
||
data_ptr, place.get_device_id()
|
||
)
|
||
)
|
||
|
||
|
||
def manual_seed_all(seed: int) -> None:
|
||
"""
|
||
|
||
Sets the seed for global default generator, which manages the random number generation.
|
||
|
||
Args:
|
||
seed(int): The random seed to set.
|
||
|
||
Returns:
|
||
None
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
>>> paddle.device.manual_seed_all(102)
|
||
|
||
"""
|
||
paddle.seed(seed)
|
||
|
||
|
||
class _AutocastMode:
|
||
@staticmethod
|
||
def autocast(
|
||
enabled=True, dtype=paddle.float16, cache_enabled=True
|
||
) -> AbstractContextManager:
|
||
"""
|
||
Create a context which enables auto-mixed-precision(AMP) of operators executed in dynamic graph mode.
|
||
If enabled, the input data type (float32, float16 or bfloat16) of each operator is decided
|
||
by autocast algorithm for better performance.
|
||
|
||
Commonly, it is used together with `GradScaler` and `decorator` to achieve Auto-Mixed-Precision in
|
||
imperative mode.
|
||
|
||
Args:
|
||
device_type(str, optional): Device type. But because the paddle does not distinguish between devices, this parameter does not work.
|
||
enable(bool, optional): Enable auto-mixed-precision or not. Default is True.
|
||
dtype(str, optional): Whether to use 'float16' or 'bfloat16'. Default is 'float16'.
|
||
cache_enabled(bool, optional): whether to enable cache or not. Default is True. But this parameter is not used
|
||
|
||
Note:
|
||
paddle.cuda.amp.
|
||
|
||
Examples:
|
||
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
|
||
>>> conv2d = paddle.nn.Conv2D(3, 2, 3, bias_attr=False)
|
||
>>> data = paddle.rand([10, 3, 32, 32])
|
||
|
||
>>> with paddle.device.amp.autocast():
|
||
... conv = conv2d(data)
|
||
... print(conv.dtype)
|
||
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
||
paddle.float16
|
||
>>> # doctest: -SKIP
|
||
|
||
>>> with paddle.device.amp.autocast(enabled=False):
|
||
... conv = conv2d(data)
|
||
... print(conv.dtype)
|
||
>>> # doctest: +SKIP("This has diff in xdoctest env")
|
||
paddle.float32
|
||
>>> # doctest: -SKIP
|
||
|
||
"""
|
||
return _autocast(device_type='cuda', enabled=enabled, dtype=dtype)
|
||
|
||
|
||
class amp:
|
||
"""Namespace for amp marker operations."""
|
||
|
||
autocast = staticmethod(_AutocastMode.autocast)
|
||
autocast_mode = _AutocastMode()
|
||
GradScaler = _GradScaler
|
||
|
||
|
||
class nvtx:
|
||
"""Namespace for NVTX marker operations."""
|
||
|
||
@staticmethod
|
||
@contextmanager
|
||
def range(
|
||
msg: str, *args: Any, **kwargs: Any
|
||
) -> Generator[None, None, None]:
|
||
"""
|
||
Context manager/decorator that pushes and pops an NVTX range.
|
||
|
||
Args:
|
||
msg (str): The name of the NVTX range.
|
||
*args: Arguments used to format ``msg``.
|
||
**kwargs: Keyword arguments used to format ``msg``.
|
||
"""
|
||
nvtx.range_push(msg.format(*args, **kwargs))
|
||
try:
|
||
yield
|
||
finally:
|
||
nvtx.range_pop()
|
||
|
||
@staticmethod
|
||
def range_push(msg: str):
|
||
"""
|
||
Push an NVTX range marker with the given message.
|
||
|
||
Args:
|
||
msg (str): The name of the NVTX range.
|
||
Example:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
>>> # paddle.device.nvtx.range_push("test") is equivalent to paddle.cuda.nvtx.range_push("test")
|
||
>>> paddle.device.nvtx.range_push("test")
|
||
|
||
"""
|
||
paddle.base.core.nvprof_nvtx_push(msg)
|
||
|
||
@staticmethod
|
||
def range_pop():
|
||
"""
|
||
Pop the most recent NVTX range marker.
|
||
Example:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
>>> # paddle.device.nvtx.range_pop("test") is equivalent to paddle.cuda.nvtx.range_pop("test")
|
||
>>> paddle.device.nvtx.range_pop()
|
||
"""
|
||
paddle.base.core.nvprof_nvtx_pop()
|
||
|
||
|
||
def reset_peak_memory_stats(device: PlaceLike | int | None = None) -> None:
|
||
"""
|
||
Resets all devices' peak memory statistics.
|
||
|
||
This method resets the peak memory usage recorded for each device during the execution of the program.
|
||
It sets the peak memory usage back to zero for all devices.
|
||
|
||
Example:
|
||
.. code-block:: pycon
|
||
|
||
>>> # doctest: +REQUIRES(env:GPU)
|
||
>>> import paddle
|
||
>>> paddle.device.set_device('gpu') # or '<custom_device>'
|
||
|
||
>>> # paddle.cuda.reset_max_memory_allocated() is equivalent to paddle.device.reset_max_memory_allocated()
|
||
|
||
>>> paddle.device.reset_max_memory_allocated(paddle.CUDAPlace(0))
|
||
>>> paddle.device.reset_max_memory_allocated(0)
|
||
>>> paddle.device.reset_max_memory_allocated("gpu:0")
|
||
"""
|
||
reset_max_memory_allocated()
|
||
|
||
|
||
class Device(str):
|
||
"""
|
||
Paddle computing device.
|
||
|
||
This class represents a computing device in Paddle, such as CPU, GPU (CUDA), or XPU,
|
||
and can be passed directly to Paddle tensor creation APIs.
|
||
|
||
Note:
|
||
- Only device types "cpu", "gpu", "cuda", and "xpu" are supported.
|
||
- The string representation of the device (e.g., "cuda:0") can be used directly
|
||
in Paddle APIs that accept a device argument.
|
||
- This class supports context manager usage to temporarily set the default device.
|
||
|
||
Args:
|
||
type (str|int, optional): The device type or a legacy device index.
|
||
- str: "cpu", "cuda", "cuda:0", "gpu:1", "xpu:0"
|
||
- int: legacy, interpreted as the default GPU device index
|
||
index (int, optional): The device index, used with `type` string. Ignored for CPU.
|
||
|
||
Attributes:
|
||
type (str): Device type ("cpu", "cuda", "gpu", "xpu").
|
||
index (int|None): Device index. None for CPU.
|
||
|
||
Examples:
|
||
.. code-block:: pycon
|
||
|
||
>>> import paddle
|
||
|
||
# String initialization
|
||
>>> d1 = paddle.device("cpu")
|
||
>>> d2 = paddle.device("cuda:0")
|
||
>>> d3 = paddle.device("xpu", 1)
|
||
|
||
# Type + index initialization
|
||
>>> d4 = paddle.device(type="cuda", index=0)
|
||
|
||
# Legacy int initialization
|
||
>>> d5 = paddle.device(0) # equivalent to paddle.device("cuda", 0)
|
||
|
||
# Copy from another device
|
||
>>> d6 = paddle.device(d2)
|
||
|
||
# Using as context manager
|
||
>>> with paddle.device("cuda:1"):
|
||
... x = paddle.zeros([2, 3]) # created on CUDA device 1
|
||
|
||
>>> print(d2.type) # "cuda"
|
||
>>> print(d2.index) # 0
|
||
>>> print(d1) # "cpu"
|
||
>>> print(d2) # "cuda:0"
|
||
"""
|
||
|
||
_DEFAULT_DEVICE_STACK = []
|
||
_SUPPORTED_TYPES = {"cpu", "gpu", "cuda", "xpu"}
|
||
|
||
def __new__(
|
||
cls, type: PlaceLike | int | None = None, index: int | None = None
|
||
):
|
||
if isinstance(type, paddle.base.libpaddle.Place):
|
||
if type.is_cpu_place():
|
||
dev_type = 'cpu'
|
||
dev_index = None
|
||
elif type.is_gpu_place():
|
||
dev_type = 'cuda'
|
||
dev_index = type.gpu_device_id()
|
||
elif type.is_xpu_place():
|
||
dev_type = 'xpu'
|
||
dev_index = type.gpu_device_id()
|
||
elif type.is_custom_place():
|
||
dev_type = type.custom_device_type()
|
||
dev_index = type.custom_device_id()
|
||
else:
|
||
raise ValueError(f"Unknown place type: {type}")
|
||
|
||
elif isinstance(type, str):
|
||
t = type.lower()
|
||
if ":" in t:
|
||
dev_type, idx = t.split(":")
|
||
dev_type = dev_type.lower()
|
||
if (
|
||
dev_type not in cls._SUPPORTED_TYPES
|
||
and dev_type not in core.get_all_custom_device_type()
|
||
):
|
||
raise ValueError(f"Unsupported device type: {dev_type}")
|
||
dev_index = int(idx)
|
||
elif t in cls._SUPPORTED_TYPES:
|
||
dev_type = t
|
||
dev_index = (
|
||
index if (index is not None and t != "cpu") else None
|
||
)
|
||
elif t in core.get_all_custom_device_type():
|
||
dev_type = t
|
||
dev_index = index
|
||
else:
|
||
raise ValueError(f"Unsupported device type: {t}")
|
||
|
||
elif isinstance(type, int):
|
||
dev_type = "cuda"
|
||
dev_index = type
|
||
|
||
elif type is None and index is not None:
|
||
raise ValueError("Device type must be specified if index is given")
|
||
|
||
else:
|
||
raise TypeError(f"Unsupported type for Device: {type}")
|
||
cond = None
|
||
s = f"{dev_type}:{dev_index}" if dev_index != cond else dev_type
|
||
if dev_type == 'cpu':
|
||
s = 'cpu'
|
||
obj = str.__new__(cls, s)
|
||
obj._dev_type = dev_type
|
||
obj._index = dev_index
|
||
return obj
|
||
|
||
@property
|
||
def type(self):
|
||
return self._dev_type
|
||
|
||
@property
|
||
def index(self):
|
||
return self._index
|
||
|
||
def _to_place(self) -> core.Place:
|
||
if self.type == "cpu":
|
||
return core.CPUPlace()
|
||
elif self.type in {"gpu", "cuda"}:
|
||
return core.CUDAPlace(self.index)
|
||
elif self.type == "xpu":
|
||
return core.XPUPlace(self.index)
|
||
elif self.type in core.get_all_custom_device_type():
|
||
return core.CustomPlace(self.type, self.index)
|
||
else:
|
||
raise ValueError(f"Unsupported device type: {self.type}")
|
||
|
||
def __repr__(self) -> str:
|
||
if self.type == "cpu" or self.index is None:
|
||
return f"device(type='{self.type}')"
|
||
return f"device(type='{self.type}', index={self.index})"
|
||
|
||
def __dlpack_device__(self) -> tuple[int, int]:
|
||
return self._to_place().__dlpack_device__()
|
||
|
||
def __enter__(self):
|
||
current_device = paddle.get_device()
|
||
Device._DEFAULT_DEVICE_STACK.append(current_device)
|
||
paddle.set_device(str(self))
|
||
return self
|
||
|
||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||
previous_device = Device._DEFAULT_DEVICE_STACK.pop()
|
||
paddle.set_device(previous_device)
|
||
|
||
def __getattr__(self, name: str):
|
||
return getattr(self._to_place(), name)
|
||
|
||
|
||
class _DeviceModule(types.ModuleType):
|
||
"""A callable package module: paddle.device(...) -> Device(...)"""
|
||
|
||
def __call__(self, *args, **kwargs) -> Device:
|
||
return Device(*args, **kwargs)
|
||
|
||
def __getattr__(self, name: str):
|
||
# support lazy import submodeule：paddle.device.cuda / paddle.device.xpu / ...
|
||
try:
|
||
mod = importlib.import_module(f"{self.__name__}.{name}")
|
||
setattr(self, name, mod)
|
||
return mod
|
||
except ModuleNotFoundError as e:
|
||
raise AttributeError(name) from e
|
||
|
||
|
||
_self = sys.modules[__name__]
|
||
_proxy = _DeviceModule(__name__, _self.__doc__)
|
||
_proxy.__dict__.update(_self.__dict__)
|
||
sys.modules[__name__] = _proxy
|