# mypy: allow-untyped-defsr"""This package enables an interface for accessing MTIA backend in python"""importthreadingimportwarningsfromtypingimportAny,Callable,Dict,List,Optional,Tuple,Unionimporttorchfromtorchimportdeviceas_device,Tensorfromtorch._utilsimport_dummy_type,_LazySeedTracker,classpropertyfromtorch.typesimportDevicefrom._utilsimport_get_device_index_device_t=Union[_device,str,int]# torch.mtia.Event/Stream is alias of torch.Event/StreamEvent=torch.EventStream=torch.Stream_initialized=False_queued_calls:List[Tuple[Callable[[],None],List[str]]]=[]# don't invoke these until initialization occurs_tls=threading.local()_initialization_lock=threading.Lock()_lazy_seed_tracker=_LazySeedTracker()
[docs]defis_initialized():r"""Return whether PyTorch's MTIA state has been initialized."""return_initializedandnot_is_in_bad_fork()
def_is_in_bad_fork()->bool:returntorch._C._mtia_isInBadFork()def_lazy_init()->None:global_initialized,_queued_callsifis_initialized()orhasattr(_tls,"is_initializing"):returnwith_initialization_lock:# We be double-checking locking, boys! This is OK because# the above test was GIL protected anyway. The inner test# is for when a thread blocked on some other thread which was# doing the initialization; when they get the lock, they will# find there is nothing left to do.ifis_initialized():return# It is important to prevent other threads from entering _lazy_init# immediately, while we are still guaranteed to have the GIL, because some# of the C calls we make below will release the GILif_is_in_bad_fork():raiseRuntimeError("Cannot re-initialize MTIA in forked subprocess. To use MTIA with ""multiprocessing, you must use the 'spawn' start method")ifnot_is_compiled():raiseAssertionError("Torch not compiled with MTIA enabled. ""Ensure you have `import mtia.host_runtime.torch_mtia` in your python ""src file and include `//mtia/host_runtime/torch_mtia:torch_mtia` as ""your target dependency!")torch._C._mtia_init()# Some of the queued calls may reentrantly call _lazy_init();# we need to just return without initializing in that case.# However, we must not let any *other* threads in!_tls.is_initializing=True_queued_calls.extend(callsforcallsin_lazy_seed_tracker.get_calls()ifcalls)try:forqueued_call,orig_tracebackin_queued_calls:try:queued_call()exceptExceptionase:msg=(f"MTIA call failed lazily at initialization with error: {str(e)}\n\n"f"MTIA call was originally invoked at:\n\n{''.join(orig_traceback)}")raiseDeferredMtiaCallError(msg)fromefinally:delattr(_tls,"is_initializing")_initialized=True
def_is_compiled()->bool:r"""Return true if compiled with MTIA support."""returntorch._C._mtia_isBuilt()
[docs]defis_available()->bool:r"""Return true if MTIA device is available"""ifnot_is_compiled():returnFalse# MTIA has to init devices first to know if there is any devices available.returndevice_count()>0
[docs]defsynchronize(device:Optional[_device_t]=None)->None:r"""Waits for all jobs in all streams on a MTIA device to complete."""withtorch.mtia.device(device):returntorch._C._mtia_deviceSynchronize()
[docs]defdevice_count()->int:r"""Return the number of MTIA devices available."""returntorch._C._accelerator_hooks_device_count()
[docs]defcurrent_device()->int:r"""Return the index of a currently selected device."""returntorch._C._accelerator_hooks_get_current_device()
[docs]defcurrent_stream(device:Optional[_device_t]=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.mtia.current_device`, if :attr:`device` is ``None`` (default). """returntorch._C._mtia_getCurrentStream(_get_device_index(device,optional=True))
[docs]defdefault_stream(device:Optional[_device_t]=None)->Stream:r"""Return the default :class:`Stream` for a given device. Args: device (torch.device or int, optional): selected device. Returns the default :class:`Stream` for the current device, given by :func:`~torch.mtia.current_device`, if :attr:`device` is ``None`` (default). """returntorch._C._mtia_getDefaultStream(_get_device_index(device,optional=True))
[docs]defget_device_capability(device:Optional[_device_t]=None)->Tuple[int,int]:r"""Return capability of a given device as a tuple of (major version, minor version). Args: device (torch.device or int, optional) selected device. Returns statistics for the current device, given by current_device(), if device is None (default). """returntorch._C._mtia_getDeviceCapability(_get_device_index(device,optional=True))
[docs]defempty_cache()->None:r"""Empty the MTIA device cache."""returntorch._C._mtia_emptyCache()
[docs]defset_stream(stream:Stream):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``. """ifstreamisNone:returntorch._C._mtia_setCurrentStream(stream)
[docs]defset_device(device:_device_t)->None:r"""Set the current device. Args: device (torch.device or int): selected device. This function is a no-op if this argument is negative. """device=_get_device_index(device)ifdevice>=0:torch._C._accelerator_hooks_set_current_device(device)
[docs]classdevice:r"""Context-manager that changes the selected device. Args: device (torch.device or int): device index to select. It's a no-op if this argument is a negative integer or ``None``. """def__init__(self,device:Any):self.idx=_get_device_index(device,optional=True)self.prev_idx=-1def__enter__(self):self.prev_idx=torch._C._accelerator_hooks_maybe_exchange_device(self.idx)def__exit__(self,type:Any,value:Any,traceback:Any):self.idx=torch._C._accelerator_hooks_maybe_exchange_device(self.prev_idx)returnFalse
[docs]classStreamContext:r"""Context-manager that selects a given stream. All MTIA 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:Optional["torch.mtia.Stream"]def__init__(self,stream:Optional["torch.mtia.Stream"]):self.cur_stream=Noneself.stream=streamself.idx=_get_device_index(None,True)ifnottorch.jit.is_scripting():ifself.idxisNone:self.idx=-1self.src_prev_stream=(Noneifnottorch.jit.is_scripting()elsetorch.mtia.default_stream(None))self.dst_prev_stream=(Noneifnottorch.jit.is_scripting()elsetorch.mtia.default_stream(None))def__enter__(self):# Local cur_stream variable for type refinementcur_stream=self.stream# Return if stream is None or MTIA device not availableifcur_streamisNoneorself.idx==-1:returnself.src_prev_stream=torch.mtia.current_stream(None)# If the stream is not on the current device, then# set the current stream on the deviceifself.src_prev_stream.device!=cur_stream.device:withdevice(cur_stream.device):self.dst_prev_stream=torch.mtia.current_stream(cur_stream.device)torch.mtia.set_stream(cur_stream)def__exit__(self,type:Any,value:Any,traceback:Any):# Local cur_stream variable for type refinementcur_stream=self.stream# If stream is None or no MTIA device available, returnifcur_streamisNoneorself.idx==-1:return# Reset the stream on the original device# and destination deviceifself.src_prev_stream.device!=cur_stream.device:# type: ignore[union-attr]torch.mtia.set_stream(self.dst_prev_stream)# type: ignore[arg-type]torch.mtia.set_stream(self.src_prev_stream)# type: ignore[arg-type]
[docs]defstream(stream:Optional["torch.mtia.Stream"])->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``. ..Note:: In eager mode stream is of type Stream class while in JIT it doesn't support torch.mtia.stream """returnStreamContext(stream)
[docs]defget_rng_state(device:Union[int,str,torch.device]="mtia")->Tensor:r"""Returns the random number generator state as a ByteTensor. Args: device (torch.device or int, optional): The device to return the RNG state of. Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device). """warnings.warn("get_rng_state is not implemented in torch.mtia",UserWarning,stacklevel=2,)returntorch.zeros([1],dtype=torch.uint8,device=device)
[docs]defset_rng_state(new_state:Tensor,device:Union[int,str,torch.device]="mtia")->None:r"""Sets the random number generator state. Args: new_state (torch.ByteTensor): The desired state device (torch.device or int, optional): The device to set the RNG state. Default: ``'mtia'`` (i.e., ``torch.device('mtia')``, the current mtia device). """warnings.warn("set_rng_state is not implemented in torch.mtia",UserWarning,stacklevel=2,)
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.