690 lines
25 KiB
Python
690 lines
25 KiB
Python
# Copyright (c) 2022 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
|
|
|
|
from typing import TYPE_CHECKING, TypeAlias
|
|
|
|
import paddle
|
|
from paddle.base import core
|
|
from paddle.utils import deprecated
|
|
|
|
from .streams import Event, Stream, create_event, create_stream # noqa: F401
|
|
|
|
if TYPE_CHECKING:
|
|
from paddle import XPUPlace
|
|
from paddle.base.libpaddle import _gpuDeviceProperties
|
|
|
|
_XPUPlaceLike: TypeAlias = (
|
|
XPUPlace
|
|
| str # some str like 'xpu:0', 'xpu:1', etc.
|
|
| int # some int like 0, 1, etc.
|
|
)
|
|
|
|
__all__ = [
|
|
'Stream',
|
|
'Event',
|
|
'synchronize',
|
|
'device_count',
|
|
'set_debug_level',
|
|
'empty_cache',
|
|
'max_memory_allocated',
|
|
'max_memory_reserved',
|
|
'reset_max_memory_allocated',
|
|
'reset_max_memory_reserved',
|
|
'memory_allocated',
|
|
'memory_reserved',
|
|
'memory_total', # memory managed by runtime, not paddle
|
|
'memory_used', # memory managed by runtime, not paddle
|
|
'get_device_properties',
|
|
]
|
|
|
|
|
|
def current_stream(device: _XPUPlaceLike | None = None) -> core.XPUStream:
|
|
'''
|
|
Return the current XPU stream by the device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace()|int|None, optional): The device or the ID of the device which want to get stream from.
|
|
If device is None, the device is the current device. Default: None.
|
|
|
|
Returns:
|
|
XPUStream: the stream to the device.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> s1 = paddle.device.xpu.current_stream()
|
|
|
|
>>> s2 = paddle.device.xpu.current_stream(0)
|
|
|
|
>>> s3 = paddle.device.xpu.current_stream(paddle.XPUPlace(0))
|
|
|
|
'''
|
|
|
|
device_id = -1
|
|
|
|
if device is not None:
|
|
if isinstance(device, int):
|
|
device_id = device
|
|
elif isinstance(device, core.XPUPlace):
|
|
device_id = device.get_device_id()
|
|
elif isinstance(device, str):
|
|
place = paddle.device._convert_to_place(device)
|
|
device_id = place.get_device_id()
|
|
else:
|
|
raise ValueError("device type must be int or paddle.XPUPlace")
|
|
|
|
return core._xpu_get_current_stream(device_id)
|
|
|
|
|
|
def extract_xpu_device_id(device: _XPUPlaceLike, op_name: str) -> int:
|
|
'''
|
|
Return the id of the given xpu device. It is just a utility that will not be exposed to users.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace or int or str): The device, the id of the device or
|
|
the string name of device like 'xpu: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_xpu_current_device_id()
|
|
|
|
if isinstance(device, int):
|
|
device_id = device
|
|
elif isinstance(device, core.XPUPlace):
|
|
device_id = device.get_device_id()
|
|
elif isinstance(device, str):
|
|
if device.startswith('xpu:'):
|
|
device_id = int(device[4:])
|
|
else:
|
|
raise ValueError(
|
|
f"The current string {device} is not expected. Because {op_name} only support string which is like 'xpu:x'. "
|
|
"Please input appropriate string again!"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"The device type {device} is not expected. Because {op_name} only support int, str or paddle.XPUPlace. "
|
|
"Please input appropriate device again!"
|
|
)
|
|
|
|
assert device_id >= 0, (
|
|
f"The device id must be not less than 0, but got id = {device_id}."
|
|
)
|
|
assert device_id < device_count(), (
|
|
f"The device id {device_id} exceeds xpu card number {device_count()}"
|
|
)
|
|
return device_id
|
|
|
|
|
|
@deprecated(
|
|
since="2.5.0",
|
|
update_to="paddle.device.synchronize",
|
|
level=1,
|
|
reason="synchronize in paddle.device.xpu will be removed in future",
|
|
)
|
|
def synchronize(device: _XPUPlaceLike | None = None) -> int:
|
|
"""
|
|
Wait for the compute on the given XPU device to finish.
|
|
|
|
Parameters:
|
|
device(paddle.XPUPlace()|int, optional): The device or the ID of the device.
|
|
If device is None, the device is the current device. Default: None.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
>>> paddle.device.xpu.synchronize()
|
|
>>> paddle.device.xpu.synchronize(0)
|
|
>>> paddle.device.xpu.synchronize(paddle.XPUPlace(0))
|
|
|
|
"""
|
|
|
|
device_id = -1
|
|
|
|
if device is not None:
|
|
if isinstance(device, int):
|
|
device_id = device
|
|
elif isinstance(device, core.XPUPlace):
|
|
device_id = device.get_device_id()
|
|
elif isinstance(device, str):
|
|
if device.startswith('xpu:'):
|
|
device_id = int(device[4:])
|
|
elif device == 'xpu':
|
|
device_id = 0
|
|
else:
|
|
raise ValueError(
|
|
f"The current string {device} is not expected. Because paddle.device.cuda."
|
|
"synchronize only support string which is like 'xpu:x' or 'xpu'. "
|
|
"Please input appropriate string again!"
|
|
)
|
|
else:
|
|
raise ValueError("device type must be int or paddle.XPUPlace")
|
|
|
|
return core._xpu_device_synchronize(device_id)
|
|
|
|
|
|
def device_count() -> int:
|
|
'''
|
|
Return the number of XPUs available.
|
|
|
|
Returns:
|
|
int: the number of XPUs available.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
|
|
>>> paddle.device.xpu.device_count()
|
|
|
|
'''
|
|
|
|
num_xpus = (
|
|
core.get_xpu_device_count()
|
|
if hasattr(core, 'get_xpu_device_count')
|
|
else 0
|
|
)
|
|
|
|
return num_xpus
|
|
|
|
|
|
def set_debug_level(level: int = 0) -> None:
|
|
'''
|
|
Set the debug level of XPUs' api. The default level is 0, which means no debug info.
|
|
|
|
Args:
|
|
int: Debug level of XPUs available.
|
|
Level 0x1 for trace (Print the invocation of the interface),
|
|
0x10 for checksum (Print the checksum of the tensor),
|
|
0x100 for dump (Save the tensor as a file in npy format),
|
|
0x1000 for profiling (Record the execution time of each operator).
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> paddle.device.xpu.set_debug_level(1)
|
|
'''
|
|
name = "paddle.device.xpu.set_debug_level"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
else:
|
|
core.set_xpu_debug_level(level)
|
|
|
|
|
|
def empty_cache() -> None:
|
|
'''
|
|
Releases idle cached memory held by the allocator so that those can be used in other XPU
|
|
application and visible in `xpu-smi`. In most cases you don't need to use this function,
|
|
Paddle does not release the memory back to the OS when you remove Tensors on the XPU,
|
|
Because it keeps xpu memory in a pool so that next allocations can be done much faster.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> tensor = paddle.randn([512, 512, 512], "float64")
|
|
>>> del tensor
|
|
>>> paddle.device.xpu.empty_cache()
|
|
'''
|
|
name = "paddle.device.xpu.empty_cache"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
else:
|
|
core.xpu_empty_cache()
|
|
|
|
|
|
def get_device_properties(
|
|
device: _XPUPlaceLike | None = None,
|
|
) -> _gpuDeviceProperties:
|
|
'''
|
|
Return the properties of given device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x' which to get the properties of the
|
|
device from. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Returns:
|
|
_gpuDeviceProperties: 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
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
>>> paddle.device.xpu.get_device_properties()
|
|
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
|
|
|
>>> paddle.device.xpu.get_device_properties(0)
|
|
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
|
|
|
>>> paddle.device.xpu.get_device_properties('xpu:0')
|
|
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
|
|
|
>>> paddle.device.xpu.get_device_properties(paddle.XPUPlace(0))
|
|
>>> # _gpuDeviceProperties(name='GPU', major=8, minor=6, total_memory=98304MB, multi_processor_count=8)
|
|
|
|
'''
|
|
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
"The API paddle.device.xpu.get_device_properties is not supported in "
|
|
"CPU-only PaddlePaddle. Please reinstall PaddlePaddle with XPU support "
|
|
"to call this API."
|
|
)
|
|
|
|
if device is not None:
|
|
if isinstance(device, int):
|
|
device_id = device
|
|
elif isinstance(device, core.XPUPlace):
|
|
device_id = device.get_device_id()
|
|
elif isinstance(device, str):
|
|
if device.startswith('xpu:'):
|
|
device_id = int(device[4:])
|
|
elif device == 'xpu':
|
|
device_id = 0
|
|
else:
|
|
raise ValueError(
|
|
f"The current string {device} is not expected. Because paddle.device."
|
|
"xpu.get_device_properties only support string which is like 'xpu:x' or 'xpu'. "
|
|
"Please input appropriate string again!"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"The device type {device} is not expected. Because paddle.device.xpu."
|
|
"get_device_properties only support int, str or paddle.XPUPlace. "
|
|
"Please input appropriate device again!"
|
|
)
|
|
else:
|
|
device_id = -1
|
|
|
|
return core.get_device_properties(device_id)
|
|
|
|
|
|
def max_memory_allocated(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the peak size of xpu memory that is allocated to tensor of the given device.
|
|
|
|
Note:
|
|
The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may larger than the memory size that tensor actually need.
|
|
For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The peak size of xpu memory that is allocated to tensor of the given device, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(paddle.XPUPlace(0))
|
|
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated(0)
|
|
>>> max_memory_allocated_size = paddle.device.xpu.max_memory_allocated("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.max_memory_allocated"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.device_memory_stat_peak_value("Allocated", device_id)
|
|
|
|
|
|
def max_memory_reserved(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the peak size of XPU memory that is held by the allocator of the given device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The peak size of XPU memory that is held by the allocator of the given device, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(paddle.XPUPlace(0))
|
|
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved(0)
|
|
>>> max_memory_reserved_size = paddle.device.xpu.max_memory_reserved("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.max_memory_reserved"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.device_memory_stat_peak_value("Reserved", device_id)
|
|
|
|
|
|
def reset_max_memory_allocated(device: _XPUPlaceLike | None = None) -> None:
|
|
'''
|
|
Reset the peak size of XPU memory that is allocated to tensor of the given device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> paddle.device.xpu.reset_max_memory_allocated(paddle.XPUPlace(0))
|
|
>>> paddle.device.xpu.reset_max_memory_allocated(0)
|
|
>>> paddle.device.xpu.reset_max_memory_allocated("xpu:0")
|
|
'''
|
|
|
|
name = "paddle.device.xpu.reset_max_memory_allocated"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
core.device_memory_stat_reset_peak_value("Allocated", device_id)
|
|
|
|
|
|
def reset_max_memory_reserved(device: _XPUPlaceLike | None = None) -> None:
|
|
'''
|
|
Reset the peak size of XPU memory that is held by the allocator of the given device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> paddle.device.xpu.reset_max_memory_reserved(paddle.XPUPlace(0))
|
|
>>> paddle.device.xpu.reset_max_memory_reserved(0)
|
|
>>> paddle.device.xpu.reset_max_memory_reserved("xpu:0")
|
|
'''
|
|
|
|
name = "paddle.device.xpu.reset_max_memory_reserved"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
core.device_memory_stat_reset_peak_value("Reserved", device_id)
|
|
|
|
|
|
def memory_allocated(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the current size of xpu memory that is allocated to tensor of the given device.
|
|
|
|
Note:
|
|
The size of XPU memory allocated to tensor is 256-byte aligned in Paddle, which may be larger than the memory size that tensor actually need.
|
|
For instance, a float32 0-D Tensor with shape [] in XPU will take up 256 bytes memory, even though storing a float32 data requires only 4 bytes.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The current size of xpu memory that is allocated to tensor of the given device, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> memory_allocated_size = paddle.device.xpu.memory_allocated(paddle.XPUPlace(0))
|
|
>>> memory_allocated_size = paddle.device.xpu.memory_allocated(0)
|
|
>>> memory_allocated_size = paddle.device.xpu.memory_allocated("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.memory_allocated"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.device_memory_stat_current_value("Allocated", device_id)
|
|
|
|
|
|
def memory_reserved(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the current size of XPU memory that is held by the allocator of the given device.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The current size of XPU memory that is held by the allocator of the given device, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> memory_reserved_size = paddle.device.xpu.memory_reserved(paddle.XPUPlace(0))
|
|
>>> memory_reserved_size = paddle.device.xpu.memory_reserved(0)
|
|
>>> memory_reserved_size = paddle.device.xpu.memory_reserved("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.memory_reserved"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.device_memory_stat_current_value("Reserved", device_id)
|
|
|
|
|
|
def memory_total(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the total size of XPU memory of the given device that is held by the XPU Runtime.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The total size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> memory_total_size = paddle.device.xpu.memory_total(paddle.XPUPlace(0))
|
|
>>> memory_total_size = paddle.device.xpu.memory_total(0)
|
|
>>> memory_total_size = paddle.device.xpu.memory_total("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.memory_total"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.get_xpu_device_total_memory(device_id)
|
|
|
|
|
|
def memory_used(device: _XPUPlaceLike | None = None) -> int:
|
|
'''
|
|
Return the used size of XPU memory of the given device that is held by the XPU Runtime.
|
|
|
|
Args:
|
|
device(paddle.XPUPlace|int|str|None, optional): The device, the id of the device or
|
|
the string name of device like 'xpu:x'. If device is None, the device is the current device.
|
|
Default: None.
|
|
|
|
Return:
|
|
int: The used size of XPU memory of the given device that is held by the XPU Runtime, in bytes.
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.set_device('xpu')
|
|
|
|
>>> memory_used_size = paddle.device.xpu.memory_used(paddle.XPUPlace(0))
|
|
>>> memory_used_size = paddle.device.xpu.memory_used(0)
|
|
>>> memory_used_size = paddle.device.xpu.memory_used("xpu:0")
|
|
'''
|
|
name = "paddle.device.xpu.memory_used"
|
|
if not core.is_compiled_with_xpu():
|
|
raise ValueError(
|
|
f"The API {name} is only supported in XPU PaddlePaddle. Please reinstall PaddlePaddle with XPU support to call this API."
|
|
)
|
|
device_id = extract_xpu_device_id(device, op_name=name)
|
|
return core.get_xpu_device_used_memory(device_id)
|
|
|
|
|
|
def get_rng_state(device: _XPUPlaceLike | None = None) -> core.GeneratorState:
|
|
'''
|
|
Get the random state for the default generator.
|
|
|
|
Returns:
|
|
Tensor: The random state tensor.
|
|
|
|
Examples:
|
|
|
|
.. code-block:: pycon
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.get_rng_state()
|
|
|
|
'''
|
|
place = paddle.device.device_to_place(device)
|
|
if isinstance(place, core.CPUPlace):
|
|
return core.default_cpu_generator().get_state()
|
|
return core.default_xpu_generator(place.get_device_id()).get_state()
|
|
|
|
|
|
def set_rng_state(
|
|
new_state: core.GeneratorState, device: _XPUPlaceLike | None = None
|
|
) -> None:
|
|
"""
|
|
Set the random number generator state of the specified device.
|
|
|
|
Args:
|
|
new_state (core.GeneratorState): The desired RNG state to set.
|
|
This should be a state object previously obtained from ``get_rng_state()``.
|
|
device (DeviceLike, optional): The device to set the RNG state for.
|
|
If not specified, uses the current default device (as returned by ``paddle.framework._current_expected_place_()``).
|
|
Can be a device object, integer device ID, or device string.
|
|
|
|
Returns:
|
|
None
|
|
|
|
Examples:
|
|
.. code-block:: pycon
|
|
|
|
>>> import paddle
|
|
>>> # Save RNG state
|
|
>>> state = paddle.device.get_rng_state()
|
|
>>> # Do some random operations
|
|
>>> x = paddle.randn([2, 3])
|
|
>>> # Restore RNG state
|
|
>>> paddle.device.set_rng_state(state)
|
|
"""
|
|
place = paddle.device.device_to_place(device)
|
|
if isinstance(place, core.CPUPlace):
|
|
core.default_cpu_generator().set_state(new_state)
|
|
else:
|
|
core.default_xpu_generator(place.get_device_id()).set_state(new_state)
|
|
|
|
|
|
def manual_seed(seed: int) -> None:
|
|
r"""Set the seed for generating random numbers for the current Device.
|
|
|
|
.. warning::
|
|
If you are working with a multi-Device model, this function is insufficient
|
|
to get determinism. To seed all Devices, use :func:`manual_seed_all`.
|
|
If current Device is CPU, this function will set the seed of the default CPU generator.
|
|
|
|
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
|
|
|
|
>>> # doctest: +REQUIRES(env:XPU)
|
|
>>> import paddle
|
|
>>> paddle.device.manual_seed(102)
|
|
>>> # paddle.cuda.manual_seed(102) is equivalent to paddle.device.manual_seed(102)
|
|
>>> paddle.cuda.manual_seed(102)
|
|
|
|
"""
|
|
seed = int(seed)
|
|
place = paddle.framework._current_expected_place_()
|
|
if isinstance(place, core.CPUPlace):
|
|
core.default_cpu_generator().manual_seed(seed)
|
|
else:
|
|
core.default_xpu_generator(place.get_device_id()).manual_seed(seed)
|