kolandr/gridbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

Browse Source
This commit is contained in:
Kolp
2026-09-24 13:22:23 +07:00
commit 642cc11a9f
18968 changed files with 5683248 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kronos-venv/lib/python3.12/site-packages/torch/xpu/__init__.py
+650
View File
@@ -0,0 +1,650 @@
# mypy: allow-untyped-defs
r"""
This package introduces support for the XPU backend, specifically tailored for
Intel GPU optimization.
This package is lazily initialized, so you can always import it, and use
:func:`is_available()` to determine if your system supports XPU.
"""
from __future__ import annotations
import threading
import traceback
from functools import lru_cache
from typing import Any, NewType, TYPE_CHECKING
import torch
import torch._C
from torch._utils import _dummy_type, _LazySeedTracker
if TYPE_CHECKING:
from collections.abc import Callable
from torch.types import Device
from ._utils import _get_device_index
from .graphs import (
graph,
graph_pool_handle,
is_current_stream_capturing,
make_graphed_callables,
XPUGraph,
)
from .streams import Event, Stream
_initialized = False
_tls = threading.local()
_initialization_lock = threading.Lock()
_queued_calls: list[
tuple[Callable[[], None], list[str]]
] = [] # don't invoke these until initialization occurs
_is_in_bad_fork = getattr(torch._C, "_xpu_isInBadFork", lambda: False)
_lazy_seed_tracker = _LazySeedTracker()
default_generators: tuple[torch._C.Generator] = () # type: ignore[assignment]
def _is_compiled() -> bool:
r"""Return true if compile with XPU support."""
return torch._C._has_xpu
if _is_compiled():
_XpuDeviceProperties = torch._C._XpuDeviceProperties
_exchange_device = torch._C._xpu_exchangeDevice
_maybe_exchange_device = torch._C._xpu_maybeExchangeDevice
else:
# Define dummy if PyTorch was compiled without XPU
_XpuDeviceProperties = _dummy_type("_XpuDeviceProperties") # type: ignore[assignment, misc]
def _exchange_device(device: int) -> int:
raise NotImplementedError("PyTorch was compiled without XPU support")
def _maybe_exchange_device(device: int) -> int:
raise NotImplementedError("PyTorch was compiled without XPU support")
@lru_cache(maxsize=1)
def device_count() -> int:
r"""Return the number of XPU device available."""
if not _is_compiled():
return 0
return torch._C._xpu_getDeviceCount()
def is_available() -> bool:
r"""Return a bool indicating if XPU is currently available."""
# This function never throws.
return device_count() > 0
def is_bf16_supported(including_emulation: bool = True) -> bool:
r"""Return a bool indicating if the current XPU device supports dtype bfloat16."""
if not is_available():
return False
return (
including_emulation
or torch.xpu.get_device_properties().has_bfloat16_conversions
)
def is_tf32_supported() -> bool:
r"""Return a bool indicating if the current XPU device supports dtype tf32."""
if not is_available():
return False
# On Intel Xe architecture and newer, TF32 operations can be accelerated
# through DPAS (Dot Product Accumulate Systolic) instructions. Therefore,
# TF32 support can be determined by checking whether the device supports
# subgroup matrix multiply-accumulate operations.
return torch.xpu.get_device_properties().has_subgroup_matrix_multiply_accumulate
def is_initialized():
r"""Return whether PyTorch's XPU state has been initialized."""
return _initialized and not _is_in_bad_fork()
def _lazy_call(callable, **kwargs) -> None:
if is_initialized():
callable()
else:
global _lazy_seed_tracker
if kwargs.get("seed_all", False):
_lazy_seed_tracker.queue_seed_all(callable, traceback.format_stack())
elif kwargs.get("seed", False):
_lazy_seed_tracker.queue_seed(callable, traceback.format_stack())
else:
# Don't store the actual traceback to avoid memory cycle
_queued_calls.append((callable, traceback.format_stack()))
def init() -> None:
r"""Initialize PyTorch's XPU state.
This is a Python API about lazy initialization that avoids initializing
XPU until the first time it is accessed. Does nothing if the XPU state is
already initialized.
"""
_lazy_init()
def _lazy_init() -> None:
global _initialized, _queued_calls
if is_initialized() or hasattr(_tls, "is_initializing"):
return
with _initialization_lock:
# This test was was protected via GIL. Double-check whether XPU has
# already been initialized.
if is_initialized():
return
# Stop promptly upon encountering a bad fork error.
if _is_in_bad_fork():
raise RuntimeError(
"Cannot re-initialize XPU in forked subprocess. To use XPU with "
"multiprocessing, you must use the 'spawn' start method"
)
if not _is_compiled():
raise AssertionError("Torch not compiled with XPU enabled")
# This function inits XPU backend and detects bad fork processing.
torch._C._xpu_init()
# Some of the queued calls may reentrantly call _lazy_init(); We need to
# just return without initializing in that case.
_tls.is_initializing = True
_queued_calls.extend(calls for calls in _lazy_seed_tracker.get_calls() if calls)
try:
for queued_call, orig_traceback in _queued_calls:
try:
queued_call()
except Exception as e:
msg = (
f"XPU call failed lazily at initialization with error: {str(e)}\n\n"
f"XPU call was originally invoked at:\n\n{''.join(orig_traceback)}"
)
raise Exception(msg) from e # noqa: TRY002
finally:
delattr(_tls, "is_initializing")
_initialized = True
class _DeviceGuard:
def __init__(self, index: int) -> None:
self.idx = index
self.prev_idx = -1
def __enter__(self):
self.prev_idx = torch.xpu._exchange_device(self.idx)
def __exit__(self, type: Any, value: Any, traceback: Any):
self.idx = torch.xpu._maybe_exchange_device(self.prev_idx)
return False
class device:
r"""Context-manager that changes the selected device.
Args:
device (torch.device or int or str): device index to select. It's a no-op if
this argument is a negative integer or ``None``.
"""
def __init__(self, device: Any) -> None:
self.idx = _get_device_index(device, optional=True)
self.prev_idx = -1
def __enter__(self):
self.prev_idx = torch.xpu._exchange_device(self.idx)
def __exit__(self, type: Any, value: Any, traceback: Any):
self.idx = torch.xpu._maybe_exchange_device(self.prev_idx)
return False
class device_of(device):
r"""Context-manager that changes the current device to that of given object.
You can use both tensors and storages as arguments. If a given object is
not allocated on a XPU, this is a no-op.
Args:
obj (Tensor or Storage): object allocated on the selected device.
"""
def __init__(self, obj) -> None:
idx = obj.get_device() if obj.is_xpu else -1
super().__init__(idx)
def set_device(device: Device) -> None:
r"""Set the current device.
Args:
device (torch.device or int or str): selected device. This function is a
no-op if this argument is negative.
"""
_lazy_init()
device = _get_device_index(device)
if device >= 0:
torch._C._xpu_setDevice(device)
def get_device_name(device: Device = None) -> str:
r"""Get the name of a device.
Args:
device (torch.device or int or str, optional): device for which to
return the name. This function is a no-op if this argument is a
negative integer. It uses the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
Returns:
str: the name of the device
"""
return get_device_properties(device).name
@lru_cache(None)
def get_device_capability(device: Device = None) -> dict[str, Any]:
r"""Get the xpu capability of a device.
Args:
device (torch.device or int or str, optional): device for which to
return the device capability. This function is a no-op if this
argument is a negative integer. It uses the current device, given by
:func:`~torch.xpu.current_device`, if :attr:`device` is ``None``
(default).
Returns:
dict[str, Any]: the xpu capability dictionary of the device
"""
props = get_device_properties(device)
# Only keep attributes that are safe for dictionary serialization.
serializable_types = (int, float, bool, str, type(None), list, tuple, dict)
return {
key: value
for key in dir(props)
if not key.startswith("__")
and isinstance((value := getattr(props, key)), serializable_types)
}
def get_device_properties(
device: Device = None,
) -> _XpuDeviceProperties:
r"""Get the properties of a device. Returns _XpuDeviceProperties containing the following device properties:
- ``name`` (str): device name.
- ``platform_name`` (str): SYCL platform name.
- ``vendor`` (str): device vendor.
- ``device_id`` (int): device identifier (product ID).
- ``driver_version`` (str): driver version.
- ``version`` (str): runtime version.
- ``max_compute_units`` (int): number of parallel compute units.
- ``gpu_eu_count`` (int): number of EUs (Execution Unit).
- ``max_work_group_size``: (int): maximum number of work-items permitted in a work-group.
- ``max_num_sub_groups`` (int): maximum number of sub-groups supported in a work-group.
- ``memory_clock_rate`` (int) maximum clock rate of device's global memory in MHz.
- ``memory_bus_width`` (int) maximum bus width between device and memory in bits.
- ``sub_group_sizes``: (list[int]): a list of supported sub-group sizes.
- ``local_mem_size`` (int): device local memory capacity that can be allocated per work-group in bytes.
- ``has_fp16`` (bool): whether float16 dtype is supported.
- ``has_fp64`` (bool): whether float64 dtype is supported.
- ``has_atomic64`` (bool): whether 64-bit atomic operations are supported.
- ``has_bfloat16_conversions`` (bool): whether bfloat16 conversions are supported.
- ``has_subgroup_matrix_multiply_accumulate`` (bool): whether DPAS (Dot Product Accumulate Systolic) is supported.
- ``has_subgroup_matrix_multiply_accumulate_tensor_float32`` (bool): whether DPAS with tf32 inputs is supported.
- ``has_subgroup_2d_block_io`` (bool): whether 2D block I/O for efficient matrix multiplication is supported.
- ``total_memory`` (int): device global memory in bytes.
- ``gpu_subslice_count`` (int): number of subslice.
- ``architecture`` (int): device architecture identifier (experimental).
- ``type`` (str): device type, e.g. 'cpu', 'gpu', accelerator', 'host', 'unknown'.
- ``uuid`` (Any): device UUID (Universal Unique ID), 16 bytes.
Args:
device (torch.device or int or str): device for which to return the
properties of the device.
Returns:
_XpuDeviceProperties: the properties of the device
"""
_lazy_init()
device = _get_device_index(device, optional=True)
return _get_device_properties(device) # type: ignore[name-defined] # noqa: F821
def current_device() -> int:
r"""Return the index of a currently selected device."""
_lazy_init()
return torch._C._xpu_getDevice()
def _get_device(device: int | str | torch.device) -> torch.device:
r"""Return the torch.device type object from the passed in device.
Args:
device (torch.device or int or str): selected device.
"""
if isinstance(device, str):
device = torch.device(device)
elif isinstance(device, int):
device = torch.device("xpu", device)
return device
def can_device_access_peer(device: Device, peer: Device) -> bool:
r"""Query whether a device can access a peer device's memory.
Args:
device (torch.device or int or str): selected device.
peer (torch.device or int or str): peer device to query access to.
Returns:
bool: ``True`` if ``device`` can access ``peer``, ``False`` otherwise.
"""
_lazy_init()
device = _get_device_index(device, optional=True)
peer = _get_device_index(peer, optional=True)
return torch._C._xpu_canDeviceAccessPeer(device, peer)
class StreamContext:
r"""Context-manager that selects a given stream.
All XPU kernels queued within its context will be enqueued on a selected
stream.
Args:
Stream (Stream): selected stream. This manager is a no-op if it's
``None``.
.. note:: Streams are per-device.
"""
cur_stream: torch.xpu.Stream | None
def __init__(self, stream: torch.xpu.Stream | None) -> None:
self.stream = stream
self.idx = _get_device_index(None, True)
if self.idx is None:
self.idx = -1 # pyrefly: ignore [bad-assignment]
def __enter__(self):
cur_stream = self.stream
if cur_stream is None or self.idx == -1:
return
self.src_prev_stream = torch.xpu.current_stream(None)
# If the stream is not on the current device, then set the current stream on the device
if self.src_prev_stream.device != cur_stream.device:
with device(cur_stream.device):
self.dst_prev_stream = torch.xpu.current_stream(cur_stream.device)
torch.xpu.set_stream(cur_stream)
def __exit__(self, type: Any, value: Any, traceback: Any):
cur_stream = self.stream
if cur_stream is None or self.idx == -1:
return
# Reset the stream on the original device and destination device
if self.src_prev_stream.device != cur_stream.device:
torch.xpu.set_stream(self.dst_prev_stream)
torch.xpu.set_stream(self.src_prev_stream)
def stream(stream: torch.xpu.Stream | None) -> StreamContext:
r"""Wrap around the Context-manager StreamContext that selects a given stream.
Arguments:
stream (Stream): selected stream. This manager is a no-op if it's ``None``.
"""
return StreamContext(stream)
def _set_stream_by_id(stream_id, device_index, device_type) -> None:
r"""set stream specified by the stream id, device index and device type
Args: stream_id (int): not visible to the user, used to assigned to the specific stream.
device_index (int): selected device index.
device_type (int): selected device type.
"""
torch._C._xpu_setStream(
stream_id=stream_id,
device_index=device_index,
device_type=device_type,
)
def set_stream(stream: Stream) -> None:
r"""Set the current stream. This is a wrapper API to set the stream.
Usage of this function is discouraged in favor of the ``stream``
context manager.
Args:
stream (Stream): selected stream. This function is a no-op
if this argument is ``None``.
"""
if stream is None:
return
_lazy_init()
_set_stream_by_id(
stream_id=stream.stream_id,
device_index=stream.device_index,
device_type=stream.device_type,
)
def current_stream(device: Device = None) -> Stream:
r"""Return the currently selected :class:`Stream` for a given device.
Args:
device (torch.device or int, optional): selected device. Returns
the currently selected :class:`Stream` for the current device, given
by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None``
(default).
"""
_lazy_init()
streamdata = torch._C._xpu_getCurrentStream(
_get_device_index(device, optional=True)
)
return Stream(
stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2]
)
def get_stream_from_external(data_ptr: int, device: Device = None) -> Stream:
r"""Return a :class:`Stream` from an external SYCL queue.
This function is used to wrap SYCL queue created in other libraries in order
to facilitate data exchange and multi-library interactions.
.. note:: This function doesn't manage the queue life-cycle, it is the user
responsibility to keep the referenced queue alive while this returned stream is
being used. The different SYCL queue pointers will result in distinct
:class:`Stream` objects, even if the SYCL queues they dereference are equivalent.
Args:
data_ptr(int): Integer representation of the `sycl::queue*` value passed externally.
device(torch.device or int, optional): the device where the queue was originally created.
It is the user responsibility to ensure the device is specified correctly.
"""
_lazy_init()
streamdata = torch._C._xpu_getStreamFromExternal(
data_ptr, _get_device_index(device, optional=True)
)
return Stream(
stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2]
)
def synchronize(device: Device = None) -> None:
r"""Wait for all kernels in all streams on a XPU device to complete.
Args:
device (torch.device or int, optional): device for which to synchronize.
It uses the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
_lazy_init()
device = _get_device_index(device, optional=True)
return torch._C._xpu_synchronize(device)
def get_arch_list() -> list[str]:
r"""Return list XPU architectures this library was compiled for."""
if not _is_compiled():
return []
arch_flags = torch._C._xpu_getArchFlags()
if arch_flags is None:
return []
return arch_flags.split()
def get_gencode_flags() -> str:
r"""Return XPU AOT(ahead-of-time) build flags this library was compiled with."""
arch_list = get_arch_list()
if len(arch_list) == 0:
return ""
return f"-device {','.join(arch for arch in arch_list)}"
def _get_generator(device: torch.device) -> torch._C.Generator:
r"""Return the XPU Generator object for the given device.
Args:
device (torch.device): selected device.
"""
idx = device.index
if idx is None:
idx = current_device()
return torch.xpu.default_generators[idx]
def _set_rng_state_offset(
offset: int, device: int | str | torch.device = "xpu"
) -> None:
r"""Set the random number generator state offset of the specified GPU.
Args:
offset (int): The desired offset
device (torch.device or int, optional): The device to set the RNG state.
Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device).
"""
final_device = _get_device(device)
def cb() -> None:
default_generator = _get_generator(final_device)
default_generator.set_offset(offset)
_lazy_call(cb)
def _get_rng_state_offset(device: int | str | torch.device = "xpu") -> int:
r"""Return the random number generator state offset of the specified GPU.
Args:
device (torch.device or int, optional): The device to return the RNG state offset of.
Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device).
.. warning::
This function eagerly initializes XPU.
"""
_lazy_init()
final_device = _get_device(device)
default_generator = _get_generator(final_device)
return default_generator.get_offset()
# import here to avoid circular import
from .memory import (
change_current_allocator,
empty_cache,
get_per_process_memory_fraction,
max_memory_allocated,
max_memory_reserved,
mem_get_info,
memory_allocated,
memory_reserved,
memory_snapshot,
memory_stats,
memory_stats_as_nested_dict,
MemPool,
reset_accumulated_memory_stats,
reset_peak_memory_stats,
set_per_process_memory_fraction,
use_mem_pool,
XPUPluggableAllocator,
)
from .random import (
get_rng_state,
get_rng_state_all,
initial_seed,
manual_seed,
manual_seed_all,
seed,
seed_all,
set_rng_state,
set_rng_state_all,
)
_POOL_HANDLE = NewType("_POOL_HANDLE", tuple[int, int])
__all__ = [
"Event",
"Stream",
"StreamContext",
"XPUPluggableAllocator",
"XPUGraph",
"can_device_access_peer",
"change_current_allocator",
"current_device",
"current_stream",
"default_generators",
"device",
"device_of",
"device_count",
"empty_cache",
"get_arch_list",
"get_device_capability",
"get_device_name",
"get_device_properties",
"get_gencode_flags",
"get_per_process_memory_fraction",
"get_rng_state",
"get_rng_state_all",
"get_stream_from_external",
"graph",
"graph_pool_handle",
"init",
"initial_seed",
"is_available",
"is_bf16_supported",
"is_current_stream_capturing",
"is_initialized",
"is_tf32_supported",
"make_graphed_callables",
"manual_seed",
"manual_seed_all",
"max_memory_allocated",
"max_memory_reserved",
"mem_get_info",
"memory_allocated",
"memory_reserved",
"memory_snapshot",
"memory_stats",
"memory_stats_as_nested_dict",
"MemPool",
"use_mem_pool",
"reset_accumulated_memory_stats",
"reset_peak_memory_stats",
"seed",
"seed_all",
"set_device",
"set_per_process_memory_fraction",
"set_rng_state",
"set_rng_state_all",
"set_stream",
"stream",
"streams",
"synchronize",
]
kronos-venv/lib/python3.12/site-packages/torch/xpu/_gpu_trace.py
+69
View File
@@ -0,0 +1,69 @@
from collections.abc import Callable
from torch._utils import CallbackRegistry
EventCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry("XPU event creation")
EventDeletionCallbacks: "CallbackRegistry[int]" = CallbackRegistry("XPU event deletion")
EventRecordCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry(
"XPU event record"
)
EventWaitCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry("XPU event wait")
MemoryAllocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
"XPU memory allocation"
)
MemoryDeallocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
"XPU memory deallocation"
)
StreamCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
"XPU stream creation"
)
DeviceSynchronizationCallbacks: "CallbackRegistry[[]]" = CallbackRegistry(
"XPU device synchronization"
)
StreamSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
"XPU stream synchronization"
)
EventSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry(
"XPU event synchronization"
)
def register_callback_for_event_creation(cb: Callable[[int], None]) -> None:
EventCreationCallbacks.add_callback(cb)
def register_callback_for_event_deletion(cb: Callable[[int], None]) -> None:
EventDeletionCallbacks.add_callback(cb)
def register_callback_for_event_record(cb: Callable[[int, int], None]) -> None:
EventRecordCallbacks.add_callback(cb)
def register_callback_for_event_wait(cb: Callable[[int, int], None]) -> None:
EventWaitCallbacks.add_callback(cb)
def register_callback_for_memory_allocation(cb: Callable[[int], None]) -> None:
MemoryAllocationCallbacks.add_callback(cb)
def register_callback_for_memory_deallocation(cb: Callable[[int], None]) -> None:
MemoryDeallocationCallbacks.add_callback(cb)
def register_callback_for_stream_creation(cb: Callable[[int], None]) -> None:
StreamCreationCallbacks.add_callback(cb)
def register_callback_for_device_synchronization(cb: Callable[[], None]) -> None:
DeviceSynchronizationCallbacks.add_callback(cb)
def register_callback_for_stream_synchronization(cb: Callable[[int], None]) -> None:
StreamSynchronizationCallbacks.add_callback(cb)
def register_callback_for_event_synchronization(cb: Callable[[int], None]) -> None:
EventSynchronizationCallbacks.add_callback(cb)
kronos-venv/lib/python3.12/site-packages/torch/xpu/_utils.py
+39
View File
@@ -0,0 +1,39 @@
from typing import Any
import torch
# The _get_device_index has been moved to torch.utils._get_device_index
from torch._utils import _get_device_index as _torch_get_device_index
def _get_device_index(
device: Any, optional: bool = False, allow_cpu: bool = False
) -> int:
r"""Get the device index from :attr:`device`, which can be a torch.device
object, a Python integer, or ``None``.
If :attr:`device` is a torch.device object, returns the device index if it
is a XPU device. Note that for a XPU device without a specified index,
i.e., ``torch.device('xpu')``, this will return the current default XPU
device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``,
CPU devices will be accepted and ``-1`` will be returned in this case.
If :attr:`device` is a Python integer, it is returned as is.
If :attr:`device` is ``None``, this will return the current default XPU
device if :attr:`optional` is ``True``.
"""
if isinstance(device, int):
return device
if isinstance(device, str):
device = torch.device(device)
if isinstance(device, torch.device):
if allow_cpu:
if device.type not in ["xpu", "cpu"]:
raise ValueError(f"Expected a xpu or cpu device, but got: {device}")
elif device.type != "xpu":
raise ValueError(f"Expected a xpu device, but got: {device}")
if not torch.jit.is_scripting():
if isinstance(device, torch.xpu.device):
return device.idx
return _torch_get_device_index(device, optional, allow_cpu)
kronos-venv/lib/python3.12/site-packages/torch/xpu/graphs.py
+529
View File
@@ -0,0 +1,529 @@
from __future__ import annotations
import typing
from collections.abc import Callable
from typing import overload, TYPE_CHECKING, TypeAlias
from typing_extensions import ParamSpec, Self, TypeVar
import torch
from torch import Tensor
if TYPE_CHECKING:
from torch.xpu import _POOL_HANDLE
from .._utils import _dummy_type
__all__ = [
"is_current_stream_capturing",
"graph_pool_handle",
"XPUGraph",
"graph",
"make_graphed_callables",
]
_R = TypeVar("_R")
_P = ParamSpec("_P")
if not hasattr(torch._C, "_XpuStreamBase"):
# Define dummy base classes
torch._C.__dict__["_XPUGraph"] = _dummy_type("_XPUGraph")
torch._C.__dict__["_xpu_graph_pool_handle"] = _dummy_type("_xpu_graph_pool_handle")
torch._C.__dict__["_xpu_isCurrentStreamCapturing"] = _dummy_type(
"_xpu_isCurrentStreamCapturing"
)
from torch._C import _xpu_graph_pool_handle, _xpu_isCurrentStreamCapturing, _XPUGraph
def is_current_stream_capturing() -> bool:
r"""Return True if XPU graph capture is underway on the current XPU stream, False otherwise.
If a XPU context does not exist on the current device, returns False without initializing the context.
"""
return _xpu_isCurrentStreamCapturing()
def graph_pool_handle() -> _POOL_HANDLE:
r"""Return an opaque token representing the id of a graph memory pool."""
return torch.xpu._POOL_HANDLE(_xpu_graph_pool_handle())
class XPUGraph(_XPUGraph):
r"""Wrapper around a XPU graph.
Arguments:
keep_graph (bool, optional): If ``keep_graph=False``, the
executable command graph will be instantiated on GPU at the end of
``capture_end`` and the underlying modifiable command graph will be
destroyed. Note that the executable command graph will not be
instantiated at the end of ``capture_end`` in this
case. Instead, it will be instantiated via an explicit called
to ``instantiate`` or automatically on the first call to
``replay`` if ``instantiate`` was not already called. Calling
``instantiate`` manually before ``replay`` is recommended to
prevent increased latency on the first call to ``replay``.
"""
def __new__(cls, keep_graph: bool = False) -> Self:
return super().__new__(cls, keep_graph)
def capture_begin(self, pool: _POOL_HANDLE | None = None) -> None:
r"""Begin capturing XPU work on the current xpu stream.
Typically, you shouldn't call ``capture_begin`` yourself.
Use :class:`~torch.xpu.graph`, which call ``capture_begin`` internally.
Arguments:
pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) that hints this graph may share memory
with the indicated pool.
"""
super().capture_begin(pool=pool)
def capture_end(self) -> None:
r"""End XPU graph capture on the current stream.
After ``capture_end``, ``replay`` may be called on this instance.
Typically, you shouldn't call ``capture_end`` yourself.
Use :class:`~torch.xpu.graph`, which call ``capture_end`` internally.
"""
super().capture_end()
def instantiate(self) -> None:
r"""Instantiate the XPU graph. Will be called by
``capture_end`` if ``keep_graph=False``, or by ``replay`` if
``keep_graph=True`` and ``instantiate`` has not already been
explicitly called. Does not destroy the xpu modify command graph returned
by ``raw_xpu_graph``.
"""
super().instantiate()
def replay(self) -> None:
r"""Replay the XPU work captured by this graph."""
super().replay()
def reset(self) -> None:
r"""Delete the graph currently held by this instance."""
super().reset()
def pool(self) -> _POOL_HANDLE:
r"""Return an opaque token representing the id of this graph's memory pool.
This id can optionally be passed to another graph's ``capture_begin``,
which hints the other graph may share the same memory pool.
"""
return super().pool()
def enable_debug_mode(self) -> None:
r"""Enable debugging mode for XPUGraph.debug_dump."""
return super().enable_debug_mode()
def debug_dump(self, debug_path: str) -> None:
r"""
Arguments:
debug_path (required): Path to dump the graph to.
Calls a debugging function to dump the graph if the debugging is
enabled via XPUGraph.enable_debug_mode()
"""
return super().debug_dump(debug_path)
def raw_xpu_graph(self) -> int:
r"""Returns the underlying xpuGraph_t. ``keep_graph`` must be True.
XPU doesn't provide APIs to manipulate this object.
""" # noqa: B950
return super().raw_xpu_graph()
def raw_xpu_graph_exec(self) -> int:
r"""Returns the underlying xpuGraphExec_t. ``instantiate`` must have been called if ``keep_graph`` is True, or ``capture_end`` must have been called if ``keep_graph`` is False. If you call ``instantiate()`` after ``raw_xpu_graph_exec()``, the previously returned xpuGraphExec_t will be destroyed. It is your responsibility not to use this object after destruction.
XPU doesn't provide APIs to manipulate this object.
""" # noqa: B950
return super().raw_xpu_graph_exec()
class graph:
r"""Context-manager that captures XPU work into a :class:`torch.xpu.XPUGraph` object for later replay.
Arguments:
xpu_graph (torch.xpu.XPUGraph): Graph object used for capture.
pool (optional): Opaque token (returned by a call to :func:`~torch.xpu.graph_pool_handle()` or
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) hinting this graph's capture
may share memory from the specified pool.
stream (torch.xpu.Stream, optional): If supplied, will be set as the current stream in the context.
If not supplied, ``graph`` sets its own internal side stream as the current stream in the context.
.. note::
For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture
used an explicit ``stream`` argument, you should pass the same ``stream`` argument to this capture.
""" # noqa: B950
default_capture_stream: torch.xpu.Stream | None = None
def __init__(
self,
xpu_graph: XPUGraph,
pool: _POOL_HANDLE | None = None,
stream: torch.xpu.Stream | None = None,
):
# Lazy-init of default_capture_stream helps avoid circular-import errors.
# Not thread safe, but graphs already have the general (explicitly documented)
# restriction that only one capture may be underway at a time in the process.
if self.__class__.default_capture_stream is None:
self.__class__.default_capture_stream = torch.xpu.Stream()
self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,)
self.capture_stream = (
stream if stream is not None else self.__class__.default_capture_stream
)
if self.capture_stream is None:
raise AssertionError("capture_stream must not be None")
self.stream_ctx = self.capture_stream
self.xpu_graph = xpu_graph
def __enter__(self) -> None:
# Free as much memory as we can for the graph
torch.xpu.synchronize()
torch.xpu.empty_cache()
self.stream_ctx.__enter__()
self.xpu_graph.capture_begin(*self.pool)
def __exit__(self, *args: object) -> None:
self.xpu_graph.capture_end()
self.stream_ctx.__exit__(*args)
_ModuleOrCallable: TypeAlias = torch.nn.Module | Callable[..., object]
@overload
def make_graphed_callables(
callables: _ModuleOrCallable,
sample_args: tuple[Tensor, ...],
num_warmup_iters: int = 3,
allow_unused_input: bool = False,
pool: _POOL_HANDLE | None = None,
) -> _ModuleOrCallable: ...
@overload
def make_graphed_callables(
callables: tuple[_ModuleOrCallable, ...],
sample_args: tuple[tuple[Tensor, ...], ...],
num_warmup_iters: int = 3,
allow_unused_input: bool = False,
pool: _POOL_HANDLE | None = None,
) -> tuple[_ModuleOrCallable, ...]: ...
def make_graphed_callables(
callables: _ModuleOrCallable | tuple[_ModuleOrCallable, ...],
sample_args: tuple[Tensor, ...] | tuple[tuple[Tensor, ...], ...],
num_warmup_iters: int = 3,
allow_unused_input: bool = False,
pool: _POOL_HANDLE | None = None,
) -> _ModuleOrCallable | tuple[_ModuleOrCallable, ...]:
r"""Accept callables (functions or :class:`nn.Module<torch.nn.Module>`\ s) and returns graphed versions.
Each graphed callable's forward pass runs its source callable's
forward XPU work as a XPU graph inside a single autograd node.
The graphed callable's forward pass also appends
a backward node to the autograd graph. During backward, this node runs the
callable's backward work as a XPU graph.
Therefore, each graphed callable should be a drop-in replacement for its source callable
in an autograd-enabled training loop.
See :ref:`Partial-network capture<partial-network-capture>` for detailed use and constraints.
If you pass a tuple of several callables, their captures will use the same memory pool.
Arguments:
callables (torch.nn.Module or Python function, or tuple of these): Callable or callables to graph.
If you pass a tuple of callables, their order in the tuple must be the same order they'll run
in the live workload.
sample_args (tuple of Tensors, or tuple of tuples of Tensors): Samples args for each callable.
If a single callable was passed, ``sample_args`` must be a single tuple of argument Tensors.
If a tuple of callables was passed, ``sample_args`` must be tuple of tuples of argument Tensors.
num_warmup_iters (int): The number of warmup iterations. Currently, ``DataDistributedParallel`` needs
11 iterations for warm up. Default: ``3``.
allow_unused_input (bool): If False, specifying inputs that were not used when computing outputs
(and therefore their grad is always zero) is an error. Defaults to False.
pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) that hints this graph may share memory
with the indicated pool.
.. note::
The ``requires_grad`` state of each Tensor in ``sample_args`` must match the state
that's expected for the corresponding real input in the training loop.
.. warning::
This API is in beta and may change in future releases.
.. warning::
``sample_args`` for each callable must contain only Tensors. Other types are not allowed.
.. warning::
Returned callables do not support higher order differentiation (e.g., double backward).
.. warning::
In any :class:`~torch.nn.Module` passed to :func:`~make_graphed_callables`, only parameters
may be trainable. Buffers must have ``requires_grad=False``.
.. warning::
After you pass a :class:`torch.nn.Module` through :func:`~make_graphed_callables`,
you may not add or remove any of that Module's parameters or buffers.
.. warning::
:class:`torch.nn.Module`\s passed to :func:`~torch.xpu.make_graphed_callables` must not have module hooks
registered on them at the time they are passed. However, registering hooks on modules *after* passing them
through :func:`~torch.xpu.make_graphed_callables` is allowed.
.. warning::
When running a graphed callable, you must pass its arguments in the same order and format
they appeared in that callable's ``sample_args``.
.. warning::
The automatic mixed precision is supported in :func:`~torch.xpu.make_graphed_callables` only with disabled
caching. The context manager `torch.amp.autocast()` must have `cache_enabled=False`.
"""
if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled():
raise RuntimeError(
"make_graphed_callables does not support the autocast caching. Please set `cache_enabled=False`."
)
just_one_callable = False
_sample_args: tuple[tuple[Tensor, ...], ...]
if not isinstance(callables, tuple):
just_one_callable = True
callables = (callables,)
_sample_args = (typing.cast(tuple[Tensor, ...], sample_args),)
else:
_sample_args = typing.cast(tuple[tuple[Tensor, ...], ...], sample_args)
flatten_sample_args = []
for c, args in zip(callables, _sample_args):
if isinstance(c, torch.nn.Module):
if not (
len(c._backward_hooks) == 0
and len(c._forward_hooks) == 0
and len(c._forward_pre_hooks) == 0
):
raise RuntimeError(
"Modules must not have hooks registered at the time they are passed. However, registering hooks "
+ "on modules after passing them through make_graphed_callables is allowed."
)
if not all(b.requires_grad is False for b in c.buffers()):
raise RuntimeError(
"In any :class:`~torch.nn.Module` passed to "
+ ":func:`~make_graphed_callables`, only parameters may be trainable. All buffers must have "
+ "``requires_grad=False``."
)
flatten_arg = torch.utils._pytree.arg_tree_leaves(*args)
flatten_sample_args.append(tuple(flatten_arg))
if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg):
raise TypeError(
"In the beta API, sample_args "
+ "for each callable must contain only Tensors. Other types are not allowed."
)
# If a callable is an nn.Module, its graph's full input surface is the args the user explicitly
# passes to forward (ie, its sample_args) AND the module's parameter attributes.
per_callable_len_user_args = [len(args) for args in flatten_sample_args]
per_callable_module_params = [
tuple(c.parameters()) if isinstance(c, torch.nn.Module) else ()
for c in callables
]
per_callable_static_input_surfaces = [
flatten_sample_args[i] + per_callable_module_params[i]
for i in range(len(callables))
]
fwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))]
bwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))]
mempool = graph_pool_handle() if pool is None else pool
# Warmup
torch.xpu.synchronize()
with torch.xpu.stream(torch.xpu.Stream()):
for func, args, static_input_surface in zip(
callables, _sample_args, per_callable_static_input_surfaces
):
grad_inputs, outputs, outputs_grad = None, None, None
for _ in range(num_warmup_iters):
outputs = torch.utils._pytree.tree_leaves(func(*args))
outputs_grad = tuple(o for o in outputs if o.requires_grad)
if len(outputs_grad) > 0:
grad_inputs = torch.autograd.grad(
outputs=outputs_grad,
inputs=tuple(
i for i in static_input_surface if i.requires_grad
),
grad_outputs=tuple(
torch.empty_like(o) for o in outputs if o.requires_grad
),
only_inputs=True,
allow_unused=allow_unused_input,
)
for v in [outputs, outputs_grad, grad_inputs]:
del v
torch.xpu.synchronize()
# Capture forward graphs
per_callable_static_outputs = []
per_callable_output_unflatten_spec = []
for func, args, fwd_graph in zip(callables, _sample_args, fwd_graphs):
# each graph uses the same mempool
with torch.xpu.graph(fwd_graph, pool=mempool):
func_outputs = func(*args)
flatten_outputs, spec = torch.utils._pytree.tree_flatten(func_outputs)
per_callable_static_outputs.append(tuple(flatten_outputs))
per_callable_output_unflatten_spec.append(spec)
# Capture backward graphs in reverse order
per_callable_static_grad_outputs = []
per_callable_static_grad_inputs = []
for static_input_surface, static_outputs, bwd_graph in zip(
reversed(per_callable_static_input_surfaces),
reversed(per_callable_static_outputs),
reversed(bwd_graphs),
):
static_grad_outputs = tuple(
torch.empty_like(o) if o.requires_grad else None for o in static_outputs
)
outputs_grad = tuple(o for o in static_outputs if o.requires_grad)
grad_inputs = None
if len(outputs_grad) > 0:
with torch.xpu.graph(bwd_graph, pool=mempool):
grad_inputs = torch.autograd.grad(
outputs=outputs_grad,
inputs=tuple(i for i in static_input_surface if i.requires_grad),
grad_outputs=tuple(o for o in static_grad_outputs if o is not None),
only_inputs=True,
allow_unused=allow_unused_input,
)
static_grad_inputs = []
grad_idx = 0
for arg in static_input_surface:
if arg.requires_grad and grad_inputs is not None:
static_grad_inputs.append(grad_inputs[grad_idx])
grad_idx += 1
else:
static_grad_inputs.append(None) # type: ignore[arg-type]
static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment]
per_callable_static_grad_outputs.append(static_grad_outputs)
per_callable_static_grad_inputs.append(static_grad_inputs)
# Reverses the most recent two lists
per_callable_static_grad_outputs.reverse()
per_callable_static_grad_inputs.reverse()
def make_graphed_autograd_function(
fwd_graph: XPUGraph,
bwd_graph: XPUGraph,
module_params: tuple[torch.nn.Parameter, ...],
len_user_args: int,
output_unflatten_spec: torch.utils._pytree.TreeSpec,
static_input_surface: tuple[Tensor, ...],
static_outputs: tuple[Tensor, ...],
static_grad_outputs: tuple[Tensor | None, ...],
static_grad_inputs: tuple[Tensor, ...],
) -> Callable[..., object]:
class Graphed(torch.autograd.Function):
@staticmethod
# pyrefly: ignore [bad-override]
def forward(ctx: object, *inputs: Tensor) -> tuple[Tensor, ...]:
# At this stage, only the user args may (potentially) be new tensors.
for i in range(len_user_args):
if static_input_surface[i].data_ptr() != inputs[i].data_ptr():
static_input_surface[i].copy_(inputs[i])
fwd_graph.replay()
if not isinstance(static_outputs, tuple):
raise RuntimeError("static_outputs must be a tuple")
return tuple(o.detach() for o in static_outputs)
@staticmethod
@torch.autograd.function.once_differentiable
# pyrefly: ignore [bad-override]
def backward(ctx: object, *grads: Tensor) -> tuple[Tensor, ...]:
if len(grads) != len(static_grad_outputs):
raise RuntimeError(
f"Expected {len(static_grad_outputs)} gradients but got {len(grads)}"
)
for g, grad in zip(static_grad_outputs, grads):
if g is not None:
if g.data_ptr() != grad.data_ptr():
g.copy_(grad)
bwd_graph.replay()
if not isinstance(static_grad_inputs, tuple):
raise RuntimeError("static_grad_inputs must be a tuple")
return tuple(
b.detach() if b is not None else b for b in static_grad_inputs
)
def functionalized(*user_args: object) -> object:
# Runs the new autograd function which replays the XPU graphs
flatten_user_args = torch.utils._pytree.arg_tree_leaves(*user_args)
out = Graphed.apply(*(tuple(flatten_user_args) + module_params))
return torch.utils._pytree.tree_unflatten(out, output_unflatten_spec)
return functionalized
ret: list[_ModuleOrCallable] = []
for i, func in enumerate(callables):
graphed = make_graphed_autograd_function(
fwd_graphs[i],
bwd_graphs[i],
per_callable_module_params[i],
per_callable_len_user_args[i],
per_callable_output_unflatten_spec[i],
per_callable_static_input_surfaces[i],
per_callable_static_outputs[i],
per_callable_static_grad_outputs[i],
per_callable_static_grad_inputs[i],
)
if isinstance(func, torch.nn.Module):
def make_graphed_forward(
func: torch.nn.Module,
graph_training_state: bool,
graphed: Callable[_P, _R],
orig_fwd: Callable[_P, _R],
) -> Callable[_P, _R]:
def new_fwd(*user_args: _P.args, **user_kwargs: _P.kwargs) -> _R:
if func.training == graph_training_state:
return graphed(*user_args, **user_kwargs)
else:
return orig_fwd(*user_args, **user_kwargs)
return new_fwd
func.forward = make_graphed_forward(
func, func.training, graphed, func.forward
)
ret.append(func)
else:
ret.append(graphed)
if just_one_callable:
return ret[0]
return tuple(ret)
kronos-venv/lib/python3.12/site-packages/torch/xpu/memory.py
+643
View File
@@ -0,0 +1,643 @@
import collections
import contextlib
import ctypes
import pickle
import sys
from typing import Any, Literal
import torch
from torch._utils import _augment_memory_snapshot_stack_traces, _dummy_type
from torch.types import Device
from . import _get_device_index, _is_compiled, _lazy_init, is_initialized
if not _is_compiled():
# Define dummy base classes
torch._C.__dict__["_xpu_XPUAllocator"] = _dummy_type("_xpu_XPUAllocator")
torch._C.__dict__["_XPUMemPool"] = _dummy_type("_XPUMemPool")
torch._C.__dict__["_xpu_beginAllocateCurrentThreadToPool"] = _dummy_type(
"_xpu_beginAllocateCurrentThreadToPool"
)
torch._C.__dict__["_xpu_endAllocateToPool"] = _dummy_type("_xpu_endAllocateToPool")
torch._C.__dict__["_xpu_releasePool"] = _dummy_type("_xpu_releasePool")
def empty_cache() -> None:
r"""Release all unoccupied cached memory currently held by the caching
allocator so that those can be used in other XPU application.
.. note::
:func:`~torch.xpu.empty_cache` doesn't increase the amount of XPU
memory available for PyTorch. However, it may help reduce fragmentation
of XPU memory in certain cases.
"""
if is_initialized():
torch._C._xpu_emptyCache()
def reset_peak_memory_stats(device: Device = None) -> None:
r"""Reset the "peak" stats tracked by the XPU memory allocator.
See :func:`~torch.xpu.memory_stats` for details. Peak stats correspond to the
`"peak"` key in each individual stat dict.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
device = _get_device_index(device, optional=True)
return torch._C._xpu_resetPeakMemoryStats(device)
def reset_accumulated_memory_stats(device: Device = None) -> None:
r"""Reset the "accumulated" (historical) stats tracked by the XPU memory allocator.
See :func:`~torch.xpu.memory_stats` for details. Accumulated stats correspond to
the `"allocated"` and `"freed"` keys in each individual stat dict.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
device = _get_device_index(device, optional=True)
return torch._C._xpu_resetAccumulatedMemoryStats(device)
def memory_stats_as_nested_dict(device: Device = None) -> dict[str, Any]:
r"""Return the result of :func:`~torch.xpu.memory_stats` as a nested dictionary."""
if not is_initialized():
return {}
device = _get_device_index(device, optional=True)
return torch._C._xpu_memoryStats(device)
def memory_stats(device: Device = None) -> dict[str, Any]:
r"""Return a dictionary of XPU memory allocator statistics for a given device.
The return value of this function is a dictionary of statistics, each of
which is a non-negative integer.
Core statistics:
- ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
amount of allocated memory.
- ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
amount of reserved memory.
- ``"active_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
amount of active memory.
- ``"requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
memory requested by client code, compare this with allocated_bytes to check if
allocation rounding adds too much overhead.
For these core statistics, values are broken down as follows.
Pool type:
- ``all``: combined statistics across all memory pools.
- ``large_pool``: statistics for the large allocation pool (for size >= 1MB allocations).
- ``small_pool``: statistics for the small allocation pool (for size < 1MB allocations).
Metric type:
- ``current``: current value of this metric.
- ``peak``: maximum value of this metric.
- ``allocated``: historical total increase in this metric.
- ``freed``: historical total decrease in this metric.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistics for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
result = []
def _recurse_add_to_result(prefix: str, obj: Any) -> None:
if isinstance(obj, dict):
if len(prefix) > 0:
prefix += "."
for k, v in obj.items():
_recurse_add_to_result(prefix + k, v)
else:
result.append((prefix, obj))
stats = memory_stats_as_nested_dict(device=device)
_recurse_add_to_result("", stats)
result.sort()
return collections.OrderedDict(result)
def memory_allocated(device: Device = None) -> int:
r"""Return the current GPU memory occupied by tensors in bytes for a given device.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
.. note::
This is likely less than the amount shown in `xpu-smi` since some
unused memory can be held by the caching allocator and some context
needs to be created on GPU.
"""
return memory_stats(device=device).get("allocated_bytes.all.current", 0)
def max_memory_allocated(device: Device = None) -> int:
r"""Return the maximum GPU memory occupied by tensors in bytes for a given device.
By default, this returns the peak allocated memory since the beginning of
this program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to
reset the starting point in tracking this metric. For example, these two
functions can measure the peak allocated memory usage of each iteration in a
training loop.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
return memory_stats(device=device).get("allocated_bytes.all.peak", 0)
def memory_reserved(device: Device = None) -> int:
r"""Return the current GPU memory managed by the caching allocator in bytes for a given device.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
return memory_stats(device=device).get("reserved_bytes.all.current", 0)
def max_memory_reserved(device: Device = None) -> int:
r"""Return the maximum GPU memory managed by the caching allocator in bytes for a given device.
By default, this returns the peak cached memory since the beginning of this
program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to reset
the starting point in tracking this metric. For example, these two functions
can measure the peak cached memory amount of each iteration in a training
loop.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
"""
return memory_stats(device=device).get("reserved_bytes.all.peak", 0)
def mem_get_info(device: Device = None) -> tuple[int, int]:
r"""Return the global free and total GPU memory for a given device.
Args:
device (torch.device or int or str, optional): selected device. Returns
statistic for the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
Returns:
tuple[int, int]: a tuple of two integers (free_memory, total_memory) in bytes.
The first value is the free memory on the device (available across all processes and applications),
The second value is the device's total hardware memory capacity.
"""
_lazy_init()
device = _get_device_index(device, optional=True)
return torch._C._xpu_getMemoryInfo(device)
def get_per_process_memory_fraction(device: Device = None) -> float:
r"""
Retrieve the memory fraction currently set for a process on a given XPU device.
This fraction represents the portion of the total device memory that
the caching allocator is allowed to use. The allowed memory is calculated as:
.. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction}
Args:
device (torch.device or int or str, optional): selected device. It uses the current device,
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
Returns:
float: The memory fraction in the range 0.0 to 1.0.
"""
_lazy_init()
device = _get_device_index(device, optional=True)
return torch._C._xpu_getMemoryFraction(device)
def set_per_process_memory_fraction(fraction: float, device: Device = None) -> None:
r"""
Set the memory fraction for a single process on XPU device.
This function limits the amount of memory that the caching allocator can allocate
on the specified XPU device. The allowed memory is computed as:
.. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction}
If the process attempts to allocate more than this allowed memory,
an out-of-memory error will be raised by the allocator.
Arguments:
fraction (float): Range: 0~1. Allowed memory equals total_memory * fraction.
device (torch.device or int or str, optional): selected device. It uses the current device,
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
.. note:: In general, the total available free memory is less than the total capacity.
"""
_lazy_init()
device = _get_device_index(device, optional=True)
if not isinstance(fraction, float):
raise TypeError("Invalid type for fraction argument, must be `float`")
torch._C._xpu_setMemoryFraction(fraction, device)
def memory_snapshot(
mempool_id: tuple[int, int] | None = None,
) -> list[dict[str, Any]]:
r"""
Return a snapshot of the XPU memory allocator state across all devices.
Provides detailed information for each memory segment managed by the allocator
including its size, owning pool, associated stream, call stack traces, and other relevant attributes.
Arguments:
mempool_id (tuple[int, int] or None, optional): The memory pool id. If None, the default memory pool is used.
Returns:
list[dict[str, Any]]: List of memory segments and their attributes.
"""
if not is_initialized():
return []
return torch._C._xpu_memorySnapshot(mempool_id)["segments"]
def _snapshot(device: Device = None, augment_with_fx_traces: bool = False):
"""
Capture a snapshot of the XPU memory state at the time this function is called.
The returned snapshot is a dictionary with the following structure.
.. code-block:: python
class Snapshot(TypedDict):
segments: List[Segment]
device_traces: List[List[TraceEntry]]
class Segment(TypedDict):
# A Segment represents a contiguous memory region returned by the SYCL runtime.
#
# All reserved memory is composed of these segments. Segments are
# cached and reused by the allocator. When allocations are smaller
# than the segment, the segment may be split into multiple Blocks.
#
# Calling :func:`~torch.xpu.memory.empty_cache` releases segments that are entirely inactive.
address: int
total_size: int # total size of segment
stream: int
segment_type: Literal["small", "large"] # 'large' (>1MB)
allocated_size: int # size of memory in use
active_size: int # size of memory in use or in active_awaiting_free state
blocks: List[Block]
class Block(TypedDict):
# A sub-region of a Segment, either currently allocated or cached for reuse.
size: int
requested_size: int # Original requested size (may be smaller than `size`)
address: int
state: Literal[
"active_allocated", # used by a tensor
"active_awaiting_free", # waiting for another stream synchronization, then become free
"inactive", # free for reuse
]
frames: List[Frame] # stack trace from where the allocation occurred
class Frame(TypedDict):
filename: str
line: int
name: str
# Optional fields when `augment_with_fx_traces=True` and the frame
# corresponds to FX-generated code.
fx_node_op: str # FX node operation type (e.g., 'call_function', 'output')
fx_node_name: str # FX node name (e.g., 'linear', 'relu_1')
fx_original_trace: str # Original model source code stack trace
class TraceEntry(TypedDict):
# Trace entries are recorded only when :func:`~torch.xpu.memory._record_memory_history` is enabled.
action: Literal[
"alloc" # memory allocated
"free_requested", # received a call to free memory
"free_completed", # memory reclaimed and reusable
"segment_alloc", # ask SYCL runtime for more memory
"segment_free", # called SYCL runtime to return memory to XPU
"segment_map", # ask SYCL runtime to map memory
"segment_unmap", # called SYCL runtime to unmap memory
"snapshot", # snapshot taken
"oom", # threw an OOM exception
]
addr: int # not present for OOM
frames: List[Frame]
size: int
stream: int
device_free: int # only present for OOM, the amount of free memory reported by the device
Arguments:
device (torch.device or int or str, optional): selected device. It uses the current device,
given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default).
augment_with_fx_traces (bool, optional): If True, augment stack trace frames with FX debug information
that maps generated FX code back to original model source code. This adds the FX-related
fields (fx_node_op, fx_node_name, fx_original_trace) to Frame objects. Default is ``False``.
Returns:
The Snapshot dictionary object
"""
s = torch._C._xpu_memorySnapshot(None)
if augment_with_fx_traces:
s = _augment_memory_snapshot_stack_traces(s) # type: ignore[assignment, arg-type]
return s
def _dump_snapshot(
filename: str = "dump_snapshot.pickle", augment_with_fx_traces: bool = False
) -> None:
"""
Save a pickled version of the `torch.memory._snapshot()` dictionary to a file.
This file can be opened by the interactive snapshot viewer at pytorch.org/memory_viz
Snapshot file sizes scale with `max_entries` and stack trace depth per entry,
with several KB per entry. These can easily be in the GB range for longer running
workflows with large `max_entries`.
Arguments:
filename (str, optional): Name of the file to create. Defaults to "dump_snapshot.pickle".
augment_with_fx_traces (bool, optional): If True, augment the snapshot with FX debug information
before dumping. This maps generated FX code stack traces back to original model
source code. Defaults to ``False``.
"""
s = _snapshot(augment_with_fx_traces=augment_with_fx_traces)
with open(filename, "wb") as f:
pickle.dump(s, f)
def _record_memory_history(
enabled: Literal["state", "all"] | None = "all",
context: Literal["state", "alloc", "all"] | None = "all",
stacks: Literal["python", "all"] = "all",
max_entries: int = sys.maxsize,
clear_history: bool = False,
skip_actions: list[str] | None = None,
) -> None:
"""
Enable recording of stack traces associated with memory allocations, so you can
tell what allocated any piece of memory in :func:`~torch.xpu.memory._snapshot()`.
In addition to keeping stack traces with each current allocation and free,
this will also enable recording of a history of all alloc/free events.
Use :func:`~torch.xpu.memory._snapshot()` to retrieve this information,
and the tools in `_memory_viz.py` to visualize snapshots.
Buffer behavior
---------------
This will store up to `max_entries` instances of `TraceEntry` when enabled.
Python trace collection defaults to `sys.maxsize`, meaning long-running
or indefinitely running jobs should set a reasonable limit to avoid excessive
memory use. Expect each entry to be several KB.
Longer running workflows or those with smaller `max_entries` values will only
store the last accumulated `max_entries` entries, meaning new entries overwrite
older entries, reference to ring buffer behavior.
Latency impact
--------------
The Python trace collection is fast (2us per trace), so you may consider
enabling this on production jobs if you anticipate ever having to debug
memory issues.
C++ trace collection is also fast (~50ns/frame), which for many typical programs
works out to ~2us per trace, but can vary depending on stack depth.
Arguments:
enabled (Literal["state", "all"], optional):
`None`, disable recording memory history.
`"state"`, keep information for currently allocated memory.
`"all"`, additionally keep a history of all alloc/free calls.
Defaults to "all".
context (Literal["state", "alloc", "all"], optional):
`None`, Do not record any tracebacks.
`"state"`, Record tracebacks for currently allocated memory.
`"alloc"`, additionally keep tracebacks for alloc calls.
`"all"`, additionally keep tracebacks for free calls.
Defaults to "all".
stacks (Literal["python", "all"], optional):
`"python"`, include Python, TorchScript, and inductor frames in tracebacks.
`"all"`, additionally include C++ frames.
Defaults to "all".
max_entries (int, optional): Keep a maximum of `max_entries`
alloc/free events in the recorded history recorded.
clear_history (bool, optional): Clear history when enabling, defaults to ``False``.
skip_actions (list[str], optional): List of action types to skip when recording
memory history. This can be used to reduce memory overhead by excluding
certain types of events from being recorded. Valid action types are:
- `"alloc"`: Memory allocation events
- `"free_requested"`: Free requests (memory marked for freeing)
- `"free_completed"`: Completed free operations (memory actually freed)
- `"segment_alloc"`: Segment allocation from SYCL runtime
- `"segment_free"`: Segment freed back to XPU via SYCL runtime
- `"segment_map"`: Segment map events
- `"segment_unmap"`: Segment unmap events
- `"snapshot"`: Memory snapshot generation events
- `"oom"`: Out-of-memory exceptions
For example, to skip recording free_requested events:
`skip_actions=["free_requested"]`
Defaults to ``None`` (record all actions).
"""
torch._C._xpu_recordMemoryHistory(
enabled,
context,
stacks,
max_entries,
clear_history,
skip_actions if skip_actions is not None else [],
)
class _XPUAllocator:
r"""Wrapper over internal XPU memory allocators."""
def __init__(self, allocator: torch._C._xpu_XPUAllocator):
self._allocator = allocator
def allocator(self):
return self._allocator
class XPUPluggableAllocator(_XPUAllocator):
r"""
XPU memory allocator loaded dynamically from a shared library.
This lets users provide custom allocation and free functions implemented
in a separate shared library. The allocator is registered and could become
available for use via :func:`~torch.xpu.memory.change_current_allocator`.
Arguments:
path_to_lib_file (str):
Filesystem path to the shared library file containing the allocation
and free functions.
alloc_fn_name (str):
Name of the allocation function exported from the shared library.
The function must have the signature:
``void* alloc_fn(size_t size, int device, sycl::queue* queue);``
free_fn_name (str):
Name of the free function exported from the shared library.
The function must have the signature:
``void free_fn(void* ptr, size_t size, int device, sycl::queue* queue);``
"""
def __init__(self, path_to_lib_file: str, alloc_fn_name: str, free_fn_name: str):
allocator_lib = ctypes.CDLL(path_to_lib_file)
alloc_fn_ptr = getattr(allocator_lib, alloc_fn_name)
free_fn_ptr = getattr(allocator_lib, free_fn_name)
alloc_fn_addr = ctypes.cast(alloc_fn_ptr, ctypes.c_void_p).value
free_fn_addr = ctypes.cast(free_fn_ptr, ctypes.c_void_p).value
if alloc_fn_addr is None or free_fn_addr is None:
raise RuntimeError(
"Failed to load allocator symbols from the shared library."
)
self._allocator = torch._C._xpu_customAllocator(alloc_fn_addr, free_fn_addr)
def change_current_allocator(allocator: _XPUAllocator) -> None:
r"""Change the currently used memory allocator to be the one provided.
.. note::
If the current allocator has already been used/initialized, this function will error.
Arguments:
allocator (torch.xpu.memory._XPUAllocator): allocator to be set as the active one.
"""
torch._C._xpu_changeCurrentAllocator(allocator.allocator())
def _get_current_allocator() -> _XPUAllocator:
r"""Return the allocator being currently used.
Returns:
_XPUAllocator: the allocator being currently used.
"""
return _XPUAllocator(torch._C._xpu_getAllocator())
class MemPool(torch._C._XPUMemPool):
r"""MemPool represents a pool of memory in a caching allocator. Currently,
it's just the ID of the pool object maintained in the XPUCachingAllocator.
Args:
allocator(torch._C._xpu_XPUAllocator, optional): a
torch._C._xpu_XPUAllocator object that can be used to
define how memory gets allocated in the pool. If :attr:`allocator`
is ``None`` (default), memory allocation follows the default/
current configuration of the XPUCachingAllocator.
use_on_oom(bool): a bool that indicates if this pool can be used
as a last resort if a memory allocation outside of the pool fails due
to Out Of Memory. This is ``False`` by default.
"""
def __init__(
self,
allocator: torch._C._xpu_XPUAllocator | None = None,
use_on_oom: bool = False,
):
super().__init__(allocator, True, use_on_oom)
@property
def id(self) -> tuple[int, int]:
r"""Returns the ID of this pool as a tuple of two ints."""
return super().id
@property
def allocator(self) -> torch._C._xpu_XPUAllocator | None:
r"""Returns the allocator this MemPool routes allocations to."""
return super().allocator
def use_count(self) -> int:
r"""Returns the reference count of this pool."""
return super().use_count()
def snapshot(self):
r"""Return a snapshot of the XPU memory allocator pool state across all
devices.
Interpreting the output of this function requires familiarity with the
memory allocator internals.
"""
snapshot = torch.xpu.memory_snapshot(self.id)
return snapshot
@contextlib.contextmanager
def use_mem_pool(pool: MemPool, device: "Device" = None):
r"""A context manager that routes allocations to a given pool.
Args:
pool(torch.xpu.MemPool): a :class:`MemPool` object to be made active so that
allocations route to this pool.
device (torch.device or int, optional): selected device. Uses :class:`MemPool on
the current device, given by :func:`~torch.xpu.current_device`,
if :attr:`device` is ``None`` (default).
.. note::
This context manager makes only current thread's allocations route to
the given pool. If a new thread is spawned inside the context manager
(e.g. by calling backward) the allocations in that thread will not
route to the given pool.
"""
device_index = (
torch.xpu.current_device() if device is None else _get_device_index(device)
)
torch._C._xpu_beginAllocateCurrentThreadToPool(device_index, pool.id)
try:
yield
finally:
torch._C._xpu_endAllocateToPool(device_index, pool.id)
torch._C._xpu_releasePool(device_index, pool.id)
__all__ = [
"MemPool",
"XPUPluggableAllocator",
"change_current_allocator",
"empty_cache",
"get_per_process_memory_fraction",
"max_memory_allocated",
"max_memory_reserved",
"mem_get_info",
"memory_allocated",
"memory_reserved",
"memory_snapshot",
"memory_stats",
"memory_stats_as_nested_dict",
"reset_accumulated_memory_stats",
"reset_peak_memory_stats",
"set_per_process_memory_fraction",
"use_mem_pool",
]
kronos-venv/lib/python3.12/site-packages/torch/xpu/random.py
+176
View File
@@ -0,0 +1,176 @@
# mypy: allow-untyped-defs
from collections.abc import Iterable
import torch
from torch import Tensor
from . import _lazy_call, _lazy_init, current_device, device_count, is_initialized
def get_rng_state(device: int | str | torch.device = "xpu") -> Tensor:
r"""Return the random number generator state of the specified GPU as a ByteTensor.
Args:
device (torch.device or int, optional): The device to return the RNG state of.
Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device).
.. warning::
This function eagerly initializes XPU.
"""
_lazy_init()
if isinstance(device, str):
device = torch.device(device)
elif isinstance(device, int):
device = torch.device("xpu", device)
idx = device.index
if idx is None:
idx = current_device()
default_generator = torch.xpu.default_generators[idx]
return default_generator.get_state()
def get_rng_state_all() -> list[Tensor]:
r"""Return a list of ByteTensor representing the random number states of all devices."""
results = [get_rng_state(i) for i in range(device_count())]
return results
def set_rng_state(new_state: Tensor, device: int | str | torch.device = "xpu") -> None:
r"""Set the random number generator state of the specified GPU.
Args:
new_state (torch.ByteTensor): The desired state
device (torch.device or int, optional): The device to set the RNG state.
Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device).
"""
if not is_initialized():
with torch._C._DisableFuncTorch():
new_state = new_state.clone(memory_format=torch.contiguous_format)
if isinstance(device, str):
device = torch.device(device)
elif isinstance(device, int):
device = torch.device("xpu", device)
def cb() -> None:
idx = device.index
if idx is None:
idx = current_device()
default_generator = torch.xpu.default_generators[idx]
default_generator.set_state(new_state)
_lazy_call(cb)
def set_rng_state_all(new_states: Iterable[Tensor]) -> None:
r"""Set the random number generator state of all devices.
Args:
new_states (Iterable of torch.ByteTensor): The desired state for each device.
"""
for i, state in enumerate(new_states):
set_rng_state(state, i)
def manual_seed(seed: int) -> None:
r"""Set the seed for generating random numbers for the current GPU.
It's safe to call this function if XPU is not available; in that case, it is silently ignored.
Args:
seed (int): The desired seed.
.. warning::
If you are working with a multi-GPU model, this function is insufficient
to get determinism. To seed all GPUs, use :func:`manual_seed_all`.
"""
seed = int(seed)
def cb() -> None:
idx = current_device()
default_generator = torch.xpu.default_generators[idx]
default_generator.manual_seed(seed)
_lazy_call(cb, seed=True)
def manual_seed_all(seed: int) -> None:
r"""Set the seed for generating random numbers on all GPUs.
It's safe to call this function if XPU is not available; in that case, it is silently ignored.
Args:
seed (int): The desired seed.
"""
seed = int(seed)
def cb() -> None:
for i in range(device_count()):
default_generator = torch.xpu.default_generators[i]
default_generator.manual_seed(seed)
_lazy_call(cb, seed_all=True)
def seed() -> None:
r"""Set the seed for generating random numbers to a random number for the current GPU.
It's safe to call this function if XPU is not available; in that case, it is silently ignored.
.. warning::
If you are working with a multi-GPU model, this function will only initialize
the seed on one GPU. To initialize all GPUs, use :func:`seed_all`.
"""
def cb() -> None:
idx = current_device()
default_generator = torch.xpu.default_generators[idx]
default_generator.seed()
_lazy_call(cb)
def seed_all() -> None:
r"""Set the seed for generating random numbers to a random number on all GPUs.
It's safe to call this function if XPU is not available; in that case, it is silently ignored.
"""
def cb() -> None:
random_seed = 0
seeded = False
for i in range(device_count()):
default_generator = torch.xpu.default_generators[i]
if not seeded:
default_generator.seed()
random_seed = default_generator.initial_seed()
seeded = True
else:
default_generator.manual_seed(random_seed)
_lazy_call(cb)
def initial_seed() -> int:
r"""Return the current random seed of the current GPU.
.. warning::
This function eagerly initializes XPU.
"""
_lazy_init()
idx = current_device()
default_generator = torch.xpu.default_generators[idx]
return default_generator.initial_seed()
__all__ = [
"get_rng_state",
"get_rng_state_all",
"set_rng_state",
"set_rng_state_all",
"manual_seed",
"manual_seed_all",
"seed",
"seed_all",
"initial_seed",
]
kronos-venv/lib/python3.12/site-packages/torch/xpu/streams.py
+181
View File
@@ -0,0 +1,181 @@
# mypy: allow-untyped-defs
# pylint: disable=useless-parent-delegation
from __future__ import annotations
import ctypes
import torch
from torch._utils import _dummy_type
if not hasattr(torch._C, "_XpuStreamBase"):
# Define dummy base classes
torch._C.__dict__["_XpuStreamBase"] = _dummy_type("_XpuStreamBase")
torch._C.__dict__["_XpuEventBase"] = _dummy_type("_XpuEventBase")
class Stream(torch._C._XpuStreamBase):
r"""Wrapper around a XPU stream.
A XPU stream is a linear sequence of execution that belongs to a specific
device, independent from other streams. It supports with statement as a
context manager to ensure the operators within the with block are running
on the corresponding stream.
Args:
device(torch.device or int, optional): a device on which to allocate
the stream. If :attr:`device` is ``None`` (default) or a negative
integer, this will use the current device.
priority(int, optional): priority of the stream, which can be positive, 0, or negative.
A lower number indicates a higher priority. By default, the priority is set to 0.
If the value falls outside of the allowed priority range, it will automatically be
mapped to the nearest valid priority (lowest for large positive numbers or
highest for large negative numbers).
"""
def __new__(cls, device=None, priority=0, **kwargs):
# setting device manager is expensive, so we avoid it unless necessary
if device is None or ("stream_id" in kwargs and "device_index" in kwargs):
return super().__new__(cls, priority=priority, **kwargs)
else:
with torch.xpu.device(device):
return super().__new__(cls, priority=priority, **kwargs)
def wait_event(self, event: Event | torch.Event) -> None:
r"""Make all future work submitted to the stream wait for an event.
Args:
event (Event, torch.Event): an event to wait for.
"""
event.wait(self)
def wait_stream(self, stream: Stream | torch.Stream) -> None:
r"""Synchronize 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, torch.Stream): a stream to synchronize.
"""
self.wait_event(stream.record_event())
def record_event(self, event: Event | torch.Event | None = None):
r"""Record an event.
Args:
event (Event, torch.Event, optional): event to record. If not given, a new one
will be allocated.
Returns:
Recorded event.
"""
if event is None:
event = Event()
event.record(self)
return event
def query(self) -> bool:
r"""Check if all the work submitted has been completed.
Returns:
A boolean indicating if all kernels in this stream are completed.
"""
return super().query()
def synchronize(self) -> None:
r"""Wait for all the kernels in this stream to complete."""
super().synchronize()
@property
def _as_parameter_(self):
return ctypes.c_void_p(self.sycl_queue)
def __eq__(self, o):
if isinstance(o, Stream):
return super().__eq__(o)
return False
def __hash__(self):
return hash((self.sycl_queue, self.device))
def __repr__(self) -> str:
return f"torch.xpu.Stream(device={self.device} sycl_queue={self.sycl_queue:#x})"
class Event(torch._C._XpuEventBase):
r"""Wrapper around a XPU event.
XPU events are synchronization markers that can be used to monitor the
device's progress, and to synchronize XPU streams.
The underlying XPU events are lazily initialized when the event is first
recorded. After creation, only streams on the same device may record the
event. However, streams on any device can wait on the event.
Args:
enable_timing (bool, optional): indicates if the event should measure time
(default: ``False``)
"""
def __new__(cls, enable_timing=False):
return super().__new__(cls, enable_timing=enable_timing)
def record(self, stream: Stream | torch.Stream | None = None) -> None:
r"""Record the event in a given stream.
Args:
stream (Stream, torch.Stream, optional): Uses ``torch.xpu.current_stream()`` if no stream is specified.
The stream's device must match the event's device.
"""
if stream is None:
stream = torch.xpu.current_stream()
super().record(stream)
def wait(self, stream: Stream | torch.Stream | None = None) -> None:
r"""Make all future work submitted to the given stream wait for this event.
Args:
stream (Stream, torch.Stream, optional): Uses ``torch.xpu.current_stream()`` if no stream is specified.
"""
if stream is None:
stream = torch.xpu.current_stream()
super().wait(stream)
def query(self) -> bool:
r"""Check if all work currently captured by event has completed.
Returns:
A boolean indicating if all work currently captured by event has
completed.
"""
return super().query()
def elapsed_time(self, end_event: Event):
r"""Return the time elapsed.
Time reported in milliseconds after the event was recorded and
before the end_event was recorded.
Args:
end_event (Event): the end event.
"""
return super().elapsed_time(end_event)
def synchronize(self) -> None:
r"""Wait 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.
"""
super().synchronize()
@property
def _as_parameter_(self):
return ctypes.c_void_p(self.sycl_event)
def __repr__(self) -> str:
if self.sycl_event:
return f"torch.xpu.Event(sycl_event={self.sycl_event:#x})"
else:
return "torch.xpu.Event(uninitialized)"